[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
2025-10-29 16:31:23 -07:00
committed by GitHub
Unverified
parent f6eadd412e
commit 943d92674e
54 changed files with 7532 additions and 6684 deletions
@@ -1,87 +1,59 @@
# Copyright (c) Microsoft. All rights reserved.
import asyncio
import uuid
from dataclasses import dataclass
from typing import Any
from typing import Literal
from agent_framework import (
Executor,
RequestInfoExecutor,
RequestInfoMessage,
RequestResponse,
RequestInfoEvent,
SubWorkflowRequestMessage,
SubWorkflowResponseMessage,
Workflow,
WorkflowBuilder,
WorkflowContext,
WorkflowExecutor,
handler,
response_handler,
)
from typing_extensions import Never
"""
Sample: Sub-workflow with parallel request handling by specialized interceptors
This sample demonstrates how to handle multiple parallel requests from a sub-workflow to
different executors in the main workflow.
This sample demonstrates how different parent executors can handle different types of requests
from the same sub-workflow using regular @handler methods for RequestInfoMessage subclasses.
Prerequisite:
- Understanding of sub-workflows.
- Understanding of requests and responses.
Prerequisites:
- No external services required (external handling simulated via `RequestInfoExecutor`).
This pattern is useful when a sub-workflow needs to interact with multiple external systems
or services.
Key architectural principles:
1. Specialized interceptors: Each parent executor handles only specific request types
2. Type-based routing: ResourceCache handles ResourceRequest, PolicyEngine handles PolicyCheckRequest
3. Automatic type filtering: Each interceptor only receives requests with matching types
4. Fallback forwarding: Unhandled requests are forwarded to external services
This sample implements a resource request distribution system where:
1. A sub-workflow generates requests for computing resources and policy checks.
2. The main workflow has executors that handle resource allocation and policy checking.
3. Responses are routed back to the sub-workflow, which collects and processes them.
The example simulates a resource allocation system where:
- Sub-workflow makes mixed requests for resources (CPU, memory) and policy checks
- ResourceCache executor intercepts ResourceRequest messages, serves from cache or forwards
- PolicyEngine executor intercepts PolicyCheckRequest messages, applies rules or forwards
- Each interceptor uses typed @handler methods for automatic filtering
The sub-workflow sends two types of requests:
- ResourceRequest: Requests for computing resources (e.g., CPU, memory).
- PolicyRequest: Requests to check resource allocation policies.
Flow visualization:
Coordinator
|
| Mixed list[resource + policy requests]
v
[ Sub-workflow: WorkflowExecutor(ResourceRequester) ]
|
| Emits different RequestInfoMessage types:
| - ResourceRequest
| - PolicyCheckRequest
v
Parent workflow routes to specialized handlers:
| |
| ResourceCache.handle_resource_request | PolicyEngine.handle_policy_request
| (@handler ResourceRequest) | (@handler PolicyCheckRequest)
v v
Cache hit/miss decision Policy allow/deny decision
| |
| RequestResponse OR forward | RequestResponse OR forward
v v
Back to sub-workflow <----------> External RequestInfoExecutor
|
v
External responses route back
The main workflow contains:
- ResourceAllocator: Simulates a system that allocates computing resources.
- PolicyEngine: Simulates a policy engine that approves or denies resource requests.
"""
# 1. Define domain-specific request/response types
@dataclass
class ResourceRequest(RequestInfoMessage):
class ComputingResourceRequest:
"""Request for computing resources."""
resource_type: str = "cpu" # cpu, memory, disk, etc.
amount: int = 1
priority: str = "normal" # low, normal, high
@dataclass
class PolicyCheckRequest(RequestInfoMessage):
"""Request to check resource allocation policy."""
resource_type: str = ""
amount: int = 0
policy_type: str = "quota" # quota, compliance, security
request_type: Literal["resource", "policy"]
resource_type: Literal["cpu", "memory", "disk", "gpu"]
amount: int
priority: Literal["low", "normal", "high"] | None = None
policy_type: Literal["quota", "security"] | None = None
@dataclass
@@ -102,340 +74,291 @@ class PolicyResponse:
@dataclass
class RequestFinished:
pass
class ResourceRequest:
"""Request for computing resources."""
resource_type: Literal["cpu", "memory", "disk", "gpu"]
amount: int
priority: Literal["low", "normal", "high"]
id: str = str(uuid.uuid4())
# 2. Implement the sub-workflow executor - makes resource and policy requests
class ResourceRequester(Executor):
"""Simple executor that requests resources and checks policies."""
@dataclass
class PolicyRequest:
"""Request to check resource allocation policy."""
def __init__(self):
super().__init__(id="resource_requester")
self._request_count = 0
policy_type: Literal["quota", "security"]
resource_type: Literal["cpu", "memory", "disk", "gpu"]
amount: int
id: str = str(uuid.uuid4())
def build_resource_request_distribution_workflow() -> Workflow:
class RequestDistribution(Executor):
"""Distributes computing resource requests to appropriate executors."""
@handler
async def distribute_requests(
self,
requests: list[ComputingResourceRequest],
ctx: WorkflowContext[ResourceRequest | PolicyRequest | int],
) -> None:
for req in requests:
if req.request_type == "resource":
if req.priority is None:
raise ValueError("Priority must be set for resource requests")
await ctx.send_message(ResourceRequest(req.resource_type, req.amount, req.priority))
elif req.request_type == "policy":
if req.policy_type is None:
raise ValueError("Policy type must be set for policy requests")
await ctx.send_message(PolicyRequest(req.policy_type, req.resource_type, req.amount))
else:
raise ValueError(f"Unknown request type: {req.request_type}")
# Notify the collector about the number of requests sent
await ctx.send_message(len(requests))
class ResourceRequester(Executor):
"""Handles resource allocation requests."""
@handler
async def run(self, request: ResourceRequest, ctx: WorkflowContext) -> None:
await ctx.request_info(request, ResourceRequest, ResourceResponse)
@response_handler
async def handle_response(
self, original_request: ResourceRequest, response: ResourceResponse, ctx: WorkflowContext[ResourceResponse]
) -> None:
print(f"Resource allocated: {response.allocated} {response.resource_type} from {response.source}")
await ctx.send_message(response)
class PolicyChecker(Executor):
"""Handles policy check requests."""
@handler
async def run(self, request: PolicyRequest, ctx: WorkflowContext) -> None:
await ctx.request_info(request, PolicyRequest, PolicyResponse)
@response_handler
async def handle_response(
self, original_request: PolicyRequest, response: PolicyResponse, ctx: WorkflowContext[PolicyResponse]
) -> None:
print(f"Policy check result: {response.approved} - {response.reason}")
await ctx.send_message(response)
class ResultCollector(Executor):
"""Collects and processes all responses."""
def __init__(self, id: str) -> None:
super().__init__(id)
self._request_count = 0
self._responses: list[ResourceResponse | PolicyResponse] = []
@handler
async def set_request_count(self, count: int, ctx: WorkflowContext) -> None:
if count <= 0:
raise ValueError("Request count must be positive")
self._request_count = count
@handler
async def collect(self, response: ResourceResponse | PolicyResponse, ctx: WorkflowContext[Never, str]) -> None:
self._responses.append(response)
print(f"Collected {len(self._responses)}/{self._request_count} responses")
if len(self._responses) == self._request_count:
# All responses received, process them
await ctx.yield_output(f"All {self._request_count} requests processed.")
elif len(self._responses) > self._request_count:
raise ValueError("Received more responses than expected")
orchestrator = RequestDistribution("orchestrator")
resource_requester = ResourceRequester("resource_requester")
policy_checker = PolicyChecker("policy_checker")
result_collector = ResultCollector("result_collector")
return (
WorkflowBuilder()
.set_start_executor(orchestrator)
.add_edge(orchestrator, resource_requester)
.add_edge(orchestrator, policy_checker)
.add_edge(resource_requester, result_collector)
.add_edge(policy_checker, result_collector)
.add_edge(orchestrator, result_collector) # For request count
.build()
)
class ResourceAllocator(Executor):
"""Simulates a system that allocates computing resources."""
def __init__(self, id: str) -> None:
super().__init__(id)
self._cache: dict[str, int] = {"cpu": 10, "memory": 50, "disk": 100}
# Record pending requests to match responses
self._pending_requests: dict[str, RequestInfoEvent] = {}
async def _handle_resource_request(self, request: ResourceRequest) -> ResourceResponse | None:
"""Allocates resources based on request and available cache."""
available = self._cache.get(request.resource_type, 0)
if available >= request.amount:
self._cache[request.resource_type] -= request.amount
return ResourceResponse(request.resource_type, request.amount, "cache")
return None
@handler
async def request_resources(
self,
requests: list[dict[str, Any]],
ctx: WorkflowContext[ResourceRequest | PolicyCheckRequest],
async def handle_subworkflow_request(
self, request: SubWorkflowRequestMessage, ctx: WorkflowContext[SubWorkflowResponseMessage]
) -> None:
"""Process a list of resource requests."""
print(f"🏭 Sub-workflow processing {len(requests)} requests")
self._request_count += len(requests)
for req_data in requests:
req_type = req_data.get("request_type", "resource")
request: ResourceRequest | PolicyCheckRequest
if req_type == "resource":
print(f" πŸ“¦ Requesting resource: {req_data.get('type', 'cpu')} x{req_data.get('amount', 1)}")
request = ResourceRequest(
resource_type=req_data.get("type", "cpu"),
amount=req_data.get("amount", 1),
priority=req_data.get("priority", "normal"),
)
# Send to parent workflow for interception - not to target_id
await ctx.send_message(request)
elif req_type == "policy":
print(
f" πŸ›‘οΈ Checking policy: {req_data.get('type', 'cpu')} x{req_data.get('amount', 1)} "
f"({req_data.get('policy_type', 'quota')})"
)
request = PolicyCheckRequest(
resource_type=req_data.get("type", "cpu"),
amount=req_data.get("amount", 1),
policy_type=req_data.get("policy_type", "quota"),
)
# Send to parent workflow for interception - not to target_id
await ctx.send_message(request)
@handler
async def handle_resource_response(
self,
response: RequestResponse[ResourceRequest, ResourceResponse],
ctx: WorkflowContext[Never, RequestFinished],
) -> None:
"""Handle resource allocation response."""
if response.data:
source_icon = "πŸͺ" if response.data.source == "cache" else "🌐"
print(
f"πŸ“¦ {source_icon} Sub-workflow received: {response.data.allocated} {response.data.resource_type} "
f"from {response.data.source}"
)
if self._collect_results():
# Yield completion result to the parent workflow.
await ctx.yield_output(RequestFinished())
@handler
async def handle_policy_response(
self,
response: RequestResponse[PolicyCheckRequest, PolicyResponse],
ctx: WorkflowContext[Never, RequestFinished],
) -> None:
"""Handle policy check response."""
if response.data:
status_icon = "βœ…" if response.data.approved else "❌"
print(
f"πŸ›‘οΈ {status_icon} Sub-workflow received policy response: "
f"{response.data.approved} - {response.data.reason}"
)
if self._collect_results():
# Yield completion result to the parent workflow.
await ctx.yield_output(RequestFinished())
def _collect_results(self) -> bool:
"""Collect and summarize results."""
self._request_count -= 1
print(f"πŸ“Š Sub-workflow completed request ({self._request_count} remaining)")
return self._request_count == 0
# 3. Implement the Resource Cache - Uses typed handler for ResourceRequest
class ResourceCache(Executor):
"""Interceptor that handles RESOURCE requests from cache using typed routing."""
# Use class attributes to avoid Pydantic assignment restrictions
cache: dict[str, int] = {"cpu": 10, "memory": 50, "disk": 100}
results: list[ResourceResponse] = []
def __init__(self):
super().__init__(id="resource_cache")
# Instance initialization only; state kept in class attributes as above
@handler
async def handle_resource_request(
self, request: ResourceRequest, ctx: WorkflowContext[RequestResponse[ResourceRequest, Any] | ResourceRequest]
) -> None:
"""Handle RESOURCE requests from sub-workflows and check cache first."""
resource_request = request
print(f"πŸͺ CACHE interceptor checking: {resource_request.amount} {resource_request.resource_type}")
available = self.cache.get(resource_request.resource_type, 0)
if available >= resource_request.amount:
# We can satisfy from cache
self.cache[resource_request.resource_type] -= resource_request.amount
response_data = ResourceResponse(
resource_type=resource_request.resource_type, allocated=resource_request.amount, source="cache"
)
print(f" βœ… Cache satisfied: {resource_request.amount} {resource_request.resource_type}")
self.results.append(response_data)
# Send response back to sub-workflow
response = RequestResponse(data=response_data, original_request=request, request_id=request.request_id)
await ctx.send_message(response, target_id=request.source_executor_id)
else:
# Cache miss - forward to external
print(f" ❌ Cache miss: need {resource_request.amount}, have {available} {resource_request.resource_type}")
await ctx.send_message(request)
@handler
async def collect_result(
self, response: RequestResponse[ResourceRequest, ResourceResponse], ctx: WorkflowContext
) -> None:
"""Collect results from external requests that were forwarded."""
if response.data and response.data.source != "cache": # Don't double-count our own results
self.results.append(response.data)
print(
f"πŸͺ 🌐 Cache received external response: {response.data.allocated} {response.data.resource_type} "
f"from {response.data.source}"
)
# 4. Implement the Policy Engine - Uses typed handler for PolicyCheckRequest
class PolicyEngine(Executor):
"""Interceptor that handles POLICY requests using typed routing."""
# Use class attributes for simple sample state
quota: dict[str, int] = {
"cpu": 5, # Only allow up to 5 CPU units
"memory": 20, # Only allow up to 20 memory units
"disk": 1000, # Liberal disk policy
}
results: list[PolicyResponse] = []
def __init__(self):
super().__init__(id="policy_engine")
# Instance initialization only; state kept in class attributes as above
@handler
async def handle_policy_request(
self,
request: PolicyCheckRequest,
ctx: WorkflowContext[RequestResponse[PolicyCheckRequest, Any] | PolicyCheckRequest],
) -> None:
"""Handle POLICY requests from sub-workflows and apply rules."""
policy_request = request
print(
f"πŸ›‘οΈ POLICY interceptor checking: {policy_request.amount} {policy_request.resource_type}, policy={policy_request.policy_type}"
)
quota_limit = self.quota.get(policy_request.resource_type, 0)
if policy_request.policy_type == "quota":
if policy_request.amount <= quota_limit:
response_data = PolicyResponse(approved=True, reason=f"Within quota ({quota_limit})")
print(f" βœ… Policy approved: {policy_request.amount} <= {quota_limit}")
self.results.append(response_data)
# Send response back to sub-workflow
response = RequestResponse(data=response_data, original_request=request, request_id=request.request_id)
await ctx.send_message(response, target_id=request.source_executor_id)
return
# Exceeds quota - forward to external for review
print(f" ❌ Policy exceeds quota: {policy_request.amount} > {quota_limit}, forwarding to external")
await ctx.send_message(request)
"""Handles requests from sub-workflows."""
source_event: RequestInfoEvent = request.source_event
if not isinstance(source_event.data, ResourceRequest):
return
# Unknown policy type - forward to external
print(f" ❓ Unknown policy type: {policy_request.policy_type}, forwarding")
await ctx.send_message(request)
request_payload: ResourceRequest = source_event.data
response = await self._handle_resource_request(request_payload)
if response:
await ctx.send_message(request.create_response(response))
else:
# Request cannot be fulfilled via cache, forward the request to external
self._pending_requests[request_payload.id] = source_event
await ctx.request_info(request_payload, ResourceRequest, ResourceResponse)
@handler
async def collect_policy_result(
self, response: RequestResponse[PolicyCheckRequest, PolicyResponse], ctx: WorkflowContext
@response_handler
async def handle_external_response(
self,
original_request: ResourceRequest,
response: ResourceResponse,
ctx: WorkflowContext[SubWorkflowResponseMessage],
) -> None:
"""Collect policy results from external requests that were forwarded."""
if response.data:
self.results.append(response.data)
print(f"πŸ›‘οΈ 🌐 Policy received external response: {response.data.approved} - {response.data.reason}")
"""Handles responses from external systems and routes them to the sub-workflow."""
print(f"External resource allocated: {response.allocated} {response.resource_type} from {response.source}")
source_event = self._pending_requests.pop(original_request.id, None)
if source_event is None:
raise ValueError("No matching pending request found for the resource response")
await ctx.send_message(SubWorkflowResponseMessage(data=response, source_event=source_event))
class Coordinator(Executor):
def __init__(self):
super().__init__(id="coordinator")
class PolicyEngine(Executor):
"""Simulates a policy engine that approves or denies resource requests."""
def __init__(self, id: str) -> None:
super().__init__(id)
self._quota: dict[str, int] = {
"cpu": 5, # Only allow up to 5 CPU units
"memory": 20, # Only allow up to 20 memory units
"disk": 1000, # Liberal disk policy
}
# Record pending requests to match responses
self._pending_requests: dict[str, RequestInfoEvent] = {}
@handler
async def start(self, requests: list[dict[str, Any]], ctx: WorkflowContext[list[dict[str, Any]]]) -> None:
"""Start the resource allocation process."""
await ctx.send_message(requests, target_id="resource_workflow")
async def handle_subworkflow_request(
self, request: SubWorkflowRequestMessage, ctx: WorkflowContext[SubWorkflowResponseMessage]
) -> None:
"""Handles requests from sub-workflows."""
source_event: RequestInfoEvent = request.source_event
if not isinstance(source_event.data, PolicyRequest):
return
@handler
async def handle_completion(self, completion: RequestFinished, ctx: WorkflowContext) -> None:
"""Handle sub-workflow completion.
request_payload: PolicyRequest = source_event.data
# Simple policy logic for demonstration
if request_payload.policy_type == "quota":
allowed_amount = self._quota.get(request_payload.resource_type, 0)
if request_payload.amount <= allowed_amount:
response = PolicyResponse(True, "Within quota limits")
else:
response = PolicyResponse(False, "Exceeds quota limits")
await ctx.send_message(request.create_response(response))
else:
# For other policy types, forward to external system
self._pending_requests[request_payload.id] = source_event
await ctx.request_info(request_payload, PolicyRequest, PolicyResponse)
It comes from the sub-workflow yielded output.
"""
print("🎯 Main workflow received completion.")
@response_handler
async def handle_external_response(
self,
original_request: PolicyRequest,
response: PolicyResponse,
ctx: WorkflowContext[SubWorkflowResponseMessage],
) -> None:
"""Handles responses from external systems and routes them to the sub-workflow."""
print(f"External policy check result: {response.approved} - {response.reason}")
source_event = self._pending_requests.pop(original_request.id, None)
if source_event is None:
raise ValueError("No matching pending request found for the policy response")
await ctx.send_message(SubWorkflowResponseMessage(data=response, source_event=source_event))
async def main() -> None:
"""Demonstrate parallel request interception patterns."""
print("πŸš€ Starting Sub-Workflow Parallel Request Interception Demo...")
print("=" * 60)
# Create executors in the main workflow
sub_workflow = build_resource_request_distribution_workflow()
resource_allocator = ResourceAllocator("resource_allocator")
policy_engine = PolicyEngine("policy_engine")
# 5. Create the sub-workflow
resource_requester = ResourceRequester()
sub_request_info = RequestInfoExecutor(id="sub_request_info")
sub_workflow = (
WorkflowBuilder()
.set_start_executor(resource_requester)
.add_edge(resource_requester, sub_request_info)
.add_edge(sub_request_info, resource_requester)
.build()
# Create the WorkflowExecutor for the sub-workflow
# Setting allow_direct_output=True to let the sub-workflow output directly.
# This is because the sub-workflow is the both the entry point and the exit
# point of the main workflow.
sub_workflow_executor = WorkflowExecutor(
sub_workflow,
"sub_workflow_executor",
allow_direct_output=True,
)
# 6. Create parent workflow with PROPER interceptor pattern
cache = ResourceCache() # Intercepts ResourceRequest
policy = PolicyEngine() # Intercepts PolicyCheckRequest (different type!)
workflow_executor = WorkflowExecutor(sub_workflow, id="resource_workflow")
main_request_info = RequestInfoExecutor(id="main_request_info")
# Create a simple coordinator that starts the process
coordinator = Coordinator()
# TYPED ROUTING: Each executor handles specific typed RequestInfoMessage messages
# Build the main workflow
main_workflow = (
WorkflowBuilder()
.set_start_executor(coordinator)
.add_edge(coordinator, workflow_executor) # Start sub-workflow
.add_edge(workflow_executor, coordinator) # Sub-workflow completion back to coordinator
.add_edge(workflow_executor, cache) # WorkflowExecutor sends ResourceRequest to cache
.add_edge(workflow_executor, policy) # WorkflowExecutor sends PolicyCheckRequest to policy
.add_edge(cache, workflow_executor) # Cache sends RequestResponse back
.add_edge(policy, workflow_executor) # Policy sends RequestResponse back
.add_edge(cache, main_request_info) # Cache forwards ResourceRequest to external
.add_edge(policy, main_request_info) # Policy forwards PolicyCheckRequest to external
.add_edge(main_request_info, workflow_executor) # External responses back to sub-workflow
.set_start_executor(sub_workflow_executor)
.add_edge(sub_workflow_executor, resource_allocator)
.add_edge(resource_allocator, sub_workflow_executor)
.add_edge(sub_workflow_executor, policy_engine)
.add_edge(policy_engine, sub_workflow_executor)
.build()
)
# 7. Test with various requests (mixed resource and policy)
# Test requests
test_requests = [
{"request_type": "resource", "type": "cpu", "amount": 2, "priority": "normal"}, # Cache hit
{"request_type": "policy", "type": "cpu", "amount": 3, "policy_type": "quota"}, # Policy hit
{"request_type": "resource", "type": "memory", "amount": 15, "priority": "normal"}, # Cache hit
{"request_type": "policy", "type": "memory", "amount": 100, "policy_type": "quota"}, # Policy miss -> external
{"request_type": "resource", "type": "gpu", "amount": 1, "priority": "high"}, # Cache miss -> external
{"request_type": "policy", "type": "disk", "amount": 500, "policy_type": "quota"}, # Policy hit
{"request_type": "policy", "type": "cpu", "amount": 1, "policy_type": "security"}, # Unknown policy -> external
ComputingResourceRequest("resource", "cpu", 2, priority="normal"), # cache hit
ComputingResourceRequest("policy", "cpu", 3, policy_type="quota"), # policy hit
ComputingResourceRequest("resource", "memory", 15, priority="normal"), # cache hit
ComputingResourceRequest("policy", "memory", 100, policy_type="quota"), # policy miss -> external
ComputingResourceRequest("resource", "gpu", 1, priority="high"), # cache miss -> external
ComputingResourceRequest("policy", "disk", 500, policy_type="quota"), # policy hit
ComputingResourceRequest("policy", "cpu", 1, policy_type="security"), # unknown policy -> external
]
print(f"πŸ§ͺ Testing with {len(test_requests)} mixed requests:")
for i, req in enumerate(test_requests, 1):
req_icon = "πŸ“¦" if req["request_type"] == "resource" else "πŸ›‘οΈ"
print(
f" {i}. {req_icon} {req['type']} x{req['amount']} "
f"({req.get('priority', req.get('policy_type', 'default'))})"
)
print("=" * 70)
# Run the workflow
print(f"πŸ§ͺ Testing with {len(test_requests)} mixed requests.")
print("πŸš€ Starting main workflow...")
run_result = await main_workflow.run(test_requests)
# 8. Run the workflow
print("🎬 Running workflow...")
events = await main_workflow.run(test_requests)
# Handle request info events
request_info_events = run_result.get_request_info_events()
if request_info_events:
print(f"\nπŸ” Handling {len(request_info_events)} request info events...\n")
# 9. Handle any external requests that couldn't be intercepted
request_events = events.get_request_info_events()
if request_events:
print(f"\n🌐 Handling {len(request_events)} external request(s)...")
external_responses: dict[str, Any] = {}
for event in request_events:
responses: dict[str, ResourceResponse | PolicyResponse] = {}
for event in request_info_events:
if isinstance(event.data, ResourceRequest):
# Handle ResourceRequest - create ResourceResponse
# Simulate external resource allocation
resource_response = ResourceResponse(
resource_type=event.data.resource_type, allocated=event.data.amount, source="external_provider"
)
external_responses[event.request_id] = resource_response
print(f" 🏭 External provider: {resource_response.allocated} {resource_response.resource_type}")
elif isinstance(event.data, PolicyCheckRequest):
# Handle PolicyCheckRequest - create PolicyResponse
policy_response = PolicyResponse(approved=True, reason="External policy service approved")
external_responses[event.request_id] = policy_response
print(f" πŸ”’ External policy: {'βœ… APPROVED' if policy_response.approved else '❌ DENIED'}")
responses[event.request_id] = resource_response
elif isinstance(event.data, PolicyRequest):
# Simulate external policy check
response = PolicyResponse(True, "External system approved")
responses[event.request_id] = response
else:
print(f"Unknown request info event data type: {type(event.data)}")
await main_workflow.send_responses(external_responses)
run_result = await main_workflow.send_responses(responses)
outputs = run_result.get_outputs()
if outputs:
print("\nWorkflow completed with outputs:")
for output in outputs:
print(f"- {output}")
else:
print("\n🎯 All requests were intercepted internally!")
# 10. Show results and analysis
print("\n" + "=" * 70)
print("πŸ“Š RESULTS ANALYSIS")
print("=" * 70)
print(f"\nπŸͺ Cache Results ({len(cache.results)} handled):")
for result in cache.results:
print(f" βœ… {result.allocated} {result.resource_type} from {result.source}")
print(f"\nπŸ›‘οΈ Policy Results ({len(policy.results)} handled):")
for result in policy.results:
status_icon = "βœ…" if result.approved else "❌"
print(f" {status_icon} Approved: {result.approved} - {result.reason}")
print("\nπŸ’Ύ Final Cache State:")
for resource, amount in cache.cache.items():
print(f" πŸ“¦ {resource}: {amount} remaining")
print("\nπŸ“ˆ Summary:")
print(f" 🎯 Total requests: {len(test_requests)}")
print(f" πŸͺ Resource requests handled: {len(cache.results)}")
print(f" πŸ›‘οΈ Policy requests handled: {len(policy.results)}")
print(f" 🌐 External requests: {len(request_events) if request_events else 0}")
print("\n" + "=" * 70)
raise RuntimeError("Workflow did not produce an output.")
if __name__ == "__main__":
@@ -5,289 +5,307 @@ from dataclasses import dataclass
from agent_framework import (
Executor,
RequestInfoExecutor,
RequestInfoMessage,
RequestResponse,
SubWorkflowRequestMessage,
SubWorkflowResponseMessage,
Workflow,
WorkflowBuilder,
WorkflowContext,
WorkflowExecutor,
WorkflowOutputEvent,
handler,
response_handler,
)
from typing_extensions import Never
"""
Sample: Sub-Workflows with Request Interception
This sample demonstrates how to handle request from the sub-workflow in the main workflow.
This sample shows how to:
1. Create workflows that execute other workflows as sub-workflows
2. Intercept requests from sub-workflows using an executor with @handler for RequestInfoMessage subclasses
3. Conditionally handle or forward requests using RequestResponse messages
4. Handle external requests that are forwarded by the parent workflow
5. Proper request/response correlation for concurrent processing
Prerequisite:
- Understanding of sub-workflows.
- Understanding of requests and responses.
The example simulates an email validation system where:
- Sub-workflows validate multiple email addresses concurrently
- Parent workflows can intercept domain check requests for optimization
- Known domains (example.com, company.com) are approved locally
- Unknown domains (unknown.org) are forwarded to external services
- Request correlation ensures each email gets the correct domain check response
- External domain check requests are processed and responses routed back correctly
This pattern is useful when you want to reuse a workflow that makes requests to an external system,
but you want to intercept those requests in the main workflow and handle them without further propagation
to the external system.
Key concepts demonstrated:
- WorkflowExecutor: Wraps a workflow to make it behave as an executor
- RequestInfoMessage handler: @handler method to intercept sub-workflow requests
- Request correlation: Using request_id and source_executor_id to match responses with original requests
- Concurrent processing: Multiple emails processed simultaneously without interference
- External request routing: RequestInfoExecutor handles forwarded external requests
- Sub-workflow isolation: Sub-workflows work normally without knowing they're nested
- Sub-workflows complete by yielding outputs when validation is finished
Prerequisites:
- No external services required (external calls are simulated via `RequestInfoExecutor`).
Simple flow visualization:
Parent Orchestrator (handles DomainCheckRequest)
|
| EmailValidationRequest(email) x3 (concurrent)
v
[ Sub-workflow: WorkflowExecutor(EmailValidator) ]
|
| DomainCheckRequest(domain) with request_id and source_executor_id
v
Interception? yes -> handled locally with RequestResponse(data=True)
no -> forwarded to RequestInfoExecutor -> external service
|
v
Response routed back to sub-workflow using source_executor_id
This sample implements a smart email delivery system that validates email addresses before sending emails.
1. We will start by creating a workflow that validates email addresses in a sequential manner. The validation
consists of three steps: sanitization, format validation, and domain validation. The domain validation
step will involve checking if the email domain is valid by making a request to an external system.
2. Then we will create a main workflow that uses the email validation workflow as a sub-workflow. The main
workflow will intercept the domain validation requests from the sub-workflow and handle them internally
without propagating them to an external system.
3. Once the email address is validated, the main workflow will proceed to send the email if the address is valid,
or block the email if the address is invalid.
"""
# 1. Define domain-specific message types
@dataclass
class EmailValidationRequest:
"""Request to validate an email address."""
class SanitizedEmailResult:
"""Result of email sanitization and validation.
email: str
The properties get built up as the email address goes through
the validation steps in the workflow.
"""
@dataclass
class DomainCheckRequest(RequestInfoMessage):
"""Request to check if a domain is approved."""
domain: str = ""
@dataclass
class ValidationResult:
"""Result of email validation."""
email: str
original: str
sanitized: str
is_valid: bool
reason: str
# 2. Implement the sub-workflow executor (completely standard)
class EmailValidator(Executor):
"""Validates email addresses - doesn't know it's in a sub-workflow."""
def build_email_address_validation_workflow() -> Workflow:
"""Build an email address validation workflow.
def __init__(self) -> None:
"""Initialize the EmailValidator executor."""
super().__init__(id="email_validator")
# Use a dict to track multiple pending emails by request_id
self._pending_emails: dict[str, str] = {}
This workflow consists of three steps (each is represented by an executor):
1. Sanitize the email address, such as removing leading/trailing spaces.
2. Validate the email address format, such as checking for "@" and domain.
3. Extract the domain from the email address and request domain validation,
after which it completes with the final result.
"""
@handler
async def validate_request(
self,
request: EmailValidationRequest,
ctx: WorkflowContext[DomainCheckRequest | ValidationResult, ValidationResult],
) -> None:
"""Validate an email address."""
print(f"πŸ” Sub-workflow validating email: {request.email}")
class EmailSanitizer(Executor):
"""Sanitize email address by trimming spaces."""
# Extract domain
domain = request.email.split("@")[1] if "@" in request.email else ""
@handler
async def handle(self, email_address: str, ctx: WorkflowContext[SanitizedEmailResult]) -> None:
"""Trim leading and trailing spaces from the email address.
if not domain:
print(f"❌ Invalid email format: {request.email}")
result = ValidationResult(email=request.email, is_valid=False, reason="Invalid email format")
await ctx.yield_output(result)
return
This executor doesn't produce any workflow output, but sends the sanitized
email address to the next executor in the workflow.
"""
sanitized = email_address.strip()
print(f"βœ‚οΈ Sanitized email address: '{sanitized}'")
await ctx.send_message(SanitizedEmailResult(original=email_address, sanitized=sanitized, is_valid=False))
print(f"🌐 Sub-workflow requesting domain check for: {domain}")
# Request domain check
domain_check = DomainCheckRequest(domain=domain)
# Store the pending email with the request_id for correlation
self._pending_emails[domain_check.request_id] = request.email
await ctx.send_message(domain_check, target_id="email_request_info")
class EmailFormatValidator(Executor):
"""Validate email address format."""
@handler
async def handle_domain_response(
self,
response: RequestResponse[DomainCheckRequest, bool],
ctx: WorkflowContext[ValidationResult, ValidationResult],
) -> None:
"""Handle domain check response from RequestInfo with correlation."""
approved = bool(response.data)
domain = (
response.original_request.domain
if (hasattr(response, "original_request") and response.original_request)
else "unknown"
)
print(f"πŸ“¬ Sub-workflow received domain response for '{domain}': {approved}")
@handler
async def handle(
self,
partial_result: SanitizedEmailResult,
ctx: WorkflowContext[SanitizedEmailResult, SanitizedEmailResult],
) -> None:
"""Validate the email address format.
# Find the corresponding email using the request_id
request_id = (
response.original_request.request_id
if (hasattr(response, "original_request") and response.original_request)
else None
)
if request_id and request_id in self._pending_emails:
email = self._pending_emails.pop(request_id) # Remove from pending
result = ValidationResult(
email=email,
is_valid=approved,
reason="Domain approved" if approved else "Domain not approved",
This executor can potentially produce a workflow output (False if the format is invalid).
When the format is valid, it sends the validated email address to the next executor in the workflow.
"""
if "@" not in partial_result.sanitized or "." not in partial_result.sanitized.split("@")[-1]:
print(f"❌ Invalid email format: '{partial_result.sanitized}'")
await ctx.yield_output(
SanitizedEmailResult(
original=partial_result.original, sanitized=partial_result.sanitized, is_valid=False
)
)
return
print(f"βœ… Validated email format: '{partial_result.sanitized}'")
await ctx.send_message(
SanitizedEmailResult(
original=partial_result.original, sanitized=partial_result.sanitized, is_valid=False
)
)
print(f"βœ… Sub-workflow completing validation for: {email}")
await ctx.yield_output(result)
class DomainValidator(Executor):
"""Validate email domain."""
def __init__(self, id: str):
super().__init__(id=id)
self._pending_domains: dict[str, SanitizedEmailResult] = {}
@handler
async def handle(self, partial_result: SanitizedEmailResult, ctx: WorkflowContext) -> None:
"""Extract the domain from the email address and request domain validation.
This executor doesn't produce any workflow output, but sends a domain validation request
to an external system to user for validation.
"""
domain = partial_result.sanitized.split("@")[-1]
print(f"πŸ” Validating domain: '{domain}'")
self._pending_domains[domain] = partial_result
# Send a request to the external system via the request_info mechanism
await ctx.request_info(domain, str, bool)
@response_handler
async def handle_domain_validation_response(
self, original_request: str, is_valid: bool, ctx: WorkflowContext[Never, SanitizedEmailResult]
) -> None:
"""Handle the domain validation response.
This method receives the response from the external system and yields the final
validation result (True if both format and domain are valid, False otherwise).
"""
if original_request not in self._pending_domains:
raise ValueError(f"Received response for unknown domain: '{original_request}'")
partial_result = self._pending_domains.pop(original_request)
if is_valid:
print(f"βœ… Domain '{original_request}' is valid.")
await ctx.yield_output(
SanitizedEmailResult(
original=partial_result.original, sanitized=partial_result.sanitized, is_valid=True
)
)
else:
print(f"❌ Domain '{original_request}' is invalid.")
await ctx.yield_output(
SanitizedEmailResult(
original=partial_result.original, sanitized=partial_result.sanitized, is_valid=False
)
)
# Build the workflow
sanitizer = EmailSanitizer(id="email_sanitizer")
format_validator = EmailFormatValidator(id="email_format_validator")
domain_validator = DomainValidator(id="domain_validator")
return (
WorkflowBuilder()
.set_start_executor(sanitizer)
.add_edge(sanitizer, format_validator)
.add_edge(format_validator, domain_validator)
.build()
)
@dataclass
class Email:
recipient: str
subject: str
body: str
# 3. Implement the parent workflow with request interception
class SmartEmailOrchestrator(Executor):
"""Parent orchestrator that can intercept domain checks."""
"""Orchestrates email address validation using a sub-workflow."""
approved_domains: set[str] = set()
def __init__(self, approved_domains: set[str] | None = None):
"""Initialize the SmartEmailOrchestrator with approved domains.
def __init__(self, id: str, approved_domains: set[str]):
"""Initialize the orchestrator with a set of approved domains.
Args:
approved_domains: Set of pre-approved domains, defaults to example.com, test.org, company.com
id: The executor ID.
approved_domains: A set of domains that are considered valid.
"""
super().__init__(id="email_orchestrator", approved_domains=approved_domains)
self._results: list[ValidationResult] = []
super().__init__(id=id)
self._approved_domains = approved_domains
# Keep track of previously approved and disapproved recipients
self._approved_recipients: set[str] = set()
self._disapproved_recipients: set[str] = set()
# Record pending emails waiting for validation results
self._pending_emails: dict[str, Email] = {}
@handler
async def start_validation(self, emails: list[str], ctx: WorkflowContext[EmailValidationRequest]) -> None:
"""Start validating a batch of emails."""
print(f"πŸ“§ Starting validation of {len(emails)} email addresses")
print("=" * 60)
for email in emails:
print(f"πŸ“€ Sending '{email}' to sub-workflow for validation")
request = EmailValidationRequest(email=email)
await ctx.send_message(request, target_id="email_validator_workflow")
async def run(self, email: Email, ctx: WorkflowContext[Email | str, bool]) -> None:
"""Start the email delivery process.
This handler receives an Email object. If the recipient has been previously approved,
it sends the email object to the next executor to handle delivery. If the recipient
has been previously disapproved, it yields False as the final result. Otherwise,
it sends the recipient email address to the sub-workflow for validation.
"""
recipient = email.recipient
if recipient in self._approved_recipients:
print(f"πŸ“§ Recipient '{recipient}' has been previously approved.")
await ctx.send_message(email)
return
if recipient in self._disapproved_recipients:
print(f"🚫 Blocking email to previously disapproved recipient: '{recipient}'")
await ctx.yield_output(False)
return
print(f"πŸ” Validating new recipient email address: '{recipient}'")
self._pending_emails[recipient] = email
await ctx.send_message(recipient)
@handler
async def handle_domain_request(
self,
request: DomainCheckRequest,
ctx: WorkflowContext[RequestResponse[DomainCheckRequest, bool] | DomainCheckRequest],
async def handler_domain_validation_request(
self, request: SubWorkflowRequestMessage, ctx: WorkflowContext[SubWorkflowResponseMessage]
) -> None:
"""Handle requests from sub-workflows."""
print(f"πŸ” Parent intercepting domain check for: {request.domain}")
"""Handle requests from the sub-workflow for domain validation.
if request.domain in self.approved_domains:
print(f"βœ… Domain '{request.domain}' is pre-approved locally!")
# 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:
print(f"❓ Domain '{request.domain}' unknown, forwarding to external service...")
# Forward to external handler
await ctx.send_message(request)
Note that the message type must be SubWorkflowRequestMessage to intercept the request. And
the response must be sent back using SubWorkflowResponseMessage to route the response
back to the sub-workflow.
"""
if not isinstance(request.source_event.data, str):
raise TypeError(f"Expected domain string, got {type(request.source_event.data)}")
domain = request.source_event.data
is_valid = domain in self._approved_domains
print(f"🌐 External domain validation for '{domain}': {'valid' if is_valid else 'invalid'}")
await ctx.send_message(request.create_response(is_valid), target_id=request.executor_id)
@handler
async def collect_result(self, result: ValidationResult, ctx: WorkflowContext) -> None:
"""Collect validation results. It comes from the sub-workflow yielded output."""
status_icon = "βœ…" if result.is_valid else "❌"
print(f"πŸ“₯ {status_icon} Validation result: {result.email} -> {result.reason}")
self._results.append(result)
async def handle_validation_result(self, result: SanitizedEmailResult, ctx: WorkflowContext[Email, bool]) -> None:
"""Handle the email address validation result.
@property
def results(self) -> list[ValidationResult]:
"""Get the collected validation results."""
return self._results
This handler receives the validation result from the sub-workflow.
If the email address is valid, it adds the recipient to the approved list
and sends the email object to the next executor to handle delivery.
If the email address is invalid, it adds the recipient to the disapproved list
and yields False as the final result.
"""
email = self._pending_emails.pop(result.original)
email.recipient = result.sanitized # Use the sanitized email address
if result.is_valid:
print(f"βœ… Email address '{result.original}' is valid.")
self._approved_recipients.add(result.original)
await ctx.send_message(email)
else:
print(f"🚫 Email address '{result.original}' is invalid. Blocking email.")
self._disapproved_recipients.add(result.original)
await ctx.yield_output(False)
async def run_example() -> None:
"""Run the sub-workflow example."""
print("πŸš€ Setting up sub-workflow with request interception...")
print()
class EmailDelivery(Executor):
"""Simulates email delivery."""
# 4. Build the sub-workflow
email_validator = EmailValidator()
# Match the target_id used in EmailValidator ("email_request_info")
request_info = RequestInfoExecutor(id="email_request_info")
@handler
async def handle(self, email: Email, ctx: WorkflowContext[Never, bool]) -> None:
"""Simulate sending the email and yield True as the final result."""
print(f"πŸ“€ Sending email to '{email.recipient}' with subject '{email.subject}'")
await asyncio.sleep(1) # Simulate network delay
print(f"βœ… Email sent to '{email.recipient}' successfully.")
await ctx.yield_output(True)
validation_workflow = (
WorkflowBuilder()
.set_start_executor(email_validator)
.add_edge(email_validator, request_info)
.add_edge(request_info, email_validator)
.build()
)
# 5. Build the parent workflow with interception
orchestrator = SmartEmailOrchestrator(approved_domains={"example.com", "company.com"})
workflow_executor = WorkflowExecutor(validation_workflow, id="email_validator_workflow")
# Add a RequestInfoExecutor to handle forwarded external requests
main_request_info = RequestInfoExecutor(id="main_request_info")
async def main() -> None:
# A list of approved domains
approved_domains = {"example.com", "company.com"}
main_workflow = (
# Create executors in the main workflow
orchestrator = SmartEmailOrchestrator(id="smart_email_orchestrator", approved_domains=approved_domains)
email_delivery = EmailDelivery(id="email_delivery")
# Create the sub-workflow for email address validation
validation_workflow = build_email_address_validation_workflow()
validation_workflow_executor = WorkflowExecutor(validation_workflow, id="email_validation_workflow")
# Build the main workflow
workflow = (
WorkflowBuilder()
.set_start_executor(orchestrator)
.add_edge(orchestrator, workflow_executor)
.add_edge(workflow_executor, orchestrator) # For ValidationResult collection and request interception
# Add edges for external request handling
.add_edge(orchestrator, main_request_info)
.add_edge(main_request_info, workflow_executor) # Route external responses to sub-workflow
.add_edge(orchestrator, validation_workflow_executor)
.add_edge(validation_workflow_executor, orchestrator)
.add_edge(orchestrator, email_delivery)
.build()
)
# 6. Prepare test inputs: known domain, unknown domain
test_emails = [
"user@example.com", # Should be intercepted and approved
"admin@company.com", # Should be intercepted and approved
"guest@unknown.org", # Should be forwarded externally
Email(recipient="user1@example.com", subject="Hello User1", body="This is a test email."),
Email(recipient=" user2@invalid", subject="Hello User2", body="This is a test email."),
Email(recipient=" user3@company.com ", subject="Hello User3", body="This is a test email."),
Email(recipient="user4@unknown.com", subject="Hello User4", body="This is a test email."),
# Re-send to an approved recipient
Email(recipient="user1@example.com", subject="Hello User1", body="This is a test email."),
# Re-send to a disapproved recipient
Email(recipient=" user2@invalid", subject="Hello User2", body="This is a test email."),
]
# 7. Run the workflow
result = await main_workflow.run(test_emails)
# 8. Handle any external requests
request_events = result.get_request_info_events()
if request_events:
print(f"\n🌐 Handling {len(request_events)} external request(s)...")
for event in request_events:
if event.data and hasattr(event.data, "domain"):
print(f"πŸ” External domain check needed for: {event.data.domain}")
# Simulate external responses
external_responses: dict[str, bool] = {}
for event in request_events:
# Simulate external domain checking
if event.data and hasattr(event.data, "domain"):
domain = event.data.domain
# Let's say unknown.org is actually approved externally
approved = domain == "unknown.org"
print(f"🌐 External service response for '{domain}': {'APPROVED' if approved else 'REJECTED'}")
external_responses[event.request_id] = approved
# 9. Send external responses
await main_workflow.send_responses(external_responses)
else:
print("\n🎯 All requests were intercepted and handled locally!")
# 10. Display final summary
print("\nπŸ“Š Final Results Summary:")
print("=" * 60)
for result in orchestrator.results:
status = "βœ… VALID" if result.is_valid else "❌ INVALID"
print(f"{status} {result.email}: {result.reason}")
print(f"\n🏁 Processed {len(orchestrator.results)} emails total")
# Execute the workflow
for email in test_emails:
print(f"\nπŸš€ Processing email to '{email.recipient}'")
async for event in workflow.run_stream(email):
if isinstance(event, WorkflowOutputEvent):
print(f"πŸŽ‰ Final result for '{email.recipient}': {'Delivered' if event.data else 'Blocked'}")
if __name__ == "__main__":
asyncio.run(run_example())
asyncio.run(main())