improvement: add cache details to usage for (vertex ai) gemini and anthropic models #1373
Conversation
|
DescriptionSummary By MatterAI
🔄 What ChangedThis Pull Request introduces a new 🔍 Impact of the ChangeThis enhancement provides more granular visibility into token consumption by explicitly reporting cached token usage. This improves cost analysis, observability, and understanding of how caching mechanisms affect token usage for these AI models. 📁 Total Files Changed
🧪 Test AddedN/A 🔒Security VulnerabilitiesN/A MotivationTo provide comprehensive and granular token usage information, including cached tokens, for Vertex AI Gemini and Anthropic models, thereby improving cost analysis, observability, and understanding of model interactions. Type of Change
How Has This Been Tested?
Screenshots (if applicable)N/A Checklist
Related IssuesN/A Tip Quality Recommendations
Tanka Poem ♫
Sequence DiagramsequenceDiagram
participant Client
participant AnthropicChatComplete as Anthropic Chat Complete
participant GoogleVertexAIChatComplete as Google Vertex AI Chat Complete
participant AnthropicAPI as Anthropic API
participant GoogleVertexAIApi as Google Vertex AI API
Client->>AnthropicChatComplete: chatComplete(request)
AnthropicChatComplete->>AnthropicAPI: POST /messages (payload)
AnthropicAPI-->>AnthropicChatComplete: API Response (usage, output_tokens, cache_read_input_tokens)
Note over AnthropicChatComplete: Calculate total_tokens, prompt_tokens, completion_tokens, and populate usage.prompt_tokens_details.cached_tokens
AnthropicChatComplete-->>Client: Chat Completion Response (with usage.prompt_tokens_details.cached_tokens)
Client->>GoogleVertexAIChatComplete: chatComplete(request)
GoogleVertexAIChatComplete->>GoogleVertexAIApi: POST /generateContent (payload)
GoogleVertexAIApi-->>GoogleVertexAIChatComplete: API Response (usage, output_tokens, cache_read_input_tokens)
Note over GoogleVertexAIChatComplete: Calculate total_tokens, prompt_tokens, completion_tokens, and populate usage.prompt_tokens_details.cached_tokens
GoogleVertexAIChatComplete-->>Client: Chat Completion Response (with usage.prompt_tokens_details.cached_tokens)
Note over AnthropicChatComplete,GoogleVertexAIChatComplete: Also applies to streaming responses, where usage is updated per chunk.
|
…th vertex and gemini
![matter-code-review[bot]](https://avatars.githubusercontent.com/in/474549?s=60&v=4){kind=small}
| audio_tokens: | ||
| parsedChunk.usageMetadata?.candidatesTokensDetails?.reduce( | ||
| (acc, curr) => { | ||
| if (curr.modality === VERTEX_MODALITY.AUDIO) | ||
| return acc + curr.tokenCount; | ||
| return acc; | ||
| }, | ||
| 0 | ||
| ), |
There was a problem hiding this comment.
🔴 Security Issue
Issue: Potential null dereference when accessing parsedChunk.usageMetadata?.candidatesTokensDetails without validating its existence.
Fix: Add explicit null check before destructuring.
Impact: Prevents runtime errors when usage data is missing.
| audio_tokens: | |
| parsedChunk.usageMetadata?.candidatesTokensDetails?.reduce( | |
| (acc, curr) => { | |
| if (curr.modality === VERTEX_MODALITY.AUDIO) | |
| return acc + curr.tokenCount; | |
| return acc; | |
| }, | |
| 0 | |
| ), | |
| audio_tokens: | |
| (parsedChunk.usageMetadata?.candidatesTokensDetails ?? []).reduce( | |
| (acc, curr) => { | |
| if (curr?.modality === VERTEX_MODALITY.AUDIO) | |
| return acc + (curr?.tokenCount ?? 0); | |
| return acc; | |
| }, | |
| 0 | |
| ), |
|
@miuosz I've replicated your changes in the references:
|
…th vertex and gemini
| audio_tokens: parsedChunk.usageMetadata?.promptTokensDetails?.reduce( | ||
| (acc, curr) => { | ||
| if (curr.modality === VERTEX_MODALITY.AUDIO) | ||
| return acc + curr.tokenCount; | ||
| return acc; | ||
| }, | ||
| 0 | ||
| ), | ||
| }, | ||
| }; | ||
| } |
There was a problem hiding this comment.
🔴 Security Issue
Issue: Potential null dereference when accessing parsedChunk.usageMetadata?.promptTokensDetails without validating its existence.
Fix: Add explicit null check before destructuring.
Impact: Prevents runtime errors when usage data is missing.
| audio_tokens: parsedChunk.usageMetadata?.promptTokensDetails?.reduce( | |
| (acc, curr) => { | |
| if (curr.modality === VERTEX_MODALITY.AUDIO) | |
| return acc + curr.tokenCount; | |
| return acc; | |
| }, | |
| 0 | |
| ), | |
| }, | |
| }; | |
| } | |
| audio_tokens: (parsedChunk.usageMetadata?.promptTokensDetails ?? []).reduce( | |
| ( | |
| acc: number, | |
| curr: { modality: VERTEX_MODALITY; tokenCount: number } | |
| ) => { | |
| if (curr?.modality === VERTEX_MODALITY.AUDIO) | |
| return acc + (curr?.tokenCount ?? 0); | |
| return acc; | |
| }, | |
| 0 | |
| ), |
| ...streamState.usage, | ||
| completion_tokens: parsedChunk.usage?.output_tokens, | ||
| total_tokens: totalTokens, | ||
| prompt_tokens_details: { | ||
| cached_tokens: streamState.usage?.cache_read_input_tokens ?? 0, | ||
| }, | ||
| }, |
There was a problem hiding this comment.
🟡 Code Quality Issue
Issue: Inconsistent handling of cache_read_input_tokens in usage object. In one place it's included in the spread operator, in another it's explicitly set.
Fix: Standardize the usage object construction to avoid redundancy.
Impact: Improves code maintainability and reduces potential confusion.
| ...streamState.usage, | |
| completion_tokens: parsedChunk.usage?.output_tokens, | |
| total_tokens: totalTokens, | |
| prompt_tokens_details: { | |
| cached_tokens: streamState.usage?.cache_read_input_tokens ?? 0, | |
| }, | |
| }, | |
| usage: { | |
| ...streamState.usage, | |
| completion_tokens: parsedChunk.usage?.output_tokens, | |
| total_tokens: totalTokens, | |
| prompt_tokens_details: { | |
| cached_tokens: streamState.usage?.cache_read_input_tokens ?? 0, | |
| }, | |
| ...(shouldSendCacheUsage && { | |
| cache_read_input_tokens: cache_read_input_tokens, | |
| cache_creation_input_tokens: cache_creation_input_tokens, | |
| }), | |
| }, |
| ...streamState.usage, | ||
| completion_tokens: parsedChunk.usage?.output_tokens, | ||
| prompt_tokens: streamState.usage?.prompt_tokens, | ||
| total_tokens: | ||
| (streamState.usage?.prompt_tokens || 0) + | ||
| (parsedChunk.usage?.output_tokens || 0), | ||
| total_tokens: totalTokens, | ||
| prompt_tokens_details: { | ||
| cached_tokens: streamState.usage?.cache_read_input_tokens ?? 0, | ||
| }, | ||
| }, |
There was a problem hiding this comment.
🟡 Code Quality Issue
Issue: Inconsistent handling of cache_read_input_tokens in usage object. In one place it's included in the spread operator, in another it's explicitly set.
Fix: Standardize the usage object construction to avoid redundancy.
Impact: Improves code maintainability and reduces potential confusion.
| ...streamState.usage, | |
| completion_tokens: parsedChunk.usage?.output_tokens, | |
| prompt_tokens: streamState.usage?.prompt_tokens, | |
| total_tokens: | |
| (streamState.usage?.prompt_tokens || 0) + | |
| (parsedChunk.usage?.output_tokens || 0), | |
| total_tokens: totalTokens, | |
| prompt_tokens_details: { | |
| cached_tokens: streamState.usage?.cache_read_input_tokens ?? 0, | |
| }, | |
| }, | |
| usage: { | |
| ...streamState.usage, | |
| completion_tokens: parsedChunk.usage?.output_tokens, | |
| total_tokens: totalTokens, | |
| prompt_tokens_details: { | |
| cached_tokens: streamState.usage?.cache_read_input_tokens ?? 0, | |
| }, | |
| ...(shouldSendCacheUsage && { | |
| cache_read_input_tokens: cache_read_input_tokens, | |
| cache_creation_input_tokens: cache_creation_input_tokens, | |
| }), | |
| }, |
2 participants
Description
Motivation
Without proper
usagedetails, it's impossible to accurately count tokens or estimate the cost.1. claude-opus-4
before
after
2. gemini-2.5-flash
before
after
Type of Change
How Has This Been Tested?
Screenshots (if applicable)
Checklist
Related Issues