From 685f0cb20316581a065443be7047bdbc03f21618 Mon Sep 17 00:00:00 2001 From: 123XieC <123XieC@users.noreply.github.com> Date: Tue, 3 Nov 2026 11:14:49 +0800 Subject: [PATCH] fix(python): include cache tokens in input_token_count for Bedrock and Claude --- .../agent_framework_bedrock/_chat_client.py | 28 +++++++++++++------ .../bedrock/tests/test_bedrock_client.py | 3 +- .../claude/agent_framework_claude/_agent.py | 23 +++++++++++---- .../claude/tests/test_claude_agent.py | 6 ++-- 4 files changed, 44 insertions(+), 16 deletions(-) diff --git a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py index 16a49b41cb0..5865b8b0b8d 100644 --- a/python/packages/bedrock/agent_framework_bedrock/_chat_client.py +++ b/python/packages/bedrock/agent_framework_bedrock/_chat_client.py @@ -868,17 +868,29 @@ def _parse_usage(self, usage: dict[str, Any] | None) -> UsageDetails | None: if not usage: return None details: UsageDetails = {} - if (input_tokens := usage.get("inputTokens")) is not None: - details["input_token_count"] = input_tokens - if (output_tokens := usage.get("outputTokens")) is not None: - details["output_token_count"] = output_tokens - if (total_tokens := usage.get("totalTokens")) is not None: - details["total_token_count"] = total_tokens + input_tokens = usage.get("inputTokens") + output_tokens = usage.get("outputTokens") + cache_read = usage.get("cacheReadInputTokens") + cache_write = usage.get("cacheWriteInputTokens") # Bedrock Converse reports these when prompt caching is active. - if (cache_read := usage.get("cacheReadInputTokens")) is not None: + if cache_read is not None: details["cache_read_input_token_count"] = cache_read - if (cache_write := usage.get("cacheWriteInputTokens")) is not None: + if cache_write is not None: details["cache_creation_input_token_count"] = cache_write + # inputTokens covers only the uncached part when prompt caching is active; + # total input (and gen_ai.usage.input_tokens) must include cache tokens, + # matching other connectors and the GenAI semconv. + input_total = input_tokens + if cache_read is not None: + input_total = cache_read if input_total is None else input_total + cache_read + if cache_write is not None: + input_total = cache_write if input_total is None else input_total + cache_write + if input_total is not None: + details["input_token_count"] = input_total + if output_tokens is not None: + details["output_token_count"] = output_tokens + if (total_tokens := usage.get("totalTokens")) is not None: + details["total_token_count"] = total_tokens return details or None def _parse_message_contents(self, content_blocks: Sequence[dict[str, Any]]) -> list[Any]: diff --git a/python/packages/bedrock/tests/test_bedrock_client.py b/python/packages/bedrock/tests/test_bedrock_client.py index 9f753e8835d..e05128c151d 100644 --- a/python/packages/bedrock/tests/test_bedrock_client.py +++ b/python/packages/bedrock/tests/test_bedrock_client.py @@ -620,7 +620,8 @@ def test_parse_usage_surfaces_cache_tokens() -> None: }) assert details is not None - assert details["input_token_count"] == 10 + # inputTokens (10) + cacheReadInputTokens (8) + cacheWriteInputTokens (3) = 21 + assert details["input_token_count"] == 21 assert details["cache_read_input_token_count"] == 8 assert details["cache_creation_input_token_count"] == 3 diff --git a/python/packages/claude/agent_framework_claude/_agent.py b/python/packages/claude/agent_framework_claude/_agent.py index bed9ba7f0f2..036cb83a7d3 100644 --- a/python/packages/claude/agent_framework_claude/_agent.py +++ b/python/packages/claude/agent_framework_claude/_agent.py @@ -988,19 +988,32 @@ async def _stream_run( usage = message.usage or {} input_tokens = usage.get("input_tokens") output_tokens = usage.get("output_tokens") + cache_creation_tokens = usage.get("cache_creation_input_tokens") + cache_read_tokens = usage.get("cache_read_input_tokens") + # Anthropic's input_tokens excludes cache creation/read; total input + # (and gen_ai.usage.input_tokens) must include them, matching other + # connectors and the GenAI semconv. + input_total = input_tokens + for cache_tokens in (cache_creation_tokens, cache_read_tokens): + if isinstance(cache_tokens, int): + input_total = ( + cache_tokens + if input_total is None + else input_total + cache_tokens + ) total_token_count = ( - input_tokens + output_tokens - if isinstance(input_tokens, int) and isinstance(output_tokens, int) + input_total + output_tokens + if isinstance(input_total, int) and isinstance(output_tokens, int) else None ) usage_details = UsageDetails(**{ key: value for key, value in { - "input_token_count": input_tokens, + "input_token_count": input_total, "output_token_count": output_tokens, "total_token_count": total_token_count, - "cache_creation_input_token_count": usage.get("cache_creation_input_tokens"), - "cache_read_input_token_count": usage.get("cache_read_input_tokens"), + "cache_creation_input_token_count": cache_creation_tokens, + "cache_read_input_token_count": cache_read_tokens, }.items() if isinstance(value, int) }) diff --git a/python/packages/claude/tests/test_claude_agent.py b/python/packages/claude/tests/test_claude_agent.py index 51c5b5b5aed..cc876f0b269 100644 --- a/python/packages/claude/tests/test_claude_agent.py +++ b/python/packages/claude/tests/test_claude_agent.py @@ -353,9 +353,11 @@ async def test_run_captures_result_message_usage_and_finish_reason( assert response.finish_reason == expected_finish_reason assert response.usage_details == { - "input_token_count": 42, + # input_tokens (42) + cache_creation (3) + cache_read (5) = 50 + "input_token_count": 50, "output_token_count": 18, - "total_token_count": 60, + # 50 + 18 = 68 + "total_token_count": 68, "cache_creation_input_token_count": 3, "cache_read_input_token_count": 5, }