admin/CC-Switch
1
0
Fork 0
mirror of https://github.com/farion1231/cc-switch.git synced 2026-07-27 08:14:33 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(proxy/usage): enhance token parser and logger for multi-format support

Browse Source 
Parser enhancements:
- Add OpenAI Chat Completions format parsing (prompt_tokens/completion_tokens)
- Add model field to TokenUsage for actual model name extraction
- Add from_codex_response_adjusted() for proper cache token handling
- Add debug logging for better stream event tracing

Logger enhancements:
- Add first_token_ms, provider_type, is_streaming, cost_multiplier fields
- Extend RequestLog struct with full metadata tracking
- Update log_with_calculation() signature for new fields

Calculator: Update tests with model field in TokenUsage.
This commit is contained in:
YoVinchen
2025-12-03 21:09:51 +08:00
parent 1e95a3f16b
commit 087622bf6d
3 changed files with 224 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src-tauri/src/proxy/usage/calculator.rs
+4
View File
@@ -105,6 +105,7 @@ mod tests {
output_tokens: 500,
cache_read_tokens: 200,
cache_creation_tokens: 100,
model: None,
};
let pricing = ModelPricing::from_strings("3.0", "15.0", "0.3", "3.75").unwrap();
@@ -134,6 +135,7 @@ mod tests {
output_tokens: 0,
cache_read_tokens: 0,
cache_creation_tokens: 0,
model: None,
};
let pricing = ModelPricing::from_strings("3.0", "15.0", "0", "0").unwrap();
@@ -153,6 +155,7 @@ mod tests {
output_tokens: 500,
cache_read_tokens: 0,
cache_creation_tokens: 0,
model: None,
};
let multiplier = Decimal::from_str("1.0").unwrap();
@@ -168,6 +171,7 @@ mod tests {
output_tokens: 1,
cache_read_tokens: 1,
cache_creation_tokens: 1,
model: None,
};
let pricing = ModelPricing::from_strings("0.075", "0.3", "0.01875", "0.075").unwrap();
src-tauri/src/proxy/usage/logger.rs
+28 -2
View File
@@ -18,9 +18,16 @@ pub struct RequestLog {
pub usage: TokenUsage,
pub cost: Option<CostBreakdown>,
pub latency_ms: u64,
pub first_token_ms: Option<u64>,
pub status_code: u16,
pub error_message: Option<String>,
pub session_id: Option<String>,
/// 供应商类型 (claude, claude_auth, codex, gemini, gemini_cli, openrouter)
pub provider_type: Option<String>,
/// 是否为流式请求
pub is_streaming: bool,
/// 成本倍数
pub cost_multiplier: String,
}
/// 使用量记录器
@@ -66,8 +73,9 @@ impl<'a> UsageLogger<'a> {
request_id, provider_id, app_type, model,
input_tokens, output_tokens, cache_read_tokens, cache_creation_tokens,
input_cost_usd, output_cost_usd, cache_read_cost_usd, cache_creation_cost_usd, total_cost_usd,
latency_ms, status_code, error_message, session_id, created_at
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18)",
latency_ms, first_token_ms, status_code, error_message, session_id,
provider_type, is_streaming, cost_multiplier, created_at
) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9, ?10, ?11, ?12, ?13, ?14, ?15, ?16, ?17, ?18, ?19, ?20, ?21, ?22)",
rusqlite::params![
log.request_id,
log.provider_id,
@@ -83,9 +91,13 @@ impl<'a> UsageLogger<'a> {
cache_creation_cost,
total_cost,
log.latency_ms as i64,
log.first_token_ms.map(|v| v as i64),
log.status_code as i64,
log.error_message,
log.session_id,
log.provider_type,
log.is_streaming as i64,
log.cost_multiplier,
created_at,
],
)
@@ -114,9 +126,13 @@ impl<'a> UsageLogger<'a> {
usage: TokenUsage::default(),
cost: None,
latency_ms,
first_token_ms: None,
status_code,
error_message: Some(error_message),
session_id: None,
provider_type: None,
is_streaming: false,
cost_multiplier: "1.0".to_string(),
};
self.log_request(&log)
@@ -147,8 +163,11 @@ impl<'a> UsageLogger<'a> {
usage: TokenUsage,
cost_multiplier: Decimal,
latency_ms: u64,
first_token_ms: Option<u64>,
status_code: u16,
session_id: Option<String>,
provider_type: Option<String>,
is_streaming: bool,
) -> Result<(), AppError> {
let pricing = self.get_model_pricing(&model)?;
@@ -166,9 +185,13 @@ impl<'a> UsageLogger<'a> {
usage,
cost,
latency_ms,
first_token_ms,
status_code,
error_message: None,
session_id,
provider_type,
is_streaming,
cost_multiplier: cost_multiplier.to_string(),
};
self.log_request(&log)
@@ -201,6 +224,7 @@ mod tests {
output_tokens: 500,
cache_read_tokens: 0,
cache_creation_tokens: 0,
model: None,
};
logger.log_with_calculation(
@@ -213,6 +237,8 @@ mod tests {
100,
200,
None,
Some("claude".to_string()),
false,
)?;
// 验证记录已插入
src-tauri/src/proxy/usage/parser.rs
+192 -3
View File
@@ -16,6 +16,8 @@ pub struct TokenUsage {
pub output_tokens: u32,
pub cache_read_tokens: u32,
pub cache_creation_tokens: u32,
/// 从响应中提取的实际模型名称（如果可用）
pub model: Option<String>,
}
/// API 类型
@@ -43,6 +45,7 @@ impl TokenUsage {
.get("cache_creation_input_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32,
model: None,
})
}
@@ -96,15 +99,33 @@ impl TokenUsage {
output_tokens: usage.get("completion_tokens")?.as_u64()? as u32,
cache_read_tokens: 0,
cache_creation_tokens: 0,
model: None,
})
}
/// 从 Codex API 非流式响应解析
pub fn from_codex_response(body: &Value) -> Option<Self> {
let usage = body.get("usage")?;
let usage = body.get("usage");
if usage.is_none() {
log::debug!(
"[Codex] 响应中没有 usage 字段，body keys: {:?}",
body.as_object().map(|o| o.keys().collect::<Vec<_>>())
);
return None;
}
let usage = usage?;
let input_tokens = usage.get("input_tokens").and_then(|v| v.as_u64());
let output_tokens = usage.get("output_tokens").and_then(|v| v.as_u64());
if input_tokens.is_none() || output_tokens.is_none() {
log::debug!("[Codex] usage 字段缺少 input_tokens 或 output_tokens，usage: {usage:?}");
return None;
}
Some(Self {
input_tokens: usage.get("input_tokens")?.as_u64()? as u32,
output_tokens: usage.get("output_tokens")?.as_u64()? as u32,
input_tokens: input_tokens? as u32,
output_tokens: output_tokens? as u32,
cache_read_tokens: usage
.get("cache_read_input_tokens")
.and_then(|v| v.as_u64())
@@ -113,27 +134,110 @@ impl TokenUsage {
.get("cache_creation_input_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32,
model: None,
})
}
/// 从 Codex API 响应解析并调整 input_tokens
///
/// Codex 的 input_tokens 需要减去 cached_tokens 以获得实际计费的 token 数
/// 公式: adjusted_input = max(input_tokens - cached_tokens, 0)
#[allow(dead_code)]
pub fn from_codex_response_adjusted(body: &Value) -> Option<Self> {
let usage = body.get("usage")?;
let input_tokens = usage.get("input_tokens")?.as_u64()? as u32;
let output_tokens = usage.get("output_tokens")?.as_u64()? as u32;
// 获取 cached_tokens (可能在 input_tokens_details 中)
let cached_tokens = usage
.get("input_tokens_details")
.and_then(|d| d.get("cached_tokens"))
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32;
// 调整 input_tokens: 减去 cached_tokens
let adjusted_input = input_tokens.saturating_sub(cached_tokens);
Some(Self {
input_tokens: adjusted_input,
output_tokens,
cache_read_tokens: cached_tokens,
cache_creation_tokens: usage
.get("cache_creation_input_tokens")
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32,
model: None,
})
}
/// 从 Codex API 流式响应解析
#[allow(dead_code)]
pub fn from_codex_stream_events(events: &[Value]) -> Option<Self> {
log::debug!("[Codex] 解析流式事件，共 {} 个事件", events.len());
for event in events {
if let Some(event_type) = event.get("type").and_then(|v| v.as_str()) {
log::debug!("[Codex] 事件类型: {event_type}");
if event_type == "response.completed" {
if let Some(response) = event.get("response") {
log::debug!("[Codex] 找到 response.completed 事件，解析 usage");
return Self::from_codex_response(response);
}
}
}
}
log::debug!("[Codex] 未找到 response.completed 事件");
None
}
/// 从 OpenAI Chat Completions API 响应解析 (prompt_tokens, completion_tokens)
pub fn from_openai_response(body: &Value) -> Option<Self> {
let usage = body.get("usage")?;
// OpenAI 使用 prompt_tokens 和 completion_tokens
let prompt_tokens = usage.get("prompt_tokens").and_then(|v| v.as_u64())?;
let completion_tokens = usage.get("completion_tokens").and_then(|v| v.as_u64())?;
// 获取 cached_tokens (可能在 prompt_tokens_details 中)
let cached_tokens = usage
.get("prompt_tokens_details")
.and_then(|d| d.get("cached_tokens"))
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32;
Some(Self {
input_tokens: prompt_tokens as u32,
output_tokens: completion_tokens as u32,
cache_read_tokens: cached_tokens,
cache_creation_tokens: 0,
model: None,
})
}
/// 从 OpenAI Chat Completions API 流式响应解析
pub fn from_openai_stream_events(events: &[Value]) -> Option<Self> {
log::debug!("[Codex] 解析 OpenAI 流式事件，共 {} 个事件", events.len());
// OpenAI 流式响应在最后一个 chunk 中包含 usage
for event in events.iter().rev() {
if let Some(usage) = event.get("usage") {
if !usage.is_null() {
log::debug!("[Codex] 找到 usage: {usage:?}");
return Self::from_openai_response(event);
}
}
}
log::debug!("[Codex] 未找到 usage 信息");
None
}
/// 从 Gemini API 非流式响应解析
pub fn from_gemini_response(body: &Value) -> Option<Self> {
let usage = body.get("usageMetadata")?;
// 提取实际使用的模型名称（modelVersion 字段）
let model = body
.get("modelVersion")
.and_then(|v| v.as_str())
.map(|s| s.to_string());
Some(Self {
input_tokens: usage.get("promptTokenCount")?.as_u64()? as u32,
output_tokens: usage.get("candidatesTokenCount")?.as_u64()? as u32,
@@ -142,6 +246,7 @@ impl TokenUsage {
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32,
cache_creation_tokens: 0,
model,
})
}
@@ -151,6 +256,7 @@ impl TokenUsage {
let mut total_input = 0u32;
let mut total_output = 0u32;
let mut total_cache_read = 0u32;
let mut model: Option<String> = None;
for chunk in chunks {
if let Some(usage) = chunk.get("usageMetadata") {
@@ -167,6 +273,13 @@ impl TokenUsage {
.and_then(|v| v.as_u64())
.unwrap_or(0) as u32;
}
// 提取实际使用的模型名称（modelVersion 字段）
if model.is_none() {
if let Some(model_version) = chunk.get("modelVersion").and_then(|v| v.as_str()) {
model = Some(model_version.to_string());
}
}
}
if total_input > 0 || total_output > 0 {
@@ -175,6 +288,7 @@ impl TokenUsage {
output_tokens: total_output,
cache_read_tokens: total_cache_read,
cache_creation_tokens: 0,
model,
})
} else {
None
@@ -252,6 +366,7 @@ mod tests {
#[test]
fn test_gemini_response_parsing() {
let response = json!({
"modelVersion": "gemini-3-pro-high",
"usageMetadata": {
"promptTokenCount": 100,
"candidatesTokenCount": 50,
@@ -264,5 +379,79 @@ mod tests {
assert_eq!(usage.output_tokens, 50);
assert_eq!(usage.cache_read_tokens, 20);
assert_eq!(usage.cache_creation_tokens, 0);
assert_eq!(usage.model, Some("gemini-3-pro-high".to_string()));
}
#[test]
fn test_gemini_response_parsing_no_model() {
// 测试没有 modelVersion 字段的情况
let response = json!({
"usageMetadata": {
"promptTokenCount": 100,
"candidatesTokenCount": 50,
"cachedContentTokenCount": 20
}
});
let usage = TokenUsage::from_gemini_response(&response).unwrap();
assert_eq!(usage.input_tokens, 100);
assert_eq!(usage.output_tokens, 50);
assert_eq!(usage.cache_read_tokens, 20);
assert_eq!(usage.cache_creation_tokens, 0);
assert_eq!(usage.model, None);
}
#[test]
fn test_codex_response_adjusted() {
let response = json!({
"usage": {
"input_tokens": 1000,
"output_tokens": 500,
"input_tokens_details": {
"cached_tokens": 300
}
}
});
let usage = TokenUsage::from_codex_response_adjusted(&response).unwrap();
// input_tokens 应该被调整: 1000 - 300 = 700
assert_eq!(usage.input_tokens, 700);
assert_eq!(usage.output_tokens, 500);
assert_eq!(usage.cache_read_tokens, 300);
}
#[test]
fn test_codex_response_adjusted_no_cache() {
let response = json!({
"usage": {
"input_tokens": 1000,
"output_tokens": 500
}
});
let usage = TokenUsage::from_codex_response_adjusted(&response).unwrap();
// 没有 cached_tokens，input_tokens 保持不变
assert_eq!(usage.input_tokens, 1000);
assert_eq!(usage.output_tokens, 500);
assert_eq!(usage.cache_read_tokens, 0);
}
#[test]
fn test_codex_response_adjusted_saturating_sub() {
// 测试 cached_tokens > input_tokens 的边界情况
let response = json!({
"usage": {
"input_tokens": 100,
"output_tokens": 50,
"input_tokens_details": {
"cached_tokens": 200
}
}
});
let usage = TokenUsage::from_codex_response_adjusted(&response).unwrap();
// saturating_sub 确保不会下溢
assert_eq!(usage.input_tokens, 0);
assert_eq!(usage.cache_read_tokens, 200);
}
}