mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Add Entity State Providers for DurableTask Package (#2981)
* Add Entity State Providers * address comments * Fix tests * Fix tests * Revert unrelated changes and remove thread_id * Revert unrelated files
This commit is contained in:
Laveesh Rohra
committed by
GitHub
Unverified
GitHub
parent
87a38bc7da
commit
a02527f00a
@@ -859,7 +859,6 @@ class AgentFunctionApp(DFAppBase):
|
||||
request_response_format=request_response_format,
|
||||
response_format=req_body.get("response_format"),
|
||||
enable_tool_calls=enable_tool_calls,
|
||||
thread_id=thread_id,
|
||||
correlation_id=correlation_id,
|
||||
created_at=datetime.now(timezone.utc),
|
||||
).to_dict()
|
||||
|
||||
@@ -8,346 +8,41 @@ allows for long-running agent conversations.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from collections.abc import AsyncIterable, Callable
|
||||
from collections.abc import Callable
|
||||
from typing import Any, cast
|
||||
|
||||
import azure.durable_functions as df
|
||||
from agent_framework import (
|
||||
AgentProtocol,
|
||||
AgentRunResponse,
|
||||
AgentRunResponseUpdate,
|
||||
ChatMessage,
|
||||
ErrorContent,
|
||||
Role,
|
||||
get_logger,
|
||||
)
|
||||
from agent_framework import AgentProtocol, get_logger
|
||||
from agent_framework_durabletask import (
|
||||
AgentCallbackContext,
|
||||
AgentEntity,
|
||||
AgentEntityStateProviderMixin,
|
||||
AgentResponseCallbackProtocol,
|
||||
DurableAgentState,
|
||||
DurableAgentStateData,
|
||||
DurableAgentStateEntry,
|
||||
DurableAgentStateRequest,
|
||||
DurableAgentStateResponse,
|
||||
RunRequest,
|
||||
)
|
||||
|
||||
logger = get_logger("agent_framework.azurefunctions.entities")
|
||||
|
||||
|
||||
class AgentEntity:
|
||||
"""Durable entity that manages agent execution and conversation state.
|
||||
class AzureFunctionEntityStateProvider(AgentEntityStateProviderMixin):
|
||||
"""Azure Functions Durable Entity state provider for AgentEntity.
|
||||
|
||||
This entity:
|
||||
- Maintains conversation history
|
||||
- Executes agent with messages
|
||||
- Stores agent responses
|
||||
- Handles tool execution
|
||||
|
||||
Operations:
|
||||
- run: Execute the agent with a message
|
||||
- run_agent: (Deprecated) Execute the agent with a message
|
||||
- reset: Clear conversation history
|
||||
|
||||
Attributes:
|
||||
agent: The AgentProtocol instance
|
||||
state: The DurableAgentState managing conversation history
|
||||
This class utilizes the Durable Entity context from `azure-functions-durable` package
|
||||
to get and set the state of the agent entity.
|
||||
"""
|
||||
|
||||
agent: AgentProtocol
|
||||
state: DurableAgentState
|
||||
def __init__(self, context: df.DurableEntityContext) -> None:
|
||||
self._context = context
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
agent: AgentProtocol,
|
||||
callback: AgentResponseCallbackProtocol | None = None,
|
||||
):
|
||||
"""Initialize the agent entity.
|
||||
def _get_state_dict(self) -> dict[str, Any]:
|
||||
raw_state = self._context.get_state(lambda: {})
|
||||
if not isinstance(raw_state, dict):
|
||||
return {}
|
||||
return cast(dict[str, Any], raw_state)
|
||||
|
||||
Args:
|
||||
agent: The Microsoft Agent Framework agent instance (must implement AgentProtocol)
|
||||
callback: Optional callback invoked during streaming updates and final responses
|
||||
"""
|
||||
self.agent = agent
|
||||
self.state = DurableAgentState()
|
||||
self.callback = callback
|
||||
def _set_state_dict(self, state: dict[str, Any]) -> None:
|
||||
self._context.set_state(state)
|
||||
|
||||
logger.debug(f"[AgentEntity] Initialized with agent type: {type(agent).__name__}")
|
||||
|
||||
def _is_error_response(self, entry: DurableAgentStateEntry) -> bool:
|
||||
"""Check if a conversation history entry is an error response.
|
||||
|
||||
Error responses should be kept in history for tracking but not sent to the agent
|
||||
since Azure OpenAI doesn't support 'error' content type.
|
||||
|
||||
Args:
|
||||
entry: A conversation history entry (DurableAgentStateEntry or dict)
|
||||
|
||||
Returns:
|
||||
True if the entry is a response containing error content, False otherwise
|
||||
"""
|
||||
if isinstance(entry, DurableAgentStateResponse):
|
||||
return entry.is_error
|
||||
return False
|
||||
|
||||
async def run_agent(
|
||||
self,
|
||||
context: df.DurableEntityContext,
|
||||
request: RunRequest | dict[str, Any] | str,
|
||||
) -> AgentRunResponse:
|
||||
"""(Deprecated) Execute the agent with a message directly in the entity.
|
||||
|
||||
Args:
|
||||
context: Entity context
|
||||
request: RunRequest object, dict, or string message (for backward compatibility)
|
||||
|
||||
Returns:
|
||||
AgentRunResponse enriched with execution metadata.
|
||||
"""
|
||||
return await self.run(context, request)
|
||||
|
||||
async def run(
|
||||
self,
|
||||
context: df.DurableEntityContext,
|
||||
request: RunRequest | dict[str, Any] | str,
|
||||
) -> AgentRunResponse:
|
||||
"""Execute the agent with a message directly in the entity.
|
||||
|
||||
Args:
|
||||
context: Entity context
|
||||
request: RunRequest object, dict, or string message (for backward compatibility)
|
||||
|
||||
Returns:
|
||||
AgentRunResponse enriched with execution metadata.
|
||||
"""
|
||||
if isinstance(request, str):
|
||||
run_request = RunRequest(message=request, role=Role.USER)
|
||||
elif isinstance(request, dict):
|
||||
run_request = RunRequest.from_dict(request)
|
||||
else:
|
||||
run_request = request
|
||||
|
||||
message = run_request.message
|
||||
thread_id = run_request.thread_id
|
||||
correlation_id = run_request.correlation_id
|
||||
if not thread_id:
|
||||
raise ValueError("RunRequest must include 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
|
||||
|
||||
logger.debug(f"[AgentEntity.run] Received Message: {run_request}")
|
||||
|
||||
state_request = DurableAgentStateRequest.from_run_request(run_request)
|
||||
self.state.data.conversation_history.append(state_request)
|
||||
|
||||
try:
|
||||
# Build messages from conversation history, excluding error responses
|
||||
# Error responses are kept in history for tracking but not sent to the agent
|
||||
chat_messages: list[ChatMessage] = [
|
||||
m.to_chat_message()
|
||||
for entry in self.state.data.conversation_history
|
||||
if not self._is_error_response(entry)
|
||||
for m in entry.messages
|
||||
]
|
||||
|
||||
run_kwargs: dict[str, Any] = {"messages": chat_messages}
|
||||
if not enable_tool_calls:
|
||||
run_kwargs["tools"] = None
|
||||
if response_format:
|
||||
run_kwargs["response_format"] = response_format
|
||||
|
||||
agent_run_response: AgentRunResponse = await self._invoke_agent(
|
||||
run_kwargs=run_kwargs,
|
||||
correlation_id=correlation_id,
|
||||
thread_id=thread_id,
|
||||
request_message=message,
|
||||
)
|
||||
|
||||
logger.debug(
|
||||
"[AgentEntity.run] Agent invocation completed - response type: %s",
|
||||
type(agent_run_response).__name__,
|
||||
)
|
||||
|
||||
try:
|
||||
response_text = agent_run_response.text if agent_run_response.text else "No response"
|
||||
logger.debug(f"Response: {response_text[:100]}...")
|
||||
except Exception as extraction_error:
|
||||
logger.error(
|
||||
"Error extracting response text: %s",
|
||||
extraction_error,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
state_response = DurableAgentStateResponse.from_run_response(correlation_id, agent_run_response)
|
||||
self.state.data.conversation_history.append(state_response)
|
||||
|
||||
logger.debug("[AgentEntity.run] AgentRunResponse stored in conversation history")
|
||||
|
||||
return agent_run_response
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception("[AgentEntity.run] Agent execution failed.")
|
||||
|
||||
# Create error message
|
||||
error_message = ChatMessage(
|
||||
role=Role.ASSISTANT, contents=[ErrorContent(message=str(exc), error_code=type(exc).__name__)]
|
||||
)
|
||||
|
||||
error_response = AgentRunResponse(messages=[error_message])
|
||||
|
||||
# Create and store error response in conversation history
|
||||
error_state_response = DurableAgentStateResponse.from_run_response(correlation_id, error_response)
|
||||
error_state_response.is_error = True
|
||||
self.state.data.conversation_history.append(error_state_response)
|
||||
|
||||
return error_response
|
||||
|
||||
async def _invoke_agent(
|
||||
self,
|
||||
run_kwargs: dict[str, Any],
|
||||
correlation_id: str,
|
||||
thread_id: str,
|
||||
request_message: str,
|
||||
) -> AgentRunResponse:
|
||||
"""Execute the agent, preferring streaming when available."""
|
||||
callback_context: AgentCallbackContext | None = None
|
||||
if self.callback is not None:
|
||||
callback_context = self._build_callback_context(
|
||||
correlation_id=correlation_id,
|
||||
thread_id=thread_id,
|
||||
request_message=request_message,
|
||||
)
|
||||
|
||||
run_stream_callable = getattr(self.agent, "run_stream", None)
|
||||
if callable(run_stream_callable):
|
||||
try:
|
||||
stream_candidate = run_stream_callable(**run_kwargs)
|
||||
if inspect.isawaitable(stream_candidate):
|
||||
stream_candidate = await stream_candidate
|
||||
|
||||
return await self._consume_stream(
|
||||
stream=cast(AsyncIterable[AgentRunResponseUpdate], stream_candidate),
|
||||
callback_context=callback_context,
|
||||
)
|
||||
except TypeError as type_error:
|
||||
if "__aiter__" not in str(type_error):
|
||||
raise
|
||||
logger.debug(
|
||||
"run_stream returned a non-async result; falling back to run(): %s",
|
||||
type_error,
|
||||
)
|
||||
except Exception as stream_error:
|
||||
logger.warning(
|
||||
"run_stream failed; falling back to run(): %s",
|
||||
stream_error,
|
||||
exc_info=True,
|
||||
)
|
||||
else:
|
||||
logger.debug("Agent does not expose run_stream; falling back to run().")
|
||||
|
||||
agent_run_response = await self._invoke_non_stream(run_kwargs)
|
||||
await self._notify_final_response(agent_run_response, callback_context)
|
||||
return agent_run_response
|
||||
|
||||
async def _consume_stream(
|
||||
self,
|
||||
stream: AsyncIterable[AgentRunResponseUpdate],
|
||||
callback_context: AgentCallbackContext | None = None,
|
||||
) -> AgentRunResponse:
|
||||
"""Consume streaming responses and build the final AgentRunResponse."""
|
||||
updates: list[AgentRunResponseUpdate] = []
|
||||
|
||||
async for update in stream:
|
||||
updates.append(update)
|
||||
await self._notify_stream_update(update, callback_context)
|
||||
|
||||
if updates:
|
||||
response = AgentRunResponse.from_agent_run_response_updates(updates)
|
||||
else:
|
||||
logger.debug("[AgentEntity] No streaming updates received; creating empty response")
|
||||
response = AgentRunResponse(messages=[])
|
||||
|
||||
await self._notify_final_response(response, callback_context)
|
||||
return response
|
||||
|
||||
async def _invoke_non_stream(self, run_kwargs: dict[str, Any]) -> AgentRunResponse:
|
||||
"""Invoke the agent without streaming support."""
|
||||
run_callable = getattr(self.agent, "run", None)
|
||||
if run_callable is None or not callable(run_callable):
|
||||
raise AttributeError("Agent does not implement run() method")
|
||||
|
||||
result = run_callable(**run_kwargs)
|
||||
if inspect.isawaitable(result):
|
||||
result = await result
|
||||
|
||||
if not isinstance(result, AgentRunResponse):
|
||||
raise TypeError(f"Agent run() must return an AgentRunResponse instance; received {type(result).__name__}")
|
||||
|
||||
return result
|
||||
|
||||
async def _notify_stream_update(
|
||||
self,
|
||||
update: AgentRunResponseUpdate,
|
||||
context: AgentCallbackContext | None,
|
||||
) -> None:
|
||||
"""Invoke the streaming callback if one is registered."""
|
||||
if self.callback is None or context is None:
|
||||
return
|
||||
|
||||
try:
|
||||
callback_result = self.callback.on_streaming_response_update(update, context)
|
||||
if inspect.isawaitable(callback_result):
|
||||
await callback_result
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"[AgentEntity] Streaming callback raised an exception: %s",
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
async def _notify_final_response(
|
||||
self,
|
||||
response: AgentRunResponse,
|
||||
context: AgentCallbackContext | None,
|
||||
) -> None:
|
||||
"""Invoke the final response callback if one is registered."""
|
||||
if self.callback is None or context is None:
|
||||
return
|
||||
|
||||
try:
|
||||
callback_result = self.callback.on_agent_response(response, context)
|
||||
if inspect.isawaitable(callback_result):
|
||||
await callback_result
|
||||
except Exception as exc:
|
||||
logger.warning(
|
||||
"[AgentEntity] Response callback raised an exception: %s",
|
||||
exc,
|
||||
exc_info=True,
|
||||
)
|
||||
|
||||
def _build_callback_context(
|
||||
self,
|
||||
correlation_id: str,
|
||||
thread_id: str,
|
||||
request_message: str,
|
||||
) -> AgentCallbackContext:
|
||||
"""Create the callback context provided to consumers."""
|
||||
agent_name = getattr(self.agent, "name", None) or type(self.agent).__name__
|
||||
return AgentCallbackContext(
|
||||
agent_name=agent_name,
|
||||
correlation_id=correlation_id,
|
||||
thread_id=thread_id,
|
||||
request_message=request_message,
|
||||
)
|
||||
|
||||
def reset(self, context: df.DurableEntityContext) -> None:
|
||||
"""Reset the entity state (clear conversation history)."""
|
||||
logger.debug("[AgentEntity.reset] Resetting entity state")
|
||||
self.state.data = DurableAgentStateData(conversation_history=[])
|
||||
logger.debug("[AgentEntity.reset] State reset complete")
|
||||
def _get_thread_id_from_entity(self) -> str:
|
||||
return self._context.entity_key
|
||||
|
||||
|
||||
def create_agent_entity(
|
||||
@@ -368,19 +63,10 @@ def create_agent_entity(
|
||||
"""Async handler that executes the entity operations."""
|
||||
try:
|
||||
logger.debug("[entity_function] Entity triggered")
|
||||
logger.debug(f"[entity_function] Operation: {context.operation_name}")
|
||||
logger.debug("[entity_function] Operation: %s", context.operation_name)
|
||||
|
||||
current_state = context.get_state(lambda: None)
|
||||
logger.debug("Retrieved state: %s", str(current_state)[:100])
|
||||
entity = AgentEntity(agent, callback)
|
||||
|
||||
if current_state is not None:
|
||||
entity.state = DurableAgentState.from_dict(current_state)
|
||||
logger.debug(
|
||||
"[entity_function] Restored entity from state (message_count: %s)", entity.state.message_count
|
||||
)
|
||||
else:
|
||||
logger.debug("[entity_function] Created new entity instance")
|
||||
state_provider = AzureFunctionEntityStateProvider(context)
|
||||
entity = AgentEntity(agent, callback, state_provider=state_provider)
|
||||
|
||||
operation = context.operation_name
|
||||
|
||||
@@ -394,21 +80,18 @@ def create_agent_entity(
|
||||
# Fall back to treating input as message string
|
||||
request = "" if input_data is None else str(cast(object, input_data))
|
||||
|
||||
result = await entity.run(context, request)
|
||||
result = await entity.run(request)
|
||||
context.set_result(result.to_dict())
|
||||
|
||||
elif operation == "reset":
|
||||
entity.reset(context)
|
||||
entity.reset()
|
||||
context.set_result({"status": "reset"})
|
||||
|
||||
else:
|
||||
logger.error("[entity_function] Unknown operation: %s", operation)
|
||||
context.set_result({"error": f"Unknown operation: {operation}"})
|
||||
|
||||
serialized_state = entity.state.to_dict()
|
||||
logger.debug("State dict: %s", serialized_state)
|
||||
context.set_state(serialized_state)
|
||||
logger.info(f"[entity_function] Operation {operation} completed successfully")
|
||||
logger.info("[entity_function] Operation %s completed successfully", operation)
|
||||
|
||||
except Exception as exc:
|
||||
logger.exception("[entity_function] Error executing entity operation %s", exc)
|
||||
|
||||
@@ -278,9 +278,9 @@ class DurableAIAgent(AgentProtocol):
|
||||
message=message_str,
|
||||
enable_tool_calls=enable_tool_calls,
|
||||
correlation_id=correlation_id,
|
||||
thread_id=session_id.key,
|
||||
response_format=response_format,
|
||||
orchestration_id=self.context.instance_id,
|
||||
created_at=self.context.current_utc_datetime,
|
||||
)
|
||||
|
||||
logger.debug("[DurableAIAgent] Calling entity %s with message: %s", entity_id, message_str[:100])
|
||||
|
||||
@@ -2,6 +2,8 @@
|
||||
|
||||
"""Unit tests for AgentFunctionApp."""
|
||||
|
||||
# pyright: reportPrivateUsage=false
|
||||
|
||||
import json
|
||||
from collections.abc import Awaitable, Callable
|
||||
from typing import Any, TypeVar
|
||||
@@ -17,15 +19,36 @@ from agent_framework_durabletask import (
|
||||
THREAD_ID_HEADER,
|
||||
WAIT_FOR_RESPONSE_FIELD,
|
||||
WAIT_FOR_RESPONSE_HEADER,
|
||||
AgentEntity,
|
||||
AgentEntityStateProviderMixin,
|
||||
DurableAgentState,
|
||||
)
|
||||
|
||||
from agent_framework_azurefunctions import AgentFunctionApp
|
||||
from agent_framework_azurefunctions._entities import AgentEntity, create_agent_entity
|
||||
from agent_framework_azurefunctions._entities import create_agent_entity
|
||||
|
||||
TFunc = TypeVar("TFunc", bound=Callable[..., Any])
|
||||
|
||||
|
||||
def _identity_decorator(func: TFunc) -> TFunc:
|
||||
return func
|
||||
|
||||
|
||||
class _InMemoryStateProvider(AgentEntityStateProviderMixin):
|
||||
def __init__(self, *, thread_id: str = "test-thread", initial_state: dict[str, Any] | None = None) -> None:
|
||||
self._thread_id = thread_id
|
||||
self._state_dict: dict[str, Any] = initial_state or {}
|
||||
|
||||
def _get_state_dict(self) -> dict[str, Any]:
|
||||
return self._state_dict
|
||||
|
||||
def _set_state_dict(self, state: dict[str, Any]) -> None:
|
||||
self._state_dict = state
|
||||
|
||||
def _get_thread_id_from_entity(self) -> str:
|
||||
return self._thread_id
|
||||
|
||||
|
||||
class TestAgentFunctionAppInit:
|
||||
"""Test suite for AgentFunctionApp initialization."""
|
||||
|
||||
@@ -89,7 +112,7 @@ class TestAgentFunctionAppInit:
|
||||
app.add_agent(mock_agent, callback=specific_callback)
|
||||
|
||||
setup_mock.assert_called_once()
|
||||
_, _, passed_callback, enable_http_endpoint, enable_mcp_tool_trigger = setup_mock.call_args[0]
|
||||
_, _, passed_callback, enable_http_endpoint, _enable_mcp_tool_trigger = setup_mock.call_args[0]
|
||||
assert passed_callback is specific_callback
|
||||
assert enable_http_endpoint is True
|
||||
|
||||
@@ -105,7 +128,7 @@ class TestAgentFunctionAppInit:
|
||||
app.add_agent(mock_agent)
|
||||
|
||||
setup_mock.assert_called_once()
|
||||
_, _, passed_callback, enable_http_endpoint, enable_mcp_tool_trigger = setup_mock.call_args[0]
|
||||
_, _, passed_callback, enable_http_endpoint, _enable_mcp_tool_trigger = setup_mock.call_args[0]
|
||||
assert passed_callback is default_callback
|
||||
assert enable_http_endpoint is True
|
||||
|
||||
@@ -120,7 +143,7 @@ class TestAgentFunctionAppInit:
|
||||
AgentFunctionApp(agents=[mock_agent], default_callback=default_callback)
|
||||
|
||||
setup_mock.assert_called_once()
|
||||
_, _, passed_callback, enable_http_endpoint, enable_mcp_tool_trigger = setup_mock.call_args[0]
|
||||
_, _, passed_callback, enable_http_endpoint, _enable_mcp_tool_trigger = setup_mock.call_args[0]
|
||||
assert passed_callback is default_callback
|
||||
assert enable_http_endpoint is True
|
||||
|
||||
@@ -336,13 +359,12 @@ class TestAgentEntityOperations:
|
||||
return_value=AgentRunResponse(messages=[ChatMessage(role="assistant", text="Test response")])
|
||||
)
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="test-conv-123"))
|
||||
|
||||
result = await entity.run(
|
||||
mock_context,
|
||||
{"message": "Test message", "thread_id": "test-conv-123", "correlationId": "corr-app-entity-1"},
|
||||
)
|
||||
result = await entity.run({
|
||||
"message": "Test message",
|
||||
"correlationId": "corr-app-entity-1",
|
||||
})
|
||||
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert result.text == "Test response"
|
||||
@@ -355,22 +377,17 @@ class TestAgentEntityOperations:
|
||||
return_value=AgentRunResponse(messages=[ChatMessage(role="assistant", text="Response 1")])
|
||||
)
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
|
||||
|
||||
# Send first message
|
||||
await entity.run(
|
||||
mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-app-entity-2"}
|
||||
)
|
||||
await entity.run({"message": "Message 1", "correlationId": "corr-app-entity-2"})
|
||||
|
||||
# Each conversation turn creates 2 entries: request and response
|
||||
history = entity.state.data.conversation_history[0].messages # Request entry
|
||||
assert len(history) == 1 # Just the user message
|
||||
|
||||
# Send second message
|
||||
await entity.run(
|
||||
mock_context, {"message": "Message 2", "thread_id": "conv-2", "correlationId": "corr-app-entity-2b"}
|
||||
)
|
||||
await entity.run({"message": "Message 2", "correlationId": "corr-app-entity-2b"})
|
||||
|
||||
# Now we have 4 entries total (2 requests + 2 responses)
|
||||
# Access the first request entry
|
||||
@@ -394,32 +411,26 @@ class TestAgentEntityOperations:
|
||||
return_value=AgentRunResponse(messages=[ChatMessage(role="assistant", text="Response")])
|
||||
)
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
|
||||
|
||||
assert len(entity.state.data.conversation_history) == 0
|
||||
|
||||
await entity.run(
|
||||
mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-app-entity-3a"}
|
||||
)
|
||||
await entity.run({"message": "Message 1", "correlationId": "corr-app-entity-3a"})
|
||||
assert len(entity.state.data.conversation_history) == 2
|
||||
|
||||
await entity.run(
|
||||
mock_context, {"message": "Message 2", "thread_id": "conv-1", "correlationId": "corr-app-entity-3b"}
|
||||
)
|
||||
await entity.run({"message": "Message 2", "correlationId": "corr-app-entity-3b"})
|
||||
assert len(entity.state.data.conversation_history) == 4
|
||||
|
||||
def test_entity_reset(self) -> None:
|
||||
"""Test that entity reset clears state."""
|
||||
mock_agent = Mock()
|
||||
entity = AgentEntity(mock_agent)
|
||||
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider())
|
||||
|
||||
# Set some state
|
||||
entity.state = DurableAgentState()
|
||||
|
||||
# Reset
|
||||
mock_context = Mock()
|
||||
entity.reset(mock_context)
|
||||
entity.reset()
|
||||
|
||||
assert len(entity.state.data.conversation_history) == 0
|
||||
|
||||
@@ -448,7 +459,6 @@ class TestAgentEntityFactory:
|
||||
mock_context.operation_name = "run"
|
||||
mock_context.get_input.return_value = {
|
||||
"message": "Test message",
|
||||
"thread_id": "conv-123",
|
||||
"correlationId": "corr-app-factory-1",
|
||||
}
|
||||
mock_context.get_state.return_value = None
|
||||
@@ -476,7 +486,6 @@ class TestAgentEntityFactory:
|
||||
mock_context.operation_name = "run_agent"
|
||||
mock_context.get_input.return_value = {
|
||||
"message": "Test message",
|
||||
"thread_id": "conv-123",
|
||||
"correlationId": "corr-app-factory-1",
|
||||
}
|
||||
mock_context.get_state.return_value = None
|
||||
@@ -596,7 +605,11 @@ class TestAgentEntityFactory:
|
||||
}
|
||||
|
||||
mock_context = Mock()
|
||||
mock_context.operation_name = "reset"
|
||||
mock_context.operation_name = "run"
|
||||
mock_context.get_input.return_value = {
|
||||
"message": "Test message",
|
||||
"correlationId": "corr-restore-1",
|
||||
}
|
||||
mock_context.get_state.return_value = existing_state
|
||||
|
||||
with patch.object(DurableAgentState, "from_dict", wraps=DurableAgentState.from_dict) as from_dict_mock:
|
||||
@@ -613,12 +626,12 @@ class TestErrorHandling:
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(side_effect=Exception("Agent error"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
entity = AgentEntity(mock_agent, state_provider=_InMemoryStateProvider(thread_id="conv-1"))
|
||||
|
||||
result = await entity.run(
|
||||
mock_context, {"message": "Test message", "thread_id": "conv-1", "correlationId": "corr-app-error-1"}
|
||||
)
|
||||
result = await entity.run({
|
||||
"message": "Test message",
|
||||
"correlationId": "corr-app-error-1",
|
||||
})
|
||||
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert len(result.messages) == 1
|
||||
@@ -711,7 +724,7 @@ class TestIncomingRequestParsing:
|
||||
|
||||
request = Mock()
|
||||
request.params = {"thread_id": "query-thread"}
|
||||
req_body = {}
|
||||
req_body: dict[str, Any] = {}
|
||||
|
||||
thread_id = app._resolve_thread_id(request, req_body)
|
||||
|
||||
@@ -778,7 +791,7 @@ class TestHttpRunRoute:
|
||||
|
||||
assert run_request["message"] == "Plain text via HTTP"
|
||||
assert run_request["role"] == "user"
|
||||
assert "thread_id" in run_request
|
||||
assert "thread_id" not in run_request
|
||||
|
||||
async def test_http_run_accept_header_returns_json(self) -> None:
|
||||
"""Test that Accept header requesting JSON results in JSON response."""
|
||||
@@ -914,9 +927,9 @@ class TestMCPToolEndpoint:
|
||||
patch.object(app, "durable_client_input") as client_mock,
|
||||
):
|
||||
# Setup mock decorator chain
|
||||
func_name_mock.return_value = lambda f: f
|
||||
mcp_trigger_mock.return_value = lambda f: f
|
||||
client_mock.return_value = lambda f: f
|
||||
func_name_mock.return_value = _identity_decorator
|
||||
mcp_trigger_mock.return_value = _identity_decorator
|
||||
client_mock.return_value = _identity_decorator
|
||||
|
||||
app._setup_mcp_tool_trigger(mock_agent.name, mock_agent.description)
|
||||
|
||||
@@ -939,11 +952,11 @@ class TestMCPToolEndpoint:
|
||||
app = AgentFunctionApp()
|
||||
|
||||
with (
|
||||
patch.object(app, "function_name", return_value=lambda f: f),
|
||||
patch.object(app, "function_name", return_value=_identity_decorator),
|
||||
patch.object(app, "mcp_tool_trigger") as mcp_trigger_mock,
|
||||
patch.object(app, "durable_client_input", return_value=lambda f: f),
|
||||
patch.object(app, "durable_client_input", return_value=_identity_decorator),
|
||||
):
|
||||
mcp_trigger_mock.return_value = lambda f: f
|
||||
mcp_trigger_mock.return_value = _identity_decorator
|
||||
|
||||
app._setup_mcp_tool_trigger(mock_agent.name, None)
|
||||
|
||||
@@ -1065,10 +1078,10 @@ class TestMCPToolEndpoint:
|
||||
app = AgentFunctionApp(agents=[mock_agent], enable_mcp_tool_trigger=True)
|
||||
|
||||
# Capture the health check handler function
|
||||
captured_handler = None
|
||||
captured_handler: Callable[[func.HttpRequest], func.HttpResponse] | None = None
|
||||
|
||||
def capture_decorator(*args, **kwargs):
|
||||
def decorator(func):
|
||||
def capture_decorator(*args: Any, **kwargs: Any) -> Callable[[TFunc], TFunc]:
|
||||
def decorator(func: TFunc) -> TFunc:
|
||||
nonlocal captured_handler
|
||||
captured_handler = func
|
||||
return func
|
||||
|
||||
@@ -1,42 +1,22 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
"""Unit tests for AgentEntity and entity operations.
|
||||
"""Unit tests for create_agent_entity factory function.
|
||||
|
||||
Run with: pytest tests/test_entities.py -v
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
from collections.abc import AsyncIterator, Callable
|
||||
from datetime import datetime
|
||||
from collections.abc import Callable
|
||||
from typing import Any, TypeVar
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
from unittest.mock import AsyncMock, Mock
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentRunResponse, AgentRunResponseUpdate, ChatMessage, ErrorContent, Role
|
||||
from agent_framework_durabletask import (
|
||||
DurableAgentState,
|
||||
DurableAgentStateData,
|
||||
DurableAgentStateMessage,
|
||||
DurableAgentStateRequest,
|
||||
DurableAgentStateTextContent,
|
||||
RunRequest,
|
||||
)
|
||||
from pydantic import BaseModel
|
||||
from agent_framework import AgentRunResponse, ChatMessage
|
||||
|
||||
from agent_framework_azurefunctions._entities import AgentEntity, create_agent_entity
|
||||
from agent_framework_azurefunctions._entities import create_agent_entity
|
||||
|
||||
TFunc = TypeVar("TFunc", bound=Callable[..., Any])
|
||||
|
||||
|
||||
def _role_value(chat_message: DurableAgentStateMessage) -> str:
|
||||
"""Helper to extract the string role from a ChatMessage."""
|
||||
role = getattr(chat_message, "role", None)
|
||||
role_value = getattr(role, "value", role)
|
||||
if role_value is None:
|
||||
return ""
|
||||
return str(role_value)
|
||||
|
||||
|
||||
def _agent_response(text: str | None) -> AgentRunResponse:
|
||||
"""Create an AgentRunResponse with a single assistant message."""
|
||||
message = (
|
||||
@@ -45,379 +25,6 @@ def _agent_response(text: str | None) -> AgentRunResponse:
|
||||
return AgentRunResponse(messages=[message])
|
||||
|
||||
|
||||
class RecordingCallback:
|
||||
"""Callback implementation capturing streaming and final responses for assertions."""
|
||||
|
||||
def __init__(self):
|
||||
self.stream_mock = AsyncMock()
|
||||
self.response_mock = AsyncMock()
|
||||
|
||||
async def on_streaming_response_update(
|
||||
self,
|
||||
update: AgentRunResponseUpdate,
|
||||
context: Any,
|
||||
) -> None:
|
||||
await self.stream_mock(update, context)
|
||||
|
||||
async def on_agent_response(self, response: AgentRunResponse, context: Any) -> None:
|
||||
await self.response_mock(response, context)
|
||||
|
||||
|
||||
class EntityStructuredResponse(BaseModel):
|
||||
answer: float
|
||||
|
||||
|
||||
class TestAgentEntityInit:
|
||||
"""Test suite for AgentEntity initialization."""
|
||||
|
||||
def test_init_creates_entity(self) -> None:
|
||||
"""Test that AgentEntity initializes correctly."""
|
||||
mock_agent = Mock()
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
|
||||
assert entity.agent == mock_agent
|
||||
assert len(entity.state.data.conversation_history) == 0
|
||||
assert entity.state.data.extension_data is None
|
||||
assert entity.state.schema_version == DurableAgentState.SCHEMA_VERSION
|
||||
|
||||
def test_init_stores_agent_reference(self) -> None:
|
||||
"""Test that the agent reference is stored correctly."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.name = "TestAgent"
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
|
||||
assert entity.agent.name == "TestAgent"
|
||||
|
||||
def test_init_with_different_agent_types(self) -> None:
|
||||
"""Test initialization with different agent types."""
|
||||
agent1 = Mock()
|
||||
agent1.__class__.__name__ = "AzureOpenAIAgent"
|
||||
|
||||
agent2 = Mock()
|
||||
agent2.__class__.__name__ = "CustomAgent"
|
||||
|
||||
entity1 = AgentEntity(agent1)
|
||||
entity2 = AgentEntity(agent2)
|
||||
|
||||
assert entity1.agent.__class__.__name__ == "AzureOpenAIAgent"
|
||||
assert entity2.agent.__class__.__name__ == "CustomAgent"
|
||||
|
||||
|
||||
class TestAgentEntityRunAgent:
|
||||
"""Test suite for the run_agent operation."""
|
||||
|
||||
async def test_run_executes_agent(self) -> None:
|
||||
"""Test that run executes the agent."""
|
||||
mock_agent = Mock()
|
||||
mock_response = _agent_response("Test response")
|
||||
mock_agent.run = AsyncMock(return_value=mock_response)
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
result = await entity.run(
|
||||
mock_context, {"message": "Test message", "thread_id": "conv-123", "correlationId": "corr-entity-1"}
|
||||
)
|
||||
|
||||
# Verify agent.run was called
|
||||
mock_agent.run.assert_called_once()
|
||||
_, kwargs = mock_agent.run.call_args
|
||||
sent_messages: list[Any] = kwargs.get("messages")
|
||||
assert len(sent_messages) == 1
|
||||
sent_message = sent_messages[0]
|
||||
assert isinstance(sent_message, ChatMessage)
|
||||
assert getattr(sent_message, "text", None) == "Test message"
|
||||
assert getattr(sent_message.role, "value", sent_message.role) == "user"
|
||||
|
||||
# Verify result
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert result.text == "Test response"
|
||||
|
||||
async def test_run_agent_executes_agent(self) -> None:
|
||||
"""Test that run_agent executes the agent."""
|
||||
mock_agent = Mock()
|
||||
mock_response = _agent_response("Test response")
|
||||
mock_agent.run = AsyncMock(return_value=mock_response)
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
result = await entity.run_agent(
|
||||
mock_context, {"message": "Test message", "thread_id": "conv-123", "correlationId": "corr-entity-1"}
|
||||
)
|
||||
|
||||
# Verify agent.run was called
|
||||
mock_agent.run.assert_called_once()
|
||||
_, kwargs = mock_agent.run.call_args
|
||||
sent_messages: list[Any] = kwargs.get("messages")
|
||||
assert len(sent_messages) == 1
|
||||
sent_message = sent_messages[0]
|
||||
assert isinstance(sent_message, ChatMessage)
|
||||
assert getattr(sent_message, "text", None) == "Test message"
|
||||
assert getattr(sent_message.role, "value", sent_message.role) == "user"
|
||||
|
||||
# Verify result
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert result.text == "Test response"
|
||||
|
||||
async def test_run_agent_streaming_callbacks_invoked(self) -> None:
|
||||
"""Ensure streaming updates trigger callbacks and run() is not used."""
|
||||
|
||||
updates = [
|
||||
AgentRunResponseUpdate(text="Hello"),
|
||||
AgentRunResponseUpdate(text=" world"),
|
||||
]
|
||||
|
||||
async def update_generator() -> AsyncIterator[AgentRunResponseUpdate]:
|
||||
for update in updates:
|
||||
yield update
|
||||
|
||||
mock_agent = Mock()
|
||||
mock_agent.name = "StreamingAgent"
|
||||
mock_agent.run_stream = Mock(return_value=update_generator())
|
||||
mock_agent.run = AsyncMock(side_effect=AssertionError("run() should not be called when streaming succeeds"))
|
||||
|
||||
callback = RecordingCallback()
|
||||
entity = AgentEntity(mock_agent, callback=callback)
|
||||
mock_context = Mock()
|
||||
|
||||
result = await entity.run(
|
||||
mock_context,
|
||||
{
|
||||
"message": "Tell me something",
|
||||
"thread_id": "session-1",
|
||||
"correlationId": "corr-stream-1",
|
||||
},
|
||||
)
|
||||
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert "Hello" in result.text
|
||||
assert callback.stream_mock.await_count == len(updates)
|
||||
assert callback.response_mock.await_count == 1
|
||||
mock_agent.run.assert_not_called()
|
||||
|
||||
# Validate callback arguments
|
||||
stream_calls = callback.stream_mock.await_args_list
|
||||
for expected_update, recorded_call in zip(updates, stream_calls, strict=True):
|
||||
assert recorded_call.args[0] is expected_update
|
||||
context = recorded_call.args[1]
|
||||
assert context.agent_name == "StreamingAgent"
|
||||
assert context.correlation_id == "corr-stream-1"
|
||||
assert context.thread_id == "session-1"
|
||||
assert context.request_message == "Tell me something"
|
||||
|
||||
final_call = callback.response_mock.await_args
|
||||
assert final_call is not None
|
||||
final_response, final_context = final_call.args
|
||||
assert final_context.agent_name == "StreamingAgent"
|
||||
assert final_context.correlation_id == "corr-stream-1"
|
||||
assert final_context.thread_id == "session-1"
|
||||
assert final_context.request_message == "Tell me something"
|
||||
assert getattr(final_response, "text", "").strip()
|
||||
|
||||
async def test_run_agent_final_callback_without_streaming(self) -> None:
|
||||
"""Ensure the final callback fires even when streaming is unavailable."""
|
||||
|
||||
mock_agent = Mock()
|
||||
mock_agent.name = "NonStreamingAgent"
|
||||
mock_agent.run_stream = None
|
||||
agent_response = _agent_response("Final response")
|
||||
mock_agent.run = AsyncMock(return_value=agent_response)
|
||||
|
||||
callback = RecordingCallback()
|
||||
entity = AgentEntity(mock_agent, callback=callback)
|
||||
mock_context = Mock()
|
||||
|
||||
result = await entity.run(
|
||||
mock_context,
|
||||
{
|
||||
"message": "Hi",
|
||||
"thread_id": "session-2",
|
||||
"correlationId": "corr-final-1",
|
||||
},
|
||||
)
|
||||
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert result.text == "Final response"
|
||||
assert callback.stream_mock.await_count == 0
|
||||
assert callback.response_mock.await_count == 1
|
||||
|
||||
final_call = callback.response_mock.await_args
|
||||
assert final_call is not None
|
||||
assert final_call.args[0] is agent_response
|
||||
final_context = final_call.args[1]
|
||||
assert final_context.agent_name == "NonStreamingAgent"
|
||||
assert final_context.correlation_id == "corr-final-1"
|
||||
assert final_context.thread_id == "session-2"
|
||||
assert final_context.request_message == "Hi"
|
||||
|
||||
async def test_run_agent_updates_conversation_history(self) -> None:
|
||||
"""Test that run_agent updates the conversation history."""
|
||||
mock_agent = Mock()
|
||||
mock_response = _agent_response("Agent response")
|
||||
mock_agent.run = AsyncMock(return_value=mock_response)
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
await entity.run(
|
||||
mock_context, {"message": "User message", "thread_id": "conv-1", "correlationId": "corr-entity-2"}
|
||||
)
|
||||
|
||||
# Should have 1 entry: user message + assistant response
|
||||
user_history = entity.state.data.conversation_history[0].messages
|
||||
assistant_history = entity.state.data.conversation_history[1].messages
|
||||
|
||||
assert len(user_history) == 1
|
||||
|
||||
user_msg = user_history[0]
|
||||
assert _role_value(user_msg) == "user"
|
||||
assert user_msg.text == "User message"
|
||||
|
||||
assistant_msg = assistant_history[0]
|
||||
assert _role_value(assistant_msg) == "assistant"
|
||||
assert assistant_msg.text == "Agent response"
|
||||
|
||||
async def test_run_agent_increments_message_count(self) -> None:
|
||||
"""Test that run_agent increments the message count."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
assert len(entity.state.data.conversation_history) == 0
|
||||
|
||||
await entity.run(
|
||||
mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-entity-3a"}
|
||||
)
|
||||
assert len(entity.state.data.conversation_history) == 2
|
||||
|
||||
await entity.run(
|
||||
mock_context, {"message": "Message 2", "thread_id": "conv-1", "correlationId": "corr-entity-3b"}
|
||||
)
|
||||
assert len(entity.state.data.conversation_history) == 4
|
||||
|
||||
await entity.run(
|
||||
mock_context, {"message": "Message 3", "thread_id": "conv-1", "correlationId": "corr-entity-3c"}
|
||||
)
|
||||
assert len(entity.state.data.conversation_history) == 6
|
||||
|
||||
async def test_run_agent_with_none_thread_id(self) -> None:
|
||||
"""Test run_agent with a None thread identifier."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
with pytest.raises(ValueError, match="thread_id"):
|
||||
await entity.run(mock_context, {"message": "Message", "thread_id": None, "correlationId": "corr-entity-5"})
|
||||
|
||||
async def test_run_agent_multiple_conversations(self) -> None:
|
||||
"""Test that run_agent maintains history across multiple messages."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
# Send multiple messages
|
||||
await entity.run(
|
||||
mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-entity-8a"}
|
||||
)
|
||||
await entity.run(
|
||||
mock_context, {"message": "Message 2", "thread_id": "conv-1", "correlationId": "corr-entity-8b"}
|
||||
)
|
||||
await entity.run(
|
||||
mock_context, {"message": "Message 3", "thread_id": "conv-1", "correlationId": "corr-entity-8c"}
|
||||
)
|
||||
|
||||
history = entity.state.data.conversation_history
|
||||
assert len(history) == 6
|
||||
assert entity.state.message_count == 6
|
||||
|
||||
|
||||
class TestAgentEntityReset:
|
||||
"""Test suite for the reset operation."""
|
||||
|
||||
def test_reset_clears_conversation_history(self) -> None:
|
||||
"""Test that reset clears the conversation history."""
|
||||
mock_agent = Mock()
|
||||
entity = AgentEntity(mock_agent)
|
||||
|
||||
# Add some history with proper DurableAgentStateEntry objects
|
||||
entity.state.data.conversation_history = [
|
||||
DurableAgentStateRequest(
|
||||
correlation_id="test-1",
|
||||
created_at=datetime.now(),
|
||||
messages=[
|
||||
DurableAgentStateMessage(
|
||||
role="user",
|
||||
contents=[DurableAgentStateTextContent(text="msg1")],
|
||||
)
|
||||
],
|
||||
),
|
||||
]
|
||||
|
||||
mock_context = Mock()
|
||||
entity.reset(mock_context)
|
||||
|
||||
assert entity.state.data.conversation_history == []
|
||||
|
||||
def test_reset_with_extension_data(self) -> None:
|
||||
"""Test that reset works when entity has extension data."""
|
||||
mock_agent = Mock()
|
||||
entity = AgentEntity(mock_agent)
|
||||
|
||||
# Set up some initial state with conversation history
|
||||
entity.state.data = DurableAgentStateData(conversation_history=[], extension_data={"some_key": "some_value"})
|
||||
|
||||
mock_context = Mock()
|
||||
entity.reset(mock_context)
|
||||
|
||||
assert len(entity.state.data.conversation_history) == 0
|
||||
|
||||
def test_reset_clears_message_count(self) -> None:
|
||||
"""Test that reset clears the message count."""
|
||||
mock_agent = Mock()
|
||||
entity = AgentEntity(mock_agent)
|
||||
|
||||
mock_context = Mock()
|
||||
entity.reset(mock_context)
|
||||
|
||||
assert len(entity.state.data.conversation_history) == 0
|
||||
|
||||
async def test_reset_after_conversation(self) -> None:
|
||||
"""Test reset after a full conversation."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
# Have a conversation
|
||||
await entity.run(
|
||||
mock_context, {"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-entity-10a"}
|
||||
)
|
||||
await entity.run(
|
||||
mock_context, {"message": "Message 2", "thread_id": "conv-1", "correlationId": "corr-entity-10b"}
|
||||
)
|
||||
|
||||
# Verify state before reset
|
||||
assert entity.state.message_count == 4
|
||||
assert len(entity.state.data.conversation_history) == 4
|
||||
|
||||
# Reset
|
||||
entity.reset(mock_context)
|
||||
|
||||
# Verify state after reset
|
||||
assert entity.state.message_count == 0
|
||||
assert len(entity.state.data.conversation_history) == 0
|
||||
|
||||
|
||||
class TestCreateAgentEntity:
|
||||
"""Test suite for the create_agent_entity factory function."""
|
||||
|
||||
@@ -439,9 +46,9 @@ class TestCreateAgentEntity:
|
||||
# Mock context
|
||||
mock_context = Mock()
|
||||
mock_context.operation_name = "run"
|
||||
mock_context.entity_key = "conv-123"
|
||||
mock_context.get_input.return_value = {
|
||||
"message": "Test message",
|
||||
"thread_id": "conv-123",
|
||||
"correlationId": "corr-entity-factory",
|
||||
}
|
||||
mock_context.get_state.return_value = None
|
||||
@@ -535,7 +142,7 @@ class TestCreateAgentEntity:
|
||||
assert state["data"] == {"conversationHistory": []}
|
||||
|
||||
def test_entity_function_restores_existing_state(self) -> None:
|
||||
"""Test that the entity function restores existing state."""
|
||||
"""Test that the entity function can operate when existing state is present."""
|
||||
mock_agent = Mock()
|
||||
|
||||
entity_function = create_agent_entity(mock_agent)
|
||||
@@ -584,482 +191,14 @@ class TestCreateAgentEntity:
|
||||
mock_context.operation_name = "reset"
|
||||
mock_context.get_state.return_value = existing_state
|
||||
|
||||
with patch.object(DurableAgentState, "from_dict", wraps=DurableAgentState.from_dict) as from_dict_mock:
|
||||
entity_function(mock_context)
|
||||
|
||||
from_dict_mock.assert_called_once_with(existing_state)
|
||||
|
||||
|
||||
class TestErrorHandling:
|
||||
"""Test suite for error handling in entities."""
|
||||
|
||||
async def test_run_agent_handles_agent_exception(self) -> None:
|
||||
"""Test that run_agent handles agent exceptions."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(side_effect=Exception("Agent failed"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
result = await entity.run(
|
||||
mock_context, {"message": "Message", "thread_id": "conv-1", "correlationId": "corr-entity-error-1"}
|
||||
)
|
||||
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert len(result.messages) == 1
|
||||
content = result.messages[0].contents[0]
|
||||
assert isinstance(content, ErrorContent)
|
||||
assert "Agent failed" in (content.message or "")
|
||||
assert content.error_code == "Exception"
|
||||
|
||||
async def test_run_agent_handles_value_error(self) -> None:
|
||||
"""Test that run_agent handles ValueError instances."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(side_effect=ValueError("Invalid input"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
result = await entity.run(
|
||||
mock_context, {"message": "Message", "thread_id": "conv-1", "correlationId": "corr-entity-error-2"}
|
||||
)
|
||||
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert len(result.messages) == 1
|
||||
content = result.messages[0].contents[0]
|
||||
assert isinstance(content, ErrorContent)
|
||||
assert content.error_code == "ValueError"
|
||||
assert "Invalid input" in str(content.message)
|
||||
|
||||
async def test_run_agent_handles_timeout_error(self) -> None:
|
||||
"""Test that run_agent handles TimeoutError instances."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(side_effect=TimeoutError("Request timeout"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
result = await entity.run(
|
||||
mock_context, {"message": "Message", "thread_id": "conv-1", "correlationId": "corr-entity-error-3"}
|
||||
)
|
||||
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert len(result.messages) == 1
|
||||
content = result.messages[0].contents[0]
|
||||
assert isinstance(content, ErrorContent)
|
||||
assert content.error_code == "TimeoutError"
|
||||
|
||||
def test_entity_function_handles_exception_in_operation(self) -> None:
|
||||
"""Test that the entity function handles exceptions gracefully."""
|
||||
mock_agent = Mock()
|
||||
|
||||
entity_function = create_agent_entity(mock_agent)
|
||||
|
||||
mock_context = Mock()
|
||||
mock_context.operation_name = "run"
|
||||
mock_context.get_input.side_effect = Exception("Input error")
|
||||
mock_context.get_state.return_value = None
|
||||
|
||||
# Execute - should not raise
|
||||
entity_function(mock_context)
|
||||
|
||||
# Verify error was set
|
||||
assert mock_context.set_result.called
|
||||
result = mock_context.set_result.call_args[0][0]
|
||||
assert "error" in result
|
||||
|
||||
async def test_run_agent_preserves_message_on_error(self) -> None:
|
||||
"""Test that run_agent preserves message information on error."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(side_effect=Exception("Error"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
result = await entity.run(
|
||||
mock_context,
|
||||
{"message": "Test message", "thread_id": "conv-123", "correlationId": "corr-entity-error-4"},
|
||||
)
|
||||
|
||||
# Even on error, message info should be preserved
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert len(result.messages) == 1
|
||||
content = result.messages[0].contents[0]
|
||||
assert isinstance(content, ErrorContent)
|
||||
|
||||
|
||||
class TestConversationHistory:
|
||||
"""Test suite for conversation history tracking."""
|
||||
|
||||
async def test_conversation_history_has_timestamps(self) -> None:
|
||||
"""Test that conversation history entries include timestamps."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
await entity.run(
|
||||
mock_context, {"message": "Message", "thread_id": "conv-1", "correlationId": "corr-entity-history-1"}
|
||||
)
|
||||
|
||||
# Check both user and assistant messages have timestamps
|
||||
for entry in entity.state.data.conversation_history:
|
||||
timestamp = entry.created_at
|
||||
assert timestamp is not None
|
||||
# Verify timestamp is in ISO format
|
||||
datetime.fromisoformat(str(timestamp))
|
||||
|
||||
async def test_conversation_history_ordering(self) -> None:
|
||||
"""Test that conversation history maintains the correct order."""
|
||||
mock_agent = Mock()
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
# Send multiple messages with different responses
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response 1"))
|
||||
await entity.run(
|
||||
mock_context,
|
||||
{"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-entity-history-2a"},
|
||||
)
|
||||
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response 2"))
|
||||
await entity.run(
|
||||
mock_context,
|
||||
{"message": "Message 2", "thread_id": "conv-1", "correlationId": "corr-entity-history-2b"},
|
||||
)
|
||||
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response 3"))
|
||||
await entity.run(
|
||||
mock_context,
|
||||
{"message": "Message 3", "thread_id": "conv-1", "correlationId": "corr-entity-history-2c"},
|
||||
)
|
||||
|
||||
# Verify order
|
||||
history = entity.state.data.conversation_history
|
||||
# Each conversation turn creates 2 entries: request and response
|
||||
assert history[0].messages[0].text == "Message 1" # Request 1
|
||||
assert history[1].messages[0].text == "Response 1" # Response 1
|
||||
assert history[2].messages[0].text == "Message 2" # Request 2
|
||||
assert history[3].messages[0].text == "Response 2" # Response 2
|
||||
assert history[4].messages[0].text == "Message 3" # Request 3
|
||||
assert history[5].messages[0].text == "Response 3" # Response 3
|
||||
|
||||
async def test_conversation_history_role_alternation(self) -> None:
|
||||
"""Test that conversation history alternates between user and assistant roles."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
await entity.run(
|
||||
mock_context,
|
||||
{"message": "Message 1", "thread_id": "conv-1", "correlationId": "corr-entity-history-3a"},
|
||||
)
|
||||
await entity.run(
|
||||
mock_context,
|
||||
{"message": "Message 2", "thread_id": "conv-1", "correlationId": "corr-entity-history-3b"},
|
||||
)
|
||||
|
||||
# Check role alternation
|
||||
history = entity.state.data.conversation_history
|
||||
# Each conversation turn creates 2 entries: request and response
|
||||
assert history[0].messages[0].role == "user" # Request 1
|
||||
assert history[1].messages[0].role == "assistant" # Response 1
|
||||
assert history[2].messages[0].role == "user" # Request 2
|
||||
assert history[3].messages[0].role == "assistant" # Response 2
|
||||
|
||||
|
||||
class TestRunRequestSupport:
|
||||
"""Test suite for RunRequest support in entities."""
|
||||
|
||||
async def test_run_agent_with_run_request_object(self) -> None:
|
||||
"""Test run_agent with a RunRequest object."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
request = RunRequest(
|
||||
message="Test message",
|
||||
thread_id="conv-123",
|
||||
role=Role.USER,
|
||||
enable_tool_calls=True,
|
||||
correlation_id="corr-runreq-1",
|
||||
)
|
||||
|
||||
result = await entity.run(mock_context, request)
|
||||
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert result.text == "Response"
|
||||
|
||||
async def test_run_agent_with_dict_request(self) -> None:
|
||||
"""Test run_agent with a dictionary request."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
request_dict = {
|
||||
"message": "Test message",
|
||||
"thread_id": "conv-456",
|
||||
"role": "system",
|
||||
"enable_tool_calls": False,
|
||||
"correlationId": "corr-runreq-2",
|
||||
}
|
||||
|
||||
result = await entity.run(mock_context, request_dict)
|
||||
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert result.text == "Response"
|
||||
|
||||
async def test_run_agent_with_string_raises_without_correlation(self) -> None:
|
||||
"""Test that run_agent rejects legacy string input without correlation ID."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
await entity.run(mock_context, "Simple message")
|
||||
|
||||
async def test_run_agent_stores_role_in_history(self) -> None:
|
||||
"""Test that run_agent stores the role in conversation history."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
# Send as system role
|
||||
request = RunRequest(
|
||||
message="System message",
|
||||
thread_id="conv-runreq-3",
|
||||
role=Role.SYSTEM,
|
||||
correlation_id="corr-runreq-3",
|
||||
)
|
||||
|
||||
await entity.run(mock_context, request)
|
||||
|
||||
# Check that system role was stored
|
||||
history = entity.state.data.conversation_history
|
||||
assert history[0].messages[0].role == "system"
|
||||
assert history[0].messages[0].text == "System message"
|
||||
|
||||
async def test_run_agent_with_response_format(self) -> None:
|
||||
"""Test run_agent with a JSON response format."""
|
||||
mock_agent = Mock()
|
||||
# Return JSON response
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response('{"answer": 42}'))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
request = RunRequest(
|
||||
message="What is the answer?",
|
||||
thread_id="conv-runreq-4",
|
||||
response_format=EntityStructuredResponse,
|
||||
correlation_id="corr-runreq-4",
|
||||
)
|
||||
|
||||
result = await entity.run(mock_context, request)
|
||||
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert result.text == '{"answer": 42}'
|
||||
assert result.value is None
|
||||
|
||||
async def test_run_agent_disable_tool_calls(self) -> None:
|
||||
"""Test run_agent with tool calls disabled."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
request = RunRequest(
|
||||
message="Test", thread_id="conv-runreq-5", enable_tool_calls=False, correlation_id="corr-runreq-5"
|
||||
)
|
||||
|
||||
result = await entity.run(mock_context, request)
|
||||
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
# Agent should have been called (tool disabling is framework-dependent)
|
||||
mock_agent.run.assert_called_once()
|
||||
|
||||
async def test_entity_function_with_run_request_dict(self) -> None:
|
||||
"""Test that the entity function handles the RunRequest dict format."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response("Response"))
|
||||
|
||||
entity_function = create_agent_entity(mock_agent)
|
||||
|
||||
mock_context = Mock()
|
||||
mock_context.operation_name = "run"
|
||||
mock_context.get_input.return_value = {
|
||||
"message": "Test message",
|
||||
"thread_id": "conv-789",
|
||||
"role": "user",
|
||||
"enable_tool_calls": True,
|
||||
"correlationId": "corr-runreq-6",
|
||||
}
|
||||
mock_context.get_state.return_value = None
|
||||
|
||||
await asyncio.to_thread(entity_function, mock_context)
|
||||
|
||||
# Verify result was set
|
||||
assert mock_context.set_result.called
|
||||
result = mock_context.set_result.call_args[0][0]
|
||||
assert isinstance(result, dict)
|
||||
|
||||
# Check if messages are present
|
||||
assert "messages" in result
|
||||
assert len(result["messages"]) > 0
|
||||
message = result["messages"][0]
|
||||
|
||||
# Check for text in various possible locations
|
||||
text_found = False
|
||||
if "text" in message and message["text"] == "Response":
|
||||
text_found = True
|
||||
elif "contents" in message:
|
||||
for content in message["contents"]:
|
||||
if isinstance(content, dict) and content.get("text") == "Response":
|
||||
text_found = True
|
||||
break
|
||||
|
||||
assert text_found, f"Response text not found in message: {message}"
|
||||
|
||||
|
||||
class TestDurableAgentStateRequestOrchestrationId:
|
||||
"""Test suite for DurableAgentStateRequest orchestration_id field."""
|
||||
|
||||
def test_request_with_orchestration_id(self) -> None:
|
||||
"""Test creating a request with an orchestration_id."""
|
||||
request = DurableAgentStateRequest(
|
||||
correlation_id="corr-123",
|
||||
created_at=datetime.now(),
|
||||
messages=[
|
||||
DurableAgentStateMessage(
|
||||
role="user",
|
||||
contents=[DurableAgentStateTextContent(text="test")],
|
||||
)
|
||||
],
|
||||
orchestration_id="orch-456",
|
||||
)
|
||||
|
||||
assert request.orchestration_id == "orch-456"
|
||||
|
||||
def test_request_to_dict_includes_orchestration_id(self) -> None:
|
||||
"""Test that to_dict includes orchestrationId when set."""
|
||||
request = DurableAgentStateRequest(
|
||||
correlation_id="corr-123",
|
||||
created_at=datetime.now(),
|
||||
messages=[
|
||||
DurableAgentStateMessage(
|
||||
role="user",
|
||||
contents=[DurableAgentStateTextContent(text="test")],
|
||||
)
|
||||
],
|
||||
orchestration_id="orch-789",
|
||||
)
|
||||
|
||||
data = request.to_dict()
|
||||
|
||||
assert "orchestrationId" in data
|
||||
assert data["orchestrationId"] == "orch-789"
|
||||
|
||||
def test_request_to_dict_excludes_orchestration_id_when_none(self) -> None:
|
||||
"""Test that to_dict excludes orchestrationId when not set."""
|
||||
request = DurableAgentStateRequest(
|
||||
correlation_id="corr-123",
|
||||
created_at=datetime.now(),
|
||||
messages=[
|
||||
DurableAgentStateMessage(
|
||||
role="user",
|
||||
contents=[DurableAgentStateTextContent(text="test")],
|
||||
)
|
||||
],
|
||||
)
|
||||
|
||||
data = request.to_dict()
|
||||
|
||||
assert "orchestrationId" not in data
|
||||
|
||||
def test_request_from_dict_with_orchestration_id(self) -> None:
|
||||
"""Test from_dict correctly parses orchestrationId."""
|
||||
data = {
|
||||
"$type": "request",
|
||||
"correlationId": "corr-123",
|
||||
"createdAt": "2024-01-01T00:00:00Z",
|
||||
"messages": [{"role": "user", "contents": [{"$type": "text", "text": "test"}]}],
|
||||
"orchestrationId": "orch-from-dict",
|
||||
}
|
||||
|
||||
request = DurableAgentStateRequest.from_dict(data)
|
||||
|
||||
assert request.orchestration_id == "orch-from-dict"
|
||||
|
||||
def test_request_from_run_request_with_orchestration_id(self) -> None:
|
||||
"""Test from_run_request correctly transfers orchestration_id."""
|
||||
run_request = RunRequest(
|
||||
message="test message",
|
||||
correlation_id="corr-run",
|
||||
orchestration_id="orch-from-run-request",
|
||||
)
|
||||
|
||||
durable_request = DurableAgentStateRequest.from_run_request(run_request)
|
||||
|
||||
assert durable_request.orchestration_id == "orch-from-run-request"
|
||||
|
||||
def test_request_from_run_request_without_orchestration_id(self) -> None:
|
||||
"""Test from_run_request correctly handles missing orchestration_id."""
|
||||
run_request = RunRequest(
|
||||
message="test message",
|
||||
correlation_id="corr-run",
|
||||
)
|
||||
|
||||
durable_request = DurableAgentStateRequest.from_run_request(run_request)
|
||||
|
||||
assert durable_request.orchestration_id is None
|
||||
|
||||
|
||||
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.
|
||||
|
||||
When a RunRequest has no created_at value, the resulting DurableAgentStateMessage
|
||||
should also have None for created_at, not default to current UTC time.
|
||||
"""
|
||||
run_request = RunRequest(
|
||||
message="test message",
|
||||
correlation_id="corr-run",
|
||||
created_at=None, # Explicitly None
|
||||
)
|
||||
|
||||
durable_message = DurableAgentStateMessage.from_run_request(run_request)
|
||||
|
||||
assert durable_message.created_at is 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."""
|
||||
run_request = RunRequest(
|
||||
message="test message",
|
||||
correlation_id="corr-run",
|
||||
created_at="2024-01-15T10:30:00Z",
|
||||
)
|
||||
|
||||
durable_message = DurableAgentStateMessage.from_run_request(run_request)
|
||||
|
||||
assert durable_message.created_at is not None
|
||||
assert durable_message.created_at.year == 2024
|
||||
assert durable_message.created_at.month == 1
|
||||
assert durable_message.created_at.day == 15
|
||||
# Reset should clear history and persist via set_state
|
||||
assert mock_context.set_state.called
|
||||
persisted_state = mock_context.set_state.call_args[0][0]
|
||||
assert persisted_state["data"]["conversationHistory"] == []
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -150,20 +150,18 @@ class TestRunRequest:
|
||||
|
||||
def test_init_with_defaults(self) -> None:
|
||||
"""Test RunRequest initialization with defaults."""
|
||||
request = RunRequest(message="Hello", thread_id="thread-default")
|
||||
request = RunRequest(message="Hello")
|
||||
|
||||
assert request.message == "Hello"
|
||||
assert request.role == Role.USER
|
||||
assert request.response_format is None
|
||||
assert request.enable_tool_calls is True
|
||||
assert request.thread_id == "thread-default"
|
||||
|
||||
def test_init_with_all_fields(self) -> None:
|
||||
"""Test RunRequest initialization with all fields."""
|
||||
schema = ModuleStructuredResponse
|
||||
request = RunRequest(
|
||||
message="Hello",
|
||||
thread_id="thread-123",
|
||||
role=Role.SYSTEM,
|
||||
response_format=schema,
|
||||
enable_tool_calls=False,
|
||||
@@ -173,31 +171,29 @@ class TestRunRequest:
|
||||
assert request.role == Role.SYSTEM
|
||||
assert request.response_format is schema
|
||||
assert request.enable_tool_calls is False
|
||||
assert request.thread_id == "thread-123"
|
||||
|
||||
def test_init_coerces_string_role(self) -> None:
|
||||
"""Ensure string role values are coerced into Role instances."""
|
||||
request = RunRequest(message="Hello", thread_id="thread-str-role", role="system") # type: ignore[arg-type]
|
||||
request = RunRequest(message="Hello", role="system") # type: ignore[arg-type]
|
||||
|
||||
assert request.role == Role.SYSTEM
|
||||
|
||||
def test_to_dict_with_defaults(self) -> None:
|
||||
"""Test to_dict with default values."""
|
||||
request = RunRequest(message="Test message", thread_id="thread-to-dict")
|
||||
request = RunRequest(message="Test message")
|
||||
data = request.to_dict()
|
||||
|
||||
assert data["message"] == "Test message"
|
||||
assert data["enable_tool_calls"] is True
|
||||
assert data["role"] == "user"
|
||||
assert "response_format" not in data or data["response_format"] is None
|
||||
assert data["thread_id"] == "thread-to-dict"
|
||||
assert "thread_id" not in data
|
||||
|
||||
def test_to_dict_with_all_fields(self) -> None:
|
||||
"""Test to_dict with all fields."""
|
||||
schema = ModuleStructuredResponse
|
||||
request = RunRequest(
|
||||
message="Hello",
|
||||
thread_id="thread-456",
|
||||
role=Role.ASSISTANT,
|
||||
response_format=schema,
|
||||
enable_tool_calls=False,
|
||||
@@ -210,17 +206,22 @@ class TestRunRequest:
|
||||
assert data["response_format"]["module"] == schema.__module__
|
||||
assert data["response_format"]["qualname"] == schema.__qualname__
|
||||
assert data["enable_tool_calls"] is False
|
||||
assert data["thread_id"] == "thread-456"
|
||||
assert "thread_id" not in data
|
||||
|
||||
def test_from_dict_with_defaults(self) -> None:
|
||||
"""Test from_dict with minimal data."""
|
||||
data = {"message": "Hello", "thread_id": "thread-from-dict"}
|
||||
data = {"message": "Hello"}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.message == "Hello"
|
||||
assert request.role == Role.USER
|
||||
assert request.enable_tool_calls is True
|
||||
assert request.thread_id == "thread-from-dict"
|
||||
|
||||
def test_from_dict_ignores_thread_id_field(self) -> None:
|
||||
"""Ensure legacy thread_id input does not break RunRequest parsing."""
|
||||
request = RunRequest.from_dict({"message": "Hello", "thread_id": "ignored"})
|
||||
|
||||
assert request.message == "Hello"
|
||||
|
||||
def test_from_dict_with_all_fields(self) -> None:
|
||||
"""Test from_dict with all fields."""
|
||||
@@ -233,7 +234,6 @@ class TestRunRequest:
|
||||
"qualname": ModuleStructuredResponse.__qualname__,
|
||||
},
|
||||
"enable_tool_calls": False,
|
||||
"thread_id": "thread-789",
|
||||
}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
@@ -241,11 +241,10 @@ class TestRunRequest:
|
||||
assert request.role == Role.SYSTEM
|
||||
assert request.response_format is ModuleStructuredResponse
|
||||
assert request.enable_tool_calls is False
|
||||
assert request.thread_id == "thread-789"
|
||||
|
||||
def test_from_dict_with_unknown_role_preserves_value(self) -> None:
|
||||
"""Test from_dict keeps custom roles intact."""
|
||||
data = {"message": "Test", "role": "reviewer", "thread_id": "thread-with-custom-role"}
|
||||
data = {"message": "Test", "role": "reviewer"}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.role.value == "reviewer"
|
||||
@@ -253,18 +252,15 @@ class TestRunRequest:
|
||||
|
||||
def test_from_dict_empty_message(self) -> None:
|
||||
"""Test from_dict with empty message."""
|
||||
data = {"thread_id": "thread-empty"}
|
||||
request = RunRequest.from_dict(data)
|
||||
request = RunRequest.from_dict({})
|
||||
|
||||
assert request.message == ""
|
||||
assert request.role == Role.USER
|
||||
assert request.thread_id == "thread-empty"
|
||||
|
||||
def test_round_trip_dict_conversion(self) -> None:
|
||||
"""Test round-trip to_dict and from_dict."""
|
||||
original = RunRequest(
|
||||
message="Test message",
|
||||
thread_id="thread-123",
|
||||
role=Role.SYSTEM,
|
||||
response_format=ModuleStructuredResponse,
|
||||
enable_tool_calls=False,
|
||||
@@ -277,13 +273,11 @@ class TestRunRequest:
|
||||
assert restored.role == original.role
|
||||
assert restored.response_format is ModuleStructuredResponse
|
||||
assert restored.enable_tool_calls == original.enable_tool_calls
|
||||
assert restored.thread_id == original.thread_id
|
||||
|
||||
def test_round_trip_with_pydantic_response_format(self) -> None:
|
||||
"""Ensure Pydantic response formats serialize and deserialize properly."""
|
||||
original = RunRequest(
|
||||
message="Structured",
|
||||
thread_id="thread-pydantic",
|
||||
response_format=ModuleStructuredResponse,
|
||||
)
|
||||
|
||||
@@ -298,14 +292,14 @@ class TestRunRequest:
|
||||
|
||||
def test_init_with_correlationId(self) -> None:
|
||||
"""Test RunRequest initialization with correlationId."""
|
||||
request = RunRequest(message="Test message", thread_id="thread-corr-init", correlation_id="corr-123")
|
||||
request = RunRequest(message="Test message", correlation_id="corr-123")
|
||||
|
||||
assert request.message == "Test message"
|
||||
assert request.correlation_id == "corr-123"
|
||||
|
||||
def test_to_dict_with_correlationId(self) -> None:
|
||||
"""Test to_dict includes correlationId."""
|
||||
request = RunRequest(message="Test", thread_id="thread-corr-to-dict", correlation_id="corr-456")
|
||||
request = RunRequest(message="Test", correlation_id="corr-456")
|
||||
data = request.to_dict()
|
||||
|
||||
assert data["message"] == "Test"
|
||||
@@ -313,18 +307,16 @@ class TestRunRequest:
|
||||
|
||||
def test_from_dict_with_correlationId(self) -> None:
|
||||
"""Test from_dict with correlationId."""
|
||||
data = {"message": "Test", "correlationId": "corr-789", "thread_id": "thread-corr-from-dict"}
|
||||
data = {"message": "Test", "correlationId": "corr-789"}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.message == "Test"
|
||||
assert request.correlation_id == "corr-789"
|
||||
assert request.thread_id == "thread-corr-from-dict"
|
||||
|
||||
def test_round_trip_with_correlationId(self) -> None:
|
||||
"""Test round-trip to_dict and from_dict with correlationId."""
|
||||
original = RunRequest(
|
||||
message="Test message",
|
||||
thread_id="thread-123",
|
||||
role=Role.SYSTEM,
|
||||
correlation_id="corr-123",
|
||||
)
|
||||
@@ -335,13 +327,11 @@ class TestRunRequest:
|
||||
assert restored.message == original.message
|
||||
assert restored.role == original.role
|
||||
assert restored.correlation_id == original.correlation_id
|
||||
assert restored.thread_id == original.thread_id
|
||||
|
||||
def test_init_with_orchestration_id(self) -> None:
|
||||
"""Test RunRequest initialization with orchestration_id."""
|
||||
request = RunRequest(
|
||||
message="Test message",
|
||||
thread_id="thread-orch-init",
|
||||
orchestration_id="orch-123",
|
||||
)
|
||||
|
||||
@@ -352,7 +342,6 @@ class TestRunRequest:
|
||||
"""Test to_dict includes orchestrationId."""
|
||||
request = RunRequest(
|
||||
message="Test",
|
||||
thread_id="thread-orch-to-dict",
|
||||
orchestration_id="orch-456",
|
||||
)
|
||||
data = request.to_dict()
|
||||
@@ -364,7 +353,6 @@ class TestRunRequest:
|
||||
"""Test to_dict excludes orchestrationId when not set."""
|
||||
request = RunRequest(
|
||||
message="Test",
|
||||
thread_id="thread-orch-none",
|
||||
)
|
||||
data = request.to_dict()
|
||||
|
||||
@@ -375,19 +363,16 @@ class TestRunRequest:
|
||||
data = {
|
||||
"message": "Test",
|
||||
"orchestrationId": "orch-789",
|
||||
"thread_id": "thread-orch-from-dict",
|
||||
}
|
||||
request = RunRequest.from_dict(data)
|
||||
|
||||
assert request.message == "Test"
|
||||
assert request.orchestration_id == "orch-789"
|
||||
assert request.thread_id == "thread-orch-from-dict"
|
||||
|
||||
def test_round_trip_with_orchestration_id(self) -> None:
|
||||
"""Test round-trip to_dict and from_dict with orchestration_id."""
|
||||
original = RunRequest(
|
||||
message="Test message",
|
||||
thread_id="thread-123",
|
||||
role=Role.SYSTEM,
|
||||
correlation_id="corr-123",
|
||||
orchestration_id="orch-123",
|
||||
@@ -400,20 +385,17 @@ class TestRunRequest:
|
||||
assert restored.role == original.role
|
||||
assert restored.correlation_id == original.correlation_id
|
||||
assert restored.orchestration_id == original.orchestration_id
|
||||
assert restored.thread_id == original.thread_id
|
||||
|
||||
|
||||
class TestModelIntegration:
|
||||
"""Test suite for integration between models."""
|
||||
|
||||
def test_run_request_with_session_id(self) -> None:
|
||||
"""Test using RunRequest with AgentSessionId."""
|
||||
def test_run_request_with_session_id_string(self) -> None:
|
||||
"""AgentSessionId string can still be used by callers, but is not stored on RunRequest."""
|
||||
session_id = AgentSessionId.with_random_key("AgentEntity")
|
||||
request = RunRequest(message="Test message", thread_id=str(session_id))
|
||||
session_id_str = str(session_id)
|
||||
|
||||
assert request.thread_id is not None
|
||||
assert request.thread_id == str(session_id)
|
||||
assert request.thread_id.startswith("@AgentEntity@")
|
||||
assert session_id_str.startswith("@AgentEntity@")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -300,8 +300,7 @@ class TestDurableAIAgent:
|
||||
assert request["enable_tool_calls"] is True
|
||||
assert "correlationId" in request
|
||||
assert request["correlationId"] == "correlation-guid"
|
||||
assert "thread_id" in request
|
||||
assert request["thread_id"] == "thread-guid"
|
||||
assert "thread_id" not in request
|
||||
# Verify orchestration ID is set from context.instance_id
|
||||
assert "orchestrationId" in request
|
||||
assert request["orchestrationId"] == "test-instance-001"
|
||||
|
||||
Reference in New Issue
Block a user