models - litellm - capture usage (strands-agents#73)

pgrayy · web-flow · commit d0a57a9c46fc · 2025-05-21T12:08:35.000-04:00
diff --git a/src/strands/models/litellm.py b/src/strands/models/litellm.py
@@ -334,6 +334,8 @@ def stream(self, request: dict[str, Any]) -> Iterable[dict[str, Any]]:
 
         yield {"chunk_type": "message_stop", "data": choice.finish_reason}
 
-        event = next(response)
-        if hasattr(event, "usage"):
-            yield {"chunk_type": "metadata", "data": event.usage}
+        # Skip remaining events as we don't have use for anything except the final usage payload
+        for event in response:
+            _ = event
+
+        yield {"chunk_type": "metadata", "data": event.usage}
diff --git a/tests/strands/models/test_litellm.py b/tests/strands/models/test_litellm.py
@@ -545,12 +545,14 @@ def test_stream(litellm_client, model):
 
 def test_stream_empty(litellm_client, model):
     mock_delta = unittest.mock.Mock(content=None, tool_calls=None)
+    mock_usage = unittest.mock.Mock(prompt_tokens=0, completion_tokens=0, total_tokens=0)
 
     mock_event_1 = unittest.mock.Mock(choices=[unittest.mock.Mock(finish_reason=None, delta=mock_delta)])
     mock_event_2 = unittest.mock.Mock(choices=[unittest.mock.Mock(finish_reason="stop")])
-    mock_event_3 = unittest.mock.Mock(spec=[])
+    mock_event_3 = unittest.mock.Mock()
+    mock_event_4 = unittest.mock.Mock(usage=mock_usage)
 
-    litellm_client.chat.completions.create.return_value = iter([mock_event_1, mock_event_2, mock_event_3])
+    litellm_client.chat.completions.create.return_value = iter([mock_event_1, mock_event_2, mock_event_3, mock_event_4])
 
     request = {"model": "m1", "messages": [{"role": "user", "content": []}]}
     response = model.stream(request)
@@ -561,6 +563,7 @@ def test_stream_empty(litellm_client, model):
         {"chunk_type": "content_start", "data_type": "text"},
         {"chunk_type": "content_stop", "data_type": "text"},
         {"chunk_type": "message_stop", "data": "stop"},
+        {"chunk_type": "metadata", "data": mock_usage},
     ]
 
     assert tru_events == exp_events
