From 2b58df57fc5bae08155ffd7198643ed739e6a463 Mon Sep 17 00:00:00 2001 From: hugocasa Date: Mon, 20 Jul 2026 18:47:25 +0200 Subject: [PATCH] feat(ai): extract prompt cache token usage from OpenAI and Azure providers (#10214) * feat(ai): extract prompt cache token usage from OpenAI and Azure providers Parse the nested cache token details OpenAI returns and thread them into TokenUsage, matching the Anthropic and Bedrock providers. - sse.rs: add OpenAIPromptTokensDetails / OpenAIInputTokensDetails and the optional prompt_tokens_details / input_tokens_details fields. - other.rs (Chat Completions) and openai.rs (Responses): populate cache_read via .with_cache(cached_tokens, None). OpenAI's prompt_tokens/input_tokens already include cached tokens (cached is a subset), so total/prompt are unchanged; cache_read is recorded separately for reporting. For the same reason the frontend token-usage conversions are left as-is (adding cached would double-count); optional cache fields and a clarifying comment are added to prevent a future incorrect Anthropic-style fix. Fixes WIN-2207 Co-Authored-By: Claude Opus 4.8 (1M context) * test(ai): pin OpenAI/Azure cache-token deserialization paths Add regression tests deserializing the real Chat Completions and Responses usage payloads, guarding the prompt_tokens_details.cached_tokens / input_tokens_details.cached_tokens paths against a silent rename. Co-Authored-By: Claude Opus 4.8 (1M context) * refactor(ai): extract to_token_usage() and test the cache mapping Address review nit: move the usage->TokenUsage conversion into OpenAIChatUsage::to_token_usage / OpenAIResponsesUsage::to_token_usage so the providers call one method and the tests exercise the real mapping. Tests now assert cache_read is populated while input/total are unchanged. Co-Authored-By: Claude Opus 4.8 (1M context) --------- Co-authored-by: Claude Opus 4.8 (1M context) --- backend/windmill-ai/src/providers/openai.rs | 4 +- backend/windmill-ai/src/providers/other.rs | 3 +- backend/windmill-ai/src/sse.rs | 77 ++++++++++++++++++- .../lib/components/copilot/chat/tokenUsage.ts | 6 ++ 4 files changed, 84 insertions(+), 6 deletions(-) diff --git a/backend/windmill-ai/src/providers/openai.rs b/backend/windmill-ai/src/providers/openai.rs index 8a90d84658..8ddd752b70 100644 --- a/backend/windmill-ai/src/providers/openai.rs +++ b/backend/windmill-ai/src/providers/openai.rs @@ -518,9 +518,7 @@ impl QueryBuilder for OpenAIQueryBuilder { parser.parse_events(response).await?; // Convert OpenAI Responses usage to TokenUsage - let usage = parser - .usage - .map(|u| TokenUsage::new(u.input_tokens, u.output_tokens, u.total_tokens)); + let usage = parser.usage.map(|u| u.to_token_usage()); Ok(ParsedResponse::Text { content: if parser.accumulated_content.is_empty() { diff --git a/backend/windmill-ai/src/providers/other.rs b/backend/windmill-ai/src/providers/other.rs index f641daec64..790bcb6c28 100644 --- a/backend/windmill-ai/src/providers/other.rs +++ b/backend/windmill-ai/src/providers/other.rs @@ -269,8 +269,7 @@ impl QueryBuilder for OtherQueryBuilder { } // Convert OpenAI Chat Completions usage to TokenUsage - let usage = openai_usage - .map(|u| TokenUsage::new(u.prompt_tokens, u.completion_tokens, u.total_tokens)); + let usage = openai_usage.map(|u| u.to_token_usage()); Ok(ParsedResponse::Text { content: if accumulated_content.is_empty() { diff --git a/backend/windmill-ai/src/sse.rs b/backend/windmill-ai/src/sse.rs index 66ff7da9de..96692d0173 100644 --- a/backend/windmill-ai/src/sse.rs +++ b/backend/windmill-ai/src/sse.rs @@ -13,7 +13,7 @@ use crate::{ AnthropicExtraContent, ExtraContent, GoogleExtraContent, OpenAIFunction, OpenAIToolCall, }, query_builder::StreamEventSink, - types::StreamingEvent, + types::{StreamingEvent, TokenUsage}, }; #[derive(Deserialize)] @@ -44,6 +44,14 @@ pub struct OpenAIChoice { pub delta: Option, } +/// Nested prompt token details returned by the Chat Completions API. +/// `cached_tokens` is the portion of `prompt_tokens` served from cache (a subset, not additive). +#[derive(Deserialize, Debug, Clone, Default)] +pub struct OpenAIPromptTokensDetails { + #[serde(default)] + pub cached_tokens: Option, +} + /// OpenAI Chat Completions API usage information (from final chunk with stream_options.include_usage) #[derive(Deserialize, Debug, Clone, Default)] pub struct OpenAIChatUsage { @@ -53,6 +61,24 @@ pub struct OpenAIChatUsage { pub completion_tokens: Option, #[serde(default)] pub total_tokens: Option, + #[serde(default)] + pub prompt_tokens_details: Option, +} + +impl OpenAIChatUsage { + /// cached_tokens is a subset of prompt_tokens, so input/total are reported as-is + /// and only recorded as cache_read for reporting. + pub fn to_token_usage(self) -> TokenUsage { + TokenUsage::new( + self.prompt_tokens, + self.completion_tokens, + self.total_tokens, + ) + .with_cache( + self.prompt_tokens_details.and_then(|d| d.cached_tokens), + None, + ) + } } #[derive(Deserialize)] @@ -684,6 +710,14 @@ pub struct OpenAIUrlCitationEvent { pub title: Option, } +/// Nested input token details returned by the Responses API. +/// `cached_tokens` is the portion of `input_tokens` served from cache (a subset, not additive). +#[derive(Deserialize, Debug, Clone, Default)] +pub struct OpenAIInputTokensDetails { + #[serde(default)] + pub cached_tokens: Option, +} + /// OpenAI Responses API usage information #[derive(Deserialize, Debug, Clone)] pub struct OpenAIResponsesUsage { @@ -693,6 +727,19 @@ pub struct OpenAIResponsesUsage { pub output_tokens: Option, #[serde(default)] pub total_tokens: Option, + #[serde(default)] + pub input_tokens_details: Option, +} + +impl OpenAIResponsesUsage { + /// cached_tokens is a subset of input_tokens, so input/total are reported as-is + /// and only recorded as cache_read for reporting. + pub fn to_token_usage(self) -> TokenUsage { + TokenUsage::new(self.input_tokens, self.output_tokens, self.total_tokens).with_cache( + self.input_tokens_details.and_then(|d| d.cached_tokens), + None, + ) + } } /// OpenAI Responses API response object (from response.completed event) @@ -927,6 +974,34 @@ mod tests { assert_eq!(json["content"], "hmm"); } + #[test] + fn openai_chat_usage_maps_cached_prompt_tokens() { + // Payload shape returned by OpenAI and Azure OpenAI Chat Completions. + // cached_tokens lives under prompt_tokens_details and is a subset of prompt_tokens, + // so it must land in cache_read while input/total stay as the provider reported them. + let usage: OpenAIChatUsage = serde_json::from_str( + r#"{"prompt_tokens":4819,"completion_tokens":1,"total_tokens":4820,"prompt_tokens_details":{"cached_tokens":4736,"audio_tokens":0}}"#, + ) + .unwrap(); + let token_usage = usage.to_token_usage(); + assert_eq!(token_usage.cache_read_input_tokens, Some(4736)); + assert_eq!(token_usage.input_tokens, Some(4819)); + assert_eq!(token_usage.total_tokens, Some(4820)); + } + + #[test] + fn openai_responses_usage_maps_cached_input_tokens() { + // Payload shape returned by the OpenAI Responses API. + let usage: OpenAIResponsesUsage = serde_json::from_str( + r#"{"input_tokens":4819,"input_tokens_details":{"cache_write_tokens":0,"cached_tokens":4736},"output_tokens":2,"total_tokens":4821}"#, + ) + .unwrap(); + let token_usage = usage.to_token_usage(); + assert_eq!(token_usage.cache_read_input_tokens, Some(4736)); + assert_eq!(token_usage.input_tokens, Some(4819)); + assert_eq!(token_usage.total_tokens, Some(4821)); + } + #[test] fn openai_delta_parses_reasoning_content() { // DeepSeek and similar stream reasoning under `reasoning_content`. diff --git a/frontend/src/lib/components/copilot/chat/tokenUsage.ts b/frontend/src/lib/components/copilot/chat/tokenUsage.ts index 51290f9269..0091e11227 100644 --- a/frontend/src/lib/components/copilot/chat/tokenUsage.ts +++ b/frontend/src/lib/components/copilot/chat/tokenUsage.ts @@ -70,12 +70,15 @@ export function anthropicUsageToChatTokenUsage( } } +// Unlike Anthropic, OpenAI's input_tokens already includes cached tokens +// (input_tokens_details.cached_tokens is a subset), so it must not be added again. export function openAIResponsesUsageToChatTokenUsage( usage: | { input_tokens?: number | null output_tokens?: number | null total_tokens?: number | null + input_tokens_details?: { cached_tokens?: number | null } | null } | null | undefined @@ -90,12 +93,15 @@ export function openAIResponsesUsageToChatTokenUsage( } } +// Unlike Anthropic, OpenAI's prompt_tokens already includes cached tokens +// (prompt_tokens_details.cached_tokens is a subset), so it must not be added again. export function openAICompletionsUsageToChatTokenUsage( usage: | { prompt_tokens?: number | null completion_tokens?: number | null total_tokens?: number | null + prompt_tokens_details?: { cached_tokens?: number | null } | null } | null | undefined