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

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
+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");