Add model_name to ModelResponse (#701)

Co-authored-by: David Montague <[email protected]>

Author: sydney-runkle

docs/agents.md

@@ -445,6 +445,7 @@ with capture_run_messages() as messages:  # (2)!
                         part_kind='tool-call',
                     )
                 ],
+                model_name='function:model_logic',
                 timestamp=datetime.datetime(...),
                 kind='response',
             ),
@@ -469,6 +470,7 @@ with capture_run_messages() as messages:  # (2)!
                         part_kind='tool-call',
                     )
                 ],
+                model_name='function:model_logic',
                 timestamp=datetime.datetime(...),
                 kind='response',
             ),

docs/message-history.md

@@ -62,6 +62,7 @@ print(result.all_messages())
                 part_kind='text',
             )
         ],
+        model_name='function:model_logic',
         timestamp=datetime.datetime(...),
         kind='response',
     ),
@@ -135,6 +136,7 @@ async def main():
                         part_kind='text',
                     )
                 ],
+                model_name='function:stream_model_logic',
                 timestamp=datetime.datetime(...),
                 kind='response',
             ),
@@ -191,6 +193,7 @@ print(result2.all_messages())
                 part_kind='text',
             )
         ],
+        model_name='function:model_logic',
         timestamp=datetime.datetime(...),
         kind='response',
     ),
@@ -211,6 +214,7 @@ print(result2.all_messages())
                 part_kind='text',
             )
         ],
+        model_name='function:model_logic',
         timestamp=datetime.datetime(...),
         kind='response',
     ),
@@ -265,6 +269,7 @@ print(result2.all_messages())
                 part_kind='text',
             )
         ],
+        model_name='function:model_logic',
         timestamp=datetime.datetime(...),
         kind='response',
     ),
@@ -285,6 +290,7 @@ print(result2.all_messages())
                 part_kind='text',
             )
         ],
+        model_name='function:model_logic',
         timestamp=datetime.datetime(...),
         kind='response',
     ),

docs/testing-evals.md

@@ -151,6 +151,7 @@ async def test_forecast():
                     tool_call_id=None,
                 )
             ],
+            model_name='test',
             timestamp=IsNow(tz=timezone.utc),
         ),
         ModelRequest(
@@ -169,6 +170,7 @@ async def test_forecast():
                     content='{"weather_forecast":"Sunny with a chance of rain"}',
                 )
             ],
+            model_name='test',
             timestamp=IsNow(tz=timezone.utc),
         ),
     ]

docs/tools.md

@@ -92,6 +92,7 @@ print(dice_result.all_messages())
                 part_kind='tool-call',
             )
         ],
+        model_name='function:model_logic',
         timestamp=datetime.datetime(...),
         kind='response',
     ),
@@ -116,6 +117,7 @@ print(dice_result.all_messages())
                 part_kind='tool-call',
             )
         ],
+        model_name='function:model_logic',
         timestamp=datetime.datetime(...),
         kind='response',
     ),
@@ -138,6 +140,7 @@ print(dice_result.all_messages())
                 part_kind='text',
             )
         ],
+        model_name='function:model_logic',
         timestamp=datetime.datetime(...),
         kind='response',
     ),

pydantic_ai_slim/pydantic_ai/messages.py

@@ -252,6 +252,9 @@ class ModelResponse:
     parts: list[ModelResponsePart]
     """The parts of the model message."""
 
