diff --git a/src/strands/models/mistral.py b/src/strands/models/mistral.py index 504e81c92..f44a11d30 100644 --- a/src/strands/models/mistral.py +++ b/src/strands/models/mistral.py @@ -496,8 +496,8 @@ async def stream( yield self.format_chunk({"chunk_type": "message_stop", "data": choice.finish_reason}) - if hasattr(chunk, "usage"): - yield self.format_chunk({"chunk_type": "metadata", "data": chunk.usage}) + if hasattr(chunk, "data") and hasattr(chunk.data, "usage") and chunk.data.usage: + yield self.format_chunk({"chunk_type": "metadata", "data": chunk.data.usage}) except Exception as e: if "rate" in str(e).lower() or "429" in str(e): diff --git a/tests/strands/models/test_mistral.py b/tests/strands/models/test_mistral.py index ad74bae89..338c66b19 100644 --- a/tests/strands/models/test_mistral.py +++ b/tests/strands/models/test_mistral.py @@ -451,9 +451,9 @@ async def test_stream(mistral_client, model, agenerator, alist, captured_warning delta=unittest.mock.Mock(content="test stream", tool_calls=None), finish_reason="end_turn", ) - ] + ], + usage=mock_usage, ), - usage=mock_usage, ) mistral_client.chat.stream_async = unittest.mock.AsyncMock(return_value=agenerator([mock_event])) @@ -492,9 +492,9 @@ async def test_tool_choice_not_supported_warns(mistral_client, model, agenerator delta=unittest.mock.Mock(content="test stream", tool_calls=None), finish_reason="end_turn", ) - ] + ], + usage=mock_usage, ), - usage=mock_usage, ) mistral_client.chat.stream_async = unittest.mock.AsyncMock(return_value=agenerator([mock_event]))