diff --git a/python/packages/main/agent_framework/_workflow/__init__.py b/python/packages/main/agent_framework/_workflow/__init__.py index c97927c0d5..86a19fdaf6 100644 --- a/python/packages/main/agent_framework/_workflow/__init__.py +++ b/python/packages/main/agent_framework/_workflow/__init__.py @@ -51,10 +51,7 @@ from ._executor import ( RequestInfoExecutor, RequestInfoMessage, RequestResponse, - SubWorkflowRequestInfo, - SubWorkflowResponse, handler, - intercepts_request, ) from ._function_executor import FunctionExecutor, executor from ._magentic import ( @@ -157,8 +154,6 @@ __all__ = [ "SharedState", "SingleEdgeGroup", "StandardMagenticManager", - "SubWorkflowRequestInfo", - "SubWorkflowResponse", "SwitchCaseEdgeGroup", "SwitchCaseEdgeGroupCase", "SwitchCaseEdgeGroupDefault", @@ -185,7 +180,6 @@ __all__ = [ "create_edge_runner", "executor", "handler", - "intercepts_request", "validate_workflow_graph", ] diff --git a/python/packages/main/agent_framework/_workflow/__init__.pyi b/python/packages/main/agent_framework/_workflow/__init__.pyi index 349e6463a3..36104b5fc4 100644 --- a/python/packages/main/agent_framework/_workflow/__init__.pyi +++ b/python/packages/main/agent_framework/_workflow/__init__.pyi @@ -47,10 +47,7 @@ from ._executor import ( RequestInfoExecutor, RequestInfoMessage, RequestResponse, - SubWorkflowRequestInfo, - SubWorkflowResponse, handler, - intercepts_request, ) from ._function_executor import FunctionExecutor, executor from ._magentic import ( @@ -153,8 +150,6 @@ __all__ = [ "SharedState", "SingleEdgeGroup", "StandardMagenticManager", - "SubWorkflowRequestInfo", - "SubWorkflowResponse", "SwitchCaseEdgeGroup", "SwitchCaseEdgeGroupCase", "SwitchCaseEdgeGroupDefault", @@ -181,6 +176,5 @@ __all__ = [ "create_edge_runner", "executor", "handler", - "intercepts_request", "validate_workflow_graph", ] diff --git a/python/packages/main/agent_framework/_workflow/_executor.py b/python/packages/main/agent_framework/_workflow/_executor.py index 161adedf4b..4d3f794edb 100644 --- a/python/packages/main/agent_framework/_workflow/_executor.py +++ b/python/packages/main/agent_framework/_workflow/_executor.py @@ -9,7 +9,7 @@ import uuid from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from dataclasses import asdict, dataclass, field, fields, is_dataclass from textwrap import shorten -from typing import Any, ClassVar, Generic, TypeVar, cast, overload +from typing import Any, ClassVar, Generic, TypeVar, cast from pydantic import Field @@ -104,20 +104,6 @@ class Executor(AFBaseModel): ``` Access via the `workflow_output_types` property. - ### Request Types - The types of sub-workflow requests an executor can intercept via `@intercepts_request`: - ```python - class ParentExecutor(Executor): - @intercepts_request - async def handle_request( - self, - request: MyRequest, - ctx: WorkflowContext, - ) -> RequestResponse[MyRequest, str]: - # Can intercept 'MyRequest' from sub-workflows - ``` - Access via the `request_types` property. - ## Handler Discovery Executors discover their capabilities through decorated methods: @@ -130,17 +116,21 @@ class Executor(AFBaseModel): await ctx.send_message(message.upper()) ``` - ### @intercepts_request Decorator - Marks methods that intercept sub-workflow requests: + ### Sub-workflow Request Interception + Use @handler methods to intercept sub-workflow requests: ```python class ParentExecutor(Executor): - @intercepts_request - async def check_domain( - self, request: DomainRequest, ctx: WorkflowContext - ) -> RequestResponse[DomainRequest, bool]: + @handler + async def handle_domain_request( + self, + request: DomainRequest, # Subclass of RequestInfoMessage + ctx: WorkflowContext[RequestResponse[RequestInfoMessage, Any] | DomainRequest], + ) -> None: if self.is_allowed(request.domain): - return RequestResponse.handled(True) - return RequestResponse.forward() + response = RequestResponse(data=True, original_request=request, request_id=request.request_id) + await ctx.send_message(response, target_id=request.source_executor_id) + else: + await ctx.send_message(request) # Forward to external ``` ## Context Types @@ -196,7 +186,7 @@ class Executor(AFBaseModel): ## Implementation Notes - Do not call `execute()` directly - it's invoked by the workflow engine - Do not override `execute()` - define handlers using decorators instead - - Each executor must have at least one `@handler` or `@intercepts_request` method + - Each executor must have at least one `@handler` method - Handler method signatures are validated at initialization time """ @@ -226,15 +216,13 @@ class Executor(AFBaseModel): super().__init__(**kwargs) self._handlers: dict[type, Callable[[Any, WorkflowContext[Any, Any]], Awaitable[None]]] = {} - self._request_interceptors: dict[type | str, list[dict[str, Any]]] = {} self._handler_specs: list[dict[str, Any]] = [] self._discover_handlers() - if not self._handlers and not self._request_interceptors: + if not self._handlers: raise ValueError( f"Executor {self.__class__.__name__} has no handlers defined. " - "Please define at least one handler using the @handler decorator " - "or @intercepts_request decorator." + "Please define at least one handler using the @handler decorator." ) async def execute( @@ -276,7 +264,6 @@ class Executor(AFBaseModel): source_span_ids=source_span_ids, ): # Find the handler and handler spec that matches the message type. - # This includes the self._handle_sub_workflow_request handler for SubWorkflowRequestInfo. handler: Callable[[Any, WorkflowContext[Any, Any]], Awaitable[None]] | None = None ctx_annotation = None for message_type in self._handlers: @@ -351,202 +338,39 @@ class Executor(AFBaseModel): ) def _discover_handlers(self) -> None: - """Discover message handlers and request interceptors in the executor class.""" + """Discover message handlers in the executor class.""" # Use __class__.__dict__ to avoid accessing pydantic's dynamic attributes for attr_name in dir(self.__class__): try: attr = getattr(self.__class__, attr_name) - if callable(attr): - # Discover @handler methods - if hasattr(attr, "_handler_spec"): - handler_spec = attr._handler_spec # type: ignore - message_type = handler_spec["message_type"] + # Discover @handler methods + if callable(attr) and hasattr(attr, "_handler_spec"): + handler_spec = attr._handler_spec # type: ignore + message_type = handler_spec["message_type"] - # Keep full generic types for handler registration to avoid conflicts - # Different RequestResponse[T, U] specializations are distinct handler types + # Keep full generic types for handler registration to avoid conflicts + # Different RequestResponse[T, U] specializations are distinct handler types - if self._handlers.get(message_type) is not None: - raise ValueError(f"Duplicate handler for type {message_type} in {self.__class__.__name__}") + if self._handlers.get(message_type) is not None: + raise ValueError(f"Duplicate handler for type {message_type} in {self.__class__.__name__}") - # RequestInfoExecutor is allowed to handle SubWorkflowRequestInfo directly - # but other executors should use @intercepts_request - if message_type is SubWorkflowRequestInfo and not isinstance(self, RequestInfoExecutor): - raise ValueError( - f"Executor {self.__class__.__name__} cannot define a handler for " - "SubWorkflowRequestInfo directly. " - f"Use @intercepts_request decorator to intercept sub-workflow requests." - ) + # Get the bound method + bound_method = getattr(self, attr_name) + self._handlers[message_type] = bound_method - # Get the bound method - bound_method = getattr(self, attr_name) - self._handlers[message_type] = bound_method - - # Add to unified handler specs list - self._handler_specs.append({ - "name": handler_spec["name"], - "message_type": message_type, - "output_types": handler_spec.get("output_types", []), - "workflow_output_types": handler_spec.get("workflow_output_types", []), - "ctx_annotation": handler_spec.get("ctx_annotation"), - "source": "class_method", # Distinguish from instance handlers if needed - }) - - # Discover @intercepts_request methods - if hasattr(attr, "_intercepts_request"): - # Get the bound method for interceptors - bound_method = getattr(self, attr_name) - interceptor_info = { - "method": bound_method, - "from_workflow": getattr(attr, "_from_workflow", None), - "condition": getattr(attr, "_intercept_condition", None), - } - request_type = attr._intercepts_request # type: ignore - if request_type not in self._request_interceptors: - self._request_interceptors[request_type] = [] - self._request_interceptors[request_type].append(interceptor_info) - - # We need to register the handler for SubWorkflowRequestInfo - # if not already registered. - if SubWorkflowRequestInfo not in self._handlers: - self._handlers[SubWorkflowRequestInfo] = self._handle_sub_workflow_request - self._handler_specs.append({ - "name": attr_name, - "message_type": SubWorkflowRequestInfo, - "output_types": [SubWorkflowResponse, SubWorkflowRequestInfo, RequestInfoMessage], - "workflow_output_types": [], - "ctx_annotation": WorkflowContext[ - SubWorkflowRequestInfo | SubWorkflowResponse | RequestInfoMessage - ], - "source": "class_method", - }) + # Add to unified handler specs list + self._handler_specs.append({ + "name": handler_spec["name"], + "message_type": message_type, + "output_types": handler_spec.get("output_types", []), + "workflow_output_types": handler_spec.get("workflow_output_types", []), + "ctx_annotation": handler_spec.get("ctx_annotation"), + "source": "class_method", # Distinguish from instance handlers if needed + }) except AttributeError: # Skip attributes that may not be accessible continue - async def _handle_sub_workflow_request( - self, - request: "SubWorkflowRequestInfo", - ctx: WorkflowContext["SubWorkflowResponse | SubWorkflowRequestInfo | RequestInfoMessage"], - ) -> None: - """Automatic routing to @intercepts_request methods. - - This is only active for executors that have @intercepts_request methods. - """ - # Try to find a matching interceptor for the request and execute it - for request_type, interceptor_list in self._request_interceptors.items(): - if self._does_request_match_type(request.data, request_type): - for interceptor_info in interceptor_list: - if self._does_interceptor_apply_to_request(request, interceptor_info): - logger.debug( - f"Executor {self.id} intercepting request {request.request_id} " - f"of type {type(request.data).__name__} from sub-workflow {request.sub_workflow_id}" - ) - await self._execute_interceptor(request, interceptor_info, ctx) - return # Only the first matching interceptor is executed - logger.debug( - f"Executor {self.id} has no matching interceptor for request {request.request_id} " - f"of type {type(request.data).__name__} from sub-workflow {request.sub_workflow_id}; forwarding " - "inner request to RequestInfoExecutor for external handling." - ) - # No interceptor found - forward inner request to RequestInfoExecutor - await ctx.send_message(request.data) - - def _does_request_match_type(self, request_data: Any, request_type: type | str) -> bool: - """Check if request data matches the expected type. - - Args: - request_data: The request data to check - request_type: The type to match against (can be a type or string) - - Returns: - True if the request matches the type, False otherwise. - """ - if isinstance(request_type, type): - return is_instance_of(request_data, request_type) - - return ( - isinstance(request_type, str) - and hasattr(request_data, "__class__") - and request_data.__class__.__name__ == request_type - ) - - def _does_interceptor_apply_to_request( - self, request: "SubWorkflowRequestInfo", interceptor_info: dict[str, Any] - ) -> bool: - """Check if an interceptor applies to the given request. - - Args: - request: The sub-workflow request - interceptor_info: Information about the interceptor - - Returns: - True if the interceptor should handle this request, False otherwise. - """ - # Check workflow scope if specified - from_workflow = interceptor_info["from_workflow"] - if from_workflow and request.sub_workflow_id != from_workflow: - return False - - # Check additional condition - condition = interceptor_info["condition"] - return not (condition and not condition(request)) - - async def _execute_interceptor( - self, - request: "SubWorkflowRequestInfo", - interceptor_info: dict[str, Any], - ctx: WorkflowContext["SubWorkflowResponse | SubWorkflowRequestInfo | RequestInfoMessage"], - ) -> None: - """Execute a single interceptor method. - - Args: - request: The sub-workflow request - interceptor_info: Information about the interceptor method - ctx: The workflow context - """ - method = interceptor_info["method"] - response = await method(request.data, ctx) - - if not isinstance(response, RequestResponse): - raise RuntimeError( - f"Interceptor method {method.__name__} must return RequestResponse, got {type(response)}" - ) - - # Add automatic correlation info to the response - correlated_response = RequestResponse[RequestInfoMessage, Any].with_correlation( - response, - request.data, - request.request_id, - ) - - if correlated_response.is_handled: - logger.debug( - f"Executor {self.id}'s interceptor handled request {request.request_id} " - f"of type {type(request.data).__name__} from sub-workflow {request.sub_workflow_id}. " - f"Sending response back to sub-workflow." - ) - # Send response back to sub-workflow that made the request - await ctx.send_message( - SubWorkflowResponse( - request_id=request.request_id, - data=correlated_response.data, - ), - target_id=request.sub_workflow_id, - ) - else: - logger.debug( - f"Executor {self.id}'s interceptor did not handle request {request.request_id} " - f"of type {type(request.data).__name__} from sub-workflow {request.sub_workflow_id}. " - "Forwarding to RequestInfoExecutor for external handling." - ) - # Forward the request with potentially modified data - # Update the data if interceptor provided a modified request - if correlated_response.forward_request: - request.data = correlated_response.forward_request - - # Send the inner request to RequestInfoExecutor to create external request - await ctx.send_message(request) - def can_handle(self, message: Any) -> bool: """Check if the executor can handle a given message type. @@ -577,13 +401,6 @@ class Executor(AFBaseModel): output_types: List of output types for send_message() workflow_output_types: List of workflow output types for yield_output() """ - if message_type is SubWorkflowRequestInfo: - raise ValueError( - f"Executor {self.__class__.__name__} cannot define a handler for " - "SubWorkflowRequestInfo directly. " - f"Use @intercepts_request decorator to intercept sub-workflow requests." - ) - if message_type in self._handlers: raise ValueError(f"Handler for type {message_type} already registered in {self.__class__.__name__}") @@ -638,21 +455,6 @@ class Executor(AFBaseModel): return list(output_types) - @property - def request_types(self) -> list[type[Any]]: - """Get the list of request types that this executor can intercept via @intercepts_request. - - Returns: - A list of the request types that this executor's interceptors can handle. - """ - request_types: list[type[Any]] = [] - - for request_type in self._request_interceptors: - if isinstance(request_type, type): - request_types.append(request_type) - - return request_types - # endregion: Executor @@ -738,6 +540,11 @@ class RequestInfoMessage: """ request_id: str = field(default_factory=lambda: str(uuid.uuid4())) + """Unique identifier for correlating requests and responses.""" + + source_executor_id: str | None = None + """ID of the executor expecting a response to this request. + May differ from the executor that sent the request if intercepted and forwarded.""" TRequest = TypeVar("TRequest", bound="RequestInfoMessage") @@ -746,217 +553,26 @@ TResponse = TypeVar("TResponse") @dataclass class RequestResponse(Generic[TRequest, TResponse]): - """Response from @intercepts_request methods with automatic correlation support. + """Response type for request/response correlation in workflows. - This type allows intercepting executors to indicate whether they handled - a request or whether it should be forwarded to external sources. When handled, - the framework automatically adds correlation info to link responses to requests. + This type is used by RequestInfoExecutor to create correlated responses + that include the original request context for proper message routing. """ - is_handled: bool - data: TResponse | None = None - forward_request: TRequest | None = None - original_request: TRequest | None = None # Added for automatic correlation - request_id: str | None = None # Added for tracking + data: TResponse + """The response data returned from handling the request.""" - @classmethod - def handled(cls, data: TResponse) -> "RequestResponse[TRequest, TResponse]": - """Create a response indicating the request was handled. + original_request: TRequest + """The original request that this response corresponds to.""" - Correlation info (original_request, request_id) will be added automatically - by the framework when processing intercepted requests. - """ - return cls(is_handled=True, data=data) - - @classmethod - def forward(cls, modified_request: Any = None) -> "RequestResponse[TRequest, TResponse]": - """Create a response indicating the request should be forwarded.""" - return cls(is_handled=False, forward_request=modified_request) - - @staticmethod - def with_correlation( - original_response: "RequestResponse[TRequest, TResponse]", - original_request: TRequest, - request_id: str, - ) -> "RequestResponse[TRequest, TResponse]": - """Add correlation info to a response. - - This is called automatically by the framework when processing intercepted requests. - """ - return RequestResponse( - is_handled=original_response.is_handled, - data=original_response.data, - forward_request=original_response.forward_request, - original_request=original_request, - request_id=request_id, - ) - - -@dataclass -class SubWorkflowRequestInfo: - """Routes requests from sub-workflows to parent workflows. - - This message type wraps requests from sub-workflows to add routing context, - allowing parent workflows to intercept and potentially handle the request. - """ - - request_id: str # Original request ID from sub-workflow - sub_workflow_id: str # ID of the WorkflowExecutor that sent this - data: RequestInfoMessage # The actual request data - - -@dataclass -class SubWorkflowResponse: - """Routes responses back to sub-workflows. - - This message type is used to send responses back to sub-workflows that - made requests, ensuring the response reaches the correct sub-workflow. - """ - - request_id: str # Matches the original request ID - data: Any # The actual response data + request_id: str + """The ID of the original request.""" # endregion: Request/Response Types -# region Intercepts Request Decorator - -# TypeVar for request type that must be a RequestInfoMessage subclass -RequestInfoMessageT = TypeVar("RequestInfoMessageT", bound="RequestInfoMessage") - - -@overload -def intercepts_request( - func: Callable[ - [Any, RequestInfoMessageT, WorkflowContext[Any, Any]], Awaitable[RequestResponse[RequestInfoMessageT, Any]] - ], -) -> Callable[ - [Any, RequestInfoMessageT, WorkflowContext[Any, Any]], Awaitable[RequestResponse[RequestInfoMessageT, Any]] -]: ... - - -@overload -def intercepts_request( - *, - from_workflow: str | None = None, - condition: Callable[[Any], bool] | None = None, -) -> Callable[ - [ - Callable[ - [Any, RequestInfoMessageT, WorkflowContext[Any, Any]], Awaitable[RequestResponse[RequestInfoMessageT, Any]] - ] - ], - Callable[ - [Any, RequestInfoMessageT, WorkflowContext[Any, Any]], Awaitable[RequestResponse[RequestInfoMessageT, Any]] - ], -]: ... - - -def intercepts_request( - func: Callable[..., Any] | None = None, - *, - from_workflow: str | None = None, - condition: Callable[[Any], bool] | None = None, -) -> Callable[..., Any]: - """Decorator to mark methods that intercept sub-workflow requests. - - The request type is automatically inferred from the method's second parameter type hint. - The type must be a subclass of RequestInfoMessage. - - This decorator allows executors in parent workflows to intercept and handle requests from - sub-workflows before they are sent to external sources. - - Args: - func: The function being decorated (automatically passed when used without parentheses). - from_workflow: Optional ID of specific sub-workflow to intercept from. - condition: Optional callable that must return True for interception. - - Returns: - The decorated function with interception metadata. - - Example: - @intercepts_request - async def check_domain( - self, request: DomainCheckRequest, ctx: WorkflowContext - ) -> RequestResponse[DomainCheckRequest, bool]: - # Type automatically inferred as DomainCheckRequest from parameter annotation - if request.domain in self.approved_domains: - return RequestResponse.handled(True) - return RequestResponse.forward() - - @intercepts_request(from_workflow="email_validator") - async def handle_specific( - self, request: EmailRequest, ctx: WorkflowContext - ) -> RequestResponse[EmailRequest, str]: - # Only intercepts EmailRequest from the "email_validator" workflow - return RequestResponse.handled("handled by parent") - """ - - def decorator(func: Callable[..., Any]) -> Callable[..., Any]: - # Extract request type from method signature - sig = inspect.signature(func) - params = list(sig.parameters.values()) - - if len(params) < 2: - raise ValueError(f"Interceptor method '{func.__name__}' must have at least 2 parameters (self, request)") - - request_param = params[1] # Second parameter (after self) - request_type = request_param.annotation - - if request_type is inspect.Parameter.empty: - raise ValueError(f"Interceptor method '{func.__name__}' request parameter must have a type annotation") - - # Runtime validation that it's a RequestInfoMessage subclass - if isinstance(request_type, type): - # We need to check if RequestInfoMessage is defined yet, since this runs at import time - try: - # Try to get RequestInfoMessage from the current module's globals - request_info_message_class = None - func_module = inspect.getmodule(func) - if func_module and hasattr(func_module, "RequestInfoMessage"): - request_info_message_class = func_module.RequestInfoMessage - else: - # Look in the current module (where this decorator is defined) - import sys - - current_module = sys.modules[__name__] - if hasattr(current_module, "RequestInfoMessage"): - request_info_message_class = current_module.RequestInfoMessage - - if request_info_message_class and not issubclass(request_type, request_info_message_class): - raise TypeError( - f"Interceptor method '{func.__name__}' can only handle RequestInfoMessage subclasses, " - f"got {request_type}. Make sure your request type inherits from RequestInfoMessage." - ) - except (AttributeError, NameError): - # RequestInfoMessage might not be defined yet at import time, skip validation - # This will be caught later when the interceptor is actually called - pass - - @functools.wraps(func) - async def wrapper(self: Any, request: RequestInfoMessage, ctx: WorkflowContext[Any, Any]) -> Any: - return await func(self, request, ctx) - - # Add metadata for discovery - store the inferred type - wrapper._intercepts_request = request_type # type: ignore - wrapper._from_workflow = from_workflow # type: ignore - wrapper._intercept_condition = condition # type: ignore - - return wrapper - - # If func is provided, we're being called without parentheses: @intercepts_request - if func is not None: - return decorator(func) - - # Otherwise, we're being called with parentheses: @intercepts_request(from_workflow="...") - return decorator - - -# endregion: Intercepts Request Decorator # region Request Info Executor - - class RequestInfoExecutor(Executor): """Built-in executor that handles request/response patterns in workflows. @@ -975,12 +591,12 @@ class RequestInfoExecutor(Executor): """ super().__init__(id=id) self._request_events: dict[str, RequestInfoEvent] = {} - self._sub_workflow_contexts: dict[str, dict[str, str]] = {} @handler async def run(self, message: RequestInfoMessage, ctx: WorkflowContext) -> None: """Run the RequestInfoExecutor with the given message.""" - source_executor_id = ctx.get_source_executor_id() + # Use source_executor_id from message if available, otherwise fall back to context + source_executor_id = message.source_executor_id or ctx.get_source_executor_id() event = RequestInfoEvent( request_id=message.request_id, @@ -992,41 +608,11 @@ class RequestInfoExecutor(Executor): await self._record_pending_request_snapshot(message, source_executor_id, ctx) await ctx.add_event(event) - @handler - async def handle_sub_workflow_request( - self, - message: SubWorkflowRequestInfo, - ctx: WorkflowContext, - ) -> None: - """Handle forwarded sub-workflow request. - - This method handles requests that were forwarded from parent workflows - because they couldn't be handled locally. - """ - # When called directly from runner, we need to use the sub_workflow_id as the source - source_executor_id = message.sub_workflow_id - - # Store context for routing response back - self._sub_workflow_contexts[message.request_id] = { - "sub_workflow_id": message.sub_workflow_id, - "parent_executor_id": source_executor_id, - } - - # Create event for external handling - preserve the SubWorkflowRequestInfo wrapper - event = RequestInfoEvent( - request_id=message.request_id, # Use original request ID - source_executor_id=source_executor_id, - request_type=type(message.data), # Type of the wrapped data # type: ignore - request_data=message.data, # The wrapped request data - ) - self._request_events[message.request_id] = event - await ctx.add_event(event) - async def handle_response( self, response_data: Any, request_id: str, - ctx: WorkflowContext[SubWorkflowResponse | RequestResponse[RequestInfoMessage, Any]], + ctx: WorkflowContext[RequestResponse[RequestInfoMessage, Any]], ) -> None: """Handle a response to a request. @@ -1043,31 +629,11 @@ class RequestInfoExecutor(Executor): self._request_events.pop(request_id, None) - # Check if this was a forwarded sub-workflow request - if request_id in self._sub_workflow_contexts: - context = self._sub_workflow_contexts.pop(request_id) - - # Send back to sub-workflow that made the original request - await ctx.send_message( - SubWorkflowResponse( - request_id=request_id, - data=response_data, - ), - target_id=context["sub_workflow_id"], - ) - else: - # Regular response - send directly back to source - # Create a correlated response that includes both the response data and original request - if not isinstance(event.data, RequestInfoMessage): - raise TypeError(f"Expected RequestInfoMessage, got {type(event.data)}") - correlated_response = RequestResponse[RequestInfoMessage, Any].handled(response_data) - correlated_response = RequestResponse[RequestInfoMessage, Any].with_correlation( - correlated_response, - event.data, - request_id, - ) - - await ctx.send_message(correlated_response, target_id=event.source_executor_id) + # Create a correlated response that includes both the response data and original request + if not isinstance(event.data, RequestInfoMessage): + raise TypeError(f"Expected RequestInfoMessage, got {type(event.data)}") + correlated_response = RequestResponse(data=response_data, original_request=event.data, request_id=request_id) + await ctx.send_message(correlated_response, target_id=event.source_executor_id) await self._clear_pending_request_snapshot(request_id, ctx) @@ -1194,7 +760,7 @@ class RequestInfoExecutor(Executor): return repr(value) async def has_pending_request(self, request_id: str, ctx: WorkflowContext[Any]) -> bool: - if request_id in self._request_events or request_id in self._sub_workflow_contexts: + if request_id in self._request_events: return True snapshot = await self._get_pending_request_snapshot(request_id, ctx) return snapshot is not None diff --git a/python/packages/main/agent_framework/_workflow/_runner.py b/python/packages/main/agent_framework/_workflow/_runner.py index 844ca95618..d37615e254 100644 --- a/python/packages/main/agent_framework/_workflow/_runner.py +++ b/python/packages/main/agent_framework/_workflow/_runner.py @@ -22,7 +22,6 @@ from ._runner_context import ( _decode_checkpoint_value, ) from ._shared_state import SharedState -from ._typing_utils import is_instance_of logger = logging.getLogger(__name__) @@ -160,90 +159,6 @@ class Runner: async def _deliver_messages(source_executor_id: str, messages: list[Message]) -> None: """Outer loop to concurrently deliver messages from all sources to their targets.""" - # Special handling for SubWorkflowRequestInfo messages - async def _deliver_sub_workflow_requests(messages: list[Message]) -> None: - from ._executor import SubWorkflowRequestInfo - - # Handle SubWorkflowRequestInfo messages - only process those not already targeted - sub_workflow_messages: list[Message] = [] - for msg in messages: - # Skip messages sent directly to RequestInfoExecutor - they are already forwarded - if self._is_message_to_request_info_executor(msg): - continue - - if isinstance(msg.data, SubWorkflowRequestInfo): - sub_workflow_messages.append(msg) - - for message in sub_workflow_messages: - # message.data is guaranteed to be SubWorkflowRequestInfo via filtering above - sub_request = message.data # type: ignore[assignment] - - # Find executor that can intercept the wrapped type - interceptor_found = False - for executor in self._executors.values(): - interceptors = getattr(executor, "_request_interceptors", []) - if interceptors and executor.id != message.source_id: - for registered_type in interceptors: # type: ignore[assignment] - # Check type matching - handle both type and string cases - matched = False - if ( - isinstance(registered_type, type) - and is_instance_of(sub_request.data, registered_type) - ) or ( - isinstance(registered_type, str) - and hasattr(sub_request.data, "__class__") - and sub_request.data.__class__.__name__ == registered_type - ): - matched = True - - if matched: - # Send directly to the intercepting executor - logger.info( - f"Sending sub-workflow request of type '{sub_request.data.__class__.__name__}' " - f"from sub-workflow '{sub_request.sub_workflow_id}' " - f"to executor '{executor.id}' for interception." - ) - await executor.execute( - sub_request, - [message.source_id], # source_executor_ids - self._shared_state, # shared_state - self._ctx, # runner_context - trace_contexts=[message.trace_context] if message.trace_context else None, - source_span_ids=[message.source_span_id] if message.source_span_id else None, - ) - interceptor_found = True - break - if interceptor_found: - break - - if not interceptor_found: - # No interceptor found - send directly to RequestInfoExecutor if available. - - # Find the RequestInfoExecutor instance - request_info_executor = self._find_request_info_executor() - - if request_info_executor: - logger.info( - f"Sending sub-workflow request of type '{sub_request.data.__class__.__name__}' " - f"from sub-workflow '{sub_request.sub_workflow_id}' to RequestInfoExecutor " - f"'{request_info_executor.id}'" - ) - await request_info_executor.execute( - sub_request, - [message.source_id], # source_executor_ids - self._shared_state, # shared_state - self._ctx, # runner_context - trace_contexts=[message.trace_context] if message.trace_context else None, - source_span_ids=[message.source_span_id] if message.source_span_id else None, - ) - else: - logger.warning( - f"Sub-workflow request of type '{sub_request.data.__class__.__name__}' " - f"from sub-workflow '{sub_request.sub_workflow_id}' could not be handled: " - f"no RequestInfoExecutor found in the workflow. Add a RequestInfoExecutor " - f"to handle external requests or add an @intercepts_request handler." - ) - async def _deliver_message_inner(edge_runner: EdgeRunner, message: Message) -> bool: """Inner loop to deliver a single message through an edge runner.""" return await edge_runner.send_message(message, self._shared_state, self._ctx) @@ -261,28 +176,9 @@ class Runner: return message.data = decoded - # Handle SubWorkflowRequestInfo messages specially - await _deliver_sub_workflow_requests(messages) - - # Filter out SubWorkflowRequestInfo messages from normal edge routing - # since they were handled specially - from ._executor import SubWorkflowRequestInfo - - non_sub_workflow_messages: list[Message] = [] - for msg in messages: - # Keep messages sent directly to RequestInfoExecutor (forwarded messages) - if self._is_message_to_request_info_executor(msg): - non_sub_workflow_messages.append(msg) - continue - - # Skip SubWorkflowRequestInfo messages (handled by special routing) - if isinstance(msg.data, SubWorkflowRequestInfo): - continue - - non_sub_workflow_messages.append(msg) - + # Route all messages through normal workflow edges associated_edge_runners = self._edge_runner_map.get(source_executor_id, []) - for message in non_sub_workflow_messages: + for message in messages: _normalize_message_payload(message) # Deliver a message through all edge runners associated with the source executor concurrently. tasks = [_deliver_message_inner(edge_runner, message) for edge_runner in associated_edge_runners] diff --git a/python/packages/main/agent_framework/_workflow/_validation.py b/python/packages/main/agent_framework/_workflow/_validation.py index f2ece7e3c0..1a1ae5faab 100644 --- a/python/packages/main/agent_framework/_workflow/_validation.py +++ b/python/packages/main/agent_framework/_workflow/_validation.py @@ -10,7 +10,6 @@ from typing import Any, Union, get_args, get_origin from ._edge import Edge, EdgeGroup, FanInEdgeGroup from ._executor import Executor, RequestInfoExecutor -from ._workflow_executor import WorkflowExecutor logger = logging.getLogger(__name__) @@ -187,7 +186,6 @@ class WorkflowGraphValidator: self._validate_self_loops() self._validate_dead_ends() self._validate_cycles() - self._validate_interceptor_uniqueness() def _validate_handler_output_annotations(self) -> None: """Validate that each handler's ctx parameter is annotated with WorkflowContext[T]. @@ -273,9 +271,6 @@ class WorkflowGraphValidator: # Get output types from source executor source_output_types = list(source_executor.output_types) - # Also include intercepted request types as potential outputs - # since @intercepts_request methods can forward requests - source_output_types.extend(source_executor.request_types) # Get input types from target executor target_input_types = target_executor.input_types @@ -482,52 +477,6 @@ class WorkflowGraphValidator: "Ensure proper termination conditions exist to prevent infinite loops." ) - def _validate_interceptor_uniqueness(self) -> None: - """Validate that only one executor intercepts a given request type from a specific sub-workflow. - - This prevents non-deterministic behavior where multiple executors could intercept - the same request type from the same sub-workflow. - """ - # Find all WorkflowExecutor instances in the workflow - workflow_executors: dict[str, WorkflowExecutor] = {} - for executor_id, executor in self._executors.items(): - if isinstance(executor, WorkflowExecutor): - workflow_executors[executor_id] = executor - - # For each WorkflowExecutor, check which executors can intercept its requests - for workflow_id, _workflow_executor in workflow_executors.items(): - # Map of request_type -> list of intercepting executor IDs - interceptors_by_type: dict[type | str, list[str]] = {} - - # Find all executors that have edges from this WorkflowExecutor - # These are potential interceptors - for edge in self._edges: - if edge.source_id == workflow_id: - target_executor = self._executors.get(edge.target_id) - if target_executor and hasattr(target_executor, "_request_interceptors"): - # Check what request types this executor intercepts - for request_type, interceptor_list in target_executor._request_interceptors.items(): - # Check if any interceptor is scoped to this workflow or unscoped - for interceptor_info in interceptor_list: - from_workflow = interceptor_info.get("from_workflow") - # If unscoped or specifically scoped to this workflow - if from_workflow is None or from_workflow == workflow_id: - if request_type not in interceptors_by_type: - interceptors_by_type[request_type] = [] - interceptors_by_type[request_type].append(edge.target_id) - - # Check for duplicates - for request_type, executor_ids in interceptors_by_type.items(): - unique_executors = list(set(executor_ids)) # Remove duplicates from same executor - if len(unique_executors) > 1: - type_name = request_type.__name__ if isinstance(request_type, type) else str(request_type) - raise InterceptorConflictError( - f"Multiple executors intercept the same request type '{type_name}' " - f"from sub-workflow '{workflow_id}': {', '.join(unique_executors)}. " - f"Only one executor should intercept a given request type from a specific sub-workflow " - f"to ensure deterministic behavior." - ) - # endregion # region Type Compatibility Utilities @@ -572,6 +521,7 @@ class WorkflowGraphValidator: for s_arg, t_arg in zip(source_args, target_args, strict=True) ) + # No other special compatibility cases return False # endregion diff --git a/python/packages/main/agent_framework/_workflow/_workflow_executor.py b/python/packages/main/agent_framework/_workflow/_workflow_executor.py index 3f9ded3681..0dd714ae1c 100644 --- a/python/packages/main/agent_framework/_workflow/_workflow_executor.py +++ b/python/packages/main/agent_framework/_workflow/_workflow_executor.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import inspect import logging import uuid from dataclasses import dataclass @@ -19,10 +20,10 @@ from ._executor import ( Executor, RequestInfoExecutor, RequestInfoMessage, - SubWorkflowRequestInfo, - SubWorkflowResponse, + RequestResponse, handler, ) +from ._typing_utils import is_instance_of from ._workflow_context import WorkflowContext logger = logging.getLogger(__name__) @@ -76,9 +77,11 @@ class WorkflowExecutor(Executor): request = MyDataRequest(query="user info") # RequestInfoExecutor emits RequestInfoEvent - # WorkflowExecutor wraps and forwards to parent - wrapped = SubWorkflowRequestInfo(request_id="...", sub_workflow_id="child_workflow", data=request) - # Parent workflow can intercept via @intercepts_request + # WorkflowExecutor sets source_executor_id and forwards to parent + request.source_executor_id = "child_workflow_executor_id" + # Parent workflow can handle via @handler for RequestInfoMessage subclasses, + # or directly forward to external source via a RequestInfoExecutor in the parent + # workflow. ``` ### State Management @@ -103,8 +106,8 @@ class WorkflowExecutor(Executor): Combines sub-workflow outputs with request coordination types: ```python # Includes all sub-workflow output types - # Plus SubWorkflowRequestInfo if sub-workflow can make requests - output_types = workflow.output_types + [SubWorkflowRequestInfo] # if applicable + # Plus RequestInfoMessage if sub-workflow can make requests + output_types = workflow.output_types + [RequestInfoMessage] # if applicable ``` ## Error Handling @@ -169,14 +172,20 @@ class WorkflowExecutor(Executor): Parent workflows can intercept sub-workflow requests: ```python class ParentExecutor(Executor): - @intercepts_request - async def handle_child_request( - self, request: MyDataRequest, ctx: WorkflowContext[Any] - ) -> RequestResponse[MyDataRequest, str]: + @handler + async def handle_request( + self, + request: MyRequestType, # Subclass of RequestInfoMessage + ctx: WorkflowContext[RequestResponse[RequestInfoMessage, Any] | RequestInfoMessage], + ) -> None: # Handle request locally or forward to external source if self.can_handle_locally(request): - return RequestResponse.handled("local result") - return RequestResponse.forward() # Send to external handler + # Send response back to sub-workflow + response = RequestResponse(data="local result", original_request=request, request_id=request.request_id) + await ctx.send_message(response, target_id=request.source_executor_id) + else: + # Forward to external handler + await ctx.send_message(request) ``` ## Implementation Notes @@ -208,12 +217,18 @@ class WorkflowExecutor(Executor): @property def input_types(self) -> list[type[Any]]: - """Get the input types based on the underlying workflow's input types. + """Get the input types based on the underlying workflow's input types plus WorkflowExecutor-specific types. Returns: - A list of input types that the underlying workflow can accept. + A list of input types that the WorkflowExecutor can accept. """ - return self.workflow.input_types + input_types = list(self.workflow.input_types) + + # WorkflowExecutor can also handle RequestResponse for sub-workflow responses + if RequestResponse not in input_types: + input_types.append(RequestResponse) + + return input_types @property def output_types(self) -> list[type[Any]]: @@ -221,20 +236,43 @@ class WorkflowExecutor(Executor): Returns: A list of output types that the underlying workflow can produce. - Includes SubWorkflowRequestInfo if the sub-workflow contains RequestInfoExecutor. + Includes specific RequestInfoMessage subtypes if the sub-workflow contains RequestInfoExecutor. """ output_types = list(self.workflow.output_types) # Check if the sub-workflow contains a RequestInfoExecutor - # If so, this WorkflowExecutor can also output SubWorkflowRequestInfo messages - for executor in self.workflow.executors.values(): - if isinstance(executor, RequestInfoExecutor): - if SubWorkflowRequestInfo not in output_types: - output_types.append(SubWorkflowRequestInfo) - break + # If so, collect the specific RequestInfoMessage subtypes from all executors + has_request_info_executor = any( + isinstance(executor, RequestInfoExecutor) for executor in self.workflow.executors.values() + ) + + if has_request_info_executor: + # Collect all RequestInfoMessage subtypes from executor output types + for executor in self.workflow.executors.values(): + for output_type in executor.output_types: + # Check if this is a RequestInfoMessage subclass + if ( + inspect.isclass(output_type) + and issubclass(output_type, RequestInfoMessage) + and output_type not in output_types + ): + output_types.append(output_type) return output_types + def can_handle(self, message: Any) -> bool: + """Override can_handle to only accept messages that the wrapped workflow can handle. + + This prevents the WorkflowExecutor from accepting messages that should go to other + executors (like RequestInfoExecutor). + """ + # Always handle RequestResponse (for the handle_response handler) + if isinstance(message, RequestResponse): + return True + + # For other messages, only handle if the wrapped workflow can accept them as input + return any(is_instance_of(message, input_type) for input_type in self.workflow.input_types) + @handler # No output_types - can send any completion data type async def process_workflow(self, input_data: object, ctx: WorkflowContext[Any]) -> None: """Execute the sub-workflow with raw input data. @@ -246,8 +284,8 @@ class WorkflowExecutor(Executor): input_data: The input data to send to the sub-workflow. ctx: The workflow context from the parent. """ - # Skip SubWorkflowResponse and SubWorkflowRequestInfo - they have specific handlers - if isinstance(input_data, (SubWorkflowResponse, SubWorkflowRequestInfo)): + # Skip RequestResponse - it has a specific handler + if isinstance(input_data, RequestResponse): logger.debug(f"WorkflowExecutor {self.id} ignoring input of type {type(input_data)}") return @@ -318,15 +356,12 @@ class WorkflowExecutor(Executor): execution_context.pending_requests[event.request_id] = event.data # Map request to execution for response routing self._request_to_execution[event.request_id] = execution_context.execution_id - # Wrap request with routing context and send to parent + # Set source_executor_id for response routing and send to parent if not isinstance(event.data, RequestInfoMessage): raise TypeError(f"Expected RequestInfoMessage, got {type(event.data)}") - wrapped_request = SubWorkflowRequestInfo( - request_id=event.request_id, - sub_workflow_id=self.id, - data=event.data, - ) - await ctx.send_message(wrapped_request) + # Set the source_executor_id to this WorkflowExecutor's ID for response routing + event.data.source_executor_id = self.id + await ctx.send_message(event.data) # Update expected response count for this execution execution_context.expected_response_count = len(request_info_events) @@ -375,7 +410,7 @@ class WorkflowExecutor(Executor): @handler async def handle_response( self, - response: SubWorkflowResponse, + response: RequestResponse[RequestInfoMessage, Any], ctx: WorkflowContext[Any], ) -> None: """Handle response from parent for a forwarded request. diff --git a/python/packages/main/tests/workflow/test_checkpoint_decode.py b/python/packages/main/tests/workflow/test_checkpoint_decode.py index 7c528a5149..04d78f8e18 100644 --- a/python/packages/main/tests/workflow/test_checkpoint_decode.py +++ b/python/packages/main/tests/workflow/test_checkpoint_decode.py @@ -14,11 +14,10 @@ class SampleRequest(RequestInfoMessage): def test_decode_dataclass_with_nested_request() -> None: - original = RequestResponse[SampleRequest, str].handled("approve") - original = RequestResponse[SampleRequest, str].with_correlation( - original, - SampleRequest(request_id="abc", prompt="prompt"), - "abc", + original = RequestResponse[SampleRequest, str]( + data="approve", + original_request=SampleRequest(request_id="abc", prompt="prompt"), + request_id="abc", ) encoded = _encode_checkpoint_value(original) @@ -32,11 +31,10 @@ def test_decode_dataclass_with_nested_request() -> None: def test_is_instance_of_coerces_request_response_original_request_dict() -> None: - response = RequestResponse[SampleRequest, str].handled("approve") - response = RequestResponse[SampleRequest, str].with_correlation( - response, - SampleRequest(request_id="req-1", prompt="prompt"), - "req-1", + response = RequestResponse[SampleRequest, str]( + data="approve", + original_request=SampleRequest(request_id="req-1", prompt="prompt"), + request_id="req-1", ) # Simulate checkpoint decode fallback leaving a dict diff --git a/python/packages/main/tests/workflow/test_request_info_executor_rehydrate.py b/python/packages/main/tests/workflow/test_request_info_executor_rehydrate.py index b43d427477..36cfffc957 100644 --- a/python/packages/main/tests/workflow/test_request_info_executor_rehydrate.py +++ b/python/packages/main/tests/workflow/test_request_info_executor_rehydrate.py @@ -171,11 +171,10 @@ def test_pending_requests_from_checkpoint_and_summary() -> None: request = SimpleApproval(prompt="Review draft", draft="Draft text", iteration=3) request.request_id = "req-42" - response = RequestResponse[SimpleApproval, str].handled("approve") - response = RequestResponse[SimpleApproval, str].with_correlation( - response, - request, - request.request_id, + response = RequestResponse[SimpleApproval, str]( + data="approve", + original_request=request, + request_id=request.request_id, ) encoded_response = _encode_checkpoint_value(response) diff --git a/python/packages/main/tests/workflow/test_simple_sub_workflow.py b/python/packages/main/tests/workflow/test_simple_sub_workflow.py deleted file mode 100644 index f57203a1cb..0000000000 --- a/python/packages/main/tests/workflow/test_simple_sub_workflow.py +++ /dev/null @@ -1,108 +0,0 @@ -# Copyright (c) Microsoft. All rights reserved. - -import asyncio -from dataclasses import dataclass - -from typing_extensions import Never - -from agent_framework import ( - Executor, - WorkflowBuilder, - WorkflowContext, - WorkflowExecutor, - handler, -) - - -@dataclass -class SimpleRequest: - """Simple request for testing.""" - - text: str - - -@dataclass -class SimpleResponse: - """Simple response for testing.""" - - result: str - - -class SimpleSubExecutor(Executor): - """Simple executor for sub-workflow.""" - - def __init__(self): - super().__init__(id="simple_sub") - - @handler - async def process(self, request: SimpleRequest, ctx: WorkflowContext[Never, SimpleResponse]) -> None: - """Process a simple request.""" - # Just echo back with prefix and complete - response = SimpleResponse(result=f"processed: {request.text}") - await ctx.yield_output(response) - - -class SimpleParent(Executor): - """Simple parent executor.""" - - result: SimpleResponse | None = None - - def __init__(self): - super().__init__(id="simple_parent") - - @handler - async def start(self, text: str, ctx: WorkflowContext[SimpleRequest]) -> None: - """Start the process.""" - request = SimpleRequest(text=text) - await ctx.send_message(request, target_id="sub_workflow") - - @handler - async def collect(self, response: SimpleResponse, ctx: WorkflowContext) -> None: - """Collect the result.""" - self.result = response - - -async def test_simple_sub_workflow(): - """Test the simplest possible sub-workflow.""" - # Create sub-workflow with dummy executor to satisfy validation - sub_executor = SimpleSubExecutor() - - class DummyExecutor(Executor): - def __init__(self): - super().__init__(id="dummy") - - @handler - async def process(self, message: object, ctx: WorkflowContext) -> None: - pass # Do nothing - - dummy = DummyExecutor() - sub_workflow = ( - WorkflowBuilder() - .set_start_executor(sub_executor) - .add_edge(sub_executor, dummy) # Add edge to satisfy validation - .build() - ) - - # Create parent workflow - parent = SimpleParent() - workflow_executor = WorkflowExecutor(sub_workflow, id="sub_workflow") - - main_workflow = ( - WorkflowBuilder() - .set_start_executor(parent) - .add_edge(parent, workflow_executor) - .add_edge(workflow_executor, parent) - .build() - ) - - # Run the workflow - await main_workflow.run("hello world") - - # Check result - assert parent.result is not None - assert parent.result.result == "processed: hello world" - - -if __name__ == "__main__": - # Run the simple test - asyncio.run(test_simple_sub_workflow()) diff --git a/python/packages/main/tests/workflow/test_sub_workflow.py b/python/packages/main/tests/workflow/test_sub_workflow.py index 2067734734..5979746224 100644 --- a/python/packages/main/tests/workflow/test_sub_workflow.py +++ b/python/packages/main/tests/workflow/test_sub_workflow.py @@ -1,6 +1,5 @@ # Copyright (c) Microsoft. All rights reserved. -import asyncio from dataclasses import dataclass from typing import Any @@ -12,11 +11,11 @@ from agent_framework import ( RequestInfoExecutor, RequestInfoMessage, RequestResponse, + Workflow, WorkflowBuilder, WorkflowContext, WorkflowExecutor, handler, - intercepts_request, ) @@ -45,6 +44,61 @@ class ValidationResult: reason: str +# Test helper functions +def create_email_validation_workflow() -> Workflow: + """Create a standard email validation workflow.""" + email_validator = EmailValidator() + email_request_info = RequestInfoExecutor(id="email_request_info") + + return ( + WorkflowBuilder() + .set_start_executor(email_validator) + .add_edge(email_validator, email_request_info) + .add_edge(email_request_info, email_validator) + .build() + ) + + +class BasicParent(Executor): + """Basic parent executor for simple sub-workflow tests.""" + + result: ValidationResult | None = Field(default=None) + cache: dict[str, bool] = Field(default_factory=dict) + + def __init__(self, cache: dict[str, bool] | None = None, **kwargs: Any): + if cache is not None: + kwargs["cache"] = cache + super().__init__(id="basic_parent", **kwargs) + + @handler + async def start(self, email: str, ctx: WorkflowContext[EmailValidationRequest]) -> None: + request = EmailValidationRequest(email=email) + await ctx.send_message(request, target_id="email_workflow") + + @handler + async def handle_domain_request( + self, + request: DomainCheckRequest, + ctx: WorkflowContext[RequestResponse[DomainCheckRequest, Any] | DomainCheckRequest], + ) -> None: + """Handle requests from sub-workflows with optional caching.""" + domain_request = request + + if domain_request.domain in self.cache: + # Return cached result + response = RequestResponse( + data=self.cache[domain_request.domain], original_request=request, request_id=request.request_id + ) + await ctx.send_message(response, target_id=request.source_executor_id) + else: + # Not in cache, forward to external + await ctx.send_message(request) + + @handler + async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None: + self.result = result + + # Test executors class EmailValidator(Executor): """Validates email addresses in a sub-workflow.""" @@ -54,7 +108,7 @@ class EmailValidator(Executor): @handler async def validate_request( - self, request: EmailValidationRequest, ctx: WorkflowContext[RequestInfoMessage, ValidationResult] + self, request: EmailValidationRequest, ctx: WorkflowContext[DomainCheckRequest, ValidationResult] ) -> None: """Validate an email address.""" # Extract domain and check if it's approved @@ -101,17 +155,23 @@ class ParentOrchestrator(Executor): request = EmailValidationRequest(email=email) await ctx.send_message(request, target_id="email_workflow") - @intercepts_request - async def check_domain( - self, request: DomainCheckRequest, ctx: WorkflowContext[Any] - ) -> RequestResponse[DomainCheckRequest, bool]: - """Intercept domain check requests from sub-workflows.""" - # Check if we know this domain - if request.domain in self.approved_domains: - return RequestResponse[DomainCheckRequest, bool].handled(True) + @handler + async def handle_domain_request( + self, + request: DomainCheckRequest, + ctx: WorkflowContext[RequestResponse[DomainCheckRequest, Any] | DomainCheckRequest], + ) -> None: + """Handle requests from sub-workflows.""" + domain_request = request - # We don't know this domain, forward to external - return RequestResponse[DomainCheckRequest, bool].forward() + # Check if we know this domain + if domain_request.domain in self.approved_domains: + # Send response back to sub-workflow + response = RequestResponse(data=True, original_request=request, request_id=request.request_id) + await ctx.send_message(response, target_id=request.source_executor_id) + else: + # We don't know this domain, forward to external + await ctx.send_message(request) @handler async def collect_result(self, result: ValidationResult, ctx: WorkflowContext) -> None: @@ -122,34 +182,10 @@ class ParentOrchestrator(Executor): async def test_basic_sub_workflow() -> None: """Test basic sub-workflow execution without interception.""" # Create sub-workflow - email_validator = EmailValidator() - email_request_info = RequestInfoExecutor(id="email_request_info") - - validation_workflow = ( - WorkflowBuilder() - .set_start_executor(email_validator) - .add_edge(email_validator, email_request_info) - .add_edge(email_request_info, email_validator) - .build() - ) + validation_workflow = create_email_validation_workflow() # Create parent workflow without interception - class SimpleParent(Executor): - result: ValidationResult | None = Field(default=None) - - def __init__(self, **kwargs: Any): - super().__init__(id="simple_parent", **kwargs) - - @handler - async def start(self, email: str, ctx: WorkflowContext[EmailValidationRequest]) -> None: - request = EmailValidationRequest(email=email) - await ctx.send_message(request, target_id="email_workflow") - - @handler - async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None: - self.result = result - - parent = SimpleParent() + parent = BasicParent() workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow") main_request_info = RequestInfoExecutor(id="main_request_info") @@ -159,7 +195,7 @@ async def test_basic_sub_workflow() -> None: .add_edge(parent, workflow_executor) .add_edge(workflow_executor, parent) .add_edge(workflow_executor, main_request_info) - .add_edge(main_request_info, workflow_executor) # CRITICAL: For SubWorkflowResponse routing + .add_edge(main_request_info, workflow_executor) # CRITICAL: For RequestResponse routing .build() ) @@ -184,21 +220,12 @@ async def test_basic_sub_workflow() -> None: async def test_sub_workflow_with_interception(): - """Test sub-workflow with parent interception of requests.""" + """Test sub-workflow with parent interception and conditional forwarding.""" # Create sub-workflow - email_validator = EmailValidator() - email_request_info = RequestInfoExecutor(id="email_request_info") + validation_workflow = create_email_validation_workflow() - validation_workflow = ( - WorkflowBuilder() - .set_start_executor(email_validator) - .add_edge(email_validator, email_request_info) - .add_edge(email_request_info, email_validator) - .build() - ) - - # Create parent workflow with interception - parent = ParentOrchestrator(approved_domains={"example.com", "internal.org"}) + # Create parent workflow with interception cache + parent = BasicParent(cache={"example.com": True, "internal.org": True}) workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow") parent_request_info = RequestInfoExecutor(id="request_info") @@ -208,29 +235,23 @@ async def test_sub_workflow_with_interception(): .add_edge(parent, workflow_executor) .add_edge(workflow_executor, parent) .add_edge(parent, parent_request_info) # For forwarded requests - .add_edge(parent_request_info, workflow_executor) # For SubWorkflowResponse routing + .add_edge(parent_request_info, workflow_executor) # For RequestResponse routing .build() ) - # Test 1: Email with known domain (intercepted) - result = await main_workflow.run(["user@example.com"]) - - # Should complete without external requests + # Test 1: Email with cached domain (intercepted) + result = await main_workflow.run("user@example.com") request_events = result.get_request_info_events() - assert len(request_events) == 0 # No external requests, handled internally + assert len(request_events) == 0 # No external requests, handled from cache + assert parent.result is not None + assert parent.result.email == "user@example.com" + assert parent.result.is_valid is True - assert len(parent.results) == 1 - assert parent.results[0].email == "user@example.com" - assert parent.results[0].is_valid is True - assert parent.results[0].reason == "Domain approved" - - # Test 2: Email with unknown domain (forwarded) - parent.results.clear() - result = await main_workflow.run(["user@unknown.com"]) - - # Should have external request + # Test 2: Email with unknown domain (forwarded to external) + parent.result = None + result = await main_workflow.run("user@unknown.com") request_events = result.get_request_info_events() - assert len(request_events) == 1 + assert len(request_events) == 1 # Forwarded to external assert isinstance(request_events[0].data, DomainCheckRequest) assert request_events[0].data.domain == "unknown.com" @@ -238,89 +259,18 @@ async def test_sub_workflow_with_interception(): await main_workflow.send_responses({ request_events[0].request_id: False # Domain not approved }) + assert parent.result is not None + assert parent.result.email == "user@unknown.com" + assert parent.result.is_valid is False - assert len(parent.results) == 1 - assert parent.results[0].email == "user@unknown.com" - assert parent.results[0].is_valid is False - assert parent.results[0].reason == "Domain not approved" - - -async def test_conditional_forwarding() -> None: - """Test conditional forwarding with RequestResponse.forward().""" - - class ConditionalParent(Executor): - """Parent that conditionally handles requests.""" - - cache: dict[str, bool] = Field(default_factory=lambda: {"cached.com": True}) - result: ValidationResult | None = Field(default=None) - - def __init__(self, **kwargs: Any): - super().__init__(id="conditional_parent", **kwargs) - - @handler - async def start(self, email: str, ctx: WorkflowContext[EmailValidationRequest]) -> None: - request = EmailValidationRequest(email=email) - await ctx.send_message(request, target_id="email_workflow") - - @intercepts_request - async def check_domain( - self, request: DomainCheckRequest, ctx: WorkflowContext[Any] - ) -> RequestResponse[DomainCheckRequest, bool]: - """Check cache first, then forward if not found.""" - if request.domain in self.cache: - # Return cached result - return RequestResponse[DomainCheckRequest, bool].handled(self.cache[request.domain]) - - # Not in cache, forward to external - return RequestResponse[DomainCheckRequest, bool].forward() - - @handler - async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None: - self.result = result - - # Setup workflows - email_validator = EmailValidator() - request_info = RequestInfoExecutor(id="request_info") - - validation_workflow = ( - WorkflowBuilder() - .set_start_executor(email_validator) - .add_edge(email_validator, request_info) - .add_edge(request_info, email_validator) - .build() - ) - - parent = ConditionalParent() - workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow") - parent_request_info = RequestInfoExecutor(id="request_info") - - main_workflow = ( - WorkflowBuilder() - .set_start_executor(parent) - .add_edge(parent, workflow_executor) - .add_edge(workflow_executor, parent) - .add_edge(parent, parent_request_info) - .add_edge(parent_request_info, workflow_executor) # For SubWorkflowResponse routing - .build() - ) - - # Test cached domain - result = await main_workflow.run("user@cached.com") + # Test 3: Another cached domain + parent.result = None + result = await main_workflow.run("user@internal.org") request_events = result.get_request_info_events() assert len(request_events) == 0 # Handled from cache assert parent.result is not None assert parent.result.is_valid is True - # Test uncached domain - parent.result = None - result = await main_workflow.run("user@new.com") - request_events = result.get_request_info_events() - assert len(request_events) == 1 # Forwarded to external - - await main_workflow.send_responses({request_events[0].request_id: True}) - assert parent.result is not None - assert parent.result.is_valid is True - async def test_workflow_scoped_interception() -> None: """Test interception scoped to specific sub-workflows.""" @@ -339,42 +289,41 @@ async def test_workflow_scoped_interception() -> None: await ctx.send_message(EmailValidationRequest(email=data["email1"]), target_id="workflow_a") await ctx.send_message(EmailValidationRequest(email=data["email2"]), target_id="workflow_b") - @intercepts_request(from_workflow="workflow_a") - async def check_domain_a( - self, request: DomainCheckRequest, ctx: WorkflowContext[Any] - ) -> RequestResponse[DomainCheckRequest, bool]: - """Strict rules for workflow A.""" - if request.domain == "strict.com": - return RequestResponse[DomainCheckRequest, bool].handled(True) - return RequestResponse[DomainCheckRequest, bool].forward() + @handler + async def handle_domain_request( + self, + request: DomainCheckRequest, + ctx: WorkflowContext[RequestResponse[DomainCheckRequest, Any] | DomainCheckRequest], + ) -> None: + domain_request = request - @intercepts_request(from_workflow="workflow_b") - async def check_domain_b( - self, request: DomainCheckRequest, ctx: WorkflowContext[Any] - ) -> RequestResponse[DomainCheckRequest, bool]: - """Lenient rules for workflow B.""" - if request.domain.endswith(".com"): - return RequestResponse[DomainCheckRequest, bool].handled(True) - return RequestResponse[DomainCheckRequest, bool].forward() + if request.source_executor_id == "workflow_a": + # Strict rules for workflow A + if domain_request.domain == "strict.com": + response = RequestResponse(data=True, original_request=request, request_id=request.request_id) + await ctx.send_message(response, target_id=request.source_executor_id) + else: + # Forward to external + await ctx.send_message(request) + elif request.source_executor_id == "workflow_b": + # Lenient rules for workflow B + if domain_request.domain.endswith(".com"): + response = RequestResponse(data=True, original_request=request, request_id=request.request_id) + await ctx.send_message(response, target_id=request.source_executor_id) + else: + # Forward to external + await ctx.send_message(request) + else: + # Unknown source, forward to external + await ctx.send_message(request) @handler async def collect(self, result: ValidationResult, ctx: WorkflowContext) -> None: self.results[result.email] = result # Create two identical sub-workflows - def create_validation_workflow(): - validator = EmailValidator() - request_info = RequestInfoExecutor(id="request_info") - return ( - WorkflowBuilder() - .set_start_executor(validator) - .add_edge(validator, request_info) - .add_edge(request_info, validator) - .build() - ) - - workflow_a = create_validation_workflow() - workflow_b = create_validation_workflow() + workflow_a = create_email_validation_workflow() + workflow_b = create_email_validation_workflow() parent = MultiWorkflowParent() executor_a = WorkflowExecutor(workflow_a, "workflow_a") @@ -389,8 +338,8 @@ async def test_workflow_scoped_interception() -> None: .add_edge(executor_a, parent) .add_edge(executor_b, parent) .add_edge(parent, parent_request_info) - .add_edge(parent_request_info, executor_a) # For SubWorkflowResponse routing - .add_edge(parent_request_info, executor_b) # For SubWorkflowResponse routing + .add_edge(parent_request_info, executor_a) # For RequestResponse routing + .add_edge(parent_request_info, executor_b) # For RequestResponse routing .build() ) @@ -432,16 +381,7 @@ async def test_concurrent_sub_workflow_execution() -> None: self.results.append(result) # Create sub-workflow for email validation - email_validator = EmailValidator() - email_request_info = RequestInfoExecutor(id="email_request_info") - - validation_workflow = ( - WorkflowBuilder() - .set_start_executor(email_validator) - .add_edge(email_validator, email_request_info) - .add_edge(email_request_info, email_validator) - .build() - ) + validation_workflow = create_email_validation_workflow() # Create parent workflow processor = ConcurrentProcessor() @@ -454,7 +394,7 @@ async def test_concurrent_sub_workflow_execution() -> None: .add_edge(processor, workflow_executor) .add_edge(workflow_executor, processor) .add_edge(workflow_executor, parent_request_info) # For external requests - .add_edge(parent_request_info, workflow_executor) # For SubWorkflowResponse routing + .add_edge(parent_request_info, workflow_executor) # For RequestResponse routing .build() ) @@ -497,12 +437,3 @@ async def test_concurrent_sub_workflow_execution() -> None: # Verify that concurrent executions were properly isolated # (This is implicitly tested by the fact that we got correct results for all emails) - - -if __name__ == "__main__": - # Run tests - asyncio.run(test_basic_sub_workflow()) - asyncio.run(test_sub_workflow_with_interception()) - asyncio.run(test_conditional_forwarding()) - asyncio.run(test_workflow_scoped_interception()) - asyncio.run(test_concurrent_sub_workflow_execution()) diff --git a/python/packages/main/tests/workflow/test_typing_utils.py b/python/packages/main/tests/workflow/test_typing_utils.py index 157225096a..db2d4da4a7 100644 --- a/python/packages/main/tests/workflow/test_typing_utils.py +++ b/python/packages/main/tests/workflow/test_typing_utils.py @@ -95,7 +95,8 @@ def test_request_response_type() -> None: """Test RequestResponse generic type checking.""" request_instance = RequestResponse[RequestInfoMessage, str]( - is_handled=False, + data="approve", + request_id="req-1", original_request=RequestInfoMessage(), ) diff --git a/python/samples/getting_started/workflow/README.md b/python/samples/getting_started/workflow/README.md index 2fdb3f2389..7eba2c95a9 100644 --- a/python/samples/getting_started/workflow/README.md +++ b/python/samples/getting_started/workflow/README.md @@ -49,8 +49,8 @@ Once comfortable with these, explore the rest of the samples below. | Sample | File | Concepts | |---|---|---| | Sub-Workflow (Basics) | [composition/sub_workflow_basics.py](./composition/sub_workflow_basics.py) | Wrap a workflow as an executor and orchestrate sub-workflows | -| Sub-Workflow: Request Interception | [composition/sub_workflow_request_interception.py](./composition/sub_workflow_request_interception.py) | Intercept and forward requests with decorators and request handling | -| Sub-Workflow: Parallel Requests | [composition/sub_workflow_parallel_requests.py](./composition/sub_workflow_parallel_requests.py) | Multi-type interception and external forwarding patterns | +| Sub-Workflow: Request Interception | [composition/sub_workflow_request_interception.py](./composition/sub_workflow_request_interception.py) | Intercept and forward sub-workflow requests using @handler for RequestInfoMessage subclasses | +| Sub-Workflow: Parallel Requests | [composition/sub_workflow_parallel_requests.py](./composition/sub_workflow_parallel_requests.py) | Multiple specialized interceptors handling different request types from same sub-workflow | ### control-flow | Sample | File | Concepts | diff --git a/python/samples/getting_started/workflow/composition/sub_workflow_parallel_requests.py b/python/samples/getting_started/workflow/composition/sub_workflow_parallel_requests.py index 22f3f37b64..c7d03f76b5 100644 --- a/python/samples/getting_started/workflow/composition/sub_workflow_parallel_requests.py +++ b/python/samples/getting_started/workflow/composition/sub_workflow_parallel_requests.py @@ -9,71 +9,60 @@ from typing_extensions import Never from agent_framework import ( Executor, RequestInfoExecutor, + RequestInfoMessage, + RequestResponse, WorkflowBuilder, WorkflowContext, + WorkflowExecutor, handler, ) -# Import the new sub-workflow types directly from the implementation package -try: - from agent_framework import ( - RequestInfoMessage, - RequestResponse, - WorkflowExecutor, - intercepts_request, - ) -except ImportError: - import os - import sys - - sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..", "..", "packages", "workflow")) - from agent_framework import ( - RequestInfoMessage, - RequestResponse, - WorkflowExecutor, - intercepts_request, - ) - """ -Sample: Sub-workflow with parallel requests +Sample: Sub-workflow with parallel request handling by specialized interceptors -This sample demonstrates the PROPER pattern for request interception. +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. Prerequisites: - No external services required (external handling simulated via `RequestInfoExecutor`). -Key principles: -1. Only ONE executor intercepts a given request type from a specific sub-workflow -2. Different executors can intercept DIFFERENT request types from the same sub-workflow -3. The same executor can intercept the same request type from DIFFERENT sub-workflows - -This ensures: -- Deterministic behavior -- Clear responsibility boundaries -- Easier debugging and maintenance +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 The example simulates a resource allocation system where: -- Sub-workflow requests resources (CPU, memory, etc.) -- A single Cache executor intercepts and handles resource requests -- The Cache can either satisfy from cache or forward to external +- 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 -Simple flow visualization: +Flow visualization: Coordinator | - | list[resource/policy requests] + | Mixed list[resource + policy requests] v [ Sub-workflow: WorkflowExecutor(ResourceRequester) ] - | | - | ResourceRequest | PolicyCheckRequest - v v - ResourceCache (@intercepts) PolicyEngine (@intercepts) - | handled/forward | handled/forward - v v - RequestInfo (external) <----- forwarded when not handled - | responses + | + | Emits different RequestInfoMessage types: + | - ResourceRequest + | - PolicyCheckRequest v - Back to sub-workflow -> completion -> results collected + 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 """ @@ -201,9 +190,9 @@ class ResourceRequester(Executor): return self._request_count == 0 -# 3. Implement the Resource Cache - ONLY intercepts ResourceRequest +# 3. Implement the Resource Cache - Uses typed handler for ResourceRequest class ResourceCache(Executor): - """Interceptor that handles RESOURCE requests from cache.""" + """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} @@ -213,26 +202,32 @@ class ResourceCache(Executor): super().__init__(id="resource_cache") # Instance initialization only; state kept in class attributes as above - @intercepts_request - async def check_cache( - self, request: ResourceRequest, ctx: WorkflowContext - ) -> RequestResponse[ResourceRequest, ResourceResponse]: - """Intercept RESOURCE requests and check cache first.""" - print(f"🏪 CACHE interceptor checking: {request.amount} {request.resource_type}") + @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(request.resource_type, 0) + available = self.cache.get(resource_request.resource_type, 0) - if available >= request.amount: + if available >= resource_request.amount: # We can satisfy from cache - self.cache[request.resource_type] -= request.amount - response = ResourceResponse(resource_type=request.resource_type, allocated=request.amount, source="cache") - print(f" ✅ Cache satisfied: {request.amount} {request.resource_type}") - self.results.append(response) - return RequestResponse[ResourceRequest, ResourceResponse].handled(response) + 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) - # Cache miss - forward to external - print(f" ❌ Cache miss: need {request.amount}, have {available} {request.resource_type}") - return RequestResponse[ResourceRequest, ResourceResponse].forward() + # 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( @@ -247,9 +242,9 @@ class ResourceCache(Executor): ) -# 4. Implement the Policy Engine - ONLY intercepts PolicyCheckRequest (different type!) +# 4. Implement the Policy Engine - Uses typed handler for PolicyCheckRequest class PolicyEngine(Executor): - """Interceptor that handles POLICY requests.""" + """Interceptor that handles POLICY requests using typed routing.""" # Use class attributes for simple sample state quota: dict[str, int] = { @@ -263,28 +258,35 @@ class PolicyEngine(Executor): super().__init__(id="policy_engine") # Instance initialization only; state kept in class attributes as above - @intercepts_request - async def check_policy( - self, request: PolicyCheckRequest, ctx: WorkflowContext - ) -> RequestResponse[PolicyCheckRequest, PolicyResponse]: - """Intercept POLICY requests and apply rules.""" - print(f"🛡️ POLICY interceptor checking: {request.amount} {request.resource_type}, policy={request.policy_type}") + @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(request.resource_type, 0) + 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 - if request.policy_type == "quota": - if request.amount <= quota_limit: - response = PolicyResponse(approved=True, reason=f"Within quota ({quota_limit})") - print(f" ✅ Policy approved: {request.amount} <= {quota_limit}") - self.results.append(response) - return RequestResponse[PolicyCheckRequest, PolicyResponse].handled(response) # Exceeds quota - forward to external for review - print(f" ❌ Policy exceeds quota: {request.amount} > {quota_limit}, forwarding to external") - return RequestResponse[PolicyCheckRequest, PolicyResponse].forward() + print(f" ❌ Policy exceeds quota: {policy_request.amount} > {quota_limit}, forwarding to external") + await ctx.send_message(request) + return # Unknown policy type - forward to external - print(f" ❓ Unknown policy type: {request.policy_type}, forwarding") - return RequestResponse[PolicyCheckRequest, PolicyResponse].forward() + print(f" ❓ Unknown policy type: {policy_request.policy_type}, forwarding") + await ctx.send_message(request) @handler async def collect_policy_result( @@ -340,20 +342,19 @@ async def main() -> None: # Create a simple coordinator that starts the process coordinator = Coordinator() - # PROPER PATTERN: Each executor intercepts DIFFERENT request types + # TYPED ROUTING: Each executor handles specific typed RequestInfoMessage messages 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) # Cache intercepts ResourceRequest - .add_edge(cache, workflow_executor) # Cache handles ResourceRequest - .add_edge(workflow_executor, policy) # Policy handles PolicyCheckRequest - .add_edge(policy, workflow_executor) # Policy intercepts PolicyCheckRequest - .add_edge(cache, main_request_info) # Cache forwards to external - .add_edge(policy, main_request_info) # Policy forwards to external - .add_edge(main_request_info, workflow_executor) # External responses back - .add_edge(workflow_executor, main_request_info) # Sub-workflow forwards to main + .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 .build() ) diff --git a/python/samples/getting_started/workflow/composition/sub_workflow_request_interception.py b/python/samples/getting_started/workflow/composition/sub_workflow_request_interception.py index dfa1dc57b6..9c1bbd58eb 100644 --- a/python/samples/getting_started/workflow/composition/sub_workflow_request_interception.py +++ b/python/samples/getting_started/workflow/composition/sub_workflow_request_interception.py @@ -12,7 +12,6 @@ from agent_framework import ( WorkflowContext, WorkflowExecutor, handler, - intercepts_request, ) """ @@ -20,8 +19,8 @@ Sample: Sub-Workflows with Request Interception This sample shows how to: 1. Create workflows that execute other workflows as sub-workflows -2. Intercept requests from sub-workflows in parent workflows using @intercepts_request -3. Conditionally handle or forward requests using RequestResponse.handled() and RequestResponse.forward() +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 @@ -35,9 +34,8 @@ The example simulates an email validation system where: Key concepts demonstrated: - WorkflowExecutor: Wraps a workflow to make it behave as an executor -- @intercepts_request: Decorator for parent workflows to handle sub-workflow requests -- RequestResponse: Enables conditional handling vs forwarding of requests -- Request correlation: Using request_id to match responses with original requests +- 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 @@ -48,19 +46,19 @@ Prerequisites: Simple flow visualization: - Parent Orchestrator (@intercepts_request) + Parent Orchestrator (handles DomainCheckRequest) | | EmailValidationRequest(email) x3 (concurrent) v [ Sub-workflow: WorkflowExecutor(EmailValidator) ] | - | DomainCheckRequest(domain) with request_id correlation + | DomainCheckRequest(domain) with request_id and source_executor_id v - Interception? yes -> handled locally with RequestResponse.handled(True) + Interception? yes -> handled locally with RequestResponse(data=True) no -> forwarded to RequestInfoExecutor -> external service | v - Response routed back to sub-workflow using request_id + Response routed back to sub-workflow using source_executor_id """ @@ -92,7 +90,7 @@ class ValidationResult: class EmailValidator(Executor): """Validates email addresses - doesn't know it's in a sub-workflow.""" - def __init__(self): + def __init__(self) -> None: """Initialize the EmailValidator executor.""" super().__init__(id="email_validator") # Use a dict to track multiple pending emails by request_id @@ -180,17 +178,28 @@ class SmartEmailOrchestrator(Executor): request = EmailValidationRequest(email=email) await ctx.send_message(request, target_id="email_validator_workflow") - @intercepts_request - async def check_domain( - self, request: DomainCheckRequest, ctx: WorkflowContext - ) -> RequestResponse[DomainCheckRequest, bool]: - """Intercept domain check requests from sub-workflows.""" + @handler + async def handle_domain_request( + self, + request: DomainCheckRequest, + ctx: WorkflowContext[RequestResponse[DomainCheckRequest, bool] | DomainCheckRequest] + ) -> None: + """Handle requests from sub-workflows.""" print(f"🔍 Parent intercepting domain check for: {request.domain}") + if request.domain in self.approved_domains: print(f"✅ Domain '{request.domain}' is pre-approved locally!") - return RequestResponse[DomainCheckRequest, bool].handled(True) - print(f"❓ Domain '{request.domain}' unknown, forwarding to external service...") - return RequestResponse[DomainCheckRequest, bool].forward() + # 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) @handler async def collect_result(self, result: ValidationResult, ctx: WorkflowContext) -> None: @@ -233,7 +242,7 @@ async def run_example() -> None: WorkflowBuilder() .set_start_executor(orchestrator) .add_edge(orchestrator, workflow_executor) - .add_edge(workflow_executor, orchestrator) + .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