diff --git a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py index 8b6f49e821..699ba0e722 100644 --- a/python/packages/anthropic/agent_framework_anthropic/_chat_client.py +++ b/python/packages/anthropic/agent_framework_anthropic/_chat_client.py @@ -673,7 +673,10 @@ async def _stream_turn( # Anthropic streams cumulative usage snapshots (message_start seeds it, # each message_delta carries the running total), so thread a per-stream # accumulator to _process_stream_event to emit increments instead. + # raw_usage_carry keeps last-seen input/cache components when a later + # snapshot omits them (cache_* often only appear on message_start). emitted_usage: dict[str, int] = {} + raw_usage_carry: dict[str, int] = {} local_function_call_started = False open_local_function_call_blocks: set[int] = set() message_stop_received = False @@ -683,6 +686,7 @@ async def _stream_turn( chunk, emitted_usage, request_state=turn_request_state, + raw_usage_carry=raw_usage_carry, ) if chunk.type == "content_block_start" and parsed_chunk is not None: if any( @@ -1487,6 +1491,7 @@ def _process_stream_event( event: BetaRawMessageStreamEvent, emitted_usage: dict[str, int] | None = None, request_state: _AnthropicRequestState | None = None, + raw_usage_carry: dict[str, int] | None = None, ) -> ChatResponseUpdate | None: """Process a streaming event from the Anthropic client. @@ -1497,6 +1502,10 @@ def _process_stream_event( increments (see ``_incremental_usage``). Pass ``None`` for a one-off event to keep the snapshot unchanged. request_state: State used to parse this stream. + raw_usage_carry: Per-stream last-seen Anthropic input/cache token + components. Streaming deltas often omit ``cache_*`` fields that were + only present on ``message_start``; carrying them forward keeps + ``input_token_count`` inclusive of cache tokens across the stream. Returns: A ChatResponseUpdate object containing the processed update. @@ -1504,7 +1513,11 @@ def _process_stream_event( match event.type: case "message_start": usage_details: list[Content] = [] - if event.message.usage and (details := self._parse_usage_from_anthropic(event.message.usage)): + if event.message.usage and ( + details := self._parse_usage_from_anthropic( + event.message.usage, raw_usage_carry=raw_usage_carry + ) + ): usage_details.append( Content.from_usage( usage_details=self._incremental_usage(details, emitted_usage), @@ -1529,7 +1542,7 @@ def _process_stream_event( raw_representation=event, ) case "message_delta": - usage = self._parse_usage_from_anthropic(event.usage) + usage = self._parse_usage_from_anthropic(event.usage, raw_usage_carry=raw_usage_carry) return ChatResponseUpdate( contents=[ Content.from_usage( @@ -1571,19 +1584,60 @@ def _process_stream_event( logger.debug(f"Ignoring unsupported event type: {event.type}") return None - def _parse_usage_from_anthropic(self, usage: BetaUsage | BetaMessageDeltaUsage | None) -> UsageDetails | None: - """Parse usage details from the Anthropic message usage.""" + def _parse_usage_from_anthropic( + self, + usage: BetaUsage | BetaMessageDeltaUsage | None, + *, + raw_usage_carry: dict[str, int] | None = None, + ) -> UsageDetails | None: + """Parse usage details from the Anthropic message usage. + + When ``raw_usage_carry`` is provided (streaming), last-seen ``input_tokens`` / + ``cache_*`` components are reused if the current snapshot omits them. Anthropic + often reports cache tokens only on ``message_start``; without the carry, + a later ``message_delta`` with a higher ``input_tokens`` and omitted cache + fields would under-count ``input_token_count`` after incremental conversion. + """ if not usage: return None usage_details = UsageDetails(output_token_count=usage.output_tokens) - if usage.input_tokens is not None: - usage_details["input_token_count"] = usage.input_tokens - if usage.cache_creation_input_tokens is not None: - usage_details["anthropic.cache_creation_input_tokens"] = usage.cache_creation_input_tokens - usage_details["cache_creation_input_token_count"] = usage.cache_creation_input_tokens - if usage.cache_read_input_tokens is not None: - usage_details["anthropic.cache_read_input_tokens"] = usage.cache_read_input_tokens - usage_details["cache_read_input_token_count"] = usage.cache_read_input_tokens + # Anthropic's input_tokens excludes cache_creation and cache_read; total + # input (and gen_ai.usage.input_tokens) must include them, matching other + # connectors and the GenAI semconv. + input_tokens = usage.input_tokens + cache_creation = usage.cache_creation_input_tokens + cache_read = usage.cache_read_input_tokens + if raw_usage_carry is not None: + if input_tokens is not None: + raw_usage_carry["input_tokens"] = input_tokens + else: + input_tokens = raw_usage_carry.get("input_tokens") + if cache_creation is not None: + raw_usage_carry["cache_creation_input_tokens"] = cache_creation + else: + cache_creation = raw_usage_carry.get("cache_creation_input_tokens") + if cache_read is not None: + raw_usage_carry["cache_read_input_tokens"] = cache_read + else: + cache_read = raw_usage_carry.get("cache_read_input_tokens") + + input_total = input_tokens if input_tokens is not None else None + if cache_creation is not None: + usage_details["anthropic.cache_creation_input_tokens"] = cache_creation + usage_details["cache_creation_input_token_count"] = cache_creation + if input_total is not None: + input_total += cache_creation + else: + input_total = cache_creation + if cache_read is not None: + usage_details["anthropic.cache_read_input_tokens"] = cache_read + usage_details["cache_read_input_token_count"] = cache_read + if input_total is not None: + input_total += cache_read + else: + input_total = cache_read + if input_total is not None: + usage_details["input_token_count"] = input_total return usage_details @staticmethod diff --git a/python/packages/anthropic/tests/test_anthropic_client.py b/python/packages/anthropic/tests/test_anthropic_client.py index 2a864e32ef..3a15b7ca93 100644 --- a/python/packages/anthropic/tests/test_anthropic_client.py +++ b/python/packages/anthropic/tests/test_anthropic_client.py @@ -3203,7 +3203,13 @@ def test_process_stream_event_message_start_sets_assistant_role(mock_anthropic_c assert result.role == "assistant" -def _usage_message_start_event(*, input_tokens: int, output_tokens: int) -> MagicMock: +def _usage_message_start_event( + *, + input_tokens: int, + output_tokens: int, + cache_creation_input_tokens: int | None = None, + cache_read_input_tokens: int | None = None, +) -> MagicMock: event = MagicMock() event.type = "message_start" event.message.id = "msg_usage" @@ -3211,7 +3217,12 @@ def _usage_message_start_event(*, input_tokens: int, output_tokens: int) -> Magi event.message.model = "claude-3-5-sonnet-20241022" event.message.content = [] event.message.stop_reason = None - event.message.usage = BetaUsage(input_tokens=input_tokens, output_tokens=output_tokens) + event.message.usage = BetaUsage( + input_tokens=input_tokens, + output_tokens=output_tokens, + cache_creation_input_tokens=cache_creation_input_tokens, + cache_read_input_tokens=cache_read_input_tokens, + ) return event @@ -3282,6 +3293,48 @@ def test_streaming_usage_delta_input_not_double_counted(mock_anthropic_client: M assert response.usage_details["input_token_count"] == 12 +def test_streaming_usage_preserves_cache_tokens_across_delta(mock_anthropic_client: MagicMock) -> None: + """message_delta omitting cache_* must keep cache tokens in input_token_count. + + Anthropic often reports cache_creation/cache_read only on message_start. A later + message_delta with a higher input_tokens and omitted cache fields must still include + the last-seen cache components so incremental conversion reconstructs + input + cache_creation + cache_read (e.g. 120+20+30=170), not bare input (120). + """ + client = create_test_anthropic_client(mock_anthropic_client) + emitted: dict[str, int] = {} + raw_carry: dict[str, int] = {} + updates = [ + u + for u in ( + client._process_stream_event( + _usage_message_start_event( + input_tokens=100, + output_tokens=1, + cache_creation_input_tokens=20, + cache_read_input_tokens=30, + ), + emitted, + raw_usage_carry=raw_carry, + ), + client._process_stream_event( + _usage_message_delta_event(input_tokens=120, output_tokens=25), + emitted, + raw_usage_carry=raw_carry, + ), + ) + if u is not None + ] + + response = ChatResponse.from_updates(updates) + + assert response.usage_details is not None + assert response.usage_details["output_token_count"] == 25 + assert response.usage_details["input_token_count"] == 170 + assert response.usage_details["cache_creation_input_token_count"] == 20 + assert response.usage_details["cache_read_input_token_count"] == 30 + + def test_process_stream_event_message_start_role_prevents_tool_use_collapse() -> None: """Regression test: tool_use blocks must not end up in a user-role message. @@ -4886,7 +4939,8 @@ def test_parse_usage_with_cache_tokens(mock_anthropic_client: MagicMock) -> None assert result is not None result_dict = cast("dict[str, Any]", result) assert result["output_token_count"] == 50 - assert result["input_token_count"] == 100 + # 100 input + 20 cache_creation + 30 cache_read + assert result["input_token_count"] == 150 assert result_dict["anthropic.cache_creation_input_tokens"] == 20 assert result_dict["anthropic.cache_read_input_tokens"] == 30 assert result["cache_creation_input_token_count"] == 20