lxy/DevFlow
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

新增: Phase2 阶段收尾（Sprint 1-20）

Browse Source 
重构：删 5 零引用 crate（df-evolve/plugin/stages/task/traceability）+ 清死模块、ai.rs 拆 11 子 module、ai.ts 拆 6 composable、i18n 拆目录
功能：知识库全栈（df-project/scan + CRUD + 时间线 + 前端）、Settings 拆分、appSettings KV 迁移、模型池、LLM 并发 Semaphore
修复：审批持久化根治、ConditionEngine 默认拒绝、NodeRegistry unimplemented 清除、promote 补偿删除、工具结果截断 50KB、路径校验防 symlink 逃逸
文档：B-03 人工审批设计、决策记录三分档、规格契约自检、经验记录、todo 看板、PROGRESS 更新

详见 PROGRESS.md。src-tauri/儿童每日打卡应用/ 与本项目无关，已排除。
This commit is contained in:
绝尘
2026-06-14 14:08:20 +08:00
parent 98393b4908
commit cf017f81e2
167 changed files with 19549 additions and 6886 deletions
Download Patch File Download Diff File Expand all files Collapse all files

13
crates/df-ai/src/ai_tools.rs
View File

@@ -96,6 +96,19 @@ impl AiToolRegistry {
self.tools.get(name)
}
/// 执行指定工具 — handler 是唯一执行路径（schema+risk+实现同源，消除双轨）
pub async fn execute(
&self,
name: &str,
args: serde_json::Value,
) -> anyhow::Result<serde_json::Value> {
let tool = self
.tools
.get(name)
.ok_or_else(|| anyhow::anyhow!("未知工具: {}", name))?;
(tool.handler)(args).await
}
/// 获取所有已注册工具名称
pub fn tool_names(&self) -> Vec<String> {
self.tools.keys().cloned().collect()

375
crates/df-ai/src/anthropic_compat.rs
View File

@@ -8,7 +8,7 @@
use async_trait::async_trait;
use eventsource_stream::Eventsource;
use futures::{Stream, StreamExt};
use futures::StreamExt;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use tracing::{debug, error, warn};
@@ -79,6 +79,126 @@ struct AnthropicUsage {
output_tokens: u32,
}
// ============================================================
// SSE 解析纯函数（与 HTTP 解耦，便于单测）
// ============================================================
/// 将一条 Anthropic Messages SSE 事件 data 解析为 StreamChunk，并按需更新 usage 累加器。
///
/// 按 `type` 字段分发：
/// - `message_start` → 用 `message.usage.input_tokens` 初始化累加器（output 置 0）。
/// - `message_delta` → **output_tokens 是累计值（非增量）**，直接覆盖 `completion_tokens` 并重算 `total`。
/// - `content_block_delta` (text_delta/input_json_delta) → 文本/工具入参增量。
/// - `content_block_start` (tool_use) → 工具块开始，带 id+name。
/// - `message_stop` → 返回 `finished=true` 终态 chunk，`usage` 取自累加器（`take()`）。
/// - `error` → 返回 `finished=true` 终态空 chunk。
/// - 其它（content_block_stop / ping 等）→ 空 chunk。
///
/// 等价性：content_block / message_stop / error 等事件分支与原 stream() 闭包逐字一致；
/// usage 透传（message_stop 终态 take() 带出、message_delta 的 output_tokens 按累计值覆盖）
/// 为本次新增能力，对应 StreamChunk 新增的 usage 字段。
pub(crate) fn apply_anthropic_event(data: &str, usage_accum: &mut Option<TokenUsage>) -> StreamChunk {
// 解析 data 中的 JSON，按 type 字段决定如何转 StreamChunk
let v: serde_json::Value = match serde_json::from_str(data) {
Ok(v) => v,
Err(_) => {
return StreamChunk { delta: String::new(), finished: false, tool_calls: None, usage: None }
}
};
let ty = v.get("type").and_then(|t| t.as_str()).unwrap_or("");
match ty {
// 消息开始：取 input_tokens 初始化累积器（output 此时未知，置 0）
"message_start" => {
if let Some(inp) = v
.get("message")
.and_then(|m| m.get("usage"))
.and_then(|u| u.get("input_tokens"))
.and_then(|t| t.as_u64())
{
*usage_accum = Some(TokenUsage {
prompt_tokens: inp as u32,
completion_tokens: 0,
total_tokens: inp as u32,
});
}
StreamChunk { delta: String::new(), finished: false, tool_calls: None, usage: None }
}
// 消息增量：output_tokens 是累计值（非增量），直接覆盖 completion + 重算 total
"message_delta" => {
if let Some(out) = v.get("usage").and_then(|u| u.get("output_tokens")).and_then(|t| t.as_u64()) {
let acc = usage_accum
.get_or_insert(TokenUsage { prompt_tokens: 0, completion_tokens: 0, total_tokens: 0 });
acc.completion_tokens = out as u32;
acc.total_tokens = acc.prompt_tokens + acc.completion_tokens;
}
StreamChunk { delta: String::new(), finished: false, tool_calls: None, usage: None }
}
// 文本增量
"content_block_delta" => {
if let Some(delta) = v.get("delta") {
if delta.get("type").and_then(|t| t.as_str()) == Some("text_delta") {
let text = delta.get("text").and_then(|t| t.as_str()).unwrap_or("").to_string();
return StreamChunk { delta: text, finished: false, tool_calls: None, usage: None };
}
// 工具入参增量
if delta.get("type").and_then(|t| t.as_str()) == Some("input_json_delta") {
let partial = delta.get("partial_json").and_then(|t| t.as_str()).unwrap_or("").to_string();
let idx = v.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as u32;
return StreamChunk {
delta: String::new(),
finished: false,
tool_calls: Some(vec![ToolCallDelta {
index: idx,
id: None,
function_name: None,
function_arguments: Some(partial),
}]),
usage: None,
};
}
}
StreamChunk { delta: String::new(), finished: false, tool_calls: None, usage: None }
}
// 工具块开始：带 id + name
"content_block_start" => {
if let Some(cb) = v.get("content_block") {
if cb.get("type").and_then(|t| t.as_str()) == Some("tool_use") {
let idx = v.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as u32;
let id = cb.get("id").and_then(|t| t.as_str()).map(|s| s.to_string());
let name = cb.get("name").and_then(|t| t.as_str()).map(|s| s.to_string());
return StreamChunk {
delta: String::new(),
finished: false,
tool_calls: Some(vec![ToolCallDelta {
index: idx,
id,
function_name: name,
function_arguments: None,
}]),
usage: None,
};
}
}
StreamChunk { delta: String::new(), finished: false, tool_calls: None, usage: None }
}
// 消息结束：带出累积 usage
"message_stop" => StreamChunk {
delta: String::new(),
finished: true,
tool_calls: None,
usage: usage_accum.take(),
},
// 错误事件
"error" => {
let msg = v.get("error").and_then(|e| e.get("message")).and_then(|m| m.as_str()).unwrap_or("stream error");
error!(%msg, "Anthropic 流式错误事件");
StreamChunk { delta: String::new(), finished: true, tool_calls: None, usage: None }
}
// content_block_stop / ping 等不产出 chunk
_ => StreamChunk { delta: String::new(), finished: false, tool_calls: None, usage: None },
}
}
// ============================================================
// Provider 实现
// ============================================================
@@ -332,77 +452,15 @@ impl LlmProvider for AnthropicCompatProvider {
anyhow::bail!("Anthropic 流式 API 错误 {}: {}", status, text);
}
// 流式解析：eventsource 逐事件处理，按 type 字段分发转 StreamChunk
// 流式解析：eventsource 逐事件处理，按 type 字段分发转 StreamChunk。
// 事件解析/usage 累积逻辑抽到 apply_anthropic_event 纯函数，便于单测；此处闭包只负责传 data。
// usage 累积：message_start 给 input_tokens，message_delta 给累计 output_tokens（非增量），message_stop 带出。
let mut usage_accum: Option<TokenUsage> = None;
let stream = resp
.bytes_stream()
.eventsource()
.map(move |event| match event {
Ok(ev) => {
// 解析 data 中的 JSON，按 type 字段决定如何转 StreamChunk
let v: serde_json::Value = match serde_json::from_str(&ev.data) {
Ok(v) => v,
Err(_) => return Ok(StreamChunk { delta: String::new(), finished: false, tool_calls: None }),
};
let ty = v.get("type").and_then(|t| t.as_str()).unwrap_or("");
match ty {
// 文本增量
"content_block_delta" => {
if let Some(delta) = v.get("delta") {
if delta.get("type").and_then(|t| t.as_str()) == Some("text_delta") {
let text = delta.get("text").and_then(|t| t.as_str()).unwrap_or("").to_string();
return Ok(StreamChunk { delta: text, finished: false, tool_calls: None });
}
// 工具入参增量
if delta.get("type").and_then(|t| t.as_str()) == Some("input_json_delta") {
let partial = delta.get("partial_json").and_then(|t| t.as_str()).unwrap_or("").to_string();
let idx = v.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as u32;
return Ok(StreamChunk {
delta: String::new(),
finished: false,
tool_calls: Some(vec![ToolCallDelta {
index: idx,
id: None,
function_name: None,
function_arguments: Some(partial),
}]),
});
}
}
Ok(StreamChunk { delta: String::new(), finished: false, tool_calls: None })
}
// 工具块开始：带 id + name
"content_block_start" => {
if let Some(cb) = v.get("content_block") {
if cb.get("type").and_then(|t| t.as_str()) == Some("tool_use") {
let idx = v.get("index").and_then(|i| i.as_u64()).unwrap_or(0) as u32;
let id = cb.get("id").and_then(|t| t.as_str()).map(|s| s.to_string());
let name = cb.get("name").and_then(|t| t.as_str()).map(|s| s.to_string());
return Ok(StreamChunk {
delta: String::new(),
finished: false,
tool_calls: Some(vec![ToolCallDelta {
index: idx,
id,
function_name: name,
function_arguments: None,
}]),
});
}
}
Ok(StreamChunk { delta: String::new(), finished: false, tool_calls: None })
}
// 消息结束
"message_stop" => Ok(StreamChunk { delta: String::new(), finished: true, tool_calls: None }),
// 错误事件
"error" => {
let msg = v.get("error").and_then(|e| e.get("message")).and_then(|m| m.as_str()).unwrap_or("stream error");
error!(%msg, "Anthropic 流式错误事件");
Ok(StreamChunk { delta: String::new(), finished: true, tool_calls: None })
}
// message_start / content_block_stop / message_delta 等不产出 chunk
_ => Ok(StreamChunk { delta: String::new(), finished: false, tool_calls: None }),
}
}
Ok(ev) => Ok(apply_anthropic_event(&ev.data, &mut usage_accum)),
Err(e) => {
error!(error = %e, "Anthropic SSE 事件流错误");
Err(anyhow::anyhow!("Anthropic SSE 错误: {}", e))
@@ -424,3 +482,184 @@ impl LlmProvider for AnthropicCompatProvider {
}
}
}
// ============================================================
// 单测（不发真实 HTTP，喂构造的 SSE data 字符串序列）
// ============================================================
#[cfg(test)]
mod tests {
use super::*;
/// 辅助：构造 message_start 事件
fn message_start(input_tokens: u32) -> String {
format!(
r#"{{"type":"message_start","message":{{"usage":{{"input_tokens":{},"output_tokens":0}}}}}}"#,
input_tokens
)
}
/// 辅助：构造 message_delta 事件（output_tokens 为累计值）
fn message_delta(output_tokens: u32) -> String {
format!(
r#"{{"type":"message_delta","delta":{{"stop_reason":"end_turn"}},"usage":{{"output_tokens":{}}}}}"#,
output_tokens
)
}
/// 辅助：构造文本增量 content_block_delta
fn text_delta(text: &str) -> String {
format!(
r#"{{"type":"content_block_delta","index":0,"delta":{{"type":"text_delta","text":"{}"}}}}"#,
text
)
}
/// 辅助：构造 message_stop 事件
fn message_stop() -> &'static str {
r#"{"type":"message_stop"}"#
}
/// 完整流：message_start 初始化 input + 多次 message_delta 累计覆盖 output + message_stop 带出
#[test]
fn anthropic_full_stream_accumulates_usage() {
let mut acc: Option<TokenUsage> = None;
// 1) message_start：input=42，output=0
let c = apply_anthropic_event(&message_start(42), &mut acc);
assert!(!c.finished);
assert!(c.usage.is_none());
let a = acc.as_ref().expect("message_start 应初始化累加器");
assert_eq!(a.prompt_tokens, 42);
assert_eq!(a.completion_tokens, 0);
assert_eq!(a.total_tokens, 42);
// 2) 文本增量不影响 usage
let c = apply_anthropic_event(&text_delta("Hello"), &mut acc);
assert_eq!(c.delta, "Hello");
assert!(!c.finished);
assert_eq!(acc.as_ref().unwrap().completion_tokens, 0, "文本增量不应改 output");
// 3) message_delta：output_tokens=10（累计值，覆盖）
let c = apply_anthropic_event(&message_delta(10), &mut acc);
assert!(!c.finished);
let a = acc.as_ref().unwrap();
assert_eq!(a.prompt_tokens, 42, "input 保持");
assert_eq!(a.completion_tokens, 10, "output 被覆盖为累计值");
assert_eq!(a.total_tokens, 52, "total 重算 = input+output");
// 4) 再次 message_delta：output_tokens=30（更大累计值，再覆盖）
let _ = apply_anthropic_event(&message_delta(30), &mut acc);
let a = acc.as_ref().unwrap();
assert_eq!(a.completion_tokens, 30, "后续累计值覆盖前值");
assert_eq!(a.total_tokens, 72);
// 5) message_stop：带出累积 usage，finished=true，累加器清空
let c = apply_anthropic_event(message_stop(), &mut acc);
assert!(c.finished);
let u = c.usage.expect("message_stop 应带出累积 usage");
assert_eq!(u.prompt_tokens, 42);
assert_eq!(u.completion_tokens, 30);
assert_eq!(u.total_tokens, 72);
assert!(acc.is_none(), "take() 后累加器应清空");
}
/// message_delta 在没有 message_start 时也能补全累加器（get_or_insert 兜底）
#[test]
fn anthropic_message_delta_without_start_uses_default_input() {
let mut acc: Option<TokenUsage> = None;
let _ = apply_anthropic_event(&message_delta(15), &mut acc);
let a = acc.as_ref().unwrap();
assert_eq!(a.prompt_tokens, 0, "无 message_start 时 input 兜底为 0");
assert_eq!(a.completion_tokens, 15);
assert_eq!(a.total_tokens, 15);
}
/// message_delta 的 output_tokens 必须是累计覆盖（非累加）：连续两个 delta 5 和 8，结果应是 8 不是 13
#[test]
fn anthropic_message_delta_output_is_cumulative_not_incremental() {
let mut acc: Option<TokenUsage> = None;
apply_anthropic_event(&message_start(100), &mut acc);
apply_anthropic_event(&message_delta(5), &mut acc);
apply_anthropic_event(&message_delta(8), &mut acc);
let c = apply_anthropic_event(message_stop(), &mut acc);
let u = c.usage.unwrap();
assert_eq!(u.completion_tokens, 8, "output_tokens 是累计值，覆盖而非累加");
assert_eq!(u.total_tokens, 108);
}
/// 无 usage 字段的流：message_stop 时 usage 为 None
#[test]
fn anthropic_message_stop_without_any_usage() {
let mut acc: Option<TokenUsage> = None;
let _ = apply_anthropic_event(&text_delta("hi"), &mut acc);
assert!(acc.is_none(), "文本增量不初始化累加器");
let c = apply_anthropic_event(message_stop(), &mut acc);
assert!(c.finished);
assert!(c.usage.is_none(), "无 usage 时 message_stop usage 为 None");
}
/// content_block_start (tool_use) 带 id+name
#[test]
fn anthropic_content_block_start_tool_use() {
let mut acc: Option<TokenUsage> = None;
let data = r#"{"type":"content_block_start","index":1,"content_block":{"type":"tool_use","id":"tool_1","name":"get_weather"}}"#;
let c = apply_anthropic_event(data, &mut acc);
assert!(acc.is_none(), "content_block_start 不动 usage");
let tcs = c.tool_calls.expect("应有 tool_calls");
assert_eq!(tcs.len(), 1);
assert_eq!(tcs[0].index, 1);
assert_eq!(tcs[0].id.as_deref(), Some("tool_1"));
assert_eq!(tcs[0].function_name.as_deref(), Some("get_weather"));
assert!(tcs[0].function_arguments.is_none());
assert!(!c.finished);
}
/// content_block_delta (input_json_delta) → 工具入参增量
#[test]
fn anthropic_content_block_delta_input_json() {
let mut acc: Option<TokenUsage> = None;
let data = r#"{"type":"content_block_delta","index":2,"delta":{"type":"input_json_delta","partial_json":"{\"q\":"}}"#;
let c = apply_anthropic_event(data, &mut acc);
let tcs = c.tool_calls.expect("应有 tool_calls 增量");
assert_eq!(tcs[0].index, 2);
assert_eq!(tcs[0].function_arguments.as_deref(), Some("{\"q\":"));
assert!(tcs[0].id.is_none());
assert_eq!(c.delta, "");
assert!(!c.finished);
}
/// error 事件 → finished=true 终态空 chunk
#[test]
fn anthropic_error_event_finishes_stream() {
let mut acc: Option<TokenUsage> = None;
apply_anthropic_event(&message_start(10), &mut acc);
let c = apply_anthropic_event(r#"{"type":"error","error":{"message":"overloaded"}}"#, &mut acc);
assert!(c.finished, "error 应终止流");
assert!(c.usage.is_none(), "error 不带出 usage");
assert!(acc.is_some(), "error 不应清空已累积的 usage（与原实现一致）");
}
/// ping / content_block_stop 等事件 → 空且不 finished
#[test]
fn anthropic_ping_and_block_stop_yield_empty_chunk() {
let mut acc: Option<TokenUsage> = None;
let c = apply_anthropic_event(r#"{"type":"ping"}"#, &mut acc);
assert!(!c.finished);
assert_eq!(c.delta, "");
assert!(acc.is_none());
let c = apply_anthropic_event(r#"{"type":"content_block_stop","index":0}"#, &mut acc);
assert!(!c.finished);
assert_eq!(c.delta, "");
}
/// 非法 JSON → 空 chunk，不 panic
#[test]
fn anthropic_malformed_json_yields_empty_chunk() {
let mut acc: Option<TokenUsage> = None;
let c = apply_anthropic_event("not json", &mut acc);
assert!(!c.finished);
assert_eq!(c.delta, "");
assert!(acc.is_none());
}
}

494
crates/df-ai/src/context.rs
View File

@@ -1,59 +1,511 @@
//! 上下文管理器 — 管理对话上下文和 token 预算
use std::collections::VecDeque;
//!
//! 职责：
//! - 维护消息历史及其 token 计数缓存
//! - 提供预算感知的消息裁剪（保护工具调用三元组）
//! - 为 run_agentic_loop 提供受控的消息视图
//!
//! 裁剪策略与模型选择是正交维度：本模块只管「窗口多大、怎么裁」，
//! 用哪个 model / 是否启用 reasoning 由调用方在 CompletionRequest 层决定。
use serde::{Deserialize, Serialize};
use crate::provider::ChatMessage;
use crate::provider::{ChatMessage, MessageRole};
// ============================================================
// Token 估算器（零依赖粗估）
// ============================================================
/// Token 粗估器 — 字符级近似计数，无 tokenizer 依赖
///
/// 用于发送前预算控制，误差 ±15% 完全可接受（保守估计，宁可多算）。
#[derive(Debug, Clone)]
pub struct TokenEstimator {
/// 字符 → token 转换系数（默认 0.35，即 ~2.8 字符/token，中英混合偏保守）
pub chars_ratio: f32,
/// 每条消息固定开销（role 标记 + 格式）
pub per_message_overhead: u32,
/// 每个 tool_call 的额外开销（name + arguments JSON 结构）
pub per_tool_call_overhead: u32,
}
impl Default for TokenEstimator {
fn default() -> Self {
Self {
chars_ratio: 0.35,
per_message_overhead: 4,
per_tool_call_overhead: 30,
}
}
}
impl TokenEstimator {
/// 估算单条消息的 token 数（保守估计）
pub fn estimate_message(&self, msg: &ChatMessage) -> u32 {
let content_tokens = (msg.content.chars().count() as f32 * self.chars_ratio).ceil() as u32;
let mut total = content_tokens + self.per_message_overhead;
// tool_calls 的 JSON 结构开销（role=Assistant 时可能有）
if let Some(ref calls) = msg.tool_calls {
for call in calls {
total += self.per_tool_call_overhead;
total += (call.function.name.chars().count() as f32 * self.chars_ratio).ceil() as u32;
total += (call.function.arguments.chars().count() as f32 * self.chars_ratio).ceil() as u32;
}
}
// tool_call_id 开销（role=Tool 时有）
if msg.tool_call_id.is_some() {
total += 3;
}
total
}
/// 估算纯文本字符串的 token 数（用于 system prompt）
pub fn estimate_text(&self, text: &str) -> u32 {
(text.chars().count() as f32 * self.chars_ratio).ceil() as u32
}
}
// ============================================================
// 上下文窗口配置
// ============================================================
/// 上下文窗口配置
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ContextConfig {
/// 最大 token 数
/// 窗口上限 token 数（默认 128_000）
pub max_tokens: u32,
/// 保留的系统提示 token 数
pub system_reserve: u32,
/// 输出预留 token 数（窗口中留给模型生成的部分，默认 8_192）
pub output_reserve: u32,
/// 安全系数 0.0~1.0（默认 0.85，留 15% 余量）
pub safety_ratio: f32,
}
impl Default for ContextConfig {
fn default() -> Self {
Self {
max_tokens: 128_000,
system_reserve: 4_000,
output_reserve: 8_192,
safety_ratio: 0.85,
}
}
}
impl ContextConfig {
/// 预算上限 = (max_tokens - output_reserve) × safety_ratio
pub fn budget_limit(&self) -> u32 {
(self.max_tokens.saturating_sub(self.output_reserve) as f32 * self.safety_ratio) as u32
}
}
// ============================================================
// 消息分组（淘汰时保持工具调用三元组原子性）
// ============================================================
/// 消息在逻辑上的分组标签，用于淘汰时保持原子性
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
enum MessageGroup {
/// 普通 User / Assistant 文本消息（可独立淘汰）
Standalone,
/// Assistant 带 tool_calls，是三元组的头
ToolCallHead,
/// Tool 结果消息，是三元组的尾
ToolResultTail,
}
/// 带有 token 缓存和分组信息的消息条目
struct TrackedMessage {
message: ChatMessage,
token_count: u32,
group: MessageGroup,
}
fn classify_group(msg: &ChatMessage) -> MessageGroup {
match msg.role {
MessageRole::Tool => MessageGroup::ToolResultTail,
MessageRole::Assistant => {
if msg.tool_calls.as_ref().is_some_and(|c| !c.is_empty()) {
MessageGroup::ToolCallHead
} else {
MessageGroup::Standalone
}
}
_ => MessageGroup::Standalone,
}
}
// ============================================================
// 上下文管理器
// ============================================================
/// 上下文管理器
///
/// 唯一的消息真相来源（替代原来的 `Vec<ChatMessage>`）。
/// 裁剪仅影响发送视图（`build_for_request`），不影响持久化（`all_messages_clone`）。
pub struct ContextManager {
/// 消息历史
messages: VecDeque<ChatMessage>,
/// 配置
messages: Vec<TrackedMessage>,
/// 当前历史总 token 数（不含 system prompt）
history_tokens: u32,
config: ContextConfig,
estimator: TokenEstimator,
}
/// 保护区大小：最后 N 条消息永不淘汰（≈ 最近 2 个完整用户轮次）
const PROTECT_COUNT: usize = 6;
impl ContextManager {
/// 创建上下文管理器
pub fn new(config: ContextConfig) -> Self {
Self {
messages: VecDeque::new(),
messages: Vec::new(),
history_tokens: 0,
config,
estimator: TokenEstimator::default(),
}
}
/// 添加消息
/// 追加消息（自动计算 token 并更新缓存）
///
/// 不在此处淘汰——push 可能发生在 agentic loop 中间（追加 tool_result），
/// 此时不应裁剪正在使用的活跃消息。裁剪在 `build_for_request` 时统一处理。
pub fn push(&mut self, message: ChatMessage) {
self.messages.push_back(message);
// TODO: 当超过 token 预算时，淘汰旧消息
let tokens = self.estimator.estimate_message(&message);
let group = classify_group(&message);
self.history_tokens += tokens;
self.messages.push(TrackedMessage {
message,
token_count: tokens,
group,
});
}
/// 获取当前消息列表
pub fn messages(&self) -> &VecDeque<ChatMessage> {
&self.messages
}
/// 清空上下文
/// 清空所有消息
pub fn clear(&mut self) {
self.messages.clear();
self.history_tokens = 0;
}
/// 消息数量
pub fn len(&self) -> usize {
self.messages.len()
}
pub fn is_empty(&self) -> bool {
self.messages.is_empty()
}
/// 当前历史占用的 token 数（不含 system prompt）
pub fn history_tokens(&self) -> u32 {
self.history_tokens
}
/// 预算上限
pub fn budget_limit(&self) -> u32 {
self.config.budget_limit()
}
// ── 核心：构建请求消息（受控裁剪版本）──
/// 构建发送给 LLM 的消息列表
///
/// `sys_tokens` 为调用方已估算好的 system prompt token 数。
/// 超预算时自动裁剪旧消息（保护工具调用三元组 + 最近 PROTECT_COUNT 条）。
/// 返回 (消息列表, 是否发生了裁剪)。
pub fn build_for_request(&self, sys_tokens: u32) -> (Vec<ChatMessage>, bool) {
let budget = self.budget_limit();
let available = budget.saturating_sub(sys_tokens);
// system prompt 自身超预算：裁剪无法缓解(仍返回保护区兜底),warn 便于诊断
if sys_tokens > budget {
tracing::warn!(
"system prompt (~{} tokens) 超过上下文预算 ({}),裁剪无法缓解",
sys_tokens, budget
);
}
// 未超预算 → 直接返回全量
if self.history_tokens <= available {
return (self.all_messages_clone(), false);
}
// 超预算 → 视图裁剪（不修改 self.messages，保证 all_messages_clone 仍返回全量）
let protect_start = self.messages.len().saturating_sub(PROTECT_COUNT);
let units = self.build_eviction_units(protect_start);
let mut removed: u64 = 0;
let mut trim_end = 0;
for unit in &units {
if self.history_tokens.saturating_sub(removed as u32) <= available {
break;
}
removed += unit.token_sum as u64;
trim_end = unit.end;
}
if trim_end == 0 {
tracing::warn!(
"history (~{} tokens) 超预算 ({}) 但无可淘汰单元(全在保护区 {} 条),发送兜底可能触发 provider 超限",
self.history_tokens, available, PROTECT_COUNT
);
return (self.all_messages_clone(), false);
}
let msgs: Vec<ChatMessage> = self.messages[trim_end..]
.iter()
.map(|t| t.message.clone())
.collect();
tracing::info!(
"context_trimmed: skip {} messages, ~{} tokens (view-only, full history retained)",
trim_end, removed
);
(msgs, true)
}
/// 全量克隆（持久化 save_conversation / build_for_request 未裁剪分支，不受裁剪影响）
pub fn all_messages_clone(&self) -> Vec<ChatMessage> {
self.messages.iter().map(|t| t.message.clone()).collect()
}
/// 从 Vec 恢复（兼容从 DB 加载）
pub fn restore_from_messages(&mut self, messages: Vec<ChatMessage>) {
self.clear();
for msg in messages {
self.push(msg);
}
}
/// 就地替换某条 tool_result 的内容（兼容审批 replace_tool_result）
/// 返回 true 如果找到并替换了
pub fn replace_tool_result_content(&mut self, tool_call_id: &str, new_content: &str) -> bool {
let pos = self.messages.iter().position(|t| {
matches!(t.message.role, MessageRole::Tool)
&& t.message.tool_call_id.as_deref() == Some(tool_call_id)
});
let Some(i) = pos else { return false };
// 先更新 content，再重估 token 并校正总量
let old_tokens = self.messages[i].token_count;
self.messages[i].message.content = new_content.to_string();
let new_tokens = self.estimator.estimate_message(&self.messages[i].message);
self.messages[i].token_count = new_tokens;
self.history_tokens = self.history_tokens.saturating_sub(old_tokens).saturating_add(new_tokens);
true
}
/// 只读迭代（兼容 ensure_conversation_title 的 .iter().filter() 等）
pub fn iter(&self) -> impl Iterator<Item = &ChatMessage> {
self.messages.iter().map(|t| &t.message)
}
// ── 内部方法 ──
/// 构建淘汰单元列表
///
/// 每个单元是连续消息范围 [start, end)，保证：
/// - 工具调用三元组（ToolCallHead + ToolResultTail* + 紧随的文本 Assistant）在同一单元
/// - 保护区内的消息不纳入任何单元
fn build_eviction_units(&self, protect_start: usize) -> Vec<EvictionUnit> {
let mut units = Vec::new();
let mut i = 0usize;
while i < protect_start {
let mut token_sum = 0u32;
if self.messages[i].group == MessageGroup::ToolCallHead {
// 收集完整三元组：Head + 后续所有 ToolResultTail + 紧随的文本 Assistant
token_sum += self.messages[i].token_count;
i += 1;
while i < protect_start && self.messages[i].group == MessageGroup::ToolResultTail {
token_sum += self.messages[i].token_count;
i += 1;
}
// 紧随的 Standalone Assistant（工具调用的最终文本回复）
if i < protect_start
&& self.messages[i].group == MessageGroup::Standalone
&& matches!(self.messages[i].message.role, MessageRole::Assistant)
{
token_sum += self.messages[i].token_count;
i += 1;
}
} else {
// Standalone / ToolResultTail（理论上孤立 Tail 不该出现，按单条处理）
token_sum += self.messages[i].token_count;
i += 1;
}
units.push(EvictionUnit { end: i, token_sum });
}
units
}
}
/// 淘汰单元：连续消息范围 [..end) + token 总和
struct EvictionUnit {
end: usize,
token_sum: u32,
}
#[cfg(test)]
mod tests {
use super::*;
use crate::provider::ToolCall;
fn cfg(max_tokens: u32) -> ContextConfig {
ContextConfig {
max_tokens,
output_reserve: 0,
safety_ratio: 1.0,
}
}
#[test]
fn short_history_no_trim() {
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::user("你好"));
mgr.push(ChatMessage::assistant("你好啊"));
let (msgs, trimmed) = mgr.build_for_request(10);
assert!(!trimmed);
assert_eq!(msgs.len(), 2);
}
#[test]
fn over_budget_trims_old() {
// 小预算强制裁剪：20 条超预算，触发裁剪且保留保护区
let mut mgr = ContextManager::new(cfg(200));
for i in 0..20 {
mgr.push(ChatMessage::user(&format!("这是第 {} 条较长的消息用于撑爆预算", i)));
}
let (msgs, trimmed) = mgr.build_for_request(0);
assert!(trimmed, "超预算应触发裁剪");
assert!(msgs.len() < 20, "应裁掉部分旧消息, 实际 {}", msgs.len());
// 保护区：最新一条必保留
assert_eq!(
msgs.last().unwrap().content,
"这是第 19 条较长的消息用于撑爆预算",
"保护区最新消息被误裁"
);
// 裁剪是视图：内存全量不变
assert_eq!(mgr.all_messages_clone().len(), 20, "裁剪污染了内存全量");
}
#[test]
fn tool_triplet_kept_atomic() {
// 三元组不可分离：Head 与 Tail 同进同出，永不从中间切断
// 布局：6 旧(淘汰区) + 三元组(裁剪边界) + 6 新(保护区) = 15 条
let mut mgr = ContextManager::new(cfg(95));
for i in 0..6 {
mgr.push(ChatMessage::user(&format!("旧消息 {}", i)));
}
mgr.push(ChatMessage::assistant_with_tools(
"调工具",
vec![ToolCall::new("tc1", "read_file", "{}")],
));
mgr.push(ChatMessage::tool_result("tc1", "文件内容"));
mgr.push(ChatMessage::assistant("完成"));
for i in 0..6 {
mgr.push(ChatMessage::user(&format!("新消息 {}", i)));
}
// 分支一：预算宽松，三元组整体保留 → Head 在则 Tail 在
let (msgs_keep, trimmed1) = mgr.build_for_request(0);
assert!(trimmed1, "分支一应触发裁剪");
assert_eq!(
has_head(&msgs_keep),
has_tail(&msgs_keep),
"分支一三元组被切断: head={} tail={}",
has_head(&msgs_keep),
has_tail(&msgs_keep)
);
// 分支二：预算紧张，三元组整体丢弃 → Head 不在则 Tail 也不在
let (msgs_drop, trimmed2) = mgr.build_for_request(40);
assert!(trimmed2, "分支二应触发裁剪");
assert_eq!(
has_head(&msgs_drop),
has_tail(&msgs_drop),
"分支二三元组被切断: head={} tail={}",
has_head(&msgs_drop),
has_tail(&msgs_drop)
);
// 裁剪是视图：两次 build 都不应改变内存全量
assert_eq!(
mgr.all_messages_clone().len(),
15,
"裁剪污染了内存全量"
);
}
fn has_head(msgs: &[ChatMessage]) -> bool {
msgs.iter()
.any(|m| matches!(m.role, MessageRole::Assistant) && m.tool_calls.is_some())
}
fn has_tail(msgs: &[ChatMessage]) -> bool {
msgs.iter().any(|m| matches!(m.role, MessageRole::Tool))
}
#[test]
fn replace_tool_result_updates_tokens() {
let mut mgr = ContextManager::new(cfg(100_000));
mgr.push(ChatMessage::tool_result("tc1", "短"));
let before = mgr.history_tokens();
assert!(mgr.replace_tool_result_content("tc1", "这是一个明显更长的替换内容用于验证 token 重估"));
let after = mgr.history_tokens();
assert!(after > before);
}
#[test]
fn restore_rebuilds_token_cache() {
let mut mgr = ContextManager::new(cfg(100_000));
let src = vec![
ChatMessage::user("测试消息一"),
ChatMessage::assistant("回复一"),
ChatMessage::user("测试消息二"),
];
mgr.restore_from_messages(src);
assert!(mgr.history_tokens() > 0);
assert_eq!(mgr.len(), 3);
}
#[test]
fn empty_history_returns_empty() {
let mgr = ContextManager::new(cfg(100_000));
let (msgs, trimmed) = mgr.build_for_request(10);
assert!(!trimmed);
assert!(msgs.is_empty(), "空历史应返回空列表");
}
#[test]
fn protect_zone_returns_full_when_untrimmable() {
// 消息全在保护区(< PROTECT_COUNT 条)且超预算 → 无可淘汰单元,走 trim_end==0 兜底返回全量
let mut mgr = ContextManager::new(cfg(10));
mgr.push(ChatMessage::user("撑爆小预算的长消息内容"));
mgr.push(ChatMessage::user("第二条撑爆预算的长消息"));
let (msgs, trimmed) = mgr.build_for_request(0);
assert!(!trimmed, "无可淘汰单元应返回 false(兜底)");
assert_eq!(msgs.len(), 2, "兜底返回全部保护区消息");
}
#[test]
fn system_over_budget_trims_to_protect_zone() {
// system prompt 吃光预算 → history 仍尝试裁剪到保护区,不 panic
let mut mgr = ContextManager::new(cfg(200));
for i in 0..10 {
mgr.push(ChatMessage::user(&format!("消息 {} 撑量", i)));
}
let (msgs, _trimmed) = mgr.build_for_request(195);
assert!(
msgs.len() <= PROTECT_COUNT,
"system 超预算时裁剪后至多保留保护区 {} 条,实际 {}",
PROTECT_COUNT,
msgs.len()
);
}
}

2
crates/df-ai/src/coordinator.rs
View File

@@ -1,4 +1,6 @@
//! Agent 协调器 — 管理多 Agent 协作
//!
//! ⚠ B 路线占位：当前单链 ReAct 够用，多 Agent 协作待 B 路线立项。有意保留空壳，勿删。
/// Agent 协调器
///

24
crates/df-ai/src/lib.rs
View File

@@ -8,3 +8,27 @@ pub mod openai_compat;
pub mod provider;
pub mod router;
pub mod stream;
use provider::LlmProvider;
/// 按 provider_type 构建 LLM Provider 实例（统一选择逻辑，消除调用方重复 match）
///
/// `anthropic` 协议走 AnthropicCompatProvider（GLM 订阅端点 / Claude 官方），
/// 其余（openai / glm / deepseek 等 OpenAI 兼容）走 OpenAICompatProvider。
/// 调用方（AI Chat 的 run_agentic_loop、df-nodes 的 AiNode）统一引用此工厂，
/// 新增 provider 只改这一处。
pub fn build_provider(
provider_type: &str,
base_url: &str,
api_key: &str,
model: &str,
) -> Box<dyn LlmProvider> {
match provider_type {
"anthropic" => Box::new(anthropic_compat::AnthropicCompatProvider::new(
base_url, api_key, model,
)),
_ => Box::new(openai_compat::OpenAICompatProvider::new(
base_url, api_key, model,
)),
}
}

373
crates/df-ai/src/openai_compat.rs
View File

@@ -3,14 +3,12 @@
//! 覆盖: OpenAI / GLM (open.bigmodel.cn) / DeepSeek / Claude OpenAI 兼容模式
//! 支持: 同步调用 + SSE 流式 + Function Calling / Tool Use
use std::pin::Pin;
use async_trait::async_trait;
use eventsource_stream::Eventsource;
use futures::{Stream, StreamExt};
use futures::StreamExt;
use reqwest::Client;
use serde::{Deserialize, Serialize};
use tracing::{debug, error, info, warn};
use tracing::{debug, error, warn};
use crate::provider::{
CompletionRequest, CompletionResponse, LlmProvider, ProviderFeatures, StreamChunk, StreamResult,
@@ -35,6 +33,9 @@ struct OpenAiRequest {
tools: Option<Vec<serde_json::Value>>,
#[serde(skip_serializing_if = "Option::is_none")]
tool_choice: Option<serde_json::Value>,
/// 流式时请求末 chunk 携带 usage（OpenAI 官方 + DeepSeek/GLM 兼容）
#[serde(skip_serializing_if = "Option::is_none")]
stream_options: Option<serde_json::Value>,
}
/// OpenAI 消息格式
@@ -93,6 +94,9 @@ struct OpenAiUsage {
#[derive(Debug, Deserialize)]
struct OpenAiStreamChunk {
choices: Vec<OpenAiStreamChoice>,
/// 末 chunk（choices 为空）携带的累计 usage
#[serde(default)]
usage: Option<OpenAiUsage>,
}
#[derive(Debug, Deserialize)]
@@ -120,6 +124,87 @@ struct OpenAiStreamFunction {
arguments: Option<String>,
}
// ============================================================
// SSE 解析纯函数（与 HTTP 解耦，便于单测）
// ============================================================
/// 将一条 OpenAI 兼容 SSE 事件 data 解析为 StreamChunk，并按需更新 usage 累加器。
///
/// - `[DONE]` → 返回 `finished=true` 的终态 chunk，`usage` 取自累加器（`take()`）。
/// - 普通文本/工具增量 chunk → 返回对应 `StreamChunk`，usage 字段恒为 None（usage 仅在终态带出）。
/// - usage（`stream_options.include_usage` 时末段或 usage-only chunk 携带）→ 覆盖累加器（覆盖语义保对）。
/// - 解析失败 → 返回空 chunk（与原内联实现一致）。
///
/// 等价性：delta / tool_calls / finished / 解析失败等分支与原 stream() 闭包逐字一致；
/// usage 透传（[DONE] 终态 take() 带出、usage chunk 覆盖累加器）为本次新增能力，
/// 对应 StreamChunk 新增的 usage 字段 + 请求体新增 stream_options.include_usage。
pub(crate) fn apply_openai_sse(data: &str, usage_accum: &mut Option<TokenUsage>) -> StreamChunk {
// OpenAI 发送 "data: [DONE]" 表示流结束，带出累积 usage
if data == "[DONE]" {
return StreamChunk {
delta: String::new(),
finished: true,
tool_calls: None,
usage: usage_accum.take(),
};
}
match serde_json::from_str::<OpenAiStreamChunk>(data) {
Ok(chunk) => {
// 提取 usage（带 include_usage 时末段 chunk 携带，覆盖累积）
if let Some(u) = chunk.usage {
*usage_accum = Some(TokenUsage {
prompt_tokens: u.prompt_tokens,
completion_tokens: u.completion_tokens,
total_tokens: u.total_tokens,
});
}
if let Some(choice) = chunk.choices.into_iter().next() {
let delta_text = choice.delta.content.unwrap_or_default();
// "length" = max_tokens 截断，属正常终止（非断连），纳入 finished
let finished = choice.finish_reason.as_deref() == Some("stop")
|| choice.finish_reason.as_deref() == Some("tool_calls")
|| choice.finish_reason.as_deref() == Some("length");
let tool_calls = choice.delta.tool_calls.map(|tcs| {
tcs.into_iter()
.map(|tc| ToolCallDelta {
index: tc.index,
id: tc.id,
function_name: tc.function.as_ref().and_then(|f| f.name.clone()),
function_arguments: tc.function.and_then(|f| f.arguments),
})
.collect()
});
StreamChunk {
delta: delta_text,
finished,
tool_calls,
usage: None,
}
} else {
// choices 为空 = usage-only chunk，不输出文本（usage 已累积）
StreamChunk {
delta: String::new(),
finished: false,
tool_calls: None,
usage: None,
}
}
}
Err(e) => {
debug!("SSE 数据解析失败: {} — data: {}", e, data);
StreamChunk {
delta: String::new(),
finished: false,
tool_calls: None,
usage: None,
}
}
}
}
// ============================================================
// OpenAI Compat Provider
// ============================================================
@@ -184,6 +269,18 @@ impl OpenAICompatProvider {
}
}
/// 构建 embeddings API URL（与 chat_url 同套智能拼接规则）
fn embed_url(&self) -> String {
let base = self.base_url.trim_end_matches('/');
if base.ends_with("/embeddings") {
return base.to_string();
}
if Self::ends_with_version(base) {
return format!("{}/embeddings", base);
}
format!("{}/v1/embeddings", base)
}
/// 将通用请求转换为 OpenAI 格式
fn convert_request(&self, req: CompletionRequest) -> OpenAiRequest {
let model = if req.model.is_empty() {
@@ -240,6 +337,12 @@ impl OpenAICompatProvider {
stream: req.stream,
tools,
tool_choice: req.tool_choice,
// 流式请求末 chunk 带 usage（同步调用 complete 不需要）
stream_options: if req.stream {
Some(serde_json::json!({ "include_usage": true }))
} else {
None
},
}
}
@@ -254,6 +357,34 @@ impl OpenAICompatProvider {
#[async_trait]
impl LlmProvider for OpenAICompatProvider {
/// 文本嵌入: POST /v1/embeddings(OpenAI 兼容,智谱/阿里百炼/OpenAI 通用)
async fn embed(&self, model: &str, texts: Vec<String>) -> anyhow::Result<Vec<Vec<f32>>> {
#[derive(serde::Deserialize)]
struct EmbedData { embedding: Vec<f32>, index: usize }
#[derive(serde::Deserialize)]
struct EmbedResponse { data: Vec<EmbedData> }
let resp = self
.client
.post(self.embed_url())
.header("Authorization", format!("Bearer {}", self.api_key))
.header("Content-Type", "application/json")
.json(&serde_json::json!({ "model": model, "input": texts }))
.send()
.await?;
if !resp.status().is_success() {
let status = resp.status();
let body = resp.text().await.unwrap_or_default();
anyhow::bail!("Embedding API 错误 {}: {}", status, body);
}
let mut body: EmbedResponse = resp.json().await?;
// 按 index 排序保证与输入顺序一致(API 不保证返回顺序)
body.data.sort_by_key(|d| d.index);
Ok(body.data.into_iter().map(|d| d.embedding).collect())
}
async fn complete(&self, request: CompletionRequest) -> anyhow::Result<CompletionResponse> {
let mut req = request;
req.stream = false;
@@ -328,68 +459,18 @@ impl LlmProvider for OpenAICompatProvider {
anyhow::bail!("LLM 流式 API 错误 {}: {}", status, body);
}
// 累积流式 usage：开 include_usage 后，末段正常 chunk（finish_reason）及额外 usage-only chunk（choices=[]）都带 usage。
// usage 解析/累积逻辑抽到 apply_openai_sse 纯函数，便于单测；此处闭包只负责传 data 与传递 last_usage。
let mut last_usage: Option<TokenUsage> = None;
let stream = resp
.bytes_stream()
.eventsource()
.map(move |event| {
match event {
Ok(event) => {
// OpenAI 发送 "data: [DONE]" 表示流结束
if event.data == "[DONE]" {
return Ok(StreamChunk {
delta: String::new(),
finished: true,
tool_calls: None,
});
}
match serde_json::from_str::<OpenAiStreamChunk>(&event.data) {
Ok(chunk) => {
if let Some(choice) = chunk.choices.into_iter().next() {
let delta_text = choice.delta.content.unwrap_or_default();
// "length" = max_tokens 截断，属正常终止（非断连），纳入 finished
let finished = choice.finish_reason.as_deref() == Some("stop")
|| choice.finish_reason.as_deref() == Some("tool_calls")
|| choice.finish_reason.as_deref() == Some("length");
let tool_calls = choice.delta.tool_calls.map(|tcs| {
tcs.into_iter()
.map(|tc| ToolCallDelta {
index: tc.index,
id: tc.id,
function_name: tc.function.as_ref().and_then(|f| f.name.clone()),
function_arguments: tc.function.and_then(|f| f.arguments),
})
.collect()
});
Ok(StreamChunk {
delta: delta_text,
finished,
tool_calls,
})
} else {
Ok(StreamChunk {
delta: String::new(),
finished: false,
tool_calls: None,
})
}
}
Err(e) => {
debug!("SSE 数据解析失败: {} — data: {}", e, event.data);
Ok(StreamChunk {
delta: String::new(),
finished: false,
tool_calls: None,
})
}
}
}
Err(e) => {
error!("SSE 流错误: {}", e);
Err(anyhow::anyhow!("SSE 流错误: {}", e))
}
.map(move |event| match event {
Ok(event) => Ok(apply_openai_sse(&event.data, &mut last_usage)),
Err(e) => {
error!("SSE 流错误: {}", e);
Err(anyhow::anyhow!("SSE 流错误: {}", e))
}
});
@@ -408,3 +489,169 @@ impl LlmProvider for OpenAICompatProvider {
}
}
}
// ============================================================
// 单测（不发真实 HTTP，喂构造的 SSE data 字符串序列）
// ============================================================
#[cfg(test)]
mod tests {
use super::*;
/// 辅助：构造普通文本 delta chunk 的 SSE data
fn text_chunk(content: &str, finish_reason: Option<&str>) -> String {
let fr = match finish_reason {
Some(r) => format!(", \"finish_reason\": \"{}\"", r),
None => String::from(", \"finish_reason\": null"),
};
format!(
r#"{{"choices":[{{"delta":{{"content":"{}"}}{}}}]}}"#,
content, fr
)
}
/// 辅助：构造带 usage 的 chunk（choices 为空 → usage-only 末 chunk，对应 include_usage）
fn usage_only_chunk(prompt: u32, completion: u32) -> String {
format!(
r#"{{"choices":[],"usage":{{"prompt_tokens":{},"completion_tokens":{},"total_tokens":{}}}}}"#,
prompt,
completion,
prompt + completion
)
}
/// 辅助：构造既有 content 又带 usage 的末段 chunk（部分兼容端点会把 usage 挂到正常末 chunk 上）
fn text_chunk_with_usage(content: &str, finish_reason: &str, prompt: u32, completion: u32) -> String {
format!(
r#"{{"choices":[{{"delta":{{"content":"{}"}},"finish_reason":"{}"}}],"usage":{{"prompt_tokens":{},"completion_tokens":{},"total_tokens":{}}}}}"#,
content,
finish_reason,
prompt,
completion,
prompt + completion
)
}
/// 多 chunk 文本流后，末 chunk 携带 usage（include_usage 覆盖语义）
#[test]
fn openai_sse_multi_chunk_with_terminal_usage() {
let mut acc: Option<TokenUsage> = None;
// 1) 首个文本增量，无 usage
let c = apply_openai_sse(&text_chunk("Hello", None), &mut acc);
assert_eq!(c.delta, "Hello");
assert!(!c.finished);
assert!(c.usage.is_none());
assert!(acc.is_none(), "无 usage 的 chunk 不应改累加器");
// 2) 第二个文本增量
let c = apply_openai_sse(&text_chunk(" world", None), &mut acc);
assert_eq!(c.delta, " world");
assert!(!c.finished);
assert!(acc.is_none());
// 3) 末段正常 chunk 带 finish_reason=stop（仍是文本 delta，不带 usage）
let c = apply_openai_sse(&text_chunk("", Some("stop")), &mut acc);
assert!(c.finished);
assert_eq!(c.delta, "");
assert!(acc.is_none(), "此 chunk 无 usage 字段，累加器仍为 None");
// 4) usage-only chunk（choices=[]）携带累计 usage → 覆盖累加器
let c = apply_openai_sse(&usage_only_chunk(12, 34), &mut acc);
assert!(!c.finished);
assert!(c.usage.is_none(), "非 [DONE] chunk 不带出 usage");
let acc = acc.expect("累加器应已被 usage-only chunk 覆盖写入");
assert_eq!(acc.prompt_tokens, 12);
assert_eq!(acc.completion_tokens, 34);
assert_eq!(acc.total_tokens, 46);
}
/// usage 挂在正常末段 chunk（含 finish_reason）上，而非独立 usage-only chunk
#[test]
fn openai_sse_usage_on_terminal_text_chunk() {
let mut acc: Option<TokenUsage> = None;
let c = apply_openai_sse(&text_chunk_with_usage("", "stop", 100, 200), &mut acc);
assert!(c.finished);
assert!(c.usage.is_none(), "非 [DONE] 不带出 usage，仅覆盖累加器");
let acc = acc.expect("末段 chunk 的 usage 应已覆盖累加器");
assert_eq!(acc.prompt_tokens, 100);
assert_eq!(acc.completion_tokens, 200);
assert_eq!(acc.total_tokens, 300);
}
/// [DONE] 时 take() 带出累积 usage，且取走后累加器清空
#[test]
fn openai_sse_done_takes_accumulated_usage() {
let mut acc: Option<TokenUsage> = None;
apply_openai_sse(&text_chunk("x", None), &mut acc);
apply_openai_sse(&usage_only_chunk(5, 7), &mut acc);
let c = apply_openai_sse("[DONE]", &mut acc);
assert!(c.finished);
let u = c.usage.expect("[DONE] 应带出累积 usage");
assert_eq!(u.prompt_tokens, 5);
assert_eq!(u.completion_tokens, 7);
assert_eq!(u.total_tokens, 12);
assert!(acc.is_none(), "take() 后累加器应清空");
}
/// 无 usage 的流：[DONE] 时 usage 字段为 None
#[test]
fn openai_sse_done_without_usage() {
let mut acc: Option<TokenUsage> = None;
apply_openai_sse(&text_chunk("hi", None), &mut acc);
let c = apply_openai_sse("[DONE]", &mut acc);
assert!(c.finished);
assert!(c.usage.is_none(), "全程无 usage 时 [DONE] usage 应为 None");
assert!(acc.is_none());
}
/// 后续 usage chunk 覆盖先前 usage（多轮 / 重发场景）
#[test]
fn openai_sse_later_usage_overrides_earlier() {
let mut acc: Option<TokenUsage> = None;
apply_openai_sse(&usage_only_chunk(1, 1), &mut acc);
apply_openai_sse(&usage_only_chunk(50, 60), &mut acc);
let c = apply_openai_sse("[DONE]", &mut acc);
let u = c.usage.unwrap();
assert_eq!(u.prompt_tokens, 50, "末 usage 应覆盖前值");
assert_eq!(u.completion_tokens, 60);
assert_eq!(u.total_tokens, 110);
}
/// finish_reason=length（max_tokens 截断）按正常终止处理
#[test]
fn openai_sse_length_finish_reason_treated_as_finished() {
let mut acc: Option<TokenUsage> = None;
let c = apply_openai_sse(&text_chunk("...", Some("length")), &mut acc);
assert!(c.finished, "length 应视为正常终止");
assert!(acc.is_none());
}
/// 非法 JSON data → 返回空 chunk，不 panic、不改累加器
#[test]
fn openai_sse_malformed_json_yields_empty_chunk() {
let mut acc: Option<TokenUsage> = None;
let c = apply_openai_sse("not a json", &mut acc);
assert_eq!(c.delta, "");
assert!(!c.finished);
assert!(c.usage.is_none());
assert!(acc.is_none());
}
/// tool_calls 增量解析
#[test]
fn openai_sse_tool_call_delta() {
let mut acc: Option<TokenUsage> = None;
let data = r#"{"choices":[{"delta":{"tool_calls":[{"index":0,"id":"call_1","function":{"name":"get_weather","arguments":"{\"q\":"}}]}}]}"#;
let c = apply_openai_sse(data, &mut acc);
assert!(acc.is_none());
let tcs = c.tool_calls.expect("应有 tool_calls 增量");
assert_eq!(tcs.len(), 1);
assert_eq!(tcs[0].index, 0);
assert_eq!(tcs[0].id.as_deref(), Some("call_1"));
assert_eq!(tcs[0].function_name.as_deref(), Some("get_weather"));
assert_eq!(tcs[0].function_arguments.as_deref(), Some("{\"q\":"));
assert!(!c.finished);
}
}

26
crates/df-ai/src/provider.rs
View File

@@ -45,23 +45,26 @@ pub struct ChatMessage {
/// 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>,
}
impl ChatMessage {
pub fn system(content: impl Into<String>) -> Self {
Self { role: MessageRole::System, content: content.into(), tool_call_id: None, tool_calls: None }
Self { role: MessageRole::System, content: content.into(), tool_call_id: None, tool_calls: None, model: None }
}
pub fn user(content: impl Into<String>) -> Self {
Self { role: MessageRole::User, content: content.into(), tool_call_id: None, tool_calls: None }
Self { role: MessageRole::User, content: content.into(), tool_call_id: None, tool_calls: None, model: None }
}
pub fn assistant(content: impl Into<String>) -> Self {
Self { role: MessageRole::Assistant, content: content.into(), tool_call_id: None, tool_calls: None }
Self { role: MessageRole::Assistant, content: content.into(), tool_call_id: None, tool_calls: None, model: 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) }
Self { role: MessageRole::Assistant, content: content.into(), tool_call_id: None, tool_calls: Some(tool_calls), model: 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 }
Self { role: MessageRole::Tool, content: content.into(), tool_call_id: Some(call_id.into()), tool_calls: None, model: None }
}
}
@@ -141,7 +144,7 @@ pub struct CompletionResponse {
}
/// Token 用量
#[derive(Debug, Clone, Serialize, Deserialize)]
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct TokenUsage {
pub prompt_tokens: u32,
pub completion_tokens: u32,
@@ -166,6 +169,9 @@ pub struct StreamChunk {
/// 工具调用增量（如有）
#[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>,
}
/// 工具调用增量（流式中的片段）
@@ -199,6 +205,14 @@ pub trait LlmProvider: Send + Sync {
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;