refacotr _map_model_response with context

Author: ajac-zero

pydantic_ai_slim/pydantic_ai/models/openai.py

@@ -707,42 +707,105 @@ def _get_web_search_options(self, model_request_parameters: ModelRequestParamete
                     f'`{tool.__class__.__name__}` is not supported by `OpenAIChatModel`. If it should be, please file an issue.'
                 )
 
+    @dataclass
+    class _MapModelResposeContext:
+        """Context object for mapping a `ModelResponse` to OpenAI chat completion parameters.
+
+        This class is designed to be subclassed to add new fields for custom logic,
+        collecting various parts of the model response (like text and tool calls)
+        to form a single assistant message.
+        """
+        texts: list[str] = field(default_factory=list)
+        tool_calls: list[ChatCompletionMessageFunctionToolCallParam] = field(default_factory=list)
+
+        def into_message_param(self) -> chat.ChatCompletionAssistantMessageParam:
+            """Converts the collected texts and tool calls into a single OpenAI `ChatCompletionAssistantMessageParam`.
+
+            This method serves as a hook that can be overridden by subclasses
+            to implement custom logic for how collected parts are transformed into the final message parameter.
+
+            Returns:
+                An OpenAI `ChatCompletionAssistantMessageParam` object representing the assistant's response.
+            """
+            message_param = chat.ChatCompletionAssistantMessageParam(role='assistant')
+            if self.texts:
+                # Note: model responses from this model should only have one text item, so the following
+                # shouldn't merge multiple texts into one unless you switch models between runs:
+                message_param['content'] = '\n\n'.join(self.texts)
+            else:
+                message_param['content'] = None
+            if self.tool_calls:
+                message_param['tool_calls'] = self.tool_calls
+            return message_param
+
+    def _map_response_text_part(self, ctx: _MapModelResposeContext, item: TextPart) -> None:
+        """Maps a `TextPart` to the response context.
+
+        This method serves as a hook that can be overridden by subclasses
+        to implement custom logic for handling text parts.
+        """
+        ctx.texts.append(item.content)
+
+    def _map_response_thinking_part(self, ctx: _MapModelResposeContext, item: ThinkingPart) -> None:
+        """Maps a `ThinkingPart` to the response context.
+
+        This method serves as a hook that can be overridden by subclasses
+        to implement custom logic for handling thinking parts.
+        """
+        # NOTE: DeepSeek `reasoning_content` field should NOT be sent back per https://api-docs.deepseek.com/guides/reasoning_model,
+        # but we currently just send it in `<think>` tags anyway as we don't want DeepSeek-specific checks here.
+        # If you need this changed, please file an issue.
+        start_tag, end_tag = self.profile.thinking_tags
+        ctx.texts.append('\n'.join([start_tag, item.content, end_tag]))
+
+    def _map_response_tool_call_part(self, ctx: _MapModelResposeContext, item: ToolCallPart) -> None:
+        """Maps a `ToolCallPart` to the response context.
+
+        This method serves as a hook that can be overridden by subclasses
+        to implement custom logic for handling tool call parts.
+        """
+        ctx.tool_calls.append(self._map_tool_call(item))
+
+    def _map_response_builtin_part(
+        self, ctx: _MapModelResposeContext, item: BuiltinToolCallPart | BuiltinToolReturnPart
+    ) -> None:
+        """Maps a built-in tool call or return part to the response context.
+
+        This method serves as a hook that can be overridden by subclasses
+        to implement custom logic for handling built-in tool parts.
+        """
+        # OpenAI doesn't return built-in tool calls
+        pass
+
+    def _map_response_file_part(self, ctx: _MapModelResposeContext, item: FilePart) -> None:
+        """Maps a `FilePart` to the response context.
+
+        This method serves as a hook that can be overridden by subclasses
+        to implement custom logic for handling file parts.
+        """
+        # Files generated by models are not sent back to models that don't themselves generate files.
+        pass
+
     def _map_model_response(self, message: ModelResponse) -> chat.ChatCompletionMessageParam:
         """Hook that determines how `ModelResponse` is mapped into `ChatCompletionMessageParam` objects before sending.
 
         Subclasses of `OpenAIChatModel` may override this method to provide their own mapping logic.
         """
-        texts: list[str] = []
-        tool_calls: list[ChatCompletionMessageFunctionToolCallParam] = []
+        ctx = self._MapModelResposeContext()
         for item in message.parts:
             if isinstance(item, TextPart):
-                texts.append(item.content)
+                self._map_response_text_part(ctx, item)
             elif isinstance(item, ThinkingPart):
-                # NOTE: DeepSeek `reasoning_content` field should NOT be sent back per https://api-docs.deepseek.com/guides/reasoning_model,
-                # but we currently just send it in `<think>` tags anyway as we don't want DeepSeek-specific checks here.
-                # If you need this changed, please file an issue.
-                start_tag, end_tag = self.profile.thinking_tags
-                texts.append('\n'.join([start_tag, item.content, end_tag]))
+                self._map_response_thinking_part(ctx, item)
             elif isinstance(item, ToolCallPart):