+    model_name: str | None = None
+    """The name of the model that generated the response."""
+
     timestamp: datetime = field(default_factory=_now_utc)
     """The timestamp of the response.
 
@@ -262,14 +265,14 @@ class ModelResponse:
     """Message type identifier, this is available on all parts as a discriminator."""
 
     @classmethod
-    def from_text(cls, content: str, timestamp: datetime | None = None) -> Self:
+    def from_text(cls, content: str, model_name: str | None = None, timestamp: datetime | None = None) -> Self:
         """Create a `ModelResponse` containing a single `TextPart`."""
-        return cls([TextPart(content=content)], timestamp=timestamp or _now_utc())
+        return cls([TextPart(content=content)], model_name=model_name, timestamp=timestamp or _now_utc())
 
     @classmethod
-    def from_tool_call(cls, tool_call: ToolCallPart) -> Self:
+    def from_tool_call(cls, tool_call: ToolCallPart, model_name: str | None = None) -> Self:
         """Create a `ModelResponse` containing a single `ToolCallPart`."""
-        return cls([tool_call])
+        return cls([tool_call], model_name=model_name)
 
 
 ModelMessage = Annotated[Union[ModelRequest, ModelResponse], pydantic.Discriminator('kind')]

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -161,6 +161,7 @@ async def request_stream(
 class StreamedResponse(ABC):
     """Streamed response from an LLM when calling a tool."""
 
+    _model_name: str
     _usage: Usage = field(default_factory=Usage, init=False)
     _parts_manager: ModelResponsePartsManager = field(default_factory=ModelResponsePartsManager, init=False)
     _event_iterator: AsyncIterator[ModelResponseStreamEvent] | None = field(default=None, init=False)
@@ -184,7 +185,13 @@ async def _get_event_iterator(self) -> AsyncIterator[ModelResponseStreamEvent]:
 
     def get(self) -> ModelResponse:
         """Build a [`ModelResponse`][pydantic_ai.messages.ModelResponse] from the data received from the stream so far."""
-        return ModelResponse(parts=self._parts_manager.get_parts(), timestamp=self.timestamp())
+        return ModelResponse(
+            parts=self._parts_manager.get_parts(), model_name=self._model_name, timestamp=self.timestamp()
+        )
+
+    def model_name(self) -> str:
+        """Get the model name of the response."""
+        return self._model_name
 
     def usage(self) -> Usage:
         """Get the usage of the response so far. This will not be the final usage until the stream is exhausted."""

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -152,7 +152,7 @@ class AnthropicAgentModel(AgentModel):
     """Implementation of `AgentModel` for Anthropic models."""
 
     client: AsyncAnthropic
-    model_name: str
+    model_name: AnthropicModelName
     allow_text_result: bool
     tools: list[ToolParam]
 
@@ -210,8 +210,7 @@ async def _messages_create(
             timeout=model_settings.get('timeout', NOT_GIVEN),
         )
 
-    @staticmethod
-    def _process_response(response: AnthropicMessage) -> ModelResponse:
+    def _process_response(self, response: AnthropicMessage) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         items: list[ModelResponsePart] = []
         for item in response.content:
@@ -227,7 +226,7 @@ def _process_response(response: AnthropicMessage) -> ModelResponse:
                     )
                 )
 
-        return ModelResponse(items)
+        return ModelResponse(items, model_name=self.model_name)
 
     @staticmethod
     async def _process_streamed_response(response: AsyncStream[RawMessageStreamEvent]) -> StreamedResponse:

pydantic_ai_slim/pydantic_ai/models/cohere.py

@@ -172,8 +172,7 @@ async def _chat(
             p=model_settings.get('top_p', OMIT),
         )
 
-    @staticmethod
-    def _process_response(response: ChatResponse) -> ModelResponse:
+    def _process_response(self, response: ChatResponse) -> ModelResponse:
         """Process a non-streamed response, and prepare a message to return."""
         parts: list[ModelResponsePart] = []
         if response.message.content is not None and len(response.message.content) > 0:
@@ -190,7 +189,7 @@ def _process_response(response: ChatResponse) -> ModelResponse:
                         tool_call_id=c.id,
                     )
                 )
-        return ModelResponse(parts=parts)
+        return ModelResponse(parts=parts, model_name=self.model_name)
 
     @classmethod
     def _map_message(cls, message: ModelMessage) -> Iterable[ChatMessageV2]:

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -71,16 +71,15 @@ async def agent_model(
         result_tools: list[ToolDefinition],
     ) -> AgentModel:
         return FunctionAgentModel(
-            self.function, self.stream_function, AgentInfo(function_tools, allow_text_result, result_tools, None)
+            self.function,
+            self.stream_function,
+            AgentInfo(function_tools, allow_text_result, result_tools, None),
         )
 
     def name(self) -> str:
-        labels: list[str] = []
-        if self.function is not None:
-            labels.append(self.function.__name__)
-        if self.stream_function is not None:
-            labels.append(f'stream-{self.stream_function.__name__}')
-        return f'function:{",".join(labels)}'
+        function_name = self.function.__name__ if self.function is not None else ''
+        stream_function_name = self.stream_function.__name__ if self.stream_function is not None else ''
+        return f'function:{function_name}:{stream_function_name}'
 
 
 @dataclass(frozen=True)
@@ -147,12 +146,15 @@ async def request(
         agent_info = replace(self.agent_info, model_settings=model_settings)
 
         assert self.function is not None, 'FunctionModel must receive a `function` to support non-streamed requests'
+        model_name = f'function:{self.function.__name__}'
+
         if inspect.iscoroutinefunction(self.function):
             response = await self.function(messages, agent_info)
         else:
             response_ = await _utils.run_in_executor(self.function, messages, agent_info)
             assert isinstance(response_, ModelResponse), response_
             response = response_
+        response.model_name = model_name
         # TODO is `messages` right here? Should it just be new messages?
         return response, _estimate_usage(chain(messages, [response]))
 
@@ -163,13 +165,15 @@ async def request_stream(
         assert (
             self.stream_function is not None
         ), 'FunctionModel must receive a `stream_function` to support streamed requests'
+        model_name = f'function:{self.stream_function.__name__}'
+
         response_stream = PeekableAsyncStream(self.stream_function(messages, self.agent_info))
 
         first = await response_stream.peek()
         if isinstance(first, _utils.Unset):
             raise ValueError('Stream function must return at least one item')
 
-        yield FunctionStreamedResponse(response_stream)
+        yield FunctionStreamedResponse(_model_name=model_name, _iter=response_stream)
 
 
 @dataclass

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -229,15 +229,13 @@ async def _make_request(
                 raise exceptions.UnexpectedModelBehavior(f'Unexpected response from gemini {r.status_code}', r.text)
             yield r
 
-    @staticmethod
-    def _process_response(response: _GeminiResponse) -> ModelResponse:
+    def _process_response(self, response: _GeminiResponse) -> ModelResponse:
         if len(response['candidates']) != 1:
             raise UnexpectedModelBehavior('Expected exactly one candidate in Gemini response')
         parts = response['candidates'][0]['content']['parts']
-        return _process_response_from_parts(parts)
+        return _process_response_from_parts(parts, model_name=self.model_name)
 
-    @staticmethod
-    async def _process_streamed_response(http_response: HTTPResponse) -> StreamedResponse:
+    async def _process_streamed_response(self, http_response: HTTPResponse) -> StreamedResponse:
         """Process a streamed response, and prepare a streaming response to return."""
         aiter_bytes = http_response.aiter_bytes()
         start_response: _GeminiResponse | None = None
@@ -258,7 +256,7 @@ async def _process_streamed_response(http_response: HTTPResponse) -> StreamedRes
         if start_response is None:
             raise UnexpectedModelBehavior('Streamed response ended without content or tool calls')
 
-        return GeminiStreamedResponse(_content=content, _stream=aiter_bytes)
+        return GeminiStreamedResponse(_model_name=self.model_name, _content=content, _stream=aiter_bytes)
 
     @classmethod
     def _message_to_gemini_content(
@@ -432,7 +430,9 @@ def _function_call_part_from_call(tool: ToolCallPart) -> _GeminiFunctionCallPart
     return _GeminiFunctionCallPart(function_call=_GeminiFunctionCall(name=tool.tool_name, args=tool.args_as_dict()))
 
 
-def _process_response_from_parts(parts: Sequence[_GeminiPartUnion], timestamp: datetime | None = None) -> ModelResponse:
+def _process_response_from_parts(
+    parts: Sequence[_GeminiPartUnion], model_name: GeminiModelName, timestamp: datetime | None = None
+) -> ModelResponse:
     items: list[ModelResponsePart] = []
     for part in parts:
         if 'text' in part:
@@ -448,7 +448,7 @@ def _process_response_from_parts(parts: Sequence[_GeminiPartUnion], timestamp: d
             raise exceptions.UnexpectedModelBehavior(
                 f'Unsupported response from Gemini, expected all parts to be function calls or text, got: {part!r}'
             )
-    return ModelResponse(items, timestamp=timestamp or _utils.now_utc())
+    return ModelResponse(parts=items, model_name=model_name, timestamp=timestamp or _utils.now_utc())
 
 
 class _GeminiFunctionCall(TypedDict):