diff --git a/pydantic_ai_slim/pydantic_ai/models/mistral.py b/pydantic_ai_slim/pydantic_ai/models/mistral.py index 4a29c0b7d..5372dbb08 100644 --- a/pydantic_ai_slim/pydantic_ai/models/mistral.py +++ b/pydantic_ai_slim/pydantic_ai/models/mistral.py @@ -45,6 +45,8 @@ check_allow_model_requests, get_user_agent, ) +from mistralai import CompletionChunk as MistralCompletionChunk +from mistralai.models import ChatCompletionResponse as MistralChatCompletionResponse try: from mistralai import ( @@ -679,13 +681,9 @@ def _validate_required_json_schema(json_dict: dict[str, Any], json_schema: dict[ def _map_usage(response: MistralChatCompletionResponse | MistralCompletionChunk) -> Usage: """Maps a Mistral Completion Chunk or Chat Completion Response to a Usage.""" - if response.usage: - return Usage( - request_tokens=response.usage.prompt_tokens, - response_tokens=response.usage.completion_tokens, - total_tokens=response.usage.total_tokens, - details=None, - ) + usage = response.usage + if usage: + return Usage(usage.prompt_tokens, usage.completion_tokens, usage.total_tokens, None) else: return Usage() # pragma: no cover