Add _DurableModelStub

Author: AnatoliB

samples-v2/openai_agents/durable_model_stub.py

@@ -0,0 +1,141 @@
+from __future__ import annotations
+import azure.durable_functions as df
+
+import logging
+from typing import Optional
+
+logger = logging.getLogger(__name__)
+
+from typing import Any, AsyncIterator, Union, cast
+
+from agents import (
+    AgentOutputSchema,
+    AgentOutputSchemaBase,
+    CodeInterpreterTool,
+    FileSearchTool,
+    FunctionTool,
+    Handoff,
+    HostedMCPTool,
+    ImageGenerationTool,
+    Model,
+    ModelResponse,
+    ModelSettings,
+    ModelTracing,
+    Tool,
+    TResponseInputItem,
+    WebSearchTool,
+)
+from agents.items import TResponseStreamEvent
+from openai.types.responses.response_prompt_param import ResponsePromptParam
+
+
+class _DurableModelStub(Model):
+    def __init__(
+        self,
+        model_name: Optional[str],
+        durable_orchestration_context: df.DurableOrchestrationContext,
+    ) -> None:
+        self.model_name = model_name
+        self.durable_orchestration_context = durable_orchestration_context
+
+    async def get_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,
+        *,
+        previous_response_id: Optional[str],
+        prompt: Optional[ResponsePromptParam],
+    ) -> ModelResponse:
+        def make_tool_info(tool: Tool) -> ToolInput:
+            if isinstance(
+                tool,
+                (
+                    FileSearchTool,
+                    WebSearchTool,
+                    ImageGenerationTool,
+                    CodeInterpreterTool,
+                ),
+            ):
+                return tool
+            elif isinstance(tool, HostedMCPTool):
+                return HostedMCPToolInput(tool_config=tool.tool_config)
+            elif isinstance(tool, FunctionTool):
+                return FunctionToolInput(
+                    name=tool.name,
+                    description=tool.description,
+                    params_json_schema=tool.params_json_schema,
+                    strict_json_schema=tool.strict_json_schema,
+                )
+            else:
+                raise ValueError(f"Unsupported tool type: {tool.name}")
+
+        tool_infos = [make_tool_info(x) for x in tools]
+        handoff_infos = [
+            HandoffInput(
+                tool_name=x.tool_name,
+                tool_description=x.tool_description,
+                input_json_schema=x.input_json_schema,
+                agent_name=x.agent_name,
+                strict_json_schema=x.strict_json_schema,
+            )
+            for x in handoffs
+        ]
+        if output_schema is not None and not isinstance(
+            output_schema, AgentOutputSchema
+        ):
+            raise TypeError(
+                f"Only AgentOutputSchema is supported by Durable Model, got {type(output_schema).__name__}"
+            )
+        agent_output_schema = output_schema
+        output_schema_input = (
+            None
+            if agent_output_schema is None
+            else AgentOutputSchemaInput(
+                output_type_name=agent_output_schema.name(),
+                is_wrapped=agent_output_schema._is_wrapped,
+                output_schema=agent_output_schema.json_schema()
+                if not agent_output_schema.is_plain_text()
+                else None,
+                strict_json_schema=agent_output_schema.is_strict_json_schema(),
+            )
+        )
+
+        activity_input = ActivityModelInput(
+            model_name=self.model_name,
+            system_instructions=system_instructions,
+            input=cast(Union[str, list[TResponseInputItem]], input),
+            model_settings=model_settings,
+            tools=tool_infos,
+            output_schema=output_schema_input,
+            handoffs=handoff_infos,
+            tracing=ModelTracingInput(tracing.value),
+            previous_response_id=previous_response_id,
+            prompt=prompt,
+        )
+
+        activity_output =self.durable_orchestration_context.call_activity(
+            "InvokeModelActivity",
+            activity_input
+        )
+
+        return activity_output
+
+    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,
+        *,
+        previous_response_id: Optional[str],
+        prompt: ResponsePromptParam | None,
+    ) -> AsyncIterator[TResponseStreamEvent]:
+        raise NotImplementedError("Durable model doesn't support streams yet")

samples-v2/openai_agents/durable_openai_runner.py

@@ -2,6 +2,7 @@
 import logging
 from typing import Any, Callable, Optional
 import azure.functions as func
+import azure.durable_functions as df
 
 import json
 import typing
@@ -21,11 +22,14 @@
 from agents.run import DEFAULT_AGENT_RUNNER, DEFAULT_MAX_TURNS, AgentRunner
 from pydantic_core import to_json
 
+from durable_model_stub import _DurableModelStub
+
 logger = logging.getLogger(__name__)
 
 class DurableOpenAIRunner:
-    def __init__(self) -> None:
+    def __init__(self, durable_orchestration_context: df.DurableOrchestrationContext) -> None:
         self._runner = DEFAULT_AGENT_RUNNER or AgentRunner()
+        self.durable_orchestration_context = durable_orchestration_context
 
     def run_sync(
         self,
@@ -53,13 +57,14 @@ def run_sync(
             raise ValueError(
                 "Durable Functions require a model name to be a string in the run config and/or agent."
             )
-        # updated_run_config = replace(
-        #     run_config,
-        #     model=_TemporalModelStub(
-        #         model_name=model_name
-        #     ),
-        # )
-        updated_run_config = run_config
+        
+        updated_run_config = replace(
+            run_config,
+            model=_DurableModelStub(
+                model_name=model_name,
+                durable_orchestration_context=self.durable_orchestration_context
+            ),
+        )
 
         return self._runner.run_sync(
             starting_agent=starting_agent,
@@ -87,3 +92,18 @@ def run_streamed(
         **kwargs: Any,
     ) -> RunResultStreaming:
         raise RuntimeError("Durable Functions do not support streaming.")
+    def run(
+        self,
+        starting_agent: Agent[TContext],
+        input: Union[str, list[TResponseInputItem]],
+        **kwargs: Any,
+    ) -> RunResult:
+        raise RuntimeError("Durable Functions do not support asynchronous runs.")
+
+    def run_streamed(
+        self,
+        starting_agent: Agent[TContext],
+        input: Union[str, list[TResponseInputItem]],
+        **kwargs: Any,
+    ) -> RunResultStreaming:
+        raise RuntimeError("Durable Functions do not support streaming.")

samples-v2/openai_agents/function_app.py

@@ -4,9 +4,6 @@
 from agents.run import set_default_agent_runner
 from durable_openai_runner import DurableOpenAIRunner
 
-set_default_agent_runner(DurableOpenAIRunner())
-
-
 app = func.FunctionApp(http_auth_level=func.AuthLevel.FUNCTION)
 
 @app.route(route="orchestrators/{functionName}")
@@ -20,7 +17,8 @@ async def hello_orchestration_starter(req: func.HttpRequest, client):
 
 @app.orchestration_trigger(context_name="context")
 def basic_hello_world_orchestrator(context):
+    set_default_agent_runner(DurableOpenAIRunner(durable_orchestration_context=context))
+
     from basic.hello_world import main
     result = main()
     return result
-