OpenAI tool context (#942)

* Toying with context input to tools

* Doc update

* Remove import

Author: tconley1428

temporalio/contrib/openai_agents/temporal_tools.py

@@ -41,7 +41,8 @@ def activity_as_tool(
     This function takes a Temporal activity function and converts it into an
     OpenAI agent tool that can be used by the agent to execute the activity
     during workflow execution. The tool will automatically handle the conversion
-    of inputs and outputs between the agent and the activity.
+    of inputs and outputs between the agent and the activity. Note that if you take a context,
+    mutation will not be persisted, as the activity may not be running in the same location.
 
     Args:
         fn: A Temporal activity function to convert to a tool.
@@ -85,6 +86,11 @@ async def run_activity(ctx: RunContextWrapper[Any], input: str) -> Any:
 
         # Activities don't support keyword only arguments, so we can ignore the kwargs_dict return
         args, _ = schema.to_call_args(schema.params_pydantic_model(**json_data))
+
+        # Add the context to the arguments if it takes that
+        if schema.takes_context:
+            args = [ctx] + args
+
         result = await workflow.execute_activity(
             fn,
             args=args,

tests/contrib/openai_agents/test_openai.py

@@ -221,6 +221,14 @@ async def get_weather_object(input: WeatherInput) -> Weather:
     )
 
 
+@activity.defn
+async def get_weather_context(ctx: RunContextWrapper[str], city: str) -> Weather:
+    """
+    Get the weather for a given city.
+    """
+    return Weather(city=city, temperature_range="14-20C", conditions=ctx.context)
+
+
 class TestWeatherModel(TestModel):
     responses = [
         ModelResponse(
@@ -265,6 +273,20 @@ class TestWeatherModel(TestModel):
             usage=Usage(),
             response_id=None,
         ),
+        ModelResponse(
+            output=[
+                ResponseFunctionToolCall(
+                    arguments='{"city":"Tokyo"}',
+                    call_id="call",
+                    name="get_weather_context",
+                    type="function_call",
+                    id="id",
+                    status="completed",
+                )
+            ],
+            usage=Usage(),
+            response_id=None,
+        ),
         ModelResponse(
             output=[
                 ResponseOutputMessage(
@@ -304,9 +326,14 @@ async def run(self, question: str) -> str:
                 activity_as_tool(
                     get_weather_country, start_to_close_timeout=timedelta(seconds=10)
                 ),
+                activity_as_tool(
+                    get_weather_context, start_to_close_timeout=timedelta(seconds=10)
+                ),
             ],
         )  # type: Agent
-        result = await Runner.run(starting_agent=agent, input=question)
+        result = await Runner.run(
+            starting_agent=agent, input=question, context="Stormy"
+        )
         return result.final_output
 
 
@@ -337,6 +364,7 @@ async def test_tool_workflow(client: Client, use_local_model: bool):
                 get_weather,
                 get_weather_object,
                 get_weather_country,
+                get_weather_context,
             ],
             interceptors=[OpenAIAgentsTracingInterceptor()],
         ) as worker:
@@ -357,7 +385,7 @@ async def test_tool_workflow(client: Client, use_local_model: bool):
                     if e.HasField("activity_task_completed_event_attributes"):
                         events.append(e)
 
-                assert len(events) == 7
+                assert len(events) == 9
                 assert (
                     "function_call"
                     in events[0]
@@ -395,11 +423,23 @@ async def test_tool_workflow(client: Client, use_local_model: bool):
                     .data.decode()
                 )
                 assert (
-                    "Test weather result"
+                    "function_call"
                     in events[6]
                     .activity_task_completed_event_attributes.result.payloads[0]
                     .data.decode()
                 )
+                assert (
+                    "Stormy"
+                    in events[7]
+                    .activity_task_completed_event_attributes.result.payloads[0]
+                    .data.decode()
+                )
+                assert (
+                    "Test weather result"
+                    in events[8]
+                    .activity_task_completed_event_attributes.result.payloads[0]
+                    .data.decode()
+                )
 
 
 class TestPlannerModel(OpenAIResponsesModel):