-                tool_calls.append(self._map_tool_call(item))
-            # OpenAI doesn't return built-in tool calls
+                self._map_response_tool_call_part(ctx, item)
             elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):  # pragma: no cover
-                pass
+                self._map_response_builtin_part(ctx, item)
             elif isinstance(item, FilePart):  # pragma: no cover
-                # Files generated by models are not sent back to models that don't themselves generate files.
-                pass
+                self._map_response_file_part(ctx, item)
             else:
                 assert_never(item)
-        message_param = chat.ChatCompletionAssistantMessageParam(role='assistant')
-        if texts:
-            # Note: model responses from this model should only have one text item, so the following
-            # shouldn't merge multiple texts into one unless you switch models between runs:
-            message_param['content'] = '\n\n'.join(texts)
-        else:
-            message_param['content'] = None
-        if tool_calls:
-            message_param['tool_calls'] = tool_calls
-        return message_param
+        return ctx.into_message_param()
 
     def _map_finish_reason(
         self, key: Literal['stop', 'length', 'tool_calls', 'content_filter', 'function_call']

pydantic_ai_slim/pydantic_ai/models/openrouter.py

@@ -1,23 +1,17 @@
 from __future__ import annotations as _annotations
 
 from collections.abc import Iterable
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Any, Literal, cast
 
 from pydantic import BaseModel
 from typing_extensions import TypedDict, assert_never, override
 
 from ..exceptions import ModelHTTPError, UnexpectedModelBehavior
 from ..messages import (
-    BuiltinToolCallPart,
-    BuiltinToolReturnPart,
-    FilePart,
     FinishReason,
-    ModelResponse,
     ModelResponseStreamEvent,
-    TextPart,
     ThinkingPart,
-    ToolCallPart,
 )
 from ..profiles import ModelProfileSpec
 from ..providers import Provider
@@ -522,40 +516,26 @@ def _process_provider_details(self, response: chat.ChatCompletion) -> dict[str,
         provider_details.update(_map_openrouter_provider_details(response))
         return provider_details
 
+    @dataclass
+    class _MapModelResposeContext(OpenAIChatModel._MapModelResposeContext):  # type: ignore[reportPrivateUsage]
+        reasoning_details: list[dict[str, Any]] = field(default_factory=list)
+
+        def into_message_param(self) -> chat.ChatCompletionAssistantMessageParam:
+            message_param = super().into_message_param()
+            if self.reasoning_details:
+                message_param['reasoning_details'] = self.reasoning_details  # type: ignore[reportGeneralTypeIssues]
+            return message_param
+
     @override
-    def _map_model_response(self, message: ModelResponse) -> chat.ChatCompletionMessageParam:
-        texts: list[str] = []
-        tool_calls: list[chat.ChatCompletionMessageFunctionToolCallParam] = []
-        reasoning_details: list[dict[str, Any]] = []
-        for item in message.parts:
-            if isinstance(item, TextPart):
-                texts.append(item.content)
-            elif isinstance(item, ThinkingPart):
-                if item.provider_name == self.system:
-                    reasoning_details.append(_into_reasoning_detail(item).model_dump())
-                elif content := item.content:  # pragma: lax no cover
-                    start_tag, end_tag = self.profile.thinking_tags
-                    texts.append('\n'.join([start_tag, content, end_tag]))
-                else:
-                    pass
-            elif isinstance(item, ToolCallPart):
-                tool_calls.append(self._map_tool_call(item))
-            elif isinstance(item, BuiltinToolCallPart | BuiltinToolReturnPart):  # pragma: no cover
-                pass
-            elif isinstance(item, FilePart):  # pragma: no cover
-                pass
-            else:
-                assert_never(item)
-        message_param = chat.ChatCompletionAssistantMessageParam(role='assistant')
-        if texts:
-            message_param['content'] = '\n\n'.join(texts)
+    def _map_response_thinking_part(self, ctx: OpenAIChatModel._MapModelResposeContext, item: ThinkingPart) -> None:
+        assert isinstance(ctx, self._MapModelResposeContext)
+        if item.provider_name == self.system:
+            ctx.reasoning_details.append(_into_reasoning_detail(item).model_dump())
+        elif content := item.content:  # pragma: lax no cover
+            start_tag, end_tag = self.profile.thinking_tags
+            ctx.texts.append('\n'.join([start_tag, content, end_tag]))
         else:
-            message_param['content'] = None
-        if tool_calls:
-            message_param['tool_calls'] = tool_calls
-        if reasoning_details:
-            message_param['reasoning_details'] = reasoning_details  # type: ignore[reportGeneralTypeIssues]
-        return message_param
+            pass
 
     @property
     @override