From 6da29d14df88041cb8e54e1a7be7e07de1d3f346 Mon Sep 17 00:00:00 2001 From: King Star Date: Tue, 8 Sep 2026 22:49:40 +0800 Subject: [PATCH] fix(anthropic): emit prompt cache usage attributes --- .../providers/anthropic/attributes/message.py | 69 +++++++++---------- .../providers/anthropic/stream_wrapper.py | 27 ++------ .../anthropic/test_attributes.py | 30 ++++++++ .../anthropic/test_stream_wrapper.py | 42 ++++++++++- 4 files changed, 108 insertions(+), 60 deletions(-) diff --git a/agentops/instrumentation/providers/anthropic/attributes/message.py b/agentops/instrumentation/providers/anthropic/attributes/message.py index 4346a2673..4cfaef7d1 100644 --- a/agentops/instrumentation/providers/anthropic/attributes/message.py +++ b/agentops/instrumentation/providers/anthropic/attributes/message.py @@ -26,6 +26,35 @@ ) +def get_usage_attributes(usage: Any) -> AttributeMap: + """Extract Anthropic token usage, including prompt-cache breakdowns. + + Anthropic reports cache read and cache creation tokens separately from + native ``input_tokens``. Keep the existing native total unchanged and + expose those fields independently so consumers can choose their accounting + policy without double-counting cache tokens. + """ + attributes = {} + + usage_fields = { + SpanAttributes.LLM_USAGE_PROMPT_TOKENS: "input_tokens", + SpanAttributes.LLM_USAGE_COMPLETION_TOKENS: "output_tokens", + SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS: "cache_read_input_tokens", + SpanAttributes.LLM_USAGE_CACHE_CREATION_INPUT_TOKENS: "cache_creation_input_tokens", + } + for attribute, field in usage_fields.items(): + value = getattr(usage, field, None) + if value is not None: + attributes[attribute] = value + + input_tokens = getattr(usage, "input_tokens", None) + output_tokens = getattr(usage, "output_tokens", None) + if input_tokens is not None and output_tokens is not None: + attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = input_tokens + output_tokens + + return attributes + + def get_message_attributes( args: Optional[Tuple] = None, kwargs: Optional[Dict] = None, return_value: Any = None ) -> AttributeMap: @@ -318,18 +347,7 @@ def get_message_response_attributes(response: "Message") -> AttributeMap: # Extract usage information if hasattr(response, "usage"): - usage = response.usage - if hasattr(usage, "input_tokens"): - input_tokens = usage.input_tokens - attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] = input_tokens - - if hasattr(usage, "output_tokens"): - output_tokens = usage.output_tokens - attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] = output_tokens - - if hasattr(usage, "input_tokens") and hasattr(usage, "output_tokens"): - total_tokens = usage.input_tokens + usage.output_tokens - attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = total_tokens + attributes.update(get_usage_attributes(response.usage)) # Extract stop reason if available if hasattr(response, "stop_reason"): @@ -435,19 +453,7 @@ def get_completion_response_attributes(response: "Completion") -> AttributeMap: # Extract usage information (newer versions have this) if hasattr(response, "usage"): - usage = response.usage - if hasattr(usage, "input_tokens"): - input_tokens = usage.input_tokens - attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] = input_tokens - - if hasattr(usage, "output_tokens"): - output_tokens = usage.output_tokens - attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] = output_tokens - - # Calculate total tokens if we have both input and output - if hasattr(usage, "input_tokens") and hasattr(usage, "output_tokens"): - total_tokens = usage.input_tokens + usage.output_tokens - attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = total_tokens + attributes.update(get_usage_attributes(response.usage)) return attributes @@ -518,18 +524,7 @@ def get_stream_event_attributes(event: Any) -> AttributeMap: elif event_type == "RawMessageStartEvent": if hasattr(event, "message"): if hasattr(event.message, "usage"): - usage = event.message.usage - if hasattr(usage, "input_tokens"): - input_tokens = usage.input_tokens - attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] = input_tokens - - if hasattr(usage, "output_tokens"): - output_tokens = usage.output_tokens - attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] = output_tokens - - if hasattr(usage, "input_tokens") and hasattr(usage, "output_tokens"): - total_tokens = usage.input_tokens + usage.output_tokens - attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] = total_tokens + attributes.update(get_usage_attributes(event.message.usage)) elif event_type == "RawMessageDeltaEvent": if hasattr(event, "delta"): diff --git a/agentops/instrumentation/providers/anthropic/stream_wrapper.py b/agentops/instrumentation/providers/anthropic/stream_wrapper.py index c42084c19..f4262c014 100644 --- a/agentops/instrumentation/providers/anthropic/stream_wrapper.py +++ b/agentops/instrumentation/providers/anthropic/stream_wrapper.py @@ -17,6 +17,7 @@ from agentops.instrumentation.providers.anthropic.attributes.message import ( get_message_request_attributes, get_stream_attributes, + get_usage_attributes, ) from agentops.instrumentation.providers.anthropic.event_handler_wrapper import EventHandleWrapper @@ -165,16 +166,8 @@ def __exit__(self, exc_type, exc_val, exc_tb): span.set_attribute(MessageAttributes.COMPLETION_CONTENT.format(i=0), content_text) if hasattr(final_message, "usage"): - usage = final_message.usage - if hasattr(usage, "input_tokens"): - span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, usage.input_tokens) - - if hasattr(usage, "output_tokens"): - span.set_attribute(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, usage.output_tokens) - - if hasattr(usage, "input_tokens") and hasattr(usage, "output_tokens"): - total_tokens = usage.input_tokens + usage.output_tokens - span.set_attribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, total_tokens) + for key, value in get_usage_attributes(final_message.usage).items(): + span.set_attribute(key, value) except Exception as e: logger.debug(f"Failed to extract final message data: {e}") finally: @@ -404,18 +397,8 @@ async def __aexit__(self, exc_type, exc_val, exc_tb): ) if hasattr(final_message, "usage"): - usage = final_message.usage - if hasattr(usage, "input_tokens"): - span.set_attribute(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, usage.input_tokens) - - if hasattr(usage, "output_tokens"): - span.set_attribute( - SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, usage.output_tokens - ) - - if hasattr(usage, "input_tokens") and hasattr(usage, "output_tokens"): - total_tokens = usage.input_tokens + usage.output_tokens - span.set_attribute(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, total_tokens) + for key, value in get_usage_attributes(final_message.usage).items(): + span.set_attribute(key, value) except Exception as e: logger.debug(f"Failed to extract final async message data: {e}") finally: diff --git a/tests/unit/instrumentation/anthropic/test_attributes.py b/tests/unit/instrumentation/anthropic/test_attributes.py index a7d634d9c..d14991fc7 100644 --- a/tests/unit/instrumentation/anthropic/test_attributes.py +++ b/tests/unit/instrumentation/anthropic/test_attributes.py @@ -1,5 +1,7 @@ """Tests for Anthropic attribute extraction functionality.""" +from types import SimpleNamespace + from agentops.semconv import ( InstrumentationAttributes, SpanAttributes, @@ -12,6 +14,7 @@ extract_request_attributes, ) from agentops.instrumentation.providers.anthropic.attributes.message import ( + get_message_response_attributes, get_message_request_attributes, get_stream_attributes, get_stream_event_attributes, @@ -72,6 +75,29 @@ def test_get_message_request_attributes(): assert MessageAttributes.PROMPT_CONTENT.format(i=1) in attributes +def test_get_message_response_attributes_exposes_cache_usage_separately(): + """Anthropic cache usage is reported without changing native token totals.""" + response = SimpleNamespace( + id="msg_cache", + model="claude-sonnet", + content=[], + usage=SimpleNamespace( + input_tokens=100, + output_tokens=20, + cache_read_input_tokens=500, + cache_creation_input_tokens=50, + ), + ) + + attributes = get_message_response_attributes(response) + + assert attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 100 + assert attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 20 + assert attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 120 + assert attributes[SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS] == 500 + assert attributes[SpanAttributes.LLM_USAGE_CACHE_CREATION_INPUT_TOKENS] == 50 + + # Stream Attributes Tests def test_get_stream_attributes(): """Test extraction of stream attributes.""" @@ -108,6 +134,8 @@ class MockUsage: def __init__(self): self.input_tokens = 10 self.output_tokens = 5 + self.cache_read_input_tokens = 40 + self.cache_creation_input_tokens = 30 class MockMessage: def __init__(self): @@ -124,6 +152,8 @@ def __init__(self): assert attributes[SpanAttributes.LLM_USAGE_PROMPT_TOKENS] == 10 assert attributes[SpanAttributes.LLM_USAGE_COMPLETION_TOKENS] == 5 assert attributes[SpanAttributes.LLM_USAGE_TOTAL_TOKENS] == 15 + assert attributes[SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS] == 40 + assert attributes[SpanAttributes.LLM_USAGE_CACHE_CREATION_INPUT_TOKENS] == 30 def test_get_stream_event_attributes_raw_message_delta(): diff --git a/tests/unit/instrumentation/anthropic/test_stream_wrapper.py b/tests/unit/instrumentation/anthropic/test_stream_wrapper.py index c20028234..4367e5499 100644 --- a/tests/unit/instrumentation/anthropic/test_stream_wrapper.py +++ b/tests/unit/instrumentation/anthropic/test_stream_wrapper.py @@ -1,4 +1,5 @@ import pytest +from types import SimpleNamespace from unittest.mock import MagicMock from opentelemetry.trace import SpanKind @@ -108,7 +109,12 @@ def test_stream_final_message_attributes(mock_tracer, mock_stream_manager): final_message = MagicMock() final_message.content = [MagicMock(text="Final response")] - final_message.usage = MagicMock(input_tokens=10, output_tokens=20) + final_message.usage = SimpleNamespace( + input_tokens=10, + output_tokens=20, + cache_read_input_tokens=40, + cache_creation_input_tokens=30, + ) mock_stream_manager._MessageStreamManager__stream._MessageStream__final_message_snapshot = final_message result = wrapper(wrapped, None, [], {}) @@ -123,3 +129,37 @@ def test_stream_final_message_attributes(mock_tracer, mock_stream_manager): span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, 10) span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, 20) span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, 30) + span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS, 40) + span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_CACHE_CREATION_INPUT_TOKENS, 30) + + +@pytest.mark.asyncio +async def test_async_stream_final_message_attributes_include_cache_usage(mock_tracer, mock_async_stream_manager): + """Async stream final usage includes both Anthropic cache token fields.""" + wrapper = messages_stream_async_wrapper(mock_tracer) + wrapped = MagicMock(return_value=mock_async_stream_manager) + final_message = SimpleNamespace( + content=[SimpleNamespace(text="Final response")], + usage=SimpleNamespace( + input_tokens=10, + output_tokens=20, + cache_read_input_tokens=40, + cache_creation_input_tokens=30, + ), + ) + mock_async_stream_manager._AsyncMessageStreamManager__stream._AsyncMessageStream__final_message_snapshot = ( + final_message + ) + + result = wrapper(wrapped, None, [], {}) + + async with result as stream: + async for _ in stream.text_stream: + pass + + span = mock_tracer.start_span.return_value + span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_PROMPT_TOKENS, 10) + span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_COMPLETION_TOKENS, 20) + span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_TOTAL_TOKENS, 30) + span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_CACHE_READ_INPUT_TOKENS, 40) + span.set_attribute.assert_any_call(SpanAttributes.LLM_USAGE_CACHE_CREATION_INPUT_TOKENS, 30)