Populate ModelResponse.model_name from responses (#883)

Author: alexmojaki

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -236,7 +236,7 @@ def _process_response(self, response: AnthropicMessage) -> ModelResponse:
                     )
                 )
 
-        return ModelResponse(items, model_name=self._model_name)
+        return ModelResponse(items, model_name=response.model)
 
     async def _process_streamed_response(self, response: AsyncStream[RawMessageStreamEvent]) -> StreamedResponse:
         peekable_response = _utils.PeekableAsyncStream(response)

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -233,7 +233,7 @@ def _process_response(self, response: _GeminiResponse) -> ModelResponse:
             else:
                 raise UnexpectedModelBehavior('Content field missing from Gemini response', str(response))
         parts = response['candidates'][0]['content']['parts']
-        return _process_response_from_parts(parts, model_name=self._model_name)
+        return _process_response_from_parts(parts, model_name=response.get('model_version', self._model_name))
 
     async def _process_streamed_response(self, http_response: HTTPResponse) -> StreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
@@ -610,6 +610,7 @@ class _GeminiResponse(TypedDict):
     # usageMetadata appears to be required by both APIs but is omitted when streaming responses until the last response
     usage_metadata: NotRequired[Annotated[_GeminiUsageMetaData, pydantic.Field(alias='usageMetadata')]]
     prompt_feedback: NotRequired[Annotated[_GeminiPromptFeedback, pydantic.Field(alias='promptFeedback')]]
+    model_version: NotRequired[Annotated[str, pydantic.Field(alias='modelVersion')]]
 
 
 class _GeminiCandidates(TypedDict):

pydantic_ai_slim/pydantic_ai/models/groq.py

@@ -212,7 +212,7 @@ def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
                 items.append(ToolCallPart(tool_name=c.function.name, args=c.function.arguments, tool_call_id=c.id))
-        return ModelResponse(items, model_name=self._model_name, timestamp=timestamp)
+        return ModelResponse(items, model_name=response.model, timestamp=timestamp)
 
     async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> GroqStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""

pydantic_ai_slim/pydantic_ai/models/mistral.py

@@ -296,7 +296,7 @@ def _process_response(self, response: MistralChatCompletionResponse) -> ModelRes
                 tool = self._map_mistral_to_pydantic_tool_call(tool_call=tool_call)
                 parts.append(tool)
 
