mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
[BREAKING] Python: Replace RequestInfoExecutor with request_info API and @response_handler (#1466)
* Prototype: Add request_info API and @response_handler * Add original_request as a parameter to the response handler * Prototype: request interception in sub workflows * Prototype: request interception in sub workflows 2 * WIP: Make checkpointing work * checkpointing with sub workflow * Fix function executor * Allow sub-workflow to output directly * Remove ReqeustInfoExecutor and related classes; Debugging checkpoint_with_human_in_the_loop * Fix Handoff and sample * fix pending requests in checkpoint * Fix unit tests * Fix formatting * Resolve comments * Address comment * Add checkpoint tests * Add tests * misc * fix mypy * fix mypy * Use request type as part of the key * Log warning if there is not response handler for a request * Update Internal edge group comments * REcord message type in executor processing span * Update sample * Improve tests
This commit is contained in:
Tao Chen
committed by
GitHub
Unverified
GitHub
parent
f6eadd412e
commit
943d92674e
@@ -20,6 +20,7 @@ def test_workflow_checkpoint_default_values():
|
||||
assert checkpoint.timestamp != ""
|
||||
assert checkpoint.messages == {}
|
||||
assert checkpoint.shared_state == {}
|
||||
assert checkpoint.pending_request_info_events == {}
|
||||
assert checkpoint.iteration_count == 0
|
||||
assert checkpoint.metadata == {}
|
||||
assert checkpoint.version == "1.0"
|
||||
@@ -32,6 +33,7 @@ def test_workflow_checkpoint_custom_values():
|
||||
workflow_id="test-workflow-456",
|
||||
timestamp=custom_timestamp,
|
||||
messages={"executor1": [{"data": "test"}]},
|
||||
pending_request_info_events={"req123": {"data": "test"}},
|
||||
shared_state={"key": "value"},
|
||||
iteration_count=5,
|
||||
metadata={"test": True},
|
||||
@@ -43,6 +45,7 @@ def test_workflow_checkpoint_custom_values():
|
||||
assert checkpoint.timestamp == custom_timestamp
|
||||
assert checkpoint.messages == {"executor1": [{"data": "test"}]}
|
||||
assert checkpoint.shared_state == {"key": "value"}
|
||||
assert checkpoint.pending_request_info_events == {"req123": {"data": "test"}}
|
||||
assert checkpoint.iteration_count == 5
|
||||
assert checkpoint.metadata == {"test": True}
|
||||
assert checkpoint.version == "2.0"
|
||||
@@ -50,7 +53,11 @@ def test_workflow_checkpoint_custom_values():
|
||||
|
||||
async def test_memory_checkpoint_storage_save_and_load():
|
||||
storage = InMemoryCheckpointStorage()
|
||||
checkpoint = WorkflowCheckpoint(workflow_id="test-workflow", messages={"executor1": [{"data": "hello"}]})
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
workflow_id="test-workflow",
|
||||
messages={"executor1": [{"data": "hello"}]},
|
||||
pending_request_info_events={"req123": {"data": "test"}},
|
||||
)
|
||||
|
||||
# Save checkpoint
|
||||
saved_id = await storage.save_checkpoint(checkpoint)
|
||||
@@ -62,6 +69,7 @@ async def test_memory_checkpoint_storage_save_and_load():
|
||||
assert loaded_checkpoint.checkpoint_id == checkpoint.checkpoint_id
|
||||
assert loaded_checkpoint.workflow_id == checkpoint.workflow_id
|
||||
assert loaded_checkpoint.messages == checkpoint.messages
|
||||
assert loaded_checkpoint.pending_request_info_events == checkpoint.pending_request_info_events
|
||||
|
||||
|
||||
async def test_memory_checkpoint_storage_load_nonexistent():
|
||||
@@ -152,6 +160,7 @@ async def test_file_checkpoint_storage_save_and_load():
|
||||
workflow_id="test-workflow",
|
||||
messages={"executor1": [{"data": "hello", "source_id": "test", "target_id": None}]},
|
||||
shared_state={"key": "value"},
|
||||
pending_request_info_events={"req123": {"data": "test"}},
|
||||
)
|
||||
|
||||
# Save checkpoint
|
||||
@@ -169,6 +178,7 @@ async def test_file_checkpoint_storage_save_and_load():
|
||||
assert loaded_checkpoint.workflow_id == checkpoint.workflow_id
|
||||
assert loaded_checkpoint.messages == checkpoint.messages
|
||||
assert loaded_checkpoint.shared_state == checkpoint.shared_state
|
||||
assert loaded_checkpoint.pending_request_info_events == checkpoint.pending_request_info_events
|
||||
|
||||
|
||||
async def test_file_checkpoint_storage_load_nonexistent():
|
||||
@@ -284,6 +294,7 @@ async def test_file_checkpoint_storage_json_serialization():
|
||||
workflow_id="complex-workflow",
|
||||
messages={"executor1": [{"data": {"nested": {"value": 42}}, "source_id": "test", "target_id": None}]},
|
||||
shared_state={"list": [1, 2, 3], "dict": {"a": "b", "c": {"d": "e"}}, "bool": True, "null": None},
|
||||
pending_request_info_events={"req123": {"data": "test"}},
|
||||
)
|
||||
|
||||
# Save and load
|
||||
@@ -303,6 +314,7 @@ async def test_file_checkpoint_storage_json_serialization():
|
||||
assert data["shared_state"]["list"] == [1, 2, 3]
|
||||
assert data["shared_state"]["bool"] is True
|
||||
assert data["shared_state"]["null"] is None
|
||||
assert data["pending_request_info_events"]["req123"]["data"] == "test"
|
||||
|
||||
|
||||
def test_checkpoint_storage_protocol_compliance():
|
||||
|
||||
@@ -3,7 +3,6 @@
|
||||
from dataclasses import dataclass # noqa: I001
|
||||
from typing import Any, cast
|
||||
|
||||
from agent_framework._workflows._request_info_executor import RequestInfoMessage, RequestResponse
|
||||
from agent_framework._workflows._checkpoint_encoding import (
|
||||
decode_checkpoint_value,
|
||||
encode_checkpoint_value,
|
||||
@@ -11,30 +10,45 @@ from agent_framework._workflows._checkpoint_encoding import (
|
||||
from agent_framework._workflows._typing_utils import is_instance_of
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class SampleRequest(RequestInfoMessage):
|
||||
@dataclass
|
||||
class SampleRequest:
|
||||
"""Sample request message for testing checkpoint encoding/decoding."""
|
||||
|
||||
request_id: str
|
||||
prompt: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class SampleResponse:
|
||||
"""Sample response message for testing checkpoint encoding/decoding."""
|
||||
|
||||
data: str
|
||||
original_request: SampleRequest
|
||||
request_id: str
|
||||
|
||||
|
||||
def test_decode_dataclass_with_nested_request() -> None:
|
||||
original = RequestResponse[SampleRequest, str](
|
||||
"""Test that dataclass with nested dataclass fields can be encoded and decoded correctly."""
|
||||
original = SampleResponse(
|
||||
data="approve",
|
||||
original_request=SampleRequest(request_id="abc", prompt="prompt"),
|
||||
request_id="abc",
|
||||
)
|
||||
|
||||
encoded = encode_checkpoint_value(original)
|
||||
decoded = cast(RequestResponse[SampleRequest, str], decode_checkpoint_value(encoded))
|
||||
decoded = cast(SampleResponse, decode_checkpoint_value(encoded))
|
||||
|
||||
assert isinstance(decoded, RequestResponse)
|
||||
assert isinstance(decoded, SampleResponse)
|
||||
assert decoded.data == "approve"
|
||||
assert decoded.request_id == "abc"
|
||||
assert isinstance(decoded.original_request, SampleRequest)
|
||||
assert decoded.original_request.prompt == "prompt"
|
||||
assert decoded.original_request.request_id == "abc"
|
||||
|
||||
|
||||
def test_is_instance_of_coerces_request_response_original_request_dict() -> None:
|
||||
response = RequestResponse[SampleRequest, str](
|
||||
def test_is_instance_of_coerces_nested_dataclass_dict() -> None:
|
||||
"""Test that is_instance_of can handle nested structures with dict conversion."""
|
||||
response = SampleResponse(
|
||||
data="approve",
|
||||
original_request=SampleRequest(request_id="req-1", prompt="prompt"),
|
||||
request_id="req-1",
|
||||
@@ -49,5 +63,66 @@ def test_is_instance_of_coerces_request_response_original_request_dict() -> None
|
||||
},
|
||||
)
|
||||
|
||||
assert is_instance_of(response, RequestResponse[SampleRequest, str])
|
||||
assert is_instance_of(response, SampleResponse)
|
||||
assert isinstance(response.original_request, dict)
|
||||
|
||||
# Verify the dict contains expected values
|
||||
dict_request = cast(dict[str, Any], response.original_request)
|
||||
assert dict_request["request_id"] == "req-1"
|
||||
assert dict_request["prompt"] == "prompt"
|
||||
|
||||
|
||||
def test_encode_decode_simple_dataclass() -> None:
|
||||
"""Test encoding and decoding of a simple dataclass."""
|
||||
original = SampleRequest(request_id="test-123", prompt="test prompt")
|
||||
|
||||
encoded = encode_checkpoint_value(original)
|
||||
decoded = cast(SampleRequest, decode_checkpoint_value(encoded))
|
||||
|
||||
assert isinstance(decoded, SampleRequest)
|
||||
assert decoded.request_id == "test-123"
|
||||
assert decoded.prompt == "test prompt"
|
||||
|
||||
|
||||
def test_encode_decode_nested_structures() -> None:
|
||||
"""Test encoding and decoding of complex nested structures."""
|
||||
nested_data = {
|
||||
"requests": [
|
||||
SampleRequest(request_id="req-1", prompt="first prompt"),
|
||||
SampleRequest(request_id="req-2", prompt="second prompt"),
|
||||
],
|
||||
"responses": {
|
||||
"req-1": SampleResponse(
|
||||
data="first response",
|
||||
original_request=SampleRequest(request_id="req-1", prompt="first prompt"),
|
||||
request_id="req-1",
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
encoded = encode_checkpoint_value(nested_data)
|
||||
decoded = decode_checkpoint_value(encoded)
|
||||
|
||||
assert isinstance(decoded, dict)
|
||||
assert "requests" in decoded
|
||||
assert "responses" in decoded
|
||||
|
||||
# Check the requests list
|
||||
requests = cast(list[Any], decoded["requests"])
|
||||
assert isinstance(requests, list)
|
||||
assert len(requests) == 2
|
||||
assert all(isinstance(req, SampleRequest) for req in requests)
|
||||
first_request = cast(SampleRequest, requests[0])
|
||||
second_request = cast(SampleRequest, requests[1])
|
||||
assert first_request.request_id == "req-1"
|
||||
assert second_request.request_id == "req-2"
|
||||
|
||||
# Check the responses dict
|
||||
responses = cast(dict[str, Any], decoded["responses"])
|
||||
assert isinstance(responses, dict)
|
||||
assert "req-1" in responses
|
||||
response = cast(SampleResponse, responses["req-1"])
|
||||
assert isinstance(response, SampleResponse)
|
||||
assert response.data == "first response"
|
||||
assert isinstance(response.original_request, SampleRequest)
|
||||
assert response.original_request.request_id == "req-1"
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework import Executor, WorkflowContext, handler
|
||||
from agent_framework import Executor, Message, WorkflowContext, handler
|
||||
|
||||
|
||||
def test_executor_without_id():
|
||||
@@ -64,9 +64,9 @@ def test_executor_with_valid_handlers():
|
||||
executor = MockExecutorWithValidHandlers(id="test")
|
||||
assert executor.id is not None
|
||||
assert len(executor._handlers) == 2 # type: ignore
|
||||
assert executor.can_handle("text") is True
|
||||
assert executor.can_handle(42) is True
|
||||
assert executor.can_handle(3.14) is False
|
||||
assert executor.can_handle(Message(data="text", source_id="mock")) is True
|
||||
assert executor.can_handle(Message(data=42, source_id="mock")) is True
|
||||
assert executor.can_handle(Message(data=3.14, source_id="mock")) is False
|
||||
|
||||
|
||||
def test_executor_handlers_with_output_types():
|
||||
|
||||
@@ -7,6 +7,7 @@ from typing_extensions import Never
|
||||
|
||||
from agent_framework import (
|
||||
FunctionExecutor,
|
||||
Message,
|
||||
WorkflowBuilder,
|
||||
WorkflowContext,
|
||||
executor,
|
||||
@@ -230,9 +231,9 @@ class TestFunctionExecutor:
|
||||
async def string_processor(text: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(text)
|
||||
|
||||
assert string_processor.can_handle("hello")
|
||||
assert not string_processor.can_handle(123)
|
||||
assert not string_processor.can_handle([])
|
||||
assert string_processor.can_handle(Message(data="hello", source_id="Mock"))
|
||||
assert not string_processor.can_handle(Message(data=123, source_id="Mock"))
|
||||
assert not string_processor.can_handle(Message(data=[], source_id="Mock"))
|
||||
|
||||
def test_duplicate_handler_registration(self):
|
||||
"""Test that registering duplicate handlers raises an error."""
|
||||
@@ -309,9 +310,9 @@ class TestFunctionExecutor:
|
||||
async def int_processor(value: int):
|
||||
return value * 2
|
||||
|
||||
assert int_processor.can_handle(42)
|
||||
assert not int_processor.can_handle("hello")
|
||||
assert not int_processor.can_handle([])
|
||||
assert int_processor.can_handle(Message(data=42, source_id="mock"))
|
||||
assert not int_processor.can_handle(Message(data="hello", source_id="mock"))
|
||||
assert not int_processor.can_handle(Message(data=[], source_id="mock"))
|
||||
|
||||
async def test_single_parameter_execution(self):
|
||||
"""Test that single-parameter functions can be executed properly."""
|
||||
@@ -325,7 +326,7 @@ class TestFunctionExecutor:
|
||||
WorkflowBuilder().set_start_executor(double_value).build()
|
||||
|
||||
# For testing purposes, we can check that the handler is registered correctly
|
||||
assert double_value.can_handle(5)
|
||||
assert double_value.can_handle(Message(data=5, source_id="mock"))
|
||||
assert int in double_value._handlers
|
||||
|
||||
def test_sync_function_basic(self):
|
||||
@@ -369,9 +370,9 @@ class TestFunctionExecutor:
|
||||
def string_handler(text: str):
|
||||
return text.strip()
|
||||
|
||||
assert string_handler.can_handle("hello")
|
||||
assert not string_handler.can_handle(123)
|
||||
assert not string_handler.can_handle([])
|
||||
assert string_handler.can_handle(Message(data="hello", source_id="mock"))
|
||||
assert not string_handler.can_handle(Message(data=123, source_id="mock"))
|
||||
assert not string_handler.can_handle(Message(data=[], source_id="mock"))
|
||||
|
||||
def test_sync_function_validation(self):
|
||||
"""Test validation for synchronous functions."""
|
||||
@@ -413,8 +414,8 @@ class TestFunctionExecutor:
|
||||
assert isinstance(async_func, FunctionExecutor)
|
||||
|
||||
# Both should handle strings
|
||||
assert sync_func.can_handle("test")
|
||||
assert async_func.can_handle("test")
|
||||
assert sync_func.can_handle(Message(data="test", source_id="mock"))
|
||||
assert async_func.can_handle(Message(data="test", source_id="mock"))
|
||||
|
||||
# Both should be different instances
|
||||
assert sync_func is not async_func
|
||||
@@ -443,8 +444,8 @@ class TestFunctionExecutor:
|
||||
assert async_spec["workflow_output_types"] == [str] # Second parameter is str
|
||||
|
||||
# Verify the executors can handle their input types
|
||||
assert to_upper_sync.can_handle("hello")
|
||||
assert reverse_async.can_handle("HELLO")
|
||||
assert to_upper_sync.can_handle(Message(data="hello", source_id="mock"))
|
||||
assert reverse_async.can_handle(Message(data="HELLO", source_id="mock"))
|
||||
|
||||
# For integration testing, we mainly verify that the handlers are properly registered
|
||||
# and the functions are wrapped correctly
|
||||
|
||||
@@ -312,7 +312,6 @@ async def test_magentic_checkpoint_resume_round_trip():
|
||||
async for ev in wf.run_stream(task_text):
|
||||
if isinstance(ev, RequestInfoEvent) and ev.request_type is MagenticPlanReviewRequest:
|
||||
req_event = ev
|
||||
break
|
||||
assert req_event is not None
|
||||
|
||||
checkpoints = await storage.list_checkpoints()
|
||||
@@ -334,10 +333,16 @@ async def test_magentic_checkpoint_resume_round_trip():
|
||||
|
||||
reply = MagenticPlanReviewReply(decision=MagenticPlanReviewDecision.APPROVE)
|
||||
completed: WorkflowOutputEvent | None = None
|
||||
req_event = None
|
||||
async for event in wf_resume.run_stream_from_checkpoint(
|
||||
resume_checkpoint.checkpoint_id,
|
||||
responses={req_event.request_id: reply},
|
||||
):
|
||||
if isinstance(event, RequestInfoEvent) and event.request_type is MagenticPlanReviewRequest:
|
||||
req_event = event
|
||||
assert req_event is not None
|
||||
|
||||
responses = {req_event.request_id: reply}
|
||||
async for event in wf_resume.send_responses_streaming(responses=responses):
|
||||
if isinstance(event, WorkflowOutputEvent):
|
||||
completed = event
|
||||
assert completed is not None
|
||||
@@ -666,7 +671,6 @@ async def test_magentic_checkpoint_resume_rejects_participant_renames():
|
||||
async for event in workflow.run_stream("task"):
|
||||
if isinstance(event, RequestInfoEvent) and event.request_type is MagenticPlanReviewRequest:
|
||||
req_event = event
|
||||
break
|
||||
|
||||
assert req_event is not None
|
||||
|
||||
@@ -685,7 +689,6 @@ async def test_magentic_checkpoint_resume_rejects_participant_renames():
|
||||
with pytest.raises(ValueError, match="Workflow graph has changed"):
|
||||
async for _ in renamed_workflow.run_stream_from_checkpoint(
|
||||
target_checkpoint.checkpoint_id, # type: ignore[reportUnknownMemberType]
|
||||
responses={req_event.request_id: MagenticPlanReviewReply(decision=MagenticPlanReviewDecision.APPROVE)},
|
||||
):
|
||||
pass
|
||||
|
||||
|
||||
@@ -0,0 +1,413 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
from agent_framework import (
|
||||
FileCheckpointStorage,
|
||||
RequestInfoEvent,
|
||||
WorkflowBuilder,
|
||||
WorkflowContext,
|
||||
WorkflowRunState,
|
||||
WorkflowStatusEvent,
|
||||
handler,
|
||||
response_handler,
|
||||
)
|
||||
from agent_framework._workflows._executor import Executor
|
||||
from agent_framework._workflows._request_info_mixin import RequestInfoMixin
|
||||
|
||||
|
||||
@dataclass
|
||||
class UserApprovalRequest:
|
||||
"""A request for user approval with context."""
|
||||
|
||||
prompt: str
|
||||
context: str
|
||||
request_id: str = ""
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.request_id:
|
||||
import uuid
|
||||
|
||||
self.request_id = str(uuid.uuid4())
|
||||
|
||||
|
||||
@dataclass
|
||||
class CalculationRequest:
|
||||
"""A request for a complex calculation."""
|
||||
|
||||
operation: str
|
||||
operands: list[float]
|
||||
request_id: str = ""
|
||||
|
||||
def __post_init__(self):
|
||||
if not self.request_id:
|
||||
import uuid
|
||||
|
||||
self.request_id = str(uuid.uuid4())
|
||||
|
||||
|
||||
class ApprovalRequiredExecutor(Executor, RequestInfoMixin):
|
||||
"""Executor that requires approval before proceeding."""
|
||||
|
||||
def __init__(self, id: str):
|
||||
super().__init__(id=id)
|
||||
self.approval_received = False
|
||||
self.final_result = None
|
||||
|
||||
@handler
|
||||
async def start_process(self, message: str, ctx: WorkflowContext) -> None:
|
||||
"""Start a process that requires approval."""
|
||||
# Request approval from external system
|
||||
approval_request = UserApprovalRequest(
|
||||
prompt=f"Please approve the operation: {message}",
|
||||
context="This is a critical operation that requires human approval.",
|
||||
)
|
||||
await ctx.request_info(approval_request, UserApprovalRequest, bool)
|
||||
|
||||
@response_handler
|
||||
async def handle_approval_response(
|
||||
self, original_request: UserApprovalRequest, approved: bool, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
"""Handle the approval response."""
|
||||
self.approval_received = True
|
||||
|
||||
if approved:
|
||||
self.final_result = f"Operation approved: {original_request.prompt}"
|
||||
await ctx.send_message(f"APPROVED: {original_request.context}")
|
||||
else:
|
||||
self.final_result = "Operation denied by user"
|
||||
await ctx.send_message("DENIED: Operation was not approved")
|
||||
|
||||
|
||||
class CalculationExecutor(Executor, RequestInfoMixin):
|
||||
"""Executor that delegates complex calculations to external services."""
|
||||
|
||||
def __init__(self, id: str):
|
||||
super().__init__(id=id)
|
||||
self.calculations_performed: list[tuple[str, list[float], float]] = []
|
||||
|
||||
@handler
|
||||
async def process_calculation(self, message: str, ctx: WorkflowContext[str]) -> None:
|
||||
"""Process a calculation request."""
|
||||
# Parse the message to extract operation
|
||||
parts = message.split()
|
||||
if len(parts) >= 3:
|
||||
operation = parts[0]
|
||||
try:
|
||||
operands = [float(x) for x in parts[1:]]
|
||||
calc_request = CalculationRequest(operation=operation, operands=operands)
|
||||
await ctx.request_info(calc_request, CalculationRequest, float)
|
||||
except ValueError:
|
||||
await ctx.send_message("Invalid calculation format")
|
||||
else:
|
||||
await ctx.send_message("Insufficient parameters for calculation")
|
||||
|
||||
@response_handler
|
||||
async def handle_calculation_response(
|
||||
self, original_request: CalculationRequest, result: float, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
"""Handle the calculation response."""
|
||||
self.calculations_performed.append((original_request.operation, original_request.operands, result))
|
||||
operands_str = ", ".join(map(str, original_request.operands))
|
||||
await ctx.send_message(f"Calculation complete: {original_request.operation}({operands_str}) = {result}")
|
||||
|
||||
|
||||
class MultiRequestExecutor(Executor, RequestInfoMixin):
|
||||
"""Executor that makes multiple requests and waits for all responses."""
|
||||
|
||||
def __init__(self, id: str):
|
||||
super().__init__(id=id)
|
||||
self.responses_received: list[tuple[str, bool | float]] = []
|
||||
|
||||
@handler
|
||||
async def start_multi_request(self, message: str, ctx: WorkflowContext) -> None:
|
||||
"""Start multiple requests simultaneously."""
|
||||
# Request approval
|
||||
approval_request = UserApprovalRequest(
|
||||
prompt="Approve batch operation", context="Multiple operations will be performed"
|
||||
)
|
||||
await ctx.request_info(approval_request, UserApprovalRequest, bool)
|
||||
|
||||
# Request calculation
|
||||
calc_request = CalculationRequest(operation="multiply", operands=[10.0, 5.0])
|
||||
await ctx.request_info(calc_request, CalculationRequest, float)
|
||||
|
||||
@response_handler
|
||||
async def handle_approval_response(
|
||||
self, original_request: UserApprovalRequest, approved: bool, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
"""Handle approval response."""
|
||||
self.responses_received.append(("approval", approved))
|
||||
await self._check_completion(ctx)
|
||||
|
||||
@response_handler
|
||||
async def handle_calculation_response(
|
||||
self, original_request: CalculationRequest, result: float, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
"""Handle calculation response."""
|
||||
self.responses_received.append(("calculation", result))
|
||||
await self._check_completion(ctx)
|
||||
|
||||
async def _check_completion(self, ctx: WorkflowContext[str]) -> None:
|
||||
"""Check if all responses are received and send final result."""
|
||||
if len(self.responses_received) == 2:
|
||||
approval_result = next((r[1] for r in self.responses_received if r[0] == "approval"), None)
|
||||
calc_result = next((r[1] for r in self.responses_received if r[0] == "calculation"), None)
|
||||
|
||||
if approval_result and calc_result is not None:
|
||||
await ctx.send_message(f"All operations complete. Calculation result: {calc_result}")
|
||||
else:
|
||||
await ctx.send_message("Operations completed with mixed results")
|
||||
|
||||
|
||||
class OutputCollector(Executor):
|
||||
"""Simple executor that collects outputs for testing."""
|
||||
|
||||
def __init__(self, id: str):
|
||||
super().__init__(id=id)
|
||||
self.collected_outputs: list[str] = []
|
||||
|
||||
@handler
|
||||
async def collect_output(self, message: str, ctx: WorkflowContext) -> None:
|
||||
"""Collect output messages."""
|
||||
self.collected_outputs.append(message)
|
||||
|
||||
|
||||
class TestRequestInfoAndResponse:
|
||||
"""Test cases for end-to-end request info and response handling at the workflow level."""
|
||||
|
||||
async def test_approval_workflow(self):
|
||||
"""Test end-to-end workflow with approval request."""
|
||||
executor = ApprovalRequiredExecutor(id="approval_executor")
|
||||
workflow = WorkflowBuilder().set_start_executor(executor).build()
|
||||
|
||||
# First run the workflow until it emits a request
|
||||
request_info_event: RequestInfoEvent | None = None
|
||||
async for event in workflow.run_stream("test operation"):
|
||||
if isinstance(event, RequestInfoEvent):
|
||||
request_info_event = event
|
||||
|
||||
assert request_info_event is not None
|
||||
assert isinstance(request_info_event.data, UserApprovalRequest)
|
||||
assert request_info_event.data.prompt == "Please approve the operation: test operation"
|
||||
|
||||
# Send response and continue workflow
|
||||
completed = False
|
||||
async for event in workflow.send_responses_streaming({request_info_event.request_id: True}):
|
||||
if isinstance(event, WorkflowStatusEvent) and event.state == WorkflowRunState.IDLE:
|
||||
completed = True
|
||||
|
||||
assert completed
|
||||
assert executor.approval_received is True
|
||||
assert executor.final_result == "Operation approved: Please approve the operation: test operation"
|
||||
|
||||
async def test_calculation_workflow(self):
|
||||
"""Test end-to-end workflow with calculation request."""
|
||||
executor = CalculationExecutor(id="calc_executor")
|
||||
workflow = WorkflowBuilder().set_start_executor(executor).build()
|
||||
|
||||
# First run the workflow until it emits a calculation request
|
||||
request_info_event: RequestInfoEvent | None = None
|
||||
async for event in workflow.run_stream("multiply 15.5 2.0"):
|
||||
if isinstance(event, RequestInfoEvent):
|
||||
request_info_event = event
|
||||
|
||||
assert request_info_event is not None
|
||||
assert isinstance(request_info_event.data, CalculationRequest)
|
||||
assert request_info_event.data.operation == "multiply"
|
||||
assert request_info_event.data.operands == [15.5, 2.0]
|
||||
|
||||
# Send response with calculated result
|
||||
calculated_result = 31.0
|
||||
completed = False
|
||||
async for event in workflow.send_responses_streaming({request_info_event.request_id: calculated_result}):
|
||||
if isinstance(event, WorkflowStatusEvent) and event.state == WorkflowRunState.IDLE:
|
||||
completed = True
|
||||
|
||||
assert completed
|
||||
assert len(executor.calculations_performed) == 1
|
||||
assert executor.calculations_performed[0] == ("multiply", [15.5, 2.0], calculated_result)
|
||||
|
||||
async def test_multiple_requests_workflow(self):
|
||||
"""Test workflow with multiple concurrent requests."""
|
||||
executor = MultiRequestExecutor(id="multi_executor")
|
||||
workflow = WorkflowBuilder().set_start_executor(executor).build()
|
||||
|
||||
# Collect all request events by running the full stream
|
||||
request_events: list[RequestInfoEvent] = []
|
||||
async for event in workflow.run_stream("start batch"):
|
||||
if isinstance(event, RequestInfoEvent):
|
||||
request_events.append(event)
|
||||
|
||||
assert len(request_events) == 2
|
||||
|
||||
# Find the approval and calculation requests
|
||||
approval_event: RequestInfoEvent | None = next(
|
||||
(e for e in request_events if isinstance(e.data, UserApprovalRequest)), None
|
||||
)
|
||||
calc_event: RequestInfoEvent | None = next(
|
||||
(e for e in request_events if isinstance(e.data, CalculationRequest)), None
|
||||
)
|
||||
|
||||
assert approval_event is not None
|
||||
assert calc_event is not None
|
||||
|
||||
# Send responses for both requests
|
||||
responses = {approval_event.request_id: True, calc_event.request_id: 50.0}
|
||||
completed = False
|
||||
async for event in workflow.send_responses_streaming(responses):
|
||||
if isinstance(event, WorkflowStatusEvent) and event.state == WorkflowRunState.IDLE:
|
||||
completed = True
|
||||
|
||||
assert completed
|
||||
assert len(executor.responses_received) == 2
|
||||
|
||||
async def test_denied_approval_workflow(self):
|
||||
"""Test workflow when approval is denied."""
|
||||
executor = ApprovalRequiredExecutor(id="approval_executor")
|
||||
workflow = WorkflowBuilder().set_start_executor(executor).build()
|
||||
|
||||
# First run the workflow until it emits a request
|
||||
request_info_event: RequestInfoEvent | None = None
|
||||
async for event in workflow.run_stream("sensitive operation"):
|
||||
if isinstance(event, RequestInfoEvent):
|
||||
request_info_event = event
|
||||
|
||||
assert request_info_event is not None
|
||||
|
||||
# Deny the request
|
||||
completed = False
|
||||
async for event in workflow.send_responses_streaming({request_info_event.request_id: False}):
|
||||
if isinstance(event, WorkflowStatusEvent) and event.state == WorkflowRunState.IDLE:
|
||||
completed = True
|
||||
|
||||
assert completed
|
||||
assert executor.approval_received is True
|
||||
assert executor.final_result == "Operation denied by user"
|
||||
|
||||
async def test_workflow_state_with_pending_requests(self):
|
||||
"""Test workflow state when waiting for responses."""
|
||||
executor = ApprovalRequiredExecutor(id="approval_executor")
|
||||
workflow = WorkflowBuilder().set_start_executor(executor).build()
|
||||
|
||||
# Run workflow until idle with pending requests
|
||||
request_info_event: RequestInfoEvent | None = None
|
||||
idle_with_pending = False
|
||||
async for event in workflow.run_stream("test operation"):
|
||||
if isinstance(event, RequestInfoEvent):
|
||||
request_info_event = event
|
||||
elif isinstance(event, WorkflowStatusEvent) and event.state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS:
|
||||
idle_with_pending = True
|
||||
|
||||
assert request_info_event is not None
|
||||
assert idle_with_pending
|
||||
|
||||
# Continue with response
|
||||
completed = False
|
||||
async for event in workflow.send_responses_streaming({request_info_event.request_id: True}):
|
||||
if isinstance(event, WorkflowStatusEvent) and event.state == WorkflowRunState.IDLE:
|
||||
completed = True
|
||||
|
||||
assert completed
|
||||
|
||||
async def test_invalid_calculation_input(self):
|
||||
"""Test workflow handling of invalid calculation input."""
|
||||
executor = CalculationExecutor(id="calc_executor")
|
||||
workflow = WorkflowBuilder().set_start_executor(executor).build()
|
||||
|
||||
# Send invalid input (no numbers)
|
||||
completed = False
|
||||
async for event in workflow.run_stream("invalid input"):
|
||||
if isinstance(event, WorkflowStatusEvent) and event.state == WorkflowRunState.IDLE:
|
||||
completed = True
|
||||
|
||||
assert completed
|
||||
# Should not have any calculations performed due to invalid input
|
||||
assert len(executor.calculations_performed) == 0
|
||||
|
||||
async def test_checkpoint_with_pending_request_info_events(self):
|
||||
"""Test that request info events are properly serialized in checkpoints and can be restored."""
|
||||
import tempfile
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
# Use file-based storage to test full serialization
|
||||
storage = FileCheckpointStorage(temp_dir)
|
||||
|
||||
# Create workflow with checkpointing enabled
|
||||
executor = ApprovalRequiredExecutor(id="approval_executor")
|
||||
workflow = WorkflowBuilder().set_start_executor(executor).with_checkpointing(storage).build()
|
||||
|
||||
# Step 1: Run workflow to completion to ensure checkpoints are created
|
||||
request_info_event: RequestInfoEvent | None = None
|
||||
async for event in workflow.run_stream("checkpoint test operation"):
|
||||
if isinstance(event, RequestInfoEvent):
|
||||
request_info_event = event
|
||||
|
||||
# Verify request was emitted
|
||||
assert request_info_event is not None
|
||||
assert isinstance(request_info_event.data, UserApprovalRequest)
|
||||
assert request_info_event.data.prompt == "Please approve the operation: checkpoint test operation"
|
||||
assert request_info_event.source_executor_id == "approval_executor"
|
||||
|
||||
# Step 2: List checkpoints to find the one with our pending request
|
||||
checkpoints = await storage.list_checkpoints()
|
||||
assert len(checkpoints) > 0, "No checkpoints were created during workflow execution"
|
||||
|
||||
# Find the checkpoint with our pending request
|
||||
checkpoint_with_request = None
|
||||
for checkpoint in checkpoints:
|
||||
if request_info_event.request_id in checkpoint.pending_request_info_events:
|
||||
checkpoint_with_request = checkpoint
|
||||
break
|
||||
|
||||
assert checkpoint_with_request is not None, "No checkpoint found with pending request info event"
|
||||
|
||||
# Step 3: Verify the pending request info event was properly serialized
|
||||
serialized_event = checkpoint_with_request.pending_request_info_events[request_info_event.request_id]
|
||||
assert "data" in serialized_event
|
||||
assert "request_id" in serialized_event
|
||||
assert "source_executor_id" in serialized_event
|
||||
assert "request_type" in serialized_event
|
||||
assert serialized_event["request_id"] == request_info_event.request_id
|
||||
assert serialized_event["source_executor_id"] == "approval_executor"
|
||||
|
||||
# Step 4: Create a fresh workflow and restore from checkpoint
|
||||
new_executor = ApprovalRequiredExecutor(id="approval_executor")
|
||||
restored_workflow = WorkflowBuilder().set_start_executor(new_executor).with_checkpointing(storage).build()
|
||||
|
||||
# Step 5: Resume from checkpoint and verify the request can be continued
|
||||
completed = False
|
||||
restored_request_event: RequestInfoEvent | None = None
|
||||
async for event in restored_workflow.run_stream_from_checkpoint(checkpoint_with_request.checkpoint_id):
|
||||
# Should re-emit the pending request info event
|
||||
if isinstance(event, RequestInfoEvent) and event.request_id == request_info_event.request_id:
|
||||
restored_request_event = event
|
||||
elif (
|
||||
isinstance(event, WorkflowStatusEvent)
|
||||
and event.state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
|
||||
):
|
||||
completed = True
|
||||
|
||||
assert completed, "Workflow should reach idle with pending requests state after restoration"
|
||||
assert restored_request_event is not None, "Restored request info event should be emitted"
|
||||
|
||||
# Verify the restored event matches the original
|
||||
assert restored_request_event.source_executor_id == request_info_event.source_executor_id
|
||||
assert isinstance(restored_request_event.data, UserApprovalRequest)
|
||||
assert restored_request_event.data.prompt == request_info_event.data.prompt
|
||||
assert restored_request_event.data.context == request_info_event.data.context
|
||||
|
||||
# Step 6: Provide response to the restored request and complete the workflow
|
||||
final_completed = False
|
||||
async for event in restored_workflow.send_responses_streaming({
|
||||
request_info_event.request_id: True # Approve the request
|
||||
}):
|
||||
if isinstance(event, WorkflowStatusEvent) and event.state == WorkflowRunState.IDLE:
|
||||
final_completed = True
|
||||
|
||||
assert final_completed, "Workflow should complete after providing response to restored request"
|
||||
|
||||
# Step 7: Verify the executor state was properly restored and response was processed
|
||||
assert new_executor.approval_received is True
|
||||
expected_result = "Operation approved: Please approve the operation: checkpoint test operation"
|
||||
assert new_executor.final_result == expected_result
|
||||
@@ -0,0 +1,169 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework import InMemoryCheckpointStorage, InProcRunnerContext
|
||||
from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value
|
||||
from agent_framework._workflows._checkpoint_summary import get_checkpoint_summary
|
||||
from agent_framework._workflows._events import RequestInfoEvent
|
||||
from agent_framework._workflows._shared_state import SharedState
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockRequest: ...
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class SimpleApproval:
|
||||
prompt: str = ""
|
||||
draft: str = ""
|
||||
iteration: int = 0
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SlottedApproval:
|
||||
note: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class TimedApproval:
|
||||
issued_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
|
||||
|
||||
async def test_rehydrate_request_info_event() -> None:
|
||||
"""Rehydration should succeed for valid request info events."""
|
||||
request_info_event = RequestInfoEvent(
|
||||
request_id="request-123",
|
||||
source_executor_id="review_gateway",
|
||||
request_type=MockRequest,
|
||||
request_data=MockRequest(),
|
||||
response_type=bool,
|
||||
)
|
||||
|
||||
runner_context = InProcRunnerContext(InMemoryCheckpointStorage())
|
||||
await runner_context.add_request_info_event(request_info_event)
|
||||
|
||||
checkpoint_id = await runner_context.create_checkpoint(SharedState(), iteration_count=1)
|
||||
checkpoint = await runner_context.load_checkpoint(checkpoint_id)
|
||||
|
||||
assert checkpoint is not None
|
||||
assert checkpoint.pending_request_info_events
|
||||
assert "request-123" in checkpoint.pending_request_info_events
|
||||
assert "request_type" in checkpoint.pending_request_info_events["request-123"]
|
||||
|
||||
# Rehydrate the context
|
||||
await runner_context.apply_checkpoint(checkpoint)
|
||||
|
||||
pending_requests = await runner_context.get_pending_request_info_events()
|
||||
assert "request-123" in pending_requests
|
||||
rehydrated_event = pending_requests["request-123"]
|
||||
assert rehydrated_event.request_id == "request-123"
|
||||
assert rehydrated_event.source_executor_id == "review_gateway"
|
||||
assert rehydrated_event.request_type is MockRequest
|
||||
assert rehydrated_event.response_type is bool
|
||||
assert isinstance(rehydrated_event.data, MockRequest)
|
||||
|
||||
|
||||
async def test_rehydrate_fails_when_request_type_missing() -> None:
|
||||
"""Rehydration should fail is the request type is missing or fails to import."""
|
||||
request_info_event = RequestInfoEvent(
|
||||
request_id="request-123",
|
||||
source_executor_id="review_gateway",
|
||||
request_type=MockRequest,
|
||||
request_data=MockRequest(),
|
||||
response_type=bool,
|
||||
)
|
||||
|
||||
runner_context = InProcRunnerContext(InMemoryCheckpointStorage())
|
||||
await runner_context.add_request_info_event(request_info_event)
|
||||
|
||||
checkpoint_id = await runner_context.create_checkpoint(SharedState(), iteration_count=1)
|
||||
checkpoint = await runner_context.load_checkpoint(checkpoint_id)
|
||||
|
||||
assert checkpoint is not None
|
||||
assert checkpoint.pending_request_info_events
|
||||
assert "request-123" in checkpoint.pending_request_info_events
|
||||
assert "request_type" in checkpoint.pending_request_info_events["request-123"]
|
||||
|
||||
# Modify the checkpoint to simulate missing request type
|
||||
checkpoint.pending_request_info_events["request-123"]["request_type"] = "nonexistent.module:MissingRequest"
|
||||
|
||||
# Rehydrate the context
|
||||
with pytest.raises(ImportError):
|
||||
await runner_context.apply_checkpoint(checkpoint)
|
||||
|
||||
|
||||
async def test_pending_requests_in_summary() -> None:
|
||||
"""Test that pending requests are correctly summarized in the checkpoint summary."""
|
||||
request_info_event = RequestInfoEvent(
|
||||
request_id="request-123",
|
||||
source_executor_id="review_gateway",
|
||||
request_type=MockRequest,
|
||||
request_data=MockRequest(),
|
||||
response_type=bool,
|
||||
)
|
||||
|
||||
runner_context = InProcRunnerContext(InMemoryCheckpointStorage())
|
||||
await runner_context.add_request_info_event(request_info_event)
|
||||
|
||||
checkpoint_id = await runner_context.create_checkpoint(SharedState(), iteration_count=1)
|
||||
checkpoint = await runner_context.load_checkpoint(checkpoint_id)
|
||||
|
||||
assert checkpoint is not None
|
||||
summary = get_checkpoint_summary(checkpoint)
|
||||
|
||||
assert summary.checkpoint_id == checkpoint_id
|
||||
assert summary.status == "awaiting request response"
|
||||
|
||||
assert len(summary.pending_request_info_events) == 1
|
||||
pending_event = summary.pending_request_info_events[0]
|
||||
assert isinstance(pending_event, RequestInfoEvent)
|
||||
assert pending_event.request_id == "request-123"
|
||||
|
||||
assert pending_event.source_executor_id == "review_gateway"
|
||||
assert pending_event.request_type is MockRequest
|
||||
assert pending_event.response_type is bool
|
||||
assert isinstance(pending_event.data, MockRequest)
|
||||
|
||||
|
||||
async def test_request_info_event_serializes_non_json_payloads() -> None:
|
||||
req_1 = RequestInfoEvent(
|
||||
request_id="req-1",
|
||||
source_executor_id="source",
|
||||
request_type=TimedApproval,
|
||||
request_data=TimedApproval(issued_at=datetime(2024, 5, 4, 12, 30, 45)),
|
||||
response_type=bool,
|
||||
)
|
||||
req_2 = RequestInfoEvent(
|
||||
request_id="req-2",
|
||||
source_executor_id="source",
|
||||
request_type=SlottedApproval,
|
||||
request_data=SlottedApproval(note="slot-based"),
|
||||
response_type=bool,
|
||||
)
|
||||
|
||||
runner_context = InProcRunnerContext(InMemoryCheckpointStorage())
|
||||
await runner_context.add_request_info_event(req_1)
|
||||
await runner_context.add_request_info_event(req_2)
|
||||
|
||||
checkpoint_id = await runner_context.create_checkpoint(SharedState(), iteration_count=1)
|
||||
checkpoint = await runner_context.load_checkpoint(checkpoint_id)
|
||||
|
||||
# Should be JSON serializable despite datetime/slots
|
||||
serialized = json.dumps(encode_checkpoint_value(checkpoint))
|
||||
deserialized = json.loads(serialized)
|
||||
|
||||
assert "value" in deserialized
|
||||
deserialized = deserialized["value"]
|
||||
|
||||
assert "pending_request_info_events" in deserialized
|
||||
pending_request_info_events = deserialized["pending_request_info_events"]
|
||||
assert "req-1" in pending_request_info_events
|
||||
assert isinstance(pending_request_info_events["req-1"]["data"]["value"]["issued_at"], str)
|
||||
|
||||
assert "req-2" in pending_request_info_events
|
||||
assert pending_request_info_events["req-2"]["data"]["value"]["note"] == "slot-based"
|
||||
@@ -1,285 +0,0 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import json
|
||||
from dataclasses import dataclass, field
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any
|
||||
|
||||
from agent_framework._workflows._checkpoint import WorkflowCheckpoint
|
||||
from agent_framework._workflows._checkpoint_encoding import encode_checkpoint_value
|
||||
from agent_framework._workflows._checkpoint_summary import get_checkpoint_summary
|
||||
from agent_framework._workflows._const import EXECUTOR_STATE_KEY
|
||||
from agent_framework._workflows._events import RequestInfoEvent, WorkflowEvent
|
||||
from agent_framework._workflows._request_info_executor import (
|
||||
PendingRequestDetails,
|
||||
PendingRequestSnapshot,
|
||||
RequestInfoExecutor,
|
||||
RequestInfoMessage,
|
||||
RequestResponse,
|
||||
)
|
||||
from agent_framework._workflows._runner_context import Message
|
||||
from agent_framework._workflows._shared_state import SharedState
|
||||
from agent_framework._workflows._workflow_context import WorkflowContext
|
||||
|
||||
PENDING_STATE_KEY = RequestInfoExecutor._PENDING_SHARED_STATE_KEY # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
|
||||
class _StubRunnerContext:
|
||||
"""Minimal runner context stub for exercising WorkflowContext helpers."""
|
||||
|
||||
async def send_message(self, message: Message) -> None: # pragma: no cover - unused in tests
|
||||
return None
|
||||
|
||||
async def drain_messages(self) -> dict[str, list[Message]]: # pragma: no cover - unused
|
||||
return {}
|
||||
|
||||
async def has_messages(self) -> bool: # pragma: no cover - unused
|
||||
return False
|
||||
|
||||
async def add_event(self, event: WorkflowEvent) -> None: # pragma: no cover - unused
|
||||
return None
|
||||
|
||||
async def drain_events(self) -> list[WorkflowEvent]: # pragma: no cover - unused
|
||||
return []
|
||||
|
||||
async def has_events(self) -> bool: # pragma: no cover - unused
|
||||
return False
|
||||
|
||||
async def next_event(self) -> WorkflowEvent: # pragma: no cover - unused
|
||||
raise RuntimeError("Not implemented in stub context")
|
||||
|
||||
def has_checkpointing(self) -> bool: # pragma: no cover - unused
|
||||
return False
|
||||
|
||||
def set_workflow_id(self, workflow_id: str) -> None: # pragma: no cover - unused
|
||||
pass
|
||||
|
||||
def reset_for_new_run(self) -> None: # pragma: no cover - unused
|
||||
pass
|
||||
|
||||
async def create_checkpoint(
|
||||
self,
|
||||
shared_state: SharedState,
|
||||
iteration_count: int,
|
||||
metadata: dict[str, Any] | None = None,
|
||||
) -> str: # pragma: no cover - unused
|
||||
raise RuntimeError("Checkpointing not supported in stub context")
|
||||
|
||||
async def load_checkpoint(self, checkpoint_id: str) -> WorkflowCheckpoint | None: # pragma: no cover - unused
|
||||
return None
|
||||
|
||||
async def apply_checkpoint(self, checkpoint: WorkflowCheckpoint) -> None: # pragma: no cover - unused
|
||||
pass
|
||||
|
||||
def set_streaming(self, streaming: bool) -> None: # pragma: no cover - unused
|
||||
pass
|
||||
|
||||
def is_streaming(self) -> bool: # pragma: no cover - unused
|
||||
return False
|
||||
|
||||
|
||||
@dataclass(kw_only=True)
|
||||
class SimpleApproval(RequestInfoMessage):
|
||||
prompt: str = ""
|
||||
draft: str = ""
|
||||
iteration: int = 0
|
||||
|
||||
|
||||
@dataclass(slots=True)
|
||||
class SlottedApproval(RequestInfoMessage):
|
||||
note: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class TimedApproval(RequestInfoMessage):
|
||||
issued_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
|
||||
|
||||
|
||||
async def test_rehydrate_falls_back_when_request_type_missing() -> None:
|
||||
"""Rehydration should succeed even if the original request type cannot be imported.
|
||||
|
||||
This simulates resuming a workflow where the HumanApprovalRequest class is unavailable
|
||||
in the current process (e.g., defined in __main__ during the original run).
|
||||
"""
|
||||
request_id = "request-123"
|
||||
snapshot = PendingRequestSnapshot(
|
||||
request_id=request_id,
|
||||
source_executor_id="review_gateway",
|
||||
request_type="nonexistent.module:MissingRequest",
|
||||
request_as_json_safe_dict={
|
||||
"request_id": request_id,
|
||||
},
|
||||
)
|
||||
|
||||
ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), _StubRunnerContext())
|
||||
await ctx.set_executor_state({PENDING_STATE_KEY: {request_id: snapshot}})
|
||||
|
||||
executor = RequestInfoExecutor(id="request_info")
|
||||
|
||||
event = await executor._rehydrate_request_event(request_id, ctx) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
assert event is not None
|
||||
assert event.request_id == request_id
|
||||
assert isinstance(event.data, RequestInfoMessage)
|
||||
|
||||
|
||||
async def test_has_pending_request_detects_snapshot() -> None:
|
||||
request_id = "request-123"
|
||||
snapshot = PendingRequestSnapshot(
|
||||
request_id=request_id,
|
||||
source_executor_id="review_gateway",
|
||||
request_type="nonexistent.module:MissingRequest",
|
||||
request_as_json_safe_dict={
|
||||
"request_id": request_id,
|
||||
},
|
||||
)
|
||||
|
||||
ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), _StubRunnerContext())
|
||||
await ctx.set_executor_state({PENDING_STATE_KEY: {request_id: snapshot}})
|
||||
|
||||
executor = RequestInfoExecutor(id="request_info")
|
||||
|
||||
assert await executor.has_pending_request(request_id, ctx)
|
||||
|
||||
|
||||
async def test_has_pending_request_false_when_snapshot_absent() -> None:
|
||||
ctx: WorkflowContext[Any] = WorkflowContext("request_info", ["workflow"], SharedState(), _StubRunnerContext())
|
||||
await ctx.set_executor_state({PENDING_STATE_KEY: {}})
|
||||
|
||||
executor = RequestInfoExecutor(id="request_info")
|
||||
|
||||
assert not await executor.has_pending_request("missing", ctx)
|
||||
|
||||
|
||||
def test_pending_requests_from_checkpoint_and_summary() -> None:
|
||||
request = SimpleApproval(prompt="Review draft", draft="Draft text", iteration=3)
|
||||
request.request_id = "req-42"
|
||||
|
||||
response = RequestResponse[SimpleApproval, str](
|
||||
data="approve",
|
||||
original_request=request,
|
||||
request_id=request.request_id,
|
||||
)
|
||||
|
||||
encoded_response = encode_checkpoint_value(response)
|
||||
|
||||
checkpoint = WorkflowCheckpoint(
|
||||
checkpoint_id="cp-1",
|
||||
workflow_id="wf",
|
||||
messages={
|
||||
"request_info": [
|
||||
{
|
||||
"data": encoded_response,
|
||||
"source_id": "request_info",
|
||||
"target_id": "review_gateway",
|
||||
}
|
||||
]
|
||||
},
|
||||
shared_state={
|
||||
PENDING_STATE_KEY: {
|
||||
request.request_id: {
|
||||
"request_id": request.request_id,
|
||||
"prompt": request.prompt,
|
||||
"draft": request.draft,
|
||||
"iteration": request.iteration,
|
||||
"source_executor_id": "review_gateway",
|
||||
}
|
||||
}
|
||||
},
|
||||
iteration_count=1,
|
||||
)
|
||||
|
||||
summary = get_checkpoint_summary(checkpoint)
|
||||
assert summary.checkpoint_id == "cp-1"
|
||||
assert summary.status == "awaiting request response"
|
||||
assert summary.pending_requests[0].request_id == "req-42"
|
||||
|
||||
pending = summary.pending_requests
|
||||
assert len(pending) == 1
|
||||
entry = pending[0]
|
||||
assert isinstance(entry, PendingRequestDetails)
|
||||
assert entry.request_id == "req-42"
|
||||
assert entry.prompt == "Review draft"
|
||||
assert entry.draft == "Draft text"
|
||||
assert entry.iteration == 3
|
||||
assert entry.original_request is not None
|
||||
|
||||
|
||||
def test_snapshot_state_serializes_non_json_payloads() -> None:
|
||||
executor = RequestInfoExecutor(id="request_info")
|
||||
|
||||
timed = TimedApproval(issued_at=datetime(2024, 5, 4, 12, 30, 45))
|
||||
timed.request_id = "timed"
|
||||
slotted = SlottedApproval(note="slot-based")
|
||||
slotted.request_id = "slotted"
|
||||
|
||||
executor._request_events = { # pyright: ignore[reportPrivateUsage]
|
||||
timed.request_id: RequestInfoEvent(
|
||||
request_id=timed.request_id,
|
||||
source_executor_id="source",
|
||||
request_type=TimedApproval,
|
||||
request_data=timed,
|
||||
),
|
||||
slotted.request_id: RequestInfoEvent(
|
||||
request_id=slotted.request_id,
|
||||
source_executor_id="source",
|
||||
request_type=SlottedApproval,
|
||||
request_data=slotted,
|
||||
),
|
||||
}
|
||||
|
||||
state = executor.snapshot_state()
|
||||
|
||||
# Should be JSON serializable despite datetime/slots
|
||||
serialized = json.dumps(state)
|
||||
assert "timed" in serialized
|
||||
timed_payload = state["request_events"][timed.request_id]["request_data"]["value"]
|
||||
assert isinstance(timed_payload["issued_at"], str)
|
||||
|
||||
|
||||
def test_restore_state_falls_back_to_base_request_type() -> None:
|
||||
executor = RequestInfoExecutor(id="request_info")
|
||||
|
||||
approval = SimpleApproval(prompt="Review", draft="Draft", iteration=1)
|
||||
approval.request_id = "req"
|
||||
executor._request_events = { # pyright: ignore[reportPrivateUsage]
|
||||
approval.request_id: RequestInfoEvent(
|
||||
request_id=approval.request_id,
|
||||
source_executor_id="source",
|
||||
request_type=SimpleApproval,
|
||||
request_data=approval,
|
||||
)
|
||||
}
|
||||
|
||||
state = executor.snapshot_state()
|
||||
state["request_events"][approval.request_id]["request_type"] = "missing.module:GhostRequest"
|
||||
|
||||
executor.restore_state(state)
|
||||
|
||||
restored = executor._request_events[approval.request_id] # pyright: ignore[reportPrivateUsage]
|
||||
assert restored.request_type is RequestInfoMessage
|
||||
assert isinstance(restored.data, RequestInfoMessage)
|
||||
|
||||
|
||||
async def test_run_persists_pending_requests_in_runner_state() -> None:
|
||||
shared_state = SharedState()
|
||||
runner_ctx = _StubRunnerContext()
|
||||
ctx: WorkflowContext[None] = WorkflowContext("request_info", ["source"], shared_state, runner_ctx)
|
||||
|
||||
executor = RequestInfoExecutor(id="request_info")
|
||||
approval = SimpleApproval(prompt="Review", draft="Draft", iteration=1)
|
||||
approval.request_id = "req-123"
|
||||
|
||||
await executor.execute(approval, ctx.source_executor_ids, shared_state, runner_ctx)
|
||||
|
||||
# Runner state should include both pending snapshot and serialized request events
|
||||
assert await shared_state.has(EXECUTOR_STATE_KEY)
|
||||
executor_state = await shared_state.get(EXECUTOR_STATE_KEY)
|
||||
assert executor.id in executor_state
|
||||
assert PENDING_STATE_KEY in executor_state[executor.id]
|
||||
assert approval.request_id in executor_state[executor.id][PENDING_STATE_KEY]
|
||||
|
||||
response_ctx: WorkflowContext[None] = WorkflowContext("request_info", ["source"], shared_state, runner_ctx)
|
||||
await executor.handle_response("approved", approval.request_id, response_ctx) # type: ignore
|
||||
|
||||
assert executor_state[executor.id][PENDING_STATE_KEY] == {}
|
||||
@@ -0,0 +1,788 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
import inspect
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework._workflows._executor import Executor, handler
|
||||
from agent_framework._workflows._request_info_mixin import response_handler
|
||||
from agent_framework._workflows._workflow_context import WorkflowContext
|
||||
|
||||
|
||||
class TestRequestInfoMixin:
|
||||
"""Test cases for RequestInfoMixin functionality."""
|
||||
|
||||
def test_request_info_mixin_initialization(self):
|
||||
"""Test that RequestInfoMixin can be initialized."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
executor = TestExecutor()
|
||||
# After calling _discover_response_handlers, it should have the attributes
|
||||
assert hasattr(executor, "_response_handlers")
|
||||
assert hasattr(executor, "_response_handler_specs")
|
||||
assert hasattr(executor, "is_request_response_capable")
|
||||
assert executor.is_request_response_capable is False
|
||||
|
||||
def test_response_handler_decorator_creates_metadata(self):
|
||||
"""Test that the response_handler decorator creates proper metadata."""
|
||||
|
||||
@response_handler
|
||||
async def test_handler(self: Any, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
"""Test handler docstring."""
|
||||
pass
|
||||
|
||||
# Check that the decorator preserves function attributes
|
||||
assert test_handler.__name__ == "test_handler"
|
||||
assert test_handler.__doc__ == "Test handler docstring."
|
||||
assert hasattr(test_handler, "_response_handler_spec")
|
||||
|
||||
# Check the spec attributes
|
||||
spec = test_handler._response_handler_spec # type: ignore[reportAttributeAccessIssue]
|
||||
assert spec["name"] == "test_handler"
|
||||
assert spec["response_type"] is int
|
||||
assert spec["request_type"] is str
|
||||
|
||||
def test_response_handler_with_workflow_context_types(self):
|
||||
"""Test response handler with different WorkflowContext type parameters."""
|
||||
|
||||
@response_handler
|
||||
async def handler_with_output_types(
|
||||
self: Any, original_request: str, response: int, ctx: WorkflowContext[str, bool]
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
spec = handler_with_output_types._response_handler_spec # type: ignore[reportAttributeAccessIssue]
|
||||
assert "output_types" in spec
|
||||
assert "workflow_output_types" in spec
|
||||
|
||||
def test_response_handler_preserves_signature(self):
|
||||
"""Test that response_handler preserves the original function signature."""
|
||||
|
||||
async def original_handler(self: Any, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
decorated = response_handler(original_handler)
|
||||
|
||||
# Check that signature is preserved
|
||||
original_sig = inspect.signature(original_handler)
|
||||
decorated_sig = inspect.signature(decorated)
|
||||
|
||||
# Both should have the same parameter names and types
|
||||
assert list(original_sig.parameters.keys()) == list(decorated_sig.parameters.keys())
|
||||
|
||||
def test_executor_with_response_handlers(self):
|
||||
"""Test an executor with valid response handlers."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_string_response(
|
||||
self, original_request: str, response: int, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_dict_response(
|
||||
self, original_request: dict[str, Any], response: bool, ctx: WorkflowContext[bool]
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
# Should be request-response capable
|
||||
assert executor.is_request_response_capable is True
|
||||
|
||||
# Should have registered handlers
|
||||
response_handlers = executor._response_handlers # type: ignore[reportAttributeAccessIssue]
|
||||
assert len(response_handlers) == 2
|
||||
assert (str, int) in response_handlers
|
||||
assert (dict[str, Any], bool) in response_handlers
|
||||
|
||||
def test_executor_without_response_handlers(self):
|
||||
"""Test an executor without response handlers."""
|
||||
|
||||
class PlainExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="plain_executor")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
executor = PlainExecutor()
|
||||
|
||||
# Should not be request-response capable
|
||||
assert executor.is_request_response_capable is False
|
||||
|
||||
# Should have empty handlers
|
||||
response_handlers = executor._response_handlers # type: ignore[reportAttributeAccessIssue]
|
||||
assert len(response_handlers) == 0
|
||||
|
||||
def test_duplicate_response_handlers_raise_error(self):
|
||||
"""Test that duplicate response handlers for the same message type raise an error."""
|
||||
|
||||
class DuplicateExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="duplicate_executor")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_first(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_second(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Duplicate response handler for request type <class 'str'> and response type <class 'int'>",
|
||||
):
|
||||
DuplicateExecutor()
|
||||
|
||||
def test_response_handler_function_callable(self):
|
||||
"""Test that response handlers can actually be called."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
self.handled_request = None
|
||||
self.handled_response = None
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_response(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
self.handled_request = original_request
|
||||
self.handled_response = response
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
# Get the handler
|
||||
response_handler_func = executor._response_handlers[(str, int)] # type: ignore[reportAttributeAccessIssue]
|
||||
|
||||
# Create a mock context - we'll just use None since the handler doesn't use it
|
||||
asyncio.run(response_handler_func("test_request", 42, None)) # type: ignore[reportArgumentType]
|
||||
|
||||
assert executor.handled_request == "test_request"
|
||||
assert executor.handled_response == 42
|
||||
|
||||
def test_inheritance_with_response_handlers(self):
|
||||
"""Test that response handlers work correctly with inheritance."""
|
||||
|
||||
class BaseExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="base_executor")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def base_handler(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
class ChildExecutor(BaseExecutor):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.id = "child_executor"
|
||||
|
||||
@response_handler
|
||||
async def child_handler(self, original_request: str, response: bool, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
child = ChildExecutor()
|
||||
|
||||
# Should have both handlers
|
||||
response_handlers = child._response_handlers # type: ignore[reportAttributeAccessIssue]
|
||||
assert len(response_handlers) == 2
|
||||
assert (str, int) in response_handlers
|
||||
assert (str, bool) in response_handlers
|
||||
assert child.is_request_response_capable is True
|
||||
|
||||
def test_response_handler_spec_attributes(self):
|
||||
"""Test that response handler specs contain expected attributes."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def test_handler(self, original_request: str, response: int, ctx: WorkflowContext[str, bool]) -> None:
|
||||
pass
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
specs = executor._response_handler_specs # type: ignore[reportAttributeAccessIssue]
|
||||
assert len(specs) == 1
|
||||
|
||||
spec = specs[0]
|
||||
assert spec["name"] == "test_handler"
|
||||
assert spec["request_type"] is str
|
||||
assert spec["response_type"] is int
|
||||
assert "output_types" in spec
|
||||
assert "workflow_output_types" in spec
|
||||
assert "ctx_annotation" in spec
|
||||
assert spec["source"] == "class_method"
|
||||
|
||||
def test_multiple_discovery_calls_raise_error(self):
|
||||
"""Test that multiple calls to _discover_response_handlers raise an error for duplicates."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def test_handler(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
# First call should work fine
|
||||
first_handlers = len(executor._response_handlers) # type: ignore[reportAttributeAccessIssue]
|
||||
|
||||
# Second call should raise an error due to duplicate registration
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match="Duplicate response handler for request type <class 'str'> and response type <class 'int'>",
|
||||
):
|
||||
executor._discover_response_handlers() # type: ignore[reportAttributeAccessIssue]
|
||||
|
||||
# Handlers count should remain the same
|
||||
assert first_handlers == 1
|
||||
|
||||
def test_non_callable_attributes_ignored(self):
|
||||
"""Test that non-callable attributes are ignored during discovery."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
|
||||
some_variable = "not_a_function"
|
||||
another_attr = 42
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def valid_handler(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
# Should only have one handler despite other attributes
|
||||
response_handlers = executor._response_handlers # type: ignore[reportAttributeAccessIssue]
|
||||
assert len(response_handlers) == 1
|
||||
assert (str, int) in response_handlers
|
||||
|
||||
def test_same_request_type_different_response_types(self):
|
||||
"""Test that handlers with same request type but different response types are distinct."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
self.str_int_handler_called = False
|
||||
self.str_bool_handler_called = False
|
||||
self.str_dict_handler_called = False
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_str_int(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
self.str_int_handler_called = True
|
||||
|
||||
@response_handler
|
||||
async def handle_str_bool(self, original_request: str, response: bool, ctx: WorkflowContext[str]) -> None:
|
||||
self.str_bool_handler_called = True
|
||||
|
||||
@response_handler
|
||||
async def handle_str_dict(
|
||||
self, original_request: str, response: dict[str, Any], ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
self.str_dict_handler_called = True
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
# Should have three distinct handlers
|
||||
response_handlers = executor._response_handlers # type: ignore[reportAttributeAccessIssue]
|
||||
assert len(response_handlers) == 3
|
||||
assert (str, int) in response_handlers
|
||||
assert (str, bool) in response_handlers
|
||||
assert (str, dict[str, Any]) in response_handlers
|
||||
|
||||
# Test that each handler can be found correctly
|
||||
str_int_handler = executor._find_response_handler("test", 42) # pyright: ignore[reportPrivateUsage]
|
||||
str_bool_handler = executor._find_response_handler("test", True) # pyright: ignore[reportPrivateUsage]
|
||||
str_dict_handler = executor._find_response_handler("test", {"key": "value"}) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
assert str_int_handler is not None
|
||||
assert str_bool_handler is not None
|
||||
assert str_dict_handler is not None
|
||||
|
||||
# Test that handlers are called correctly
|
||||
asyncio.run(str_int_handler(42, None)) # type: ignore[reportArgumentType]
|
||||
asyncio.run(str_bool_handler(True, None)) # type: ignore[reportArgumentType]
|
||||
asyncio.run(str_dict_handler({"key": "value"}, None)) # type: ignore[reportArgumentType]
|
||||
|
||||
assert executor.str_int_handler_called
|
||||
assert executor.str_bool_handler_called
|
||||
assert executor.str_dict_handler_called
|
||||
|
||||
def test_different_request_types_same_response_type(self):
|
||||
"""Test that handlers with different request types but same response type are distinct."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
self.str_int_handler_called = False
|
||||
self.dict_int_handler_called = False
|
||||
self.list_int_handler_called = False
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_str_int(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
self.str_int_handler_called = True
|
||||
|
||||
@response_handler
|
||||
async def handle_dict_int(
|
||||
self, original_request: dict[str, Any], response: int, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
self.dict_int_handler_called = True
|
||||
|
||||
@response_handler
|
||||
async def handle_list_int(
|
||||
self, original_request: list[str], response: int, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
self.list_int_handler_called = True
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
# Should have three distinct handlers
|
||||
response_handlers = executor._response_handlers # type: ignore[reportAttributeAccessIssue]
|
||||
assert len(response_handlers) == 3
|
||||
assert (str, int) in response_handlers
|
||||
assert (dict[str, Any], int) in response_handlers
|
||||
assert (list[str], int) in response_handlers
|
||||
|
||||
# Test that each handler can be found correctly
|
||||
str_int_handler = executor._find_response_handler("test", 42) # pyright: ignore[reportPrivateUsage]
|
||||
dict_int_handler = executor._find_response_handler({"key": "value"}, 42) # pyright: ignore[reportPrivateUsage]
|
||||
list_int_handler = executor._find_response_handler(["test"], 42) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
assert str_int_handler is not None
|
||||
assert dict_int_handler is not None
|
||||
assert list_int_handler is not None
|
||||
|
||||
# Test that handlers are called correctly
|
||||
asyncio.run(str_int_handler(42, None)) # type: ignore[reportArgumentType]
|
||||
asyncio.run(dict_int_handler(42, None)) # type: ignore[reportArgumentType]
|
||||
asyncio.run(list_int_handler(42, None)) # type: ignore[reportArgumentType]
|
||||
|
||||
assert executor.str_int_handler_called
|
||||
assert executor.dict_int_handler_called
|
||||
assert executor.list_int_handler_called
|
||||
|
||||
def test_complex_type_combinations(self):
|
||||
"""Test response handlers with complex type combinations."""
|
||||
|
||||
class CustomRequest:
|
||||
pass
|
||||
|
||||
class CustomResponse:
|
||||
pass
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
self.custom_custom_called = False
|
||||
self.custom_str_called = False
|
||||
self.str_custom_called = False
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_custom_custom(
|
||||
self, original_request: CustomRequest, response: CustomResponse, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
self.custom_custom_called = True
|
||||
|
||||
@response_handler
|
||||
async def handle_custom_str(
|
||||
self, original_request: CustomRequest, response: str, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
self.custom_str_called = True
|
||||
|
||||
@response_handler
|
||||
async def handle_str_custom(
|
||||
self, original_request: str, response: CustomResponse, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
self.str_custom_called = True
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
# Should have three distinct handlers
|
||||
response_handlers = executor._response_handlers # type: ignore[reportAttributeAccessIssue]
|
||||
assert len(response_handlers) == 3
|
||||
assert (CustomRequest, CustomResponse) in response_handlers
|
||||
assert (CustomRequest, str) in response_handlers
|
||||
assert (str, CustomResponse) in response_handlers
|
||||
|
||||
# Test that each handler can be found correctly
|
||||
custom_request = CustomRequest()
|
||||
custom_response = CustomResponse()
|
||||
|
||||
custom_custom_handler = executor._find_response_handler(custom_request, custom_response) # pyright: ignore[reportPrivateUsage]
|
||||
custom_str_handler = executor._find_response_handler(custom_request, "test") # pyright: ignore[reportPrivateUsage]
|
||||
str_custom_handler = executor._find_response_handler("test", custom_response) # pyright: ignore[reportPrivateUsage]
|
||||
|
||||
assert custom_custom_handler is not None
|
||||
assert custom_str_handler is not None
|
||||
assert str_custom_handler is not None
|
||||
|
||||
def test_handler_key_uniqueness(self):
|
||||
"""Test that handler keys (request_type, response_type) are truly unique."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle1(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle2(self, original_request: int, response: str, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle3(self, original_request: str, response: str, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle4(self, original_request: int, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
# Should have four distinct handlers based on different combinations
|
||||
response_handlers = executor._response_handlers # type: ignore[reportAttributeAccessIssue]
|
||||
assert len(response_handlers) == 4
|
||||
|
||||
# Verify all expected combinations exist
|
||||
expected_keys = {
|
||||
(str, int), # handle1
|
||||
(int, str), # handle2
|
||||
(str, str), # handle3
|
||||
(int, int), # handle4
|
||||
}
|
||||
|
||||
actual_keys = set(response_handlers.keys())
|
||||
assert actual_keys == expected_keys
|
||||
|
||||
def test_no_false_matches_with_similar_types(self):
|
||||
"""Test that handlers don't match with similar but different types."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_str_int(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_list_str_float(
|
||||
self, original_request: list[str], response: float, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
# Test that wrong combinations don't match
|
||||
assert executor._find_response_handler("test", 3.14) is None # pyright: ignore[reportPrivateUsage] # str request, float response - no handler
|
||||
assert executor._find_response_handler(["test"], 42) is None # pyright: ignore[reportPrivateUsage] # list request, int response - no handler
|
||||
assert executor._find_response_handler(42, "test") is None # pyright: ignore[reportPrivateUsage] # int request, str response - no handler
|
||||
|
||||
# Test that correct combinations do match
|
||||
assert executor._find_response_handler("test", 42) is not None # pyright: ignore[reportPrivateUsage] # str request, int response - has handler
|
||||
assert executor._find_response_handler(["test"], 3.14) is not None # pyright: ignore[reportPrivateUsage] # list request, float response - has handler
|
||||
|
||||
def test_is_request_supported_with_exact_matches(self):
|
||||
"""Test is_request_supported with exact type matches."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_str_int(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_dict_bool(
|
||||
self, original_request: dict[str, Any], response: bool, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
# Test exact matches
|
||||
assert executor.is_request_supported(str, int) is True
|
||||
assert executor.is_request_supported(str, bool) is True # bool and int are compatible
|
||||
assert executor.is_request_supported(dict[str, Any], bool) is True
|
||||
|
||||
# Test non-matches
|
||||
assert executor.is_request_supported(int, str) is False
|
||||
assert executor.is_request_supported(list[str], int) is False
|
||||
|
||||
def test_is_request_supported_without_handlers(self):
|
||||
"""Test is_request_supported when no handlers are registered."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
# Should return False for any type combination
|
||||
assert executor.is_request_supported(str, int) is False
|
||||
assert executor.is_request_supported(dict[str, Any], bool) is False
|
||||
assert executor.is_request_supported(int, str) is False
|
||||
|
||||
def test_is_request_supported_before_discovery(self):
|
||||
"""Test is_request_supported before response handlers are discovered."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor", defer_discovery=True)
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_str_int(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
executor = TestExecutor()
|
||||
# Don't call _discover_response_handlers()
|
||||
|
||||
# Should return False when _response_handlers attribute doesn't exist
|
||||
assert executor.is_request_supported(str, int) is False
|
||||
assert executor.is_request_supported(dict[str, Any], bool) is False
|
||||
|
||||
def test_is_request_supported_with_compatible_types(self):
|
||||
"""Test is_request_supported with type-compatible scenarios."""
|
||||
|
||||
class BaseRequest:
|
||||
pass
|
||||
|
||||
class DerivedRequest(BaseRequest):
|
||||
pass
|
||||
|
||||
class BaseResponse:
|
||||
pass
|
||||
|
||||
class DerivedResponse(BaseResponse):
|
||||
pass
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_base_base(
|
||||
self, original_request: BaseRequest, response: BaseResponse, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_str_int(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
# Test exact matches
|
||||
assert executor.is_request_supported(BaseRequest, BaseResponse) is True
|
||||
assert executor.is_request_supported(str, int) is True
|
||||
|
||||
# Test compatible derived types (depends on is_type_compatible implementation)
|
||||
# These should return True if the type compatibility function supports inheritance
|
||||
result_derived_request = executor.is_request_supported(DerivedRequest, BaseResponse)
|
||||
result_derived_response = executor.is_request_supported(BaseRequest, DerivedResponse)
|
||||
result_both_derived = executor.is_request_supported(DerivedRequest, DerivedResponse)
|
||||
|
||||
# The actual result depends on the is_type_compatible implementation
|
||||
# We'll just assert that the method doesn't raise an exception
|
||||
assert isinstance(result_derived_request, bool)
|
||||
assert isinstance(result_derived_response, bool)
|
||||
assert isinstance(result_both_derived, bool)
|
||||
|
||||
def test_is_request_supported_with_multiple_handlers(self):
|
||||
"""Test is_request_supported when multiple handlers are registered."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_str_int(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_str_bool(self, original_request: str, response: bool, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_dict_str(
|
||||
self, original_request: dict[str, Any], response: str, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_list_float(
|
||||
self, original_request: list[str], response: float, ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
# Test all registered combinations
|
||||
assert executor.is_request_supported(str, int) is True
|
||||
assert executor.is_request_supported(str, bool) is True
|
||||
assert executor.is_request_supported(dict[str, Any], str) is True
|
||||
assert executor.is_request_supported(list[str], float) is True
|
||||
|
||||
# Test combinations that don't exist
|
||||
assert executor.is_request_supported(str, float) is False
|
||||
assert executor.is_request_supported(int, str) is False
|
||||
assert executor.is_request_supported(dict[str, Any], int) is False
|
||||
assert executor.is_request_supported(list[str], bool) is False
|
||||
|
||||
def test_is_request_supported_with_complex_types(self):
|
||||
"""Test is_request_supported with complex generic types."""
|
||||
|
||||
class TestExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="test_executor")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_dict_list(
|
||||
self, original_request: dict[str, Any], response: list[int], ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def handle_list_dict(
|
||||
self, original_request: list[str], response: dict[str, bool], ctx: WorkflowContext[str]
|
||||
) -> None:
|
||||
pass
|
||||
|
||||
executor = TestExecutor()
|
||||
|
||||
# Test complex type matches
|
||||
assert executor.is_request_supported(dict[str, Any], list[int]) is True
|
||||
assert executor.is_request_supported(list[str], dict[str, bool]) is True
|
||||
|
||||
# Test non-matches with similar but different complex types
|
||||
assert executor.is_request_supported(dict[str, Any], list[str]) is False
|
||||
assert executor.is_request_supported(list[int], dict[str, bool]) is False
|
||||
assert executor.is_request_supported(dict[int, Any], list[int]) is False
|
||||
|
||||
def test_is_request_supported_with_inheritance(self):
|
||||
"""Test is_request_supported with inherited response handlers."""
|
||||
|
||||
class BaseExecutor(Executor):
|
||||
def __init__(self):
|
||||
super().__init__(id="base_executor")
|
||||
|
||||
@handler
|
||||
async def dummy_handler(self, message: str, ctx: WorkflowContext) -> None:
|
||||
pass
|
||||
|
||||
@response_handler
|
||||
async def base_handler(self, original_request: str, response: int, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
class ChildExecutor(BaseExecutor):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.id = "child_executor"
|
||||
|
||||
@response_handler
|
||||
async def child_handler(self, original_request: str, response: bool, ctx: WorkflowContext[str]) -> None:
|
||||
pass
|
||||
|
||||
child = ChildExecutor()
|
||||
|
||||
# Should support both inherited and child-defined handlers
|
||||
assert child.is_request_supported(str, int) is True # From base class
|
||||
assert child.is_request_supported(str, bool) is True # From child class
|
||||
|
||||
# Should not support unregistered combinations
|
||||
assert child.is_request_supported(str, str) is False
|
||||
assert child.is_request_supported(int, str) is False
|
||||
@@ -6,12 +6,14 @@ from typing import Any
|
||||
import pytest
|
||||
|
||||
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
|
||||
from agent_framework._workflows._const import INTERNAL_SOURCE_ID
|
||||
from agent_framework._workflows._edge import (
|
||||
Case,
|
||||
Default,
|
||||
Edge,
|
||||
FanInEdgeGroup,
|
||||
FanOutEdgeGroup,
|
||||
InternalEdgeGroup,
|
||||
SingleEdgeGroup,
|
||||
SwitchCaseEdgeGroup,
|
||||
SwitchCaseEdgeGroupCase,
|
||||
@@ -557,16 +559,32 @@ class TestSerializationWorkflowClasses:
|
||||
|
||||
# Verify edge groups contain edges
|
||||
edge_groups = data["edge_groups"]
|
||||
assert len(edge_groups) == 1, "Should have exactly one edge group"
|
||||
edge_group = edge_groups[0]
|
||||
assert "edges" in edge_group, "Edge group should contain 'edges' field"
|
||||
assert len(edge_group["edges"]) == 1, "Should have exactly one edge"
|
||||
|
||||
edge = edge_group["edges"][0]
|
||||
assert "source_id" in edge, "Edge should have source_id"
|
||||
assert "target_id" in edge, "Edge should have target_id"
|
||||
assert edge["source_id"] == "executor1", f"Expected source_id 'executor1', got {edge['source_id']}"
|
||||
assert edge["target_id"] == "executor2", f"Expected target_id 'executor2', got {edge['target_id']}"
|
||||
single_edge_groups = [SingleEdgeGroup.from_dict(eg) for eg in edge_groups if eg["type"] == "SingleEdgeGroup"]
|
||||
internal_edge_groups = [
|
||||
InternalEdgeGroup.from_dict(eg) for eg in edge_groups if eg["type"] == "InternalEdgeGroup"
|
||||
]
|
||||
|
||||
assert len(single_edge_groups) == 1, "Should have exactly one SingleEdgeGroup for the added edge"
|
||||
assert len(internal_edge_groups) == 2, (
|
||||
"Should have exactly two (one per executor) InternalEdgeGroups for request/response handling"
|
||||
)
|
||||
|
||||
for edge_group in single_edge_groups:
|
||||
assert len(edge_group.edges) == 1, "Should have exactly one edge"
|
||||
|
||||
edge = edge_group.edges[0]
|
||||
|
||||
assert edge.source_id == "executor1", f"Expected source_id 'executor1', got {edge.source_id}"
|
||||
assert edge.target_id == "executor2", f"Expected target_id 'executor2', got {edge.target_id}"
|
||||
|
||||
for edge_group in internal_edge_groups:
|
||||
assert len(edge_group.edges) == 1, "Each InternalEdgeGroup should have exactly one edge"
|
||||
|
||||
edge = edge_group.edges[0]
|
||||
|
||||
assert edge.source_id == INTERNAL_SOURCE_ID(edge.target_id)
|
||||
assert edge.target_id in [executor1.id, executor2.id]
|
||||
|
||||
# Test model_dump_json
|
||||
json_str = workflow.to_json()
|
||||
@@ -577,12 +595,21 @@ class TestSerializationWorkflowClasses:
|
||||
|
||||
# Verify edges are preserved in JSON serialization
|
||||
json_edge_groups = parsed["edge_groups"]
|
||||
assert len(json_edge_groups) == 1, "JSON should have exactly one edge group"
|
||||
json_edge_group = json_edge_groups[0]
|
||||
assert "edges" in json_edge_group, "JSON edge group should contain 'edges' field"
|
||||
json_edge = json_edge_group["edges"][0]
|
||||
assert json_edge["source_id"] == "executor1", "JSON should preserve edge source_id"
|
||||
assert json_edge["target_id"] == "executor2", "JSON should preserve edge target_id"
|
||||
assert len(json_edge_groups) == 1 + 2, "JSON should have exactly one SingleEdgeGroup and two InternalEdgeGroups"
|
||||
|
||||
for json_edge_group in json_edge_groups:
|
||||
assert "edges" in json_edge_group, "JSON edge group should contain 'edges' field"
|
||||
assert len(json_edge_group["edges"]) == 1, "Each JSON edge group should have exactly one edge"
|
||||
if json_edge_group["type"] == "SingleEdgeGroup":
|
||||
json_edge = json_edge_group["edges"][0]
|
||||
assert json_edge["source_id"] == "executor1", "JSON should preserve edge source_id"
|
||||
assert json_edge["target_id"] == "executor2", "JSON should preserve edge target_id"
|
||||
elif json_edge_group["type"] == "InternalEdgeGroup":
|
||||
json_edge = json_edge_group["edges"][0]
|
||||
assert json_edge["source_id"] == INTERNAL_SOURCE_ID(json_edge["target_id"])
|
||||
assert json_edge["target_id"] in [executor1.id, executor2.id]
|
||||
else:
|
||||
pytest.fail(f"Unexpected edge group type: {json_edge_group['type']}")
|
||||
|
||||
def test_workflow_serialization_excludes_non_serializable_fields(self) -> None:
|
||||
"""Test that non-serializable fields are excluded from serialization."""
|
||||
|
||||
@@ -1,20 +1,20 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
from dataclasses import dataclass, field
|
||||
from uuid import uuid4
|
||||
|
||||
from typing_extensions import Never
|
||||
|
||||
from agent_framework import (
|
||||
Executor,
|
||||
RequestInfoExecutor,
|
||||
RequestInfoMessage,
|
||||
RequestResponse,
|
||||
SubWorkflowRequestMessage,
|
||||
SubWorkflowResponseMessage,
|
||||
Workflow,
|
||||
WorkflowBuilder,
|
||||
WorkflowContext,
|
||||
WorkflowExecutor,
|
||||
handler,
|
||||
response_handler,
|
||||
)
|
||||
|
||||
|
||||
@@ -27,9 +27,10 @@ class EmailValidationRequest:
|
||||
|
||||
|
||||
@dataclass
|
||||
class DomainCheckRequest(RequestInfoMessage):
|
||||
class DomainCheckRequest:
|
||||
"""Request to check if a domain is approved."""
|
||||
|
||||
id: str = field(default_factory=lambda: str(uuid4()))
|
||||
domain: str = ""
|
||||
email: str = "" # Include original email for correlation
|
||||
|
||||
@@ -43,72 +44,93 @@ class ValidationResult:
|
||||
reason: str
|
||||
|
||||
|
||||
# Test helper functions
|
||||
def create_email_validation_workflow() -> Workflow:
|
||||
"""Create a standard email validation workflow."""
|
||||
email_validator = EmailValidator()
|
||||
email_request_info = RequestInfoExecutor(id="email_request_info")
|
||||
|
||||
return (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(email_validator)
|
||||
.add_edge(email_validator, email_request_info)
|
||||
.add_edge(email_request_info, email_validator)
|
||||
.build()
|
||||
)
|
||||
|
||||
|
||||
class BasicParent(Executor):
|
||||
"""Basic parent executor for simple sub-workflow tests."""
|
||||
class Coordinator(Executor):
|
||||
"""Coordinator executor in the parent workflow for simple sub-workflow tests."""
|
||||
|
||||
def __init__(self, cache: dict[str, bool] | None = None) -> None:
|
||||
super().__init__(id="basic_parent")
|
||||
self.result: ValidationResult | None = None
|
||||
self.cache: dict[str, bool] = dict(cache) if cache is not None else {}
|
||||
self._pending_sub_workflow_requests: dict[str, SubWorkflowRequestMessage] = {}
|
||||
|
||||
@handler
|
||||
async def start(self, email: str, ctx: WorkflowContext[EmailValidationRequest]) -> None:
|
||||
request = EmailValidationRequest(email=email)
|
||||
await ctx.send_message(request, target_id="email_workflow")
|
||||
await ctx.send_message(request)
|
||||
|
||||
@handler
|
||||
async def handle_domain_request(
|
||||
self,
|
||||
request: DomainCheckRequest,
|
||||
ctx: WorkflowContext[RequestResponse[DomainCheckRequest, Any] | DomainCheckRequest],
|
||||
sub_workflow_request: SubWorkflowRequestMessage,
|
||||
ctx: WorkflowContext[SubWorkflowResponseMessage],
|
||||
) -> None:
|
||||
"""Handle requests from sub-workflows with optional caching."""
|
||||
domain_request = request
|
||||
if not isinstance(sub_workflow_request.source_event.data, DomainCheckRequest):
|
||||
raise ValueError("Unexpected request type")
|
||||
|
||||
domain_request = sub_workflow_request.source_event.data
|
||||
|
||||
if domain_request.domain in self.cache:
|
||||
# Return cached result
|
||||
response = RequestResponse(
|
||||
data=self.cache[domain_request.domain], original_request=request, request_id=request.request_id
|
||||
)
|
||||
await ctx.send_message(response, target_id=request.source_executor_id)
|
||||
await ctx.send_message(sub_workflow_request.create_response(self.cache[domain_request.domain]))
|
||||
else:
|
||||
# Not in cache, forward to external
|
||||
await ctx.send_message(request)
|
||||
self._pending_sub_workflow_requests[domain_request.id] = sub_workflow_request
|
||||
await ctx.request_info(domain_request, DomainCheckRequest, bool)
|
||||
|
||||
@response_handler
|
||||
async def handle_domain_response(
|
||||
self,
|
||||
original_request: DomainCheckRequest,
|
||||
is_approved: bool,
|
||||
ctx: WorkflowContext[SubWorkflowResponseMessage],
|
||||
) -> None:
|
||||
"""Handle domain check response with correlation and send the response back to the sub-workflow."""
|
||||
if original_request.id not in self._pending_sub_workflow_requests:
|
||||
raise ValueError("No pending sub-workflow request for the given domain check response")
|
||||
|
||||
sub_workflow_request = self._pending_sub_workflow_requests.pop(original_request.id)
|
||||
await ctx.send_message(sub_workflow_request.create_response(is_approved))
|
||||
|
||||
@handler
|
||||
async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None:
|
||||
self.result = result
|
||||
|
||||
|
||||
# Test executors
|
||||
class EmailValidator(Executor):
|
||||
class EmailFormatValidator(Executor):
|
||||
"""Validates the format of an email address."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(id="email_format_validator")
|
||||
|
||||
@handler
|
||||
async def validate(
|
||||
self, request: EmailValidationRequest, ctx: WorkflowContext[DomainCheckRequest, ValidationResult]
|
||||
) -> None:
|
||||
"""Validate email format and extract domain."""
|
||||
email = request.email
|
||||
if "@" not in email:
|
||||
result = ValidationResult(email=email, is_valid=False, reason="Invalid email format")
|
||||
await ctx.yield_output(result)
|
||||
return
|
||||
|
||||
domain = email.split("@")[1]
|
||||
domain_check = DomainCheckRequest(domain=domain, email=email)
|
||||
await ctx.send_message(domain_check)
|
||||
|
||||
|
||||
class EmailDomainValidator(Executor):
|
||||
"""Validates email addresses in a sub-workflow."""
|
||||
|
||||
def __init__(self):
|
||||
super().__init__(id="email_validator")
|
||||
super().__init__(id="email_domain_validator")
|
||||
|
||||
@handler
|
||||
async def validate_request(
|
||||
self, request: EmailValidationRequest, ctx: WorkflowContext[DomainCheckRequest, ValidationResult]
|
||||
self, request: DomainCheckRequest, ctx: WorkflowContext[DomainCheckRequest, ValidationResult]
|
||||
) -> None:
|
||||
"""Validate an email address."""
|
||||
# Extract domain and check if it's approved
|
||||
domain = request.email.split("@")[1] if "@" in request.email else ""
|
||||
domain = request.domain
|
||||
|
||||
if not domain:
|
||||
result = ValidationResult(email=request.email, is_valid=False, reason="Invalid email format")
|
||||
@@ -116,62 +138,37 @@ class EmailValidator(Executor):
|
||||
return
|
||||
|
||||
# Request domain check from external source
|
||||
domain_check = DomainCheckRequest(domain=domain, email=request.email)
|
||||
await ctx.send_message(domain_check)
|
||||
await ctx.request_info(request, DomainCheckRequest, bool)
|
||||
|
||||
@handler
|
||||
@response_handler
|
||||
async def handle_domain_response(
|
||||
self, response: RequestResponse[DomainCheckRequest, bool], ctx: WorkflowContext[Never, ValidationResult]
|
||||
self,
|
||||
original_request: DomainCheckRequest,
|
||||
is_approved: bool,
|
||||
ctx: WorkflowContext[Never, ValidationResult],
|
||||
) -> None:
|
||||
"""Handle domain check response with correlation."""
|
||||
# Use the original email from the correlated response
|
||||
result = ValidationResult(
|
||||
email=response.original_request.email,
|
||||
is_valid=response.data or False,
|
||||
reason="Domain approved" if response.data else "Domain not approved",
|
||||
email=original_request.email,
|
||||
is_valid=is_approved,
|
||||
reason="Domain approved" if is_approved else "Domain not approved",
|
||||
)
|
||||
await ctx.yield_output(result)
|
||||
|
||||
|
||||
class ParentOrchestrator(Executor):
|
||||
"""Parent workflow orchestrator with domain knowledge."""
|
||||
# Test helper functions
|
||||
def create_email_validation_workflow() -> Workflow:
|
||||
"""Create a standard email validation workflow."""
|
||||
email_format_validator = EmailFormatValidator()
|
||||
email_domain_validator = EmailDomainValidator()
|
||||
|
||||
def __init__(self, approved_domains: set[str] | None = None) -> None:
|
||||
super().__init__(id="parent_orchestrator")
|
||||
self.approved_domains: set[str] = (
|
||||
set(approved_domains) if approved_domains is not None else {"example.com", "test.org"}
|
||||
)
|
||||
self.results: list[ValidationResult] = []
|
||||
|
||||
@handler
|
||||
async def start(self, emails: list[str], ctx: WorkflowContext[EmailValidationRequest]) -> None:
|
||||
"""Start processing emails."""
|
||||
for email in emails:
|
||||
request = EmailValidationRequest(email=email)
|
||||
await ctx.send_message(request, target_id="email_workflow")
|
||||
|
||||
@handler
|
||||
async def handle_domain_request(
|
||||
self,
|
||||
request: DomainCheckRequest,
|
||||
ctx: WorkflowContext[RequestResponse[DomainCheckRequest, Any] | DomainCheckRequest],
|
||||
) -> None:
|
||||
"""Handle requests from sub-workflows."""
|
||||
domain_request = request
|
||||
|
||||
# Check if we know this domain
|
||||
if domain_request.domain in self.approved_domains:
|
||||
# Send response back to sub-workflow
|
||||
response = RequestResponse(data=True, original_request=request, request_id=request.request_id)
|
||||
await ctx.send_message(response, target_id=request.source_executor_id)
|
||||
else:
|
||||
# We don't know this domain, forward to external
|
||||
await ctx.send_message(request)
|
||||
|
||||
@handler
|
||||
async def collect_result(self, result: ValidationResult, ctx: WorkflowContext) -> None:
|
||||
"""Collect validation results."""
|
||||
self.results.append(result)
|
||||
return (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(email_format_validator)
|
||||
.add_edge(email_format_validator, email_domain_validator)
|
||||
.build()
|
||||
)
|
||||
|
||||
|
||||
async def test_basic_sub_workflow() -> None:
|
||||
@@ -180,17 +177,14 @@ async def test_basic_sub_workflow() -> None:
|
||||
validation_workflow = create_email_validation_workflow()
|
||||
|
||||
# Create parent workflow without interception
|
||||
parent = BasicParent()
|
||||
workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow")
|
||||
main_request_info = RequestInfoExecutor(id="main_request_info")
|
||||
parent = Coordinator()
|
||||
workflow_executor = WorkflowExecutor(validation_workflow, "email_validation_workflow")
|
||||
|
||||
main_workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(parent)
|
||||
.add_edge(parent, workflow_executor)
|
||||
.add_edge(workflow_executor, parent)
|
||||
.add_edge(workflow_executor, main_request_info)
|
||||
.add_edge(main_request_info, workflow_executor) # CRITICAL: For RequestResponse routing
|
||||
.build()
|
||||
)
|
||||
|
||||
@@ -220,17 +214,14 @@ async def test_sub_workflow_with_interception():
|
||||
validation_workflow = create_email_validation_workflow()
|
||||
|
||||
# Create parent workflow with interception cache
|
||||
parent = BasicParent(cache={"example.com": True, "internal.org": True})
|
||||
parent = Coordinator(cache={"example.com": True, "internal.org": True})
|
||||
workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow")
|
||||
parent_request_info = RequestInfoExecutor(id="request_info")
|
||||
|
||||
main_workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(parent)
|
||||
.add_edge(parent, workflow_executor)
|
||||
.add_edge(workflow_executor, parent)
|
||||
.add_edge(parent, parent_request_info) # For forwarded requests
|
||||
.add_edge(parent_request_info, workflow_executor) # For RequestResponse routing
|
||||
.build()
|
||||
)
|
||||
|
||||
@@ -276,6 +267,7 @@ async def test_workflow_scoped_interception() -> None:
|
||||
def __init__(self) -> None:
|
||||
super().__init__(id="multi_parent")
|
||||
self.results: dict[str, ValidationResult] = {}
|
||||
self._pending_sub_workflow_requests: dict[str, SubWorkflowRequestMessage] = {}
|
||||
|
||||
@handler
|
||||
async def start(self, data: dict[str, str], ctx: WorkflowContext[EmailValidationRequest]) -> None:
|
||||
@@ -286,30 +278,47 @@ async def test_workflow_scoped_interception() -> None:
|
||||
@handler
|
||||
async def handle_domain_request(
|
||||
self,
|
||||
request: DomainCheckRequest,
|
||||
ctx: WorkflowContext[RequestResponse[DomainCheckRequest, Any] | DomainCheckRequest],
|
||||
sub_workflow_request: SubWorkflowRequestMessage,
|
||||
ctx: WorkflowContext[SubWorkflowResponseMessage],
|
||||
) -> None:
|
||||
domain_request = request
|
||||
"""Handle requests from sub-workflows with optional caching."""
|
||||
if not isinstance(sub_workflow_request.source_event.data, DomainCheckRequest):
|
||||
raise ValueError("Unexpected request type")
|
||||
|
||||
if request.source_executor_id == "workflow_a":
|
||||
domain_request = sub_workflow_request.source_event.data
|
||||
|
||||
if sub_workflow_request.executor_id == "workflow_a" and domain_request.domain == "strict.com":
|
||||
# Strict rules for workflow A
|
||||
if domain_request.domain == "strict.com":
|
||||
response = RequestResponse(data=True, original_request=request, request_id=request.request_id)
|
||||
await ctx.send_message(response, target_id=request.source_executor_id)
|
||||
else:
|
||||
# Forward to external
|
||||
await ctx.send_message(request)
|
||||
elif request.source_executor_id == "workflow_b":
|
||||
await ctx.send_message(
|
||||
sub_workflow_request.create_response(True), target_id=sub_workflow_request.executor_id
|
||||
)
|
||||
return
|
||||
if sub_workflow_request.executor_id == "workflow_b" and domain_request.domain.endswith(".com"):
|
||||
# Lenient rules for workflow B
|
||||
if domain_request.domain.endswith(".com"):
|
||||
response = RequestResponse(data=True, original_request=request, request_id=request.request_id)
|
||||
await ctx.send_message(response, target_id=request.source_executor_id)
|
||||
else:
|
||||
# Forward to external
|
||||
await ctx.send_message(request)
|
||||
else:
|
||||
# Unknown source, forward to external
|
||||
await ctx.send_message(request)
|
||||
await ctx.send_message(
|
||||
sub_workflow_request.create_response(True), target_id=sub_workflow_request.executor_id
|
||||
)
|
||||
return
|
||||
|
||||
# Unknown source, forward to external
|
||||
self._pending_sub_workflow_requests[domain_request.id] = sub_workflow_request
|
||||
await ctx.request_info(domain_request, DomainCheckRequest, bool)
|
||||
|
||||
@response_handler
|
||||
async def handle_domain_response(
|
||||
self,
|
||||
original_request: DomainCheckRequest,
|
||||
is_approved: bool,
|
||||
ctx: WorkflowContext[SubWorkflowResponseMessage],
|
||||
) -> None:
|
||||
"""Handle domain check response with correlation and send the response back to the sub-workflow."""
|
||||
if original_request.id not in self._pending_sub_workflow_requests:
|
||||
raise ValueError("No pending sub-workflow request for the given domain check response")
|
||||
|
||||
sub_workflow_request = self._pending_sub_workflow_requests.pop(original_request.id)
|
||||
await ctx.send_message(
|
||||
sub_workflow_request.create_response(is_approved), target_id=sub_workflow_request.executor_id
|
||||
)
|
||||
|
||||
@handler
|
||||
async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None:
|
||||
@@ -322,7 +331,6 @@ async def test_workflow_scoped_interception() -> None:
|
||||
parent = MultiWorkflowParent()
|
||||
executor_a = WorkflowExecutor(workflow_a, "workflow_a")
|
||||
executor_b = WorkflowExecutor(workflow_b, "workflow_b")
|
||||
parent_request_info = RequestInfoExecutor(id="request_info")
|
||||
|
||||
main_workflow = (
|
||||
WorkflowBuilder()
|
||||
@@ -331,9 +339,6 @@ async def test_workflow_scoped_interception() -> None:
|
||||
.add_edge(parent, executor_b)
|
||||
.add_edge(executor_a, parent)
|
||||
.add_edge(executor_b, parent)
|
||||
.add_edge(parent, parent_request_info)
|
||||
.add_edge(parent_request_info, executor_a) # For RequestResponse routing
|
||||
.add_edge(parent_request_info, executor_b) # For RequestResponse routing
|
||||
.build()
|
||||
)
|
||||
|
||||
@@ -359,6 +364,7 @@ async def test_concurrent_sub_workflow_execution() -> None:
|
||||
def __init__(self) -> None:
|
||||
super().__init__(id="concurrent_processor")
|
||||
self.results: list[ValidationResult] = []
|
||||
self._pending_sub_workflow_requests: dict[str, SubWorkflowRequestMessage] = {}
|
||||
|
||||
@handler
|
||||
async def start(self, emails: list[str], ctx: WorkflowContext[EmailValidationRequest]) -> None:
|
||||
@@ -366,7 +372,35 @@ async def test_concurrent_sub_workflow_execution() -> None:
|
||||
# Send all requests concurrently to the same workflow executor
|
||||
for email in emails:
|
||||
request = EmailValidationRequest(email=email)
|
||||
await ctx.send_message(request, target_id="email_workflow")
|
||||
await ctx.send_message(request)
|
||||
|
||||
@handler
|
||||
async def handle_domain_request(
|
||||
self,
|
||||
sub_workflow_request: SubWorkflowRequestMessage,
|
||||
ctx: WorkflowContext[SubWorkflowResponseMessage],
|
||||
) -> None:
|
||||
"""Handle requests from sub-workflows with optional caching."""
|
||||
if not isinstance(sub_workflow_request.source_event.data, DomainCheckRequest):
|
||||
raise ValueError("Unexpected request type")
|
||||
|
||||
domain_request = sub_workflow_request.source_event.data
|
||||
self._pending_sub_workflow_requests[domain_request.id] = sub_workflow_request
|
||||
await ctx.request_info(domain_request, DomainCheckRequest, bool)
|
||||
|
||||
@response_handler
|
||||
async def handle_domain_response(
|
||||
self,
|
||||
original_request: DomainCheckRequest,
|
||||
is_approved: bool,
|
||||
ctx: WorkflowContext[SubWorkflowResponseMessage],
|
||||
) -> None:
|
||||
"""Handle domain check response with correlation and send the response back to the sub-workflow."""
|
||||
if original_request.id not in self._pending_sub_workflow_requests:
|
||||
raise ValueError("No pending sub-workflow request for the given domain check response")
|
||||
|
||||
sub_workflow_request = self._pending_sub_workflow_requests.pop(original_request.id)
|
||||
await ctx.send_message(sub_workflow_request.create_response(is_approved))
|
||||
|
||||
@handler
|
||||
async def collect_result(self, result: ValidationResult, ctx: WorkflowContext) -> None:
|
||||
@@ -379,15 +413,12 @@ async def test_concurrent_sub_workflow_execution() -> None:
|
||||
# Create parent workflow
|
||||
processor = ConcurrentProcessor()
|
||||
workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow")
|
||||
parent_request_info = RequestInfoExecutor(id="request_info")
|
||||
|
||||
main_workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(processor)
|
||||
.add_edge(processor, workflow_executor)
|
||||
.add_edge(workflow_executor, processor)
|
||||
.add_edge(workflow_executor, parent_request_info) # For external requests
|
||||
.add_edge(parent_request_info, workflow_executor) # For RequestResponse routing
|
||||
.build()
|
||||
)
|
||||
|
||||
|
||||
@@ -3,8 +3,13 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Generic, TypeVar, Union
|
||||
|
||||
from agent_framework._workflows import RequestInfoMessage, RequestResponse
|
||||
from agent_framework._workflows._typing_utils import is_instance_of, is_type_compatible
|
||||
from agent_framework import RequestInfoEvent
|
||||
from agent_framework._workflows._typing_utils import (
|
||||
deserialize_type,
|
||||
is_instance_of,
|
||||
is_type_compatible,
|
||||
serialize_type,
|
||||
)
|
||||
|
||||
|
||||
def test_basic_types() -> None:
|
||||
@@ -91,22 +96,6 @@ def test_custom_type() -> None:
|
||||
assert not is_instance_of(instance, dict)
|
||||
|
||||
|
||||
def test_request_response_type() -> None:
|
||||
"""Test RequestResponse generic type checking."""
|
||||
|
||||
request_instance = RequestResponse[RequestInfoMessage, str](
|
||||
data="approve",
|
||||
request_id="req-1",
|
||||
original_request=RequestInfoMessage(),
|
||||
)
|
||||
|
||||
class CustomRequestInfoMessage(RequestInfoMessage):
|
||||
info: str
|
||||
|
||||
assert is_instance_of(request_instance, RequestResponse[RequestInfoMessage, str])
|
||||
assert not is_instance_of(request_instance, RequestResponse[CustomRequestInfoMessage, str])
|
||||
|
||||
|
||||
def test_custom_generic_type() -> None:
|
||||
"""Test custom generic type checking."""
|
||||
|
||||
@@ -135,12 +124,90 @@ def test_edge_cases() -> None:
|
||||
assert not is_instance_of(5, str | None) # Optional type without matching type
|
||||
|
||||
|
||||
def test_serialize_type() -> None:
|
||||
"""Test serialization of types to strings."""
|
||||
# Test built-in types
|
||||
assert serialize_type(int) == "builtins.int"
|
||||
assert serialize_type(str) == "builtins.str"
|
||||
assert serialize_type(float) == "builtins.float"
|
||||
assert serialize_type(bool) == "builtins.bool"
|
||||
assert serialize_type(list) == "builtins.list"
|
||||
assert serialize_type(dict) == "builtins.dict"
|
||||
assert serialize_type(tuple) == "builtins.tuple"
|
||||
assert serialize_type(set) == "builtins.set"
|
||||
|
||||
# Test custom class
|
||||
@dataclass
|
||||
class TestClass:
|
||||
value: int
|
||||
|
||||
# The custom class will be in the test module
|
||||
expected = f"{TestClass.__module__}.{TestClass.__qualname__}"
|
||||
assert serialize_type(TestClass) == expected
|
||||
|
||||
|
||||
def test_deserialize_type() -> None:
|
||||
"""Test deserialization of type strings back to types."""
|
||||
# Test built-in types
|
||||
assert deserialize_type("builtins.int") is int
|
||||
assert deserialize_type("builtins.str") is str
|
||||
assert deserialize_type("builtins.float") is float
|
||||
assert deserialize_type("builtins.bool") is bool
|
||||
assert deserialize_type("builtins.list") is list
|
||||
assert deserialize_type("builtins.dict") is dict
|
||||
assert deserialize_type("builtins.tuple") is tuple
|
||||
assert deserialize_type("builtins.set") is set
|
||||
|
||||
|
||||
def test_serialize_deserialize_roundtrip() -> None:
|
||||
"""Test that serialization and deserialization are inverse operations."""
|
||||
# Test built-in types
|
||||
types_to_test = [int, str, float, bool, list, dict, tuple, set]
|
||||
|
||||
for type_to_test in types_to_test:
|
||||
serialized = serialize_type(type_to_test)
|
||||
deserialized = deserialize_type(serialized)
|
||||
assert deserialized is type_to_test
|
||||
|
||||
# Test agent framework type roundtrip
|
||||
|
||||
serialized = serialize_type(RequestInfoEvent)
|
||||
deserialized = deserialize_type(serialized)
|
||||
assert deserialized is RequestInfoEvent
|
||||
|
||||
# Verify we can instantiate the deserialized type
|
||||
instance = deserialized(
|
||||
request_id="request-123",
|
||||
source_executor_id="executor_1",
|
||||
request_type=str,
|
||||
request_data="test",
|
||||
response_type=str,
|
||||
)
|
||||
assert isinstance(instance, RequestInfoEvent)
|
||||
|
||||
|
||||
def test_deserialize_type_error_handling() -> None:
|
||||
"""Test error handling in deserialize_type function."""
|
||||
import pytest
|
||||
|
||||
# Test with non-existent module
|
||||
with pytest.raises(ModuleNotFoundError):
|
||||
deserialize_type("nonexistent.module.Type")
|
||||
|
||||
# Test with non-existent type in existing module
|
||||
with pytest.raises(AttributeError):
|
||||
deserialize_type("builtins.NonExistentType")
|
||||
|
||||
|
||||
def test_type_compatibility_basic() -> None:
|
||||
"""Test basic type compatibility scenarios."""
|
||||
# Exact type match
|
||||
assert is_type_compatible(str, str)
|
||||
assert is_type_compatible(int, int)
|
||||
|
||||
# bool is a subtype of int
|
||||
assert is_type_compatible(bool, int)
|
||||
|
||||
# Any compatibility
|
||||
assert is_type_compatible(str, Any)
|
||||
assert is_type_compatible(list[int], Any)
|
||||
|
||||
@@ -8,7 +8,6 @@ import pytest
|
||||
from agent_framework import (
|
||||
EdgeDuplicationError,
|
||||
Executor,
|
||||
ExecutorDuplicationError,
|
||||
GraphConnectivityError,
|
||||
TypeCompatibilityError,
|
||||
ValidationTypeEnum,
|
||||
@@ -83,11 +82,10 @@ def test_duplicate_executor_ids_fail_validation():
|
||||
executor1 = StringExecutor(id="dup")
|
||||
executor2 = IntExecutor(id="dup")
|
||||
|
||||
with pytest.raises(ExecutorDuplicationError) as exc_info:
|
||||
with pytest.raises(ValueError) as exc_info:
|
||||
(WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build())
|
||||
|
||||
assert exc_info.value.executor_id == "dup"
|
||||
assert exc_info.value.validation_type == ValidationTypeEnum.EXECUTOR_DUPLICATION
|
||||
assert str(exc_info.value) == "Duplicate executor ID 'dup' detected in workflow."
|
||||
|
||||
|
||||
def test_edge_duplication_validation_fails():
|
||||
@@ -185,7 +183,7 @@ def test_graph_connectivity_isolated_executors():
|
||||
assert "executor3" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_start_executor_not_in_graph():
|
||||
def test_disconnected_start_executor_not_in_graph():
|
||||
executor1 = StringExecutor(id="executor1")
|
||||
executor2 = StringExecutor(id="executor2")
|
||||
executor3 = StringExecutor(id="executor3") # Not in graph
|
||||
@@ -193,7 +191,7 @@ def test_start_executor_not_in_graph():
|
||||
with pytest.raises(GraphConnectivityError) as exc_info:
|
||||
WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor3).build()
|
||||
|
||||
assert "not present in the workflow graph" in str(exc_info.value)
|
||||
assert "The following executors are unreachable from the start executor 'executor3'" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_missing_start_executor():
|
||||
|
||||
@@ -3,8 +3,9 @@
|
||||
import asyncio
|
||||
import tempfile
|
||||
from collections.abc import AsyncIterable
|
||||
from dataclasses import dataclass
|
||||
from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
from uuid import uuid4
|
||||
|
||||
import pytest
|
||||
|
||||
@@ -21,9 +22,6 @@ from agent_framework import (
|
||||
FileCheckpointStorage,
|
||||
Message,
|
||||
RequestInfoEvent,
|
||||
RequestInfoExecutor,
|
||||
RequestInfoMessage,
|
||||
RequestResponse,
|
||||
Role,
|
||||
TextContent,
|
||||
WorkflowBuilder,
|
||||
@@ -33,6 +31,7 @@ from agent_framework import (
|
||||
WorkflowRunState,
|
||||
WorkflowStatusEvent,
|
||||
handler,
|
||||
response_handler,
|
||||
)
|
||||
|
||||
|
||||
@@ -68,6 +67,14 @@ class AggregatorExecutor(Executor):
|
||||
await ctx.yield_output(sum(msg.data for msg in messages))
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockRequest:
|
||||
"""A mock request message for testing purposes."""
|
||||
|
||||
request_id: str = field(default_factory=lambda: str(uuid4()))
|
||||
prompt: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class ApprovalMessage:
|
||||
"""A mock message for approval requests."""
|
||||
@@ -79,22 +86,22 @@ class MockExecutorRequestApproval(Executor):
|
||||
"""A mock executor that simulates a request for approval."""
|
||||
|
||||
@handler
|
||||
async def mock_handler_a(self, message: NumberMessage, ctx: WorkflowContext[RequestInfoMessage]) -> None:
|
||||
async def mock_handler_a(self, message: NumberMessage, ctx: WorkflowContext) -> None:
|
||||
"""A mock handler that requests approval."""
|
||||
await ctx.set_shared_state(self.id, message.data)
|
||||
await ctx.send_message(RequestInfoMessage())
|
||||
await ctx.request_info(MockRequest(prompt="Mock approval request"), MockRequest, ApprovalMessage)
|
||||
|
||||
@handler
|
||||
@response_handler
|
||||
async def mock_handler_b(
|
||||
self,
|
||||
message: RequestResponse[RequestInfoMessage, ApprovalMessage],
|
||||
original_request: MockRequest,
|
||||
response: ApprovalMessage,
|
||||
ctx: WorkflowContext[NumberMessage, int],
|
||||
) -> None:
|
||||
"""A mock handler that processes the approval response."""
|
||||
data = await ctx.get_shared_state(self.id)
|
||||
assert isinstance(data, int)
|
||||
assert isinstance(message.data, ApprovalMessage)
|
||||
if message.data.approved:
|
||||
if response.approved:
|
||||
await ctx.yield_output(data)
|
||||
else:
|
||||
await ctx.send_message(NumberMessage(data=data))
|
||||
@@ -182,15 +189,12 @@ async def test_workflow_send_responses_streaming():
|
||||
"""Test the workflow run with approval."""
|
||||
executor_a = IncrementExecutor(id="executor_a")
|
||||
executor_b = MockExecutorRequestApproval(id="executor_b")
|
||||
request_info_executor = RequestInfoExecutor(id="request_info")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(executor_a)
|
||||
.add_edge(executor_a, executor_b)
|
||||
.add_edge(executor_b, executor_a)
|
||||
.add_edge(executor_b, request_info_executor)
|
||||
.add_edge(request_info_executor, executor_b)
|
||||
.build()
|
||||
)
|
||||
|
||||
@@ -219,15 +223,12 @@ async def test_workflow_send_responses():
|
||||
"""Test the workflow run with approval."""
|
||||
executor_a = IncrementExecutor(id="executor_a")
|
||||
executor_b = MockExecutorRequestApproval(id="executor_b")
|
||||
request_info_executor = RequestInfoExecutor(id="request_info")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(executor_a)
|
||||
.add_edge(executor_a, executor_b)
|
||||
.add_edge(executor_b, executor_a)
|
||||
.add_edge(executor_b, request_info_executor)
|
||||
.add_edge(request_info_executor, executor_b)
|
||||
.build()
|
||||
)
|
||||
|
||||
@@ -480,6 +481,15 @@ async def test_workflow_run_stream_from_checkpoint_with_responses(simple_executo
|
||||
workflow_id="test-workflow",
|
||||
messages={},
|
||||
shared_state={},
|
||||
pending_request_info_events={
|
||||
"request_123": RequestInfoEvent(
|
||||
request_id="request_123",
|
||||
source_executor_id=simple_executor.id,
|
||||
request_type=str,
|
||||
request_data="Mock",
|
||||
response_type=str,
|
||||
).to_dict(),
|
||||
},
|
||||
iteration_count=0,
|
||||
)
|
||||
checkpoint_id = await storage.save_checkpoint(test_checkpoint)
|
||||
@@ -494,17 +504,20 @@ async def test_workflow_run_stream_from_checkpoint_with_responses(simple_executo
|
||||
)
|
||||
|
||||
# Test that run_stream_from_checkpoint accepts responses parameter
|
||||
responses = {"request_123": {"data": "test_response"}}
|
||||
responses = {"request_123": "test_response"}
|
||||
|
||||
try:
|
||||
events: list[WorkflowEvent] = []
|
||||
async for event in workflow.run_stream_from_checkpoint(checkpoint_id, responses=responses):
|
||||
events.append(event)
|
||||
if len(events) >= 2: # Limit to avoid infinite loops
|
||||
break
|
||||
except Exception:
|
||||
# Expected since we have minimal setup, but method should accept the parameters
|
||||
pass
|
||||
events: list[WorkflowEvent] = []
|
||||
async for event in workflow.run_stream_from_checkpoint(checkpoint_id):
|
||||
events.append(event)
|
||||
|
||||
assert next(
|
||||
event for event in events if isinstance(event, RequestInfoEvent) and event.request_id == "request_123"
|
||||
)
|
||||
|
||||
async for event in workflow.send_responses_streaming(responses):
|
||||
events.append(event)
|
||||
|
||||
assert len(events) > 0 # Just ensure we processed some events
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -737,7 +750,7 @@ async def test_workflow_concurrent_execution_prevention_streaming():
|
||||
|
||||
# Create an async generator that will consume the stream slowly
|
||||
async def consume_stream_slowly():
|
||||
result = []
|
||||
result: list[WorkflowEvent] = []
|
||||
async for event in workflow.run_stream(NumberMessage(data=0)):
|
||||
result.append(event)
|
||||
await asyncio.sleep(0.01) # Slow consumption
|
||||
@@ -770,7 +783,7 @@ async def test_workflow_concurrent_execution_prevention_mixed_methods():
|
||||
|
||||
# Start a streaming execution
|
||||
async def consume_stream():
|
||||
result = []
|
||||
result: list[WorkflowEvent] = []
|
||||
async for event in workflow.run_stream(NumberMessage(data=0)):
|
||||
result.append(event)
|
||||
await asyncio.sleep(0.01)
|
||||
@@ -846,6 +859,7 @@ async def test_agent_streaming_vs_non_streaming() -> None:
|
||||
assert len(agent_run_events) == 1, "Expected exactly one AgentRunEvent in non-streaming mode"
|
||||
assert len(agent_update_events) == 0, "Expected no AgentRunUpdateEvent in non-streaming mode"
|
||||
assert agent_run_events[0].executor_id == "agent_exec"
|
||||
assert agent_run_events[0].data is not None
|
||||
assert agent_run_events[0].data.messages[0].text == "Hello World"
|
||||
|
||||
# Test streaming mode with run_stream()
|
||||
@@ -866,6 +880,8 @@ async def test_agent_streaming_vs_non_streaming() -> None:
|
||||
|
||||
# Verify the updates build up to the full message
|
||||
accumulated_text = "".join(
|
||||
e.data.contents[0].text for e in stream_agent_update_events if e.data.contents and e.data.contents[0].text
|
||||
e.data.contents[0].text
|
||||
for e in stream_agent_update_events
|
||||
if e.data and e.data.contents and e.data.contents[0].text
|
||||
)
|
||||
assert accumulated_text == "Hello World", f"Expected 'Hello World', got '{accumulated_text}'"
|
||||
|
||||
@@ -14,8 +14,6 @@ from agent_framework import (
|
||||
FunctionApprovalRequestContent,
|
||||
FunctionApprovalResponseContent,
|
||||
FunctionCallContent,
|
||||
RequestInfoExecutor,
|
||||
RequestInfoMessage,
|
||||
Role,
|
||||
TextContent,
|
||||
UsageContent,
|
||||
@@ -24,6 +22,7 @@ from agent_framework import (
|
||||
WorkflowBuilder,
|
||||
WorkflowContext,
|
||||
handler,
|
||||
response_handler,
|
||||
)
|
||||
|
||||
|
||||
@@ -56,15 +55,17 @@ class SimpleExecutor(Executor):
|
||||
|
||||
|
||||
class RequestingExecutor(Executor):
|
||||
"""Executor that sends RequestInfoMessage to trigger RequestInfoEvent."""
|
||||
"""Executor that requests info."""
|
||||
|
||||
@handler
|
||||
async def handle_message(self, _: list[ChatMessage], ctx: WorkflowContext[RequestInfoMessage]) -> None:
|
||||
async def handle_message(self, _: list[ChatMessage], ctx: WorkflowContext) -> None:
|
||||
# Send a RequestInfoMessage to trigger the request info process
|
||||
await ctx.send_message(RequestInfoMessage())
|
||||
await ctx.request_info("Mock request data", str, str)
|
||||
|
||||
@handler
|
||||
async def handle_request_response(self, _: Any, ctx: WorkflowContext[ChatMessage]) -> None:
|
||||
@response_handler
|
||||
async def handle_request_response(
|
||||
self, original_request: str, response: str, ctx: WorkflowContext[ChatMessage]
|
||||
) -> None:
|
||||
# Handle the response and emit completion response
|
||||
update = AgentRunResponseUpdate(
|
||||
contents=[TextContent(text="Request completed successfully")],
|
||||
@@ -148,14 +149,11 @@ class TestWorkflowAgent:
|
||||
async def test_end_to_end_request_info_handling(self):
|
||||
"""Test end-to-end workflow with RequestInfoEvent handling."""
|
||||
# Create workflow with requesting executor -> request info executor (no cycle)
|
||||
simple_executor = SimpleExecutor(id="simple", response_text="SimpleResponse", emit_streaming=False)
|
||||
requesting_executor = RequestingExecutor(id="requester")
|
||||
request_info_executor = RequestInfoExecutor(id="request_info")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(requesting_executor)
|
||||
.add_edge(requesting_executor, request_info_executor)
|
||||
.build()
|
||||
WorkflowBuilder().set_start_executor(simple_executor).add_edge(simple_executor, requesting_executor).build()
|
||||
)
|
||||
|
||||
agent = WorkflowAgent(workflow=workflow, name="Request Test Agent")
|
||||
|
||||
@@ -100,7 +100,7 @@ def test_workflow_builder_fluent_api():
|
||||
.build()
|
||||
)
|
||||
|
||||
assert len(workflow.edge_groups) == 4
|
||||
assert len(workflow.edge_groups) == 4 + 6 # 4 defined edges + 6 internal edges for request-response handling
|
||||
assert workflow.start_executor_id == executor_a.id
|
||||
assert len(workflow.executors) == 6
|
||||
|
||||
|
||||
@@ -24,6 +24,18 @@ if TYPE_CHECKING:
|
||||
from agent_framework._workflows._runner_context import InProcRunnerContext
|
||||
|
||||
|
||||
class MockExecutor(Executor):
|
||||
"""Mock executor for testing."""
|
||||
|
||||
def __init__(self, id: str) -> None:
|
||||
super().__init__(id=id)
|
||||
|
||||
@handler
|
||||
async def handle_message(self, message: str, ctx: WorkflowContext[str]) -> None:
|
||||
"""Handle string messages."""
|
||||
...
|
||||
|
||||
|
||||
@asynccontextmanager
|
||||
async def make_context(
|
||||
executor_id: str = "exec",
|
||||
@@ -31,10 +43,11 @@ async def make_context(
|
||||
from agent_framework._workflows._runner_context import InProcRunnerContext
|
||||
from agent_framework._workflows._shared_state import SharedState
|
||||
|
||||
mock_executor = MockExecutor(executor_id)
|
||||
runner_ctx = InProcRunnerContext()
|
||||
shared_state = SharedState()
|
||||
workflow_ctx: WorkflowContext[object] = WorkflowContext(
|
||||
executor_id,
|
||||
mock_executor,
|
||||
["source"],
|
||||
shared_state,
|
||||
runner_ctx,
|
||||
|
||||
@@ -8,7 +8,7 @@ from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanE
|
||||
|
||||
from agent_framework import InMemoryCheckpointStorage, WorkflowBuilder
|
||||
from agent_framework._workflows._executor import Executor, handler
|
||||
from agent_framework._workflows._runner_context import InProcRunnerContext, Message
|
||||
from agent_framework._workflows._runner_context import InProcRunnerContext, Message, MessageType
|
||||
from agent_framework._workflows._shared_state import SharedState
|
||||
from agent_framework._workflows._workflow import Workflow
|
||||
from agent_framework._workflows._workflow_context import WorkflowContext
|
||||
@@ -127,7 +127,9 @@ async def test_span_creation_and_attributes(span_exporter: InMemorySpanExporter)
|
||||
OtelAttr.MESSAGE_DESTINATION_EXECUTOR_ID: "target-789",
|
||||
}
|
||||
with (
|
||||
create_processing_span("executor-456", "TestExecutor", "TestMessage") as processing_span,
|
||||
create_processing_span(
|
||||
"executor-456", "TestExecutor", str(MessageType.STANDARD), "TestMessage"
|
||||
) as processing_span,
|
||||
create_workflow_span(
|
||||
OtelAttr.MESSAGE_SEND_SPAN, sending_attributes, kind=trace.SpanKind.PRODUCER
|
||||
) as sending_span,
|
||||
@@ -155,7 +157,8 @@ async def test_span_creation_and_attributes(span_exporter: InMemorySpanExporter)
|
||||
assert processing_span.attributes is not None
|
||||
assert processing_span.attributes.get("executor.id") == "executor-456"
|
||||
assert processing_span.attributes.get("executor.type") == "TestExecutor"
|
||||
assert processing_span.attributes.get("message.type") == "TestMessage"
|
||||
assert processing_span.attributes.get("message.type") == str(MessageType.STANDARD)
|
||||
assert processing_span.attributes.get("message.payload_type") == "TestMessage"
|
||||
|
||||
# Check sending span
|
||||
sending_span = next(s for s in spans if s.name == "message.send")
|
||||
@@ -175,7 +178,7 @@ async def test_trace_context_handling(span_exporter: InMemorySpanExporter) -> No
|
||||
|
||||
# Test trace context propagation in messages
|
||||
workflow_ctx: WorkflowContext[str] = WorkflowContext(
|
||||
"test-executor",
|
||||
executor,
|
||||
["source"],
|
||||
shared_state,
|
||||
ctx,
|
||||
@@ -218,18 +221,20 @@ async def test_trace_context_handling(span_exporter: InMemorySpanExporter) -> No
|
||||
assert processing_span.attributes is not None
|
||||
assert processing_span.attributes.get("executor.id") == "test-executor"
|
||||
assert processing_span.attributes.get("executor.type") == "MockExecutor"
|
||||
assert processing_span.attributes.get("message.type") == "str"
|
||||
assert processing_span.attributes.get("message.type") == str(MessageType.STANDARD)
|
||||
assert processing_span.attributes.get("message.payload_type") == "str"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("enable_otel", [False], indirect=True)
|
||||
async def test_trace_context_disabled_when_tracing_disabled(enable_otel, span_exporter: InMemorySpanExporter) -> None:
|
||||
"""Test that no trace context is added when tracing is disabled."""
|
||||
# Tracing should be disabled by default
|
||||
executor = MockExecutor("test-executor")
|
||||
shared_state = SharedState()
|
||||
ctx = InProcRunnerContext()
|
||||
|
||||
workflow_ctx: WorkflowContext[str] = WorkflowContext(
|
||||
"test-executor",
|
||||
executor,
|
||||
["source"],
|
||||
shared_state,
|
||||
ctx,
|
||||
|
||||
@@ -1,7 +1,5 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
from typing_extensions import Never
|
||||
|
||||
@@ -10,8 +8,6 @@ from agent_framework import (
|
||||
ExecutorFailedEvent,
|
||||
InProcRunnerContext,
|
||||
RequestInfoEvent,
|
||||
RequestInfoExecutor,
|
||||
RequestInfoMessage,
|
||||
SharedState,
|
||||
Workflow,
|
||||
WorkflowBuilder,
|
||||
@@ -69,18 +65,26 @@ async def test_executor_failed_event_emitted_on_direct_execute():
|
||||
assert all(e.origin is WorkflowEventSource.FRAMEWORK for e in failed)
|
||||
|
||||
|
||||
class SimpleExecutor(Executor):
|
||||
"""Executor that does nothing, for testing."""
|
||||
|
||||
@handler
|
||||
async def run(self, msg: str, ctx: WorkflowContext[str]) -> None: # pragma: no cover
|
||||
await ctx.send_message(msg)
|
||||
|
||||
|
||||
class Requester(Executor):
|
||||
"""Executor that always requests external info to test idle-with-requests state."""
|
||||
|
||||
@handler
|
||||
async def ask(self, _: str, ctx: WorkflowContext[RequestInfoMessage]) -> None: # pragma: no cover
|
||||
await ctx.send_message(RequestInfoMessage())
|
||||
async def ask(self, _: str, ctx: WorkflowContext) -> None: # pragma: no cover
|
||||
await ctx.request_info("Mock request data", str, str)
|
||||
|
||||
|
||||
async def test_idle_with_pending_requests_status_streaming():
|
||||
req = Requester(id="req")
|
||||
rie = RequestInfoExecutor(id="rie")
|
||||
wf = WorkflowBuilder().set_start_executor(req).add_edge(req, rie).build()
|
||||
simple_executor = SimpleExecutor(id="simple")
|
||||
requester = Requester(id="req")
|
||||
wf = WorkflowBuilder().set_start_executor(simple_executor).add_edge(simple_executor, requester).build()
|
||||
|
||||
events = [ev async for ev in wf.run_stream("start")] # Consume stream fully
|
||||
|
||||
@@ -134,9 +138,9 @@ async def test_non_streaming_final_state_helpers():
|
||||
assert result1.get_final_state() == WorkflowRunState.IDLE
|
||||
|
||||
# Idle-with-pending-request case
|
||||
req = Requester(id="req")
|
||||
rie = RequestInfoExecutor(id="rie")
|
||||
wf2 = WorkflowBuilder().set_start_executor(req).add_edge(req, rie).build()
|
||||
simple_executor = SimpleExecutor(id="simple")
|
||||
requester = Requester(id="req")
|
||||
wf2 = WorkflowBuilder().set_start_executor(simple_executor).add_edge(simple_executor, requester).build()
|
||||
result2: WorkflowRunResult = await wf2.run("start")
|
||||
assert result2.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
|
||||
|
||||
@@ -151,32 +155,12 @@ async def test_run_includes_status_events_completed():
|
||||
|
||||
|
||||
async def test_run_includes_status_events_idle_with_requests():
|
||||
req = Requester(id="req2")
|
||||
rie = RequestInfoExecutor(id="rie2")
|
||||
wf = WorkflowBuilder().set_start_executor(req).add_edge(req, rie).build()
|
||||
simple_executor = SimpleExecutor(id="simple")
|
||||
requester = Requester(id="req2")
|
||||
wf = WorkflowBuilder().set_start_executor(simple_executor).add_edge(simple_executor, requester).build()
|
||||
result: WorkflowRunResult = await wf.run("start")
|
||||
timeline = result.status_timeline()
|
||||
assert timeline, "Expected status timeline in non-streaming run() results"
|
||||
assert len(timeline) >= 3
|
||||
assert timeline[-2].state == WorkflowRunState.IN_PROGRESS_PENDING_REQUESTS
|
||||
assert timeline[-1].state == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
|
||||
|
||||
|
||||
@dataclass
|
||||
class SnapshotRequest(RequestInfoMessage):
|
||||
prompt: str = ""
|
||||
draft: str = ""
|
||||
iteration: int = 0
|
||||
|
||||
|
||||
class SnapshotRequester(Executor):
|
||||
"""Executor that emits a rich RequestInfoMessage for persistence tests."""
|
||||
|
||||
def __init__(self, id: str, prompt: str, draft: str) -> None:
|
||||
super().__init__(id=id)
|
||||
self._prompt = prompt
|
||||
self._draft = draft
|
||||
|
||||
@handler
|
||||
async def ask(self, _: str, ctx: WorkflowContext[SnapshotRequest]) -> None: # pragma: no cover - simple helper
|
||||
await ctx.send_message(SnapshotRequest(prompt=self._prompt, draft=self._draft, iteration=1))
|
||||
|
||||
Reference in New Issue
Block a user