add Model.prepare_request to merge + customize

Author: moritzwilksch

pydantic_ai_slim/pydantic_ai/direct.py

@@ -81,7 +81,7 @@ async def main():
     return await model_instance.request(
         messages,
         model_settings,
-        model_instance.customize_request_parameters(model_request_parameters or models.ModelRequestParameters()),
+        model_request_parameters or models.ModelRequestParameters(),
     )
 
 
@@ -193,7 +193,7 @@ async def main():
     return model_instance.request_stream(
         messages,
         model_settings,
-        model_instance.customize_request_parameters(model_request_parameters or models.ModelRequestParameters()),
+        model_request_parameters or models.ModelRequestParameters(),
     )
 
 

pydantic_ai_slim/pydantic_ai/models/__init__.py

@@ -41,7 +41,7 @@
 from ..output import OutputMode
 from ..profiles import DEFAULT_PROFILE, ModelProfile, ModelProfileSpec
 from ..profiles._json_schema import JsonSchemaTransformer
-from ..settings import ModelSettings
+from ..settings import ModelSettings, merge_model_settings
 from ..tools import ToolDefinition
 from ..usage import RequestUsage
 
@@ -390,6 +390,24 @@ def customize_request_parameters(self, model_request_parameters: ModelRequestPar
 
         return model_request_parameters
 
+    def prepare_request(
+        self,
+        model_settings: ModelSettings | None,
+        model_request_parameters: ModelRequestParameters | None,
+    ) -> tuple[ModelSettings | None, ModelRequestParameters]:
+        """Prepare request inputs before they are passed to the provider.
+
+        This merges the given ``model_settings`` with the model's own ``settings`` attribute and ensures
+        ``customize_request_parameters`` is applied to the resolved
+        [`ModelRequestParameters`][pydantic_ai.models.ModelRequestParameters]. Subclasses can override this method if
+        they need to customize the preparation flow further, but most implementations should simply call
+        ``self.prepare_request(...)`` at the start of their ``request`` (and related) methods.
+        """
+        merged_settings = merge_model_settings(self.settings, model_settings)
+        resolved_parameters = model_request_parameters or ModelRequestParameters()
+        customized_parameters = self.customize_request_parameters(resolved_parameters)
+        return merged_settings, customized_parameters
+
     @property
     @abstractmethod
     def model_name(self) -> str:

pydantic_ai_slim/pydantic_ai/models/anthropic.py

@@ -39,7 +39,7 @@
 from ..profiles import ModelProfileSpec
 from ..providers import Provider, infer_provider
 from ..providers.anthropic import AsyncAnthropicClient
-from ..settings import ModelSettings, merge_model_settings
+from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import Model, ModelRequestParameters, StreamedResponse, check_allow_model_requests, download_item, get_user_agent
 
@@ -205,7 +205,10 @@ async def request(
         model_request_parameters: ModelRequestParameters,
     ) -> ModelResponse:
         check_allow_model_requests()
-        model_settings = merge_model_settings(self.settings, model_settings)
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
         response = await self._messages_create(
             messages, False, cast(AnthropicModelSettings, model_settings or {}), model_request_parameters
         )
@@ -221,7 +224,10 @@ async def request_stream(
         run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         check_allow_model_requests()
-        model_settings = merge_model_settings(self.settings, model_settings)
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
         response = await self._messages_create(
             messages, True, cast(AnthropicModelSettings, model_settings or {}), model_request_parameters
         )

pydantic_ai_slim/pydantic_ai/models/bedrock.py

@@ -43,7 +43,7 @@
 from pydantic_ai.models import Model, ModelRequestParameters, StreamedResponse, download_item
 from pydantic_ai.providers import Provider, infer_provider
 from pydantic_ai.providers.bedrock import BedrockModelProfile
-from pydantic_ai.settings import ModelSettings, merge_model_settings
+from pydantic_ai.settings import ModelSettings
 from pydantic_ai.tools import ToolDefinition
 
 if TYPE_CHECKING:
@@ -264,7 +264,10 @@ async def request(
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
     ) -> ModelResponse:
-        model_settings = merge_model_settings(self.settings, model_settings)
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
         settings = cast(BedrockModelSettings, model_settings or {})
         response = await self._messages_create(messages, False, settings, model_request_parameters)
         model_response = await self._process_response(response)
@@ -278,7 +281,10 @@ async def request_stream(
         model_request_parameters: ModelRequestParameters,
         run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
-        model_settings = merge_model_settings(self.settings, model_settings)
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
         settings = cast(BedrockModelSettings, model_settings or {})
         response = await self._messages_create(messages, True, settings, model_request_parameters)
         yield BedrockStreamedResponse(

pydantic_ai_slim/pydantic_ai/models/cohere.py

@@ -28,7 +28,7 @@
 )
 from ..profiles import ModelProfileSpec
 from ..providers import Provider, infer_provider
-from ..settings import ModelSettings, merge_model_settings
+from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import Model, ModelRequestParameters, check_allow_model_requests
 
@@ -165,7 +165,10 @@ async def request(
         model_request_parameters: ModelRequestParameters,
     ) -> ModelResponse:
         check_allow_model_requests()
-        model_settings = merge_model_settings(self.settings, model_settings)
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
         response = await self._chat(messages, cast(CohereModelSettings, model_settings or {}), model_request_parameters)
         model_response = self._process_response(response)
         return model_response

pydantic_ai_slim/pydantic_ai/models/fallback.py

@@ -11,7 +11,6 @@
 from pydantic_ai.models.instrumented import InstrumentedModel
 
 from ..exceptions import FallbackExceptionGroup, ModelHTTPError
-from ..settings import merge_model_settings
 from . import KnownModelName, Model, ModelRequestParameters, StreamedResponse, infer_model
 
 if TYPE_CHECKING:
@@ -78,10 +77,8 @@ async def request(
         exceptions: list[Exception] = []
 
         for model in self.models:
-            customized_model_request_parameters = model.customize_request_parameters(model_request_parameters)
-            merged_settings = merge_model_settings(model.settings, model_settings)
             try:
-                response = await model.request(messages, merged_settings, customized_model_request_parameters)
+                response = await model.request(messages, model_settings, model_request_parameters)
             except Exception as exc:
                 if self._fallback_on(exc):
                     exceptions.append(exc)
@@ -105,14 +102,10 @@ async def request_stream(
         exceptions: list[Exception] = []
 
         for model in self.models:
-            customized_model_request_parameters = model.customize_request_parameters(model_request_parameters)
-            merged_settings = merge_model_settings(model.settings, model_settings)
             async with AsyncExitStack() as stack:
                 try:
                     response = await stack.enter_async_context(
-                        model.request_stream(
-                            messages, merged_settings, customized_model_request_parameters, run_context
-                        )
+                        model.request_stream(messages, model_settings, model_request_parameters, run_context)
                     )
                 except Exception as exc:
                     if self._fallback_on(exc):

pydantic_ai_slim/pydantic_ai/models/function.py

@@ -32,7 +32,7 @@
     UserPromptPart,
 )
 from ..profiles import ModelProfile, ModelProfileSpec
-from ..settings import ModelSettings, merge_model_settings
+from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import Model, ModelRequestParameters, StreamedResponse
 
@@ -125,7 +125,10 @@ async def request(
         model_settings: ModelSettings | None,
         model_request_parameters: ModelRequestParameters,
     ) -> ModelResponse:
-        model_settings = merge_model_settings(self.settings, model_settings)
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
         agent_info = AgentInfo(
             function_tools=model_request_parameters.function_tools,
             allow_text_output=model_request_parameters.allow_text_output,
@@ -155,7 +158,10 @@ async def request_stream(
         model_request_parameters: ModelRequestParameters,
         run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
-        model_settings = merge_model_settings(self.settings, model_settings)
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
         agent_info = AgentInfo(
             function_tools=model_request_parameters.function_tools,
             allow_text_output=model_request_parameters.allow_text_output,

pydantic_ai_slim/pydantic_ai/models/gemini.py

@@ -38,7 +38,7 @@
 )
 from ..profiles import ModelProfileSpec
 from ..providers import Provider, infer_provider
-from ..settings import ModelSettings, merge_model_settings
+from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import Model, ModelRequestParameters, StreamedResponse, check_allow_model_requests, download_item, get_user_agent
 
@@ -155,7 +155,10 @@ async def request(
         model_request_parameters: ModelRequestParameters,
     ) -> ModelResponse:
         check_allow_model_requests()
-        model_settings = merge_model_settings(self.settings, model_settings)
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
         async with self._make_request(
             messages, False, cast(GeminiModelSettings, model_settings or {}), model_request_parameters
         ) as http_response:
@@ -172,7 +175,10 @@ async def request_stream(
         run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         check_allow_model_requests()
-        model_settings = merge_model_settings(self.settings, model_settings)
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
         async with self._make_request(
             messages, True, cast(GeminiModelSettings, model_settings or {}), model_request_parameters
         ) as http_response:

pydantic_ai_slim/pydantic_ai/models/google.py

@@ -37,7 +37,7 @@
 )
 from ..profiles import ModelProfileSpec
 from ..providers import Provider
-from ..settings import ModelSettings, merge_model_settings
+from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
     Model,
@@ -225,7 +225,10 @@ async def request(
         model_request_parameters: ModelRequestParameters,
     ) -> ModelResponse:
         check_allow_model_requests()
-        model_settings = merge_model_settings(self.settings, model_settings)
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
         model_settings = cast(GoogleModelSettings, model_settings or {})
         response = await self._generate_content(messages, False, model_settings, model_request_parameters)
         return self._process_response(response)
@@ -237,6 +240,10 @@ async def count_tokens(
         model_request_parameters: ModelRequestParameters,
     ) -> usage.RequestUsage:
         check_allow_model_requests()
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
         model_settings = cast(GoogleModelSettings, model_settings or {})
         contents, generation_config = await self._build_content_and_config(
             messages, model_settings, model_request_parameters
@@ -292,7 +299,10 @@ async def request_stream(
         run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         check_allow_model_requests()
-        model_settings = merge_model_settings(self.settings, model_settings)
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
         model_settings = cast(GoogleModelSettings, model_settings or {})
         response = await self._generate_content(messages, True, model_settings, model_request_parameters)
         yield await self._process_streamed_response(response, model_request_parameters)  # type: ignore

pydantic_ai_slim/pydantic_ai/models/groq.py

@@ -41,7 +41,7 @@
 from ..profiles import ModelProfile, ModelProfileSpec
 from ..profiles.groq import GroqModelProfile
 from ..providers import Provider, infer_provider
-from ..settings import ModelSettings, merge_model_settings
+from ..settings import ModelSettings
 from ..tools import ToolDefinition
 from . import (
     Model,
@@ -182,7 +182,10 @@ async def request(
         model_request_parameters: ModelRequestParameters,
     ) -> ModelResponse:
         check_allow_model_requests()
-        model_settings = merge_model_settings(self.settings, model_settings)
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
         try:
             response = await self._completions_create(
                 messages, False, cast(GroqModelSettings, model_settings or {}), model_request_parameters
@@ -219,7 +222,10 @@ async def request_stream(
         run_context: RunContext[Any] | None = None,
     ) -> AsyncIterator[StreamedResponse]:
         check_allow_model_requests()
-        model_settings = merge_model_settings(self.settings, model_settings)
+        model_settings, model_request_parameters = self.prepare_request(
+            model_settings,
+            model_request_parameters,
+        )
         response = await self._completions_create(
             messages, True, cast(GroqModelSettings, model_settings or {}), model_request_parameters
         )