Cleanup and fix when model is provided via run_config

Author: tconley1428

temporalio/contrib/openai_agents/_model_parameters.py

@@ -47,7 +47,12 @@ class ModelActivityParameters:
         Union[
             str,
             Callable[
-                [Agent[Any], Optional[str], Union[str, list[TResponseInputItem]]], str
+                [
+                    Optional[Agent[Any]],
+                    Optional[str],
+                    Union[str, list[TResponseInputItem]],
+                ],
+                str,
             ],
         ]
     ] = None

temporalio/contrib/openai_agents/_openai_runner.py

@@ -1,3 +1,4 @@
+import dataclasses
 import json
 import typing
 from typing import Any, Optional, Union
@@ -78,29 +79,30 @@ async def run(
         if run_config is None:
             run_config = RunConfig()
 
-        def model_name(agent: Agent[Any]) -> Optional[str]:
-            name = run_config.model or agent.model
-            if name is not None and not isinstance(name, str):
-                print("Name: ", name, " Agent: ", agent)
+        if run_config.model:
+            if not isinstance(run_config.model, str):
                 raise ValueError(
-                    "Temporal workflows require a model name to be a string in the run config and/or agent."
+                    "Temporal workflows require a model name to be a string in the run config."
                 )
-            return name
+            run_config = dataclasses.replace(
+                run_config,
+                model=_TemporalModelStub(
+                    run_config.model, model_params=self.model_params, agent=None
+                ),
+            )
 
+        # Recursively replace models in all agents
         def convert_agent(agent: Agent[Any]) -> None:
-            print("Model: ", agent.model)
-
             # Short circuit if this model was already replaced to prevent looping from circular handoffs
             if isinstance(agent.model, _TemporalModelStub):
                 return
 
-            name = model_name(agent)
+            name = _model_name(agent)
             agent.model = _TemporalModelStub(
                 model_name=name,
                 model_params=self.model_params,
                 agent=agent,
             )
-            print("Model after replace: ", agent.model)
 
             for handoff in agent.handoffs:
                 if isinstance(handoff, Agent):
@@ -159,3 +161,12 @@ def run_streamed(
                 **kwargs,
             )
         raise RuntimeError("Temporal workflows do not support streaming.")
+
+
+def _model_name(agent: Agent[Any]) -> Optional[str]:
+    name = agent.model
+    if name is not None and not isinstance(name, str):
+        raise ValueError(
+            "Temporal workflows require a model name to be a string in the agent."
+        )
+    return name

temporalio/contrib/openai_agents/_temporal_model_stub.py

@@ -51,7 +51,7 @@ def __init__(
         model_name: Optional[str],
         *,
         model_params: ModelActivityParameters,
-        agent: Agent[Any],
+        agent: Optional[Agent[Any]],
     ) -> None:
         self.model_name = model_name
         self.model_params = model_params
@@ -70,8 +70,6 @@ async def get_response(
         previous_response_id: Optional[str],
         prompt: Optional[ResponsePromptParam],
     ) -> ModelResponse:
-        print("Model stub invocation:", self.model_name)
-
         def make_tool_info(tool: Tool) -> ToolInput:
             if isinstance(
                 tool,
@@ -149,8 +147,10 @@ def make_tool_info(tool: Tool) -> ToolInput:
                     )
                 )
             )
-        else:
+        elif self.agent:
             summary = self.agent.name
+        else:
+            summary = None
 
         return await workflow.execute_activity_method(
             ModelActivity.invoke_model_activity,

tests/contrib/openai_agents/test_openai.py

@@ -32,6 +32,7 @@
     OpenAIChatCompletionsModel,
     OpenAIResponsesModel,
     OutputGuardrailTripwireTriggered,
+    RunConfig,
     RunContextWrapper,
     Runner,
     SQLiteSession,
@@ -861,7 +862,6 @@ def init_agents() -> Agent[AirlineAgentContext]:
 
     seat_booking_agent = Agent[AirlineAgentContext](
         name="Seat Booking Agent",
-        model="gpt-4o-mini",
         handoff_description="A helpful agent that can update a seat on a flight.",
         instructions=f"""{RECOMMENDED_PROMPT_PREFIX}
         You are a seat booking agent. If you are speaking to a customer, you probably were transferred to from the triage agent.
@@ -988,13 +988,14 @@ async def test_customer_service_workflow(client: Client, use_local_model: bool):
     if not use_local_model and not os.environ.get("OPENAI_API_KEY"):
         pytest.skip("No openai API key")
     new_config = client.config()
-    provider = TestModelProvider(CustomerServiceModel())
     new_config["plugins"] = [
         openai_agents.OpenAIAgentsPlugin(
             model_params=ModelActivityParameters(
                 start_to_close_timeout=timedelta(seconds=30)
             ),
-            model_provider=provider if use_local_model else None,
+            model_provider=TestModelProvider(CustomerServiceModel())
+            if use_local_model
+            else None,
         )
     ]
     client = Client(**new_config)
@@ -2066,7 +2067,7 @@ class MultipleModelsModel(StaticTestModel):
 @workflow.defn
 class MultipleModelWorkflow:
     @workflow.run
-    async def run(self):
+    async def run(self, use_run_config: bool):
         underling = Agent[None](
             name="Underling",
             instructions="You do all the work you are told.",
@@ -2081,6 +2082,7 @@ async def run(self):
         result = await Runner.run(
             starting_agent=starting_agent,
             input="Have you cleaned the store room yet?",
+            run_config=RunConfig(model="gpt-4o") if use_run_config else None,
         )
         return result.final_output
 
@@ -2104,9 +2106,40 @@ async def test_multiple_models(client: Client):
     ) as worker:
         workflow_handle = await client.start_workflow(
             MultipleModelWorkflow.run,
+            False,
             id=f"multiple-model-{uuid.uuid4()}",
             task_queue=worker.task_queue,
-            execution_timeout=timedelta(seconds=120),
+            execution_timeout=timedelta(seconds=10),
         )
         result = await workflow_handle.result()
         assert provider.model_names == {None, "gpt-4o-mini"}
+
+
+async def test_run_config_models(client: Client):
+    provider = AssertDifferentModelProvider(MultipleModelsModel())
+    new_config = client.config()
+    new_config["plugins"] = [
+        openai_agents.OpenAIAgentsPlugin(
+            model_params=ModelActivityParameters(
+                start_to_close_timeout=timedelta(seconds=120)
+            ),
+            model_provider=provider,
+        )
+    ]
+    client = Client(**new_config)
+
+    async with new_worker(
+        client,
+        MultipleModelWorkflow,
+    ) as worker:
+        workflow_handle = await client.start_workflow(
+            MultipleModelWorkflow.run,
+            True,
+            id=f"run-config-model-{uuid.uuid4()}",
+            task_queue=worker.task_queue,
+            execution_timeout=timedelta(seconds=10),
+        )
+        result = await workflow_handle.result()
+
+        # Only the model from the runconfig override is used
+        assert provider.model_names == {"gpt-4o"}