mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: [Breaking] Python: Respond with AgentRunResponse with serialized structured output (#2285)
* Respond with AgentRunResponse * Fix response_Format type * Address comments * Fix tests * Fix log * Addressed comments * Code cleanup * Use AgentTask vs Generator * Address comments * use lazy logging * fix mypy errors
This commit is contained in:
Laveesh Rohra
committed by
GitHub
Unverified
GitHub
parent
0f2c5e6cb8
commit
306c81aef8
@@ -10,7 +10,7 @@ from unittest.mock import ANY, AsyncMock, Mock, patch
|
||||
import azure.durable_functions as df
|
||||
import azure.functions as func
|
||||
import pytest
|
||||
from agent_framework import AgentRunResponse, ChatMessage
|
||||
from agent_framework import AgentRunResponse, ChatMessage, ErrorContent
|
||||
|
||||
from agent_framework_azurefunctions import AgentFunctionApp
|
||||
from agent_framework_azurefunctions._app import WAIT_FOR_RESPONSE_FIELD, WAIT_FOR_RESPONSE_HEADER
|
||||
@@ -343,10 +343,8 @@ class TestAgentEntityOperations:
|
||||
{"message": "Test message", "thread_id": "test-conv-123", "correlationId": "corr-app-entity-1"},
|
||||
)
|
||||
|
||||
assert result["status"] == "success"
|
||||
assert result["response"] == "Test response"
|
||||
assert result["message"] == "Test message"
|
||||
assert result["thread_id"] == "test-conv-123"
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert result.text == "Test response"
|
||||
assert entity.state.message_count == 2
|
||||
|
||||
async def test_entity_stores_conversation_history(self) -> None:
|
||||
@@ -591,10 +589,12 @@ class TestErrorHandling:
|
||||
mock_context, {"message": "Test message", "thread_id": "conv-1", "correlationId": "corr-app-error-1"}
|
||||
)
|
||||
|
||||
assert result["status"] == "error"
|
||||
assert "error" in result
|
||||
assert "Agent error" in result["error"]
|
||||
assert result["error_type"] == "Exception"
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert len(result.messages) == 1
|
||||
content = result.messages[0].contents[0]
|
||||
assert isinstance(content, ErrorContent)
|
||||
assert "Agent error" in (content.message or "")
|
||||
assert content.error_code == "Exception"
|
||||
|
||||
def test_entity_function_handles_exception(self) -> None:
|
||||
"""Test that the entity function handles exceptions gracefully."""
|
||||
|
||||
@@ -12,7 +12,7 @@ from typing import Any, TypeVar
|
||||
from unittest.mock import AsyncMock, Mock, patch
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentRunResponse, AgentRunResponseUpdate, ChatMessage, Role
|
||||
from agent_framework import AgentRunResponse, AgentRunResponseUpdate, ChatMessage, ErrorContent, Role
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agent_framework_azurefunctions._durable_agent_state import (
|
||||
@@ -133,10 +133,8 @@ class TestAgentEntityRunAgent:
|
||||
assert getattr(sent_message.role, "value", sent_message.role) == "user"
|
||||
|
||||
# Verify result
|
||||
assert result["status"] == "success"
|
||||
assert result["response"] == "Test response"
|
||||
assert result["message"] == "Test message"
|
||||
assert result["thread_id"] == "conv-123"
|
||||
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."""
|
||||
@@ -168,8 +166,8 @@ class TestAgentEntityRunAgent:
|
||||
},
|
||||
)
|
||||
|
||||
assert result["status"] == "success"
|
||||
assert "Hello" in result.get("response", "")
|
||||
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()
|
||||
@@ -215,8 +213,8 @@ class TestAgentEntityRunAgent:
|
||||
},
|
||||
)
|
||||
|
||||
assert result["status"] == "success"
|
||||
assert result.get("response") == "Final response"
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert result.text == "Final response"
|
||||
assert callback.stream_mock.await_count == 0
|
||||
assert callback.response_mock.await_count == 1
|
||||
|
||||
@@ -294,44 +292,6 @@ class TestAgentEntityRunAgent:
|
||||
mock_context, {"message": "Message", "thread_id": None, "correlationId": "corr-entity-5"}
|
||||
)
|
||||
|
||||
async def test_run_agent_handles_response_without_text_attribute(self) -> None:
|
||||
"""Test that run_agent handles responses without a text attribute."""
|
||||
mock_agent = Mock()
|
||||
|
||||
class NoTextResponse(AgentRunResponse):
|
||||
@property
|
||||
def text(self) -> str: # type: ignore[override]
|
||||
raise AttributeError("text attribute missing")
|
||||
|
||||
mock_response = NoTextResponse(messages=[ChatMessage(role="assistant", text="ignored")])
|
||||
mock_agent.run = AsyncMock(return_value=mock_response)
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
result = await entity.run_agent(
|
||||
mock_context, {"message": "Message", "thread_id": "conv-1", "correlationId": "corr-entity-6"}
|
||||
)
|
||||
|
||||
# Should handle gracefully
|
||||
assert result["status"] == "success"
|
||||
assert result["response"] == "Error extracting response"
|
||||
|
||||
async def test_run_agent_handles_none_response_text(self) -> None:
|
||||
"""Test that run_agent handles responses with None text."""
|
||||
mock_agent = Mock()
|
||||
mock_agent.run = AsyncMock(return_value=_agent_response(None))
|
||||
|
||||
entity = AgentEntity(mock_agent)
|
||||
mock_context = Mock()
|
||||
|
||||
result = await entity.run_agent(
|
||||
mock_context, {"message": "Message", "thread_id": "conv-1", "correlationId": "corr-entity-7"}
|
||||
)
|
||||
|
||||
assert result["status"] == "success"
|
||||
assert result["response"] == "No response"
|
||||
|
||||
async def test_run_agent_multiple_conversations(self) -> None:
|
||||
"""Test that run_agent maintains history across multiple messages."""
|
||||
mock_agent = Mock()
|
||||
@@ -621,10 +581,12 @@ class TestErrorHandling:
|
||||
mock_context, {"message": "Message", "thread_id": "conv-1", "correlationId": "corr-entity-error-1"}
|
||||
)
|
||||
|
||||
assert result["status"] == "error"
|
||||
assert "error" in result
|
||||
assert "Agent failed" in result["error"]
|
||||
assert result["error_type"] == "Exception"
|
||||
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."""
|
||||
@@ -638,9 +600,12 @@ class TestErrorHandling:
|
||||
mock_context, {"message": "Message", "thread_id": "conv-1", "correlationId": "corr-entity-error-2"}
|
||||
)
|
||||
|
||||
assert result["status"] == "error"
|
||||
assert result["error_type"] == "ValueError"
|
||||
assert "Invalid input" in result["error"]
|
||||
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."""
|
||||
@@ -654,8 +619,11 @@ class TestErrorHandling:
|
||||
mock_context, {"message": "Message", "thread_id": "conv-1", "correlationId": "corr-entity-error-3"}
|
||||
)
|
||||
|
||||
assert result["status"] == "error"
|
||||
assert result["error_type"] == "TimeoutError"
|
||||
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."""
|
||||
@@ -690,9 +658,10 @@ class TestErrorHandling:
|
||||
)
|
||||
|
||||
# Even on error, message info should be preserved
|
||||
assert result["message"] == "Test message"
|
||||
assert result["thread_id"] == "conv-123"
|
||||
assert result["status"] == "error"
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
assert len(result.messages) == 1
|
||||
content = result.messages[0].contents[0]
|
||||
assert isinstance(content, ErrorContent)
|
||||
|
||||
|
||||
class TestConversationHistory:
|
||||
@@ -800,10 +769,8 @@ class TestRunRequestSupport:
|
||||
|
||||
result = await entity.run_agent(mock_context, request)
|
||||
|
||||
assert result["status"] == "success"
|
||||
assert result["response"] == "Response"
|
||||
assert result["message"] == "Test message"
|
||||
assert result["thread_id"] == "conv-123"
|
||||
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."""
|
||||
@@ -823,9 +790,8 @@ class TestRunRequestSupport:
|
||||
|
||||
result = await entity.run_agent(mock_context, request_dict)
|
||||
|
||||
assert result["status"] == "success"
|
||||
assert result["message"] == "Test message"
|
||||
assert result["thread_id"] == "conv-456"
|
||||
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."""
|
||||
@@ -879,10 +845,9 @@ class TestRunRequestSupport:
|
||||
|
||||
result = await entity.run_agent(mock_context, request)
|
||||
|
||||
assert result["status"] == "success"
|
||||
# Should have structured_response
|
||||
if "structured_response" in result:
|
||||
assert result["structured_response"]["answer"] == 42
|
||||
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."""
|
||||
@@ -898,7 +863,7 @@ class TestRunRequestSupport:
|
||||
|
||||
result = await entity.run_agent(mock_context, request)
|
||||
|
||||
assert result["status"] == "success"
|
||||
assert isinstance(result, AgentRunResponse)
|
||||
# Agent should have been called (tool disabling is framework-dependent)
|
||||
mock_agent.run.assert_called_once()
|
||||
|
||||
@@ -925,8 +890,24 @@ class TestRunRequestSupport:
|
||||
# Verify result was set
|
||||
assert mock_context.set_result.called
|
||||
result = mock_context.set_result.call_args[0][0]
|
||||
assert result["status"] == "success"
|
||||
assert result["message"] == "Test message"
|
||||
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}"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -7,7 +7,7 @@ import pytest
|
||||
from agent_framework import Role
|
||||
from pydantic import BaseModel
|
||||
|
||||
from agent_framework_azurefunctions._models import AgentResponse, AgentSessionId, RunRequest
|
||||
from agent_framework_azurefunctions._models import AgentSessionId, RunRequest
|
||||
|
||||
|
||||
class ModuleStructuredResponse(BaseModel):
|
||||
@@ -337,107 +337,6 @@ class TestRunRequest:
|
||||
assert restored.thread_id == original.thread_id
|
||||
|
||||
|
||||
class TestAgentResponse:
|
||||
"""Test suite for AgentResponse."""
|
||||
|
||||
def test_init_with_required_fields(self) -> None:
|
||||
"""Test AgentResponse initialization with required fields."""
|
||||
response = AgentResponse(
|
||||
response="Test response", message="Test message", thread_id="thread-123", status="success"
|
||||
)
|
||||
|
||||
assert response.response == "Test response"
|
||||
assert response.message == "Test message"
|
||||
assert response.thread_id == "thread-123"
|
||||
assert response.status == "success"
|
||||
assert response.message_count == 0
|
||||
assert response.error is None
|
||||
assert response.error_type is None
|
||||
assert response.structured_response is None
|
||||
|
||||
def test_init_with_all_fields(self) -> None:
|
||||
"""Test AgentResponse initialization with all fields."""
|
||||
structured = {"answer": "42"}
|
||||
response = AgentResponse(
|
||||
response=None,
|
||||
message="What is the answer?",
|
||||
thread_id="thread-456",
|
||||
status="success",
|
||||
message_count=5,
|
||||
error=None,
|
||||
error_type=None,
|
||||
structured_response=structured,
|
||||
)
|
||||
|
||||
assert response.response is None
|
||||
assert response.structured_response == structured
|
||||
assert response.message_count == 5
|
||||
|
||||
def test_to_dict_with_text_response(self) -> None:
|
||||
"""Test to_dict with text response."""
|
||||
response = AgentResponse(
|
||||
response="Text response", message="Message", thread_id="thread-1", status="success", message_count=3
|
||||
)
|
||||
data = response.to_dict()
|
||||
|
||||
assert data["response"] == "Text response"
|
||||
assert data["message"] == "Message"
|
||||
assert data["thread_id"] == "thread-1"
|
||||
assert data["status"] == "success"
|
||||
assert data["message_count"] == 3
|
||||
assert "structured_response" not in data
|
||||
assert "error" not in data
|
||||
assert "error_type" not in data
|
||||
|
||||
def test_to_dict_with_structured_response(self) -> None:
|
||||
"""Test to_dict with structured response."""
|
||||
structured = {"answer": 42, "confidence": 0.95}
|
||||
response = AgentResponse(
|
||||
response=None,
|
||||
message="Question",
|
||||
thread_id="thread-2",
|
||||
status="success",
|
||||
structured_response=structured,
|
||||
)
|
||||
data = response.to_dict()
|
||||
|
||||
assert data["structured_response"] == structured
|
||||
assert "response" not in data
|
||||
|
||||
def test_to_dict_with_error(self) -> None:
|
||||
"""Test to_dict with error."""
|
||||
response = AgentResponse(
|
||||
response=None,
|
||||
message="Failed message",
|
||||
thread_id="thread-3",
|
||||
status="error",
|
||||
error="Something went wrong",
|
||||
error_type="ValueError",
|
||||
)
|
||||
data = response.to_dict()
|
||||
|
||||
assert data["status"] == "error"
|
||||
assert data["error"] == "Something went wrong"
|
||||
assert data["error_type"] == "ValueError"
|
||||
|
||||
def test_to_dict_prefers_structured_over_text(self) -> None:
|
||||
"""Test to_dict prefers structured_response over response."""
|
||||
structured = {"result": "structured"}
|
||||
response = AgentResponse(
|
||||
response="Text response",
|
||||
message="Message",
|
||||
thread_id="thread-4",
|
||||
status="success",
|
||||
structured_response=structured,
|
||||
)
|
||||
data = response.to_dict()
|
||||
|
||||
assert "structured_response" in data
|
||||
assert data["structured_response"] == structured
|
||||
# Text response should not be included when structured is present
|
||||
assert "response" not in data
|
||||
|
||||
|
||||
class TestModelIntegration:
|
||||
"""Test suite for integration between models."""
|
||||
|
||||
@@ -450,21 +349,6 @@ class TestModelIntegration:
|
||||
assert request.thread_id == str(session_id)
|
||||
assert request.thread_id.startswith("@AgentEntity@")
|
||||
|
||||
def test_response_from_run_request(self) -> None:
|
||||
"""Test creating AgentResponse from RunRequest."""
|
||||
request = RunRequest(message="What is 2+2?", thread_id="thread-123", role=Role.USER)
|
||||
|
||||
response = AgentResponse(
|
||||
response="4",
|
||||
message=request.message,
|
||||
thread_id=request.thread_id,
|
||||
status="success",
|
||||
message_count=1,
|
||||
)
|
||||
|
||||
assert response.message == request.message
|
||||
assert response.thread_id == request.thread_id
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
pytest.main([__file__, "-v", "--tb=short"])
|
||||
|
||||
@@ -6,10 +6,12 @@ from typing import Any
|
||||
from unittest.mock import Mock
|
||||
|
||||
import pytest
|
||||
from agent_framework import AgentThread
|
||||
from agent_framework import AgentRunResponse, AgentThread, ChatMessage
|
||||
from azure.durable_functions.models.Task import TaskBase, TaskState
|
||||
|
||||
from agent_framework_azurefunctions import AgentFunctionApp, DurableAIAgent
|
||||
from agent_framework_azurefunctions._models import AgentSessionId, DurableAgentThread
|
||||
from agent_framework_azurefunctions._orchestration import AgentTask
|
||||
|
||||
|
||||
def _app_with_registered_agents(*agent_names: str) -> AgentFunctionApp:
|
||||
@@ -21,6 +23,169 @@ def _app_with_registered_agents(*agent_names: str) -> AgentFunctionApp:
|
||||
return app
|
||||
|
||||
|
||||
class _FakeTask(TaskBase):
|
||||
"""Concrete TaskBase for testing AgentTask wiring."""
|
||||
|
||||
def __init__(self, task_id: int = 1):
|
||||
super().__init__(task_id, [])
|
||||
self._set_is_scheduled(False)
|
||||
self.action_repr = []
|
||||
self.state = TaskState.RUNNING
|
||||
|
||||
|
||||
def _create_entity_task(task_id: int = 1) -> TaskBase:
|
||||
"""Create a minimal TaskBase instance for AgentTask tests."""
|
||||
return _FakeTask(task_id)
|
||||
|
||||
|
||||
class TestAgentResponseHelpers:
|
||||
"""Tests for helper utilities that prepare AgentRunResponse values."""
|
||||
|
||||
@staticmethod
|
||||
def _create_agent_task() -> AgentTask:
|
||||
entity_task = _create_entity_task()
|
||||
return AgentTask(entity_task, None, "correlation-id")
|
||||
|
||||
def test_load_agent_response_from_instance(self) -> None:
|
||||
task = self._create_agent_task()
|
||||
response = AgentRunResponse(messages=[ChatMessage(role="assistant", text='{"foo": "bar"}')])
|
||||
|
||||
loaded = task._load_agent_response(response)
|
||||
|
||||
assert loaded is response
|
||||
assert loaded.value is None
|
||||
|
||||
def test_load_agent_response_from_serialized(self) -> None:
|
||||
task = self._create_agent_task()
|
||||
serialized = AgentRunResponse(messages=[ChatMessage(role="assistant", text="structured")]).to_dict()
|
||||
serialized["value"] = {"answer": 42}
|
||||
|
||||
loaded = task._load_agent_response(serialized)
|
||||
|
||||
assert loaded is not None
|
||||
assert loaded.value == {"answer": 42}
|
||||
loaded_dict = loaded.to_dict()
|
||||
assert loaded_dict["type"] == "agent_run_response"
|
||||
|
||||
def test_load_agent_response_rejects_none(self) -> None:
|
||||
task = self._create_agent_task()
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
task._load_agent_response(None)
|
||||
|
||||
def test_load_agent_response_rejects_unsupported_type(self) -> None:
|
||||
task = self._create_agent_task()
|
||||
|
||||
with pytest.raises(TypeError, match="Unsupported type"):
|
||||
task._load_agent_response(["invalid", "list"]) # type: ignore[arg-type]
|
||||
|
||||
def test_try_set_value_success(self) -> None:
|
||||
"""Test try_set_value correctly processes successful task completion."""
|
||||
entity_task = _create_entity_task()
|
||||
task = AgentTask(entity_task, None, "correlation-id")
|
||||
|
||||
# Simulate successful entity task completion
|
||||
entity_task.state = TaskState.SUCCEEDED
|
||||
entity_task.result = AgentRunResponse(messages=[ChatMessage(role="assistant", text="Test response")]).to_dict()
|
||||
|
||||
# Clear pending_tasks to simulate that parent has processed the child
|
||||
task.pending_tasks.clear()
|
||||
|
||||
# Call try_set_value
|
||||
task.try_set_value(entity_task)
|
||||
|
||||
# Verify task completed successfully with AgentRunResponse
|
||||
assert task.state == TaskState.SUCCEEDED
|
||||
assert isinstance(task.result, AgentRunResponse)
|
||||
assert task.result.text == "Test response"
|
||||
|
||||
def test_try_set_value_failure(self) -> None:
|
||||
"""Test try_set_value correctly handles failed task completion."""
|
||||
entity_task = _create_entity_task()
|
||||
task = AgentTask(entity_task, None, "correlation-id")
|
||||
|
||||
# Simulate failed entity task
|
||||
entity_task.state = TaskState.FAILED
|
||||
entity_task.result = Exception("Entity call failed")
|
||||
|
||||
# Call try_set_value
|
||||
task.try_set_value(entity_task)
|
||||
|
||||
# Verify task failed with the error
|
||||
assert task.state == TaskState.FAILED
|
||||
assert isinstance(task.result, Exception)
|
||||
assert str(task.result) == "Entity call failed"
|
||||
|
||||
def test_try_set_value_with_response_format(self) -> None:
|
||||
"""Test try_set_value parses structured output when response_format is provided."""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class TestSchema(BaseModel):
|
||||
answer: str
|
||||
|
||||
entity_task = _create_entity_task()
|
||||
task = AgentTask(entity_task, TestSchema, "correlation-id")
|
||||
|
||||
# Simulate successful entity task with JSON response
|
||||
entity_task.state = TaskState.SUCCEEDED
|
||||
entity_task.result = AgentRunResponse(
|
||||
messages=[ChatMessage(role="assistant", text='{"answer": "42"}')]
|
||||
).to_dict()
|
||||
|
||||
# Clear pending_tasks to simulate that parent has processed the child
|
||||
task.pending_tasks.clear()
|
||||
|
||||
# Call try_set_value
|
||||
task.try_set_value(entity_task)
|
||||
|
||||
# Verify task completed and value was parsed
|
||||
assert task.state == TaskState.SUCCEEDED
|
||||
assert isinstance(task.result, AgentRunResponse)
|
||||
assert isinstance(task.result.value, TestSchema)
|
||||
assert task.result.value.answer == "42"
|
||||
|
||||
def test_ensure_response_format_parses_value(self) -> None:
|
||||
"""Test _ensure_response_format correctly parses response value."""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class SampleSchema(BaseModel):
|
||||
name: str
|
||||
|
||||
task = self._create_agent_task()
|
||||
response = AgentRunResponse(messages=[ChatMessage(role="assistant", text='{"name": "test"}')])
|
||||
|
||||
# Value should be None initially
|
||||
assert response.value is None
|
||||
|
||||
# Parse the value
|
||||
task._ensure_response_format(SampleSchema, "test-correlation", response)
|
||||
|
||||
# Value should now be parsed
|
||||
assert isinstance(response.value, SampleSchema)
|
||||
assert response.value.name == "test"
|
||||
|
||||
def test_ensure_response_format_skips_if_already_parsed(self) -> None:
|
||||
"""Test _ensure_response_format does not re-parse if value already matches format."""
|
||||
from pydantic import BaseModel
|
||||
|
||||
class SampleSchema(BaseModel):
|
||||
name: str
|
||||
|
||||
task = self._create_agent_task()
|
||||
existing_value = SampleSchema(name="existing")
|
||||
response = AgentRunResponse(
|
||||
messages=[ChatMessage(role="assistant", text='{"name": "new"}')],
|
||||
value=existing_value,
|
||||
)
|
||||
|
||||
# Call _ensure_response_format
|
||||
task._ensure_response_format(SampleSchema, "test-correlation", response)
|
||||
|
||||
# Value should remain unchanged (not re-parsed)
|
||||
assert response.value is existing_value
|
||||
assert response.value.name == "existing"
|
||||
|
||||
|
||||
class TestDurableAIAgent:
|
||||
"""Test suite for DurableAIAgent wrapper."""
|
||||
|
||||
@@ -111,22 +276,19 @@ class TestDurableAIAgent:
|
||||
mock_context.instance_id = "test-instance-001"
|
||||
mock_context.new_uuid = Mock(side_effect=["thread-guid", "correlation-guid"])
|
||||
|
||||
# Mock call_entity to return a Task-like object
|
||||
mock_task = Mock()
|
||||
mock_task._is_scheduled = False # Task attribute that orchestration checks
|
||||
|
||||
mock_context.call_entity = Mock(return_value=mock_task)
|
||||
entity_task = _create_entity_task()
|
||||
mock_context.call_entity = Mock(return_value=entity_task)
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
# Create thread
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
# Call run() - it should return the Task directly
|
||||
# Call run() - returns AgentTask directly
|
||||
task = agent.run(messages="Test message", thread=thread, enable_tool_calls=True)
|
||||
|
||||
# Verify run() returns the Task from call_entity
|
||||
assert task == mock_task
|
||||
assert isinstance(task, AgentTask)
|
||||
assert task.children[0] == entity_task
|
||||
|
||||
# Verify call_entity was called with correct parameters
|
||||
assert mock_context.call_entity.called
|
||||
@@ -145,19 +307,18 @@ class TestDurableAIAgent:
|
||||
"""Test that run() works without explicit thread (creates unique session key)."""
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-002"
|
||||
# Two calls to new_uuid: one for session_key, one for correlationId
|
||||
mock_context.new_uuid = Mock(side_effect=["auto-generated-guid", "correlation-guid"])
|
||||
|
||||
mock_task = Mock()
|
||||
mock_task._is_scheduled = False
|
||||
mock_context.call_entity = Mock(return_value=mock_task)
|
||||
entity_task = _create_entity_task()
|
||||
mock_context.call_entity = Mock(return_value=entity_task)
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
# Call without thread
|
||||
task = agent.run(messages="Test message")
|
||||
|
||||
assert task == mock_task
|
||||
assert isinstance(task, AgentTask)
|
||||
assert task.children[0] == entity_task
|
||||
|
||||
# Verify the entity ID uses the auto-generated GUID with dafx- prefix
|
||||
call_args = mock_context.call_entity.call_args
|
||||
@@ -172,9 +333,8 @@ class TestDurableAIAgent:
|
||||
mock_context = Mock()
|
||||
mock_context.instance_id = "test-instance-003"
|
||||
|
||||
mock_task = Mock()
|
||||
mock_task._is_scheduled = False
|
||||
mock_context.call_entity = Mock(return_value=mock_task)
|
||||
entity_task = _create_entity_task()
|
||||
mock_context.call_entity = Mock(return_value=entity_task)
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
|
||||
@@ -188,7 +348,8 @@ class TestDurableAIAgent:
|
||||
|
||||
task = agent.run(messages="Test message", thread=thread, response_format=SampleSchema)
|
||||
|
||||
assert task == mock_task
|
||||
assert isinstance(task, AgentTask)
|
||||
assert task.children[0] == entity_task
|
||||
|
||||
# Verify schema was passed in the call_entity arguments
|
||||
call_args = mock_context.call_entity.call_args
|
||||
@@ -221,8 +382,8 @@ class TestDurableAIAgent:
|
||||
|
||||
mock_context = Mock()
|
||||
mock_context.new_uuid = Mock(side_effect=["thread-guid", "correlation-guid"])
|
||||
mock_task = Mock()
|
||||
mock_context.call_entity = Mock(return_value=mock_task)
|
||||
entity_task = _create_entity_task()
|
||||
mock_context.call_entity = Mock(return_value=entity_task)
|
||||
|
||||
agent = DurableAIAgent(mock_context, "TestAgent")
|
||||
thread = agent.get_new_thread()
|
||||
@@ -231,7 +392,8 @@ class TestDurableAIAgent:
|
||||
msg = ChatMessage(role="user", text="Hello")
|
||||
task = agent.run(messages=msg, thread=thread)
|
||||
|
||||
assert task == mock_task
|
||||
assert isinstance(task, AgentTask)
|
||||
assert task.children[0] == entity_task
|
||||
|
||||
# Verify message was converted to string
|
||||
call_args = mock_context.call_entity.call_args
|
||||
@@ -255,7 +417,7 @@ class TestDurableAIAgent:
|
||||
|
||||
mock_context = Mock()
|
||||
mock_context.new_uuid = Mock(return_value="test-guid-789")
|
||||
mock_context.call_entity = Mock(return_value=Mock())
|
||||
mock_context.call_entity = Mock(return_value=_create_entity_task())
|
||||
|
||||
agent = DurableAIAgent(mock_context, "WriterAgent")
|
||||
thread = agent.get_new_thread()
|
||||
@@ -314,13 +476,9 @@ class TestOrchestrationIntegration:
|
||||
# Track entity calls
|
||||
entity_calls: list[dict[str, Any]] = []
|
||||
|
||||
def mock_call_entity_side_effect(entity_id: Any, operation: str, input_data: dict[str, Any]) -> Mock:
|
||||
def mock_call_entity_side_effect(entity_id: Any, operation: str, input_data: dict[str, Any]) -> TaskBase:
|
||||
entity_calls.append({"entity_id": str(entity_id), "operation": operation, "input": input_data})
|
||||
|
||||
# Return a mock Task
|
||||
mock_task = Mock()
|
||||
mock_task._is_scheduled = False
|
||||
return mock_task
|
||||
return _create_entity_task()
|
||||
|
||||
mock_context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
|
||||
|
||||
@@ -330,13 +488,13 @@ class TestOrchestrationIntegration:
|
||||
# Create thread
|
||||
thread = agent.get_new_thread()
|
||||
|
||||
# First call - returns Task
|
||||
# First call - returns AgentTask
|
||||
task1 = agent.run("Write something", thread=thread)
|
||||
assert hasattr(task1, "_is_scheduled")
|
||||
assert isinstance(task1, AgentTask)
|
||||
|
||||
# Second call - returns Task
|
||||
# Second call - returns AgentTask
|
||||
task2 = agent.run("Improve: something", thread=thread)
|
||||
assert hasattr(task2, "_is_scheduled")
|
||||
assert isinstance(task2, AgentTask)
|
||||
|
||||
# Verify both calls used the same entity (same session key)
|
||||
assert len(entity_calls) == 2
|
||||
@@ -356,11 +514,9 @@ class TestOrchestrationIntegration:
|
||||
|
||||
entity_calls: list[str] = []
|
||||
|
||||
def mock_call_entity_side_effect(entity_id: Any, operation: str, input_data: dict[str, Any]) -> Mock:
|
||||
def mock_call_entity_side_effect(entity_id: Any, operation: str, input_data: dict[str, Any]) -> TaskBase:
|
||||
entity_calls.append(str(entity_id))
|
||||
mock_task = Mock()
|
||||
mock_task._is_scheduled = False
|
||||
return mock_task
|
||||
return _create_entity_task()
|
||||
|
||||
mock_context.call_entity = Mock(side_effect=mock_call_entity_side_effect)
|
||||
|
||||
@@ -371,12 +527,12 @@ class TestOrchestrationIntegration:
|
||||
writer_thread = writer.get_new_thread()
|
||||
editor_thread = editor.get_new_thread()
|
||||
|
||||
# Call both agents - returns Tasks
|
||||
# Call both agents - returns AgentTasks
|
||||
writer_task = writer.run("Write", thread=writer_thread)
|
||||
editor_task = editor.run("Edit", thread=editor_thread)
|
||||
|
||||
assert hasattr(writer_task, "_is_scheduled")
|
||||
assert hasattr(editor_task, "_is_scheduled")
|
||||
assert isinstance(writer_task, AgentTask)
|
||||
assert isinstance(editor_task, AgentTask)
|
||||
|
||||
# Verify different entity IDs were used
|
||||
assert len(entity_calls) == 2
|
||||
|
||||
Reference in New Issue
Block a user