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

修复: token 汇总正确性(会话级双计/is_estimated跨端/流式usage传递) + is_estimated SELECT 缺失修复

Browse Source
This commit is contained in:
lxy
2026-08-09 23:59:55 +08:00
parent 11f4978ec1
commit 6bafdcd5a5
16 changed files with 189 additions and 127 deletions
Download Patch File Download Diff File Expand all files Collapse all files
crates/df-ai-core/src/provider.rs
+9 -6
View File
@@ -42,25 +42,25 @@ impl ContentPart {
impl ChatMessage {
pub fn system(content: impl Into<String>) -> Self {
Self { id: Some(new_message_id()), role: MessageRole::System, content: content.into(), parts: None, tool_call_id: None, tool_calls: None, model: None, status: None, reasoning_content: None, prompt_tokens: None, completion_tokens: None, prompt_cache_hit_tokens: None, prompt_cache_miss_tokens: None, reasoning_tokens: None, timestamp: Some(now_millis_i64()) }
Self { id: Some(new_message_id()), role: MessageRole::System, content: content.into(), parts: None, tool_call_id: None, tool_calls: None, model: None, status: None, reasoning_content: None, prompt_tokens: None, completion_tokens: None, prompt_cache_hit_tokens: None, prompt_cache_miss_tokens: None, reasoning_tokens: None, is_estimated: None, timestamp: Some(now_millis_i64()) }
}
pub fn user(content: impl Into<String>) -> Self {
Self { id: Some(new_message_id()), role: MessageRole::User, content: content.into(), parts: None, tool_call_id: None, tool_calls: None, model: None, status: None, reasoning_content: None, prompt_tokens: None, completion_tokens: None, prompt_cache_hit_tokens: None, prompt_cache_miss_tokens: None, reasoning_tokens: None, timestamp: Some(now_millis_i64()) }
Self { id: Some(new_message_id()), role: MessageRole::User, content: content.into(), parts: None, tool_call_id: None, tool_calls: None, model: None, status: None, reasoning_content: None, prompt_tokens: None, completion_tokens: None, prompt_cache_hit_tokens: None, prompt_cache_miss_tokens: None, reasoning_tokens: None, is_estimated: None, timestamp: Some(now_millis_i64()) }
}
pub fn assistant(content: impl Into<String>) -> Self {
Self { id: Some(new_message_id()), role: MessageRole::Assistant, content: content.into(), parts: None, tool_call_id: None, tool_calls: None, model: None, status: None, reasoning_content: None, prompt_tokens: None, completion_tokens: None, prompt_cache_hit_tokens: None, prompt_cache_miss_tokens: None, reasoning_tokens: None, timestamp: Some(now_millis_i64()) }
Self { id: Some(new_message_id()), role: MessageRole::Assistant, content: content.into(), parts: None, tool_call_id: None, tool_calls: None, model: None, status: None, reasoning_content: None, prompt_tokens: None, completion_tokens: None, prompt_cache_hit_tokens: None, prompt_cache_miss_tokens: None, reasoning_tokens: None, is_estimated: None, timestamp: Some(now_millis_i64()) }
}
pub fn assistant_with_tools(content: impl Into<String>, tool_calls: Vec<ToolCall>) -> Self {
Self { id: Some(new_message_id()), role: MessageRole::Assistant, content: content.into(), parts: None, tool_call_id: None, tool_calls: Some(tool_calls), model: None, status: None, reasoning_content: None, prompt_tokens: None, completion_tokens: None, prompt_cache_hit_tokens: None, prompt_cache_miss_tokens: None, reasoning_tokens: None, timestamp: Some(now_millis_i64()) }
Self { id: Some(new_message_id()), role: MessageRole::Assistant, content: content.into(), parts: None, tool_call_id: None, tool_calls: Some(tool_calls), model: None, status: None, reasoning_content: None, prompt_tokens: None, completion_tokens: None, prompt_cache_hit_tokens: None, prompt_cache_miss_tokens: None, reasoning_tokens: None, is_estimated: None, timestamp: Some(now_millis_i64()) }
}
pub fn tool_result(call_id: impl Into<String>, content: impl Into<String>) -> Self {
Self { id: Some(new_message_id()), role: MessageRole::Tool, content: content.into(), parts: None, tool_call_id: Some(call_id.into()), tool_calls: None, model: None, status: None, reasoning_content: None, prompt_tokens: None, completion_tokens: None, prompt_cache_hit_tokens: None, prompt_cache_miss_tokens: None, reasoning_tokens: None, timestamp: Some(now_millis_i64()) }
Self { id: Some(new_message_id()), role: MessageRole::Tool, content: content.into(), parts: None, tool_call_id: Some(call_id.into()), tool_calls: None, model: None, status: None, reasoning_content: None, prompt_tokens: None, completion_tokens: None, prompt_cache_hit_tokens: None, prompt_cache_miss_tokens: None, reasoning_tokens: None, is_estimated: None, timestamp: Some(now_millis_i64()) }
}
/// 多模态 user 消息：content 文本 + parts（含 Image 片）。
/// content 作为人类可读文本（也作非 vision 端点降级载荷）；parts 透传给 vision 端点。
pub fn user_parts(content: impl Into<String>, parts: Vec<ContentPart>) -> Self {
Self { id: Some(new_message_id()), role: MessageRole::User, content: content.into(), parts: Some(parts), tool_call_id: None, tool_calls: None, model: None, status: None, reasoning_content: None, prompt_tokens: None, completion_tokens: None, prompt_cache_hit_tokens: None, prompt_cache_miss_tokens: None, reasoning_tokens: None, timestamp: Some(now_millis_i64()) }
Self { id: Some(new_message_id()), role: MessageRole::User, content: content.into(), parts: Some(parts), tool_call_id: None, tool_calls: None, model: None, status: None, reasoning_content: None, prompt_tokens: None, completion_tokens: None, prompt_cache_hit_tokens: None, prompt_cache_miss_tokens: None, reasoning_tokens: None, is_estimated: None, timestamp: Some(now_millis_i64()) }
}
/// 是否含图片片（供 provider 判定走多模态分支）。
@@ -320,6 +320,7 @@ mod tests {
prompt_cache_hit_tokens: None,
prompt_cache_miss_tokens: None,
reasoning_tokens: None,
is_estimated: None,
timestamp: None,
};
assert_eq!(m.content, "字面量构造");
@@ -377,6 +378,7 @@ mod tests {
prompt_cache_hit_tokens: None,
prompt_cache_miss_tokens: None,
reasoning_tokens: None,
is_estimated: None,
timestamp: None,
};
let json = serde_json::to_string(&m).unwrap();
@@ -437,6 +439,7 @@ mod tests {
prompt_cache_hit_tokens: None,
prompt_cache_miss_tokens: None,
reasoning_tokens: None,
is_estimated: None,
timestamp: None,
};
let json = serde_json::to_string(&m).unwrap();
crates/df-ai-core/src/types.rs
+4
View File
@@ -138,6 +138,10 @@ pub struct ChatMessage {
/// 前端仅 > 0 时显示(reason 后缀);老 JSON 反序列化为 None。
#[serde(default, skip_serializing_if = "Option::is_none")]
pub reasoning_tokens: Option<u32>,
/// 本轮 prompt 是否估算值(round_usage.prompt_tokens==0 → estimated_prompt 兜底打标)。
/// 供 reload 逐条回显「估算」标注,对齐 live 态 AiCompleted.is_estimated;老消息 None(向前兼容)。
#[serde(default, skip_serializing_if = "Option::is_none")]
pub is_estimated: Option<bool>,
}
/// 当前 Unix 毫秒(ChatMessage 打戳用;df-ai-core 不依赖 df-types,内联避免新增依赖)。
crates/df-ai/src/anthropic_helpers.rs
+5 -3
View File
@@ -165,7 +165,7 @@ pub(crate) fn apply_anthropic_event(data: &str, usage_accum: &mut Option<TokenUs
reasoning_tokens: 0,
});
}
StreamChunk { delta: String::new(), finished: false, tool_calls: None, usage: None, error: None, reasoning_content: None }
StreamChunk { delta: String::new(), finished: false, tool_calls: None, usage: usage_accum.clone(), error: None, reasoning_content: None }
}
// 消息增量：output_tokens 是累计值（非增量），直接覆盖 completion + 重算 total
"message_delta" => {
@@ -173,9 +173,11 @@ pub(crate) fn apply_anthropic_event(data: &str, usage_accum: &mut Option<TokenUs
let acc = usage_accum
.get_or_insert(TokenUsage::default());
acc.completion_tokens = out as u32;
acc.total_tokens = acc.prompt_tokens + acc.completion_tokens;
acc.total_tokens = acc.prompt_tokens.saturating_add(acc.completion_tokens);
}
StreamChunk { delta: String::new(), finished: false, tool_calls: None, usage: None, error: None, reasoning_content: None }
// 加固:usage 挂到本帧(而非仅 message_stop 带出),中途断连/端点不发 message_stop
// 时仍能拿到真实 usage,对称 openai_helpers 的修复。
StreamChunk { delta: String::new(), finished: false, tool_calls: None, usage: usage_accum.clone(), error: None, reasoning_content: None }
}
// 文本增量
"content_block_delta" => {
crates/df-ai/src/openai_helpers.rs
+5 -3
View File
@@ -256,17 +256,19 @@ pub(crate) fn apply_openai_sse(data: &str, usage_accum: &mut Option<TokenUsage>)
delta: delta_text,
finished,
tool_calls,
usage: None,
// 加固:usage 与 finish_reason 同帧的端点(部分兼容实现),此处已累积则挂上,
// 使下游 stream_recv 不必等到 [DONE] 帧即可拿到真实 usage。
usage: usage_accum.clone(),
error: None,
reasoning_content: choice.delta.reasoning_content,
}
} else {
// choices 为空 = usage-only chunk，不输出文本（usage 已累积）
// choices 为空 = usage-only chunk，不输出文本（usage 已累积），usage 一并带出
StreamChunk {
delta: String::new(),
finished: false,
tool_calls: None,
usage: None,
usage: usage_accum.clone(),
error: None,
reasoning_content: None,
}
crates/df-storage/src/crud/message_repo.rs
+16 -9
View File
@@ -46,6 +46,7 @@ fn ai_message_from_row(row: &Row<'_>) -> std::result::Result<AiMessageRecord, ru
prompt_cache_hit_tokens: row.get("prompt_cache_hit_tokens")?,
prompt_cache_miss_tokens: row.get("prompt_cache_miss_tokens")?,
reasoning_tokens: row.get("reasoning_tokens")?,
is_estimated: row.get("is_estimated")?,
})
}
@@ -86,8 +87,8 @@ impl AiMessageRepo {
(id, conversation_id, seq, role, content, parts, tool_call_id,
tool_calls, model, status, reasoning_content, timestamp, created_at,
prompt_tokens, completion_tokens,
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)",
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens, is_estimated)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19)",
)
.map_err(storage_err)?;
for rec in &records {
@@ -96,7 +97,7 @@ impl AiMessageRepo {
rec.parts, rec.tool_call_id, rec.tool_calls, rec.model, rec.status,
rec.reasoning_content, rec.timestamp, rec.created_at,
rec.prompt_tokens, rec.completion_tokens,
rec.prompt_cache_hit_tokens, rec.prompt_cache_miss_tokens, rec.reasoning_tokens
rec.prompt_cache_hit_tokens, rec.prompt_cache_miss_tokens, rec.reasoning_tokens, rec.is_estimated
])
.map_err(storage_err)?;
}
@@ -122,7 +123,8 @@ impl AiMessageRepo {
"SELECT id, conversation_id, seq, role, content, parts, tool_call_id,
tool_calls, model, status, reasoning_content, timestamp, created_at,
prompt_tokens, completion_tokens,
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens,
is_estimated
FROM ai_messages WHERE conversation_id = ?1 ORDER BY seq ASC",
)
.map_err(storage_err)?;
@@ -164,13 +166,13 @@ impl AiMessageRepo {
"SELECT id, conversation_id, seq, role, content, parts, tool_call_id,
tool_calls, model, status, reasoning_content, timestamp, created_at,
prompt_tokens, completion_tokens,
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens, is_estimated
FROM ai_messages WHERE conversation_id = ?1 AND seq < ?2 ORDER BY seq DESC LIMIT ?3"
} else {
"SELECT id, conversation_id, seq, role, content, parts, tool_call_id,
tool_calls, model, status, reasoning_content, timestamp, created_at,
prompt_tokens, completion_tokens,
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens, is_estimated
FROM ai_messages WHERE conversation_id = ?1 ORDER BY seq DESC LIMIT ?2"
};
let mut stmt = guard.prepare(sql).map_err(storage_err)?;
@@ -289,8 +291,8 @@ impl AiMessageRepo {
(id, conversation_id, seq, role, content, parts, tool_call_id,
tool_calls, model, status, reasoning_content, timestamp, created_at,
prompt_tokens, completion_tokens,
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)",
prompt_cache_hit_tokens, prompt_cache_miss_tokens, reasoning_tokens, is_estimated)
VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19)",
)
.map_err(storage_err)?;
for rec in &records {
@@ -299,7 +301,7 @@ impl AiMessageRepo {
rec.parts, rec.tool_call_id, rec.tool_calls, rec.model, rec.status,
rec.reasoning_content, rec.timestamp, rec.created_at,
rec.prompt_tokens, rec.completion_tokens,
rec.prompt_cache_hit_tokens, rec.prompt_cache_miss_tokens, rec.reasoning_tokens
rec.prompt_cache_hit_tokens, rec.prompt_cache_miss_tokens, rec.reasoning_tokens, rec.is_estimated
])
.map_err(storage_err)?;
}
@@ -373,6 +375,7 @@ mod tests {
prompt_cache_hit_tokens: None,
prompt_cache_miss_tokens: None,
reasoning_tokens: None,
is_estimated: None,
}
}
@@ -535,6 +538,7 @@ mod tests {
prompt_cache_hit_tokens: None,
prompt_cache_miss_tokens: None,
reasoning_tokens: None,
is_estimated: None,
},
AiMessageRecord {
id: "new_1".into(),
@@ -555,6 +559,7 @@ mod tests {
prompt_cache_hit_tokens: None,
prompt_cache_miss_tokens: None,
reasoning_tokens: None,
is_estimated: None,
},
];
repo.replace_conversation("conv", records).await.expect("replace");
@@ -627,6 +632,7 @@ mod tests {
prompt_cache_hit_tokens: None,
prompt_cache_miss_tokens: None,
reasoning_tokens: None,
is_estimated: None,
}],
)
.await
@@ -668,6 +674,7 @@ mod tests {
prompt_cache_hit_tokens: None,
prompt_cache_miss_tokens: None,
reasoning_tokens: None,
is_estimated: None,
};
repo.replace_conversation("c", vec![rec()]).await.expect("1st");
repo.replace_conversation("c", vec![rec()]).await.expect("2nd");
crates/df-storage/src/migrations.rs
+31 -6
View File
@@ -46,7 +46,8 @@ pub fn run(conn: &Connection) -> Result<()> {
// V33 = 审批重启恢复:ai_conversations 加 pending_approvals TEXT 列,持久化挂起审批快照,
// 重启后从 DB 恢复 pending_approvals 内存态,使待审批不丢。
// V41 = 任务关联工程模块:tasks.module_id 列(工程系统打底,项目多工程下任务落到具体 module)。
let steps: [(i32, fn(&Connection) -> Result<()>); 41] = [
// V42 = ai_messages 加 is_estimated 列(消息级估算标记,reload 逐条回显「估算」标注)。
let steps: [(i32, fn(&Connection) -> Result<()>); 42] = [
(1, migrate_v1),
(2, migrate_v2),
(3, migrate_v3),
@@ -88,6 +89,7 @@ pub fn run(conn: &Connection) -> Result<()> {
(39, migrate_v39),
(40, migrate_v40),
(41, migrate_v41),
(42, migrate_v42),
];
for (version, migrate_fn) in steps {
@@ -1257,6 +1259,22 @@ fn migrate_v41(conn: &Connection) -> Result<()> {
Ok(())
}
/// V42: ai_messages 加 is_estimated 列(消息级估算标记)
///
/// 消息级 token 持久化(V38/V39)已落 prompt/completion/cache/reasoning,但「该轮 prompt 是否
/// estimated 兜底」未存——reload 逐条回显时无法对齐 live 态 AiCompleted.is_estimated 标注。
/// 本迁移补 INTEGER 列(NULL=老消息未标记,向前兼容;0=false 真实,1=true 估算)。
/// 用 PRAGMA 探测列存在性,缺失才 ALTER(同 v38/v39 模式),对新库/老库均安全。
fn migrate_v42(conn: &Connection) -> Result<()> {
if !column_exists(conn, "ai_messages", "is_estimated") {
conn.execute("ALTER TABLE ai_messages ADD COLUMN is_estimated INTEGER", [])?;
tracing::info!("v42: ai_messages 加 is_estimated 列(消息级估算标记)");
}
conn.execute("INSERT INTO schema_version (version) VALUES (?)", [42])?;
tracing::info!("迁移 v42 完成: ai_messages 加 is_estimated 列");
Ok(())
}
/// V21 建表 SQL — 消息拆分存储 ai_messages 表
///
/// 与 V9_SQL 中的 ai_messages 镜像(V9 给新库,此 const 给老库 V21 迁移用 IF NOT EXISTS)。
@@ -1281,6 +1299,7 @@ CREATE TABLE IF NOT EXISTS ai_messages (
prompt_cache_hit_tokens INTEGER,
prompt_cache_miss_tokens INTEGER,
reasoning_tokens INTEGER,
is_estimated INTEGER,
UNIQUE(conversation_id, seq)
);
@@ -1532,6 +1551,7 @@ CREATE TABLE IF NOT EXISTS ai_messages (
prompt_cache_hit_tokens INTEGER,
prompt_cache_miss_tokens INTEGER,
reasoning_tokens INTEGER,
is_estimated INTEGER,
UNIQUE(conversation_id, seq)
);
@@ -2023,7 +2043,7 @@ mod tests {
"tasks.module_id 列缺失(V41 加)"
);
// 3. schema_version 应推进到 41(全量迁移成功落版本号)
// 3. schema_version 应推进到 42(全量迁移成功落版本号)
let max_version: i64 = conn
.query_row(
"SELECT COALESCE(MAX(version), 0) FROM schema_version",
@@ -2032,8 +2052,8 @@ mod tests {
)
.expect("查 schema_version 应成功");
assert_eq!(
max_version, 41,
"全量迁移后 schema_version 应为 41(实际 {}),说明某条 migrate_vN 链路断在中间",
max_version, 42,
"全量迁移后 schema_version 应为 42(实际 {}),说明某条 migrate_vN 链路断在中间",
max_version
);
@@ -2046,6 +2066,11 @@ mod tests {
column_exists(&conn, "project_modules", "status"),
"project_modules.status 列缺失(V40 加)"
);
// 5. V42 抽查:ai_messages.is_estimated 列存在(消息级估算标记)
assert!(
column_exists(&conn, "ai_messages", "is_estimated"),
"ai_messages.is_estimated 列缺失(V42 加)"
);
}
// ============================================================
@@ -2105,7 +2130,7 @@ mod tests {
|r| r.get(0),
)
.unwrap();
assert_eq!(max_v, 41, "首轮应推进到 41");
assert_eq!(max_v, 42, "首轮应推进到 42");
// 清空版本表强制全链第二遍(每步 execute 第二次)
conn.execute("DELETE FROM schema_version", []).unwrap();
@@ -2117,7 +2142,7 @@ mod tests {
|r| r.get(0),
)
.unwrap();
assert_eq!(max_v2, 41, "重跑后应重新推进到 41");
assert_eq!(max_v2, 42, "重跑后应重新推进到 42");
}
// ============================================================
crates/df-storage/src/models.rs
+3
View File
@@ -440,6 +440,9 @@ pub struct AiMessageRecord {
pub prompt_cache_miss_tokens: Option<u32>,
/// 思考 token(deepseek-reasoner/o1 reasoning_tokens,隐藏输出)。
pub reasoning_tokens: Option<u32>,
/// 本轮 prompt 是否估算值(round_usage.prompt_tokens==0 → estimated_prompt 兜底打标)。
/// reload 逐条回显「估算」标注;老消息 NULL → None(向前兼容)。
pub is_estimated: Option<bool>,
}
// ============================================================