Updates tests with cache tokens for the supported models

Author: Workshop Participant

tests-integ/test_bedrock_cache_point.py

@@ -16,16 +16,8 @@ def test_bedrock_cache_point():
         {"role": "assistant", "content": [{"text": "Blue!"}]},
     ]
 
-    cache_point_usage = 0
+    agent = Agent(messages=messages, load_tools_from_directory=False)
+    response = agent("What is favorite color?")
 
-    def cache_point_callback_handler(**kwargs):
-        nonlocal cache_point_usage
-        if "event" in kwargs and kwargs["event"] and "metadata" in kwargs["event"] and kwargs["event"]["metadata"]:
-            metadata = kwargs["event"]["metadata"]
-            if "usage" in metadata and metadata["usage"]:
-                if "cacheReadInputTokens" in metadata["usage"] or "cacheWriteInputTokens" in metadata["usage"]:
-                    cache_point_usage += 1
-
-    agent = Agent(messages=messages, callback_handler=cache_point_callback_handler, load_tools_from_directory=False)
-    agent("What is favorite color?")
-    assert cache_point_usage > 0
+    usage = response.metrics.accumulated_usage
+    assert usage["cacheReadInputTokens"] >= 0 or usage["cacheWriteInputTokens"] > 0  # At least one should have tokens

tests/strands/event_loop/test_streaming.py

@@ -250,7 +250,13 @@ def test_handle_message_stop():
 
 def test_extract_usage_metrics():
     event = {
-        "usage": {"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
+        "usage": {
+            "inputTokens": 0,
+            "outputTokens": 0,
+            "totalTokens": 0,
+            "cacheReadInputTokens": 0,
+            "cacheWriteInputTokens": 0,
+        },
         "metrics": {"latencyMs": 0},
     }
 
@@ -279,7 +285,13 @@ def test_extract_usage_metrics():
                 },
                 {
                     "metadata": {
-                        "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 1},
+                        "usage": {
+                            "inputTokens": 1,
+                            "outputTokens": 1,
+                            "totalTokens": 1,
+                            "cacheReadInputTokens": 1,
+                            "cacheWriteInputTokens": 1,
+                        },
                         "metrics": {"latencyMs": 1},
                     }
                 },
@@ -364,6 +376,8 @@ def test_extract_usage_metrics():
                                     "inputTokens": 1,
                                     "outputTokens": 1,
                                     "totalTokens": 1,
+                                    "cacheReadInputTokens": 1,
+                                    "cacheWriteInputTokens": 1,
                                 },
                             },
                         },
@@ -376,7 +390,13 @@ def test_extract_usage_metrics():
                             "role": "assistant",
                             "content": [{"toolUse": {"toolUseId": "123", "name": "test", "input": {"key": "value"}}}],
                         },
-                        {"inputTokens": 1, "outputTokens": 1, "totalTokens": 1},
+                        {
+                            "inputTokens": 1,
+                            "outputTokens": 1,
+                            "totalTokens": 1,
+                            "cacheReadInputTokens": 1,
+                            "cacheWriteInputTokens": 1,
+                        },
                         {"latencyMs": 1},
                     )
                 },
@@ -398,7 +418,13 @@ def test_extract_usage_metrics():
                             "role": "assistant",
                             "content": [],
                         },
-                        {"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
+                        {
+                            "inputTokens": 0,
+                            "outputTokens": 0,
+                            "totalTokens": 0,
+                            "cacheReadInputTokens": 0,
+                            "cacheWriteInputTokens": 0,
+                        },
                         {"latencyMs": 0},
                     ),
                 },
@@ -426,7 +452,13 @@ def test_extract_usage_metrics():
                 },
                 {
                     "metadata": {
-                        "usage": {"inputTokens": 1, "outputTokens": 1, "totalTokens": 1},
+                        "usage": {
+                            "inputTokens": 1,
+                            "outputTokens": 1,
+                            "totalTokens": 1,
+                            "cacheReadInputTokens": 1,
+                            "cacheWriteInputTokens": 1,
+                        },
                         "metrics": {"latencyMs": 1},
                     }
                 },
