278 lines
10 KiB
Rust
278 lines
10 KiB
Rust
//! LLM Provider trait — 统一的 LLM 调用抽象
|
||
//!
|
||
//! 支持 OpenAI 兼容 API（覆盖 OpenAI / GLM / DeepSeek / Claude 兼容模式），
|
||
//! 含 function calling / tool use 能力。
|
||
//!
|
||
//! 本文件仅含 trait + 数据结构定义（零 IO）。HTTP impl（OpenAICompatProvider /
|
||
//! AnthropicCompatProvider）+ 业务逻辑（ContextManager / AiToolRegistry /
|
||
//! build_provider 工厂）留在 df-ai crate。
|
||
|
||
use std::pin::Pin;
|
||
|
||
use async_trait::async_trait;
|
||
use futures::Stream;
|
||
use serde::{Deserialize, Serialize};
|
||
|
||
// ============================================================
|
||
// 核心数据结构
|
||
// ============================================================
|
||
|
||
/// LLM 调用请求
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct CompletionRequest {
|
||
/// 模型名称
|
||
pub model: String,
|
||
/// 提示消息列表
|
||
pub messages: Vec<ChatMessage>,
|
||
/// 温度（0.0 ~ 2.0）
|
||
pub temperature: Option<f32>,
|
||
/// 最大生成 token 数
|
||
pub max_tokens: Option<u32>,
|
||
/// 是否流式输出
|
||
pub stream: bool,
|
||
/// 可调用的工具定义
|
||
#[serde(skip_serializing_if = "Option::is_none")]
|
||
pub tools: Option<Vec<ToolDefinition>>,
|
||
/// 工具调用策略: "auto" | "none" | {"type":"function","name":"xxx"}
|
||
#[serde(skip_serializing_if = "Option::is_none")]
|
||
pub tool_choice: Option<serde_json::Value>,
|
||
}
|
||
|
||
/// 聊天消息
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct ChatMessage {
|
||
pub role: MessageRole,
|
||
pub content: String,
|
||
/// 工具调用 ID（role=Tool 时必填）
|
||
#[serde(skip_serializing_if = "Option::is_none")]
|
||
pub tool_call_id: Option<String>,
|
||
/// AI 发起的工具调用列表（role=Assistant 时可能有）
|
||
#[serde(skip_serializing_if = "Option::is_none")]
|
||
pub tool_calls: Option<Vec<ToolCall>>,
|
||
/// 生成该消息的 model（仅 assistant 消息有，消息级 model 追溯）
|
||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||
pub model: Option<String>,
|
||
/// 消息状态（UX-09 编辑重生成）：None/"active" 正常可见；
|
||
/// "truncated" 软删（编辑某条 user 消息后其后续消息标记，保留 DB 可追溯但不进 LLM 上下文、前端视图过滤）
|
||
/// 默认 None（向前兼容老 JSON 反序列化）。落库随 messages JSON 序列化，无需独立列。
|
||
#[serde(default, skip_serializing_if = "Option::is_none")]
|
||
pub status: Option<String>,
|
||
}
|
||
|
||
impl ChatMessage {
|
||
pub fn system(content: impl Into<String>) -> Self {
|
||
Self { role: MessageRole::System, content: content.into(), tool_call_id: None, tool_calls: None, model: None, status: None }
|
||
}
|
||
pub fn user(content: impl Into<String>) -> Self {
|
||
Self { role: MessageRole::User, content: content.into(), tool_call_id: None, tool_calls: None, model: None, status: None }
|
||
}
|
||
pub fn assistant(content: impl Into<String>) -> Self {
|
||
Self { role: MessageRole::Assistant, content: content.into(), tool_call_id: None, tool_calls: None, model: None, status: None }
|
||
}
|
||
pub fn assistant_with_tools(content: impl Into<String>, tool_calls: Vec<ToolCall>) -> Self {
|
||
Self { role: MessageRole::Assistant, content: content.into(), tool_call_id: None, tool_calls: Some(tool_calls), model: None, status: None }
|
||
}
|
||
pub fn tool_result(call_id: impl Into<String>, content: impl Into<String>) -> Self {
|
||
Self { role: MessageRole::Tool, content: content.into(), tool_call_id: Some(call_id.into()), tool_calls: None, model: None, status: None }
|
||
}
|
||
|
||
/// 是否处于 active 态（status 为 None 或 "active"）。其余状态一律 false。
|
||
///
|
||
/// 正面白名单（F-15 §3.2）：仅认 None / "active"，新状态
|
||
/// （如阶段2 引入的 "archived_segment" / "compressed"）自动落入不 active 分支，
|
||
/// 无需每加一个状态就来这里改。当前取值 None/Some("active")/Some("truncated")
|
||
/// 行为与旧反面排除完全等价（None=true / "active"=true / "truncated"=false）。
|
||
pub fn is_active(&self) -> bool {
|
||
matches!(self.status.as_deref(), None | Some("active"))
|
||
}
|
||
}
|
||
|
||
/// 消息角色
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
#[serde(rename_all = "lowercase")]
|
||
pub enum MessageRole {
|
||
System,
|
||
User,
|
||
Assistant,
|
||
Tool,
|
||
}
|
||
|
||
/// 工具定义
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct ToolDefinition {
|
||
#[serde(rename = "type")]
|
||
pub tool_type: String,
|
||
pub function: ToolFunction,
|
||
}
|
||
|
||
impl ToolDefinition {
|
||
pub fn function(name: impl Into<String>, description: impl Into<String>, parameters: serde_json::Value) -> Self {
|
||
Self {
|
||
tool_type: "function".into(),
|
||
function: ToolFunction { name: name.into(), description: description.into(), parameters },
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 函数定义
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct ToolFunction {
|
||
pub name: String,
|
||
pub description: String,
|
||
pub parameters: serde_json::Value,
|
||
}
|
||
|
||
/// 工具调用（AI 发起）
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct ToolCall {
|
||
pub id: String,
|
||
#[serde(rename = "type")]
|
||
pub call_type: String,
|
||
pub function: ToolCallFunction,
|
||
}
|
||
|
||
impl ToolCall {
|
||
pub fn new(id: impl Into<String>, name: impl Into<String>, arguments: impl Into<String>) -> Self {
|
||
Self {
|
||
id: id.into(),
|
||
call_type: "function".into(),
|
||
function: ToolCallFunction { name: name.into(), arguments: arguments.into() },
|
||
}
|
||
}
|
||
}
|
||
|
||
/// 工具调用函数部分
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct ToolCallFunction {
|
||
pub name: String,
|
||
pub arguments: String,
|
||
}
|
||
|
||
/// LLM 调用响应
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct CompletionResponse {
|
||
/// 生成的文本
|
||
pub text: String,
|
||
/// 使用的模型
|
||
pub model: String,
|
||
/// 消耗的 token 数
|
||
pub usage: TokenUsage,
|
||
/// AI 发起的工具调用（如有）
|
||
#[serde(skip_serializing_if = "Option::is_none")]
|
||
pub tool_calls: Option<Vec<ToolCall>>,
|
||
}
|
||
|
||
/// Token 用量
|
||
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
|
||
pub struct TokenUsage {
|
||
pub prompt_tokens: u32,
|
||
pub completion_tokens: u32,
|
||
pub total_tokens: u32,
|
||
}
|
||
|
||
/// 流式输出的 chunk
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct StreamChunk {
|
||
/// 增量文本
|
||
pub delta: String,
|
||
/// 是否结束
|
||
pub finished: bool,
|
||
/// 工具调用增量（如有）
|
||
#[serde(skip_serializing_if = "Option::is_none")]
|
||
pub tool_calls: Option<Vec<ToolCallDelta>>,
|
||
/// Token 用量（流末 chunk 携带，由 provider 解析自 SSE usage 事件）
|
||
#[serde(skip_serializing_if = "Option::is_none")]
|
||
pub usage: Option<TokenUsage>,
|
||
/// provider 流式错误事件（如 Anthropic SSE `type=="error"`）。
|
||
/// 非空表示流中途出错，不应视为正常完成（finished 路径），由 stream_llm 转 AiError。
|
||
#[serde(skip)]
|
||
pub error: Option<String>,
|
||
}
|
||
|
||
/// 工具调用增量（流式中的片段）
|
||
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||
pub struct ToolCallDelta {
|
||
/// 索引
|
||
pub index: u32,
|
||
/// 工具调用 ID（仅第一个 chunk 有）
|
||
#[serde(skip_serializing_if = "Option::is_none")]
|
||
pub id: Option<String>,
|
||
/// 函数名片段
|
||
#[serde(skip_serializing_if = "Option::is_none")]
|
||
pub function_name: Option<String>,
|
||
/// 函数参数片段
|
||
#[serde(skip_serializing_if = "Option::is_none")]
|
||
pub function_arguments: Option<String>,
|
||
}
|
||
|
||
/// 异步流类型别名
|
||
pub type StreamResult = Pin<Box<dyn Stream<Item = anyhow::Result<StreamChunk>> + Send>>;
|
||
|
||
/// LLM Provider trait
|
||
#[async_trait]
|
||
pub trait LlmProvider: Send + Sync {
|
||
/// 同步调用
|
||
async fn complete(&self, request: CompletionRequest) -> anyhow::Result<CompletionResponse>;
|
||
|
||
/// 流式调用（返回异步流）
|
||
async fn stream(
|
||
&self,
|
||
request: CompletionRequest,
|
||
) -> anyhow::Result<StreamResult>;
|
||
|
||
/// 文本嵌入：批量文本 → 语义向量(供知识库向量检索)
|
||
///
|
||
/// 默认实现返回 Err(协议不支持)。OpenAI 兼容协议覆盖实现(/v1/embeddings);
|
||
/// Anthropic 无 embedding API,保持默认。
|
||
async fn embed(&self, _model: &str, _texts: Vec<String>) -> anyhow::Result<Vec<Vec<f32>>> {
|
||
anyhow::bail!("该 Provider 不支持 embedding({})", self.name())
|
||
}
|
||
|
||
/// Provider 名称
|
||
fn name(&self) -> &str;
|
||
|
||
/// 实际请求端点（含 base_url + 关键路径，如 chat completions / messages）。
|
||
/// 默认回落 `name()`，provider 实现覆盖返真实 URL，供 401/网络错误诊断打印
|
||
/// —— 旧路径只能近似打印 provider_type，看不到实际请求端点。
|
||
fn endpoint(&self) -> String {
|
||
self.name().to_string()
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn is_active_whitelist() {
|
||
// F-15 §3.2 正面白名单：仅 None / "active" 为 true，其余一律 false。
|
||
// 零行为变化：None / "active" / "truncated" 与旧反面排除完全等价；
|
||
// "archived_segment" / "compressed"（阶段2 引入，当前代码未赋值）由白名单
|
||
// matches! 只认 None/active 自动落入 false，面向未来验证。
|
||
|
||
// None（构造默认值，向前兼容老 JSON）
|
||
let m = ChatMessage::user("hi");
|
||
assert!(m.is_active(), "None 应 active");
|
||
|
||
// "active"
|
||
let mut m = ChatMessage::user("hi");
|
||
m.status = Some("active".to_string());
|
||
assert!(m.is_active(), "Some(active) 应 active");
|
||
|
||
// "truncated" — 当前取值，与旧实现等价（false）
|
||
let mut m = ChatMessage::user("hi");
|
||
m.status = Some("truncated".to_string());
|
||
assert!(!m.is_active(), "truncated 应不 active");
|
||
|
||
// "archived_segment" — 阶段2 待引入，白名单自动隔离
|
||
let mut m = ChatMessage::user("hi");
|
||
m.status = Some("archived_segment".to_string());
|
||
assert!(!m.is_active(), "archived_segment 应不 active（白名单隔离）");
|
||
|
||
// "compressed" — 阶段2 待引入，白名单自动隔离
|
||
let mut m = ChatMessage::user("hi");
|
||
m.status = Some("compressed".to_string());
|
||
assert!(!m.is_active(), "compressed 应不 active（白名单隔离）");
|
||
}
|
||
}
|