Improve orchestration output serialization

azure/durable_functions/decorators/durable_app.py

@@ -1,6 +1,7 @@
 #  Copyright (c) Microsoft Corporation. All rights reserved.
 #  Licensed under the MIT License.
 
+
 from .metadata import OrchestrationTrigger, ActivityTrigger, EntityTrigger,\
     DurableClient
 from typing import Callable, Optional

azure/durable_functions/openai_agents/orchestrator_generator.py

@@ -1,4 +1,6 @@
 import inspect
+import json
+from typing import Any
 from agents import ModelProvider, ModelResponse
 from agents.run import set_default_agent_runner
 from azure.durable_functions.models.DurableOrchestrationContext import DurableOrchestrationContext
@@ -10,6 +12,23 @@
 from .event_loop import ensure_event_loop
 
 
+def _durable_serializer(obj: Any) -> str:
+    # Strings are already "serialized"
+    if type(obj) is str:
+        return obj
+
+    # Serialize "Durable" and OpenAI models, and typed dictionaries
+    if callable(getattr(obj, "to_json", None)):
+        return obj.to_json()
+
+    # Serialize Pydantic models
+    if callable(getattr(obj, "model_dump_json", None)):
+        return obj.model_dump_json()
+
+    # Fallback to default JSON serialization
+    return json.dumps(obj)
+
+
 async def durable_openai_agent_activity(input: str, model_provider: ModelProvider):
     """Activity logic that handles OpenAI model invocations."""
     activity_input = ActivityModelInput.from_json(input)
@@ -60,14 +79,14 @@ def durable_openai_agent_orchestrator_generator(
                         value = next(gen)
         except StopIteration as e:
             yield from durable_ai_agent_context._yield_and_clear_tasks()
-            return e.value
+            return _durable_serializer(e.value)
         except YieldException as e:
             yield from durable_ai_agent_context._yield_and_clear_tasks()
             yield e.task
     else:
         try:
             result = func(durable_ai_agent_context)
-            return result
+            return _durable_serializer(result)
         except YieldException as e:
             yield from durable_ai_agent_context._yield_and_clear_tasks()
             yield e.task

tests/orchestrator/openai_agents/test_openai_agents.py

@@ -1,6 +1,9 @@
+from typing import Optional, TypedDict
+
 import azure.durable_functions as df
 import azure.functions as func
 import json
+import pydantic
 from agents import Agent, Runner
 from azure.durable_functions.models import OrchestratorState
 from azure.durable_functions.models.actions import CallActivityAction
@@ -56,6 +59,60 @@ def openai_agent_use_tool(context):
 
     return result.final_output;
 
+@app.function_name("openai_agent_return_string_type")
+@app.orchestration_trigger(context_name="context")
+@app.durable_openai_agent_orchestrator
+def openai_agent_return_string_type(context):
+    return "Hello World"
+
+class DurableModel:
+    def __init__(self, property: str) -> None:
+        self._property = property
+
+    def to_json(self) -> str:
+        return json.dumps({"property": self._property})
+
+@app.function_name("openai_agent_return_durable_model_type")
+@app.orchestration_trigger(context_name="context")
+@app.durable_openai_agent_orchestrator
+def openai_agent_return_durable_model_type(context):
+    model = DurableModel(property="value")
+
+    return model
+
+class TypedDictionaryModel(TypedDict):
+    property: str
+
+@app.function_name("openai_agent_return_typed_dictionary_model_type")
+@app.orchestration_trigger(context_name="context")
+@app.durable_openai_agent_orchestrator
+def openai_agent_return_typed_dictionary_model_type(context):
+    model = TypedDictionaryModel(property="value")
+
+    return model
+
+class OpenAIPydanticModel(BaseModel):
+    property: str
+
+@app.function_name("openai_agent_return_openai_pydantic_model_type")
+@app.orchestration_trigger(context_name="context")
+@app.durable_openai_agent_orchestrator
+def openai_agent_return_openai_pydantic_model_type(context):
+    model = OpenAIPydanticModel(property="value")
+
+    return model
+
+class PydanticModel(pydantic.BaseModel):
+    property: str
+
+@app.function_name("openai_agent_return_pydantic_model_type")
+@app.orchestration_trigger(context_name="context")
+@app.durable_openai_agent_orchestrator
+def openai_agent_return_pydantic_model_type(context):
+    model = PydanticModel(property="value")
+
+    return model
+
 model_activity_name = "invoke_model_activity"
 
 def base_expected_state(output=None, replay_schema: ReplaySchema = ReplaySchema.V1) -> OrchestratorState:
@@ -152,3 +209,73 @@ def test_openai_agent_use_tool_analysis_completed():
 
     assert_valid_schema(result)
     assert_orchestration_state_equals(expected, result)
+
+def test_openai_agent_string_serialization():
+    context_builder = ContextBuilder('test_openai_agent_string_serialization')
+
+    result = get_orchestration_state_result(
+        context_builder, openai_agent_return_string_type, uses_pystein=True)
+
+    expected_state = base_expected_state()
+    expected_state._is_done = True
+    expected_state._output = "Hello World"
+    expected = expected_state.to_json()
+
+    assert_valid_schema(result)
+    assert_orchestration_state_equals(expected, result)
+
+def test_openai_agent_durable_model_serialization():
+    context_builder = ContextBuilder('test_openai_agent_durable_model_serialization')
+
+    result = get_orchestration_state_result(
+        context_builder, openai_agent_return_durable_model_type, uses_pystein=True)
+
+    expected_state = base_expected_state()
+    expected_state._is_done = True
+    expected_state._output = DurableModel(property="value").to_json()
+    expected = expected_state.to_json()
+
+    assert_valid_schema(result)
+    assert_orchestration_state_equals(expected, result)
+
+def test_openai_agent_typed_dictionary_model_serialization():
+    context_builder = ContextBuilder('test_openai_agent_typed_dictionary_model_serialization')
+
+    result = get_orchestration_state_result(
+        context_builder, openai_agent_return_typed_dictionary_model_type, uses_pystein=True)
+
+    expected_state = base_expected_state()
+    expected_state._is_done = True
+    expected_state._output = json.dumps(TypedDictionaryModel(property="value"))
+    expected = expected_state.to_json()
+
+    assert_valid_schema(result)
+    assert_orchestration_state_equals(expected, result)
+
+def test_openai_agent_openai_pydantic_model_serialization():
+    context_builder = ContextBuilder('test_openai_agent_openai_pydantic_model_serialization')
+
+    result = get_orchestration_state_result(
+        context_builder, openai_agent_return_openai_pydantic_model_type, uses_pystein=True)
+
+    expected_state = base_expected_state()
+    expected_state._is_done = True
+    expected_state._output = OpenAIPydanticModel(property="value").to_json()
+    expected = expected_state.to_json()
+
+    assert_valid_schema(result)
+    assert_orchestration_state_equals(expected, result)
+
+def test_openai_agent_pydantic_model_serialization():
+    context_builder = ContextBuilder('test_openai_agent_pydantic_model_serialization')
+
+    result = get_orchestration_state_result(
+        context_builder, openai_agent_return_pydantic_model_type, uses_pystein=True)
+
+    expected_state = base_expected_state()
+    expected_state._is_done = True
+    expected_state._output = PydanticModel(property="value").model_dump_json()
+    expected = expected_state.to_json()
+
+    assert_valid_schema(result)
+    assert_orchestration_state_equals(expected, result)