Repository navigation
Conversation
|
123XieC please read the following Contributor License Agreement(CLA). If you agree with the CLA, please reply with the following information.
Contributor License AgreementContribution License AgreementThis Contribution License Agreement (“Agreement”) is agreed to by the party signing below (“You”),
|
There was a problem hiding this comment.
🟢 Approval recommended
No active review issues were identified in the Bedrock and Claude cache-token accounting changes.
No open issues 🎉
What changed in this PR
Extends the cache-token accounting fix from #9859 to the Bedrock and Claude connectors, addressing #9988 and aligning input usage with other providers.
Changes:
- Include cache read and write/creation tokens in both connectors’
input_token_count, preserving separate cache counters. - Compute Claude’s
total_token_countfrom the inclusive input count plus output tokens. - Update existing cache-token tests to assert the corrected counts.
| File | Description |
|---|---|
python/packages/claude/tests/test_claude_agent.py |
Updates expected input and total counts to include cached tokens. |
python/packages/claude/agent_framework_claude/_agent.py |
Includes cache creation and read tokens in input and total usage. |
python/packages/bedrock/tests/test_bedrock_client.py |
Updates the expected input count to include cached tokens. |
python/packages/bedrock/agent_framework_bedrock/_chat_client.py |
Includes cache read and write tokens in input usage. |
Review effort: Balanced
Motivation & Context
Same root cause as #9808 (AnthropicClient, fixed by #9859), in two more connectors that #9859 did not touch.
BedrockChatClient._parse_usagemaps ConverseinputTokenstoinput_token_count. With prompt caching, Bedrock reports cache reads/writes separately incacheReadInputTokens/cacheWriteInputTokens, andinputTokenscovers only the uncached part. So on main the usage is internally inconsistent (input 12 + output 40 but total 9352).ClaudeAgentbuildsinput_token_countfrom the Claude Agent SDKResultMessage.usage["input_tokens"], which is Anthropic's uncached-only field, and computestotal_token_count = input + output, so both undercount.input_token_countis total input for the OpenAI and Gemini clients and is exported asgen_ai.usage.input_tokens, which the OTel GenAI semconv says should include cached tokens.Description & Review Guide
_parse_usage): sumcacheReadInputTokens+cacheWriteInputTokensintoinput_token_count. Cache-specific fields are unchanged.ClaudeAgent): sumcache_creation_input_tokens+cache_read_input_tokensintoinput_token_count, and derivetotal_token_countfrom the inclusive input total. Cache-specific fields are unchanged.test_parse_usage_surfaces_cache_tokens:input_token_count10 → 21.test_run_captures_result_message_usage_and_finish_reason:input_token_count42 → 50,total_token_count60 → 68.No new tests were added because the existing tests already cover the cache-token paths and their updated assertions verify the fix.
Verification
python -m pytest python/packages/bedrock/tests/test_bedrock_client.py -q→ all pass, including the updatedtest_parse_usage_surfaces_cache_tokens.python -m pytest python/packages/claude/tests/test_claude_agent.py -q→ all pass, including the updatedtest_run_captures_result_message_usage_and_finish_reason.Related Issue
Fixes #9988
AI Assistance
AI assistance details: locating the two connectors' usage parsing, and drafting the minimal change following the already-merged #9859 pattern. I reproduced the undercount locally and ran the updated tests myself.