Fix cache token accounting and tiered pricing

This commit is contained in:
elky
2026-07-10 15:13:12 +08:00
parent 736fc76345
commit 4bf5d4c044
19 changed files with 506 additions and 264 deletions
@@ -4559,7 +4559,7 @@ mod tests {
assert!(sse.contains("\"prompt_tokens\":1"));
assert!(sse.contains("\"completion_tokens\":2"));
assert!(sse.contains("\"completion_tokens_details\":{\"reasoning_tokens\":1}"));
assert!(sse.contains("\"cached_creation_tokens\":5"));
assert!(sse.contains("\"cache_write_tokens\":5"));
assert!(sse.contains("\"cached_tokens\":4"));
assert!(sse.contains("\"total_tokens\":3"));
assert!(sse.contains("data: [DONE]\n\n"));
@@ -4683,7 +4683,7 @@ mod tests {
assert!(sse.contains("\"text\":\"because\""));
assert!(sse.contains("\"output_tokens_details\":{\"reasoning_tokens\":1}"));
assert!(sse.contains("\"input_tokens_details\""));
assert!(sse.contains("\"cached_creation_tokens\":5"));
assert!(sse.contains("\"cache_write_tokens\":5"));
assert!(sse.contains("\"cached_tokens\":4"));
}
@@ -852,7 +852,11 @@ fn openai_image_usage_to_standardized_usage(value: &Value) -> Option<Standardize
.get("input_tokens_details")
.or_else(|| usage.get("prompt_tokens_details"))
.and_then(Value::as_object)
.and_then(|details| details.get("cached_creation_tokens"))
.and_then(|details| {
details
.get("cache_write_tokens")
.or_else(|| details.get("cached_creation_tokens"))
})
.and_then(Value::as_i64)
})
.unwrap_or(0);
@@ -929,7 +933,11 @@ fn openai_image_chat_usage_counts(usage: Option<&Value>) -> Option<(u64, u64, u6
.get("input_tokens_details")
.or_else(|| usage.get("prompt_tokens_details"))
.and_then(Value::as_object)
.and_then(|details| details.get("cached_creation_tokens"))
.and_then(|details| {
details
.get("cache_write_tokens")
.or_else(|| details.get("cached_creation_tokens"))
})
.and_then(Value::as_u64)
})
.unwrap_or(0);
@@ -148,7 +148,11 @@ pub fn canonical_usage_from_openai_usage(value: Option<&Value>) -> Option<Canoni
.get("input_tokens_details")
.or_else(|| usage.get("prompt_tokens_details"))
.and_then(Value::as_object)
.and_then(|details| details.get("cached_creation_tokens"))
.and_then(|details| {
details
.get("cache_write_tokens")
.or_else(|| details.get("cached_creation_tokens"))
})
.and_then(Value::as_u64)
})
.unwrap_or(0);
@@ -742,7 +746,7 @@ fn insert_openai_token_details(
}
if cache_creation_tokens > 0 {
details.insert(
"cached_creation_tokens".to_string(),
"cache_write_tokens".to_string(),
Value::from(cache_creation_tokens),
);
}
@@ -1110,7 +1110,11 @@ fn standardized_usage_from_openai_usage(value: &Value) -> Option<StandardizedUsa
.get("input_tokens_details")
.or_else(|| usage.get("prompt_tokens_details"))
.and_then(Value::as_object)
.and_then(|details| details.get("cached_creation_tokens"))
.and_then(|details| {
details
.get("cache_write_tokens")
.or_else(|| details.get("cached_creation_tokens"))
})
.and_then(Value::as_i64)
})
.unwrap_or(0);
@@ -1279,7 +1283,7 @@ mod tests {
assert!(output.contains("![generated image 2](data:image/png;base64,d29ybGQ=)"));
assert!(output.contains("\"finish_reason\":\"stop\""));
assert!(output.contains("\"cached_tokens\":20"));
assert!(output.contains("\"cached_creation_tokens\":10"));
assert!(output.contains("\"cache_write_tokens\":10"));
assert!(output.contains("data: [DONE]"));
assert!(!output.contains("image_generation.completed"));
@@ -5794,7 +5794,8 @@ pub(crate) fn openai_usage_to_canonical(value: Option<&Value>) -> Option<Canonic
.and_then(Value::as_object)
.and_then(|details| {
details
.get("cached_creation_tokens")
.get("cache_write_tokens")
.or_else(|| details.get("cached_creation_tokens"))
.or_else(|| details.get("cache_creation_tokens"))
})
.and_then(Value::as_u64)
@@ -5957,7 +5958,7 @@ pub(crate) fn canonical_usage_to_openai(value: &CanonicalUsage) -> Value {
if output.get("prompt_tokens_details").is_none() {
output["prompt_tokens_details"] = json!({});
}
output["prompt_tokens_details"]["cached_creation_tokens"] =
output["prompt_tokens_details"]["cache_write_tokens"] =
Value::from(value.cache_write_tokens);
}
output
@@ -5985,7 +5986,7 @@ pub(crate) fn canonical_usage_to_openai_responses_usage(value: &CanonicalUsage)
if output.get("input_tokens_details").is_none() {
output["input_tokens_details"] = json!({});
}
output["input_tokens_details"]["cached_creation_tokens"] =
output["input_tokens_details"]["cache_write_tokens"] =
Value::from(value.cache_write_tokens);
}
output
@@ -7732,7 +7733,10 @@ mod tests {
],
"usage": {
"input_tokens": 3,
"input_tokens_details": {"cached_tokens": 2},
"input_tokens_details": {
"cache_write_tokens": 1,
"cached_tokens": 2
},
"output_tokens": 5,
"output_tokens_details": {"reasoning_tokens": 1},
"total_tokens": 8
@@ -7763,6 +7767,7 @@ mod tests {
));
assert_eq!(canonical_response_unknown_block_count(&canonical), 2);
assert_eq!(canonical.usage.as_ref().unwrap().cache_read_tokens, 2);
assert_eq!(canonical.usage.as_ref().unwrap().cache_write_tokens, 1);
assert_eq!(canonical.usage.as_ref().unwrap().reasoning_tokens, 1);
let rebuilt_chat = canonical_to_openai_chat_response(&canonical);
@@ -7788,6 +7793,10 @@ mod tests {
assert_eq!(rebuilt["output"][5]["type"], "local_shell_call_output");
assert_eq!(rebuilt["output"][5]["call_id"], "call_shell_1");
assert_eq!(rebuilt["usage"]["input_tokens_details"]["cached_tokens"], 2);
assert_eq!(
rebuilt["usage"]["input_tokens_details"]["cache_write_tokens"],
1
);
assert_eq!(
rebuilt["usage"]["output_tokens_details"]["reasoning_tokens"],
1
@@ -8156,7 +8165,7 @@ mod tests {
3
);
assert_eq!(
rebuilt_openai["usage"]["input_tokens_details"]["cached_creation_tokens"],
rebuilt_openai["usage"]["input_tokens_details"]["cache_write_tokens"],
2
);
assert_eq!(rebuilt_openai["usage"]["total_tokens"], 23);