Switch to ActivityCallTracker

Author: AnatoliB

samples-v2/openai_agents/activity_call_tracker.py

@@ -0,0 +1,23 @@
+from azure.durable_functions.models.DurableOrchestrationContext import DurableOrchestrationContext
+from yield_exception import YieldException
+
+class ActivityCallTracker:
+    def __init__(self, context: DurableOrchestrationContext):
+        self.context = context
+        self.activities_called = 0
+        self.tasks_to_yield = []
+
+    def call_activity(self, activity_name, input: str):
+        task = self.context.call_activity(activity_name, input)
+
+        self.activities_called += 1
+
+        histories = self.context.histories
+        completed_tasks = [entry for entry in histories if entry.event_type == 5]
+        if len(completed_tasks) < self.activities_called:
+            # yield immediately
+            raise YieldException(task)
+        else:
+            # yield later
+            self.tasks_to_yield.append(task)
+            return completed_tasks[self.activities_called - 1].Result

samples-v2/openai_agents/durable_model_stub.py

@@ -8,7 +8,7 @@
     HandoffInput, 
     AgentOutputSchemaInput
 )
-from orchestrator_exceptions import OrchestratorYielded
+from activity_call_tracker import ActivityCallTracker
 import azure.durable_functions as df
 
 import json
@@ -44,10 +44,10 @@ class _DurableModelStub(Model):
     def __init__(
         self,
         model_name: Optional[str],
-        durable_orchestration_context: df.DurableOrchestrationContext,
+        activity_call_tracker: ActivityCallTracker,
     ) -> None:
         self.model_name = model_name
-        self.durable_orchestration_context = durable_orchestration_context
+        self.activity_call_tracker = activity_call_tracker
 
     async def get_response(
         self,
@@ -129,18 +129,9 @@ def make_tool_info(tool: Tool) -> ToolInput:
             prompt=prompt,
         )
 
-        # Serialize activity_input to JSON
         activity_input_json = activity_input.to_json()
 
-        task = self.durable_orchestration_context.call_activity(
-            "invoke_model_activity",
-            activity_input_json
-        )
-
-        if not self.durable_orchestration_context.is_replaying:
-            raise OrchestratorYielded(task)
-        
-        result = task.result
+        result = self.activity_call_tracker.call_activity("invoke_model_activity", activity_input_json)
         return result
 
     def stream_response(

samples-v2/openai_agents/durable_openai_runner.py

@@ -1,6 +1,7 @@
 import asyncio
 import logging
 from typing import Any, Callable, Optional
+from activity_call_tracker import ActivityCallTracker
 import azure.functions as func
 import azure.durable_functions as df
 
@@ -27,9 +28,9 @@
 logger = logging.getLogger(__name__)
 
 class DurableOpenAIRunner:
-    def __init__(self, durable_orchestration_context: df.DurableOrchestrationContext) -> None:
+    def __init__(self, activity_call_tracker: ActivityCallTracker) -> None:
         self._runner = DEFAULT_AGENT_RUNNER or AgentRunner()
-        self.durable_orchestration_context = durable_orchestration_context
+        self.activity_call_tracker = activity_call_tracker
 
     def run_sync(
         self,
@@ -60,9 +61,9 @@ def run_sync(
 
         updated_run_config = replace(
             run_config,
-            model=_DurableModelStub(
-                model_name=model_name,
-                durable_orchestration_context=self.durable_orchestration_context
+            model = _DurableModelStub(
+                model_name = model_name,
+                activity_call_tracker = self.activity_call_tracker,
             ),
         )
 

samples-v2/openai_agents/function_app.py

@@ -2,10 +2,6 @@
 import azure.functions as func
 import logging
 
-from agents.run import set_default_agent_runner
-from durable_openai_runner import DurableOpenAIRunner
-from orchestrator_exceptions import OrchestratorYielded
-
 from agents import (
     AgentOutputSchemaBase,
     CodeInterpreterTool,
@@ -26,6 +22,11 @@
     WebSearchTool,
 )
 
+from agents.run import set_default_agent_runner
+from durable_openai_runner import DurableOpenAIRunner
+from yield_exception import YieldException
+from activity_call_tracker import ActivityCallTracker
+
 app = func.FunctionApp(http_auth_level=func.AuthLevel.FUNCTION)
 
 @app.route(route="orchestrators/{functionName}")
@@ -39,14 +40,24 @@ 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))
+    activity_call_tracker = ActivityCallTracker(context)
+    durable_openai_runner = DurableOpenAIRunner(activity_call_tracker=activity_call_tracker)
+    set_default_agent_runner(durable_openai_runner)
 
     try:
         from basic.hello_world import main
         result = main()
         return result
-    except OrchestratorYielded as e:
-        yield e.activity_output
+    except YieldException as e:
+        for task in activity_call_tracker.tasks_to_yield:
+            yield task
+        activity_call_tracker.tasks_to_yield.clear()
+        yield e.task
+    finally:
+        for task in activity_call_tracker.tasks_to_yield:
+            yield task
+        activity_call_tracker.tasks_to_yield.clear()
+
 
 @app.activity_trigger(input_name="input")
 async def invoke_model_activity(input: str):

samples-v2/openai_agents/orchestrator_exceptions.py

@@ -0,0 +1,6 @@
+from azure.durable_functions.models.Task import TaskBase
+
+class YieldException(Exception):
+    def __init__(self, task: TaskBase):
+        super().__init__("Orchestrator should yield.")
+        self.task = task