@@ -506,6 +538,8 @@ def test_extract_usage_metrics():
                                     "inputTokens": 1,
                                     "outputTokens": 1,
                                     "totalTokens": 1,
+                                    "cacheReadInputTokens": 1,
+                                    "cacheWriteInputTokens": 1,
                                 },
                             },
                         },
@@ -518,7 +552,13 @@ def test_extract_usage_metrics():
                             "role": "assistant",
                             "content": [{"text": "REDACTED."}],
                         },
-                        {"inputTokens": 1, "outputTokens": 1, "totalTokens": 1},
+                        {
+                            "inputTokens": 1,
+                            "outputTokens": 1,
+                            "totalTokens": 1,
+                            "cacheReadInputTokens": 1,
+                            "cacheWriteInputTokens": 1,
+                        },
                         {"latencyMs": 1},
                     ),
                 },
@@ -584,7 +624,13 @@ async def test_stream_messages(agenerator, alist):
             "stop": (
                 "end_turn",
                 {"role": "assistant", "content": [{"text": "test"}]},
-                {"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
+                {
+                    "inputTokens": 0,
+                    "outputTokens": 0,
+                    "totalTokens": 0,
+                    "cacheReadInputTokens": 0,
+                    "cacheWriteInputTokens": 0,
+                },
                 {"latencyMs": 0},
             )
         },

tests/strands/models/test_anthropic.py

@@ -597,7 +597,12 @@ def test_format_chunk_message_stop(model):
 def test_format_chunk_metadata(model):
     event = {
         "type": "metadata",
-        "usage": {"input_tokens": 1, "output_tokens": 2},
+        "usage": {
+            "input_tokens": 1,
+            "output_tokens": 2,
+            "cache_read_input_tokens": 4,
+            "cache_creation_input_tokens": 5,
+        },
     }
 
     tru_chunk = model.format_chunk(event)
@@ -607,6 +612,8 @@ def test_format_chunk_metadata(model):
                 "inputTokens": 1,
                 "outputTokens": 2,
                 "totalTokens": 3,
