Pass DurableAIAgentContext instead of DurableOrchestrationContext to the wrapped orchestrator

Author: AnatoliB

samples-v2/openai_agents/durable_decorators.py

@@ -1,5 +1,6 @@
 from functools import wraps
 from agents.run import set_default_agent_runner
+from azure.durable_functions.models.DurableOrchestrationContext import DurableOrchestrationContext
 from durable_openai_runner import DurableOpenAIRunner
 from yield_exception import YieldException
 from durable_ai_agent_context import DurableAIAgentContext
@@ -8,19 +9,19 @@
 
 def durable_openai_agent_orchestrator(func):
     @wraps(func)
-    def wrapper(context):
+    def wrapper(durable_orchestration_context: DurableOrchestrationContext):
         event_loop.ensure_event_loop()
-        context = DurableAIAgentContext(context)
-        durable_openai_runner = DurableOpenAIRunner(context=context)
+        durable_ai_agent_context = DurableAIAgentContext(durable_orchestration_context)
+        durable_openai_runner = DurableOpenAIRunner(context=durable_ai_agent_context)
         set_default_agent_runner(durable_openai_runner)
 
         try:
-            result = func(context)
+            result = func(durable_ai_agent_context)
             return result
         except YieldException as e:
-            yield from context.yield_and_clear_tasks()
+            yield from durable_ai_agent_context.yield_and_clear_tasks()
             yield e.task
         finally:
-            yield from context.yield_and_clear_tasks()
-    
+            yield from durable_ai_agent_context.yield_and_clear_tasks()
+
     return wrapper