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
@@ -1,248 +0,0 @@
|
||||
# Design: Durable Task Provider for Agent Framework
|
||||
|
||||
## Overview
|
||||
|
||||
This package, `agent-framework-durabletask`, provides a durability layer for the Microsoft Agent Framework using the `durabletask` Python SDK. It enables stateful, reliable, and distributed agent execution on any platform (Bring Your Own Platform), decoupling the agent's durability from the Azure Functions platform.
|
||||
|
||||
## Design Decision
|
||||
|
||||
**Selected Approach: Object-Oriented Wrappers with Symmetric Factory Pattern**
|
||||
|
||||
We will use a symmetric Object-Oriented design where both the Client (external) and Orchestrator (internal) expose a consistent interface for retrieving and interacting with durable agents.
|
||||
|
||||
## Core Philosophy
|
||||
|
||||
* **Native `DurableEntity` Support**: We will leverage the `DurableEntity` support introduced in `durabletask` v1.0.0.
|
||||
* **Symmetric Factories**: `DurableAIAgentClient` (for external use) and `DurableAIAgentOrchestrator` (for internal use) both provide a `get_agent` method.
|
||||
* **Unified Interface**: `DurableAIAgent` serves as the common interface for executing agents, regardless of the context (Client vs Orchestration).
|
||||
* **Consistent Return Type**: `DurableAIAgent.run` always returns a `Task` (or compatible awaitable), ensuring consistent usage patterns.
|
||||
|
||||
## Architecture
|
||||
|
||||
### 1. Package Structure
|
||||
|
||||
```text
|
||||
packages/durabletask/
|
||||
├── pyproject.toml
|
||||
├── README.md
|
||||
├── agent_framework_durabletask/
|
||||
│ ├── __init__.py
|
||||
│ ├── _worker.py # DurableAIAgentWorker
|
||||
│ ├── _client.py # DurableAIAgentClient
|
||||
│ ├── _orchestrator.py # DurableAIAgentOrchestrator
|
||||
│ ├── _entities.py # AgentEntity implementation
|
||||
│ ├── _models.py # Data models (RunRequest, AgentResponse, etc.)
|
||||
│ ├── _durable_agent_state.py # State schema (Ported from azurefunctions)
|
||||
│ ├── _shim.py # DurableAIAgent implementation (will be ported from azurefunctions)
|
||||
│ └── _utils.py # Mixins and helpers
|
||||
└── tests/
|
||||
```
|
||||
|
||||
### 2. State Management (`_durable_agent_state.py`)
|
||||
|
||||
**Important**: This will be the state maintained in the durable entities for both `durabletask` and `azurefunctions` package.
|
||||
|
||||
### 3. The Agent Entity (`_entities.py`)
|
||||
|
||||
We will implement a class `AgentEntity` that inherits from `durabletask.entities.DurableEntity`.
|
||||
|
||||
**Important**: This will be ported from `azurefunctions` package too but with slight modifications, details TBD.
|
||||
|
||||
### 4. The Worker Wrapper (`_worker.py`)
|
||||
|
||||
The `DurableAIAgentWorker` wraps an existing `durabletask` worker instance.
|
||||
|
||||
```python
|
||||
class DurableAIAgentWorker:
|
||||
def __init__(self, worker: TaskHubGrpcWorker):
|
||||
self._worker = worker
|
||||
self._registered_agents: dict[str, AgentProtocol] = {}
|
||||
|
||||
def add_agent(self, agent: AgentProtocol) -> None:
|
||||
"""Registers an agent with the worker.
|
||||
|
||||
Uses the factory pattern to create an AgentEntity class with the agent
|
||||
instance injected, then registers it with the durabletask worker.
|
||||
"""
|
||||
# Store the agent reference
|
||||
self._registered_agents[agent.name] = agent
|
||||
|
||||
# Create a configured entity class using the factory
|
||||
entity_class = create_agent_entity(agent)
|
||||
|
||||
# Register the entity class with the worker
|
||||
# The worker.add_entity method takes a class or function
|
||||
self._worker.add_entity(entity_class)
|
||||
|
||||
def start(self):
|
||||
"""Start the worker to begin processing tasks."""
|
||||
self._worker.start()
|
||||
|
||||
def stop(self):
|
||||
"""Stop the worker gracefully."""
|
||||
self._worker.stop()
|
||||
```
|
||||
|
||||
### 5. The Mixin (`_utils.py`)
|
||||
|
||||
```python
|
||||
class GetDurableAgentMixin:
|
||||
"""Mixin to provide get_agent interface."""
|
||||
|
||||
def get_agent(self, agent_name: str) -> 'DurableAIAgent':
|
||||
raise NotImplementedError
|
||||
```
|
||||
|
||||
### 6. The Client Wrapper (`_client.py`)
|
||||
|
||||
The `DurableAIAgentClient` is for external clients (e.g., FastAPI, CLI).
|
||||
|
||||
```python
|
||||
class DurableAIAgentClient(GetDurableAgentMixin):
|
||||
def __init__(self, client: TaskHubGrpcClient):
|
||||
self._client = client
|
||||
|
||||
async def get_agent(self, agent_name: str) -> 'DurableAIAgent':
|
||||
"""Retrieves a DurableAIAgent shim.
|
||||
|
||||
Validates existence by attempting to fetch entity state/metadata.
|
||||
"""
|
||||
# Validation logic using self._client.get_entity(...)
|
||||
# ...
|
||||
return DurableAIAgent(self, agent_name)
|
||||
|
||||
def run_agent(self, agent_name: str, message: str, **kwargs) -> 'Task':
|
||||
"""Runs agent via signal + poll and returns a Task wrapper."""
|
||||
# Returns a ClientTask (wrapper around asyncio.Task)
|
||||
pass
|
||||
```
|
||||
|
||||
### 7. The Orchestration Context Wrapper (`_orchestration_context.py`)
|
||||
|
||||
The `DurableAIAgentOrchestrationContext` is for use *inside* orchestrations to get access to agents that were registered in the workers.
|
||||
|
||||
```python
|
||||
class DurableAIAgentOrchestrationContext(GetDurableAgentMixin):
|
||||
def __init__(self, context: OrchestrationContext):
|
||||
self._context = context
|
||||
|
||||
def get_agent(self, agent_name: str) -> 'DurableAIAgent':
|
||||
"""Retrieves a DurableAIAgent shim.
|
||||
|
||||
Validation is deferred or performed via call_entity if needed.
|
||||
"""
|
||||
return DurableAIAgent(self, agent_name)
|
||||
|
||||
def run_agent(self, agent_name: str, message: str, **kwargs) -> 'Task':
|
||||
"""Runs agent via call_entity and returns the Task."""
|
||||
# Returns the native durabletask.Task
|
||||
pass
|
||||
```
|
||||
|
||||
### 8. The Durable Agent Shim (`_shim.py`)
|
||||
|
||||
The `DurableAIAgent` implements `AgentProtocol` but delegates execution to the provider. This will be ported from `azurefunctions` package and updated accordingly.
|
||||
|
||||
```python
|
||||
class DurableAIAgent(AgentProtocol):
|
||||
"""A shim that delegates execution to the provider (Client or Orchestrator)."""
|
||||
|
||||
def __init__(self, provider: GetDurableAgentMixin, name: str):
|
||||
self._provider = provider
|
||||
self._name = name
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self._name
|
||||
|
||||
def run(self, message: str, **kwargs) -> 'Task':
|
||||
"""Executes the agent.
|
||||
|
||||
Returns:
|
||||
Task: A yieldable/awaitable task object.
|
||||
"""
|
||||
return self._provider.run_agent(
|
||||
agent_name=self.name,
|
||||
message=message,
|
||||
**kwargs
|
||||
)
|
||||
```
|
||||
|
||||
## Usage Experience
|
||||
|
||||
**Scenario A: Worker Side**
|
||||
```python
|
||||
# 1. Define your agent
|
||||
# The agent can be any implementation of AgentProtocol.
|
||||
# For example, a standard Agent with a model and instructions.
|
||||
my_agent = Agent(
|
||||
name="my_agent",
|
||||
instructions="You are a helpful assistant.",
|
||||
model=openai_model
|
||||
)
|
||||
|
||||
# 2. Create the worker and the agent worker wrapper
|
||||
with DurableTaskSchedulerWorker(...) as worker:
|
||||
|
||||
agent_worker = DurableAIAgentWorker(worker)
|
||||
|
||||
# 3. Register the agent
|
||||
agent_worker.add_agent(my_agent)
|
||||
|
||||
# 4. Start the worker
|
||||
worker.start()
|
||||
|
||||
# ... keep running ...
|
||||
```
|
||||
|
||||
**Scenario B: Client Side**
|
||||
```python
|
||||
# 1. Configure the Durable Task client
|
||||
client = DurableTaskSchedulerClient(...)
|
||||
|
||||
# 2. Create the Agent Client wrapper
|
||||
agent_client = DurableAIAgentClient(client)
|
||||
|
||||
# 3. Get a reference to the agent
|
||||
agent = await agent_client.get_agent("my_agent")
|
||||
|
||||
# 4. Run the agent
|
||||
# The returned object is designed to be compatible with both `await` (Client)
|
||||
# and `yield` (Orchestrator). Implementation details on this unified return type will follow.
|
||||
response = await agent.run("Hello")
|
||||
```
|
||||
|
||||
**Scenario C: Orchestration Side**
|
||||
```python
|
||||
def orchestrator(context: OrchestrationContext):
|
||||
# 1. Create the Agent Orchestrator wrapper
|
||||
agent_orch = DurableAIAgentOrchestrator(context)
|
||||
|
||||
# 2. Get a reference to the agent
|
||||
agent = agent_orch.get_agent("my_agent")
|
||||
|
||||
# 3. Run the agent (returns a Task, so we yield it)
|
||||
result = yield agent.run("Hello")
|
||||
|
||||
return result
|
||||
```
|
||||
|
||||
## Additional Styles Considered
|
||||
|
||||
### Inheritance Pattern for worker and client (like `DurableAIAgentWorker`, `DurableAIAgentClient`, etc)
|
||||
|
||||
We investigated inheriting `DurableAIAgentWorker` directly from `TaskHubGrpcWorker` (or `DurableTaskSchedulerWorker`) to provide a unified API where the agent worker *is* a durable task worker (and similarly the client).
|
||||
|
||||
**Why we chose Composition over Inheritance:**
|
||||
|
||||
1. **Initialization Divergence:** The `durabletask` package has two distinct worker classes with incompatible `__init__` signatures:
|
||||
* `TaskHubGrpcWorker`: Requires `host_address`, `metadata`, etc.
|
||||
* `DurableTaskSchedulerWorker`: Requires `host_address`, `taskhub`, `token_credential`, etc.
|
||||
|
||||
To support both via inheritance, we would need to maintain two separate classes (e.g., `DurableAIAgentGrpcWorker` and `DurableAIAgentSchedulerWorker`) or use a complex Mixin approach. This increases the API surface area and maintenance burden.
|
||||
|
||||
2. **Encapsulation:** The logic for Azure Managed DTS (authentication, routing) is currently encapsulated in an internal interceptor class within `durabletask`. Without changes to the upstream package to expose this logic, we cannot create a single "Universal" worker class that inherits from the base worker but supports Azure features.
|
||||
|
||||
3. **Flexibility:** The Composition pattern allows `DurableAIAgentWorker` to accept *any* instance of a worker that satisfies the required interface. This makes it forward-compatible with future worker implementations or custom subclasses without requiring code changes in our package.
|
||||
|
||||
4. **Simplicity:** While Composition requires a two-step setup (instantiate worker, then wrap it), it keeps the `agent-framework-durabletask` package simple, focused, and loosely coupled from the implementation details of the underlying `durabletask` workers.
|
||||
@@ -3,6 +3,7 @@
|
||||
"""Durable Task integration for Microsoft Agent Framework."""
|
||||
|
||||
from ._callbacks import AgentCallbackContext, AgentResponseCallbackProtocol
|
||||
from ._client import DurableAIAgentClient
|
||||
from ._constants import (
|
||||
DEFAULT_MAX_POLL_RETRIES,
|
||||
DEFAULT_POLL_INTERVAL_SECONDS,
|
||||
@@ -41,7 +42,12 @@ from ._durable_agent_state import (
|
||||
DurableAgentStateUsageContent,
|
||||
)
|
||||
from ._entities import AgentEntity, AgentEntityStateProviderMixin
|
||||
from ._models import RunRequest, serialize_response_format
|
||||
from ._executors import DurableAgentExecutor
|
||||
from ._models import AgentSessionId, DurableAgentThread, RunRequest
|
||||
from ._orchestration_context import DurableAIAgentOrchestrationContext
|
||||
from ._response_utils import ensure_response_format, load_agent_response
|
||||
from ._shim import DurableAIAgent
|
||||
from ._worker import DurableAIAgentWorker
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_MAX_POLL_RETRIES",
|
||||
@@ -58,8 +64,14 @@ __all__ = [
|
||||
"AgentEntity",
|
||||
"AgentEntityStateProviderMixin",
|
||||
"AgentResponseCallbackProtocol",
|
||||
"AgentSessionId",
|
||||
"ApiResponseFields",
|
||||
"ContentTypes",
|
||||
"DurableAIAgent",
|
||||
"DurableAIAgentClient",
|
||||
"DurableAIAgentOrchestrationContext",
|
||||
"DurableAIAgentWorker",
|
||||
"DurableAgentExecutor",
|
||||
"DurableAgentState",
|
||||
"DurableAgentStateContent",
|
||||
"DurableAgentStateData",
|
||||
@@ -80,7 +92,10 @@ __all__ = [
|
||||
"DurableAgentStateUriContent",
|
||||
"DurableAgentStateUsage",
|
||||
"DurableAgentStateUsageContent",
|
||||
"DurableAgentThread",
|
||||
"DurableAgentThread",
|
||||
"DurableStateFields",
|
||||
"RunRequest",
|
||||
"serialize_response_format",
|
||||
"ensure_response_format",
|
||||
"load_agent_response",
|
||||
]
|
||||
|
||||
@@ -0,0 +1,90 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Client wrapper for Durable Task Agent Framework.
|
||||
|
||||
This module provides the DurableAIAgentClient class for external clients to interact
|
||||
with durable agents via gRPC.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from agent_framework import AgentRunResponse, get_logger
|
||||
from durabletask.client import TaskHubGrpcClient
|
||||
|
||||
from ._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
|
||||
from ._executors import ClientAgentExecutor
|
||||
from ._shim import DurableAgentProvider, DurableAIAgent
|
||||
|
||||
logger = get_logger("agent_framework.durabletask.client")
|
||||
|
||||
|
||||
class DurableAIAgentClient(DurableAgentProvider[AgentRunResponse]):
|
||||
"""Client wrapper for interacting with durable agents externally.
|
||||
|
||||
This class wraps a durabletask TaskHubGrpcClient and provides a convenient
|
||||
interface for retrieving and executing durable agents from external contexts.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from durabletask import TaskHubGrpcClient
|
||||
from agent_framework_durabletask import DurableAIAgentClient
|
||||
|
||||
# Create the underlying client
|
||||
client = TaskHubGrpcClient(host_address="localhost:4001")
|
||||
|
||||
# Wrap it with the agent client
|
||||
agent_client = DurableAIAgentClient(client)
|
||||
|
||||
# Get an agent reference
|
||||
agent = agent_client.get_agent("assistant")
|
||||
|
||||
# Run the agent (synchronous call that waits for completion)
|
||||
response = agent.run("Hello, how are you?")
|
||||
print(response.text)
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: TaskHubGrpcClient,
|
||||
max_poll_retries: int = DEFAULT_MAX_POLL_RETRIES,
|
||||
poll_interval_seconds: float = DEFAULT_POLL_INTERVAL_SECONDS,
|
||||
):
|
||||
"""Initialize the client wrapper.
|
||||
|
||||
Args:
|
||||
client: The durabletask client instance to wrap
|
||||
max_poll_retries: Maximum polling attempts when waiting for responses
|
||||
poll_interval_seconds: Delay in seconds between polling attempts
|
||||
"""
|
||||
self._client = client
|
||||
|
||||
# Validate and set polling parameters
|
||||
self.max_poll_retries = max(1, max_poll_retries)
|
||||
self.poll_interval_seconds = (
|
||||
poll_interval_seconds if poll_interval_seconds > 0 else DEFAULT_POLL_INTERVAL_SECONDS
|
||||
)
|
||||
|
||||
self._executor = ClientAgentExecutor(self._client, self.max_poll_retries, self.poll_interval_seconds)
|
||||
logger.debug("[DurableAIAgentClient] Initialized with client type: %s", type(client).__name__)
|
||||
|
||||
def get_agent(self, agent_name: str) -> DurableAIAgent[AgentRunResponse]:
|
||||
"""Retrieve a DurableAIAgent shim for the specified agent.
|
||||
|
||||
This method returns a proxy object that can be used to execute the agent.
|
||||
The actual agent must be registered on a worker with the same name.
|
||||
|
||||
Args:
|
||||
agent_name: Name of the agent to retrieve (without the dafx- prefix)
|
||||
|
||||
Returns:
|
||||
DurableAIAgent instance that can be used to run the agent
|
||||
|
||||
Note:
|
||||
This method does not validate that the agent exists. Validation
|
||||
will occur when the agent is executed. If the entity doesn't exist,
|
||||
the execution will fail with an appropriate error.
|
||||
"""
|
||||
logger.debug("[DurableAIAgentClient] Creating agent proxy for: %s", agent_name)
|
||||
|
||||
return DurableAIAgent(self._executor, agent_name)
|
||||
@@ -53,7 +53,7 @@ from agent_framework import (
|
||||
)
|
||||
from dateutil import parser as date_parser
|
||||
|
||||
from ._constants import ApiResponseFields, ContentTypes, DurableStateFields
|
||||
from ._constants import ContentTypes, DurableStateFields
|
||||
from ._models import RunRequest, serialize_response_format
|
||||
|
||||
logger = get_logger("agent_framework.durabletask.durable_agent_state")
|
||||
@@ -452,7 +452,7 @@ class DurableAgentState:
|
||||
"""Get the count of conversation entries (requests + responses)."""
|
||||
return len(self.data.conversation_history)
|
||||
|
||||
def try_get_agent_response(self, correlation_id: str) -> dict[str, Any] | None:
|
||||
def try_get_agent_response(self, correlation_id: str) -> AgentRunResponse | None:
|
||||
"""Try to get an agent response by correlation ID.
|
||||
|
||||
This method searches the conversation history for a response entry matching the given
|
||||
@@ -474,14 +474,8 @@ class DurableAgentState:
|
||||
for entry in self.data.conversation_history:
|
||||
if entry.correlation_id == correlation_id and isinstance(entry, DurableAgentStateResponse):
|
||||
# Found the entry, extract response data
|
||||
# Get the text content from assistant messages only
|
||||
content = "\n".join(message.text for message in entry.messages if message.text)
|
||||
return DurableAgentStateResponse.to_run_response(entry)
|
||||
|
||||
return {
|
||||
ApiResponseFields.CONTENT: content,
|
||||
ApiResponseFields.MESSAGE_COUNT: self.message_count,
|
||||
ApiResponseFields.CORRELATION_ID: correlation_id,
|
||||
}
|
||||
return None
|
||||
|
||||
|
||||
@@ -705,6 +699,21 @@ class DurableAgentStateResponse(DurableAgentStateEntry):
|
||||
usage=DurableAgentStateUsage.from_usage(response.usage_details),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def to_run_response(
|
||||
response_entry: DurableAgentStateResponse,
|
||||
) -> AgentRunResponse:
|
||||
"""Converts a DurableAgentStateResponse back to an AgentRunResponse."""
|
||||
messages = [m.to_chat_message() for m in response_entry.messages]
|
||||
|
||||
usage_details = response_entry.usage.to_usage_details() if response_entry.usage is not None else UsageDetails()
|
||||
|
||||
return AgentRunResponse(
|
||||
created_at=response_entry.created_at.isoformat(),
|
||||
messages=messages,
|
||||
usage_details=usage_details,
|
||||
)
|
||||
|
||||
|
||||
class DurableAgentStateMessage:
|
||||
"""Represents a message within a conversation history entry.
|
||||
@@ -1214,14 +1223,24 @@ class DurableAgentStateUsage:
|
||||
input_token_count=usage.input_token_count,
|
||||
output_token_count=usage.output_token_count,
|
||||
total_token_count=usage.total_token_count,
|
||||
extensionData=usage.additional_counts,
|
||||
)
|
||||
|
||||
def to_usage_details(self) -> UsageDetails:
|
||||
# Convert back to AI SDK UsageDetails
|
||||
extension_data: dict[str, int] = {}
|
||||
if self.extensionData is not None:
|
||||
for k, v in self.extensionData.items():
|
||||
try:
|
||||
extension_data[k] = int(v)
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
|
||||
return UsageDetails(
|
||||
input_token_count=self.input_token_count,
|
||||
output_token_count=self.output_token_count,
|
||||
total_token_count=self.total_token_count,
|
||||
**extension_data,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -128,7 +128,7 @@ class AgentEntity:
|
||||
) -> AgentRunResponse:
|
||||
"""Execute the agent with a message."""
|
||||
if isinstance(request, str):
|
||||
run_request = RunRequest(message=request, role=Role.USER)
|
||||
run_request = RunRequest.from_json(request)
|
||||
elif isinstance(request, dict):
|
||||
run_request = RunRequest.from_dict(request)
|
||||
else:
|
||||
@@ -139,8 +139,6 @@ class AgentEntity:
|
||||
correlation_id = run_request.correlation_id
|
||||
if not thread_id:
|
||||
raise ValueError("Entity State Provider must provide a thread_id")
|
||||
if not correlation_id:
|
||||
raise ValueError("RunRequest must include a correlation_id")
|
||||
response_format = run_request.response_format
|
||||
enable_tool_calls = run_request.enable_tool_calls
|
||||
|
||||
|
||||
@@ -0,0 +1,460 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Provider strategies for Durable Agent execution.
|
||||
|
||||
These classes are internal execution strategies used by the DurableAIAgent shim.
|
||||
They are intentionally separate from the public client/orchestration APIs to keep
|
||||
only `get_agent` exposed to consumers. Executors implement the execution contract
|
||||
and are injected into the shim.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import time
|
||||
import uuid
|
||||
from abc import ABC, abstractmethod
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Generic, TypeVar
|
||||
|
||||
from agent_framework import AgentRunResponse, AgentThread, ChatMessage, ErrorContent, Role, get_logger
|
||||
from durabletask.client import TaskHubGrpcClient
|
||||
from durabletask.entities import EntityInstanceId
|
||||
from durabletask.task import CompositeTask, OrchestrationContext, Task
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
|
||||
from ._durable_agent_state import DurableAgentState
|
||||
from ._models import AgentSessionId, DurableAgentThread, RunRequest
|
||||
from ._response_utils import ensure_response_format, load_agent_response
|
||||
|
||||
logger = get_logger("agent_framework.durabletask.executors")
|
||||
|
||||
# TypeVar for the task type returned by executors
|
||||
TaskT = TypeVar("TaskT")
|
||||
|
||||
|
||||
class DurableAgentTask(CompositeTask[AgentRunResponse]):
|
||||
"""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: Task[Any],
|
||||
response_format: type[BaseModel] | None,
|
||||
correlation_id: str,
|
||||
):
|
||||
"""Initialize the DurableAgentTask.
|
||||
|
||||
Args:
|
||||
entity_task: The underlying entity call task
|
||||
response_format: Optional Pydantic model for response parsing
|
||||
correlation_id: Correlation ID for logging
|
||||
"""
|
||||
self._response_format = response_format
|
||||
self._correlation_id = correlation_id
|
||||
super().__init__([entity_task]) # type: ignore[misc]
|
||||
|
||||
def on_child_completed(self, task: Task[Any]) -> None:
|
||||
"""Handle completion of the underlying entity task.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
task : Task
|
||||
The entity call task that just completed
|
||||
"""
|
||||
if self.is_complete:
|
||||
return
|
||||
|
||||
if task.is_failed:
|
||||
# Propagate the failure
|
||||
self._exception = task.get_exception()
|
||||
self._is_complete = True
|
||||
if self._parent is not None:
|
||||
self._parent.on_child_completed(self)
|
||||
return
|
||||
|
||||
# Task succeeded - transform the raw result
|
||||
raw_result = task.get_result()
|
||||
logger.debug(
|
||||
"[DurableAgentTask] Converting raw result for correlation_id %s",
|
||||
self._correlation_id,
|
||||
)
|
||||
|
||||
try:
|
||||
response = load_agent_response(raw_result)
|
||||
|
||||
if self._response_format is not None:
|
||||
ensure_response_format(
|
||||
self._response_format,
|
||||
self._correlation_id,
|
||||
response,
|
||||
)
|
||||
|
||||
# Set the typed AgentRunResponse as this task's result
|
||||
self._result = response
|
||||
self._is_complete = True
|
||||
|
||||
if self._parent is not None:
|
||||
self._parent.on_child_completed(self)
|
||||
|
||||
except Exception:
|
||||
logger.exception(
|
||||
"[DurableAgentTask] Failed to convert result for correlation_id: %s",
|
||||
self._correlation_id,
|
||||
)
|
||||
raise
|
||||
|
||||
|
||||
class DurableAgentExecutor(ABC, Generic[TaskT]):
|
||||
"""Abstract base class for durable agent execution strategies.
|
||||
|
||||
Type Parameters:
|
||||
TaskT: The task type returned by this executor
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def run_durable_agent(
|
||||
self,
|
||||
agent_name: str,
|
||||
run_request: RunRequest,
|
||||
thread: AgentThread | None = None,
|
||||
) -> TaskT:
|
||||
"""Execute the durable agent.
|
||||
|
||||
Returns:
|
||||
TaskT: The task type specific to this executor implementation
|
||||
"""
|
||||
raise NotImplementedError
|
||||
|
||||
def get_new_thread(self, agent_name: str, **kwargs: Any) -> DurableAgentThread:
|
||||
"""Create a new DurableAgentThread with random session ID."""
|
||||
session_id = self._create_session_id(agent_name)
|
||||
return DurableAgentThread.from_session_id(session_id, **kwargs)
|
||||
|
||||
def _create_session_id(
|
||||
self,
|
||||
agent_name: str,
|
||||
thread: AgentThread | None = None,
|
||||
) -> AgentSessionId:
|
||||
"""Create the AgentSessionId for the execution."""
|
||||
if isinstance(thread, DurableAgentThread) and thread.session_id is not None:
|
||||
return thread.session_id
|
||||
# Create new session ID - either no thread provided or it's a regular AgentThread
|
||||
key = self.generate_unique_id()
|
||||
return AgentSessionId(name=agent_name, key=key)
|
||||
|
||||
def generate_unique_id(self) -> str:
|
||||
"""Generate a new Unique ID."""
|
||||
return uuid.uuid4().hex
|
||||
|
||||
def get_run_request(
|
||||
self,
|
||||
message: str,
|
||||
response_format: type[BaseModel] | None,
|
||||
enable_tool_calls: bool,
|
||||
) -> RunRequest:
|
||||
"""Create a RunRequest for the given parameters."""
|
||||
correlation_id = self.generate_unique_id()
|
||||
return RunRequest(
|
||||
message=message,
|
||||
response_format=response_format,
|
||||
enable_tool_calls=enable_tool_calls,
|
||||
correlation_id=correlation_id,
|
||||
)
|
||||
|
||||
|
||||
class ClientAgentExecutor(DurableAgentExecutor[AgentRunResponse]):
|
||||
"""Execution strategy for external clients.
|
||||
|
||||
Note: Returns AgentRunResponse directly since the execution
|
||||
is blocking until response is available via polling
|
||||
as per the design of TaskHubGrpcClient.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
client: TaskHubGrpcClient,
|
||||
max_poll_retries: int = DEFAULT_MAX_POLL_RETRIES,
|
||||
poll_interval_seconds: float = DEFAULT_POLL_INTERVAL_SECONDS,
|
||||
):
|
||||
self._client = client
|
||||
self.max_poll_retries = max_poll_retries
|
||||
self.poll_interval_seconds = poll_interval_seconds
|
||||
|
||||
def run_durable_agent(
|
||||
self,
|
||||
agent_name: str,
|
||||
run_request: RunRequest,
|
||||
thread: AgentThread | None = None,
|
||||
) -> AgentRunResponse:
|
||||
"""Execute the agent via the durabletask client.
|
||||
|
||||
Signals the agent entity with a message request, then polls the entity
|
||||
state to retrieve the response once processing is complete.
|
||||
|
||||
Note: This is a blocking/synchronous operation (in line with how
|
||||
TaskHubGrpcClient works) that polls until a response is available or
|
||||
timeout occurs.
|
||||
|
||||
Args:
|
||||
agent_name: Name of the agent to execute
|
||||
run_request: The run request containing message and optional response format
|
||||
thread: Optional conversation thread (creates new if not provided)
|
||||
|
||||
Returns:
|
||||
AgentRunResponse: The agent's response after execution completes
|
||||
"""
|
||||
# Signal the entity with the request
|
||||
entity_id = self._signal_agent_entity(agent_name, run_request, thread)
|
||||
|
||||
# Poll for the response
|
||||
agent_response = self._poll_for_agent_response(entity_id, run_request.correlation_id)
|
||||
|
||||
# Handle and return the result
|
||||
return self._handle_agent_response(agent_response, run_request.response_format, run_request.correlation_id)
|
||||
|
||||
def _signal_agent_entity(
|
||||
self,
|
||||
agent_name: str,
|
||||
run_request: RunRequest,
|
||||
thread: AgentThread | None,
|
||||
) -> EntityInstanceId:
|
||||
"""Signal the agent entity with a run request.
|
||||
|
||||
Args:
|
||||
agent_name: Name of the agent to execute
|
||||
run_request: The run request containing message and optional response format
|
||||
thread: Optional conversation thread
|
||||
|
||||
Returns:
|
||||
entity_id
|
||||
"""
|
||||
# Get or create session ID
|
||||
session_id = self._create_session_id(agent_name, thread)
|
||||
|
||||
# Create the entity ID
|
||||
entity_id = EntityInstanceId(
|
||||
entity=session_id.entity_name,
|
||||
key=session_id.key,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"[ClientAgentExecutor] Signaling entity '%s' (session: %s, correlation: %s)",
|
||||
agent_name,
|
||||
session_id,
|
||||
run_request.correlation_id,
|
||||
)
|
||||
|
||||
self._client.signal_entity(entity_id, "run", run_request.to_dict())
|
||||
return entity_id
|
||||
|
||||
def _poll_for_agent_response(
|
||||
self,
|
||||
entity_id: EntityInstanceId,
|
||||
correlation_id: str,
|
||||
) -> AgentRunResponse | None:
|
||||
"""Poll the entity for a response with retries.
|
||||
|
||||
Args:
|
||||
entity_id: Entity instance identifier
|
||||
correlation_id: Correlation ID to track the request
|
||||
|
||||
Returns:
|
||||
The agent response if found, None if timeout occurs
|
||||
"""
|
||||
agent_response = None
|
||||
|
||||
for attempt in range(1, self.max_poll_retries + 1):
|
||||
time.sleep(self.poll_interval_seconds)
|
||||
|
||||
agent_response = self._poll_entity_for_response(entity_id, correlation_id)
|
||||
if agent_response is not None:
|
||||
logger.info(
|
||||
"[ClientAgentExecutor] Found response (attempt %d/%d, correlation: %s)",
|
||||
attempt,
|
||||
self.max_poll_retries,
|
||||
correlation_id,
|
||||
)
|
||||
break
|
||||
|
||||
logger.debug(
|
||||
"[ClientAgentExecutor] Response not ready (attempt %d/%d)",
|
||||
attempt,
|
||||
self.max_poll_retries,
|
||||
)
|
||||
|
||||
return agent_response
|
||||
|
||||
def _handle_agent_response(
|
||||
self,
|
||||
agent_response: AgentRunResponse | None,
|
||||
response_format: type[BaseModel] | None,
|
||||
correlation_id: str,
|
||||
) -> AgentRunResponse:
|
||||
"""Handle the agent response or create an error response.
|
||||
|
||||
Args:
|
||||
agent_response: The response from polling, or None if timeout
|
||||
response_format: Optional response format for validation
|
||||
correlation_id: Correlation ID for logging
|
||||
|
||||
Returns:
|
||||
AgentRunResponse with either the agent's response or an error message
|
||||
"""
|
||||
if agent_response is not None:
|
||||
try:
|
||||
# Validate response format if specified
|
||||
if response_format is not None:
|
||||
ensure_response_format(
|
||||
response_format,
|
||||
correlation_id,
|
||||
agent_response,
|
||||
)
|
||||
|
||||
return agent_response
|
||||
|
||||
except Exception as e:
|
||||
logger.exception(
|
||||
"[ClientAgentExecutor] Error converting response for correlation: %s",
|
||||
correlation_id,
|
||||
)
|
||||
error_message = ChatMessage(
|
||||
role=Role.SYSTEM,
|
||||
contents=[
|
||||
ErrorContent(
|
||||
message=f"Error processing agent response: {e}",
|
||||
error_code="response_processing_error",
|
||||
)
|
||||
],
|
||||
)
|
||||
else:
|
||||
logger.warning(
|
||||
"[ClientAgentExecutor] Timeout after %d attempts (correlation: %s)",
|
||||
self.max_poll_retries,
|
||||
correlation_id,
|
||||
)
|
||||
error_message = ChatMessage(
|
||||
role=Role.SYSTEM,
|
||||
contents=[
|
||||
ErrorContent(
|
||||
message=f"Timeout waiting for agent response after {self.max_poll_retries} attempts",
|
||||
error_code="response_timeout",
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
return AgentRunResponse(
|
||||
messages=[error_message],
|
||||
created_at=datetime.now(timezone.utc).isoformat(),
|
||||
)
|
||||
|
||||
def _poll_entity_for_response(
|
||||
self,
|
||||
entity_id: EntityInstanceId,
|
||||
correlation_id: str,
|
||||
) -> AgentRunResponse | None:
|
||||
"""Poll the entity state for a response matching the correlation ID.
|
||||
|
||||
Args:
|
||||
entity_id: Entity instance identifier
|
||||
correlation_id: Correlation ID to search for
|
||||
|
||||
Returns:
|
||||
Response AgentRunResponse, None otherwise
|
||||
"""
|
||||
try:
|
||||
entity_metadata = self._client.get_entity(entity_id, include_state=True)
|
||||
|
||||
if entity_metadata is None:
|
||||
return None
|
||||
|
||||
state_json = entity_metadata.get_state()
|
||||
if not state_json:
|
||||
return None
|
||||
|
||||
state = DurableAgentState.from_json(state_json)
|
||||
|
||||
# Use the helper method to get response by correlation ID
|
||||
return state.try_get_agent_response(correlation_id)
|
||||
|
||||
except Exception as e:
|
||||
logger.warning(
|
||||
"[ClientAgentExecutor] Error reading entity state: %s",
|
||||
e,
|
||||
)
|
||||
return None
|
||||
|
||||
|
||||
class OrchestrationAgentExecutor(DurableAgentExecutor[DurableAgentTask]):
|
||||
"""Execution strategy for orchestrations (sync/yield)."""
|
||||
|
||||
def __init__(self, context: OrchestrationContext):
|
||||
self._context = context
|
||||
logger.debug("[OrchestrationAgentExecutor] Initialized")
|
||||
|
||||
def get_run_request(
|
||||
self,
|
||||
message: str,
|
||||
response_format: type[BaseModel] | None,
|
||||
enable_tool_calls: bool,
|
||||
) -> RunRequest:
|
||||
"""Get the current run request from the orchestration context.
|
||||
|
||||
Returns:
|
||||
RunRequest: The current run request
|
||||
"""
|
||||
request = super().get_run_request(
|
||||
message,
|
||||
response_format,
|
||||
enable_tool_calls,
|
||||
)
|
||||
request.orchestration_id = self._context.instance_id
|
||||
return request
|
||||
|
||||
def run_durable_agent(
|
||||
self,
|
||||
agent_name: str,
|
||||
run_request: RunRequest,
|
||||
thread: AgentThread | None = None,
|
||||
) -> DurableAgentTask:
|
||||
"""Execute the agent via orchestration context.
|
||||
|
||||
Calls the agent entity and returns a DurableAgentTask that can be yielded
|
||||
in orchestrations to wait for the entity's response.
|
||||
|
||||
Args:
|
||||
agent_name: Name of the agent to execute
|
||||
run_request: The run request containing message and optional response format
|
||||
thread: Optional conversation thread (creates new if not provided)
|
||||
|
||||
Returns:
|
||||
DurableAgentTask: A task wrapping the entity call that yields AgentRunResponse
|
||||
"""
|
||||
# Resolve session
|
||||
session_id = self._create_session_id(agent_name, thread)
|
||||
|
||||
# Create the entity ID
|
||||
entity_id = EntityInstanceId(
|
||||
entity=session_id.entity_name,
|
||||
key=session_id.key,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"[OrchestrationAgentExecutor] correlation_id: %s entity_id: %s session_id: %s",
|
||||
run_request.correlation_id,
|
||||
entity_id,
|
||||
session_id,
|
||||
)
|
||||
|
||||
# Call the entity and get the underlying task
|
||||
entity_task: Task[Any] = self._context.call_entity(entity_id, "run", run_request.to_dict()) # type: ignore
|
||||
|
||||
# Wrap in DurableAgentTask for response transformation
|
||||
return DurableAgentTask(
|
||||
entity_task=entity_task,
|
||||
response_format=run_request.response_format,
|
||||
correlation_id=run_request.correlation_id,
|
||||
)
|
||||
@@ -8,12 +8,15 @@ This module defines the request and response models used by the framework.
|
||||
from __future__ import annotations
|
||||
|
||||
import inspect
|
||||
import json
|
||||
import uuid
|
||||
from collections.abc import MutableMapping
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timezone
|
||||
from importlib import import_module
|
||||
from typing import TYPE_CHECKING, Any, cast
|
||||
|
||||
from agent_framework import Role
|
||||
from agent_framework import AgentThread, Role
|
||||
|
||||
from ._constants import REQUEST_RESPONSE_FORMAT_TEXT
|
||||
|
||||
@@ -101,38 +104,38 @@ class RunRequest:
|
||||
role: The role of the message sender (user, system, or assistant)
|
||||
response_format: Optional Pydantic BaseModel type describing the structured response format
|
||||
enable_tool_calls: Whether to enable tool calls for this request
|
||||
correlation_id: Optional correlation ID for tracking the response to this specific request
|
||||
correlation_id: Correlation ID for tracking the response to this specific request
|
||||
created_at: Optional timestamp when the request was created
|
||||
orchestration_id: Optional ID of the orchestration that initiated this request
|
||||
"""
|
||||
|
||||
message: str
|
||||
request_response_format: str
|
||||
correlation_id: str
|
||||
role: Role = Role.USER
|
||||
response_format: type[BaseModel] | None = None
|
||||
enable_tool_calls: bool = True
|
||||
correlation_id: str | None = None
|
||||
created_at: datetime | None = None
|
||||
orchestration_id: str | None = None
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
message: str,
|
||||
correlation_id: str,
|
||||
request_response_format: str = REQUEST_RESPONSE_FORMAT_TEXT,
|
||||
role: Role | str | None = Role.USER,
|
||||
response_format: type[BaseModel] | None = None,
|
||||
enable_tool_calls: bool = True,
|
||||
correlation_id: str | None = None,
|
||||
created_at: datetime | None = None,
|
||||
orchestration_id: str | None = None,
|
||||
) -> None:
|
||||
self.message = message
|
||||
self.correlation_id = correlation_id
|
||||
self.role = self.coerce_role(role)
|
||||
self.response_format = response_format
|
||||
self.request_response_format = request_response_format
|
||||
self.enable_tool_calls = enable_tool_calls
|
||||
self.correlation_id = correlation_id
|
||||
self.created_at = created_at
|
||||
self.created_at = created_at if created_at is not None else datetime.now(tz=timezone.utc)
|
||||
self.orchestration_id = orchestration_id
|
||||
|
||||
@staticmethod
|
||||
@@ -154,11 +157,10 @@ class RunRequest:
|
||||
"enable_tool_calls": self.enable_tool_calls,
|
||||
"role": self.role.value,
|
||||
"request_response_format": self.request_response_format,
|
||||
"correlationId": self.correlation_id,
|
||||
}
|
||||
if self.response_format:
|
||||
result["response_format"] = serialize_response_format(self.response_format)
|
||||
if self.correlation_id:
|
||||
result["correlationId"] = self.correlation_id
|
||||
if self.created_at:
|
||||
result["created_at"] = self.created_at.isoformat()
|
||||
if self.orchestration_id:
|
||||
@@ -166,6 +168,16 @@ class RunRequest:
|
||||
|
||||
return result
|
||||
|
||||
@classmethod
|
||||
def from_json(cls, data: str) -> RunRequest:
|
||||
"""Create RunRequest from JSON string."""
|
||||
try:
|
||||
dict_data = json.loads(data)
|
||||
except json.JSONDecodeError as e:
|
||||
raise ValueError("The durable agent state is not valid JSON.") from e
|
||||
|
||||
return cls.from_dict(dict_data)
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data: dict[str, Any]) -> RunRequest:
|
||||
"""Create RunRequest from dictionary."""
|
||||
@@ -176,13 +188,120 @@ class RunRequest:
|
||||
except ValueError:
|
||||
created_at = None
|
||||
|
||||
correlation_id = data.get("correlationId")
|
||||
if not correlation_id:
|
||||
raise ValueError("correlationId is required in RunRequest data")
|
||||
|
||||
return cls(
|
||||
message=data.get("message", ""),
|
||||
correlation_id=correlation_id,
|
||||
request_response_format=data.get("request_response_format", REQUEST_RESPONSE_FORMAT_TEXT),
|
||||
role=cls.coerce_role(data.get("role")),
|
||||
response_format=_deserialize_response_format(data.get("response_format")),
|
||||
enable_tool_calls=data.get("enable_tool_calls", True),
|
||||
correlation_id=data.get("correlationId"),
|
||||
created_at=created_at,
|
||||
orchestration_id=data.get("orchestrationId"),
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class AgentSessionId:
|
||||
"""Represents an agent session identifier (name + key)."""
|
||||
|
||||
name: str
|
||||
key: str
|
||||
|
||||
ENTITY_NAME_PREFIX: str = "dafx-"
|
||||
|
||||
@staticmethod
|
||||
def to_entity_name(name: str) -> str:
|
||||
return f"{AgentSessionId.ENTITY_NAME_PREFIX}{name}"
|
||||
|
||||
@staticmethod
|
||||
def with_random_key(name: str) -> AgentSessionId:
|
||||
return AgentSessionId(name=name, key=uuid.uuid4().hex)
|
||||
|
||||
@property
|
||||
def entity_name(self) -> str:
|
||||
return self.to_entity_name(self.name)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"@{self.name}@{self.key}"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"AgentSessionId(name='{self.name}', key='{self.key}')"
|
||||
|
||||
@staticmethod
|
||||
def parse(session_id_string: str) -> AgentSessionId:
|
||||
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,
|
||||
**kwargs: Any,
|
||||
) -> None:
|
||||
super().__init__(**kwargs)
|
||||
self._session_id: AgentSessionId | None = session_id
|
||||
|
||||
@property
|
||||
def session_id(self) -> AgentSessionId | None:
|
||||
return self._session_id
|
||||
|
||||
@session_id.setter
|
||||
def session_id(self, value: AgentSessionId | None) -> None:
|
||||
self._session_id = value
|
||||
|
||||
@classmethod
|
||||
def from_session_id(
|
||||
cls,
|
||||
session_id: AgentSessionId,
|
||||
**kwargs: Any,
|
||||
) -> DurableAgentThread:
|
||||
return cls(session_id=session_id, **kwargs)
|
||||
|
||||
async def serialize(self, **kwargs: Any) -> dict[str, Any]:
|
||||
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:
|
||||
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.session_id = AgentSessionId.parse(session_id_value)
|
||||
return thread
|
||||
|
||||
@@ -0,0 +1,75 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Orchestration context wrapper for Durable Task Agent Framework.
|
||||
|
||||
This module provides the DurableAIAgentOrchestrationContext class for use inside
|
||||
orchestration functions to interact with durable agents.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from agent_framework import get_logger
|
||||
from durabletask.task import OrchestrationContext
|
||||
|
||||
from ._executors import DurableAgentTask, OrchestrationAgentExecutor
|
||||
from ._shim import DurableAgentProvider, DurableAIAgent
|
||||
|
||||
logger = get_logger("agent_framework.durabletask.orchestration_context")
|
||||
|
||||
|
||||
class DurableAIAgentOrchestrationContext(DurableAgentProvider[DurableAgentTask]):
|
||||
"""Orchestration context wrapper for interacting with durable agents internally.
|
||||
|
||||
This class wraps a durabletask OrchestrationContext and provides a convenient
|
||||
interface for retrieving and executing durable agents from within orchestration
|
||||
functions.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from durabletask import Orchestration
|
||||
from agent_framework_durabletask import DurableAIAgentOrchestrationContext
|
||||
|
||||
|
||||
def my_orchestration(context: OrchestrationContext):
|
||||
# Wrap the context
|
||||
agent_context = DurableAIAgentOrchestrationContext(context)
|
||||
|
||||
# Get an agent reference
|
||||
agent = agent_context.get_agent("assistant")
|
||||
|
||||
# Run the agent (returns a Task to be yielded)
|
||||
result = yield agent.run("Hello, how are you?")
|
||||
|
||||
return result.text
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(self, context: OrchestrationContext):
|
||||
"""Initialize the orchestration context wrapper.
|
||||
|
||||
Args:
|
||||
context: The durabletask orchestration context to wrap
|
||||
"""
|
||||
self._context = context
|
||||
self._executor = OrchestrationAgentExecutor(self._context)
|
||||
logger.debug("[DurableAIAgentOrchestrationContext] Initialized")
|
||||
|
||||
def get_agent(self, agent_name: str) -> DurableAIAgent[DurableAgentTask]:
|
||||
"""Retrieve a DurableAIAgent shim for the specified agent.
|
||||
|
||||
This method returns a proxy object that can be used to execute the agent
|
||||
within an orchestration. The agent's run() method will return a Task that
|
||||
must be yielded.
|
||||
|
||||
Args:
|
||||
agent_name: Name of the agent to retrieve (without the dafx- prefix)
|
||||
|
||||
Returns:
|
||||
DurableAIAgent instance that can be used to run the agent
|
||||
|
||||
Note:
|
||||
Validation is deferred to execution time. The entity must be registered
|
||||
on a worker with the name f"dafx-{agent_name}".
|
||||
"""
|
||||
logger.debug("[DurableAIAgentOrchestrationContext] Creating agent proxy for: %s", agent_name)
|
||||
return DurableAIAgent(self._executor, agent_name)
|
||||
@@ -0,0 +1,62 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Shared utilities for handling AgentRunResponse parsing and validation."""
|
||||
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import AgentRunResponse, get_logger
|
||||
from pydantic import BaseModel
|
||||
|
||||
logger = get_logger("agent_framework.durabletask.response_utils")
|
||||
|
||||
|
||||
def load_agent_response(agent_response: AgentRunResponse | dict[str, Any] | None) -> AgentRunResponse:
|
||||
"""Convert raw payloads into AgentRunResponse instance.
|
||||
|
||||
Args:
|
||||
agent_response: The response to convert, can be an AgentRunResponse, dict, or None
|
||||
|
||||
Returns:
|
||||
AgentRunResponse: The converted response object
|
||||
|
||||
Raises:
|
||||
ValueError: If agent_response is None
|
||||
TypeError: If agent_response is an unsupported type
|
||||
"""
|
||||
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(
|
||||
response_format: type[BaseModel] | None,
|
||||
correlation_id: str,
|
||||
response: AgentRunResponse,
|
||||
) -> None:
|
||||
"""Ensure the AgentRunResponse value is parsed into the expected response_format.
|
||||
|
||||
This function modifies the response in-place by parsing its value attribute
|
||||
into the specified Pydantic model format.
|
||||
|
||||
Args:
|
||||
response_format: Optional Pydantic model class to parse the response value into
|
||||
correlation_id: Correlation ID for logging purposes
|
||||
response: The AgentRunResponse object to validate and parse
|
||||
"""
|
||||
if response_format is not None and not isinstance(response.value, response_format):
|
||||
response.try_parse_value(response_format)
|
||||
|
||||
logger.debug(
|
||||
"[ensure_response_format] Loaded AgentRunResponse.value for correlation_id %s with type: %s",
|
||||
correlation_id,
|
||||
type(response.value).__name__,
|
||||
)
|
||||
@@ -0,0 +1,185 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Durable Agent Shim for Durable Task Framework.
|
||||
|
||||
This module provides the DurableAIAgent shim that implements AgentProtocol
|
||||
and provides a consistent interface for both Client and Orchestration contexts.
|
||||
The actual execution is delegated to the context-specific providers.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from collections.abc import AsyncIterator
|
||||
from typing import Any, Generic, TypeVar
|
||||
|
||||
from agent_framework import AgentProtocol, AgentRunResponseUpdate, AgentThread, ChatMessage
|
||||
from pydantic import BaseModel
|
||||
|
||||
from ._executors import DurableAgentExecutor
|
||||
from ._models import DurableAgentThread
|
||||
|
||||
# TypeVar for the task type returned by executors
|
||||
# Covariant because TaskT only appears in return positions (output)
|
||||
TaskT = TypeVar("TaskT", covariant=True)
|
||||
|
||||
|
||||
class DurableAgentProvider(ABC, Generic[TaskT]):
|
||||
"""Abstract provider for constructing durable agent proxies.
|
||||
|
||||
Implemented by context-specific wrappers (client/orchestration) to return a
|
||||
`DurableAIAgent` shim backed by their respective `DurableAgentExecutor`
|
||||
implementation, ensuring a consistent `get_agent` entry point regardless of
|
||||
execution context.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def get_agent(self, agent_name: str) -> DurableAIAgent[TaskT]:
|
||||
"""Retrieve a DurableAIAgent shim for the specified agent.
|
||||
|
||||
Args:
|
||||
agent_name: Name of the agent to retrieve
|
||||
|
||||
Returns:
|
||||
DurableAIAgent instance that can be used to run the agent
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Must be implemented by subclasses
|
||||
"""
|
||||
raise NotImplementedError("Subclasses must implement get_agent()")
|
||||
|
||||
|
||||
class DurableAIAgent(AgentProtocol, Generic[TaskT]):
|
||||
"""A durable agent proxy that delegates execution to the provider.
|
||||
|
||||
This class implements AgentProtocol but with one critical difference:
|
||||
- AgentProtocol.run() returns a Coroutine (async, must await)
|
||||
- DurableAIAgent.run() returns TaskT (sync Task object - must yield
|
||||
or the AgentRunResponse directly in the case of TaskHubGrpcClient)
|
||||
|
||||
This represents fundamentally different execution models but maintains the same
|
||||
interface contract for all other properties and methods.
|
||||
|
||||
The underlying provider determines how execution occurs (entity calls, HTTP requests, etc.)
|
||||
and what type of Task object is returned.
|
||||
|
||||
Type Parameters:
|
||||
TaskT: The task type returned by this agent (e.g., AgentRunResponse, DurableAgentTask, AgentTask)
|
||||
"""
|
||||
|
||||
def __init__(self, executor: DurableAgentExecutor[TaskT], name: str, *, agent_id: str | None = None):
|
||||
"""Initialize the shim with a provider and agent name.
|
||||
|
||||
Args:
|
||||
executor: The execution provider (Client or OrchestrationContext)
|
||||
name: The name of the agent to execute
|
||||
agent_id: Optional unique identifier for the agent (defaults to name)
|
||||
"""
|
||||
self._executor = executor
|
||||
self._name = name
|
||||
self._id = agent_id if agent_id is not None else name
|
||||
self._display_name = name
|
||||
self._description = f"Durable agent proxy for {name}"
|
||||
|
||||
@property
|
||||
def id(self) -> str:
|
||||
"""Get the unique identifier for this agent."""
|
||||
return self._id
|
||||
|
||||
@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
|
||||
|
||||
def run( # pyright: ignore[reportIncompatibleMethodOverride]
|
||||
self,
|
||||
messages: str | ChatMessage | list[str] | list[ChatMessage] | None = None,
|
||||
*,
|
||||
thread: AgentThread | None = None,
|
||||
response_format: type[BaseModel] | None = None,
|
||||
enable_tool_calls: bool = True,
|
||||
) -> TaskT:
|
||||
"""Execute the agent via the injected provider.
|
||||
|
||||
Note:
|
||||
This method overrides AgentProtocol.run() with a different return type:
|
||||
- AgentProtocol.run() returns Coroutine[Any, Any, AgentRunResponse] (async)
|
||||
- DurableAIAgent.run() returns TaskT (Task object for yielding)
|
||||
|
||||
This is intentional to support orchestration contexts that use yield patterns
|
||||
instead of async/await patterns.
|
||||
|
||||
Returns:
|
||||
TaskT: The task type specific to the executor
|
||||
"""
|
||||
message_str = self._normalize_messages(messages)
|
||||
|
||||
run_request = self._executor.get_run_request(
|
||||
message=message_str,
|
||||
response_format=response_format,
|
||||
enable_tool_calls=enable_tool_calls,
|
||||
)
|
||||
|
||||
return self._executor.run_durable_agent(
|
||||
agent_name=self._name,
|
||||
run_request=run_request,
|
||||
thread=thread,
|
||||
)
|
||||
|
||||
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).
|
||||
|
||||
Args:
|
||||
messages: The message(s) to send to the agent
|
||||
thread: Optional agent thread for conversation context
|
||||
**kwargs: Additional arguments
|
||||
|
||||
Raises:
|
||||
NotImplementedError: Streaming is not supported for durable agents
|
||||
"""
|
||||
raise NotImplementedError("Streaming is not supported for durable agents")
|
||||
|
||||
def get_new_thread(self, **kwargs: Any) -> DurableAgentThread:
|
||||
"""Create a new agent thread via the provider."""
|
||||
return self._executor.get_new_thread(self._name, **kwargs)
|
||||
|
||||
def _normalize_messages(self, messages: str | ChatMessage | list[str] | list[ChatMessage] | None) -> str:
|
||||
"""Convert supported message inputs to a single string.
|
||||
|
||||
Args:
|
||||
messages: The messages to normalize
|
||||
|
||||
Returns:
|
||||
A single string representation of the messages
|
||||
"""
|
||||
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(messages) # type: ignore[arg-type]
|
||||
# List of ChatMessage
|
||||
return "\n".join([msg.text or "" for msg in messages]) # type: ignore[union-attr]
|
||||
return ""
|
||||
@@ -0,0 +1,218 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Worker wrapper for Durable Task Agent Framework.
|
||||
|
||||
This module provides the DurableAIAgentWorker class that wraps a durabletask worker
|
||||
and enables registration of agents as durable entities.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
from typing import Any
|
||||
|
||||
from agent_framework import AgentProtocol, get_logger
|
||||
from durabletask.worker import TaskHubGrpcWorker
|
||||
|
||||
from ._callbacks import AgentResponseCallbackProtocol
|
||||
from ._entities import AgentEntity, DurableTaskEntityStateProvider
|
||||
|
||||
logger = get_logger("agent_framework.durabletask.worker")
|
||||
|
||||
|
||||
class DurableAIAgentWorker:
|
||||
"""Wrapper for durabletask worker that enables agent registration.
|
||||
|
||||
This class wraps an existing TaskHubGrpcWorker instance and provides
|
||||
a convenient interface for registering agents as durable entities.
|
||||
|
||||
Example:
|
||||
```python
|
||||
from durabletask import TaskHubGrpcWorker
|
||||
from agent_framework import ChatAgent
|
||||
from agent_framework_durabletask import DurableAIAgentWorker
|
||||
|
||||
# Create the underlying worker
|
||||
worker = TaskHubGrpcWorker(host_address="localhost:4001")
|
||||
|
||||
# Wrap it with the agent worker
|
||||
agent_worker = DurableAIAgentWorker(worker)
|
||||
|
||||
# Register agents
|
||||
my_agent = ChatAgent(chat_client=client, name="assistant")
|
||||
agent_worker.add_agent(my_agent)
|
||||
|
||||
# Start the worker
|
||||
worker.start()
|
||||
```
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
worker: TaskHubGrpcWorker,
|
||||
callback: AgentResponseCallbackProtocol | None = None,
|
||||
):
|
||||
"""Initialize the worker wrapper.
|
||||
|
||||
Args:
|
||||
worker: The durabletask worker instance to wrap
|
||||
callback: Optional callback for agent response notifications
|
||||
"""
|
||||
self._worker = worker
|
||||
self._callback = callback
|
||||
self._registered_agents: dict[str, AgentProtocol] = {}
|
||||
logger.debug("[DurableAIAgentWorker] Initialized with worker type: %s", type(worker).__name__)
|
||||
|
||||
def add_agent(
|
||||
self,
|
||||
agent: AgentProtocol,
|
||||
callback: AgentResponseCallbackProtocol | None = None,
|
||||
) -> None:
|
||||
"""Register an agent with the worker.
|
||||
|
||||
This method creates a durable entity class for the agent and registers
|
||||
it with the underlying durabletask worker. The entity will be accessible
|
||||
by the name "dafx-{agent_name}".
|
||||
|
||||
Args:
|
||||
agent: The agent to register (must have a name)
|
||||
callback: Optional callback for this specific agent (overrides worker-level callback)
|
||||
|
||||
Raises:
|
||||
ValueError: If the agent doesn't have a name or is already registered
|
||||
"""
|
||||
agent_name = agent.name
|
||||
if not agent_name:
|
||||
raise ValueError("Agent must have a name to be registered")
|
||||
|
||||
if agent_name in self._registered_agents:
|
||||
raise ValueError(f"Agent '{agent_name}' is already registered")
|
||||
|
||||
logger.info("[DurableAIAgentWorker] Registering agent: %s as entity: dafx-%s", agent_name, agent_name)
|
||||
|
||||
# Store the agent reference
|
||||
self._registered_agents[agent_name] = agent
|
||||
|
||||
# Use agent-specific callback if provided, otherwise use worker-level callback
|
||||
effective_callback = callback or self._callback
|
||||
|
||||
# Create a configured entity class using the factory
|
||||
entity_class = self.__create_agent_entity(agent, effective_callback)
|
||||
|
||||
# Register the entity class with the worker
|
||||
# The worker.add_entity method takes a class
|
||||
entity_registered: str = self._worker.add_entity(entity_class) # pyright: ignore[reportUnknownMemberType]
|
||||
|
||||
logger.debug(
|
||||
"[DurableAIAgentWorker] Successfully registered entity class %s for agent: %s",
|
||||
entity_registered,
|
||||
agent_name,
|
||||
)
|
||||
|
||||
def start(self) -> None:
|
||||
"""Start the worker to begin processing tasks.
|
||||
|
||||
Note:
|
||||
This method delegates to the underlying worker's start method.
|
||||
The worker will block until stopped.
|
||||
"""
|
||||
logger.info("[DurableAIAgentWorker] Starting worker with %d registered agents", len(self._registered_agents))
|
||||
self._worker.start()
|
||||
|
||||
def stop(self) -> None:
|
||||
"""Stop the worker gracefully.
|
||||
|
||||
Note:
|
||||
This method delegates to the underlying worker's stop method.
|
||||
"""
|
||||
logger.info("[DurableAIAgentWorker] Stopping worker")
|
||||
self._worker.stop()
|
||||
|
||||
@property
|
||||
def registered_agent_names(self) -> list[str]:
|
||||
"""Get the names of all registered agents.
|
||||
|
||||
Returns:
|
||||
List of agent names (without the dafx- prefix)
|
||||
"""
|
||||
return list(self._registered_agents.keys())
|
||||
|
||||
def __create_agent_entity(
|
||||
self,
|
||||
agent: AgentProtocol,
|
||||
callback: AgentResponseCallbackProtocol | None = None,
|
||||
) -> type[DurableTaskEntityStateProvider]:
|
||||
"""Factory function to create a DurableEntity class configured with an agent.
|
||||
|
||||
This factory creates a new class that combines the entity state provider
|
||||
with the agent execution logic. Each agent gets its own entity class.
|
||||
|
||||
Args:
|
||||
agent: The agent instance to wrap
|
||||
callback: Optional callback for agent responses
|
||||
|
||||
Returns:
|
||||
A new DurableEntity subclass configured for this agent
|
||||
"""
|
||||
agent_name = agent.name or type(agent).__name__
|
||||
entity_name = f"dafx-{agent_name}"
|
||||
|
||||
class ConfiguredAgentEntity(DurableTaskEntityStateProvider):
|
||||
"""Durable entity configured with a specific agent instance."""
|
||||
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
# Create the AgentEntity with this state provider
|
||||
self._agent_entity = AgentEntity(
|
||||
agent=agent,
|
||||
callback=callback,
|
||||
state_provider=self,
|
||||
)
|
||||
logger.debug(
|
||||
"[ConfiguredAgentEntity] Initialized entity for agent: %s (entity name: %s)",
|
||||
agent_name,
|
||||
entity_name,
|
||||
)
|
||||
|
||||
def run(self, request: Any) -> Any:
|
||||
"""Handle run requests from clients or orchestrations.
|
||||
|
||||
Args:
|
||||
request: RunRequest as dict or string
|
||||
|
||||
Returns:
|
||||
AgentRunResponse as dict
|
||||
"""
|
||||
logger.debug("[ConfiguredAgentEntity.run] Executing agent: %s", agent_name)
|
||||
# Get or create event loop for async execution
|
||||
try:
|
||||
loop = asyncio.get_event_loop()
|
||||
except RuntimeError:
|
||||
loop = asyncio.new_event_loop()
|
||||
asyncio.set_event_loop(loop)
|
||||
|
||||
# Run the async agent execution synchronously
|
||||
if loop.is_running():
|
||||
# If loop is already running (shouldn't happen in entity context),
|
||||
# create a temporary loop
|
||||
temp_loop = asyncio.new_event_loop()
|
||||
try:
|
||||
response = temp_loop.run_until_complete(self._agent_entity.run(request))
|
||||
finally:
|
||||
temp_loop.close()
|
||||
else:
|
||||
response = loop.run_until_complete(self._agent_entity.run(request))
|
||||
|
||||
return response.to_dict()
|
||||
|
||||
def reset(self) -> None:
|
||||
"""Reset the agent's conversation history."""
|
||||
logger.debug("[ConfiguredAgentEntity.reset] Resetting agent: %s", agent_name)
|
||||
self._agent_entity.reset()
|
||||
|
||||
# Set the entity name to match the prefixed agent name
|
||||
# This is used by durabletask to register the entity
|
||||
ConfiguredAgentEntity.__name__ = entity_name
|
||||
ConfiguredAgentEntity.__qualname__ = entity_name
|
||||
|
||||
return ConfiguredAgentEntity
|
||||
@@ -0,0 +1,272 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Unit tests for AgentSessionId and DurableAgentThread."""
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentThread
|
||||
|
||||
from agent_framework_durabletask._models import AgentSessionId, DurableAgentThread
|
||||
|
||||
|
||||
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_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"
|
||||
|
||||
|
||||
class TestDurableAgentThread:
|
||||
"""Test suite for DurableAgentThread."""
|
||||
|
||||
def test_init_with_session_id(self) -> None:
|
||||
"""Test DurableAgentThread initialization with session ID."""
|
||||
session_id = AgentSessionId(name="TestAgent", key="test-key")
|
||||
thread = DurableAgentThread(session_id=session_id)
|
||||
|
||||
assert thread.session_id is not None
|
||||
assert thread.session_id == session_id
|
||||
|
||||
def test_init_without_session_id(self) -> None:
|
||||
"""Test DurableAgentThread initialization without session ID."""
|
||||
thread = DurableAgentThread()
|
||||
|
||||
assert thread.session_id is None
|
||||
|
||||
def test_session_id_setter(self) -> None:
|
||||
"""Test setting a session ID to an existing thread."""
|
||||
thread = DurableAgentThread()
|
||||
assert thread.session_id is None
|
||||
|
||||
session_id = AgentSessionId(name="TestAgent", key="test-key")
|
||||
thread.session_id = session_id
|
||||
|
||||
assert thread.session_id is not None
|
||||
assert thread.session_id == session_id
|
||||
assert thread.session_id.name == "TestAgent"
|
||||
|
||||
def test_from_session_id(self) -> None:
|
||||
"""Test creating DurableAgentThread from session ID."""
|
||||
session_id = AgentSessionId(name="TestAgent", key="test-key")
|
||||
thread = DurableAgentThread.from_session_id(session_id)
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
assert thread.session_id is not None
|
||||
assert thread.session_id == session_id
|
||||
assert thread.session_id.name == "TestAgent"
|
||||
assert thread.session_id.key == "test-key"
|
||||
|
||||
def test_from_session_id_with_service_thread_id(self) -> None:
|
||||
"""Test creating DurableAgentThread with service thread ID."""
|
||||
session_id = AgentSessionId(name="TestAgent", key="test-key")
|
||||
thread = DurableAgentThread.from_session_id(session_id, service_thread_id="service-123")
|
||||
|
||||
assert thread.session_id is not None
|
||||
assert thread.session_id == session_id
|
||||
assert thread.service_thread_id == "service-123"
|
||||
|
||||
async def test_serialize_with_session_id(self) -> None:
|
||||
"""Test serialization includes session ID."""
|
||||
session_id = AgentSessionId(name="TestAgent", key="test-key")
|
||||
thread = DurableAgentThread(session_id=session_id)
|
||||
|
||||
serialized = await thread.serialize()
|
||||
|
||||
assert isinstance(serialized, dict)
|
||||
assert "durable_session_id" in serialized
|
||||
assert serialized["durable_session_id"] == "@TestAgent@test-key"
|
||||
|
||||
async def test_serialize_without_session_id(self) -> None:
|
||||
"""Test serialization without session ID."""
|
||||
thread = DurableAgentThread()
|
||||
|
||||
serialized = await thread.serialize()
|
||||
|
||||
assert isinstance(serialized, dict)
|
||||
assert "durable_session_id" not in serialized
|
||||
|
||||
async def test_deserialize_with_session_id(self) -> None:
|
||||
"""Test deserialization restores session ID."""
|
||||
serialized = {
|
||||
"service_thread_id": "thread-123",
|
||||
"durable_session_id": "@TestAgent@test-key",
|
||||
}
|
||||
|
||||
thread = await DurableAgentThread.deserialize(serialized)
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
assert thread.session_id is not None
|
||||
assert thread.session_id.name == "TestAgent"
|
||||
assert thread.session_id.key == "test-key"
|
||||
assert thread.service_thread_id == "thread-123"
|
||||
|
||||
async def test_deserialize_without_session_id(self) -> None:
|
||||
"""Test deserialization without session ID."""
|
||||
serialized = {
|
||||
"service_thread_id": "thread-456",
|
||||
}
|
||||
|
||||
thread = await DurableAgentThread.deserialize(serialized)
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
assert thread.session_id is None
|
||||
assert thread.service_thread_id == "thread-456"
|
||||
|
||||
async def test_round_trip_serialization(self) -> None:
|
||||
"""Test round-trip serialization preserves session ID."""
|
||||
session_id = AgentSessionId(name="TestAgent", key="test-key-789")
|
||||
original = DurableAgentThread(session_id=session_id)
|
||||
|
||||
serialized = await original.serialize()
|
||||
restored = await DurableAgentThread.deserialize(serialized)
|
||||
|
||||
assert isinstance(restored, DurableAgentThread)
|
||||
assert restored.session_id is not None
|
||||
assert restored.session_id.name == session_id.name
|
||||
assert restored.session_id.key == session_id.key
|
||||
|
||||
async def test_deserialize_invalid_session_id_type(self) -> None:
|
||||
"""Test deserialization with invalid session ID type raises error."""
|
||||
serialized = {
|
||||
"service_thread_id": "thread-123",
|
||||
"durable_session_id": 12345, # Invalid type
|
||||
}
|
||||
|
||||
with pytest.raises(ValueError, match="durable_session_id must be a string"):
|
||||
await DurableAgentThread.deserialize(serialized)
|
||||
|
||||
|
||||
class TestAgentThreadCompatibility:
|
||||
"""Test suite for compatibility between AgentThread and DurableAgentThread."""
|
||||
|
||||
async def test_agent_thread_serialize(self) -> None:
|
||||
"""Test that base AgentThread can be serialized."""
|
||||
thread = AgentThread()
|
||||
|
||||
serialized = await thread.serialize()
|
||||
|
||||
assert isinstance(serialized, dict)
|
||||
assert "service_thread_id" in serialized
|
||||
|
||||
async def test_agent_thread_deserialize(self) -> None:
|
||||
"""Test that base AgentThread can be deserialized."""
|
||||
thread = AgentThread()
|
||||
serialized = await thread.serialize()
|
||||
|
||||
restored = await AgentThread.deserialize(serialized)
|
||||
|
||||
assert isinstance(restored, AgentThread)
|
||||
assert restored.service_thread_id == thread.service_thread_id
|
||||
|
||||
async def test_durable_thread_is_agent_thread(self) -> None:
|
||||
"""Test that DurableAgentThread is an AgentThread."""
|
||||
thread = DurableAgentThread()
|
||||
|
||||
assert isinstance(thread, AgentThread)
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
|
||||
|
||||
class TestModelIntegration:
|
||||
"""Test suite for integration between models."""
|
||||
|
||||
def test_session_id_string_format(self) -> None:
|
||||
"""Test that AgentSessionId string format is consistent."""
|
||||
session_id = AgentSessionId.with_random_key("AgentEntity")
|
||||
session_id_str = str(session_id)
|
||||
|
||||
assert session_id_str.startswith("@AgentEntity@")
|
||||
|
||||
async def test_thread_with_session_preserves_on_serialization(self) -> None:
|
||||
"""Test that thread with session ID preserves it through serialization."""
|
||||
session_id = AgentSessionId(name="TestAgent", key="preserved-key")
|
||||
thread = DurableAgentThread.from_session_id(session_id)
|
||||
|
||||
# Serialize and deserialize
|
||||
serialized = await thread.serialize()
|
||||
restored = await DurableAgentThread.deserialize(serialized)
|
||||
|
||||
# Session ID should be preserved
|
||||
assert restored.session_id is not None
|
||||
assert restored.session_id.name == "TestAgent"
|
||||
assert restored.session_id.key == "preserved-key"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "--tb=short"])
|
||||
@@ -0,0 +1,142 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Unit tests for DurableAIAgentClient.
|
||||
|
||||
Focuses on critical client workflows: agent retrieval, protocol compliance, and integration.
|
||||
Run with: pytest tests/test_client.py -v
|
||||
"""
|
||||
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentProtocol
|
||||
|
||||
from agent_framework_durabletask import DurableAgentThread, DurableAIAgentClient
|
||||
from agent_framework_durabletask._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
|
||||
from agent_framework_durabletask._shim import DurableAIAgent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_grpc_client() -> Mock:
|
||||
"""Create a mock TaskHubGrpcClient for testing."""
|
||||
return Mock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent_client(mock_grpc_client: Mock) -> DurableAIAgentClient:
|
||||
"""Create a DurableAIAgentClient with mock gRPC client."""
|
||||
return DurableAIAgentClient(mock_grpc_client)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent_client_with_custom_polling(mock_grpc_client: Mock) -> DurableAIAgentClient:
|
||||
"""Create a DurableAIAgentClient with custom polling parameters."""
|
||||
return DurableAIAgentClient(
|
||||
mock_grpc_client,
|
||||
max_poll_retries=15,
|
||||
poll_interval_seconds=0.5,
|
||||
)
|
||||
|
||||
|
||||
class TestDurableAIAgentClientGetAgent:
|
||||
"""Test core workflow: retrieving agents from the client."""
|
||||
|
||||
def test_get_agent_returns_durable_agent_shim(self, agent_client: DurableAIAgentClient) -> None:
|
||||
"""Verify get_agent returns a DurableAIAgent instance."""
|
||||
agent = agent_client.get_agent("assistant")
|
||||
|
||||
assert isinstance(agent, DurableAIAgent)
|
||||
assert isinstance(agent, AgentProtocol)
|
||||
|
||||
def test_get_agent_shim_has_correct_name(self, agent_client: DurableAIAgentClient) -> None:
|
||||
"""Verify retrieved agent has the correct name."""
|
||||
agent = agent_client.get_agent("my_agent")
|
||||
|
||||
assert agent.name == "my_agent"
|
||||
|
||||
def test_get_agent_multiple_times_returns_new_instances(self, agent_client: DurableAIAgentClient) -> None:
|
||||
"""Verify multiple get_agent calls return independent instances."""
|
||||
agent1 = agent_client.get_agent("assistant")
|
||||
agent2 = agent_client.get_agent("assistant")
|
||||
|
||||
assert agent1 is not agent2 # Different object instances
|
||||
|
||||
def test_get_agent_different_agents(self, agent_client: DurableAIAgentClient) -> None:
|
||||
"""Verify client can retrieve multiple different agents."""
|
||||
agent1 = agent_client.get_agent("agent1")
|
||||
agent2 = agent_client.get_agent("agent2")
|
||||
|
||||
assert agent1.name == "agent1"
|
||||
assert agent2.name == "agent2"
|
||||
|
||||
|
||||
class TestDurableAIAgentClientIntegration:
|
||||
"""Test integration scenarios between client and agent shim."""
|
||||
|
||||
def test_client_agent_has_working_run_method(self, agent_client: DurableAIAgentClient) -> None:
|
||||
"""Verify agent from client has callable run method (even if not yet implemented)."""
|
||||
agent = agent_client.get_agent("assistant")
|
||||
|
||||
assert hasattr(agent, "run")
|
||||
assert callable(agent.run)
|
||||
|
||||
def test_client_agent_can_create_threads(self, agent_client: DurableAIAgentClient) -> None:
|
||||
"""Verify agent from client can create DurableAgentThread instances."""
|
||||
agent = agent_client.get_agent("assistant")
|
||||
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
|
||||
def test_client_agent_thread_with_parameters(self, agent_client: DurableAIAgentClient) -> None:
|
||||
"""Verify agent can create threads with custom parameters."""
|
||||
agent = agent_client.get_agent("assistant")
|
||||
|
||||
thread = agent.get_new_thread(service_thread_id="client-session-123")
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
assert thread.service_thread_id == "client-session-123"
|
||||
|
||||
|
||||
class TestDurableAIAgentClientPollingConfiguration:
|
||||
"""Test polling configuration parameters for DurableAIAgentClient."""
|
||||
|
||||
def test_client_uses_default_polling_parameters(self, agent_client: DurableAIAgentClient) -> None:
|
||||
"""Verify client initializes with default polling parameters."""
|
||||
assert agent_client.max_poll_retries == DEFAULT_MAX_POLL_RETRIES
|
||||
assert agent_client.poll_interval_seconds == DEFAULT_POLL_INTERVAL_SECONDS
|
||||
|
||||
def test_client_accepts_custom_polling_parameters(
|
||||
self, agent_client_with_custom_polling: DurableAIAgentClient
|
||||
) -> None:
|
||||
"""Verify client accepts and stores custom polling parameters."""
|
||||
assert agent_client_with_custom_polling.max_poll_retries == 15
|
||||
assert agent_client_with_custom_polling.poll_interval_seconds == 0.5
|
||||
|
||||
def test_client_validates_max_poll_retries(self, mock_grpc_client: Mock) -> None:
|
||||
"""Verify client validates and normalizes max_poll_retries."""
|
||||
# Test with zero - should enforce minimum of 1
|
||||
client = DurableAIAgentClient(mock_grpc_client, max_poll_retries=0)
|
||||
assert client.max_poll_retries == 1
|
||||
|
||||
# Test with negative - should enforce minimum of 1
|
||||
client = DurableAIAgentClient(mock_grpc_client, max_poll_retries=-5)
|
||||
assert client.max_poll_retries == 1
|
||||
|
||||
def test_client_validates_poll_interval_seconds(self, mock_grpc_client: Mock) -> None:
|
||||
"""Verify client validates and normalizes poll_interval_seconds."""
|
||||
# Test with zero - should use default
|
||||
client = DurableAIAgentClient(mock_grpc_client, poll_interval_seconds=0)
|
||||
assert client.poll_interval_seconds == DEFAULT_POLL_INTERVAL_SECONDS
|
||||
|
||||
# Test with negative - should use default
|
||||
client = DurableAIAgentClient(mock_grpc_client, poll_interval_seconds=-0.5)
|
||||
assert client.poll_interval_seconds == DEFAULT_POLL_INTERVAL_SECONDS
|
||||
|
||||
# Test with valid float
|
||||
client = DurableAIAgentClient(mock_grpc_client, poll_interval_seconds=2.5)
|
||||
assert client.poll_interval_seconds == 2.5
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "--tb=short"])
|
||||
@@ -112,20 +112,21 @@ class TestDurableAgentStateMessageCreatedAt:
|
||||
"""Test suite for DurableAgentStateMessage created_at field handling."""
|
||||
|
||||
def test_message_from_run_request_without_created_at_preserves_none(self) -> None:
|
||||
"""Test from_run_request preserves None created_at instead of defaulting to current time.
|
||||
"""Test from_run_request handles auto-populated created_at from RunRequest.
|
||||
|
||||
When a RunRequest has no created_at value, the resulting DurableAgentStateMessage
|
||||
should also have None for created_at, not default to current UTC time.
|
||||
When a RunRequest is created with None for created_at, RunRequest defaults it to
|
||||
current UTC time. The resulting DurableAgentStateMessage should have this timestamp.
|
||||
"""
|
||||
run_request = RunRequest(
|
||||
message="test message",
|
||||
correlation_id="corr-run",
|
||||
created_at=None, # Explicitly None
|
||||
created_at=None, # RunRequest will default this to current time
|
||||
)
|
||||
|
||||
durable_message = DurableAgentStateMessage.from_run_request(run_request)
|
||||
|
||||
assert durable_message.created_at is None
|
||||
# RunRequest auto-populates created_at, so it should not be None
|
||||
assert durable_message.created_at is not None
|
||||
|
||||
def test_message_from_run_request_with_created_at_parses_correctly(self) -> None:
|
||||
"""Test from_run_request correctly parses a valid created_at timestamp."""
|
||||
|
||||
@@ -0,0 +1,320 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Unit tests for DurableAgentExecutor implementations.
|
||||
|
||||
Focuses on critical behavioral flows for executor strategies.
|
||||
Run with: pytest tests/test_executors.py -v
|
||||
"""
|
||||
|
||||
import time
|
||||
from typing import Any
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentRunResponse, Role
|
||||
from durabletask.entities import EntityInstanceId
|
||||
from durabletask.task import Task
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agent_framework_durabletask import DurableAgentThread
|
||||
from agent_framework_durabletask._constants import DEFAULT_MAX_POLL_RETRIES, DEFAULT_POLL_INTERVAL_SECONDS
|
||||
from agent_framework_durabletask._executors import (
|
||||
ClientAgentExecutor,
|
||||
DurableAgentTask,
|
||||
OrchestrationAgentExecutor,
|
||||
)
|
||||
from agent_framework_durabletask._models import AgentSessionId, RunRequest
|
||||
|
||||
|
||||
# Fixtures
|
||||
@pytest.fixture
|
||||
def mock_client() -> Mock:
|
||||
"""Provide a mock client for ClientAgentExecutor tests."""
|
||||
client = Mock()
|
||||
client.signal_entity = Mock()
|
||||
client.get_entity = Mock(return_value=None)
|
||||
return client
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_entity_task() -> Mock:
|
||||
"""Provide a mock entity task."""
|
||||
return Mock(spec=Task)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_orchestration_context(mock_entity_task: Mock) -> Mock:
|
||||
"""Provide a mock orchestration context with call_entity configured."""
|
||||
context = Mock()
|
||||
context.call_entity = Mock(return_value=mock_entity_task)
|
||||
return context
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def sample_run_request() -> RunRequest:
|
||||
"""Provide a sample RunRequest for tests."""
|
||||
return RunRequest(message="test message", correlation_id="test-123")
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def client_executor(mock_client: Mock) -> ClientAgentExecutor:
|
||||
"""Provide a ClientAgentExecutor with minimal polling for fast tests."""
|
||||
return ClientAgentExecutor(mock_client, max_poll_retries=1, poll_interval_seconds=0.01)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def orchestration_executor(mock_orchestration_context: Mock) -> OrchestrationAgentExecutor:
|
||||
"""Provide an OrchestrationAgentExecutor."""
|
||||
return OrchestrationAgentExecutor(mock_orchestration_context)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def successful_agent_response() -> dict[str, Any]:
|
||||
"""Provide a successful agent response dictionary."""
|
||||
return {
|
||||
"messages": [{"role": "assistant", "contents": [{"type": "text", "text": "Hello!"}]}],
|
||||
"created_at": "2025-12-30T10:00:00Z",
|
||||
}
|
||||
|
||||
|
||||
class TestExecutorThreadCreation:
|
||||
"""Test that executors properly create DurableAgentThread with parameters."""
|
||||
|
||||
def test_client_executor_creates_durable_thread(self, mock_client: Mock) -> None:
|
||||
"""Verify ClientAgentExecutor creates DurableAgentThread instances."""
|
||||
executor = ClientAgentExecutor(mock_client)
|
||||
|
||||
thread = executor.get_new_thread("test_agent")
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
|
||||
def test_client_executor_forwards_kwargs_to_thread(self, mock_client: Mock) -> None:
|
||||
"""Verify ClientAgentExecutor forwards kwargs to DurableAgentThread creation."""
|
||||
executor = ClientAgentExecutor(mock_client)
|
||||
|
||||
thread = executor.get_new_thread("test_agent", service_thread_id="client-123")
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
assert thread.service_thread_id == "client-123"
|
||||
|
||||
def test_orchestration_executor_creates_durable_thread(
|
||||
self, orchestration_executor: OrchestrationAgentExecutor
|
||||
) -> None:
|
||||
"""Verify OrchestrationAgentExecutor creates DurableAgentThread instances."""
|
||||
thread = orchestration_executor.get_new_thread("test_agent")
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
|
||||
def test_orchestration_executor_forwards_kwargs_to_thread(
|
||||
self, orchestration_executor: OrchestrationAgentExecutor
|
||||
) -> None:
|
||||
"""Verify OrchestrationAgentExecutor forwards kwargs to DurableAgentThread creation."""
|
||||
thread = orchestration_executor.get_new_thread("test_agent", service_thread_id="orch-456")
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
assert thread.service_thread_id == "orch-456"
|
||||
|
||||
|
||||
class TestClientAgentExecutorRun:
|
||||
"""Test that ClientAgentExecutor.run_durable_agent works as implemented."""
|
||||
|
||||
def test_client_executor_run_returns_response(
|
||||
self, client_executor: ClientAgentExecutor, sample_run_request: RunRequest
|
||||
) -> None:
|
||||
"""Verify ClientAgentExecutor.run_durable_agent returns AgentRunResponse (synchronous)."""
|
||||
result = client_executor.run_durable_agent("test_agent", sample_run_request)
|
||||
|
||||
# Verify it returns an AgentRunResponse (synchronous, not a coroutine)
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert result is not None
|
||||
|
||||
|
||||
class TestClientAgentExecutorPollingConfiguration:
|
||||
"""Test polling configuration parameters for ClientAgentExecutor."""
|
||||
|
||||
def test_executor_uses_default_polling_parameters(self, mock_client: Mock) -> None:
|
||||
"""Verify executor initializes with default polling parameters."""
|
||||
executor = ClientAgentExecutor(mock_client)
|
||||
|
||||
assert executor.max_poll_retries == DEFAULT_MAX_POLL_RETRIES
|
||||
assert executor.poll_interval_seconds == DEFAULT_POLL_INTERVAL_SECONDS
|
||||
|
||||
def test_executor_accepts_custom_polling_parameters(self, mock_client: Mock) -> None:
|
||||
"""Verify executor accepts and stores custom polling parameters."""
|
||||
executor = ClientAgentExecutor(mock_client, max_poll_retries=20, poll_interval_seconds=0.5)
|
||||
|
||||
assert executor.max_poll_retries == 20
|
||||
assert executor.poll_interval_seconds == 0.5
|
||||
|
||||
def test_executor_respects_custom_max_poll_retries(self, mock_client: Mock, sample_run_request: RunRequest) -> None:
|
||||
"""Verify executor respects custom max_poll_retries during polling."""
|
||||
# Create executor with only 2 retries
|
||||
executor = ClientAgentExecutor(mock_client, max_poll_retries=2, poll_interval_seconds=0.01)
|
||||
|
||||
# Run the agent
|
||||
result = executor.run_durable_agent("test_agent", sample_run_request)
|
||||
|
||||
# Verify it returns AgentRunResponse (should timeout after 2 attempts)
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
|
||||
# Verify get_entity was called 2 times (max_poll_retries)
|
||||
assert mock_client.get_entity.call_count == 2
|
||||
|
||||
def test_executor_respects_custom_poll_interval(self, mock_client: Mock, sample_run_request: RunRequest) -> None:
|
||||
"""Verify executor respects custom poll_interval_seconds during polling."""
|
||||
# Create executor with very short interval
|
||||
executor = ClientAgentExecutor(mock_client, max_poll_retries=3, poll_interval_seconds=0.01)
|
||||
|
||||
# Measure time taken
|
||||
start = time.time()
|
||||
result = executor.run_durable_agent("test_agent", sample_run_request)
|
||||
elapsed = time.time() - start
|
||||
|
||||
# Should take roughly 3 * 0.01 = 0.03 seconds (plus overhead)
|
||||
# Be generous with timing to avoid flakiness
|
||||
assert elapsed < 0.2 # Should be quick with 0.01 interval
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
|
||||
|
||||
class TestOrchestrationAgentExecutorRun:
|
||||
"""Test OrchestrationAgentExecutor.run_durable_agent implementation."""
|
||||
|
||||
def test_orchestration_executor_run_returns_durable_agent_task(
|
||||
self, orchestration_executor: OrchestrationAgentExecutor, sample_run_request: RunRequest
|
||||
) -> None:
|
||||
"""Verify OrchestrationAgentExecutor.run_durable_agent returns DurableAgentTask."""
|
||||
result = orchestration_executor.run_durable_agent("test_agent", sample_run_request)
|
||||
|
||||
assert isinstance(result, DurableAgentTask)
|
||||
|
||||
def test_orchestration_executor_calls_entity_with_correct_parameters(
|
||||
self,
|
||||
mock_orchestration_context: Mock,
|
||||
orchestration_executor: OrchestrationAgentExecutor,
|
||||
sample_run_request: RunRequest,
|
||||
) -> None:
|
||||
"""Verify call_entity is invoked with correct entity ID and request."""
|
||||
orchestration_executor.run_durable_agent("test_agent", sample_run_request)
|
||||
|
||||
# Verify call_entity was called once
|
||||
assert mock_orchestration_context.call_entity.call_count == 1
|
||||
|
||||
# Get the call arguments
|
||||
call_args = mock_orchestration_context.call_entity.call_args
|
||||
entity_id_arg = call_args[0][0]
|
||||
operation_arg = call_args[0][1]
|
||||
request_dict_arg = call_args[0][2]
|
||||
|
||||
# Verify entity ID
|
||||
assert isinstance(entity_id_arg, EntityInstanceId)
|
||||
assert entity_id_arg.entity == "dafx-test_agent"
|
||||
|
||||
# Verify operation name
|
||||
assert operation_arg == "run"
|
||||
|
||||
# Verify request dict
|
||||
assert request_dict_arg == sample_run_request.to_dict()
|
||||
|
||||
def test_orchestration_executor_uses_thread_session_id(
|
||||
self,
|
||||
mock_orchestration_context: Mock,
|
||||
orchestration_executor: OrchestrationAgentExecutor,
|
||||
sample_run_request: RunRequest,
|
||||
) -> None:
|
||||
"""Verify executor uses thread's session ID when provided."""
|
||||
# Create thread with specific session ID
|
||||
session_id = AgentSessionId(name="test_agent", key="specific-key-123")
|
||||
thread = DurableAgentThread.from_session_id(session_id)
|
||||
|
||||
result = orchestration_executor.run_durable_agent("test_agent", sample_run_request, thread=thread)
|
||||
|
||||
# Verify call_entity was called with the specific key
|
||||
call_args = mock_orchestration_context.call_entity.call_args
|
||||
entity_id_arg = call_args[0][0]
|
||||
|
||||
assert entity_id_arg.key == "specific-key-123"
|
||||
assert isinstance(result, DurableAgentTask)
|
||||
|
||||
|
||||
class TestDurableAgentTask:
|
||||
"""Test DurableAgentTask completion and response transformation."""
|
||||
|
||||
def test_durable_agent_task_transforms_successful_result(
|
||||
self, mock_entity_task: Mock, successful_agent_response: dict[str, Any]
|
||||
) -> None:
|
||||
"""Verify DurableAgentTask converts successful entity result to AgentRunResponse."""
|
||||
mock_entity_task.is_failed = False
|
||||
mock_entity_task.get_result = Mock(return_value=successful_agent_response)
|
||||
|
||||
task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")
|
||||
|
||||
# Simulate child task completion
|
||||
task.on_child_completed(mock_entity_task)
|
||||
|
||||
assert task.is_complete
|
||||
result = task.get_result()
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert len(result.messages) == 1
|
||||
assert result.messages[0].role == Role.ASSISTANT
|
||||
|
||||
def test_durable_agent_task_propagates_failure(self, mock_entity_task: Mock) -> None:
|
||||
"""Verify DurableAgentTask propagates task failures."""
|
||||
mock_entity_task.is_failed = True
|
||||
mock_entity_task.get_exception = Mock(return_value=ValueError("Entity error"))
|
||||
|
||||
task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")
|
||||
|
||||
# Simulate child task completion with failure
|
||||
task.on_child_completed(mock_entity_task)
|
||||
|
||||
assert task.is_complete
|
||||
assert task.is_failed
|
||||
exception = task.get_exception()
|
||||
assert isinstance(exception, ValueError)
|
||||
assert str(exception) == "Entity error"
|
||||
|
||||
def test_durable_agent_task_validates_response_format(self, mock_entity_task: Mock) -> None:
|
||||
"""Verify DurableAgentTask validates response format when provided."""
|
||||
mock_entity_task.is_failed = False
|
||||
mock_entity_task.get_result = Mock(
|
||||
return_value={
|
||||
"messages": [{"role": "assistant", "contents": [{"type": "text", "text": '{"answer": "42"}'}]}],
|
||||
"created_at": "2025-12-30T10:00:00Z",
|
||||
}
|
||||
)
|
||||
|
||||
class TestResponse(BaseModel):
|
||||
answer: str
|
||||
|
||||
task = DurableAgentTask(entity_task=mock_entity_task, response_format=TestResponse, correlation_id="test-123")
|
||||
|
||||
# Simulate child task completion
|
||||
task.on_child_completed(mock_entity_task)
|
||||
|
||||
assert task.is_complete
|
||||
result = task.get_result()
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
|
||||
def test_durable_agent_task_ignores_duplicate_completion(
|
||||
self, mock_entity_task: Mock, successful_agent_response: dict[str, Any]
|
||||
) -> None:
|
||||
"""Verify DurableAgentTask ignores duplicate completion calls."""
|
||||
mock_entity_task.is_failed = False
|
||||
mock_entity_task.get_result = Mock(return_value=successful_agent_response)
|
||||
|
||||
task = DurableAgentTask(entity_task=mock_entity_task, response_format=None, correlation_id="test-123")
|
||||
|
||||
# Simulate child task completion twice
|
||||
task.on_child_completed(mock_entity_task)
|
||||
first_result = task.get_result()
|
||||
|
||||
task.on_child_completed(mock_entity_task)
|
||||
second_result = task.get_result()
|
||||
|
||||
# Should be the same result, get_result should only be called once
|
||||
assert first_result is second_result
|
||||
assert mock_entity_task.get_result.call_count == 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "--tb=short"])
|
||||
@@ -18,9 +18,10 @@ class TestRunRequest:
|
||||
|
||||
def test_init_with_defaults(self) -> None:
|
||||
"""Test RunRequest initialization with defaults."""
|
||||
request = RunRequest(message="Hello")
|
||||
request = RunRequest(message="Hello", correlation_id="corr-001")
|
||||
|
||||
assert request.message == "Hello"
|
||||
assert request.correlation_id == "corr-001"
|
||||
assert request.role == Role.USER
|
||||
assert request.response_format is None
|
||||
assert request.enable_tool_calls is True
|
||||
@@ -30,30 +31,33 @@ class TestRunRequest:
|
||||
schema = ModuleStructuredResponse
|
||||
request = RunRequest(
|
||||
message="Hello",
|
||||
correlation_id="corr-002",
|
||||
role=Role.SYSTEM,
|
||||
response_format=schema,
|
||||
enable_tool_calls=False,
|
||||
)
|
||||
|
||||
assert request.message == "Hello"
|
||||
assert request.correlation_id == "corr-002"
|
||||
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]
|
||||
request = RunRequest(message="Hello", correlation_id="corr-003", 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")
|
||||
request = RunRequest(message="Test message", correlation_id="corr-004")
|
||||
data = request.to_dict()
|
||||
|
||||
assert data["message"] == "Test message"
|
||||
assert data["enable_tool_calls"] is True
|
||||
assert data["role"] == "user"
|
||||
assert data["correlationId"] == "corr-004"
|
||||
assert "response_format" not in data or data["response_format"] is None
|
||||
assert "thread_id" not in data
|
||||
|
||||
@@ -62,6 +66,7 @@ class TestRunRequest:
|
||||
schema = ModuleStructuredResponse
|
||||
request = RunRequest(
|
||||
message="Hello",
|
||||
correlation_id="corr-005",
|
||||
role=Role.ASSISTANT,
|
||||
response_format=schema,
|
||||
enable_tool_calls=False,
|
||||
@@ -69,6 +74,7 @@ class TestRunRequest:
|
||||
data = request.to_dict()
|
||||
|
||||
assert data["message"] == "Hello"
|
||||
assert data["correlationId"] == "corr-005"
|
||||
assert data["role"] == "assistant"
|
||||
assert data["response_format"]["__response_schema_type__"] == "pydantic_model"
|
||||
assert data["response_format"]["module"] == schema.__module__
|
||||
@@ -78,16 +84,17 @@ class TestRunRequest:
|
||||
|
||||
def test_from_dict_with_defaults(self) -> None:
|
||||
"""Test from_dict with minimal data."""
|
||||
data = {"message": "Hello"}
|
||||
data = {"message": "Hello", "correlationId": "corr-006"}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.message == "Hello"
|
||||
assert request.correlation_id == "corr-006"
|
||||
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"})
|
||||
request = RunRequest.from_dict({"message": "Hello", "correlationId": "corr-007", "thread_id": "ignored"})
|
||||
|
||||
assert request.message == "Hello"
|
||||
|
||||
@@ -95,6 +102,7 @@ class TestRunRequest:
|
||||
"""Test from_dict with all fields."""
|
||||
data = {
|
||||
"message": "Test",
|
||||
"correlationId": "corr-008",
|
||||
"role": "system",
|
||||
"response_format": {
|
||||
"__response_schema_type__": "pydantic_model",
|
||||
@@ -106,13 +114,14 @@ class TestRunRequest:
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.message == "Test"
|
||||
assert request.correlation_id == "corr-008"
|
||||
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:
|
||||
def test_from_dict_unknown_role_preserves_value(self) -> None:
|
||||
"""Test from_dict keeps custom roles intact."""
|
||||
data = {"message": "Test", "role": "reviewer"}
|
||||
data = {"message": "Test", "correlationId": "corr-009", "role": "reviewer"}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.role.value == "reviewer"
|
||||
@@ -120,15 +129,22 @@ class TestRunRequest:
|
||||
|
||||
def test_from_dict_empty_message(self) -> None:
|
||||
"""Test from_dict with empty message."""
|
||||
request = RunRequest.from_dict({})
|
||||
request = RunRequest.from_dict({"correlationId": "corr-010"})
|
||||
|
||||
assert request.message == ""
|
||||
assert request.correlation_id == "corr-010"
|
||||
assert request.role == Role.USER
|
||||
|
||||
def test_from_dict_missing_correlation_id_raises(self) -> None:
|
||||
"""Test from_dict raises when correlationId is missing."""
|
||||
with pytest.raises(ValueError, match="correlationId is required"):
|
||||
RunRequest.from_dict({"message": "Test"})
|
||||
|
||||
def test_round_trip_dict_conversion(self) -> None:
|
||||
"""Test round-trip to_dict and from_dict."""
|
||||
original = RunRequest(
|
||||
message="Test message",
|
||||
correlation_id="corr-011",
|
||||
role=Role.SYSTEM,
|
||||
response_format=ModuleStructuredResponse,
|
||||
enable_tool_calls=False,
|
||||
@@ -138,6 +154,7 @@ class TestRunRequest:
|
||||
restored = RunRequest.from_dict(data)
|
||||
|
||||
assert restored.message == original.message
|
||||
assert restored.correlation_id == original.correlation_id
|
||||
assert restored.role == original.role
|
||||
assert restored.response_format is ModuleStructuredResponse
|
||||
assert restored.enable_tool_calls == original.enable_tool_calls
|
||||
@@ -146,6 +163,7 @@ class TestRunRequest:
|
||||
"""Ensure Pydantic response formats serialize and deserialize properly."""
|
||||
original = RunRequest(
|
||||
message="Structured",
|
||||
correlation_id="corr-012",
|
||||
response_format=ModuleStructuredResponse,
|
||||
)
|
||||
|
||||
@@ -186,7 +204,7 @@ class TestRunRequest:
|
||||
original = RunRequest(
|
||||
message="Test message",
|
||||
role=Role.SYSTEM,
|
||||
correlation_id="corr-123",
|
||||
correlation_id="corr-124",
|
||||
)
|
||||
|
||||
data = original.to_dict()
|
||||
@@ -200,6 +218,7 @@ class TestRunRequest:
|
||||
"""Test RunRequest initialization with orchestration_id."""
|
||||
request = RunRequest(
|
||||
message="Test message",
|
||||
correlation_id="corr-125",
|
||||
orchestration_id="orch-123",
|
||||
)
|
||||
|
||||
@@ -210,6 +229,7 @@ class TestRunRequest:
|
||||
"""Test to_dict includes orchestrationId."""
|
||||
request = RunRequest(
|
||||
message="Test",
|
||||
correlation_id="corr-126",
|
||||
orchestration_id="orch-456",
|
||||
)
|
||||
data = request.to_dict()
|
||||
@@ -221,6 +241,7 @@ class TestRunRequest:
|
||||
"""Test to_dict excludes orchestrationId when not set."""
|
||||
request = RunRequest(
|
||||
message="Test",
|
||||
correlation_id="corr-127",
|
||||
)
|
||||
data = request.to_dict()
|
||||
|
||||
@@ -230,6 +251,7 @@ class TestRunRequest:
|
||||
"""Test from_dict with orchestrationId."""
|
||||
data = {
|
||||
"message": "Test",
|
||||
"correlationId": "corr-128",
|
||||
"orchestrationId": "orch-789",
|
||||
}
|
||||
request = RunRequest.from_dict(data)
|
||||
@@ -242,7 +264,7 @@ class TestRunRequest:
|
||||
original = RunRequest(
|
||||
message="Test message",
|
||||
role=Role.SYSTEM,
|
||||
correlation_id="corr-123",
|
||||
correlation_id="corr-129",
|
||||
orchestration_id="orch-123",
|
||||
)
|
||||
|
||||
|
||||
@@ -0,0 +1,98 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Unit tests for DurableAIAgentOrchestrationContext.
|
||||
|
||||
Focuses on critical orchestration workflows: agent retrieval and integration.
|
||||
Run with: pytest tests/test_orchestration_context.py -v
|
||||
"""
|
||||
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentProtocol
|
||||
|
||||
from agent_framework_durabletask import DurableAgentThread
|
||||
from agent_framework_durabletask._orchestration_context import DurableAIAgentOrchestrationContext
|
||||
from agent_framework_durabletask._shim import DurableAIAgent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_orchestration_context() -> Mock:
|
||||
"""Create a mock OrchestrationContext for testing."""
|
||||
return Mock()
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent_context(mock_orchestration_context: Mock) -> DurableAIAgentOrchestrationContext:
|
||||
"""Create a DurableAIAgentOrchestrationContext with mock context."""
|
||||
return DurableAIAgentOrchestrationContext(mock_orchestration_context)
|
||||
|
||||
|
||||
class TestDurableAIAgentOrchestrationContextGetAgent:
|
||||
"""Test core workflow: retrieving agents from orchestration context."""
|
||||
|
||||
def test_get_agent_returns_durable_agent_shim(self, agent_context: DurableAIAgentOrchestrationContext) -> None:
|
||||
"""Verify get_agent returns a DurableAIAgent instance."""
|
||||
agent = agent_context.get_agent("assistant")
|
||||
|
||||
assert isinstance(agent, DurableAIAgent)
|
||||
assert isinstance(agent, AgentProtocol)
|
||||
|
||||
def test_get_agent_shim_has_correct_name(self, agent_context: DurableAIAgentOrchestrationContext) -> None:
|
||||
"""Verify retrieved agent has the correct name."""
|
||||
agent = agent_context.get_agent("my_agent")
|
||||
|
||||
assert agent.name == "my_agent"
|
||||
|
||||
def test_get_agent_multiple_times_returns_new_instances(
|
||||
self, agent_context: DurableAIAgentOrchestrationContext
|
||||
) -> None:
|
||||
"""Verify multiple get_agent calls return independent instances."""
|
||||
agent1 = agent_context.get_agent("assistant")
|
||||
agent2 = agent_context.get_agent("assistant")
|
||||
|
||||
assert agent1 is not agent2 # Different object instances
|
||||
|
||||
def test_get_agent_different_agents(self, agent_context: DurableAIAgentOrchestrationContext) -> None:
|
||||
"""Verify context can retrieve multiple different agents."""
|
||||
agent1 = agent_context.get_agent("agent1")
|
||||
agent2 = agent_context.get_agent("agent2")
|
||||
|
||||
assert agent1.name == "agent1"
|
||||
assert agent2.name == "agent2"
|
||||
|
||||
|
||||
class TestDurableAIAgentOrchestrationContextIntegration:
|
||||
"""Test integration scenarios between orchestration context and agent shim."""
|
||||
|
||||
def test_orchestration_agent_has_working_run_method(
|
||||
self, agent_context: DurableAIAgentOrchestrationContext
|
||||
) -> None:
|
||||
"""Verify agent from context has callable run method (even if not yet implemented)."""
|
||||
agent = agent_context.get_agent("assistant")
|
||||
|
||||
assert hasattr(agent, "run")
|
||||
assert callable(agent.run)
|
||||
|
||||
def test_orchestration_agent_can_create_threads(self, agent_context: DurableAIAgentOrchestrationContext) -> None:
|
||||
"""Verify agent from context can create DurableAgentThread instances."""
|
||||
agent = agent_context.get_agent("assistant")
|
||||
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
|
||||
def test_orchestration_agent_thread_with_parameters(
|
||||
self, agent_context: DurableAIAgentOrchestrationContext
|
||||
) -> None:
|
||||
"""Verify agent can create threads with custom parameters."""
|
||||
agent = agent_context.get_agent("assistant")
|
||||
|
||||
thread = agent.get_new_thread(service_thread_id="orch-session-456")
|
||||
|
||||
assert isinstance(thread, DurableAgentThread)
|
||||
assert thread.service_thread_id == "orch-session-456"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "--tb=short"])
|
||||
@@ -0,0 +1,206 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Unit tests for DurableAIAgent shim and DurableAgentProvider.
|
||||
|
||||
Focuses on critical message normalization, delegation, and protocol compliance.
|
||||
Run with: pytest tests/test_shim.py -v
|
||||
"""
|
||||
|
||||
from typing import Any
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentProtocol, ChatMessage
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agent_framework_durabletask import DurableAgentThread
|
||||
from agent_framework_durabletask._executors import DurableAgentExecutor
|
||||
from agent_framework_durabletask._models import RunRequest
|
||||
from agent_framework_durabletask._shim import DurableAgentProvider, DurableAIAgent
|
||||
|
||||
|
||||
class ResponseFormatModel(BaseModel):
|
||||
"""Test Pydantic model for response format testing."""
|
||||
|
||||
result: str
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_executor() -> Mock:
|
||||
"""Create a mock executor for testing."""
|
||||
mock = Mock(spec=DurableAgentExecutor)
|
||||
mock.run_durable_agent = Mock(return_value=None)
|
||||
mock.get_new_thread = Mock(return_value=DurableAgentThread())
|
||||
|
||||
# Mock get_run_request to create actual RunRequest objects
|
||||
def create_run_request(
|
||||
message: str, response_format: type[BaseModel] | None = None, enable_tool_calls: bool = True
|
||||
) -> RunRequest:
|
||||
import uuid
|
||||
|
||||
return RunRequest(
|
||||
message=message,
|
||||
correlation_id=str(uuid.uuid4()),
|
||||
response_format=response_format,
|
||||
enable_tool_calls=enable_tool_calls,
|
||||
)
|
||||
|
||||
mock.get_run_request = Mock(side_effect=create_run_request)
|
||||
return mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def test_agent(mock_executor: Mock) -> DurableAIAgent[Any]:
|
||||
"""Create a test agent with mock executor."""
|
||||
return DurableAIAgent(mock_executor, "test_agent")
|
||||
|
||||
|
||||
class TestDurableAIAgentMessageNormalization:
|
||||
"""Test that DurableAIAgent properly normalizes various message input types."""
|
||||
|
||||
def test_run_accepts_string_message(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
|
||||
"""Verify run accepts and normalizes string messages."""
|
||||
test_agent.run("Hello, world!")
|
||||
|
||||
mock_executor.run_durable_agent.assert_called_once()
|
||||
# Verify agent_name and run_request were passed correctly as kwargs
|
||||
_, kwargs = mock_executor.run_durable_agent.call_args
|
||||
assert kwargs["agent_name"] == "test_agent"
|
||||
assert kwargs["run_request"].message == "Hello, world!"
|
||||
|
||||
def test_run_accepts_chat_message(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
|
||||
"""Verify run accepts and normalizes ChatMessage objects."""
|
||||
chat_msg = ChatMessage(role="user", text="Test message")
|
||||
test_agent.run(chat_msg)
|
||||
|
||||
mock_executor.run_durable_agent.assert_called_once()
|
||||
_, kwargs = mock_executor.run_durable_agent.call_args
|
||||
assert kwargs["run_request"].message == "Test message"
|
||||
|
||||
def test_run_accepts_list_of_strings(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
|
||||
"""Verify run accepts and joins list of strings."""
|
||||
test_agent.run(["First message", "Second message"])
|
||||
|
||||
mock_executor.run_durable_agent.assert_called_once()
|
||||
_, kwargs = mock_executor.run_durable_agent.call_args
|
||||
assert kwargs["run_request"].message == "First message\nSecond message"
|
||||
|
||||
def test_run_accepts_list_of_chat_messages(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
|
||||
"""Verify run accepts and joins list of ChatMessage objects."""
|
||||
messages = [
|
||||
ChatMessage(role="user", text="Message 1"),
|
||||
ChatMessage(role="assistant", text="Message 2"),
|
||||
]
|
||||
test_agent.run(messages)
|
||||
|
||||
mock_executor.run_durable_agent.assert_called_once()
|
||||
_, kwargs = mock_executor.run_durable_agent.call_args
|
||||
assert kwargs["run_request"].message == "Message 1\nMessage 2"
|
||||
|
||||
def test_run_handles_none_message(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
|
||||
"""Verify run handles None message gracefully."""
|
||||
test_agent.run(None)
|
||||
|
||||
mock_executor.run_durable_agent.assert_called_once()
|
||||
_, kwargs = mock_executor.run_durable_agent.call_args
|
||||
assert kwargs["run_request"].message == ""
|
||||
|
||||
def test_run_handles_empty_list(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
|
||||
"""Verify run handles empty list gracefully."""
|
||||
test_agent.run([])
|
||||
|
||||
mock_executor.run_durable_agent.assert_called_once()
|
||||
_, kwargs = mock_executor.run_durable_agent.call_args
|
||||
assert kwargs["run_request"].message == ""
|
||||
|
||||
|
||||
class TestDurableAIAgentParameterFlow:
|
||||
"""Test that parameters flow correctly through the shim to executor."""
|
||||
|
||||
def test_run_forwards_thread_parameter(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
|
||||
"""Verify run forwards thread parameter to executor."""
|
||||
thread = DurableAgentThread(service_thread_id="test-thread")
|
||||
test_agent.run("message", thread=thread)
|
||||
|
||||
mock_executor.run_durable_agent.assert_called_once()
|
||||
_, kwargs = mock_executor.run_durable_agent.call_args
|
||||
assert kwargs["thread"] == thread
|
||||
|
||||
def test_run_forwards_response_format(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
|
||||
"""Verify run forwards response_format parameter to executor."""
|
||||
test_agent.run("message", response_format=ResponseFormatModel)
|
||||
|
||||
mock_executor.run_durable_agent.assert_called_once()
|
||||
_, kwargs = mock_executor.run_durable_agent.call_args
|
||||
assert kwargs["run_request"].response_format == ResponseFormatModel
|
||||
|
||||
|
||||
class TestDurableAIAgentProtocolCompliance:
|
||||
"""Test that DurableAIAgent implements AgentProtocol correctly."""
|
||||
|
||||
def test_agent_implements_protocol(self, test_agent: DurableAIAgent[Any]) -> None:
|
||||
"""Verify DurableAIAgent implements AgentProtocol."""
|
||||
assert isinstance(test_agent, AgentProtocol)
|
||||
|
||||
def test_agent_has_required_properties(self, test_agent: DurableAIAgent[Any]) -> None:
|
||||
"""Verify DurableAIAgent has all required AgentProtocol properties."""
|
||||
assert hasattr(test_agent, "id")
|
||||
assert hasattr(test_agent, "name")
|
||||
assert hasattr(test_agent, "display_name")
|
||||
assert hasattr(test_agent, "description")
|
||||
|
||||
def test_agent_id_defaults_to_name(self, mock_executor: Mock) -> None:
|
||||
"""Verify agent id defaults to name when not provided."""
|
||||
agent: DurableAIAgent[Any] = DurableAIAgent(mock_executor, "my_agent")
|
||||
|
||||
assert agent.id == "my_agent"
|
||||
assert agent.name == "my_agent"
|
||||
|
||||
def test_agent_id_can_be_customized(self, mock_executor: Mock) -> None:
|
||||
"""Verify agent id can be set independently from name."""
|
||||
agent: DurableAIAgent[Any] = DurableAIAgent(mock_executor, "my_agent", agent_id="custom-id")
|
||||
|
||||
assert agent.id == "custom-id"
|
||||
assert agent.name == "my_agent"
|
||||
|
||||
|
||||
class TestDurableAIAgentThreadManagement:
|
||||
"""Test thread creation and management."""
|
||||
|
||||
def test_get_new_thread_delegates_to_executor(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
|
||||
"""Verify get_new_thread delegates to executor."""
|
||||
mock_thread = DurableAgentThread()
|
||||
mock_executor.get_new_thread.return_value = mock_thread
|
||||
|
||||
thread = test_agent.get_new_thread()
|
||||
|
||||
mock_executor.get_new_thread.assert_called_once_with("test_agent")
|
||||
assert thread == mock_thread
|
||||
|
||||
def test_get_new_thread_forwards_kwargs(self, test_agent: DurableAIAgent[Any], mock_executor: Mock) -> None:
|
||||
"""Verify get_new_thread forwards kwargs to executor."""
|
||||
mock_thread = DurableAgentThread(service_thread_id="thread-123")
|
||||
mock_executor.get_new_thread.return_value = mock_thread
|
||||
|
||||
test_agent.get_new_thread(service_thread_id="thread-123")
|
||||
|
||||
mock_executor.get_new_thread.assert_called_once()
|
||||
_, kwargs = mock_executor.get_new_thread.call_args
|
||||
assert kwargs["service_thread_id"] == "thread-123"
|
||||
|
||||
|
||||
class TestDurableAgentProviderInterface:
|
||||
"""Test that DurableAgentProvider defines the correct interface."""
|
||||
|
||||
def test_provider_cannot_be_instantiated(self) -> None:
|
||||
"""Verify DurableAgentProvider is abstract and cannot be instantiated."""
|
||||
with pytest.raises(TypeError):
|
||||
DurableAgentProvider() # type: ignore[abstract]
|
||||
|
||||
def test_provider_defines_get_agent_method(self) -> None:
|
||||
"""Verify DurableAgentProvider defines get_agent abstract method."""
|
||||
assert hasattr(DurableAgentProvider, "get_agent")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "--tb=short"])
|
||||
@@ -0,0 +1,168 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Unit tests for DurableAIAgentWorker.
|
||||
|
||||
Focuses on critical worker flows: agent registration, validation, callbacks, and lifecycle.
|
||||
"""
|
||||
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework_durabletask import DurableAIAgentWorker
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_grpc_worker() -> Mock:
|
||||
"""Create a mock TaskHubGrpcWorker for testing."""
|
||||
mock = Mock()
|
||||
mock.add_entity = Mock(return_value="dafx-test_agent")
|
||||
mock.start = Mock()
|
||||
mock.stop = Mock()
|
||||
return mock
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def mock_agent() -> Mock:
|
||||
"""Create a mock agent for testing."""
|
||||
agent = Mock()
|
||||
agent.name = "test_agent"
|
||||
return agent
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def agent_worker(mock_grpc_worker: Mock) -> DurableAIAgentWorker:
|
||||
"""Create a DurableAIAgentWorker with mock worker."""
|
||||
return DurableAIAgentWorker(mock_grpc_worker)
|
||||
|
||||
|
||||
class TestDurableAIAgentWorkerRegistration:
|
||||
"""Test agent registration behavior."""
|
||||
|
||||
def test_add_agent_accepts_agent_with_name(
|
||||
self, agent_worker: DurableAIAgentWorker, mock_agent: Mock, mock_grpc_worker: Mock
|
||||
) -> None:
|
||||
"""Verify that agents with names can be registered."""
|
||||
agent_worker.add_agent(mock_agent)
|
||||
|
||||
# Verify entity was registered with underlying worker
|
||||
mock_grpc_worker.add_entity.assert_called_once()
|
||||
# Verify agent name is tracked
|
||||
assert "test_agent" in agent_worker.registered_agent_names
|
||||
|
||||
def test_add_agent_rejects_agent_without_name(self, agent_worker: DurableAIAgentWorker) -> None:
|
||||
"""Verify that agents without names are rejected."""
|
||||
agent_no_name = Mock()
|
||||
agent_no_name.name = None
|
||||
|
||||
with pytest.raises(ValueError, match="Agent must have a name"):
|
||||
agent_worker.add_agent(agent_no_name)
|
||||
|
||||
def test_add_agent_rejects_empty_name(self, agent_worker: DurableAIAgentWorker) -> None:
|
||||
"""Verify that agents with empty names are rejected."""
|
||||
agent_empty_name = Mock()
|
||||
agent_empty_name.name = ""
|
||||
|
||||
with pytest.raises(ValueError, match="Agent must have a name"):
|
||||
agent_worker.add_agent(agent_empty_name)
|
||||
|
||||
def test_add_agent_rejects_duplicate_names(self, agent_worker: DurableAIAgentWorker, mock_agent: Mock) -> None:
|
||||
"""Verify duplicate agent names are not allowed."""
|
||||
agent_worker.add_agent(mock_agent)
|
||||
|
||||
# Try to register another agent with the same name
|
||||
duplicate_agent = Mock()
|
||||
duplicate_agent.name = "test_agent"
|
||||
|
||||
with pytest.raises(ValueError, match="already registered"):
|
||||
agent_worker.add_agent(duplicate_agent)
|
||||
|
||||
def test_registered_agent_names_tracks_multiple_agents(self, agent_worker: DurableAIAgentWorker) -> None:
|
||||
"""Verify registered_agent_names tracks all registered agents."""
|
||||
agent1 = Mock()
|
||||
agent1.name = "agent1"
|
||||
agent2 = Mock()
|
||||
agent2.name = "agent2"
|
||||
agent3 = Mock()
|
||||
agent3.name = "agent3"
|
||||
|
||||
agent_worker.add_agent(agent1)
|
||||
agent_worker.add_agent(agent2)
|
||||
agent_worker.add_agent(agent3)
|
||||
|
||||
registered = agent_worker.registered_agent_names
|
||||
assert "agent1" in registered
|
||||
assert "agent2" in registered
|
||||
assert "agent3" in registered
|
||||
assert len(registered) == 3
|
||||
|
||||
|
||||
class TestDurableAIAgentWorkerCallbacks:
|
||||
"""Test callback configuration behavior."""
|
||||
|
||||
def test_worker_level_callback_accepted(self, mock_grpc_worker: Mock) -> None:
|
||||
"""Verify worker-level callback can be set."""
|
||||
mock_callback = Mock()
|
||||
agent_worker = DurableAIAgentWorker(mock_grpc_worker, callback=mock_callback)
|
||||
|
||||
assert agent_worker is not None
|
||||
|
||||
def test_agent_level_callback_accepted(self, agent_worker: DurableAIAgentWorker, mock_agent: Mock) -> None:
|
||||
"""Verify agent-level callback can be set during registration."""
|
||||
mock_callback = Mock()
|
||||
|
||||
# Should not raise exception
|
||||
agent_worker.add_agent(mock_agent, callback=mock_callback)
|
||||
|
||||
assert "test_agent" in agent_worker.registered_agent_names
|
||||
|
||||
def test_none_callback_accepted(self, mock_grpc_worker: Mock, mock_agent: Mock) -> None:
|
||||
"""Verify None callback is valid (no callbacks required)."""
|
||||
agent_worker = DurableAIAgentWorker(mock_grpc_worker, callback=None)
|
||||
agent_worker.add_agent(mock_agent, callback=None)
|
||||
|
||||
assert "test_agent" in agent_worker.registered_agent_names
|
||||
|
||||
|
||||
class TestDurableAIAgentWorkerLifecycle:
|
||||
"""Test worker lifecycle behavior."""
|
||||
|
||||
def test_start_delegates_to_underlying_worker(
|
||||
self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock
|
||||
) -> None:
|
||||
"""Verify start() delegates to wrapped worker."""
|
||||
agent_worker.start()
|
||||
|
||||
mock_grpc_worker.start.assert_called_once()
|
||||
|
||||
def test_stop_delegates_to_underlying_worker(
|
||||
self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock
|
||||
) -> None:
|
||||
"""Verify stop() delegates to wrapped worker."""
|
||||
agent_worker.stop()
|
||||
|
||||
mock_grpc_worker.stop.assert_called_once()
|
||||
|
||||
def test_start_works_with_no_agents(self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock) -> None:
|
||||
"""Verify worker can start even with no agents registered."""
|
||||
agent_worker.start()
|
||||
|
||||
mock_grpc_worker.start.assert_called_once()
|
||||
|
||||
def test_start_works_with_multiple_agents(self, agent_worker: DurableAIAgentWorker, mock_grpc_worker: Mock) -> None:
|
||||
"""Verify worker can start with multiple agents registered."""
|
||||
agent1 = Mock()
|
||||
agent1.name = "agent1"
|
||||
agent2 = Mock()
|
||||
agent2.name = "agent2"
|
||||
|
||||
agent_worker.add_agent(agent1)
|
||||
agent_worker.add_agent(agent2)
|
||||
agent_worker.start()
|
||||
|
||||
mock_grpc_worker.start.assert_called_once()
|
||||
assert len(agent_worker.registered_agent_names) == 2
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "--tb=short"])
|
||||
Reference in New Issue
Block a user