+                "cacheReadInputTokens": 4,
+                "cacheWriteInputTokens": 5,
             },
             "metrics": {
                 "latencyMs": 0,

tests/strands/models/test_bedrock.py

@@ -497,7 +497,13 @@ async def test_converse_stream_input_guardrails(
 ):
     metadata_event = {
         "metadata": {
-            "usage": {"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
+            "usage": {
+                "inputTokens": 0,
+                "outputTokens": 0,
+                "totalTokens": 0,
+                "cacheReadInputTokens": 0,
+                "cacheWriteInputTokens": 0,
+            },
             "metrics": {"latencyMs": 245},
             "trace": {
                 "guardrail": {
@@ -552,7 +558,13 @@ async def test_converse_stream_output_guardrails(
     model.update_config(guardrail_redact_input=False, guardrail_redact_output=True)
     metadata_event = {
         "metadata": {
-            "usage": {"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
+            "usage": {
+                "inputTokens": 0,
+                "outputTokens": 0,
+                "totalTokens": 0,
+                "cacheReadInputTokens": 0,
+                "cacheWriteInputTokens": 0,
+            },
             "metrics": {"latencyMs": 245},
             "trace": {
                 "guardrail": {
@@ -609,7 +621,13 @@ async def test_converse_output_guardrails_redacts_input_and_output(
     model.update_config(guardrail_redact_output=True)
     metadata_event = {
         "metadata": {
-            "usage": {"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
+            "usage": {
+                "inputTokens": 0,
+                "outputTokens": 0,
+                "totalTokens": 0,
+                "cacheReadInputTokens": 0,
+                "cacheWriteInputTokens": 0,
+            },
             "metrics": {"latencyMs": 245},
             "trace": {
                 "guardrail": {
@@ -666,7 +684,13 @@ async def test_converse_output_no_blocked_guardrails_doesnt_redact(
 ):
     metadata_event = {
         "metadata": {
-            "usage": {"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
+            "usage": {
+                "inputTokens": 0,
+                "outputTokens": 0,
+                "totalTokens": 0,
+                "cacheReadInputTokens": 0,
+                "cacheWriteInputTokens": 0,
+            },
             "metrics": {"latencyMs": 245},
             "trace": {
                 "guardrail": {
@@ -719,7 +743,13 @@ async def test_converse_output_no_guardrail_redact(
 ):
     metadata_event = {
         "metadata": {
-            "usage": {"inputTokens": 0, "outputTokens": 0, "totalTokens": 0},
+            "usage": {
+                "inputTokens": 0,
+                "outputTokens": 0,
+                "totalTokens": 0,
+                "cacheReadInputTokens": 0,
+                "cacheWriteInputTokens": 0,
+            },
             "metrics": {"latencyMs": 245},
             "trace": {
                 "guardrail": {
@@ -909,7 +939,13 @@ async def test_stream_with_streaming_false_with_metrics_and_usage(bedrock_client
     """Test stream method with streaming=False."""
     bedrock_client.converse.return_value = {
         "output": {"message": {"role": "assistant", "content": [{"text": "test"}]}},
-        "usage": {"inputTokens": 1234, "outputTokens": 1234, "totalTokens": 2468},
+        "usage": {
+            "inputTokens": 1234,
+            "outputTokens": 1234,
+            "totalTokens": 2468,
+            "cacheReadInputTokens": 128,
+            "cacheWriteInputTokens": 512,
+        },
         "metrics": {"latencyMs": 1234},
         "stopReason": "tool_use",
     }
@@ -927,7 +963,13 @@ async def test_stream_with_streaming_false_with_metrics_and_usage(bedrock_client
         {"messageStop": {"stopReason": "tool_use", "additionalModelResponseFields": None}},
         {
             "metadata": {
-                "usage": {"inputTokens": 1234, "outputTokens": 1234, "totalTokens": 2468},
+                "usage": {
+                    "inputTokens": 1234,
+                    "outputTokens": 1234,
+                    "totalTokens": 2468,
+                    "cacheReadInputTokens": 128,
+                    "cacheWriteInputTokens": 512,
+                },
                 "metrics": {"latencyMs": 1234},
             }
         },

tests/strands/models/test_litellm.py

@@ -146,7 +146,15 @@ async def test_stream(litellm_client, model, alist):
     mock_event_3 = unittest.mock.Mock(choices=[unittest.mock.Mock(finish_reason=None, delta=mock_delta_3)])
     mock_event_4 = unittest.mock.Mock(choices=[unittest.mock.Mock(finish_reason=None, delta=mock_delta_4)])
     mock_event_5 = unittest.mock.Mock(choices=[unittest.mock.Mock(finish_reason="tool_calls", delta=mock_delta_5)])
-    mock_event_6 = unittest.mock.Mock()
+    mock_event_6 = unittest.mock.Mock(
+        usage=unittest.mock.Mock(
+            prompt_tokens_details=unittest.mock.Mock(
+                audio_tokens=None, cached_tokens=0, text_tokens=None, image_tokens=None
+            ),
+            cache_creation_input_tokens=0,
+            cache_read_input_tokens=0,
+        )
+    )
 
     litellm_client.chat.completions.create.return_value = iter(
         [mock_event_1, mock_event_2, mock_event_3, mock_event_4, mock_event_5, mock_event_6]

tests/strands/models/test_llamaapi.py

@@ -346,6 +346,8 @@ def test_format_chunk_metadata(model):
                 "inputTokens": 100,
                 "outputTokens": 50,
                 "totalTokens": 150,
+                "cacheReadInputTokens": 0,
+                "cacheWriteInputTokens": 0,
             },
             "metrics": {
                 "latencyMs": 0,

tests/strands/models/test_mistral.py

@@ -391,6 +391,8 @@ def test_format_chunk_metadata(model):
                 "inputTokens": 100,
                 "outputTokens": 50,
                 "totalTokens": 150,
+                "cacheReadInputTokens": 0,
+                "cacheWriteInputTokens": 0,
             },
             "metrics": {
                 "latencyMs": 250,
@@ -419,6 +421,8 @@ def test_format_chunk_metadata_no_latency(model):
                 "inputTokens": 100,
                 "outputTokens": 50,
                 "totalTokens": 150,
+                "cacheReadInputTokens": 0,
+                "cacheWriteInputTokens": 0,
             },
             "metrics": {
                 "latencyMs": 0,

tests/strands/models/test_ollama.py

@@ -398,6 +398,8 @@ def test_format_chunk_metadata(model):
                 "inputTokens": 100,
                 "outputTokens": 50,
                 "totalTokens": 150,
+                "cacheReadInputTokens": 0,
+                "cacheWriteInputTokens": 0,
             },
             "metrics": {
                 "latencyMs": 1.0,

tests/strands/telemetry/test_metrics.py

@@ -90,6 +90,8 @@ def usage(request):
         "inputTokens": 1,
         "outputTokens": 2,
         "totalTokens": 3,
+        "cacheReadInputTokens": 4,
+        "cacheWriteInputTokens": 5,
     }
     if hasattr(request, "param"):
         params.update(request.param)
@@ -315,17 +317,15 @@ def test_event_loop_metrics_update_usage(usage, event_loop_metrics, mock_get_met
         event_loop_metrics.update_usage(usage)
 
     tru_usage = event_loop_metrics.accumulated_usage
-    exp_usage = Usage(
-        inputTokens=3,
-        outputTokens=6,
-        totalTokens=9,
-    )
+    exp_usage = Usage(inputTokens=3, outputTokens=6, totalTokens=9, cacheReadInputTokens=12, cacheWriteInputTokens=15)
 
     assert tru_usage == exp_usage
     mock_get_meter_provider.return_value.get_meter.assert_called()
     metrics_client = event_loop_metrics._metrics_client
     metrics_client.event_loop_input_tokens.record.assert_called()
     metrics_client.event_loop_output_tokens.record.assert_called()
+    metrics_client.event_loop_input_tokens_cache_read.record.assert_called()
+    metrics_client.event_loop_input_tokens_cache_write.record.assert_called()
 
 
 def test_event_loop_metrics_update_metrics(metrics, event_loop_metrics, mock_get_meter_provider):
@@ -358,6 +358,8 @@ def test_event_loop_metrics_get_summary(trace, tool, event_loop_metrics, mock_ge
             "inputTokens": 0,
             "outputTokens": 0,
             "totalTokens": 0,
+            "cacheReadInputTokens": 0,
+            "cacheWriteInputTokens": 0,
         },
         "average_cycle_time": 0,
         "tool_usage": {
@@ -394,7 +396,7 @@ def test_event_loop_metrics_get_summary(trace, tool, event_loop_metrics, mock_ge
             {},
             "Event Loop Metrics Summary:\n"
             "├─ Cycles: total=0, avg_time=0.000s, total_time=0.000s\n"
-            "├─ Tokens: in=0, out=0, total=0\n"
+            "├─ Tokens: in=0 (cache_write=0), out=0, total=0 (cache_read=0)\n"
             "├─ Bedrock Latency: 0ms\n"
             "├─ Tool Usage:\n"
             "   └─ tool1:\n"
@@ -412,7 +414,7 @@ def test_event_loop_metrics_get_summary(trace, tool, event_loop_metrics, mock_ge
             {},
             "Event Loop Metrics Summary:\n"
             "├─ Cycles: total=0, avg_time=0.000s, total_time=0.000s\n"
-            "├─ Tokens: in=0, out=0, total=0\n"
+            "├─ Tokens: in=0 (cache_write=0), out=0, total=0 (cache_read=0)\n"
             "├─ Bedrock Latency: 0ms\n"
             "├─ Tool Usage:\n"
             "   └─ tool1:\n"