diff --git a/src/strands/telemetry/tracer.py b/src/strands/telemetry/tracer.py index 6b429393d..9e170571a 100644 --- a/src/strands/telemetry/tracer.py +++ b/src/strands/telemetry/tracer.py @@ -263,6 +263,8 @@ def end_model_invoke_span( "gen_ai.usage.completion_tokens": usage["outputTokens"], "gen_ai.usage.output_tokens": usage["outputTokens"], "gen_ai.usage.total_tokens": usage["totalTokens"], + "gen_ai.usage.cache_read_input_tokens": usage.get("cacheReadInputTokens", 0), + "gen_ai.usage.cache_write_input_tokens": usage.get("cacheWriteInputTokens", 0), } self._add_event( @@ -491,6 +493,8 @@ def end_agent_span( "gen_ai.usage.input_tokens": accumulated_usage["inputTokens"], "gen_ai.usage.output_tokens": accumulated_usage["outputTokens"], "gen_ai.usage.total_tokens": accumulated_usage["totalTokens"], + "gen_ai.usage.cache_read_input_tokens": accumulated_usage.get("cacheReadInputTokens", 0), + "gen_ai.usage.cache_write_input_tokens": accumulated_usage.get("cacheWriteInputTokens", 0), } ) diff --git a/tests/strands/telemetry/test_tracer.py b/tests/strands/telemetry/test_tracer.py index 586911bef..568fff130 100644 --- a/tests/strands/telemetry/test_tracer.py +++ b/tests/strands/telemetry/test_tracer.py @@ -177,6 +177,8 @@ def test_end_model_invoke_span(mock_span): mock_span.set_attribute.assert_any_call("gen_ai.usage.completion_tokens", 20) mock_span.set_attribute.assert_any_call("gen_ai.usage.output_tokens", 20) mock_span.set_attribute.assert_any_call("gen_ai.usage.total_tokens", 30) + mock_span.set_attribute.assert_any_call("gen_ai.usage.cache_read_input_tokens", 0) + mock_span.set_attribute.assert_any_call("gen_ai.usage.cache_write_input_tokens", 0) mock_span.add_event.assert_called_with( "gen_ai.choice", attributes={"message": json.dumps(message["content"]), "finish_reason": "end_turn"}, @@ -404,6 +406,8 @@ def test_end_agent_span(mock_span): mock_span.set_attribute.assert_any_call("gen_ai.usage.completion_tokens", 100) mock_span.set_attribute.assert_any_call("gen_ai.usage.output_tokens", 100) mock_span.set_attribute.assert_any_call("gen_ai.usage.total_tokens", 150) + mock_span.set_attribute.assert_any_call("gen_ai.usage.cache_read_input_tokens", 0) + mock_span.set_attribute.assert_any_call("gen_ai.usage.cache_write_input_tokens", 0) mock_span.add_event.assert_any_call( "gen_ai.choice", attributes={"message": "Agent response", "finish_reason": "end_turn"}, @@ -412,6 +416,64 @@ def test_end_agent_span(mock_span): mock_span.end.assert_called_once() +def test_end_model_invoke_span_with_cache_metrics(mock_span): + """Test ending a model invoke span with cache metrics.""" + tracer = Tracer() + message = {"role": "assistant", "content": [{"text": "Response"}]} + usage = Usage( + inputTokens=10, + outputTokens=20, + totalTokens=30, + cacheReadInputTokens=5, + cacheWriteInputTokens=3, + ) + stop_reason: StopReason = "end_turn" + + tracer.end_model_invoke_span(mock_span, message, usage, stop_reason) + + mock_span.set_attribute.assert_any_call("gen_ai.usage.prompt_tokens", 10) + mock_span.set_attribute.assert_any_call("gen_ai.usage.input_tokens", 10) + mock_span.set_attribute.assert_any_call("gen_ai.usage.completion_tokens", 20) + mock_span.set_attribute.assert_any_call("gen_ai.usage.output_tokens", 20) + mock_span.set_attribute.assert_any_call("gen_ai.usage.total_tokens", 30) + mock_span.set_attribute.assert_any_call("gen_ai.usage.cache_read_input_tokens", 5) + mock_span.set_attribute.assert_any_call("gen_ai.usage.cache_write_input_tokens", 3) + mock_span.set_status.assert_called_once_with(StatusCode.OK) + mock_span.end.assert_called_once() + + +def test_end_agent_span_with_cache_metrics(mock_span): + """Test ending an agent span with cache metrics.""" + tracer = Tracer() + + # Mock AgentResult with metrics including cache tokens + mock_metrics = mock.MagicMock() + mock_metrics.accumulated_usage = { + "inputTokens": 50, + "outputTokens": 100, + "totalTokens": 150, + "cacheReadInputTokens": 25, + "cacheWriteInputTokens": 10, + } + + mock_response = mock.MagicMock() + mock_response.metrics = mock_metrics + mock_response.stop_reason = "end_turn" + mock_response.__str__ = mock.MagicMock(return_value="Agent response") + + tracer.end_agent_span(mock_span, mock_response) + + mock_span.set_attribute.assert_any_call("gen_ai.usage.prompt_tokens", 50) + mock_span.set_attribute.assert_any_call("gen_ai.usage.input_tokens", 50) + mock_span.set_attribute.assert_any_call("gen_ai.usage.completion_tokens", 100) + mock_span.set_attribute.assert_any_call("gen_ai.usage.output_tokens", 100) + mock_span.set_attribute.assert_any_call("gen_ai.usage.total_tokens", 150) + mock_span.set_attribute.assert_any_call("gen_ai.usage.cache_read_input_tokens", 25) + mock_span.set_attribute.assert_any_call("gen_ai.usage.cache_write_input_tokens", 10) + mock_span.set_status.assert_called_once_with(StatusCode.OK) + mock_span.end.assert_called_once() + + def test_get_tracer_singleton(): """Test that get_tracer returns a singleton instance.""" # Reset the singleton first