From a47add0a8309e9fbcfc9b05a6a9c3d1ae0c37cce Mon Sep 17 00:00:00 2001 From: jamestotah <135163520+jamestotah@users.noreply.github.com> Date: Wed, 25 Mar 2026 06:05:40 -0700 Subject: [PATCH] fix(litellm): parse nested usage payloads (#8081) Signed-off-by: James Totah <135163520+jamestotah@users.noreply.github.com> Signed-off-by: Douwe Osinga Co-authored-by: Douwe Osinga --- crates/goose/src/providers/base.rs | 36 +++++++++++++ crates/goose/src/providers/formats/openai.rs | 55 +++++++++++++++++++- 2 files changed, 90 insertions(+), 1 deletion(-) diff --git a/crates/goose/src/providers/base.rs b/crates/goose/src/providers/base.rs index 546a9be2..eb3516ce 100644 --- a/crates/goose/src/providers/base.rs +++ b/crates/goose/src/providers/base.rs @@ -372,6 +372,8 @@ pub struct Usage { pub input_tokens: Option, pub output_tokens: Option, pub total_tokens: Option, + pub cache_read_input_tokens: Option, + pub cache_write_input_tokens: Option, } fn sum_optionals(a: Option, b: Option) -> Option @@ -395,6 +397,13 @@ impl Add for Usage { sum_optionals(self.output_tokens, other.output_tokens), sum_optionals(self.total_tokens, other.total_tokens), ) + .with_cache_tokens( + sum_optionals(self.cache_read_input_tokens, other.cache_read_input_tokens), + sum_optionals( + self.cache_write_input_tokens, + other.cache_write_input_tokens, + ), + ) } } @@ -425,8 +434,20 @@ impl Usage { input_tokens, output_tokens, total_tokens: calculated_total, + cache_read_input_tokens: None, + cache_write_input_tokens: None, } } + + pub fn with_cache_tokens( + mut self, + cache_read_input_tokens: Option, + cache_write_input_tokens: Option, + ) -> Self { + self.cache_read_input_tokens = cache_read_input_tokens; + self.cache_write_input_tokens = cache_write_input_tokens; + self + } } pub trait ProviderDef: Send + Sync { @@ -1086,4 +1107,19 @@ mod tests { assert_eq!(info.output_token_cost, Some(0.00001)); assert_eq!(info.currency, Some("$".to_string())); } + + #[test] + fn test_usage_addition_includes_cached_tokens() { + let usage_a = + Usage::new(Some(100), Some(20), Some(120)).with_cache_tokens(Some(10), Some(5)); + let usage_b = Usage::new(Some(50), Some(8), Some(58)).with_cache_tokens(Some(4), Some(1)); + + let combined = usage_a + usage_b; + + assert_eq!(combined.input_tokens, Some(150)); + assert_eq!(combined.output_tokens, Some(28)); + assert_eq!(combined.total_tokens, Some(178)); + assert_eq!(combined.cache_read_input_tokens, Some(14)); + assert_eq!(combined.cache_write_input_tokens, Some(6)); + } } diff --git a/crates/goose/src/providers/formats/openai.rs b/crates/goose/src/providers/formats/openai.rs index 28420888..92570429 100644 --- a/crates/goose/src/providers/formats/openai.rs +++ b/crates/goose/src/providers/formats/openai.rs @@ -490,6 +490,11 @@ pub fn response_to_message(response: &Value) -> anyhow::Result { } pub fn get_usage(usage: &Value) -> Usage { + let usage = usage + .get("usage") + .filter(|nested| nested.is_object()) + .unwrap_or(usage); + let input_tokens = usage .get("prompt_tokens") .and_then(|v| v.as_i64()) @@ -500,16 +505,27 @@ pub fn get_usage(usage: &Value) -> Usage { .and_then(|v| v.as_i64()) .map(|v| v as i32); + let cache_read_input_tokens = usage + .get("cache_read_input_tokens") + .and_then(|v| v.as_i64()) + .map(|v| v as i32); + + let cache_write_input_tokens = usage + .get("cache_creation_input_tokens") + .and_then(|v| v.as_i64()) + .map(|v| v as i32); + let total_tokens = usage .get("total_tokens") .and_then(|v| v.as_i64()) .map(|v| v as i32) .or_else(|| match (input_tokens, output_tokens) { - (Some(input), Some(output)) => Some(input + output), + (Some(input), Some(output)) => Some(input.saturating_add(output)), _ => None, }); Usage::new(input_tokens, output_tokens, total_tokens) + .with_cache_tokens(cache_read_input_tokens, cache_write_input_tokens) } fn extract_usage_with_output_tokens(chunk: &StreamingChunk) -> Option { @@ -1730,6 +1746,43 @@ mod tests { assert_eq!(usage.usage.total_tokens, Some(expected_total)); } + #[test] + fn test_get_usage_preserves_provider_totals_with_cache_fields() { + let usage = get_usage(&json!({ + "prompt_tokens": 120, + "completion_tokens": 30, + "total_tokens": 150, + "cache_read_input_tokens": 80, + "cache_creation_input_tokens": 20 + })); + + assert_eq!(usage.input_tokens, Some(120)); + assert_eq!(usage.output_tokens, Some(30)); + assert_eq!(usage.total_tokens, Some(150)); + assert_eq!(usage.cache_read_input_tokens, Some(80)); + assert_eq!(usage.cache_write_input_tokens, Some(20)); + } + + #[test] + fn test_get_usage_reads_nested_usage_object() { + let usage = get_usage(&json!({ + "id": "chatcmpl_test", + "usage": { + "prompt_tokens": 84, + "completion_tokens": 21, + "total_tokens": 105, + "cache_read_input_tokens": 60, + "cache_creation_input_tokens": 10 + } + })); + + assert_eq!(usage.input_tokens, Some(84)); + assert_eq!(usage.output_tokens, Some(21)); + assert_eq!(usage.total_tokens, Some(105)); + assert_eq!(usage.cache_read_input_tokens, Some(60)); + assert_eq!(usage.cache_write_input_tokens, Some(10)); + } + #[tokio::test] async fn test_streamed_multi_tool_response_to_messages() -> anyhow::Result<()> { let response_lines = r#"