Python: [Breaking] Python: Respond with AgentRunResponse with serialized structured output (#2285)

* Respond with AgentRunResponse

* Fix response_Format type

* Address comments

* Fix tests

* Fix log

* Addressed comments

* Code cleanup

* Use AgentTask vs Generator

* Address comments

* use lazy logging

* fix mypy errors
This commit is contained in:
Laveesh Rohra
2025-11-26 10:44:28 -08:00
committed by GitHub
Unverified
parent 0f2c5e6cb8
commit 306c81aef8
12 changed files with 484 additions and 432 deletions
@@ -53,7 +53,7 @@ from agent_framework import (
)
from dateutil import parser as date_parser
from ._models import RunRequest, _serialize_response_format
from ._models import RunRequest, serialize_response_format
logger = get_logger("agent_framework.azurefunctions.durable_agent_state")
@@ -494,7 +494,7 @@ class DurableAgentStateRequest(DurableAgentStateEntry):
messages=[DurableAgentStateMessage.from_run_request(request)],
created_at=datetime.now(tz=timezone.utc),
response_type=request.request_response_format,
response_schema=_serialize_response_format(request.response_format),
response_schema=serialize_response_format(request.response_format),
)
@@ -9,9 +9,7 @@ allows for long-running agent conversations.
import asyncio
import inspect
import json
from collections.abc import AsyncIterable, Callable
from datetime import datetime, timezone
from typing import Any, cast
import azure.durable_functions as df
@@ -30,11 +28,10 @@ from ._durable_agent_state import (
DurableAgentState,
DurableAgentStateData,
DurableAgentStateEntry,
DurableAgentStateMessage,
DurableAgentStateRequest,
DurableAgentStateResponse,
)
from ._models import AgentResponse, RunRequest
from ._models import RunRequest
logger = get_logger("agent_framework.azurefunctions.entities")
@@ -97,7 +94,7 @@ class AgentEntity:
self,
context: df.DurableEntityContext,
request: RunRequest | dict[str, Any] | str,
) -> dict[str, Any]:
) -> AgentRunResponse:
"""Execute the agent with a message directly in the entity.
Args:
@@ -105,13 +102,8 @@ class AgentEntity:
request: RunRequest object, dict, or string message (for backward compatibility)
Returns:
Dict with status information and response (serialized AgentResponse)
Note:
The agent returns an AgentRunResponse object which is stored in state.
This method extracts the text/structured response and returns an AgentResponse dict.
AgentRunResponse enriched with execution metadata.
"""
# Convert string or dict to RunRequest
if isinstance(request, str):
run_request = RunRequest(message=request, role=Role.USER)
elif isinstance(request, dict):
@@ -135,8 +127,6 @@ class AgentEntity:
logger.debug(f"[AgentEntity.run_agent] Received Message: {state_request}")
try:
logger.debug("[AgentEntity.run_agent] Starting agent invocation")
# Build messages from conversation history, excluding error responses
# Error responses are kept in history for tracking but not sent to the agent
chat_messages: list[ChatMessage] = [
@@ -164,83 +154,39 @@ class AgentEntity:
type(agent_run_response).__name__,
)
response_text = None
structured_response = None
response_str: str | None = None
try:
if response_format:
try:
response_str = agent_run_response.text
structured_response = json.loads(response_str)
logger.debug("Parsed structured JSON response")
except json.JSONDecodeError as decode_error:
logger.warning(f"Failed to parse JSON response: {decode_error}")
response_text = response_str
else:
raw_text = agent_run_response.text
response_text = raw_text if raw_text else "No response"
preview = response_text
logger.debug(f"Response: {preview[:100]}..." if len(preview) > 100 else f"Response: {preview}")
response_text = agent_run_response.text if agent_run_response.text else "No response"
logger.debug(f"Response: {response_text[:100]}...")
except Exception as extraction_error:
logger.error(
f"Error extracting response: {extraction_error}",
"Error extracting response text: %s",
extraction_error,
exc_info=True,
)
response_text = "Error extracting response"
state_response = DurableAgentStateResponse.from_run_response(correlation_id, agent_run_response)
self.state.data.conversation_history.append(state_response)
agent_response = AgentResponse(
response=response_text,
message=str(message),
thread_id=str(thread_id),
status="success",
message_count=len(self.state.data.conversation_history),
structured_response=structured_response,
)
result = agent_response.to_dict()
logger.debug("[AgentEntity.run_agent] AgentRunResponse stored in conversation history")
return result
return agent_run_response
except Exception as exc:
import traceback
error_traceback = traceback.format_exc()
logger.error("[AgentEntity.run_agent] Agent execution failed")
logger.error(f"Error: {exc!s}")
logger.error(f"Error type: {type(exc).__name__}")
logger.error(f"Full traceback:\n{error_traceback}")
logger.exception("[AgentEntity.run_agent] Agent execution failed.")
# Create error message
error_message = DurableAgentStateMessage.from_chat_message(
ChatMessage(
role=Role.ASSISTANT, contents=[ErrorContent(message=str(exc), error_code=type(exc).__name__)]
)
error_message = ChatMessage(
role=Role.ASSISTANT, contents=[ErrorContent(message=str(exc), error_code=type(exc).__name__)]
)
error_response = AgentRunResponse(messages=[error_message])
# Create and store error response in conversation history
error_state_response = DurableAgentStateResponse(
correlation_id=correlation_id,
created_at=datetime.now(tz=timezone.utc),
messages=[error_message],
is_error=True,
)
error_state_response = DurableAgentStateResponse.from_run_response(correlation_id, error_response)
error_state_response.is_error = True
self.state.data.conversation_history.append(error_state_response)
error_response = AgentResponse(
response=f"Error: {exc!s}",
message=str(message),
thread_id=str(thread_id),
status="error",
message_count=len(self.state.data.conversation_history),
error=str(exc),
error_type=type(exc).__name__,
)
return error_response.to_dict()
return error_response
async def _invoke_agent(
self,
@@ -432,7 +378,7 @@ def create_agent_entity(
request = "" if input_data is None else str(cast(object, input_data))
result = await entity.run_agent(context, request)
context.set_result(result)
context.set_result(result.to_dict())
elif operation == "reset":
entity.reset(context)
@@ -442,15 +388,13 @@ def create_agent_entity(
logger.error("[entity_function] Unknown operation: %s", operation)
context.set_result({"error": f"Unknown operation: {operation}"})
logger.debug("State dict: %s", entity.state.to_dict())
context.set_state(entity.state.to_dict())
serialized_state = entity.state.to_dict()
logger.debug("State dict: %s", serialized_state)
context.set_state(serialized_state)
logger.info(f"[entity_function] Operation {operation} completed successfully")
except Exception as exc:
import traceback
logger.error("[entity_function] Error in entity: %s", exc)
logger.error(f"[entity_function] Traceback:\n{traceback.format_exc()}")
logger.exception("[entity_function] Error executing entity operation %s", exc)
context.set_result({"error": str(exc), "status": "error"})
def entity_function(context: df.DurableEntityContext) -> None:
@@ -213,7 +213,7 @@ class DurableAgentThread(AgentThread):
return thread
def _serialize_response_format(response_format: type[BaseModel] | None) -> Any:
def serialize_response_format(response_format: type[BaseModel] | None) -> Any:
"""Serialize response format for transport across durable function boundaries."""
if response_format is None:
return None
@@ -339,7 +339,7 @@ class RunRequest:
"request_response_format": self.request_response_format,
}
if self.response_format:
result["response_format"] = _serialize_response_format(self.response_format)
result["response_format"] = serialize_response_format(self.response_format)
if self.thread_id:
result["thread_id"] = self.thread_id
if self.correlation_id:
@@ -362,50 +362,3 @@ class RunRequest:
correlation_id=data.get("correlationId"),
created_at=data.get("created_at"),
)
@dataclass
class AgentResponse:
"""Response from agent execution.
Attributes:
response: The agent's text response (or None for structured responses)
message: The original message sent to the agent
thread_id: The thread identifier
status: Status of the execution (success, error, etc.)
message_count: Number of messages in the conversation
error: Error message if status is error
error_type: Type of error if status is error
structured_response: Structured response if response_format was provided
"""
response: str | None
message: str
thread_id: str | None
status: str
message_count: int = 0
error: str | None = None
error_type: str | None = None
structured_response: dict[str, Any] | None = None
def to_dict(self) -> dict[str, Any]:
"""Convert to dictionary for JSON serialization."""
result: dict[str, Any] = {
"message": self.message,
"thread_id": self.thread_id,
"status": self.status,
"message_count": self.message_count,
}
# Add response or structured_response based on what's available
if self.structured_response is not None:
result["structured_response"] = self.structured_response
elif self.response is not None:
result["response"] = self.response
if self.error:
result["error"] = self.error
if self.error_type:
result["error_type"] = self.error_type
return result
@@ -6,21 +6,148 @@ This module provides support for using agents inside Durable Function orchestrat
"""
import uuid
from collections.abc import AsyncIterator
from collections.abc import AsyncIterator, Callable
from typing import TYPE_CHECKING, Any, TypeAlias, cast
from agent_framework import AgentProtocol, AgentRunResponseUpdate, AgentThread, ChatMessage, get_logger
from agent_framework import (
AgentProtocol,
AgentRunResponse,
AgentRunResponseUpdate,
AgentThread,
ChatMessage,
get_logger,
)
from azure.durable_functions.models import TaskBase
from azure.durable_functions.models.Task import CompoundTask, TaskState
from pydantic import BaseModel
from ._models import AgentSessionId, DurableAgentThread, RunRequest
logger = get_logger("agent_framework.azurefunctions.orchestration")
if TYPE_CHECKING:
from azure.durable_functions import DurableOrchestrationContext as _DurableOrchestrationContext
CompoundActionConstructor: TypeAlias = Callable[[list[Any]], Any] | None
AgentOrchestrationContextType: TypeAlias = _DurableOrchestrationContext
if TYPE_CHECKING:
from azure.durable_functions import DurableOrchestrationContext
class _TypedCompoundTask(CompoundTask): # type: ignore[misc]
_first_error: Any
def __init__(
self,
tasks: list[TaskBase],
compound_action_constructor: CompoundActionConstructor = None,
) -> None: ...
AgentOrchestrationContextType: TypeAlias = DurableOrchestrationContext
else:
AgentOrchestrationContextType = Any
_TypedCompoundTask = CompoundTask
class AgentTask(_TypedCompoundTask):
"""A custom Task that wraps entity calls and provides typed AgentRunResponse results.
This task wraps the underlying entity call task and intercepts its completion
to convert the raw result into a typed AgentRunResponse object.
"""
def __init__(
self,
entity_task: TaskBase,
response_format: type[BaseModel] | None,
correlation_id: str,
):
"""Initialize the AgentTask.
Args:
entity_task: The underlying entity call task
response_format: Optional Pydantic model for response parsing
correlation_id: Correlation ID for logging
"""
super().__init__([entity_task])
self._response_format = response_format
self._correlation_id = correlation_id
# Override action_repr to expose the inner task's action directly
# This ensures compatibility with ReplaySchema V3 which expects Action objects.
self.action_repr = entity_task.action_repr
# Also copy the task ID to match the entity task's identity
self.id = entity_task.id
def try_set_value(self, child: TaskBase) -> None:
"""Transition the AgentTask to a terminal state and set its value to `AgentRunResponse`.
Parameters
----------
child : TaskBase
The entity call task that just completed
"""
if child.state is TaskState.SUCCEEDED:
# Delegate to parent class for standard completion logic
if len(self.pending_tasks) == 0:
# Transform the raw result before setting it
raw_result = child.result
logger.debug(
"[AgentTask] Converting raw result for correlation_id %s",
self._correlation_id,
)
try:
response = self._load_agent_response(raw_result)
if self._response_format is not None:
self._ensure_response_format(
self._response_format,
self._correlation_id,
response,
)
# Set the typed AgentRunResponse as this task's result
self.set_value(is_error=False, value=response)
except Exception as e:
logger.exception(
"[AgentTask] Failed to convert result for correlation_id: %s",
self._correlation_id,
)
self.set_value(is_error=True, value=e)
else:
# If error not handled by the parent, set it explicitly.
if self._first_error is None:
self._first_error = child.result
self.set_value(is_error=True, value=self._first_error)
def _load_agent_response(self, agent_response: AgentRunResponse | dict[str, Any] | None) -> AgentRunResponse:
"""Convert raw payloads into AgentRunResponse instance."""
if agent_response is None:
raise ValueError("agent_response cannot be None")
logger.debug("[load_agent_response] Loading agent response of type: %s", type(agent_response))
if isinstance(agent_response, AgentRunResponse):
return agent_response
if isinstance(agent_response, dict):
logger.debug("[load_agent_response] Converting dict payload using AgentRunResponse.from_dict")
return AgentRunResponse.from_dict(agent_response)
raise TypeError(f"Unsupported type for agent_response: {type(agent_response)}")
def _ensure_response_format(
self,
response_format: type[BaseModel] | None,
correlation_id: str,
response: AgentRunResponse,
) -> None:
"""Ensure the AgentRunResponse value is parsed into the expected response_format."""
if response_format is not None and not isinstance(response.value, response_format):
response.try_parse_value(response_format)
logger.debug(
"[DurableAIAgent] Loaded AgentRunResponse.value for correlation_id %s with type: %s",
correlation_id,
type(response.value).__name__,
)
class DurableAIAgent(AgentProtocol):
@@ -59,7 +186,7 @@ class DurableAIAgent(AgentProtocol):
self._name = agent_name
self._display_name = agent_name
self._description = f"Durable agent proxy for {agent_name}"
logger.debug(f"[DurableAIAgent] Initialized for agent: {agent_name}")
logger.debug("[DurableAIAgent] Initialized for agent: %s", agent_name)
@property
def id(self) -> str:
@@ -81,38 +208,45 @@ class DurableAIAgent(AgentProtocol):
"""Get the description of the agent."""
return self._description
def run(
# We return an AgentTask here which is a TaskBase subclass.
# This is an intentional deviation from AgentProtocol which defines run() as async.
# The AgentTask can be yielded in Durable Functions orchestrations and will provide
# a typed AgentRunResponse result.
def run( # type: ignore[override]
self,
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
*,
thread: AgentThread | None = None,
response_format: type[BaseModel] | None = None,
**kwargs: Any,
) -> Any: # TODO(msft-team): Add a wrapper to respond correctly with `AgentRunResponse`
"""Execute the agent with messages and return a Task for orchestrations.
) -> AgentTask:
"""Execute the agent with messages and return an AgentTask for orchestrations.
This method implements AgentProtocol and returns a Task that can be yielded
in Durable Functions orchestrations.
This method implements AgentProtocol and returns an AgentTask (subclass of TaskBase)
that can be yielded in Durable Functions orchestrations. The task's result will be
a typed AgentRunResponse.
Args:
messages: The message(s) to send to the agent
thread: Optional agent thread for conversation context
**kwargs: Additional arguments (enable_tool_calls, response_format, etc.)
response_format: Optional Pydantic model for response parsing
**kwargs: Additional arguments (enable_tool_calls)
Returns:
Task that will resolve to the agent response
An AgentTask that resolves to an AgentRunResponse when yielded
Example:
@app.orchestration_trigger(context_name="context")
def my_orchestration(context):
agent = app.get_agent(context, "MyAgent")
thread = agent.get_new_thread()
result = yield agent.run("Hello", thread=thread)
response = yield agent.run("Hello", thread=thread)
# response is typed as AgentRunResponse
"""
message_str = self._normalize_messages(messages)
# Extract optional parameters from kwargs
enable_tool_calls = kwargs.get("enable_tool_calls", True)
response_format = kwargs.get("response_format")
# Get the session ID for the entity
if isinstance(thread, DurableAgentThread) and thread.session_id is not None:
@@ -122,7 +256,7 @@ class DurableAIAgent(AgentProtocol):
# This ensures each call gets its own conversation context
session_key = str(self.context.new_uuid())
session_id = AgentSessionId(name=self.agent_name, key=session_key)
logger.warning(f"[DurableAIAgent] No thread provided, created unique session_id: {session_id}")
logger.warning("[DurableAIAgent] No thread provided, created unique session_id: %s", session_id)
# Create entity ID from session ID
entity_id = session_id.to_entity_id()
@@ -130,6 +264,12 @@ class DurableAIAgent(AgentProtocol):
# Generate a deterministic correlation ID for this call
# This is required by the entity and must be unique per call
correlation_id = str(self.context.new_uuid())
logger.debug(
"[DurableAIAgent] Using correlation_id: %s for entity_id: %s for session_id: %s",
correlation_id,
entity_id,
session_id,
)
# Prepare the request using RunRequest model
run_request = RunRequest(
@@ -140,11 +280,24 @@ class DurableAIAgent(AgentProtocol):
response_format=response_format,
)
logger.debug(f"[DurableAIAgent] Calling entity {entity_id} with message: {message_str[:100]}...")
logger.debug("[DurableAIAgent] Calling entity %s with message: %s", entity_id, message_str[:100])
# Call the entity and return the Task directly
# The orchestration will yield this Task
return self.context.call_entity(entity_id, "run_agent", run_request.to_dict())
# Call the entity to get the underlying task
entity_task = self.context.call_entity(entity_id, "run_agent", run_request.to_dict())
# Wrap it in an AgentTask that will convert the result to AgentRunResponse
agent_task = AgentTask(
entity_task=entity_task,
response_format=response_format,
correlation_id=correlation_id,
)
logger.debug(
"[DurableAIAgent] Created AgentTask for correlation_id %s",
correlation_id,
)
return agent_task
def run_stream(
self,
@@ -179,7 +332,7 @@ class DurableAIAgent(AgentProtocol):
thread = DurableAgentThread.from_session_id(session_id, **kwargs)
logger.debug(f"[DurableAIAgent] Created new thread with session_id: {session_id}")
logger.debug("[DurableAIAgent] Created new thread with session_id: %s", session_id)
return thread
def _messages_to_string(self, messages: list[ChatMessage]) -> str: