mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
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:
committed by
GitHub
Unverified
parent
0f2c5e6cb8
commit
306c81aef8
+2
-2
@@ -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:
|
||||
|
||||
Reference in New Issue
Block a user