mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Complete durableagent package (#3058)
* Add worker and clients * Clean code and refactor common code * Implement sample * Add sample * Update readmes * Fix tests * Fix tests * Update requirements * Fix typo * Address comments * use response.text
This commit is contained in:
Laveesh Rohra
committed by
GitHub
Unverified
GitHub
parent
a5b36dc379
commit
e3eff65a6b
@@ -2,10 +2,9 @@
|
||||
|
||||
import importlib.metadata
|
||||
|
||||
from agent_framework_durabletask import AgentCallbackContext, AgentResponseCallbackProtocol
|
||||
from agent_framework_durabletask import AgentCallbackContext, AgentResponseCallbackProtocol, DurableAIAgent
|
||||
|
||||
from ._app import AgentFunctionApp
|
||||
from ._orchestration import DurableAIAgent
|
||||
|
||||
try:
|
||||
__version__ = importlib.metadata.version(__name__)
|
||||
|
||||
@@ -8,6 +8,7 @@ with Azure Durable Entities, enabling stateful and durable AI agent execution.
|
||||
|
||||
import json
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
@@ -28,14 +29,16 @@ from agent_framework_durabletask import (
|
||||
WAIT_FOR_RESPONSE_FIELD,
|
||||
WAIT_FOR_RESPONSE_HEADER,
|
||||
AgentResponseCallbackProtocol,
|
||||
AgentSessionId,
|
||||
ApiResponseFields,
|
||||
DurableAgentState,
|
||||
DurableAIAgent,
|
||||
RunRequest,
|
||||
)
|
||||
|
||||
from ._entities import create_agent_entity
|
||||
from ._errors import IncomingRequestError
|
||||
from ._models import AgentSessionId
|
||||
from ._orchestration import AgentOrchestrationContextType, DurableAIAgent
|
||||
from ._orchestration import AgentOrchestrationContextType, AgentTask, AzureFunctionsAgentExecutor
|
||||
|
||||
logger = get_logger("agent_framework.azurefunctions")
|
||||
|
||||
@@ -296,7 +299,7 @@ class AgentFunctionApp(DFAppBase):
|
||||
self,
|
||||
context: AgentOrchestrationContextType,
|
||||
agent_name: str,
|
||||
) -> DurableAIAgent:
|
||||
) -> DurableAIAgent[AgentTask]:
|
||||
"""Return a DurableAIAgent proxy for a registered agent.
|
||||
|
||||
Args:
|
||||
@@ -307,14 +310,15 @@ class AgentFunctionApp(DFAppBase):
|
||||
ValueError: If the requested agent has not been registered.
|
||||
|
||||
Returns:
|
||||
DurableAIAgent wrapper bound to the orchestration context.
|
||||
DurableAIAgent[AgentTask] wrapper bound to the orchestration context.
|
||||
"""
|
||||
normalized_name = str(agent_name)
|
||||
|
||||
if normalized_name not in self._agent_metadata:
|
||||
raise ValueError(f"Agent '{normalized_name}' is not registered with this app.")
|
||||
|
||||
return DurableAIAgent(context, normalized_name)
|
||||
executor = AzureFunctionsAgentExecutor(context)
|
||||
return DurableAIAgent(executor, normalized_name)
|
||||
|
||||
def _setup_agent_functions(
|
||||
self,
|
||||
@@ -377,8 +381,6 @@ class AgentFunctionApp(DFAppBase):
|
||||
"enable_tool_calls": true|false (optional, default: true)
|
||||
}
|
||||
"""
|
||||
logger.debug(f"[HTTP Trigger] Received request on route: /api/agents/{agent_name}/run")
|
||||
|
||||
request_response_format: str = REQUEST_RESPONSE_FORMAT_JSON
|
||||
thread_id: str | None = None
|
||||
|
||||
@@ -387,9 +389,9 @@ class AgentFunctionApp(DFAppBase):
|
||||
thread_id = self._resolve_thread_id(req=req, req_body=req_body)
|
||||
wait_for_response = self._should_wait_for_response(req=req, req_body=req_body)
|
||||
|
||||
logger.debug(f"[HTTP Trigger] Message: {message}")
|
||||
logger.debug(f"[HTTP Trigger] Thread ID: {thread_id}")
|
||||
logger.debug(f"[HTTP Trigger] wait_for_response: {wait_for_response}")
|
||||
logger.debug(
|
||||
f"[HTTP Trigger] Message: {message}, Thread ID: {thread_id}, wait_for_response: {wait_for_response}"
|
||||
)
|
||||
|
||||
if not message:
|
||||
logger.warning("[HTTP Trigger] Request rejected: Missing message")
|
||||
@@ -403,15 +405,18 @@ class AgentFunctionApp(DFAppBase):
|
||||
session_id = self._create_session_id(agent_name, thread_id)
|
||||
correlation_id = self._generate_unique_id()
|
||||
|
||||
logger.debug(f"[HTTP Trigger] Using session ID: {session_id}")
|
||||
logger.debug(f"[HTTP Trigger] Generated correlation ID: {correlation_id}")
|
||||
logger.debug("[HTTP Trigger] Calling entity to run agent...")
|
||||
logger.debug(
|
||||
f"[HTTP Trigger] Calling entity to run agent using session ID: {session_id} "
|
||||
f"and correlation ID: {correlation_id}"
|
||||
)
|
||||
|
||||
entity_instance_id = session_id.to_entity_id()
|
||||
entity_instance_id = df.EntityId(
|
||||
name=session_id.entity_name,
|
||||
key=session_id.key,
|
||||
)
|
||||
run_request = self._build_request_data(
|
||||
req_body,
|
||||
message,
|
||||
thread_id,
|
||||
correlation_id,
|
||||
request_response_format,
|
||||
)
|
||||
@@ -624,14 +629,16 @@ class AgentFunctionApp(DFAppBase):
|
||||
session_id = AgentSessionId.with_random_key(agent_name)
|
||||
|
||||
# Build entity instance ID
|
||||
entity_instance_id = session_id.to_entity_id()
|
||||
entity_instance_id = df.EntityId(
|
||||
name=session_id.entity_name,
|
||||
key=session_id.key,
|
||||
)
|
||||
|
||||
# Create run request
|
||||
correlation_id = self._generate_unique_id()
|
||||
run_request = self._build_request_data(
|
||||
req_body={"message": query, "role": "user"},
|
||||
message=query,
|
||||
thread_id=str(session_id),
|
||||
correlation_id=correlation_id,
|
||||
request_response_format=REQUEST_RESPONSE_FORMAT_TEXT,
|
||||
)
|
||||
@@ -783,7 +790,7 @@ class AgentFunctionApp(DFAppBase):
|
||||
agent_response = state.try_get_agent_response(correlation_id)
|
||||
if agent_response:
|
||||
result = self._build_success_result(
|
||||
response_data=agent_response,
|
||||
response_message=agent_response.text,
|
||||
message=message,
|
||||
thread_id=thread_id,
|
||||
correlation_id=correlation_id,
|
||||
@@ -829,23 +836,22 @@ class AgentFunctionApp(DFAppBase):
|
||||
)
|
||||
|
||||
def _build_success_result(
|
||||
self, response_data: dict[str, Any], message: str, thread_id: str, correlation_id: str, state: DurableAgentState
|
||||
self, response_message: str, message: str, thread_id: str, correlation_id: str, state: DurableAgentState
|
||||
) -> dict[str, Any]:
|
||||
"""Build the success result returned to the HTTP caller."""
|
||||
return self._build_response_payload(
|
||||
response=response_data.get("content"),
|
||||
response=response_message,
|
||||
message=message,
|
||||
thread_id=thread_id,
|
||||
status="success",
|
||||
correlation_id=correlation_id,
|
||||
extra_fields={"message_count": response_data.get("message_count", state.message_count)},
|
||||
extra_fields={ApiResponseFields.MESSAGE_COUNT: state.message_count},
|
||||
)
|
||||
|
||||
def _build_request_data(
|
||||
self,
|
||||
req_body: dict[str, Any],
|
||||
message: str,
|
||||
thread_id: str,
|
||||
correlation_id: str,
|
||||
request_response_format: str,
|
||||
) -> dict[str, Any]:
|
||||
@@ -912,15 +918,13 @@ class AgentFunctionApp(DFAppBase):
|
||||
|
||||
def _generate_unique_id(self) -> str:
|
||||
"""Generate a new unique identifier."""
|
||||
import uuid
|
||||
|
||||
return uuid.uuid4().hex
|
||||
|
||||
def _create_session_id(self, func_name: str, thread_id: str | None) -> AgentSessionId:
|
||||
def _create_session_id(self, agent_name: str, thread_id: str | None) -> AgentSessionId:
|
||||
"""Create a session identifier using the provided thread id or a random value."""
|
||||
if thread_id:
|
||||
return AgentSessionId(name=func_name, key=thread_id)
|
||||
return AgentSessionId.with_random_key(name=func_name)
|
||||
return AgentSessionId(name=agent_name, key=thread_id)
|
||||
return AgentSessionId.with_random_key(name=agent_name)
|
||||
|
||||
def _resolve_thread_id(self, req: func.HttpRequest, req_body: dict[str, Any]) -> str:
|
||||
"""Retrieve the thread identifier from request body or query parameters."""
|
||||
|
||||
@@ -1,201 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Azure Functions-specific data models for Durable Agent Framework.
|
||||
|
||||
This module contains Azure Functions-specific models:
|
||||
- AgentSessionId: Entity ID management for Azure Durable Entities
|
||||
- DurableAgentThread: Thread implementation that tracks AgentSessionId
|
||||
|
||||
Common models like RunRequest have been moved to agent-framework-durabletask.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import uuid
|
||||
from collections.abc import MutableMapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import azure.durable_functions as df
|
||||
from agent_framework import AgentThread
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentSessionId:
|
||||
"""Represents an agent session ID, which is used to identify a long-running agent session.
|
||||
|
||||
Attributes:
|
||||
name: The name of the agent that owns the session (case-insensitive)
|
||||
key: The unique key of the agent session (case-sensitive)
|
||||
"""
|
||||
|
||||
name: str
|
||||
key: str
|
||||
|
||||
ENTITY_NAME_PREFIX: str = "dafx-"
|
||||
|
||||
@staticmethod
|
||||
def to_entity_name(name: str) -> str:
|
||||
"""Converts an agent name to an entity name by adding the DAFx prefix.
|
||||
|
||||
Args:
|
||||
name: The agent name
|
||||
|
||||
Returns:
|
||||
The entity name with the dafx- prefix
|
||||
"""
|
||||
return f"{AgentSessionId.ENTITY_NAME_PREFIX}{name}"
|
||||
|
||||
@staticmethod
|
||||
def with_random_key(name: str) -> AgentSessionId:
|
||||
"""Creates a new AgentSessionId with the specified name and a randomly generated key.
|
||||
|
||||
Args:
|
||||
name: The name of the agent that owns the session
|
||||
|
||||
Returns:
|
||||
A new AgentSessionId with the specified name and a random GUID key
|
||||
"""
|
||||
return AgentSessionId(name=name, key=uuid.uuid4().hex)
|
||||
|
||||
def to_entity_id(self) -> df.EntityId:
|
||||
"""Converts this AgentSessionId to a Durable Functions EntityId.
|
||||
|
||||
Returns:
|
||||
EntityId for use with Durable Functions APIs
|
||||
"""
|
||||
return df.EntityId(self.to_entity_name(self.name), self.key)
|
||||
|
||||
@staticmethod
|
||||
def from_entity_id(entity_id: df.EntityId) -> AgentSessionId:
|
||||
"""Creates an AgentSessionId from a Durable Functions EntityId.
|
||||
|
||||
Args:
|
||||
entity_id: The EntityId to convert
|
||||
|
||||
Returns:
|
||||
AgentSessionId instance
|
||||
|
||||
Raises:
|
||||
ValueError: If the entity ID does not have the expected prefix
|
||||
"""
|
||||
if not entity_id.name.startswith(AgentSessionId.ENTITY_NAME_PREFIX):
|
||||
raise ValueError(
|
||||
f"'{entity_id}' is not a valid agent session ID. "
|
||||
f"Expected entity name to start with '{AgentSessionId.ENTITY_NAME_PREFIX}'"
|
||||
)
|
||||
|
||||
agent_name = entity_id.name[len(AgentSessionId.ENTITY_NAME_PREFIX) :]
|
||||
return AgentSessionId(name=agent_name, key=entity_id.key)
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""Returns a string representation in the form @name@key."""
|
||||
return f"@{self.name}@{self.key}"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""Returns a detailed string representation."""
|
||||
return f"AgentSessionId(name='{self.name}', key='{self.key}')"
|
||||
|
||||
@staticmethod
|
||||
def parse(session_id_string: str) -> AgentSessionId:
|
||||
"""Parses a string representation of an agent session ID.
|
||||
|
||||
Args:
|
||||
session_id_string: A string in the form @name@key
|
||||
|
||||
Returns:
|
||||
AgentSessionId instance
|
||||
|
||||
Raises:
|
||||
ValueError: If the string format is invalid
|
||||
"""
|
||||
if not session_id_string.startswith("@"):
|
||||
raise ValueError(f"Invalid agent session ID format: {session_id_string}")
|
||||
|
||||
parts = session_id_string[1:].split("@", 1)
|
||||
if len(parts) != 2:
|
||||
raise ValueError(f"Invalid agent session ID format: {session_id_string}")
|
||||
|
||||
return AgentSessionId(name=parts[0], key=parts[1])
|
||||
|
||||
|
||||
class DurableAgentThread(AgentThread):
|
||||
"""Durable agent thread that tracks the owning :class:`AgentSessionId`."""
|
||||
|
||||
_SERIALIZED_SESSION_ID_KEY = "durable_session_id"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
session_id: AgentSessionId | None = None,
|
||||
service_thread_id: str | None = None,
|
||||
message_store: Any = None,
|
||||
context_provider: Any = None,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
service_thread_id=service_thread_id,
|
||||
message_store=message_store,
|
||||
context_provider=context_provider,
|
||||
)
|
||||
self._session_id: AgentSessionId | None = session_id
|
||||
|
||||
@property
|
||||
def session_id(self) -> AgentSessionId | None:
|
||||
"""Returns the durable agent session identifier for this thread."""
|
||||
return self._session_id
|
||||
|
||||
def attach_session(self, session_id: AgentSessionId) -> None:
|
||||
"""Associates the thread with the provided :class:`AgentSessionId`."""
|
||||
self._session_id = session_id
|
||||
|
||||
@classmethod
|
||||
def from_session_id(
|
||||
cls,
|
||||
session_id: AgentSessionId,
|
||||
*,
|
||||
service_thread_id: str | None = None,
|
||||
message_store: Any = None,
|
||||
context_provider: Any = None,
|
||||
) -> DurableAgentThread:
|
||||
"""Creates a durable thread pre-associated with the supplied session ID."""
|
||||
return cls(
|
||||
session_id=session_id,
|
||||
service_thread_id=service_thread_id,
|
||||
message_store=message_store,
|
||||
context_provider=context_provider,
|
||||
)
|
||||
|
||||
async def serialize(self, **kwargs: Any) -> dict[str, Any]:
|
||||
"""Serializes thread state including the durable session identifier."""
|
||||
state = await super().serialize(**kwargs)
|
||||
if self._session_id is not None:
|
||||
state[self._SERIALIZED_SESSION_ID_KEY] = str(self._session_id)
|
||||
return state
|
||||
|
||||
@classmethod
|
||||
async def deserialize(
|
||||
cls,
|
||||
serialized_thread_state: MutableMapping[str, Any],
|
||||
*,
|
||||
message_store: Any = None,
|
||||
**kwargs: Any,
|
||||
) -> DurableAgentThread:
|
||||
"""Restores a durable thread, rehydrating the stored session identifier."""
|
||||
state_payload = dict(serialized_thread_state)
|
||||
session_id_value = state_payload.pop(cls._SERIALIZED_SESSION_ID_KEY, None)
|
||||
thread = await super().deserialize(
|
||||
state_payload,
|
||||
message_store=message_store,
|
||||
**kwargs,
|
||||
)
|
||||
if not isinstance(thread, DurableAgentThread):
|
||||
raise TypeError("Deserialized thread is not a DurableAgentThread instance")
|
||||
|
||||
if session_id_value is None:
|
||||
return thread
|
||||
|
||||
if not isinstance(session_id_value, str):
|
||||
raise ValueError("durable_session_id must be a string when present in serialized state")
|
||||
|
||||
thread.attach_session(AgentSessionId.parse(session_id_value))
|
||||
return thread
|
||||
@@ -5,25 +5,21 @@
|
||||
This module provides support for using agents inside Durable Function orchestrations.
|
||||
"""
|
||||
|
||||
import uuid
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from typing import TYPE_CHECKING, Any, TypeAlias, cast
|
||||
from collections.abc import Callable
|
||||
from typing import TYPE_CHECKING, Any, TypeAlias
|
||||
|
||||
from agent_framework import (
|
||||
AgentProtocol,
|
||||
AgentRunResponse,
|
||||
AgentRunResponseUpdate,
|
||||
AgentThread,
|
||||
ChatMessage,
|
||||
get_logger,
|
||||
import azure.durable_functions as df
|
||||
from agent_framework import AgentThread, get_logger
|
||||
from agent_framework_durabletask import (
|
||||
DurableAgentExecutor,
|
||||
RunRequest,
|
||||
ensure_response_format,
|
||||
load_agent_response,
|
||||
)
|
||||
from agent_framework_durabletask import RunRequest
|
||||
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
|
||||
|
||||
logger = get_logger("agent_framework.azurefunctions.orchestration")
|
||||
|
||||
CompoundActionConstructor: TypeAlias = Callable[[list[Any]], Any] | None
|
||||
@@ -96,10 +92,10 @@ class AgentTask(_TypedCompoundTask):
|
||||
)
|
||||
|
||||
try:
|
||||
response = self._load_agent_response(raw_result)
|
||||
response = load_agent_response(raw_result)
|
||||
|
||||
if self._response_format is not None:
|
||||
self._ensure_response_format(
|
||||
ensure_response_format(
|
||||
self._response_format,
|
||||
self._correlation_id,
|
||||
response,
|
||||
@@ -119,249 +115,60 @@ class AgentTask(_TypedCompoundTask):
|
||||
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))
|
||||
class AzureFunctionsAgentExecutor(DurableAgentExecutor[AgentTask]):
|
||||
"""Executor that executes durable agents inside Azure Functions orchestrations."""
|
||||
|
||||
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):
|
||||
"""A durable agent implementation that uses entity methods to interact with agent entities.
|
||||
|
||||
This class implements AgentProtocol and provides methods to work with Azure Durable Functions
|
||||
orchestrations, which use generators and yield instead of async/await.
|
||||
|
||||
Key methods:
|
||||
- get_new_thread(): Create a new conversation thread
|
||||
- run(): Execute the agent and return a Task for yielding in orchestrations
|
||||
|
||||
Note: The run() method is NOT async. It returns a Task directly that must be
|
||||
yielded in orchestrations to wait for the entity call to complete.
|
||||
|
||||
Example usage in orchestration:
|
||||
writer = app.get_agent(context, "WriterAgent")
|
||||
thread = writer.get_new_thread() # NOT yielded - returns immediately
|
||||
|
||||
response = yield writer.run( # Yielded - waits for entity call
|
||||
message="Write a haiku about coding",
|
||||
thread=thread
|
||||
)
|
||||
"""
|
||||
|
||||
def __init__(self, context: AgentOrchestrationContextType, agent_name: str):
|
||||
"""Initialize the DurableAIAgent.
|
||||
|
||||
Args:
|
||||
context: The orchestration context
|
||||
agent_name: Name of the agent (used to construct entity ID)
|
||||
"""
|
||||
def __init__(self, context: AgentOrchestrationContextType):
|
||||
self.context = context
|
||||
self.agent_name = agent_name
|
||||
self._id = str(uuid.uuid4())
|
||||
self._name = agent_name
|
||||
self._display_name = agent_name
|
||||
self._description = f"Durable agent proxy for {agent_name}"
|
||||
logger.debug("[DurableAIAgent] Initialized for agent: %s", agent_name)
|
||||
|
||||
@property
|
||||
def id(self) -> str:
|
||||
"""Get the unique identifier for this agent."""
|
||||
return self._id
|
||||
def generate_unique_id(self) -> str:
|
||||
return str(self.context.new_uuid())
|
||||
|
||||
@property
|
||||
def name(self) -> str | None:
|
||||
"""Get the name of the agent."""
|
||||
return self._name
|
||||
|
||||
@property
|
||||
def display_name(self) -> str:
|
||||
"""Get the display name of the agent."""
|
||||
return self._display_name
|
||||
|
||||
@property
|
||||
def description(self) -> str | None:
|
||||
"""Get the description of the agent."""
|
||||
return self._description
|
||||
|
||||
# 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]
|
||||
def get_run_request(
|
||||
self,
|
||||
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
|
||||
*,
|
||||
thread: AgentThread | None = None,
|
||||
response_format: type[BaseModel] | None = None,
|
||||
**kwargs: Any,
|
||||
) -> AgentTask:
|
||||
"""Execute the agent with messages and return an AgentTask for 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
|
||||
response_format: Optional Pydantic model for response parsing
|
||||
**kwargs: Additional arguments (enable_tool_calls)
|
||||
message: str,
|
||||
response_format: type[BaseModel] | None,
|
||||
enable_tool_calls: bool,
|
||||
) -> RunRequest:
|
||||
"""Get the current run request from the orchestration context.
|
||||
|
||||
Returns:
|
||||
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()
|
||||
response = yield agent.run("Hello", thread=thread)
|
||||
# response is typed as AgentRunResponse
|
||||
RunRequest: The current run request
|
||||
"""
|
||||
message_str = self._normalize_messages(messages)
|
||||
request = super().get_run_request(
|
||||
message,
|
||||
response_format,
|
||||
enable_tool_calls,
|
||||
)
|
||||
request.orchestration_id = self.context.instance_id
|
||||
return request
|
||||
|
||||
# Extract optional parameters from kwargs
|
||||
enable_tool_calls = kwargs.get("enable_tool_calls", True)
|
||||
def run_durable_agent(
|
||||
self,
|
||||
agent_name: str,
|
||||
run_request: RunRequest,
|
||||
thread: AgentThread | None = None,
|
||||
) -> AgentTask:
|
||||
|
||||
# Get the session ID for the entity
|
||||
if isinstance(thread, DurableAgentThread) and thread.session_id is not None:
|
||||
session_id = thread.session_id
|
||||
else:
|
||||
# Create a unique session ID for each call when no thread is provided
|
||||
# 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.debug("[DurableAIAgent] No thread provided, created unique session_id: %s", session_id)
|
||||
# Resolve session
|
||||
session_id = self._create_session_id(agent_name, thread)
|
||||
|
||||
# Create entity ID from session ID
|
||||
entity_id = session_id.to_entity_id()
|
||||
entity_id = df.EntityId(
|
||||
name=session_id.entity_name,
|
||||
key=session_id.key,
|
||||
)
|
||||
|
||||
# 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,
|
||||
"[AzureFunctionsAgentProvider] correlation_id: %s entity_id: %s session_id: %s",
|
||||
run_request.correlation_id,
|
||||
entity_id,
|
||||
session_id,
|
||||
)
|
||||
|
||||
# Prepare the request using RunRequest model
|
||||
# Include the orchestration's instance_id so it can be stored in the agent's entity state
|
||||
run_request = RunRequest(
|
||||
message=message_str,
|
||||
enable_tool_calls=enable_tool_calls,
|
||||
correlation_id=correlation_id,
|
||||
response_format=response_format,
|
||||
orchestration_id=self.context.instance_id,
|
||||
created_at=self.context.current_utc_datetime,
|
||||
)
|
||||
|
||||
logger.debug("[DurableAIAgent] Calling entity %s with message: %s", entity_id, message_str[:100])
|
||||
|
||||
# Call the entity to get the underlying task
|
||||
entity_task = self.context.call_entity(entity_id, "run", run_request.to_dict())
|
||||
|
||||
# Wrap it in an AgentTask that will convert the result to AgentRunResponse
|
||||
agent_task = AgentTask(
|
||||
return AgentTask(
|
||||
entity_task=entity_task,
|
||||
response_format=response_format,
|
||||
correlation_id=correlation_id,
|
||||
response_format=run_request.response_format,
|
||||
correlation_id=run_request.correlation_id,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"[DurableAIAgent] Created AgentTask for correlation_id %s",
|
||||
correlation_id,
|
||||
)
|
||||
|
||||
return agent_task
|
||||
|
||||
def run_stream(
|
||||
self,
|
||||
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
|
||||
*,
|
||||
thread: AgentThread | None = None,
|
||||
**kwargs: Any,
|
||||
) -> AsyncIterator[AgentRunResponseUpdate]:
|
||||
"""Run the agent with streaming (not supported for durable agents).
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Streaming is not supported for durable agents.
|
||||
"""
|
||||
raise NotImplementedError("Streaming is not supported for durable agents in orchestrations.")
|
||||
|
||||
def get_new_thread(self, **kwargs: Any) -> AgentThread:
|
||||
"""Create a new agent thread for this orchestration instance.
|
||||
|
||||
Each call creates a unique thread with its own conversation context.
|
||||
The session ID is deterministic (uses context.new_uuid()) to ensure
|
||||
orchestration replay works correctly.
|
||||
|
||||
Returns:
|
||||
A new AgentThread instance with a unique session ID
|
||||
"""
|
||||
# Generate a deterministic unique key for this thread
|
||||
# Using context.new_uuid() ensures the same GUID is generated during replay
|
||||
session_key = str(self.context.new_uuid())
|
||||
|
||||
# Create AgentSessionId with agent name and session key
|
||||
session_id = AgentSessionId(name=self.agent_name, key=session_key)
|
||||
|
||||
thread = DurableAgentThread.from_session_id(session_id, **kwargs)
|
||||
|
||||
logger.debug("[DurableAIAgent] Created new thread with session_id: %s", session_id)
|
||||
return thread
|
||||
|
||||
def _messages_to_string(self, messages: list[ChatMessage]) -> str:
|
||||
"""Convert a list of ChatMessage objects to a single string.
|
||||
|
||||
Args:
|
||||
messages: List of ChatMessage objects
|
||||
|
||||
Returns:
|
||||
Concatenated string of message contents
|
||||
"""
|
||||
return "\n".join([msg.text or "" for msg in messages])
|
||||
|
||||
def _normalize_messages(self, messages: str | ChatMessage | list[str] | list[ChatMessage] | None) -> str:
|
||||
"""Convert supported message inputs to a single string."""
|
||||
if messages is None:
|
||||
return ""
|
||||
if isinstance(messages, str):
|
||||
return messages
|
||||
if isinstance(messages, ChatMessage):
|
||||
return messages.text or ""
|
||||
if isinstance(messages, list):
|
||||
if not messages:
|
||||
return ""
|
||||
first_item = messages[0]
|
||||
if isinstance(first_item, str):
|
||||
return "\n".join(cast(list[str], messages))
|
||||
return self._messages_to_string(cast(list[ChatMessage], messages))
|
||||
return str(messages)
|
||||
|
||||
@@ -1,402 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Unit tests for data models (AgentSessionId, RunRequest, AgentResponse)."""
|
||||
|
||||
import azure.durable_functions as df
|
||||
import pytest
|
||||
from agent_framework import Role
|
||||
from agent_framework_durabletask import RunRequest
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agent_framework_azurefunctions._models import AgentSessionId
|
||||
|
||||
|
||||
class ModuleStructuredResponse(BaseModel):
|
||||
value: int
|
||||
|
||||
|
||||
class TestAgentSessionId:
|
||||
"""Test suite for AgentSessionId."""
|
||||
|
||||
def test_init_creates_session_id(self) -> None:
|
||||
"""Test that AgentSessionId initializes correctly."""
|
||||
session_id = AgentSessionId(name="AgentEntity", key="test-key-123")
|
||||
|
||||
assert session_id.name == "AgentEntity"
|
||||
assert session_id.key == "test-key-123"
|
||||
|
||||
def test_with_random_key_generates_guid(self) -> None:
|
||||
"""Test that with_random_key generates a GUID."""
|
||||
session_id = AgentSessionId.with_random_key(name="AgentEntity")
|
||||
|
||||
assert session_id.name == "AgentEntity"
|
||||
assert len(session_id.key) == 32 # UUID hex is 32 chars
|
||||
# Verify it's a valid hex string
|
||||
int(session_id.key, 16)
|
||||
|
||||
def test_with_random_key_unique_keys(self) -> None:
|
||||
"""Test that with_random_key generates unique keys."""
|
||||
session_id1 = AgentSessionId.with_random_key(name="AgentEntity")
|
||||
session_id2 = AgentSessionId.with_random_key(name="AgentEntity")
|
||||
|
||||
assert session_id1.key != session_id2.key
|
||||
|
||||
def test_to_entity_id_conversion(self) -> None:
|
||||
"""Test conversion to EntityId."""
|
||||
session_id = AgentSessionId(name="AgentEntity", key="test-key")
|
||||
entity_id = session_id.to_entity_id()
|
||||
|
||||
assert isinstance(entity_id, df.EntityId)
|
||||
assert entity_id.name == "dafx-AgentEntity"
|
||||
assert entity_id.key == "test-key"
|
||||
|
||||
def test_from_entity_id_conversion(self) -> None:
|
||||
"""Test creation from EntityId."""
|
||||
entity_id = df.EntityId(name="dafx-AgentEntity", key="test-key")
|
||||
session_id = AgentSessionId.from_entity_id(entity_id)
|
||||
|
||||
assert isinstance(session_id, AgentSessionId)
|
||||
assert session_id.name == "AgentEntity"
|
||||
assert session_id.key == "test-key"
|
||||
|
||||
def test_round_trip_entity_id_conversion(self) -> None:
|
||||
"""Test round-trip conversion to and from EntityId."""
|
||||
original = AgentSessionId(name="AgentEntity", key="test-key")
|
||||
entity_id = original.to_entity_id()
|
||||
restored = AgentSessionId.from_entity_id(entity_id)
|
||||
|
||||
assert restored.name == original.name
|
||||
assert restored.key == original.key
|
||||
|
||||
def test_str_representation(self) -> None:
|
||||
"""Test string representation."""
|
||||
session_id = AgentSessionId(name="AgentEntity", key="test-key-123")
|
||||
str_repr = str(session_id)
|
||||
|
||||
assert str_repr == "@AgentEntity@test-key-123"
|
||||
|
||||
def test_repr_representation(self) -> None:
|
||||
"""Test repr representation."""
|
||||
session_id = AgentSessionId(name="AgentEntity", key="test-key")
|
||||
repr_str = repr(session_id)
|
||||
|
||||
assert "AgentSessionId" in repr_str
|
||||
assert "AgentEntity" in repr_str
|
||||
assert "test-key" in repr_str
|
||||
|
||||
def test_parse_valid_session_id(self) -> None:
|
||||
"""Test parsing valid session ID string."""
|
||||
session_id = AgentSessionId.parse("@AgentEntity@test-key-123")
|
||||
|
||||
assert session_id.name == "AgentEntity"
|
||||
assert session_id.key == "test-key-123"
|
||||
|
||||
def test_parse_invalid_format_no_prefix(self) -> None:
|
||||
"""Test parsing invalid format without @ prefix."""
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
AgentSessionId.parse("AgentEntity@test-key")
|
||||
|
||||
assert "Invalid agent session ID format" in str(exc_info.value)
|
||||
|
||||
def test_parse_invalid_format_single_part(self) -> None:
|
||||
"""Test parsing invalid format with single part."""
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
AgentSessionId.parse("@AgentEntity")
|
||||
|
||||
assert "Invalid agent session ID format" in str(exc_info.value)
|
||||
|
||||
def test_parse_with_multiple_at_signs_in_key(self) -> None:
|
||||
"""Test parsing with @ signs in the key."""
|
||||
session_id = AgentSessionId.parse("@AgentEntity@key-with@symbols")
|
||||
|
||||
assert session_id.name == "AgentEntity"
|
||||
assert session_id.key == "key-with@symbols"
|
||||
|
||||
def test_parse_round_trip(self) -> None:
|
||||
"""Test round-trip parse and string conversion."""
|
||||
original = AgentSessionId(name="AgentEntity", key="test-key")
|
||||
str_repr = str(original)
|
||||
parsed = AgentSessionId.parse(str_repr)
|
||||
|
||||
assert parsed.name == original.name
|
||||
assert parsed.key == original.key
|
||||
|
||||
def test_to_entity_name_adds_prefix(self) -> None:
|
||||
"""Test that to_entity_name adds the dafx- prefix."""
|
||||
entity_name = AgentSessionId.to_entity_name("TestAgent")
|
||||
assert entity_name == "dafx-TestAgent"
|
||||
|
||||
def test_from_entity_id_strips_prefix(self) -> None:
|
||||
"""Test that from_entity_id strips the dafx- prefix."""
|
||||
entity_id = df.EntityId(name="dafx-TestAgent", key="key123")
|
||||
session_id = AgentSessionId.from_entity_id(entity_id)
|
||||
|
||||
assert session_id.name == "TestAgent"
|
||||
assert session_id.key == "key123"
|
||||
|
||||
def test_from_entity_id_raises_without_prefix(self) -> None:
|
||||
"""Test that from_entity_id raises ValueError when entity name lacks the prefix."""
|
||||
entity_id = df.EntityId(name="TestAgent", key="key123")
|
||||
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
AgentSessionId.from_entity_id(entity_id)
|
||||
|
||||
assert "not a valid agent session ID" in str(exc_info.value)
|
||||
assert "dafx-" in str(exc_info.value)
|
||||
|
||||
|
||||
class TestRunRequest:
|
||||
"""Test suite for RunRequest."""
|
||||
|
||||
def test_init_with_defaults(self) -> None:
|
||||
"""Test RunRequest initialization with defaults."""
|
||||
request = RunRequest(message="Hello")
|
||||
|
||||
assert request.message == "Hello"
|
||||
assert request.role == Role.USER
|
||||
assert request.response_format is None
|
||||
assert request.enable_tool_calls is True
|
||||
|
||||
def test_init_with_all_fields(self) -> None:
|
||||
"""Test RunRequest initialization with all fields."""
|
||||
schema = ModuleStructuredResponse
|
||||
request = RunRequest(
|
||||
message="Hello",
|
||||
role=Role.SYSTEM,
|
||||
response_format=schema,
|
||||
enable_tool_calls=False,
|
||||
)
|
||||
|
||||
assert request.message == "Hello"
|
||||
assert request.role == Role.SYSTEM
|
||||
assert request.response_format is schema
|
||||
assert request.enable_tool_calls is False
|
||||
|
||||
def test_init_coerces_string_role(self) -> None:
|
||||
"""Ensure string role values are coerced into Role instances."""
|
||||
request = RunRequest(message="Hello", role="system") # type: ignore[arg-type]
|
||||
|
||||
assert request.role == Role.SYSTEM
|
||||
|
||||
def test_to_dict_with_defaults(self) -> None:
|
||||
"""Test to_dict with default values."""
|
||||
request = RunRequest(message="Test message")
|
||||
data = request.to_dict()
|
||||
|
||||
assert data["message"] == "Test message"
|
||||
assert data["enable_tool_calls"] is True
|
||||
assert data["role"] == "user"
|
||||
assert "response_format" not in data or data["response_format"] is None
|
||||
assert "thread_id" not in data
|
||||
|
||||
def test_to_dict_with_all_fields(self) -> None:
|
||||
"""Test to_dict with all fields."""
|
||||
schema = ModuleStructuredResponse
|
||||
request = RunRequest(
|
||||
message="Hello",
|
||||
role=Role.ASSISTANT,
|
||||
response_format=schema,
|
||||
enable_tool_calls=False,
|
||||
)
|
||||
data = request.to_dict()
|
||||
|
||||
assert data["message"] == "Hello"
|
||||
assert data["role"] == "assistant"
|
||||
assert data["response_format"]["__response_schema_type__"] == "pydantic_model"
|
||||
assert data["response_format"]["module"] == schema.__module__
|
||||
assert data["response_format"]["qualname"] == schema.__qualname__
|
||||
assert data["enable_tool_calls"] is False
|
||||
assert "thread_id" not in data
|
||||
|
||||
def test_from_dict_with_defaults(self) -> None:
|
||||
"""Test from_dict with minimal data."""
|
||||
data = {"message": "Hello"}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.message == "Hello"
|
||||
assert request.role == Role.USER
|
||||
assert request.enable_tool_calls is True
|
||||
|
||||
def test_from_dict_ignores_thread_id_field(self) -> None:
|
||||
"""Ensure legacy thread_id input does not break RunRequest parsing."""
|
||||
request = RunRequest.from_dict({"message": "Hello", "thread_id": "ignored"})
|
||||
|
||||
assert request.message == "Hello"
|
||||
|
||||
def test_from_dict_with_all_fields(self) -> None:
|
||||
"""Test from_dict with all fields."""
|
||||
data = {
|
||||
"message": "Test",
|
||||
"role": "system",
|
||||
"response_format": {
|
||||
"__response_schema_type__": "pydantic_model",
|
||||
"module": ModuleStructuredResponse.__module__,
|
||||
"qualname": ModuleStructuredResponse.__qualname__,
|
||||
},
|
||||
"enable_tool_calls": False,
|
||||
}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.message == "Test"
|
||||
assert request.role == Role.SYSTEM
|
||||
assert request.response_format is ModuleStructuredResponse
|
||||
assert request.enable_tool_calls is False
|
||||
|
||||
def test_from_dict_with_unknown_role_preserves_value(self) -> None:
|
||||
"""Test from_dict keeps custom roles intact."""
|
||||
data = {"message": "Test", "role": "reviewer"}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.role.value == "reviewer"
|
||||
assert request.role != Role.USER
|
||||
|
||||
def test_from_dict_empty_message(self) -> None:
|
||||
"""Test from_dict with empty message."""
|
||||
request = RunRequest.from_dict({})
|
||||
|
||||
assert request.message == ""
|
||||
assert request.role == Role.USER
|
||||
|
||||
def test_round_trip_dict_conversion(self) -> None:
|
||||
"""Test round-trip to_dict and from_dict."""
|
||||
original = RunRequest(
|
||||
message="Test message",
|
||||
role=Role.SYSTEM,
|
||||
response_format=ModuleStructuredResponse,
|
||||
enable_tool_calls=False,
|
||||
)
|
||||
|
||||
data = original.to_dict()
|
||||
restored = RunRequest.from_dict(data)
|
||||
|
||||
assert restored.message == original.message
|
||||
assert restored.role == original.role
|
||||
assert restored.response_format is ModuleStructuredResponse
|
||||
assert restored.enable_tool_calls == original.enable_tool_calls
|
||||
|
||||
def test_round_trip_with_pydantic_response_format(self) -> None:
|
||||
"""Ensure Pydantic response formats serialize and deserialize properly."""
|
||||
original = RunRequest(
|
||||
message="Structured",
|
||||
response_format=ModuleStructuredResponse,
|
||||
)
|
||||
|
||||
data = original.to_dict()
|
||||
|
||||
assert data["response_format"]["__response_schema_type__"] == "pydantic_model"
|
||||
assert data["response_format"]["module"] == ModuleStructuredResponse.__module__
|
||||
assert data["response_format"]["qualname"] == ModuleStructuredResponse.__qualname__
|
||||
|
||||
restored = RunRequest.from_dict(data)
|
||||
assert restored.response_format is ModuleStructuredResponse
|
||||
|
||||
def test_init_with_correlationId(self) -> None:
|
||||
"""Test RunRequest initialization with correlationId."""
|
||||
request = RunRequest(message="Test message", correlation_id="corr-123")
|
||||
|
||||
assert request.message == "Test message"
|
||||
assert request.correlation_id == "corr-123"
|
||||
|
||||
def test_to_dict_with_correlationId(self) -> None:
|
||||
"""Test to_dict includes correlationId."""
|
||||
request = RunRequest(message="Test", correlation_id="corr-456")
|
||||
data = request.to_dict()
|
||||
|
||||
assert data["message"] == "Test"
|
||||
assert data["correlationId"] == "corr-456"
|
||||
|
||||
def test_from_dict_with_correlationId(self) -> None:
|
||||
"""Test from_dict with correlationId."""
|
||||
data = {"message": "Test", "correlationId": "corr-789"}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.message == "Test"
|
||||
assert request.correlation_id == "corr-789"
|
||||
|
||||
def test_round_trip_with_correlationId(self) -> None:
|
||||
"""Test round-trip to_dict and from_dict with correlationId."""
|
||||
original = RunRequest(
|
||||
message="Test message",
|
||||
role=Role.SYSTEM,
|
||||
correlation_id="corr-123",
|
||||
)
|
||||
|
||||
data = original.to_dict()
|
||||
restored = RunRequest.from_dict(data)
|
||||
|
||||
assert restored.message == original.message
|
||||
assert restored.role == original.role
|
||||
assert restored.correlation_id == original.correlation_id
|
||||
|
||||
def test_init_with_orchestration_id(self) -> None:
|
||||
"""Test RunRequest initialization with orchestration_id."""
|
||||
request = RunRequest(
|
||||
message="Test message",
|
||||
orchestration_id="orch-123",
|
||||
)
|
||||
|
||||
assert request.message == "Test message"
|
||||
assert request.orchestration_id == "orch-123"
|
||||
|
||||
def test_to_dict_with_orchestration_id(self) -> None:
|
||||
"""Test to_dict includes orchestrationId."""
|
||||
request = RunRequest(
|
||||
message="Test",
|
||||
orchestration_id="orch-456",
|
||||
)
|
||||
data = request.to_dict()
|
||||
|
||||
assert data["message"] == "Test"
|
||||
assert data["orchestrationId"] == "orch-456"
|
||||
|
||||
def test_to_dict_excludes_orchestration_id_when_none(self) -> None:
|
||||
"""Test to_dict excludes orchestrationId when not set."""
|
||||
request = RunRequest(
|
||||
message="Test",
|
||||
)
|
||||
data = request.to_dict()
|
||||
|
||||
assert "orchestrationId" not in data
|
||||
|
||||
def test_from_dict_with_orchestration_id(self) -> None:
|
||||
"""Test from_dict with orchestrationId."""
|
||||
data = {
|
||||
"message": "Test",
|
||||
"orchestrationId": "orch-789",
|
||||
}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.message == "Test"
|
||||
assert request.orchestration_id == "orch-789"
|
||||
|
||||
def test_round_trip_with_orchestration_id(self) -> None:
|
||||
"""Test round-trip to_dict and from_dict with orchestration_id."""
|
||||
original = RunRequest(
|
||||
message="Test message",
|
||||
role=Role.SYSTEM,
|
||||
correlation_id="corr-123",
|
||||
orchestration_id="orch-123",
|
||||
)
|
||||
|
||||
data = original.to_dict()
|
||||
restored = RunRequest.from_dict(data)
|
||||
|
||||
assert restored.message == original.message
|
||||
assert restored.role == original.role
|
||||
assert restored.correlation_id == original.correlation_id
|
||||
assert restored.orchestration_id == original.orchestration_id
|
||||
|
||||
|
||||
class TestModelIntegration:
|
||||
"""Test suite for integration between models."""
|
||||
|
||||
def test_run_request_with_session_id_string(self) -> None:
|
||||
"""AgentSessionId string can still be used by callers, but is not stored on RunRequest."""
|
||||
session_id = AgentSessionId.with_random_key("AgentEntity")
|
||||
session_id_str = str(session_id)
|
||||
|
||||
assert session_id_str.startswith("@AgentEntity@")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "--tb=short"])
|
||||
@@ -6,11 +6,11 @@ from typing import Any
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentRunResponse, AgentThread, ChatMessage
|
||||
from agent_framework import AgentRunResponse, ChatMessage
|
||||
from agent_framework_durabletask import DurableAIAgent
|
||||
from azure.durable_functions.models.Task import TaskBase, TaskState
|
||||
|
||||
from agent_framework_azurefunctions import AgentFunctionApp, DurableAIAgent
|
||||
from agent_framework_azurefunctions._models import AgentSessionId, DurableAgentThread
|
||||
from agent_framework_azurefunctions import AgentFunctionApp
|
||||
from agent_framework_azurefunctions._orchestration import AgentTask
|
||||
|
||||
|
||||
@@ -38,46 +38,96 @@ def _create_entity_task(task_id: int = 1) -> TaskBase:
|
||||
return _FakeTask(task_id)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_context():
|
||||
"""Create a mock orchestration context with UUID support."""
|
||||
context = Mock()
|
||||
context.instance_id = "test-instance"
|
||||
context.current_utc_datetime = Mock()
|
||||
return context
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_context_with_uuid() -> tuple[Mock, str]:
|
||||
"""Create a mock context with a single UUID."""
|
||||
from uuid import UUID
|
||||
|
||||
context = Mock()
|
||||
context.instance_id = "test-instance"
|
||||
context.current_utc_datetime = Mock()
|
||||
test_uuid = UUID("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa")
|
||||
context.new_uuid = Mock(return_value=test_uuid)
|
||||
return context, test_uuid.hex
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_context_with_multiple_uuids() -> tuple[Mock, list[str]]:
|
||||
"""Create a mock context with multiple UUIDs via side_effect."""
|
||||
from uuid import UUID
|
||||
|
||||
context = Mock()
|
||||
context.instance_id = "test-instance"
|
||||
context.current_utc_datetime = Mock()
|
||||
uuids = [
|
||||
UUID("aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"),
|
||||
UUID("bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb"),
|
||||
UUID("cccccccc-cccc-cccc-cccc-cccccccccccc"),
|
||||
]
|
||||
context.new_uuid = Mock(side_effect=uuids)
|
||||
# Return the hex versions for assertion checking
|
||||
hex_uuids = [uuid.hex for uuid in uuids]
|
||||
return context, hex_uuids
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def executor_with_uuid() -> tuple[Any, Mock, str]:
|
||||
"""Create an executor with a mocked generate_unique_id method."""
|
||||
from agent_framework_azurefunctions._orchestration import AzureFunctionsAgentExecutor
|
||||
|
||||
context = Mock()
|
||||
context.instance_id = "test-instance"
|
||||
context.current_utc_datetime = Mock()
|
||||
|
||||
executor = AzureFunctionsAgentExecutor(context)
|
||||
test_uuid_hex = "aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa"
|
||||
executor.generate_unique_id = Mock(return_value=test_uuid_hex)
|
||||
|
||||
return executor, context, test_uuid_hex
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def executor_with_multiple_uuids() -> tuple[Any, Mock, list[str]]:
|
||||
"""Create an executor with multiple mocked UUIDs."""
|
||||
from agent_framework_azurefunctions._orchestration import AzureFunctionsAgentExecutor
|
||||
|
||||
context = Mock()
|
||||
context.instance_id = "test-instance"
|
||||
context.current_utc_datetime = Mock()
|
||||
|
||||
executor = AzureFunctionsAgentExecutor(context)
|
||||
uuid_hexes = [
|
||||
"aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa",
|
||||
"bbbbbbbb-bbbb-bbbb-bbbb-bbbbbbbbbbbb",
|
||||
"cccccccc-cccc-cccc-cccc-cccccccccccc",
|
||||
"dddddddd-dddd-dddd-dddd-dddddddddddd",
|
||||
"eeeeeeee-eeee-eeee-eeee-eeeeeeeeeeee",
|
||||
]
|
||||
executor.generate_unique_id = Mock(side_effect=uuid_hexes)
|
||||
|
||||
return executor, context, uuid_hexes
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def executor_with_context(mock_context_with_uuid: tuple[Mock, str]) -> tuple[Any, Mock]:
|
||||
"""Create an executor with a mocked context."""
|
||||
from agent_framework_azurefunctions._orchestration import AzureFunctionsAgentExecutor
|
||||
|
||||
context, _ = mock_context_with_uuid
|
||||
return AzureFunctionsAgentExecutor(context), context
|
||||
|
||||
|
||||
class TestAgentResponseHelpers:
|
||||
"""Tests for helper utilities that prepare AgentRunResponse values."""
|
||||
|
||||
@staticmethod
|
||||
def _create_agent_task() -> AgentTask:
|
||||
entity_task = _create_entity_task()
|
||||
return AgentTask(entity_task, None, "correlation-id")
|
||||
|
||||
def test_load_agent_response_from_instance(self) -> None:
|
||||
task = self._create_agent_task()
|
||||
response = AgentRunResponse(messages=[ChatMessage(role="assistant", text='{"foo": "bar"}')])
|
||||
|
||||
loaded = task._load_agent_response(response)
|
||||
|
||||
assert loaded is response
|
||||
assert loaded.value is None
|
||||
|
||||
def test_load_agent_response_from_serialized(self) -> None:
|
||||
task = self._create_agent_task()
|
||||
serialized = AgentRunResponse(messages=[ChatMessage(role="assistant", text="structured")]).to_dict()
|
||||
serialized["value"] = {"answer": 42}
|
||||
|
||||
loaded = task._load_agent_response(serialized)
|
||||
|
||||
assert loaded is not None
|
||||
assert loaded.value == {"answer": 42}
|
||||
loaded_dict = loaded.to_dict()
|
||||
assert loaded_dict["type"] == "agent_run_response"
|
||||
|
||||
def test_load_agent_response_rejects_none(self) -> None:
|
||||
task = self._create_agent_task()
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
task._load_agent_response(None)
|
||||
|
||||
def test_load_agent_response_rejects_unsupported_type(self) -> None:
|
||||
task = self._create_agent_task()
|
||||
|
||||
with pytest.raises(TypeError, match="Unsupported type"):
|
||||
task._load_agent_response(["invalid", "list"]) # type: ignore[arg-type]
|
||||
"""Tests for response handling through public AgentTask API."""
|
||||
|
||||
def test_try_set_value_success(self) -> None:
|
||||
"""Test try_set_value correctly processes successful task completion."""
|
||||
@@ -144,335 +194,10 @@ class TestAgentResponseHelpers:
|
||||
assert isinstance(task.result.value, TestSchema)
|
||||
assert task.result.value.answer == "42"
|
||||
|
||||
def test_ensure_response_format_parses_value(self) -> None:
|
||||
"""Test _ensure_response_format correctly parses response value."""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class SampleSchema(BaseModel):
|
||||
name: str
|
||||
|
||||
task = self._create_agent_task()
|
||||
response = AgentRunResponse(messages=[ChatMessage(role="assistant", text='{"name": "test"}')])
|
||||
|
||||
# Value should be None initially
|
||||
assert response.value is None
|
||||
|
||||
# Parse the value
|
||||
task._ensure_response_format(SampleSchema, "test-correlation", response)
|
||||
|
||||
# Value should now be parsed
|
||||
assert isinstance(response.value, SampleSchema)
|
||||
assert response.value.name == "test"
|
||||
|
||||
def test_ensure_response_format_skips_if_already_parsed(self) -> None:
|
||||
"""Test _ensure_response_format does not re-parse if value already matches format."""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class SampleSchema(BaseModel):
|
||||
name: str
|
||||
|
||||
task = self._create_agent_task()
|
||||
existing_value = SampleSchema(name="existing")
|
||||
response = AgentRunResponse(
|
||||
messages=[ChatMessage(role="assistant", text='{"name": "new"}')],
|
||||
value=existing_value,
|
||||
)
|
||||
|
||||
# Call _ensure_response_format
|
||||
task._ensure_response_format(SampleSchema, "test-correlation", response)
|
||||
|
||||
# Value should remain unchanged (not re-parsed)
|
||||
assert response.value is existing_value
|
||||
assert response.value.name == "existing"
|
||||
|
||||
|
||||
class TestDurableAIAgent:
|
||||
"""Test suite for DurableAIAgent wrapper."""
|
||||
|
||||
def test_init(self) -> None:
|
||||
"""Test DurableAIAgent initialization."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-123"
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
assert agent.context == mock_context
|
||||
assert agent.agent_name == "TestAgent"
|
||||
|
||||
def test_implements_agent_protocol(self) -> None:
|
||||
"""Test that DurableAIAgent implements AgentProtocol."""
|
||||
from agent_framework import AgentProtocol
|
||||
|
||||
mock_context = Mock()
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
# Check that agent satisfies AgentProtocol
|
||||
assert isinstance(agent, AgentProtocol)
|
||||
|
||||
def test_has_agent_protocol_properties(self) -> None:
|
||||
"""Test that DurableAIAgent has AgentProtocol properties."""
|
||||
mock_context = Mock()
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
# AgentProtocol properties
|
||||
assert hasattr(agent, "id")
|
||||
assert hasattr(agent, "name")
|
||||
assert hasattr(agent, "description")
|
||||
assert hasattr(agent, "display_name")
|
||||
|
||||
# Verify values
|
||||
assert agent.name == "TestAgent"
|
||||
assert agent.description == "Durable agent proxy for TestAgent"
|
||||
assert agent.display_name == "TestAgent"
|
||||
assert agent.id is not None # Auto-generated UUID
|
||||
|
||||
def test_get_new_thread(self) -> None:
|
||||
"""Test creating a new agent thread."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-456"
|
||||
mock_context.new_uuid = Mock(return_value="test-guid-456")
|
||||
|
||||
agent = DurableAIAgent(mock_context, "WriterAgent")
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
assert thread.session_id is not None
|
||||
session_id = thread.session_id
|
||||
assert isinstance(session_id, AgentSessionId)
|
||||
assert session_id.name == "WriterAgent"
|
||||
assert session_id.key == "test-guid-456"
|
||||
mock_context.new_uuid.assert_called_once()
|
||||
|
||||
def test_get_new_thread_deterministic(self) -> None:
|
||||
"""Test that get_new_thread creates deterministic session IDs."""
|
||||
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-789"
|
||||
mock_context.new_uuid = Mock(side_effect=["session-guid-1", "session-guid-2"])
|
||||
|
||||
agent = DurableAIAgent(mock_context, "EditorAgent")
|
||||
|
||||
# Create multiple threads - they should have unique session IDs
|
||||
thread1 = agent.get_new_thread()
|
||||
thread2 = agent.get_new_thread()
|
||||
|
||||
assert isinstance(thread1, DurableAgentThread)
|
||||
assert isinstance(thread2, DurableAgentThread)
|
||||
|
||||
session_id1 = thread1.session_id
|
||||
session_id2 = thread2.session_id
|
||||
assert session_id1 is not None and session_id2 is not None
|
||||
assert isinstance(session_id1, AgentSessionId)
|
||||
assert isinstance(session_id2, AgentSessionId)
|
||||
assert session_id1.name == "EditorAgent"
|
||||
assert session_id2.name == "EditorAgent"
|
||||
assert session_id1.key == "session-guid-1"
|
||||
assert session_id2.key == "session-guid-2"
|
||||
assert mock_context.new_uuid.call_count == 2
|
||||
|
||||
def test_run_creates_entity_call(self) -> None:
|
||||
"""Test that run() creates proper entity call and returns a Task."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-001"
|
||||
mock_context.new_uuid = Mock(side_effect=["thread-guid", "correlation-guid"])
|
||||
|
||||
entity_task = _create_entity_task()
|
||||
mock_context.call_entity = Mock(return_value=entity_task)
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
# Create thread
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
# Call run() - returns AgentTask directly
|
||||
task = agent.run(messages="Test message", thread=thread, enable_tool_calls=True)
|
||||
|
||||
assert isinstance(task, AgentTask)
|
||||
assert task.children[0] == entity_task
|
||||
|
||||
# Verify call_entity was called with correct parameters
|
||||
assert mock_context.call_entity.called
|
||||
call_args = mock_context.call_entity.call_args
|
||||
entity_id, operation, request = call_args[0]
|
||||
|
||||
assert operation == "run"
|
||||
assert request["message"] == "Test message"
|
||||
assert request["enable_tool_calls"] is True
|
||||
assert "correlationId" in request
|
||||
assert request["correlationId"] == "correlation-guid"
|
||||
assert "thread_id" not in request
|
||||
# Verify orchestration ID is set from context.instance_id
|
||||
assert "orchestrationId" in request
|
||||
assert request["orchestrationId"] == "test-instance-001"
|
||||
|
||||
def test_run_sets_orchestration_id(self) -> None:
|
||||
"""Test that run() sets the orchestration_id from context.instance_id."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "my-orchestration-123"
|
||||
mock_context.new_uuid = Mock(side_effect=["thread-guid", "correlation-guid"])
|
||||
|
||||
entity_task = _create_entity_task()
|
||||
mock_context.call_entity = Mock(return_value=entity_task)
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
agent.run(messages="Test", thread=thread)
|
||||
|
||||
call_args = mock_context.call_entity.call_args
|
||||
request = call_args[0][2]
|
||||
|
||||
assert request["orchestrationId"] == "my-orchestration-123"
|
||||
|
||||
def test_run_without_thread(self) -> None:
|
||||
"""Test that run() works without explicit thread (creates unique session key)."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-002"
|
||||
mock_context.new_uuid = Mock(side_effect=["auto-generated-guid", "correlation-guid"])
|
||||
|
||||
entity_task = _create_entity_task()
|
||||
mock_context.call_entity = Mock(return_value=entity_task)
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
# Call without thread
|
||||
task = agent.run(messages="Test message")
|
||||
|
||||
assert isinstance(task, AgentTask)
|
||||
assert task.children[0] == entity_task
|
||||
|
||||
# Verify the entity ID uses the auto-generated GUID with dafx- prefix
|
||||
call_args = mock_context.call_entity.call_args
|
||||
entity_id = call_args[0][0]
|
||||
assert entity_id.name == "dafx-TestAgent"
|
||||
assert entity_id.key == "auto-generated-guid"
|
||||
# Should be called twice: once for session_key, once for correlationId
|
||||
assert mock_context.new_uuid.call_count == 2
|
||||
|
||||
def test_run_with_response_format(self) -> None:
|
||||
"""Test that run() passes response format correctly."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-003"
|
||||
|
||||
entity_task = _create_entity_task()
|
||||
mock_context.call_entity = Mock(return_value=entity_task)
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
from pydantic import BaseModel
|
||||
|
||||
class SampleSchema(BaseModel):
|
||||
key: str
|
||||
|
||||
# Create thread and call
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
task = agent.run(messages="Test message", thread=thread, response_format=SampleSchema)
|
||||
|
||||
assert isinstance(task, AgentTask)
|
||||
assert task.children[0] == entity_task
|
||||
|
||||
# Verify schema was passed in the call_entity arguments
|
||||
call_args = mock_context.call_entity.call_args
|
||||
input_data = call_args[0][2] # Third argument is input_data
|
||||
assert "response_format" in input_data
|
||||
assert input_data["response_format"]["__response_schema_type__"] == "pydantic_model"
|
||||
assert input_data["response_format"]["module"] == SampleSchema.__module__
|
||||
assert input_data["response_format"]["qualname"] == SampleSchema.__qualname__
|
||||
|
||||
def test_messages_to_string(self) -> None:
|
||||
"""Test converting ChatMessage list to string."""
|
||||
from agent_framework import ChatMessage
|
||||
|
||||
mock_context = Mock()
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
messages = [
|
||||
ChatMessage(role="user", text="Hello"),
|
||||
ChatMessage(role="assistant", text="Hi there"),
|
||||
ChatMessage(role="user", text="How are you?"),
|
||||
]
|
||||
|
||||
result = agent._messages_to_string(messages)
|
||||
|
||||
assert result == "Hello\nHi there\nHow are you?"
|
||||
|
||||
def test_run_with_chat_message(self) -> None:
|
||||
"""Test that run() handles ChatMessage input."""
|
||||
from agent_framework import ChatMessage
|
||||
|
||||
mock_context = Mock()
|
||||
mock_context.new_uuid = Mock(side_effect=["thread-guid", "correlation-guid"])
|
||||
entity_task = _create_entity_task()
|
||||
mock_context.call_entity = Mock(return_value=entity_task)
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
# Call with ChatMessage
|
||||
msg = ChatMessage(role="user", text="Hello")
|
||||
task = agent.run(messages=msg, thread=thread)
|
||||
|
||||
assert isinstance(task, AgentTask)
|
||||
assert task.children[0] == entity_task
|
||||
|
||||
# Verify message was converted to string
|
||||
call_args = mock_context.call_entity.call_args
|
||||
request = call_args[0][2]
|
||||
assert request["message"] == "Hello"
|
||||
|
||||
def test_run_stream_raises_not_implemented(self) -> None:
|
||||
"""Test that run_stream() method raises NotImplementedError."""
|
||||
mock_context = Mock()
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
with pytest.raises(NotImplementedError) as exc_info:
|
||||
agent.run_stream("Test message")
|
||||
|
||||
error_msg = str(exc_info.value)
|
||||
assert "Streaming is not supported" in error_msg
|
||||
|
||||
def test_entity_id_format(self) -> None:
|
||||
"""Test that EntityId is created with correct format (name, key)."""
|
||||
from azure.durable_functions import EntityId
|
||||
|
||||
mock_context = Mock()
|
||||
mock_context.new_uuid = Mock(return_value="test-guid-789")
|
||||
mock_context.call_entity = Mock(return_value=_create_entity_task())
|
||||
|
||||
agent = DurableAIAgent(mock_context, "WriterAgent")
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
# Call run() to trigger entity ID creation
|
||||
agent.run("Test", thread=thread)
|
||||
|
||||
# Verify call_entity was called with correct EntityId
|
||||
call_args = mock_context.call_entity.call_args
|
||||
entity_id = call_args[0][0]
|
||||
|
||||
# EntityId should be EntityId(name="dafx-WriterAgent", key="test-guid-789")
|
||||
# Which formats as "@dafx-writeragent@test-guid-789"
|
||||
assert isinstance(entity_id, EntityId)
|
||||
assert entity_id.name == "dafx-WriterAgent"
|
||||
assert entity_id.key == "test-guid-789"
|
||||
assert str(entity_id) == "@dafx-writeragent@test-guid-789"
|
||||
|
||||
|
||||
class TestAgentFunctionAppGetAgent:
|
||||
"""Test suite for AgentFunctionApp.get_agent."""
|
||||
|
||||
def test_get_agent_method(self) -> None:
|
||||
"""Test get_agent method creates DurableAIAgent for registered agent."""
|
||||
app = _app_with_registered_agents("MyAgent")
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-100"
|
||||
|
||||
agent = app.get_agent(mock_context, "MyAgent")
|
||||
|
||||
assert isinstance(agent, DurableAIAgent)
|
||||
assert agent.agent_name == "MyAgent"
|
||||
assert agent.context == mock_context
|
||||
|
||||
def test_get_agent_raises_for_unregistered_agent(self) -> None:
|
||||
"""Test get_agent raises ValueError when agent is not registered."""
|
||||
app = _app_with_registered_agents("KnownAgent")
|
||||
@@ -484,15 +209,9 @@ class TestAgentFunctionAppGetAgent:
|
||||
class TestOrchestrationIntegration:
|
||||
"""Integration tests for orchestration scenarios."""
|
||||
|
||||
def test_sequential_agent_calls_simulation(self) -> None:
|
||||
def test_sequential_agent_calls_simulation(self, executor_with_multiple_uuids: tuple[Any, Mock, list[str]]) -> None:
|
||||
"""Simulate sequential agent calls in an orchestration."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-orchestration-001"
|
||||
# new_uuid will be called 3 times:
|
||||
# 1. thread creation
|
||||
# 2. correlationId for first call
|
||||
# 3. correlationId for second call
|
||||
mock_context.new_uuid = Mock(side_effect=["deterministic-guid-001", "corr-1", "corr-2"])
|
||||
executor, context, uuid_hexes = executor_with_multiple_uuids
|
||||
|
||||
# Track entity calls
|
||||
entity_calls: list[dict[str, Any]] = []
|
||||
@@ -501,10 +220,10 @@ class TestOrchestrationIntegration:
|
||||
entity_calls.append({"entity_id": str(entity_id), "operation": operation, "input": input_data})
|
||||
return _create_entity_task()
|
||||
|
||||
mock_context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
|
||||
context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
|
||||
|
||||
app = _app_with_registered_agents("WriterAgent")
|
||||
agent = app.get_agent(mock_context, "WriterAgent")
|
||||
# Create agent directly with executor (not via app.get_agent)
|
||||
agent = DurableAIAgent(executor, "WriterAgent")
|
||||
|
||||
# Create thread
|
||||
thread = agent.get_new_thread()
|
||||
@@ -520,18 +239,15 @@ class TestOrchestrationIntegration:
|
||||
# Verify both calls used the same entity (same session key)
|
||||
assert len(entity_calls) == 2
|
||||
assert entity_calls[0]["entity_id"] == entity_calls[1]["entity_id"]
|
||||
# EntityId format is @dafx-writeragent@deterministic-guid-001
|
||||
assert entity_calls[0]["entity_id"] == "@dafx-writeragent@deterministic-guid-001"
|
||||
# new_uuid called 3 times: thread + 2 correlation IDs
|
||||
assert mock_context.new_uuid.call_count == 3
|
||||
# EntityId format is @dafx-writeragent@<uuid_hex>
|
||||
expected_entity_id = f"@dafx-writeragent@{uuid_hexes[0]}"
|
||||
assert entity_calls[0]["entity_id"] == expected_entity_id
|
||||
# generate_unique_id called 3 times: thread + 2 correlation IDs
|
||||
assert executor.generate_unique_id.call_count == 3
|
||||
|
||||
def test_multiple_agents_in_orchestration(self) -> None:
|
||||
def test_multiple_agents_in_orchestration(self, executor_with_multiple_uuids: tuple[Any, Mock, list[str]]) -> None:
|
||||
"""Test using multiple different agents in one orchestration."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-orchestration-002"
|
||||
# Mock new_uuid to return different GUIDs for each call
|
||||
# Order: writer thread, editor thread, writer correlation, editor correlation
|
||||
mock_context.new_uuid = Mock(side_effect=["writer-guid-001", "editor-guid-002", "writer-corr", "editor-corr"])
|
||||
executor, context, uuid_hexes = executor_with_multiple_uuids
|
||||
|
||||
entity_calls: list[str] = []
|
||||
|
||||
@@ -539,11 +255,11 @@ class TestOrchestrationIntegration:
|
||||
entity_calls.append(str(entity_id))
|
||||
return _create_entity_task()
|
||||
|
||||
mock_context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
|
||||
context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
|
||||
|
||||
app = _app_with_registered_agents("WriterAgent", "EditorAgent")
|
||||
writer = app.get_agent(mock_context, "WriterAgent")
|
||||
editor = app.get_agent(mock_context, "EditorAgent")
|
||||
# Create agents directly with executor (not via app.get_agent)
|
||||
writer = DurableAIAgent(executor, "WriterAgent")
|
||||
editor = DurableAIAgent(executor, "EditorAgent")
|
||||
|
||||
writer_thread = writer.get_new_thread()
|
||||
editor_thread = editor.get_new_thread()
|
||||
@@ -557,62 +273,11 @@ class TestOrchestrationIntegration:
|
||||
|
||||
# Verify different entity IDs were used
|
||||
assert len(entity_calls) == 2
|
||||
# EntityId format is @dafx-agentname@guid (lowercased agent name with dafx- prefix)
|
||||
assert entity_calls[0] == "@dafx-writeragent@writer-guid-001"
|
||||
assert entity_calls[1] == "@dafx-editoragent@editor-guid-002"
|
||||
|
||||
|
||||
class TestAgentThreadSerialization:
|
||||
"""Test that AgentThread can be serialized for orchestration state."""
|
||||
|
||||
async def test_agent_thread_serialize(self) -> None:
|
||||
"""Test that AgentThread can be serialized."""
|
||||
thread = AgentThread()
|
||||
|
||||
# Serialize
|
||||
serialized = await thread.serialize()
|
||||
|
||||
assert isinstance(serialized, dict)
|
||||
assert "service_thread_id" in serialized
|
||||
|
||||
async def test_agent_thread_deserialize(self) -> None:
|
||||
"""Test that AgentThread can be deserialized."""
|
||||
thread = AgentThread()
|
||||
serialized = await thread.serialize()
|
||||
|
||||
# Deserialize
|
||||
restored = await AgentThread.deserialize(serialized)
|
||||
|
||||
assert isinstance(restored, AgentThread)
|
||||
assert restored.service_thread_id == thread.service_thread_id
|
||||
|
||||
async def test_durable_agent_thread_serialization(self) -> None:
|
||||
"""Test that DurableAgentThread persists session metadata during serialization."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-999"
|
||||
mock_context.new_uuid = Mock(return_value="test-guid-999")
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
# Verify custom attribute and property exist
|
||||
assert thread.session_id is not None
|
||||
session_id = thread.session_id
|
||||
assert isinstance(session_id, AgentSessionId)
|
||||
assert session_id.name == "TestAgent"
|
||||
assert session_id.key == "test-guid-999"
|
||||
|
||||
# Standard serialization should still work
|
||||
serialized = await thread.serialize()
|
||||
assert isinstance(serialized, dict)
|
||||
assert serialized.get("durable_session_id") == str(session_id)
|
||||
|
||||
# After deserialization, we'd need to restore the custom attribute
|
||||
# This would be handled by the orchestration framework
|
||||
restored = await DurableAgentThread.deserialize(serialized)
|
||||
assert isinstance(restored, DurableAgentThread)
|
||||
assert restored.session_id == session_id
|
||||
# EntityId format is @dafx-agentname@uuid_hex (lowercased agent name with dafx- prefix)
|
||||
expected_writer_id = f"@dafx-writeragent@{uuid_hexes[0]}"
|
||||
expected_editor_id = f"@dafx-editoragent@{uuid_hexes[1]}"
|
||||
assert entity_calls[0] == expected_writer_id
|
||||
assert entity_calls[1] == expected_editor_id
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
Reference in New Issue
Block a user