203 lines
6.5 KiB
Rust
203 lines
6.5 KiB
Rust
//! ## LLM 提供商抽象（trait） + SSE 流式解析
|
||
//!
|
||
//! 定义了 LLM 层的核心抽象接口 `LlmProvider`，以及 SSE 流式响应的解析工具。
|
||
//!
|
||
//! ### 设计模式
|
||
//! - `LlmProvider` trait — 所有 LLM 提供商需实现的接口
|
||
//! - `create_provider()` — 工厂函数，根据配置创建对应的提供商实例
|
||
//! - `parse_chat_chunk()` — SSE 流式 delta 解析器
|
||
//!
|
||
//! ### 流式架构
|
||
//! ```text
|
||
//! Provider.chat_stream() → mpsc::Receiver<StreamChunk>
|
||
//! ├── StreamChunk::Text (delta 文本)
|
||
//! ├── StreamChunk::Reasoning (思考内容)
|
||
//! ├── StreamChunk::Done (完成信号 + 工具调用 + 用量)
|
||
//! └── StreamChunk::Error (错误)
|
||
//! ```
|
||
|
||
use crate::llm::types::{ConversationConfig, Message, StreamChunk};
|
||
use async_trait::async_trait;
|
||
use serde::Deserialize;
|
||
use std::sync::Arc;
|
||
use tokio::sync::mpsc;
|
||
|
||
/// 流式返回：每个 StreamChunk 是 delta 或控制信号
|
||
pub type StreamReceiver = mpsc::Receiver<StreamChunk>;
|
||
pub type StreamSender = mpsc::Sender<StreamChunk>;
|
||
|
||
/// ## LLM 提供商抽象（trait）
|
||
///
|
||
/// 这是 LLM 层的核心抽象接口，定义了"一个 LLM 提供商需要提供什么能力"。
|
||
/// 目前唯一的实现是 `DeepSeekProvider`，但通过这个 trait 可以轻松添加
|
||
/// 其他提供商（如 OpenAI、Claude 等）。
|
||
///
|
||
/// ### 关键设计
|
||
/// * 使用 `async_trait` 宏支持异步 trait 方法
|
||
/// * 返回 `StreamReceiver`（mpsc::Receiver），调用方可以按 chunk 消费流式结果
|
||
/// * `Send + Sync` 保证可以在 tokio task 间安全共享
|
||
#[async_trait]
|
||
pub trait LlmProvider: Send + Sync {
|
||
/// 提供商名称
|
||
fn name(&self) -> &str;
|
||
|
||
/// 发起流式对话
|
||
async fn chat_stream(
|
||
&self,
|
||
config: &ConversationConfig,
|
||
messages: &[Message],
|
||
) -> Result<StreamReceiver, String>;
|
||
}
|
||
|
||
// ─── 内置提供商注册 ───
|
||
|
||
pub type BoxedProvider = Arc<dyn LlmProvider>;
|
||
|
||
/// ## `create_provider` — 工厂函数：从配置创建恰当的 LLM 提供商
|
||
///
|
||
/// 目前只返回 DeepSeek，后续可以在这里根据配置 switch 到不同提供商。
|
||
pub fn create_provider(_config: &ConversationConfig) -> Result<BoxedProvider, String> {
|
||
Ok(Arc::new(super::deepseek::DeepSeekProvider::new()?))
|
||
}
|
||
|
||
// ─── 内部：SSE 解析工具 ───
|
||
|
||
/// ## 流式块解析结果（内部枚举）
|
||
///
|
||
/// 用于从 SSE 的 JSON delta 中解析出不同种类的块。
|
||
/// * `Text` — 普通文本 delta
|
||
/// * `Reasoning` — 思考内容（DeepSeek reasoning_content）
|
||
/// * `ToolCallDelta` — 工具调用的增量信息（可能需要跨多个 chunk 拼接）
|
||
/// * `FinishReason` — 结束原因（stop / tool_calls 等）
|
||
/// * `Usage` — 用量统计（通常在最后一个 chunk 中）
|
||
#[allow(dead_code)]
|
||
pub(crate) enum ParsedChunk {
|
||
Text(String),
|
||
Reasoning(String),
|
||
ToolCallDelta {
|
||
index: i32,
|
||
id: Option<String>,
|
||
name: Option<String>,
|
||
arguments: String,
|
||
},
|
||
FinishReason(String),
|
||
Usage(super::types::Usage),
|
||
}
|
||
|
||
/// ## 从 JSON body 中解析 DeepSeek/OpenAI 流式 delta
|
||
///
|
||
/// ### 输入
|
||
/// 一行 SSE 数据，格式为 `data: {json}`。
|
||
///
|
||
/// ### 处理逻辑
|
||
/// 1. 跳过非 `data: ` 前缀的行和 `[DONE]` 行
|
||
/// 2. 解析 JSON 为 `ChunkResponse`（内嵌多个 `choices`）
|
||
/// 3. 提取 usage（可能在任意一个 chunk 中）
|
||
/// 4. 对每个 choice，按优先级提取：
|
||
/// - tool_calls delta → `ToolCallDelta`
|
||
/// - reasoning_content → `Reasoning`
|
||
/// - content → `Text`
|
||
/// - finish_reason → `FinishReason`
|
||
pub(crate) fn parse_chat_chunk(line: &str) -> Option<ParsedChunk> {
|
||
if !line.starts_with("data: ") {
|
||
return None;
|
||
}
|
||
|
||
let data = line["data: ".len()..].trim();
|
||
if data == "[DONE]" {
|
||
return None;
|
||
}
|
||
|
||
#[derive(Deserialize)]
|
||
struct ToolCallDelta {
|
||
#[serde(default)]
|
||
index: Option<i32>,
|
||
#[serde(default)]
|
||
id: Option<String>,
|
||
#[serde(rename = "type", default)]
|
||
#[allow(dead_code)]
|
||
call_type: Option<String>,
|
||
#[serde(default)]
|
||
function: Option<ToolFunctionDelta>,
|
||
}
|
||
|
||
#[derive(Deserialize)]
|
||
struct ToolFunctionDelta {
|
||
#[serde(default)]
|
||
name: Option<String>,
|
||
#[serde(default)]
|
||
arguments: Option<String>,
|
||
}
|
||
|
||
#[derive(Deserialize)]
|
||
struct Delta {
|
||
#[serde(default)]
|
||
content: Option<String>,
|
||
#[serde(default)]
|
||
reasoning_content: Option<String>,
|
||
#[serde(default)]
|
||
tool_calls: Option<Vec<ToolCallDelta>>,
|
||
}
|
||
|
||
#[derive(Deserialize)]
|
||
struct ChunkChoice {
|
||
delta: Delta,
|
||
#[serde(default)]
|
||
finish_reason: Option<String>,
|
||
}
|
||
|
||
#[derive(Deserialize)]
|
||
struct ChunkResponse {
|
||
choices: Vec<ChunkChoice>,
|
||
#[serde(default)]
|
||
usage: Option<super::types::Usage>,
|
||
}
|
||
|
||
let parsed: ChunkResponse = match serde_json::from_str(data) {
|
||
Ok(p) => p,
|
||
Err(_) => return None,
|
||
};
|
||
|
||
// 提取 usage（可能在最后一个 chunk）
|
||
if let Some(ref usage) = parsed.usage
|
||
&& usage.total_tokens > 0 {
|
||
return Some(ParsedChunk::Usage(usage.clone()));
|
||
}
|
||
|
||
// SSE 流式响应中每个 chunk 只包含一个 choice 的 delta
|
||
if let Some(choice) = parsed.choices.into_iter().next() {
|
||
// 工具调用 delta
|
||
if let Some(tool_calls) = &choice.delta.tool_calls
|
||
&& let Some(tc) = tool_calls.first() {
|
||
let idx = tc.index.unwrap_or(0);
|
||
let args = tc
|
||
.function
|
||
.as_ref()
|
||
.and_then(|f| f.arguments.clone())
|
||
.unwrap_or_default();
|
||
let name = tc.function.as_ref().and_then(|f| f.name.clone());
|
||
return Some(ParsedChunk::ToolCallDelta {
|
||
index: idx,
|
||
id: tc.id.clone(),
|
||
name,
|
||
arguments: args,
|
||
});
|
||
}
|
||
|
||
if let Some(reasoning) = &choice.delta.reasoning_content
|
||
&& !reasoning.is_empty() {
|
||
return Some(ParsedChunk::Reasoning(reasoning.clone()));
|
||
}
|
||
if let Some(content) = &choice.delta.content
|
||
&& !content.is_empty() {
|
||
return Some(ParsedChunk::Text(content.clone()));
|
||
}
|
||
if let Some(reason) = &choice.finish_reason
|
||
&& !reason.is_empty() {
|
||
return Some(ParsedChunk::FinishReason(reason.clone()));
|
||
}
|
||
}
|
||
|
||
None
|
||
}
|