prompt_tokens support in message [anthropic vertex ai] #1187
Conversation
|
DescriptionSummary By MatterAI
🔄 What ChangedThis Pull Request enhances the token usage tracking for Google Vertex AI streaming responses. It introduces a mechanism to store 🔍 Impact of the ChangeThis change significantly improves the accuracy and completeness of token usage reporting in streaming scenarios, especially for Anthropic models on Vertex AI. Previously, 📁 Total Files Changed1 file changed ( 🧪 Test AddedNo specific tests were mentioned in the PR description. It is assumed that manual testing was performed to verify the token usage tracking. 🔒Security VulnerabilitiesNo security vulnerabilities were detected in the provided code changes. MotivationThis PR improves token usage tracking in Google Vertex AI streaming responses by storing prompt token information in the stream state and including token usage metrics in response chunks. Type of Change
How Has This Been Tested?
Screenshots (if applicable)N/A Checklist
Related Issues. Tip Quality Recommendations
Sequence DiagramsequenceDiagram
participant SI as Stream Input
participant VT as VertexAnthropicChatCompleteStreamChunkTransform
participant SS as Stream State
participant TO as Transformed Output
SI->>VT: parsedChunk (raw stream data)
activate VT
alt parsedChunk.type === 'message_start'
VT->>VT: Extract input_tokens from parsedChunk.message.usage
VT->>SS: Store prompt_tokens (parsedChunk.message.usage.input_tokens)
VT->>TO: Send JSON (id, type, message, usage: {prompt_tokens: parsedChunk.message.usage.input_tokens})
else parsedChunk.type === 'message_delta'
VT->>SS: Retrieve prompt_tokens
VT->>TO: Send JSON (id, type, delta, usage: {prompt_tokens: streamState.usage.prompt_tokens})
else parsedChunk.type === 'message_stop'
VT->>SS: Retrieve prompt_tokens
VT->>VT: Extract output_tokens from parsedChunk.usage
VT->>VT: Calculate total_tokens (prompt_tokens + output_tokens)
VT->>TO: Send JSON (id, type, usage: {completion_tokens, prompt_tokens, total_tokens})
end
deactivate VT
|
![matter-code-review[bot]](https://avatars.githubusercontent.com/in/474549?s=60&v=4){kind=small}
arturfromtabnine
changed the title
|
Important PR Review SkippedPR review skipped as per the configuration setting. Run a manually review by commenting /matter review 💡Tips to use Matter AICommand List
|
Description
This PR improves token usage tracking in Google Vertex AI streaming responses by storing prompt token information in the stream state and including token usage metrics in response chunks
Motivation
The current implementation doesn't properly maintain token usage information across streaming chunks, which leads missing token metrics in the chunk response. This change ensures that prompt tokens are properly stored in the stream state and that total token counts are accurately calculated and included in the response
Type of Change
How Has This Been Tested?
Screenshots (if applicable)
Checklist
Related Issues
Before vs after