mirror of
https://github.com/farion1231/cc-switch.git
synced 2026-07-27 08:14:33 +08:00
feat(proxy): extract model name from API response for accurate usage tracking
- Add model field extraction in TokenUsage parsing for Claude, OpenAI, and Codex - Prioritize response model over request model in usage logging - Update model extractors to use parsed usage.model first - Add tests for model extraction in stream and non-stream responses
This commit is contained in:
@@ -31,13 +31,26 @@ pub struct UsageParserConfig {
|
||||
// 模型提取器实现
|
||||
// ============================================================================
|
||||
|
||||
/// Claude 流式响应模型提取(直接使用请求模型)
|
||||
fn claude_model_extractor(_events: &[Value], request_model: &str) -> String {
|
||||
/// Claude 流式响应模型提取(优先使用 usage.model)
|
||||
fn claude_model_extractor(events: &[Value], request_model: &str) -> String {
|
||||
// 首先尝试从解析的 usage 中获取模型
|
||||
if let Some(usage) = TokenUsage::from_claude_stream_events(events) {
|
||||
if let Some(model) = usage.model {
|
||||
return model;
|
||||
}
|
||||
}
|
||||
request_model.to_string()
|
||||
}
|
||||
|
||||
/// OpenAI Chat Completions 流式响应模型提取
|
||||
/// OpenAI Chat Completions 流式响应模型提取(优先使用 usage.model)
|
||||
fn openai_model_extractor(events: &[Value], request_model: &str) -> String {
|
||||
// 首先尝试从解析的 usage 中获取模型
|
||||
if let Some(usage) = TokenUsage::from_openai_stream_events(events) {
|
||||
if let Some(model) = usage.model {
|
||||
return model;
|
||||
}
|
||||
}
|
||||
// 回退:从事件中直接提取
|
||||
events
|
||||
.iter()
|
||||
.find_map(|e| e.get("model")?.as_str())
|
||||
@@ -45,8 +58,15 @@ fn openai_model_extractor(events: &[Value], request_model: &str) -> String {
|
||||
.to_string()
|
||||
}
|
||||
|
||||
/// Codex Responses API 流式响应模型提取
|
||||
/// Codex Responses API 流式响应模型提取(优先使用 usage.model)
|
||||
fn codex_model_extractor(events: &[Value], request_model: &str) -> String {
|
||||
// 首先尝试从解析的 usage 中获取模型
|
||||
if let Some(usage) = TokenUsage::from_codex_stream_events(events) {
|
||||
if let Some(model) = usage.model {
|
||||
return model;
|
||||
}
|
||||
}
|
||||
// 回退:从 response.completed 事件中提取
|
||||
events
|
||||
.iter()
|
||||
.find_map(|e| {
|
||||
|
||||
@@ -93,12 +93,16 @@ pub async fn handle_non_streaming(
|
||||
|
||||
// 解析使用量
|
||||
if let Some(usage) = (parser_config.response_parser)(&json_value) {
|
||||
let model = json_value
|
||||
.get("model")
|
||||
.and_then(|m| m.as_str())
|
||||
.unwrap_or(&ctx.request_model);
|
||||
// 优先使用 usage 中解析出的模型名称,其次使用响应中的 model 字段,最后回退到请求模型
|
||||
let model = if let Some(ref m) = usage.model {
|
||||
m.clone()
|
||||
} else if let Some(m) = json_value.get("model").and_then(|m| m.as_str()) {
|
||||
m.to_string()
|
||||
} else {
|
||||
ctx.request_model.clone()
|
||||
};
|
||||
|
||||
spawn_log_usage(state, ctx, usage, model, status.as_u16(), false);
|
||||
spawn_log_usage(state, ctx, usage, &model, status.as_u16(), false);
|
||||
} else {
|
||||
log::debug!(
|
||||
"[{}] 未能解析 usage 信息,跳过记录",
|
||||
|
||||
@@ -34,6 +34,12 @@ impl TokenUsage {
|
||||
/// 从 Claude API 非流式响应解析
|
||||
pub fn from_claude_response(body: &Value) -> Option<Self> {
|
||||
let usage = body.get("usage")?;
|
||||
// 提取响应中的模型名称
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
Some(Self {
|
||||
input_tokens: usage.get("input_tokens")?.as_u64()? as u32,
|
||||
output_tokens: usage.get("output_tokens")?.as_u64()? as u32,
|
||||
@@ -45,7 +51,7 @@ impl TokenUsage {
|
||||
.get("cache_creation_input_tokens")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0) as u32,
|
||||
model: None,
|
||||
model,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -53,11 +59,20 @@ impl TokenUsage {
|
||||
#[allow(dead_code)]
|
||||
pub fn from_claude_stream_events(events: &[Value]) -> Option<Self> {
|
||||
let mut usage = Self::default();
|
||||
let mut model: Option<String> = None;
|
||||
|
||||
for event in events {
|
||||
if let Some(event_type) = event.get("type").and_then(|v| v.as_str()) {
|
||||
match event_type {
|
||||
"message_start" => {
|
||||
// 从 message_start 提取模型名称
|
||||
if model.is_none() {
|
||||
if let Some(message) = event.get("message") {
|
||||
if let Some(m) = message.get("model").and_then(|v| v.as_str()) {
|
||||
model = Some(m.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
if let Some(msg_usage) = event.get("message").and_then(|m| m.get("usage")) {
|
||||
// 从 message_start 获取 input_tokens(原生 Claude API)
|
||||
if let Some(input) =
|
||||
@@ -102,6 +117,7 @@ impl TokenUsage {
|
||||
}
|
||||
|
||||
if usage.input_tokens > 0 || usage.output_tokens > 0 {
|
||||
usage.model = model;
|
||||
Some(usage)
|
||||
} else {
|
||||
None
|
||||
@@ -141,6 +157,12 @@ impl TokenUsage {
|
||||
return None;
|
||||
}
|
||||
|
||||
// 提取响应中的模型名称
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
Some(Self {
|
||||
input_tokens: input_tokens? as u32,
|
||||
output_tokens: output_tokens? as u32,
|
||||
@@ -152,7 +174,7 @@ impl TokenUsage {
|
||||
.get("cache_creation_input_tokens")
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0) as u32,
|
||||
model: None,
|
||||
model,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -222,12 +244,18 @@ impl TokenUsage {
|
||||
.and_then(|v| v.as_u64())
|
||||
.unwrap_or(0) as u32;
|
||||
|
||||
// 提取响应中的模型名称
|
||||
let model = body
|
||||
.get("model")
|
||||
.and_then(|v| v.as_str())
|
||||
.map(|s| s.to_string());
|
||||
|
||||
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,
|
||||
model,
|
||||
})
|
||||
}
|
||||
|
||||
@@ -322,6 +350,7 @@ mod tests {
|
||||
#[test]
|
||||
fn test_claude_response_parsing() {
|
||||
let response = json!({
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"output_tokens": 50,
|
||||
@@ -335,10 +364,60 @@ mod tests {
|
||||
assert_eq!(usage.output_tokens, 50);
|
||||
assert_eq!(usage.cache_read_tokens, 20);
|
||||
assert_eq!(usage.cache_creation_tokens, 10);
|
||||
assert_eq!(usage.model, Some("claude-sonnet-4-20250514".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_claude_response_parsing_no_model() {
|
||||
let response = json!({
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"output_tokens": 50,
|
||||
"cache_read_input_tokens": 20,
|
||||
"cache_creation_input_tokens": 10
|
||||
}
|
||||
});
|
||||
|
||||
let usage = TokenUsage::from_claude_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, 10);
|
||||
assert_eq!(usage.model, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_claude_stream_parsing() {
|
||||
let events = vec![
|
||||
json!({
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"usage": {
|
||||
"input_tokens": 100,
|
||||
"cache_read_input_tokens": 20,
|
||||
"cache_creation_input_tokens": 10
|
||||
}
|
||||
}
|
||||
}),
|
||||
json!({
|
||||
"type": "message_delta",
|
||||
"usage": {
|
||||
"output_tokens": 50
|
||||
}
|
||||
}),
|
||||
];
|
||||
|
||||
let usage = TokenUsage::from_claude_stream_events(&events).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, 10);
|
||||
assert_eq!(usage.model, Some("claude-sonnet-4-20250514".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn test_claude_stream_parsing_no_model() {
|
||||
let events = vec![
|
||||
json!({
|
||||
"type": "message_start",
|
||||
@@ -363,6 +442,7 @@ mod tests {
|
||||
assert_eq!(usage.output_tokens, 50);
|
||||
assert_eq!(usage.cache_read_tokens, 20);
|
||||
assert_eq!(usage.cache_creation_tokens, 10);
|
||||
assert_eq!(usage.model, None);
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -481,6 +561,7 @@ mod tests {
|
||||
json!({
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"usage": {
|
||||
"input_tokens": 0,
|
||||
"output_tokens": 0
|
||||
@@ -502,6 +583,7 @@ mod tests {
|
||||
let usage = TokenUsage::from_claude_stream_events(&events).unwrap();
|
||||
assert_eq!(usage.input_tokens, 150);
|
||||
assert_eq!(usage.output_tokens, 75);
|
||||
assert_eq!(usage.model, Some("claude-sonnet-4-20250514".to_string()));
|
||||
}
|
||||
|
||||
#[test]
|
||||
@@ -512,6 +594,7 @@ mod tests {
|
||||
json!({
|
||||
"type": "message_start",
|
||||
"message": {
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"usage": {
|
||||
"input_tokens": 200,
|
||||
"cache_read_input_tokens": 50
|
||||
@@ -530,5 +613,6 @@ mod tests {
|
||||
assert_eq!(usage.input_tokens, 200);
|
||||
assert_eq!(usage.output_tokens, 100);
|
||||
assert_eq!(usage.cache_read_tokens, 50);
|
||||
assert_eq!(usage.model, Some("claude-sonnet-4-20250514".to_string()));
|
||||
}
|
||||
}
|
||||
|
||||
+1
-3
@@ -641,9 +641,7 @@ function App() {
|
||||
</header>
|
||||
|
||||
<main className="flex-1 pb-12 animate-fade-in ">
|
||||
<div className="pb-12">
|
||||
{renderContent()}
|
||||
</div>
|
||||
<div className="pb-12">{renderContent()}</div>
|
||||
</main>
|
||||
|
||||
<AddProviderDialog
|
||||
|
||||
@@ -51,7 +51,12 @@ export const FullScreenPanel: React.FC<FullScreenPanelProps> = ({
|
||||
>
|
||||
<div className="h-4 w-full" data-tauri-drag-region />
|
||||
<div className="mx-auto max-w-[56rem] px-6 flex items-center gap-4">
|
||||
<Button type="button" variant="outline" size="icon" onClick={onClose}>
|
||||
<Button
|
||||
type="button"
|
||||
variant="outline"
|
||||
size="icon"
|
||||
onClick={onClose}
|
||||
>
|
||||
<ArrowLeft className="h-4 w-4" />
|
||||
</Button>
|
||||
<h2 className="text-lg font-semibold text-foreground">{title}</h2>
|
||||
|
||||
Reference in New Issue
Block a user