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:
committed by
GitHub
Unverified
parent
f6eadd412e
commit
943d92674e
+277
-354
@@ -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__":
|
||||
|
||||
+244
-226
@@ -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())
|
||||
|
||||
Reference in New Issue
Block a user