Fix TemporalAgent dropping model-specific ModelSettings (e.g. openai_reasoning_effort) (#2938)

Author: DouweM

pydantic_ai_slim/pydantic_ai/durable_exec/temporal/_model.py

@@ -4,7 +4,7 @@
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from datetime import datetime
-from typing import Any
+from typing import Any, cast
 
 from pydantic import ConfigDict, with_config
 from temporalio import activity, workflow
@@ -30,7 +30,8 @@
 @with_config(ConfigDict(arbitrary_types_allowed=True))
 class _RequestParams:
     messages: list[ModelMessage]
-    model_settings: ModelSettings | None
+    # `model_settings` can't be a `ModelSettings` because Temporal would end up dropping fields only defined on its subclasses.
+    model_settings: dict[str, Any] | None
     model_request_parameters: ModelRequestParameters
     serialized_run_context: Any
 
@@ -82,7 +83,11 @@ def __init__(
 
         @activity.defn(name=f'{activity_name_prefix}__model_request')
         async def request_activity(params: _RequestParams) -> ModelResponse:
-            return await self.wrapped.request(params.messages, params.model_settings, params.model_request_parameters)
+            return await self.wrapped.request(
+                params.messages,
+                cast(ModelSettings | None, params.model_settings),
+                params.model_request_parameters,
+            )
 
         self.request_activity = request_activity
 
@@ -92,7 +97,10 @@ async def request_stream_activity(params: _RequestParams, deps: AgentDepsT) -> M
 
             run_context = self.run_context_type.deserialize_run_context(params.serialized_run_context, deps=deps)
             async with self.wrapped.request_stream(
-                params.messages, params.model_settings, params.model_request_parameters, run_context
+                params.messages,
+                cast(ModelSettings | None, params.model_settings),
+                params.model_request_parameters,
+                run_context,
             ) as streamed_response:
                 await self.event_stream_handler(run_context, streamed_response)
 
@@ -124,7 +132,7 @@ async def request(
             activity=self.request_activity,
             arg=_RequestParams(
                 messages=messages,
-                model_settings=model_settings,
+                model_settings=cast(dict[str, Any] | None, model_settings),
                 model_request_parameters=model_request_parameters,
                 serialized_run_context=None,
             ),
@@ -161,7 +169,7 @@ async def request_stream(
             args=[
                 _RequestParams(
                     messages=messages,
-                    model_settings=model_settings,
+                    model_settings=cast(dict[str, Any] | None, model_settings),
                     model_request_parameters=model_request_parameters,
                     serialized_run_context=serialized_run_context,
                 ),

tests/test_temporal.py

@@ -12,7 +12,7 @@
 import pytest
 from pydantic import BaseModel
 
-from pydantic_ai import Agent, RunContext
+from pydantic_ai import Agent, ModelSettings, RunContext
 from pydantic_ai.direct import model_request_stream
 from pydantic_ai.exceptions import ApprovalRequired, CallDeferred, ModelRetry, UserError
 from pydantic_ai.messages import (
@@ -33,6 +33,7 @@
     UserPromptPart,
 )
 from pydantic_ai.models import Model, cached_async_http_client
+from pydantic_ai.models.function import AgentInfo, FunctionModel
 from pydantic_ai.run import AgentRunResult
 from pydantic_ai.tools import DeferredToolRequests, DeferredToolResults, ToolDefinition
 from pydantic_ai.toolsets import ExternalToolset, FunctionToolset
@@ -1931,3 +1932,44 @@ async def test_temporal_agent_with_model_retry(allow_model_requests: None, clien
                 ),
             ]
         )
+
+
+class CustomModelSettings(ModelSettings, total=False):
+    custom_setting: str
+
+
+def return_settings(messages: list[ModelMessage], agent_info: AgentInfo) -> ModelResponse:
+    return ModelResponse(parts=[TextPart(str(agent_info.model_settings))])
+
+
+model_settings = CustomModelSettings(max_tokens=123, custom_setting='custom_value')
+model = FunctionModel(return_settings, settings=model_settings)
+
+settings_agent = Agent(model, name='settings_agent')
+
+# This needs to be done before the `TemporalAgent` is bound to the workflow.
+settings_temporal_agent = TemporalAgent(settings_agent, activity_config=BASE_ACTIVITY_CONFIG)
+
+
+@workflow.defn
+class SettingsAgentWorkflow:
+    @workflow.run
+    async def run(self, prompt: str) -> str:
+        result = await settings_temporal_agent.run(prompt)
+        return result.output
+
+
+async def test_custom_model_settings(allow_model_requests: None, client: Client):
+    async with Worker(
+        client,
+        task_queue=TASK_QUEUE,
+        workflows=[SettingsAgentWorkflow],
+        plugins=[AgentPlugin(settings_temporal_agent)],
+    ):
+        output = await client.execute_workflow(  # pyright: ignore[reportUnknownMemberType]
+            SettingsAgentWorkflow.run,
+            args=['Give me those settings'],
+            id=SettingsAgentWorkflow.__name__,
+            task_queue=TASK_QUEUE,
+        )
+        assert output == snapshot("{'max_tokens': 123, 'custom_setting': 'custom_value'}")