Expose some agent testing utils to users

Author: donald-pinckney

temporalio/contrib/openai_agents/__init__.py

@@ -21,15 +21,13 @@
 from temporalio.contrib.openai_agents._temporal_openai_agents import (
     OpenAIAgentsPlugin,
     OpenAIPayloadConverter,
-    TestModel,
-    TestModelProvider,
 )
 from temporalio.contrib.openai_agents._trace_interceptor import (
     OpenAIAgentsTracingInterceptor,
 )
 from temporalio.contrib.openai_agents.workflow import AgentsWorkflowError
 
-from . import workflow
+from . import test, workflow
 
 __all__ = [
     "AgentsWorkflowError",
@@ -38,7 +36,6 @@
     "OpenAIPayloadConverter",
     "StatelessMCPServerProvider",
     "StatefulMCPServerProvider",
-    "TestModel",
-    "TestModelProvider",
+    "test",
     "workflow",
 ]

temporalio/contrib/openai_agents/_temporal_openai_agents.py

@@ -6,43 +6,23 @@
 from datetime import timedelta
 from typing import AsyncIterator, Callable, Optional, Sequence, Union
 
-from agents import (
-    AgentOutputSchemaBase,
-    Handoff,
-    Model,
-    ModelProvider,
-    ModelResponse,
-    ModelSettings,
-    ModelTracing,
-    Tool,
-    TResponseInputItem,
-    set_trace_provider,
-)
-from agents.items import TResponseStreamEvent
+from agents import ModelProvider, set_trace_provider
 from agents.run import get_default_agent_runner, set_default_agent_runner
 from agents.tracing import get_trace_provider
 from agents.tracing.provider import DefaultTraceProvider
 
 from temporalio.contrib.openai_agents._invoke_model_activity import ModelActivity
 from temporalio.contrib.openai_agents._model_parameters import ModelActivityParameters
-from temporalio.contrib.openai_agents._openai_runner import (
-    TemporalOpenAIRunner,
-)
+from temporalio.contrib.openai_agents._openai_runner import TemporalOpenAIRunner
 from temporalio.contrib.openai_agents._temporal_trace_provider import (
     TemporalTraceProvider,
 )
 from temporalio.contrib.openai_agents._trace_interceptor import (
     OpenAIAgentsTracingInterceptor,
 )
 from temporalio.contrib.openai_agents.workflow import AgentsWorkflowError
-from temporalio.contrib.pydantic import (
-    PydanticPayloadConverter,
-    ToJsonOptions,
-)
-from temporalio.converter import (
-    DataConverter,
-    DefaultPayloadConverter,
-)
+from temporalio.contrib.pydantic import PydanticPayloadConverter, ToJsonOptions
+from temporalio.converter import DataConverter, DefaultPayloadConverter
 from temporalio.plugin import SimplePlugin
 from temporalio.worker import WorkflowRunner
 from temporalio.worker.workflow_sandbox import SandboxedWorkflowRunner