-        return ModelResponse(parts, model_name=self._model_name, timestamp=timestamp)
+        return ModelResponse(parts, model_name=response.model, timestamp=timestamp)
 
     async def _process_streamed_response(
         self,

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -232,7 +232,7 @@ def _process_response(self, response: chat.ChatCompletion) -> ModelResponse:
         if choice.message.tool_calls is not None:
             for c in choice.message.tool_calls:
                 items.append(ToolCallPart(c.function.name, c.function.arguments, c.id))
-        return ModelResponse(items, model_name=self._model_name, timestamp=timestamp)
+        return ModelResponse(items, model_name=response.model, timestamp=timestamp)
 
     async def _process_streamed_response(self, response: AsyncStream[ChatCompletionChunk]) -> OpenAIStreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""

tests/models/test_anthropic.py

@@ -112,7 +112,7 @@ def completion_message(content: list[ContentBlock], usage: AnthropicUsage) -> An
     return AnthropicMessage(
         id='123',
         content=content,
-        model='claude-3-5-haiku-latest',
+        model='claude-3-5-haiku-123',
         role='assistant',
         stop_reason='end_turn',
         type='message',
@@ -141,13 +141,13 @@ async def test_sync_request_text_response(allow_model_requests: None):
             ModelRequest(parts=[UserPromptPart(content='hello', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
                 parts=[TextPart(content='world')],
-                model_name='claude-3-5-haiku-latest',
+                model_name='claude-3-5-haiku-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
             ModelRequest(parts=[UserPromptPart(content='hello', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
                 parts=[TextPart(content='world')],
-                model_name='claude-3-5-haiku-latest',
+                model_name='claude-3-5-haiku-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
         ]
@@ -190,7 +190,7 @@ async def test_request_structured_response(allow_model_requests: None):
                         tool_call_id='123',
                     )
                 ],
-                model_name='claude-3-5-haiku-latest',
+                model_name='claude-3-5-haiku-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
             ModelRequest(
@@ -252,7 +252,7 @@ async def get_location(loc_name: str) -> str:
                         tool_call_id='1',
                     )
                 ],
-                model_name='claude-3-5-haiku-latest',
+                model_name='claude-3-5-haiku-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
             ModelRequest(
@@ -273,7 +273,7 @@ async def get_location(loc_name: str) -> str:
                         tool_call_id='2',
                     )
                 ],
-                model_name='claude-3-5-haiku-latest',
+                model_name='claude-3-5-haiku-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
             ModelRequest(
@@ -288,7 +288,7 @@ async def get_location(loc_name: str) -> str:
             ),
             ModelResponse(
                 parts=[TextPart(content='final response')],
-                model_name='claude-3-5-haiku-latest',
+                model_name='claude-3-5-haiku-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
         ]

tests/models/test_gemini.py

@@ -440,7 +440,7 @@ def gemini_response(content: _GeminiContent, finish_reason: Literal['STOP'] | No
     candidate = _GeminiCandidates(content=content, index=0, safety_ratings=[])
     if finish_reason:  # pragma: no cover
         candidate['finish_reason'] = finish_reason
-    return _GeminiResponse(candidates=[candidate], usage_metadata=example_usage())
+    return _GeminiResponse(candidates=[candidate], usage_metadata=example_usage(), model_version='gemini-1.5-flash-123')
 
 
 def example_usage() -> _GeminiUsageMetaData:
@@ -459,7 +459,9 @@ async def test_text_success(get_gemini_client: GetGeminiClient):
         [
             ModelRequest(parts=[UserPromptPart(content='Hello', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
-                parts=[TextPart(content='Hello world')], model_name='gemini-1.5-flash', timestamp=IsNow(tz=timezone.utc)
+                parts=[TextPart(content='Hello world')],
+                model_name='gemini-1.5-flash-123',
+                timestamp=IsNow(tz=timezone.utc),
             ),
         ]
     )
@@ -472,13 +474,13 @@ async def test_text_success(get_gemini_client: GetGeminiClient):
             ModelRequest(parts=[UserPromptPart(content='Hello', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
                 parts=[TextPart(content='Hello world')],
-                model_name='gemini-1.5-flash',
+                model_name='gemini-1.5-flash-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
             ModelRequest(parts=[UserPromptPart(content='Hello', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
                 parts=[TextPart(content='Hello world')],
-                model_name='gemini-1.5-flash',
+                model_name='gemini-1.5-flash-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
         ]
@@ -505,7 +507,7 @@ async def test_request_structured_response(get_gemini_client: GetGeminiClient):
                         args={'response': [1, 2, 123]},
                     )
                 ],
-                model_name='gemini-1.5-flash',
+                model_name='gemini-1.5-flash-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
             ModelRequest(
@@ -566,7 +568,7 @@ async def get_location(loc_name: str) -> str:
                         args={'loc_name': 'San Fransisco'},
                     )
                 ],
-                model_name='gemini-1.5-flash',
+                model_name='gemini-1.5-flash-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
             ModelRequest(
@@ -589,7 +591,7 @@ async def get_location(loc_name: str) -> str:
                         args={'loc_name': 'New York'},
                     ),
                 ],
-                model_name='gemini-1.5-flash',
+                model_name='gemini-1.5-flash-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
             ModelRequest(
@@ -604,7 +606,7 @@ async def get_location(loc_name: str) -> str:
             ),
             ModelResponse(
                 parts=[TextPart(content='final response')],
-                model_name='gemini-1.5-flash',
+                model_name='gemini-1.5-flash-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
         ]

tests/models/test_groq.py

@@ -103,7 +103,7 @@ def completion_message(message: ChatCompletionMessage, *, usage: CompletionUsage
         id='123',
         choices=[Choice(finish_reason='stop', index=0, message=message)],
         created=1704067200,  # 2024-01-01
-        model='llama-3.3-70b-versatile',
+        model='llama-3.3-70b-versatile-123',
         object='chat.completion',
         usage=usage,
     )
@@ -130,13 +130,13 @@ async def test_request_simple_success(allow_model_requests: None):
             ModelRequest(parts=[UserPromptPart(content='hello', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
                 parts=[TextPart(content='world')],
-                model_name='llama-3.3-70b-versatile',
+                model_name='llama-3.3-70b-versatile-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
             ModelRequest(parts=[UserPromptPart(content='hello', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
                 parts=[TextPart(content='world')],
-                model_name='llama-3.3-70b-versatile',
+                model_name='llama-3.3-70b-versatile-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
         ]
@@ -187,7 +187,7 @@ async def test_request_structured_response(allow_model_requests: None):
                         tool_call_id='123',
                     )
                 ],
-                model_name='llama-3.3-70b-versatile',
+                model_name='llama-3.3-70b-versatile-123',
                 timestamp=datetime(2024, 1, 1, tzinfo=timezone.utc),
             ),
             ModelRequest(
@@ -273,7 +273,7 @@ async def get_location(loc_name: str) -> str:
                         tool_call_id='1',
                     )
                 ],
-                model_name='llama-3.3-70b-versatile',
+                model_name='llama-3.3-70b-versatile-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
             ModelRequest(
@@ -294,7 +294,7 @@ async def get_location(loc_name: str) -> str:
                         tool_call_id='2',
                     )
                 ],
-                model_name='llama-3.3-70b-versatile',
+                model_name='llama-3.3-70b-versatile-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
             ModelRequest(
@@ -309,7 +309,7 @@ async def get_location(loc_name: str) -> str:
             ),
             ModelResponse(
                 parts=[TextPart(content='final response')],
-                model_name='llama-3.3-70b-versatile',
+                model_name='llama-3.3-70b-versatile-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
         ]

tests/models/test_mistral.py

@@ -123,7 +123,7 @@ def completion_message(
         id='123',
         choices=[MistralChatCompletionChoice(finish_reason='stop', index=0, message=message)],
         created=1704067200 if with_created else None,  # 2024-01-01
-        model='mistral-large-latest',
+        model='mistral-large-123',
         object='chat.completion',
         usage=usage or MistralUsageInfo(prompt_tokens=1, completion_tokens=1, total_tokens=1),
     )
@@ -217,13 +217,13 @@ async def test_multiple_completions(allow_model_requests: None):
             ModelRequest(parts=[UserPromptPart(content='hello', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
                 parts=[TextPart(content='world')],
-                model_name='mistral-large-latest',
+                model_name='mistral-large-123',
                 timestamp=IsNow(tz=timezone.utc),
             ),
             ModelRequest(parts=[UserPromptPart(content='hello again', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
                 parts=[TextPart(content='hello again')],
-                model_name='mistral-large-latest',
+                model_name='mistral-large-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
         ]
@@ -269,19 +269,19 @@ async def test_three_completions(allow_model_requests: None):
             ModelRequest(parts=[UserPromptPart(content='hello', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
                 parts=[TextPart(content='world')],
-                model_name='mistral-large-latest',
+                model_name='mistral-large-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
             ModelRequest(parts=[UserPromptPart(content='hello again', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
                 parts=[TextPart(content='hello again')],
-                model_name='mistral-large-latest',
+                model_name='mistral-large-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
             ModelRequest(parts=[UserPromptPart(content='final message', timestamp=IsNow(tz=timezone.utc))]),
             ModelResponse(
                 parts=[TextPart(content='final message')],
-                model_name='mistral-large-latest',
+                model_name='mistral-large-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
         ]
@@ -396,7 +396,7 @@ class CityLocation(BaseModel):
                         tool_call_id='123',
                     )
                 ],
-                model_name='mistral-large-latest',
+                model_name='mistral-large-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
             ModelRequest(
@@ -458,7 +458,7 @@ class CityLocation(BaseModel):
                         tool_call_id='123',
                     )
                 ],
-                model_name='mistral-large-latest',
+                model_name='mistral-large-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
             ModelRequest(
@@ -519,7 +519,7 @@ async def test_request_result_type_with_arguments_str_response(allow_model_reque
                         tool_call_id='123',
                     )
                 ],
-                model_name='mistral-large-latest',
+                model_name='mistral-large-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
             ModelRequest(
@@ -1104,7 +1104,7 @@ async def get_location(loc_name: str) -> str:
                         tool_call_id='1',
                     )
                 ],
-                model_name='mistral-large-latest',
+                model_name='mistral-large-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
             ModelRequest(
@@ -1125,7 +1125,7 @@ async def get_location(loc_name: str) -> str:
                         tool_call_id='2',
                     )
                 ],
-                model_name='mistral-large-latest',
+                model_name='mistral-large-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
             ModelRequest(
@@ -1140,7 +1140,7 @@ async def get_location(loc_name: str) -> str:
             ),
             ModelResponse(
                 parts=[TextPart(content='final response')],
-                model_name='mistral-large-latest',
+                model_name='mistral-large-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
         ]
@@ -1244,7 +1244,7 @@ async def get_location(loc_name: str) -> str:
                         tool_call_id='1',
                     )
                 ],
-                model_name='mistral-large-latest',
+                model_name='mistral-large-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
             ModelRequest(
@@ -1265,7 +1265,7 @@ async def get_location(loc_name: str) -> str:
                         tool_call_id='2',
                     )
                 ],
-                model_name='mistral-large-latest',
+                model_name='mistral-large-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
             ModelRequest(
@@ -1286,7 +1286,7 @@ async def get_location(loc_name: str) -> str:
                         tool_call_id='1',
                     )
                 ],
-                model_name='mistral-large-latest',
+                model_name='mistral-large-123',
                 timestamp=datetime(2024, 1, 1, 0, 0, tzinfo=timezone.utc),
             ),
             ModelRequest(