diff --git a/backend/windmill-ai/src/ai_google.rs b/backend/windmill-ai/src/ai_google.rs index 5f5b560055..aa2d6d0761 100644 --- a/backend/windmill-ai/src/ai_google.rs +++ b/backend/windmill-ai/src/ai_google.rs @@ -292,6 +292,22 @@ pub struct GeminiUsageMetadata { /// Thinking tokens, billed as output but counted apart from `candidatesTokenCount`. #[serde(rename = "thoughtsTokenCount", default)] pub thoughts_token_count: Option, + /// Input tokens spent on tool-use prompts, counted apart from `promptTokenCount` + /// rather than within it. + #[serde(rename = "toolUsePromptTokenCount", default)] + pub tool_use_prompt_token_count: Option, +} + +/// Input tokens as billed. Gemini reports tool-use prompts in their own field, and +/// they are disjoint from `promptTokenCount`: a live tool call returns 17 prompt + +/// 60 tool-use + 17 candidates + 52 thoughts against a `totalTokenCount` of 146, so +/// leaving them out under-reports the input of every tool-using turn. Cached tokens +/// are not added here, being already part of `promptTokenCount`. +fn gemini_prompt_tokens(usage: &GeminiUsageMetadata) -> i32 { + usage + .prompt_token_count + .unwrap_or(0) + .saturating_add(usage.tool_use_prompt_token_count.unwrap_or(0)) } /// Output tokens as billed: Gemini counts thinking apart from `candidatesTokenCount` @@ -605,7 +621,7 @@ pub fn gemini_response_to_openai(parsed: &GeminiParsedEvent, model: &str) -> ser let usage = parsed.usage.as_ref().map(|u| { serde_json::json!({ - "prompt_tokens": u.prompt_token_count.unwrap_or(0), + "prompt_tokens": gemini_prompt_tokens(u), "completion_tokens": gemini_completion_tokens(u), "total_tokens": u.total_token_count.unwrap_or(0), "prompt_tokens_details": { @@ -700,7 +716,7 @@ pub fn gemini_event_to_openai_sse_chunks( // OpenAI's `stream_options.include_usage` terminal chunk (top-level `usage`, // empty `choices`) so the frontend's `'usage' in chunk` path records them. if let Some(usage) = &parsed.usage { - let prompt_tokens = usage.prompt_token_count.unwrap_or(0); + let prompt_tokens = gemini_prompt_tokens(usage); let completion_tokens = gemini_completion_tokens(usage); let total_tokens = usage .total_token_count @@ -1003,6 +1019,7 @@ mod tests { total_token_count: Some(1120), cached_content_token_count: Some(900), thoughts_token_count: Some(100), + ..Default::default() }), ..Default::default() }; @@ -1028,6 +1045,47 @@ mod tests { assert_eq!(usage_chunk["usage"]["completion_tokens"], 120); } + #[test] + fn gemini_usage_chunk_counts_tool_use_prompt_tokens() { + // The four counts a live tool call reported, against its totalTokenCount of + // 146 — tool-use prompts are disjoint from promptTokenCount, so the input is + // 17 + 60 rather than 17. + let parsed = GeminiParsedEvent { + text: Some("Canberra".to_string()), + usage: Some(GeminiUsageMetadata { + prompt_token_count: Some(17), + candidates_token_count: Some(17), + total_token_count: Some(146), + tool_use_prompt_token_count: Some(60), + thoughts_token_count: Some(52), + ..Default::default() + }), + ..Default::default() + }; + + let mut tool_call_index = 0; + let chunks = gemini_event_to_openai_sse_chunks( + &parsed, + "chatcmpl-test", + "gemini-2.5-flash", + &mut tool_call_index, + ); + let usage_chunk = chunks + .iter() + .map(|c| parse_sse_chunk(c)) + .find(|v| v.get("usage").map(|u| !u.is_null()).unwrap_or(false)) + .expect("a chunk should carry top-level usage"); + + assert_eq!(usage_chunk["usage"]["prompt_tokens"], 77); + assert_eq!(usage_chunk["usage"]["completion_tokens"], 69); + // Everything Google billed is accounted for. + assert_eq!( + usage_chunk["usage"]["prompt_tokens"].as_i64().unwrap() + + usage_chunk["usage"]["completion_tokens"].as_i64().unwrap(), + 146 + ); + } + #[test] fn gemini_streaming_usage_total_falls_back_to_prompt_plus_completion() { let parsed = GeminiParsedEvent { diff --git a/frontend/src/lib/components/copilot/chat/tokenUsage.ts b/frontend/src/lib/components/copilot/chat/tokenUsage.ts index b568dc649d..06ef5e71d8 100644 --- a/frontend/src/lib/components/copilot/chat/tokenUsage.ts +++ b/frontend/src/lib/components/copilot/chat/tokenUsage.ts @@ -158,8 +158,10 @@ export function openAICompletionsUsageToChatTokenUsage( total_tokens?: number | null prompt_tokens_details?: { cached_tokens?: number | null - /** OpenRouter surfaces Anthropic's cache creation here; OpenAI, whose - * caching is automatic and unbilled, reports no such field. */ + /** Cache creation, reported by the providers that bill for it: OpenRouter + * passes Anthropic's through, and the Bedrock proxy folds + * `cacheWriteInputTokens` in here. OpenAI, whose caching is automatic and + * unbilled, reports no such field. */ cache_write_tokens?: number | null } | null /** OpenRouter reports what it actually charged when the request opts in. */