@@ -103,58 +83,6 @@ def set_open_ai_agent_temporal_overrides(
         set_trace_provider(previous_trace_provider or DefaultTraceProvider())
 
 
-class TestModelProvider(ModelProvider):
-    """Test model provider which simply returns the given module."""
-
-    __test__ = False
-
-    def __init__(self, model: Model):
-        """Initialize a test model provider with a model."""
-        self._model = model
-
-    def get_model(self, model_name: Union[str, None]) -> Model:
-        """Get a model from the model provider."""
-        return self._model
-
-
-class TestModel(Model):
-    """Test model for use mocking model responses."""
-
-    __test__ = False
-
-    def __init__(self, fn: Callable[[], ModelResponse]) -> None:
-        """Initialize a test model with a callable."""
-        self.fn = fn
-
-    async def get_response(
-        self,
-        system_instructions: Union[str, None],
-        input: Union[str, list[TResponseInputItem]],
-        model_settings: ModelSettings,
-        tools: list[Tool],
-        output_schema: Union[AgentOutputSchemaBase, None],
-        handoffs: list[Handoff],
-        tracing: ModelTracing,
-        **kwargs,
-    ) -> ModelResponse:
-        """Get a response from the model."""
-        return self.fn()
-
-    def stream_response(
-        self,
-        system_instructions: Optional[str],
-        input: Union[str, list[TResponseInputItem]],
-        model_settings: ModelSettings,
-        tools: list[Tool],
-        output_schema: Optional[AgentOutputSchemaBase],
-        handoffs: list[Handoff],
-        tracing: ModelTracing,
-        **kwargs,
-    ) -> AsyncIterator[TResponseStreamEvent]:
-        """Get a streamed response from the model. Unimplemented."""
-        raise NotImplementedError()
-
-
 class OpenAIPayloadConverter(PydanticPayloadConverter):
     """PayloadConverter for OpenAI agents."""
 

temporalio/contrib/openai_agents/test.py

@@ -0,0 +1,136 @@
+from typing import AsyncIterator, Callable, Optional, Union
+
+from agents import (
+    AgentOutputSchemaBase,
+    Handoff,
+    Model,
+    ModelProvider,
+    ModelResponse,
+    ModelSettings,
+    ModelTracing,
+    Tool,
+    TResponseInputItem,
+    Usage,
+)
+from agents.items import TResponseOutputItem, TResponseStreamEvent
+from openai.types.responses import (
+    ResponseFunctionToolCall,
+    ResponseOutputMessage,
+    ResponseOutputText,
+)
+
+
+class ResponseBuilders:
+    """
+    Builders for creating model responses for testing.
+    """
+
+    @staticmethod
+    def model_response(output: TResponseOutputItem) -> ModelResponse:
+        return ModelResponse(
+            output=[output],
+            usage=Usage(),
+            response_id=None,
+        )
+
+    @staticmethod
+    def response_output_message(text: str) -> ResponseOutputMessage:
+        return ResponseOutputMessage(
+            id="",
+            content=[
+                ResponseOutputText(
+                    text=text,
+                    annotations=[],
+                    type="output_text",
+                )
+            ],
+            role="assistant",
+            status="completed",
+            type="message",
+        )
+
+    @staticmethod
+    def tool_call(arguments: str, name: str) -> ModelResponse:
+        return ResponseBuilders.model_response(
+            ResponseFunctionToolCall(
+                arguments=arguments,
+                call_id="call",
+                name=name,
+                type="function_call",
+                id="id",
+                status="completed",
+            )
+        )
+
+    @staticmethod
+    def output_message(text: str) -> ModelResponse:
+        return ResponseBuilders.model_response(
+            ResponseBuilders.response_output_message(text)
+        )
+
+
+class TestModelProvider(ModelProvider):
+    """Test model provider which simply returns the given module."""
+
+    __test__ = False
+
+    def __init__(self, model: Model):
+        """Initialize a test model provider with a model."""
+        self._model = model
+
+    def get_model(self, model_name: Union[str, None]) -> Model:
+        """Get a model from the model provider."""
+        return self._model
+
+
+class TestModel(Model):
+    """Test model for use mocking model responses."""
+
+    __test__ = False
+
+    def __init__(self, fn: Callable[[], ModelResponse]) -> None:
+        """Initialize a test model with a callable."""
+        self.fn = fn
+
+    async def get_response(
+        self,
+        system_instructions: Union[str, None],
+        input: Union[str, list[TResponseInputItem]],
+        model_settings: ModelSettings,
+        tools: list[Tool],
+        output_schema: Union[AgentOutputSchemaBase, None],
+        handoffs: list[Handoff],
+        tracing: ModelTracing,
+        **kwargs,
+    ) -> ModelResponse:
+        """Get a response from the model."""
+        return self.fn()
+
+    def stream_response(
+        self,
+        system_instructions: Optional[str],
+        input: Union[str, list[TResponseInputItem]],
+        model_settings: ModelSettings,
+        tools: list[Tool],
+        output_schema: Optional[AgentOutputSchemaBase],
+        handoffs: list[Handoff],
+        tracing: ModelTracing,
+        **kwargs,
+    ) -> AsyncIterator[TResponseStreamEvent]:
+        """Get a streamed response from the model. Unimplemented."""
+        raise NotImplementedError()
+
+
+class StaticTestModel(TestModel):
+    """Static test model for use mocking model responses.
+    Set a responses attribute to a list of model responses, which will be returned sequentially.
+    """
+
+    __test__ = False
+    responses: list[ModelResponse] = []
+
+    def __init__(
+        self,
+    ) -> None:
+        self._responses = iter(self.responses)
+        super().__init__(lambda: next(self._responses))