From 6e9d35830fd6d78d12bc2680ce980f5b716e175e Mon Sep 17 00:00:00 2001 From: Eric Zhu Date: Fri, 22 Aug 2025 15:02:34 -0700 Subject: [PATCH] Python: Move Workflow, Edge, EdgeGroup and Executor to AFBaseModel (#472) * Refactor workflow to introduce EdgeRunner for edge execution. * Fix edge cases * Convert Workflow, Edge, EdgeGroup, and Executor into AFBaseModel to support object model serialization * format * remove accidental file * fix typing * Add type information to EdgeGroup and Executor subclasses * fix format * Add condition_name field to Edge * Add new fields * remove Optional * Update --- .../agent_framework_workflow/_edge.py | 412 +++++------- .../agent_framework_workflow/_edge_runner.py | 198 ++++++ .../agent_framework_workflow/_executor.py | 57 +- .../agent_framework_workflow/_runner.py | 43 +- .../agent_framework_workflow/_validation.py | 51 +- .../workflow/agent_framework_workflow/_viz.py | 28 +- .../agent_framework_workflow/_workflow.py | 132 ++-- python/packages/workflow/tests/test_edge.py | 410 +++++------ python/packages/workflow/tests/test_runner.py | 29 +- .../workflow/tests/test_serialization.py | 636 ++++++++++++++++++ .../workflow/tests/test_validation.py | 15 +- .../packages/workflow/tests/test_workflow.py | 208 ++++-- .../workflow/tests/test_workflow_builder.py | 2 +- 13 files changed, 1588 insertions(+), 633 deletions(-) create mode 100644 python/packages/workflow/agent_framework_workflow/_edge_runner.py create mode 100644 python/packages/workflow/tests/test_serialization.py diff --git a/python/packages/workflow/agent_framework_workflow/_edge.py b/python/packages/workflow/agent_framework_workflow/_edge.py index fef5e3376d..2e30c414fe 100644 --- a/python/packages/workflow/agent_framework_workflow/_edge.py +++ b/python/packages/workflow/agent_framework_workflow/_edge.py @@ -1,71 +1,74 @@ # Copyright (c) Microsoft. All rights reserved. -import asyncio import logging import uuid -from collections import defaultdict from collections.abc import Callable, Sequence from dataclasses import dataclass from typing import Any, ClassVar -from ._executor import Executor -from ._runner_context import Message, RunnerContext -from ._shared_state import SharedState -from ._workflow_context import WorkflowContext +from agent_framework._pydantic import AFBaseModel +from pydantic import Field + +from agent_framework_workflow._executor import Executor logger = logging.getLogger(__name__) -class Edge: +def _extract_function_name(func: Callable[..., Any]) -> str: + """Extract the name of any callable function for serialization. + + Args: + func: The function to extract the name from. + + Returns: + The name of the function, or a placeholder for lambda functions. + """ + if hasattr(func, "__name__"): + name = func.__name__ + # Check if it's a lambda function + if name == "": + return "" + return name + # Fallback for other callable objects + return "" + + +class Edge(AFBaseModel): """Represents a directed edge in a graph.""" ID_SEPARATOR: ClassVar[str] = "->" + source_id: str = Field(min_length=1, description="The ID of the source executor of the edge") + target_id: str = Field(min_length=1, description="The ID of the target executor of the edge") + condition_name: str | None = Field(default=None, description="The name of the condition function for serialization") + def __init__( self, - source: Executor, - target: Executor, + source_id: str, + target_id: str, condition: Callable[[Any], bool] | None = None, + **kwargs: Any, ) -> None: """Initialize the edge with a source and target node. Args: - source (Executor): The source executor of the edge. - target (Executor): The target executor of the edge. + source_id (str): The ID of the source executor of the edge. + target_id (str): The ID of the target executor of the edge. condition (Callable[[Any], bool], optional): A condition function that determines if the edge can handle the data. If None, the edge can handle any data type. Defaults to None. + kwargs: Additional keyword arguments. Unused in this implementation. """ - self.source = source - self.target = target + condition_name = _extract_function_name(condition) if condition is not None else None + kwargs.update({"source_id": source_id, "target_id": target_id, "condition_name": condition_name}) + super().__init__(**kwargs) self._condition = condition - @property - def source_id(self) -> str: - """Get the source executor ID.""" - return self.source.id - - @property - def target_id(self) -> str: - """Get the target executor ID.""" - return self.target.id - @property def id(self) -> str: """Get the unique ID of the edge.""" return f"{self.source_id}{self.ID_SEPARATOR}{self.target_id}" - def can_handle(self, message_data: Any) -> bool: - """Check if the edge can handle the given data. - - Args: - message_data (Any): The data to check. - - Returns: - bool: True if the edge can handle the data, False otherwise. - """ - return self.target.can_handle(message_data) - def should_route(self, data: Any) -> bool: """Determine if message should be routed through this edge based on the condition.""" if self._condition is None: @@ -73,69 +76,45 @@ class Edge: return self._condition(data) - async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> None: - """Send a message along this edge. - Args: - message (Message): The message to send. - shared_state (SharedState): The shared state to use for holding data. - ctx (RunnerContext): The context for the runner. - """ - if not self.can_handle(message.data): - # Caller of this method should ensure that the edge can handle the data. - raise RuntimeError(f"Edge {self.id} cannot handle data of type {type(message.data)}.") - - if self.should_route(message.data): - await self.target.execute( - message.data, WorkflowContext(self.target.id, [self.source.id], shared_state, ctx) - ) - - -class EdgeGroup: +class EdgeGroup(AFBaseModel): """Represents a group of edges that share some common properties and can be triggered together.""" - def __init__(self) -> None: + id: str = Field( + default_factory=lambda: f"EdgeGroup/{uuid.uuid4()}", description="Unique identifier for the edge group" + ) + type: str = Field(description="The type of edge group, corresponding to the class name") + edges: list[Edge] = Field(default_factory=list, description="List of edges in this group") + + def __init__(self, **kwargs: Any) -> None: """Initialize the edge group.""" - self._id = f"{self.__class__.__name__}/{uuid.uuid4()}" - - async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool: - """Send a message through the edge group. - - Args: - message (Message): The message to send. - shared_state (SharedState): The shared state to use for holding data. - ctx (RunnerContext): The context for the runner. - - Returns: - bool: True if the message was sent successfully, False if the target executor cannot handle the message. - If a message can be delivered but rejected due to a condition, it will still return True. - - Note: - Exception will not be raised if the target executor cannot handle the message. This is because - a source executor can be connected to multiple target executors, and not every target executor may - be able to handle all the messages sent by the source executor. - """ - raise NotImplementedError + if "id" not in kwargs: + kwargs["id"] = f"{self.__class__.__name__}/{uuid.uuid4()}" + if "type" not in kwargs: + kwargs["type"] = self.__class__.__name__ + super().__init__(**kwargs) @property - def id(self) -> str: - """Get the unique ID of the edge group.""" - return self._id - - @property - def source_executors(self) -> list[Executor]: + def source_executor_ids(self) -> list[str]: """Get the source executor IDs of the edges in the group.""" - raise NotImplementedError + seen = set() + result = [] + for edge in self.edges: + if edge.source_id not in seen: + result.append(edge.source_id) + seen.add(edge.source_id) + return result @property - def target_executors(self) -> list[Executor]: + def target_executor_ids(self) -> list[str]: """Get the target executor IDs of the edges in the group.""" - raise NotImplementedError - - @property - def edges(self) -> list[Edge]: - """Get the edges in the group.""" - raise NotImplementedError + seen = set() + result = [] + for edge in self.edges: + if edge.target_id not in seen: + result.append(edge.target_id) + seen.add(edge.target_id) + return result class SingleEdgeGroup(EdgeGroup): @@ -144,43 +123,22 @@ class SingleEdgeGroup(EdgeGroup): A concrete implementation of EdgeGroup that represent a group containing exactly one edge. """ - def __init__(self, source: Executor, target: Executor, condition: Callable[[Any], bool] | None = None) -> None: + def __init__( + self, source_id: str, target_id: str, condition: Callable[[Any], bool] | None = None, **kwargs: Any + ) -> None: """Initialize the single edge group with an edge. Args: - source (Executor): The source executor. - target (Executor): The target executor that the source executor can send messages to. + source_id (str): The source executor ID. + target_id (str): The target executor ID that the source executor can send messages to. condition (Callable[[Any], bool], optional): A condition function that determines - if the edge will pass the data to the target executor. If None, the edge can - will always pass the data to the target executor. + if the edge will pass the data to the target executor. If None, the edge will + always pass the data to the target executor. + kwargs: Additional keyword arguments. Unused in this implementation. """ - self._edge = Edge(source=source, target=target, condition=condition) - - async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool: - """Send a message through the single edge.""" - if message.target_id and message.target_id != self._edge.target_id: - return False - - if self._edge.can_handle(message.data): - await self._edge.send_message(message, shared_state, ctx) - return True - - return False - - @property - def source_executors(self) -> list[Executor]: - """Get the source executor of the edge.""" - return [self._edge.source] - - @property - def target_executors(self) -> list[Executor]: - """Get the target executor of the edge.""" - return [self._edge.target] - - @property - def edges(self) -> list[Edge]: - """Get the edges in the group.""" - return [self._edge] + edge = Edge(source_id=source_id, target_id=target_id, condition=condition) + kwargs["edges"] = [edge] + super().__init__(**kwargs) class FanOutEdgeGroup(EdgeGroup): @@ -190,78 +148,51 @@ class FanOutEdgeGroup(EdgeGroup): and send messages to their respective target executors. """ + selection_func_name: str | None = Field( + default=None, description="The name of the selection function for serialization" + ) + def __init__( self, - source: Executor, - targets: Sequence[Executor], + source_id: str, + target_ids: Sequence[str], selection_func: Callable[[Any, list[str]], list[str]] | None = None, + **kwargs: Any, ) -> None: """Initialize the fan-out edge group with a list of edges. Args: - source (Executor): The source executor. - targets (Sequence[Executor]): A list of target executors that the source executor can send messages to. + source_id (str): The source executor ID. + target_ids (Sequence[str]): A list of target executor IDs that the source executor can send messages to. selection_func (Callable[[Any, list[str]], list[str]], optional): A function that selects which target executors to send messages to. The function takes in the message data and a list of target executor IDs, and returns a list of selected target executor IDs. + kwargs: Additional keyword arguments. Unused in this implementation. """ - if len(targets) <= 1: + if len(target_ids) <= 1: raise ValueError("FanOutEdgeGroup must contain at least two targets.") - self._edges = [Edge(source=source, target=target) for target in targets] - self._target_ids = [edge.target_id for edge in self._edges] - self._target_map = {edge.target_id: edge for edge in self._edges} + + # Extract selection function name for serialization + selection_func_name = None + if selection_func is not None: + selection_func_name = _extract_function_name(selection_func) + + edges = [Edge(source_id=source_id, target_id=target_id) for target_id in target_ids] + kwargs.update({"edges": edges, "selection_func_name": selection_func_name}) + super().__init__(**kwargs) + + self._target_ids = list(target_ids) self._selection_func = selection_func - async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool: - """Send a message through all edges in the fan-out edge group.""" - selection_results = ( - self._selection_func(message.data, self._target_ids) if self._selection_func else self._target_ids - ) - if not self._validate_selection_result(selection_results): - raise RuntimeError( - f"Invalid selection result: {selection_results}. " - f"Expected selections to be a subset of valid target executor IDs: {self._target_ids}." - ) - - if message.target_id: - # If the target ID is specified and the selection result contains it, send the message to that edge - if message.target_id in selection_results: - edge = next((edge for edge in self._edges if edge.target_id == message.target_id), None) - if edge and edge.can_handle(message.data): - await edge.send_message(message, shared_state, ctx) - return True - return False - - # If no target ID, send the message to the selected targets - async def send_to_edge(edge: Edge) -> bool: - """Send the message to the edge at the specified index.""" - if edge.can_handle(message.data): - await edge.send_message(message, shared_state, ctx) - return True - return False - - tasks = [send_to_edge(self._target_map[target_id]) for target_id in selection_results] - results = await asyncio.gather(*tasks) - return any(results) + @property + def target_ids(self) -> list[str]: + """Get the target executor IDs for selection.""" + return self._target_ids @property - def source_executors(self) -> list[Executor]: - """Get the source executor of the edges in the group.""" - return [self._edges[0].source] - - @property - def target_executors(self) -> list[Executor]: - """Get the target executors of the edges in the group.""" - return [edge.target for edge in self._edges] - - @property - def edges(self) -> list[Edge]: - """Get the edges in the group.""" - return self._edges - - def _validate_selection_result(self, selection_results: list[str]) -> bool: - """Validate the selection results to ensure all IDs are valid target executor IDs.""" - return all(result in self._target_ids for result in selection_results) + def selection_func(self) -> Callable[[Any, list[str]], list[str]] | None: + """Get the selection function for this fan-out group.""" + return self._selection_func class FanInEdgeGroup(EdgeGroup): @@ -271,68 +202,25 @@ class FanInEdgeGroup(EdgeGroup): Messages are buffered until all edges in the group have data to send. """ - def __init__(self, sources: Sequence[Executor], target: Executor) -> None: + def __init__(self, source_ids: Sequence[str], target_id: str, **kwargs: Any) -> None: """Initialize the fan-in edge group with a list of edges. Args: - sources (Sequence[Executor]): A list of source executors that can send messages to the target executor. - target (Executor): The target executor that receives a list of messages aggregated from all sources. + source_ids (Sequence[str]): A list of source executor IDs that can send messages to the target executor. + target_id (str): The target executor ID that receives a list of messages aggregated from all sources. + kwargs: Additional keyword arguments. Unused in this implementation. """ - if len(sources) <= 1: + if len(source_ids) <= 1: raise ValueError("FanInEdgeGroup must contain at least two sources.") - self._edges = [Edge(source=source, target=target) for source in sources] - # Buffer to hold messages before sending them to the target executor - # Key is the source executor ID, value is a list of messages - self._buffer: dict[str, list[Message]] = defaultdict(list) - async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool: - """Send a message through all edges in the fan-in edge group.""" - if message.target_id and message.target_id != self._edges[0].target_id: - return False - - if self._edges[0].can_handle([message.data]): - # If the edge can handle the data, buffer the message - self._buffer[message.source_id].append(message) - else: - # If the edge cannot handle the data, return False - return False - - if self._is_ready_to_send(): - # If all edges in the group have data, send the buffered messages to the target executor - messages_to_send = [msg for edge in self._edges for msg in self._buffer[edge.source_id]] - self._buffer.clear() - # Only trigger one edge to send the messages to avoid duplicate sends - await self._edges[0].send_message( - Message([msg.data for msg in messages_to_send], self.__class__.__name__), - shared_state, - ctx, - ) - - return True - - def _is_ready_to_send(self) -> bool: - """Check if all edges in the group have data to send.""" - return all(self._buffer[edge.source_id] for edge in self._edges) - - @property - def source_executors(self) -> list[Executor]: - """Get the source executors of the edges in the group.""" - return [edge.source for edge in self._edges] - - @property - def target_executors(self) -> list[Executor]: - """Get the target executor of the edges in the group.""" - return [self._edges[0].target] - - @property - def edges(self) -> list[Edge]: - """Get the edges in the group.""" - return self._edges + edges = [Edge(source_id=source_id, target_id=target_id) for source_id in source_ids] + kwargs["edges"] = edges + super().__init__(**kwargs) @dataclass class Case: - """Represents a single case in the conditional edge group. + """Represents a single case in the switch-case edge group. Args: condition (Callable[[Any], bool]): The condition function for the case. @@ -345,7 +233,7 @@ class Case: @dataclass class Default: - """Represents the default case in the conditional edge group. + """Represents the default case in the switch-case edge group. Args: target (Executor): The target executor for the default case. @@ -354,6 +242,39 @@ class Default: target: Executor +class SwitchCaseEdgeGroupCase(AFBaseModel): + """A single case in the SwitchCaseEdgeGroup. This is used internally.""" + + target_id: str = Field(description="The target executor ID for this case") + condition_name: str | None = Field(default=None, description="The name of the condition function for serialization") + type: str = Field(default="Case", description="The type of the case") + + def __init__(self, condition: Callable[[Any], bool], target_id: str, **kwargs: Any) -> None: + """Initialize the switch case with a condition and target. + + Args: + condition: The condition function for the case. + target_id: The target executor ID for this case. + kwargs: Additional keyword arguments. + """ + condition_name = _extract_function_name(condition) + kwargs.update({"target_id": target_id, "condition_name": condition_name}) + super().__init__(**kwargs) + self._condition = condition + + @property + def condition(self) -> Callable[[Any], bool]: + """Get the condition function for this case.""" + return self._condition + + +class SwitchCaseEdgeGroupDefault(AFBaseModel): + """The default case in the SwitchCaseEdgeGroup. This is used internally.""" + + target_id: str = Field(description="The target executor ID for the default case") + type: str = Field(default="Default", description="The type of the case") + + class SwitchCaseEdgeGroup(FanOutEdgeGroup): """Represents a group of edges that assemble a conditional routing pattern. @@ -377,43 +298,52 @@ class SwitchCaseEdgeGroup(FanOutEdgeGroup): edge_4 """ + cases: list[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault] = Field( + default_factory=list, description="List of conditional cases for this switch-case group" + ) + def __init__( self, - source: Executor, - cases: Sequence[Case | Default], + source_id: str, + cases: Sequence[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault], + **kwargs: Any, ) -> None: - """Initialize the conditional edge group with a list of edges. + """Initialize the switch-case edge group with a list of edges. Args: - source (Executor): The source executor. - cases (Sequence[Case | Default]): A list of cases for the conditional edge group. + source_id (str): The source executor ID. + cases (Sequence[Case | Default]): A list of cases for the switch-case edge group. There should be exactly one default case. + kwargs: Additional keyword arguments. Unused in this implementation. """ if len(cases) < 2: raise ValueError("SwitchCaseEdgeGroup must contain at least two cases (including the default case).") - default_case = [isinstance(case, Default) for case in cases] + default_case = [isinstance(case, SwitchCaseEdgeGroupDefault) for case in cases] if sum(default_case) != 1: raise ValueError("SwitchCaseEdgeGroup must contain exactly one default case.") - if isinstance(cases[-1], Default): + if not isinstance(cases[-1], SwitchCaseEdgeGroupDefault): logger.warning( - "Default case in the conditional edge group is not the last case. " + "Default case in the switch-case edge group is not the last case. " "This will result in unexpected behavior." ) def selection_func(data: Any, targets: list[str]) -> list[str]: """Select the target executor based on the conditions.""" for index, case in enumerate(cases): - if isinstance(case, Default): - return [case.target.id] - if isinstance(case, Case): + if isinstance(case, SwitchCaseEdgeGroupDefault): + return [case.target_id] + if isinstance(case, SwitchCaseEdgeGroupCase): try: if case.condition(data): - return [case.target.id] + return [case.target_id] except Exception as e: logger.warning(f"Error occurred while evaluating condition for case {index}: {e}") raise RuntimeError("No matching case found in SwitchCaseEdgeGroup.") - super().__init__(source, [case.target for case in cases], selection_func=selection_func) + target_ids = [case.target_id for case in cases] + + kwargs.update({"cases": cases}) + super().__init__(source_id, target_ids, selection_func=selection_func, **kwargs) diff --git a/python/packages/workflow/agent_framework_workflow/_edge_runner.py b/python/packages/workflow/agent_framework_workflow/_edge_runner.py new file mode 100644 index 0000000000..184cf0b940 --- /dev/null +++ b/python/packages/workflow/agent_framework_workflow/_edge_runner.py @@ -0,0 +1,198 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import logging +from abc import ABC, abstractmethod +from collections import defaultdict +from typing import Any + +from ._edge import Edge, EdgeGroup, FanInEdgeGroup, FanOutEdgeGroup, SingleEdgeGroup, SwitchCaseEdgeGroup +from ._executor import Executor +from ._runner_context import Message, RunnerContext +from ._shared_state import SharedState +from ._workflow_context import WorkflowContext + +logger = logging.getLogger(__name__) + + +class EdgeRunner(ABC): + """Abstract base class for edge runners that handle message delivery.""" + + def __init__(self, edge_group: EdgeGroup, executors: dict[str, Executor]) -> None: + """Initialize the edge runner with an edge group and executor map. + + Args: + edge_group: The edge group to run. + executors: Map of executor IDs to executor instances. + """ + self._edge_group = edge_group + self._executors = executors + + @abstractmethod + async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool: + """Send a message through the edge group. + + Args: + message: The message to send. + shared_state: The shared state to use for holding data. + ctx: The context for the runner. + + Returns: + bool: True if the message was sent successfully, False if the target executor cannot handle the message. + """ + raise NotImplementedError + + def _can_handle(self, executor_id: str, message_data: Any) -> bool: + """Check if an executor can handle the given message data.""" + if executor_id not in self._executors: + return False + return self._executors[executor_id].can_handle(message_data) + + async def _execute_on_target( + self, target_id: str, source_id: str, message_data: Any, shared_state: SharedState, ctx: RunnerContext + ) -> None: + """Execute a message on a target executor.""" + if target_id not in self._executors: + raise RuntimeError(f"Target executor {target_id} not found.") + + target_executor = self._executors[target_id] + await target_executor.execute(message_data, WorkflowContext(target_id, [source_id], shared_state, ctx)) + + +class SingleEdgeRunner(EdgeRunner): + """Runner for single edge groups.""" + + def __init__(self, edge_group: SingleEdgeGroup, executors: dict[str, Executor]) -> None: + super().__init__(edge_group, executors) + self._edge = edge_group.edges[0] + + async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool: + """Send a message through the single edge.""" + if message.target_id and message.target_id != self._edge.target_id: + return False + + if self._can_handle(self._edge.target_id, message.data): + if self._edge.should_route(message.data): + await self._execute_on_target( + self._edge.target_id, self._edge.source_id, message.data, shared_state, ctx + ) + return True + + return False + + +class FanOutEdgeRunner(EdgeRunner): + """Runner for fan-out edge groups.""" + + def __init__(self, edge_group: FanOutEdgeGroup, executors: dict[str, Executor]) -> None: + super().__init__(edge_group, executors) + self._edges = edge_group.edges + self._target_ids = edge_group.target_ids + self._target_map = {edge.target_id: edge for edge in self._edges} + self._selection_func = edge_group.selection_func + + async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool: + """Send a message through all edges in the fan-out edge group.""" + selection_results = ( + self._selection_func(message.data, self._target_ids) if self._selection_func else self._target_ids + ) + if not self._validate_selection_result(selection_results): + raise RuntimeError( + f"Invalid selection result: {selection_results}. " + f"Expected selections to be a subset of valid target executor IDs: {self._target_ids}." + ) + + if message.target_id: + # If the target ID is specified and the selection result contains it, send the message to that edge + if message.target_id in selection_results: + edge = self._target_map.get(message.target_id) + if edge and self._can_handle(edge.target_id, message.data): + if edge.should_route(message.data): + await self._execute_on_target(edge.target_id, edge.source_id, message.data, shared_state, ctx) + return True + return False + + # If no target ID, send the message to the selected targets + async def send_to_edge(edge: Edge) -> bool: + """Send the message to the edge.""" + if self._can_handle(edge.target_id, message.data): + if edge.should_route(message.data): + await self._execute_on_target(edge.target_id, edge.source_id, message.data, shared_state, ctx) + return True + return False + + tasks = [send_to_edge(self._target_map[target_id]) for target_id in selection_results] + results = await asyncio.gather(*tasks) + return any(results) + + def _validate_selection_result(self, selection_results: list[str]) -> bool: + """Validate the selection results to ensure all IDs are valid target executor IDs.""" + return all(result in self._target_ids for result in selection_results) + + +class FanInEdgeRunner(EdgeRunner): + """Runner for fan-in edge groups.""" + + def __init__(self, edge_group: FanInEdgeGroup, executors: dict[str, Executor]) -> None: + super().__init__(edge_group, executors) + self._edges = edge_group.edges + # Buffer to hold messages before sending them to the target executor + # Key is the source executor ID, value is a list of messages + self._buffer: dict[str, list[Message]] = defaultdict(list) + + async def send_message(self, message: Message, shared_state: SharedState, ctx: RunnerContext) -> bool: + """Send a message through all edges in the fan-in edge group.""" + if message.target_id and message.target_id != self._edges[0].target_id: + return False + + # Check if target can handle list of message data (fan-in aggregates multiple messages) + if self._can_handle(self._edges[0].target_id, [message.data]): + # If the edge can handle the data, buffer the message + self._buffer[message.source_id].append(message) + else: + # If the edge cannot handle the data, return False + return False + + if self._is_ready_to_send(): + # If all edges in the group have data, send the buffered messages to the target executor + messages_to_send = [msg for edge in self._edges for msg in self._buffer[edge.source_id]] + self._buffer.clear() + # Send aggregated data to target + aggregated_data = [msg.data for msg in messages_to_send] + await self._execute_on_target( + self._edges[0].target_id, self._edge_group.__class__.__name__, aggregated_data, shared_state, ctx + ) + + return True + + def _is_ready_to_send(self) -> bool: + """Check if all edges in the group have data to send.""" + return all(self._buffer[edge.source_id] for edge in self._edges) + + +class SwitchCaseEdgeRunner(FanOutEdgeRunner): + """Runner for switch-case edge groups (inherits from FanOutEdgeRunner).""" + + def __init__(self, edge_group: SwitchCaseEdgeGroup, executors: dict[str, Executor]) -> None: + super().__init__(edge_group, executors) + + +def create_edge_runner(edge_group: EdgeGroup, executors: dict[str, Executor]) -> EdgeRunner: + """Factory function to create the appropriate edge runner for an edge group. + + Args: + edge_group: The edge group to create a runner for. + executors: Map of executor IDs to executor instances. + + Returns: + The appropriate EdgeRunner instance. + """ + if isinstance(edge_group, SingleEdgeGroup): + return SingleEdgeRunner(edge_group, executors) + if isinstance(edge_group, SwitchCaseEdgeGroup): + return SwitchCaseEdgeRunner(edge_group, executors) + if isinstance(edge_group, FanOutEdgeGroup): + return FanOutEdgeRunner(edge_group, executors) + if isinstance(edge_group, FanInEdgeGroup): + return FanInEdgeRunner(edge_group, executors) + raise ValueError(f"Unsupported edge group type: {type(edge_group)}") diff --git a/python/packages/workflow/agent_framework_workflow/_executor.py b/python/packages/workflow/agent_framework_workflow/_executor.py index a413edce3d..68aaa1bcfc 100644 --- a/python/packages/workflow/agent_framework_workflow/_executor.py +++ b/python/packages/workflow/agent_framework_workflow/_executor.py @@ -10,6 +10,8 @@ from types import UnionType from typing import Any, ClassVar, TypeVar, Union, get_args, get_origin, overload from agent_framework import AgentRunResponse, AgentRunResponseUpdate, AgentThread, AIAgent, ChatMessage +from agent_framework._pydantic import AFBaseModel +from pydantic import Field from ._events import ( AgentRunEvent, @@ -24,16 +26,33 @@ from ._workflow_context import WorkflowContext # region Executor -class Executor: +class Executor(AFBaseModel): """An executor is a component that processes messages in a workflow.""" - def __init__(self, id: str | None = None) -> None: + # Provide a default so static analyzers (e.g., pyright) don't require passing `id`. + # Runtime still sets a concrete value in __init__. + id: str = Field( + default_factory=lambda: str(uuid.uuid4()), + min_length=1, + description="Unique identifier for the executor", + ) + type: str = Field(default="", description="The type of executor, corresponding to the class name") + + def __init__(self, id: str | None = None, **kwargs: Any) -> None: """Initialize the executor with a unique identifier. Args: - id: A unique identifier for the executor. If None, a new UUID will be generated. + id: A unique identifier for the executor. If None, a new ID will be generated + following the format /. + kwargs: Additional keyword arguments. Unused in this implementation. """ - self._id = id or f"{self.__class__.__name__}/{uuid.uuid4()}" + executor_id = f"{self.__class__.__name__}/{uuid.uuid4()}" if id is None else id + + kwargs.update({"id": executor_id}) + if "type" not in kwargs: + kwargs["type"] = self.__class__.__name__ + + super().__init__(**kwargs) self._handlers: dict[type, Callable[[Any, WorkflowContext[Any]], Any]] = {} self._discover_handlers() @@ -67,22 +86,24 @@ class Executor: await handler(message, context) await context.add_event(ExecutorCompletedEvent(self.id)) - @property - def id(self) -> str: - """Get the unique identifier of the executor.""" - return self._id - def _discover_handlers(self) -> None: """Discover message handlers in the executor class.""" - for attr_name in dir(self): - attr = getattr(self, attr_name) - if callable(attr) and hasattr(attr, "_handler_spec"): - handler_spec = attr._handler_spec # type: ignore - if self._handlers.get(handler_spec["message_type"]) is not None: - raise ValueError( - f"Duplicate handler for type {handler_spec['message_type']} in {self.__class__.__name__}" - ) - self._handlers[handler_spec["message_type"]] = attr + # 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) and hasattr(attr, "_handler_spec"): + handler_spec = attr._handler_spec # type: ignore + if self._handlers.get(handler_spec["message_type"]) is not None: + raise ValueError( + f"Duplicate handler for type {handler_spec['message_type']} in {self.__class__.__name__}" + ) + # Get the bound method + bound_method = getattr(self, attr_name) + self._handlers[handler_spec["message_type"]] = bound_method + except AttributeError: + # Skip attributes that may not be accessible + continue def can_handle(self, message: Any) -> bool: """Check if the executor can handle a given message type. diff --git a/python/packages/workflow/agent_framework_workflow/_runner.py b/python/packages/workflow/agent_framework_workflow/_runner.py index 4235e6842d..1b054187de 100644 --- a/python/packages/workflow/agent_framework_workflow/_runner.py +++ b/python/packages/workflow/agent_framework_workflow/_runner.py @@ -7,6 +7,7 @@ from collections.abc import AsyncIterable, Sequence from typing import Any from ._edge import EdgeGroup +from ._edge_runner import EdgeRunner, create_edge_runner from ._events import WorkflowEvent from ._executor import Executor from ._runner_context import Message, RunnerContext @@ -21,6 +22,7 @@ class Runner: def __init__( self, edge_groups: Sequence[EdgeGroup], + executors: dict[str, Executor], shared_state: SharedState, ctx: RunnerContext, max_iterations: int = 100, @@ -30,12 +32,15 @@ class Runner: Args: edge_groups: The edge groups of the workflow. + executors: Map of executor IDs to executor instances. shared_state: The shared state for the workflow. ctx: The runner context for the workflow. max_iterations: The maximum number of iterations to run. workflow_id: The workflow ID for checkpointing. """ - self._edge_group_map = self._parse_edge_groups(edge_groups) + self._executors = executors + self._edge_runners = [create_edge_runner(group, executors) for group in edge_groups] + self._edge_runner_map = self._parse_edge_runners(self._edge_runners) self._ctx = ctx self._iteration = 0 self._max_iterations = max_iterations @@ -127,14 +132,14 @@ 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.""" - async def _deliver_message_inner(edge_group: EdgeGroup, message: Message) -> bool: - """Inner loop to deliver a single message through an edge group.""" - return await edge_group.send_message(message, self._shared_state, self._ctx) + 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) - associated_edge_groups = self._edge_group_map.get(source_executor_id, []) + associated_edge_runners = self._edge_runner_map.get(source_executor_id, []) for message in messages: - # Deliver a message through all edge groups associated with the source executor concurrently. - tasks = [_deliver_message_inner(edge_group, message) for edge_group in associated_edge_groups] + # 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] results = await asyncio.gather(*tasks) if not any(results): logger.warning( @@ -175,13 +180,7 @@ class Runner: - Else if it has a plain attribute `state` that is a dict, use that. Only JSON-serializable dicts should be provided by executors. """ - executors: dict[str, Executor] = {} - for edge_groups in self._edge_group_map.values(): - for edge_group in edge_groups: - for edge in edge_group.edges: - executors[edge.source.id] = edge.source - executors[edge.target.id] = edge.target - for exec_id, executor in executors.items(): + for exec_id, executor in self._executors.items(): state_dict: dict[str, Any] | None = None snapshot = getattr(executor, "snapshot_state", None) try: @@ -268,18 +267,18 @@ class Runner: except Exception as e: logger.warning(f"Failed to restore shared state from context: {e}") - def _parse_edge_groups(self, edge_groups: Sequence[EdgeGroup]) -> dict[str, list[EdgeGroup]]: - """Parse the edge groups of the workflow into a mapping where each source executor ID maps to its edge groups. + def _parse_edge_runners(self, edge_runners: list[EdgeRunner]) -> dict[str, list[EdgeRunner]]: + """Parse the edge runners of the workflow into a mapping where each source executor ID maps to its edge runners. Args: - edge_groups: A list of edge groups in the workflow. + edge_runners: A list of edge runners in the workflow. Returns: - A dictionary mapping each source executor ID to a list of edge groups. + A dictionary mapping each source executor ID to a list of edge runners. """ - parsed: defaultdict[str, list[EdgeGroup]] = defaultdict(list) - for group in edge_groups: - for source_executor in group.source_executors: - parsed[source_executor.id].append(group) + parsed: defaultdict[str, list[EdgeRunner]] = defaultdict(list) + for runner in edge_runners: + for source_executor_id in runner._edge_group.source_executor_ids: + parsed[source_executor_id].append(runner) return parsed diff --git a/python/packages/workflow/agent_framework_workflow/_validation.py b/python/packages/workflow/agent_framework_workflow/_validation.py index 46bd60d54e..2f3fe06077 100644 --- a/python/packages/workflow/agent_framework_workflow/_validation.py +++ b/python/packages/workflow/agent_framework_workflow/_validation.py @@ -112,17 +112,20 @@ class WorkflowGraphValidator: self._executors: dict[str, Executor] = {} # region Core Validation Methods - def validate_workflow(self, edge_groups: Sequence[EdgeGroup], start_executor: Executor | str) -> None: + def validate_workflow( + self, edge_groups: Sequence[EdgeGroup], executors: dict[str, Executor], start_executor: Executor | str + ) -> None: """Validate the entire workflow graph. Args: edge_groups: list of edge groups in the workflow + executors: Map of executor IDs to executor instances start_executor: The starting executor (can be instance or ID) Raises: WorkflowValidationError: If any validation fails """ - self._executors = self._build_executor_map(edge_groups) + self._executors = executors self._edges = [edge for group in edge_groups for edge in group.edges] self._edge_groups = edge_groups @@ -143,15 +146,6 @@ class WorkflowGraphValidator: self._validate_dead_ends() self._validate_cycles() - def _build_executor_map(self, edge_groups: Sequence[EdgeGroup]) -> dict[str, Executor]: - """Build a map of executor IDs to executor instances.""" - executors: dict[str, Executor] = {} - for group in edge_groups: - for executor in group.source_executors + group.target_executors: - executors[executor.id] = executor - - return executors - def _validate_handler_output_annotations(self) -> None: """Validate that each handler's ctx parameter is annotated with WorkflowContext[T]. @@ -165,14 +159,16 @@ class WorkflowGraphValidator: # Iterate over all registered executors in the workflow graph for executor_id, executor in self._executors.items(): - for attr_name in dir(executor): + for attr_name in dir(executor.__class__): + if attr_name.startswith("_"): + continue # Retrieve attributes without binding (so the first parameter remains 'self'). # This ensures inspect.signature sees all three parameters: (self, message, ctx). attr = None from contextlib import suppress with suppress(Exception): - attr = inspect.getattr_static(executor, attr_name) + attr = inspect.getattr_static(executor.__class__, attr_name) if attr is None: continue # Consider only callables that were decorated with @handler @@ -298,8 +294,8 @@ class WorkflowGraphValidator: Raises: TypeCompatibilityError: If type incompatibility is detected """ - source_executor = edge.source - target_executor = edge.target + source_executor = self._executors[edge.source_id] + target_executor = self._executors[edge.target_id] # Get output types from source executor source_output_types = self._get_executor_output_types(source_executor) @@ -367,12 +363,18 @@ class WorkflowGraphValidator: """ output_types: list[type[Any]] = [] - for attr_name in dir(executor): - attr = getattr(executor, attr_name) - if callable(attr) and hasattr(attr, "_handler_spec"): - handler_spec = attr._handler_spec # type: ignore - handler_output_types = handler_spec.get("output_types", []) - output_types.extend(handler_output_types) + for attr_name in dir(executor.__class__): + if attr_name.startswith("_"): + continue + try: + attr = getattr(executor.__class__, attr_name) + if callable(attr) and hasattr(attr, "_handler_spec"): + handler_spec = attr._handler_spec # type: ignore + handler_output_types = handler_spec.get("output_types", []) + output_types.extend(handler_output_types) + except AttributeError: + # Skip attributes that may not be accessible + continue return output_types @@ -621,15 +623,18 @@ class WorkflowGraphValidator: # endregion -def validate_workflow_graph(edge_groups: Sequence[EdgeGroup], start_executor: Executor | str) -> None: +def validate_workflow_graph( + edge_groups: Sequence[EdgeGroup], executors: dict[str, Executor], start_executor: Executor | str +) -> None: """Convenience function to validate a workflow graph. Args: edge_groups: list of edge groups in the workflow + executors: Map of executor IDs to executor instances start_executor: The starting executor (can be instance or ID) Raises: WorkflowValidationError: If any validation fails """ validator = WorkflowGraphValidator() - validator.validate_workflow(edge_groups, start_executor) + validator.validate_workflow(edge_groups, executors, start_executor) diff --git a/python/packages/workflow/agent_framework_workflow/_viz.py b/python/packages/workflow/agent_framework_workflow/_viz.py index 77260ae10f..3d7350e21f 100644 --- a/python/packages/workflow/agent_framework_workflow/_viz.py +++ b/python/packages/workflow/agent_framework_workflow/_viz.py @@ -36,13 +36,13 @@ class WorkflowViz: lines.append("") # Add start executor with special styling - start_executor = self._workflow.start_executor - lines.append(f' "{start_executor.id}" [fillcolor=lightgreen, label="{start_executor.id}\\n(Start)"];') + start_executor_id = self._workflow.start_executor_id + lines.append(f' "{start_executor_id}" [fillcolor=lightgreen, label="{start_executor_id}\\n(Start)"];') # Add all other executors - for executor in self._workflow.executors: - if executor.id != start_executor.id: - lines.append(f' "{executor.id}" [label="{executor.id}"];') + for executor_id in self._workflow.executors: + if executor_id != start_executor_id: + lines.append(f' "{executor_id}" [label="{executor_id}"];') # Build shared structures fan_in_nodes = self._compute_fan_in_descriptors() # (node_id, sources, target) @@ -179,16 +179,16 @@ class WorkflowViz: lines: list[str] = ["flowchart TD"] # Nodes - start_executor = self._workflow.start_executor - start_id = _san(start_executor.id) + start_executor_id = self._workflow.start_executor_id + start_id = _san(start_executor_id) # End statements with semicolons for better compatibility and quote labels for special chars - lines.append(f' {start_id}["{start_executor.id} (Start)"];') + lines.append(f' {start_id}["{start_executor_id} (Start)"];') - for executor in self._workflow.executors: - if executor.id == start_executor.id: + for executor_id in self._workflow.executors: + if executor_id == start_executor_id: continue - eid = _san(executor.id) - lines.append(f' {eid}["{executor.id}"];') + eid = _san(executor_id) + lines.append(f' {eid}["{executor_id}"];') # Build shared structures fan_in_nodes_dot = self._compute_fan_in_descriptors() # uses DOT node ids @@ -235,8 +235,8 @@ class WorkflowViz: result: list[tuple[str, list[str], str]] = [] for group in self._workflow.edge_groups: if isinstance(group, FanInEdgeGroup): - target = group.target_executors[0].id - sources = [src.id for src in group.source_executors] + target = group.target_executor_ids[0] + sources = list(group.source_executor_ids) digest = self._fan_in_digest(target, sources) node_id = f"fan_in::{target}::{digest}" result.append((node_id, sorted(sources), target)) diff --git a/python/packages/workflow/agent_framework_workflow/_workflow.py b/python/packages/workflow/agent_framework_workflow/_workflow.py index f2b80ce953..10d7bb7146 100644 --- a/python/packages/workflow/agent_framework_workflow/_workflow.py +++ b/python/packages/workflow/agent_framework_workflow/_workflow.py @@ -7,6 +7,9 @@ import uuid from collections.abc import AsyncIterable, Callable, Sequence from typing import Any +from agent_framework._pydantic import AFBaseModel +from pydantic import Field + from ._checkpoint import CheckpointStorage from ._const import DEFAULT_MAX_ITERATIONS from ._edge import ( @@ -17,6 +20,8 @@ from ._edge import ( FanOutEdgeGroup, SingleEdgeGroup, SwitchCaseEdgeGroup, + SwitchCaseEdgeGroupCase, + SwitchCaseEdgeGroupDefault, ) from ._events import RequestInfoEvent, WorkflowCompletedEvent, WorkflowEvent from ._executor import Executor, RequestInfoExecutor @@ -66,62 +71,84 @@ class WorkflowRunResult(list[WorkflowEvent]): # region Workflow -class Workflow: +class Workflow(AFBaseModel): """A class representing a workflow that can be executed. This class is a placeholder for the workflow logic and does not implement any specific functionality. It serves as a base class for more complex workflows that can be defined in subclasses. """ + edge_groups: list[EdgeGroup] = Field( + default_factory=list, description="List of edge groups that define the workflow edges" + ) + executors: dict[str, Executor] = Field( + default_factory=dict, description="Dictionary mapping executor IDs to Executor instances" + ) + start_executor_id: str = Field(min_length=1, description="The ID of the starting executor for the workflow") + max_iterations: int = Field( + default=DEFAULT_MAX_ITERATIONS, description="Maximum number of iterations the workflow will run" + ) + workflow_id: str = Field( + default_factory=lambda: str(uuid.uuid4()), description="Unique identifier for this workflow instance" + ) + def __init__( self, edge_groups: list[EdgeGroup], + executors: dict[str, Executor], start_executor: Executor | str, runner_context: RunnerContext, - max_iterations: int, + max_iterations: int = DEFAULT_MAX_ITERATIONS, + **kwargs: Any, ): """Initialize the workflow with a list of edges. Args: edge_groups: A list of EdgeGroup instances that define the workflow edges. + executors: A dictionary mapping executor IDs to Executor instances. start_executor: The starting executor for the workflow, which can be an Executor instance or its ID. runner_context: The RunnerContext instance to be used during workflow execution. max_iterations: The maximum number of iterations the workflow will run for convergence. + kwargs: Additional keyword arguments. Unused in this implementation. """ - self._edge_groups = edge_groups - self._executors = self._build_executor_map(edge_groups) - self._start_executor = start_executor + # Convert start_executor to string ID if it's an Executor instance + start_executor_id = start_executor.id if isinstance(start_executor, Executor) else start_executor - self._shared_state = SharedState() workflow_id = str(uuid.uuid4()) + + kwargs.update({ + "edge_groups": edge_groups, + "executors": executors, + "start_executor_id": start_executor_id, + "max_iterations": max_iterations, + "workflow_id": workflow_id, + }) + + super().__init__(**kwargs) + + # Store non-serializable runtime objects as private attributes + self._runner_context = runner_context + self._shared_state = SharedState() self._runner = Runner( - self._edge_groups, + self.edge_groups, + self.executors, self._shared_state, runner_context, max_iterations=max_iterations, workflow_id=workflow_id, ) - @property - def edge_groups(self) -> list[EdgeGroup]: - """Get the list of edge groups in the workflow.""" - return self._edge_groups - - @property - def start_executor(self) -> Executor: + def get_start_executor(self) -> Executor: """Get the starting executor of the workflow. Returns: - The starting executor, which can be an Executor instance or its ID. + The starting executor instance. """ - if isinstance(self._start_executor, str): - return self._get_executor_by_id(self._start_executor) - return self._start_executor + return self.executors[self.start_executor_id] - @property - def executors(self) -> list[Executor]: + def get_executors_list(self) -> list[Executor]: """Get the list of executors in the workflow.""" - return list(self._executors.values()) + return list(self.executors.values()) async def run_streaming(self, message: Any) -> AsyncIterable[WorkflowEvent]: """Run the workflow with a starting message and stream events. @@ -135,9 +162,7 @@ class Workflow: # Reset context for a new run if supported self._runner.context.reset_for_new_run(self._shared_state) - executor = self._start_executor - if isinstance(executor, str): - executor = self._get_executor_by_id(executor) + executor = self.get_start_executor() await executor.execute( message, @@ -303,25 +328,9 @@ class Workflow: Returns: The Executor instance corresponding to the given ID. """ - if executor_id not in self._executors: + if executor_id not in self.executors: raise ValueError(f"Executor with ID {executor_id} not found.") - return self._executors[executor_id] - - def _build_executor_map(self, edge_groups: list[EdgeGroup]) -> dict[str, Executor]: - """Build the executor map from edge groups. - - Args: - edge_groups: A list of EdgeGroup instances. - - Returns: - A dictionary mapping executor IDs to Executor instances. - """ - executors: dict[str, Executor] = {} - for group in edge_groups: - for executor in group.source_executors + group.target_executors: - executors[executor.id] = executor - - return executors + return self.executors[executor_id] async def _restore_from_external_checkpoint( self, checkpoint_id: str, checkpoint_storage: CheckpointStorage @@ -431,10 +440,16 @@ class WorkflowBuilder: def __init__(self, max_iterations: int = DEFAULT_MAX_ITERATIONS): """Initialize the WorkflowBuilder with an empty list of edges and no starting executor.""" self._edge_groups: list[EdgeGroup] = [] + self._executors: dict[str, Executor] = {} self._start_executor: Executor | str | None = None self._checkpoint_storage: CheckpointStorage | None = None self._max_iterations: int = max_iterations + def _add_executor(self, executor: Executor) -> str: + """Add an executor to the map and return its ID.""" + self._executors[executor.id] = executor + return executor.id + def add_edge( self, source: Executor, @@ -452,7 +467,9 @@ class WorkflowBuilder: should be traversed based on the message type. """ # TODO(@taochen): Support executor factories for lazy initialization - self._edge_groups.append(SingleEdgeGroup(source, target, condition)) + source_id = self._add_executor(source) + target_id = self._add_executor(target) + self._edge_groups.append(SingleEdgeGroup(source_id, target_id, condition)) return self def add_fan_out_edges(self, source: Executor, targets: Sequence[Executor]) -> "Self": @@ -464,7 +481,9 @@ class WorkflowBuilder: source: The source executor of the edges. targets: A list of target executors for the edges. """ - self._edge_groups.append(FanOutEdgeGroup(source, targets)) + source_id = self._add_executor(source) + target_ids = [self._add_executor(target) for target in targets] + self._edge_groups.append(FanOutEdgeGroup(source_id, target_ids)) return self @@ -486,7 +505,16 @@ class WorkflowBuilder: source: The source executor of the edges. cases: A list of case objects that determine the target executor for each message. """ - self._edge_groups.append(SwitchCaseEdgeGroup(source, cases)) + source_id = self._add_executor(source) + # Convert case data types to internal types that only uses target_id. + internal_cases: list[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault] = [] + for case in cases: + self._add_executor(case.target) + if isinstance(case, Default): + internal_cases.append(SwitchCaseEdgeGroupDefault(target_id=case.target.id)) + else: + internal_cases.append(SwitchCaseEdgeGroupCase(condition=case.condition, target_id=case.target.id)) + self._edge_groups.append(SwitchCaseEdgeGroup(source_id, internal_cases)) return self @@ -510,7 +538,9 @@ class WorkflowBuilder: targets: A list of target executors for the edges. selection_func: A function that selects target executors for messages. """ - self._edge_groups.append(FanOutEdgeGroup(source, targets, selection_func)) + source_id = self._add_executor(source) + target_ids = [self._add_executor(target) for target in targets] + self._edge_groups.append(FanOutEdgeGroup(source_id, target_ids, selection_func)) return self @@ -548,7 +578,9 @@ class WorkflowBuilder: sources: A list of source executors for the edges. target: The target executor for the edges. """ - self._edge_groups.append(FanInEdgeGroup(sources, target)) + source_ids = [self._add_executor(source) for source in sources] + target_id = self._add_executor(target) + self._edge_groups.append(FanInEdgeGroup(source_ids, target_id)) return self @@ -608,13 +640,13 @@ class WorkflowBuilder: TypeCompatibilityError, and GraphConnectivityError subclasses). """ if not self._start_executor: - raise ValueError("Starting executor must be set before building the workflow.") + raise ValueError("Starting executor must be set using set_start_executor before building the workflow.") - validate_workflow_graph(self._edge_groups, self._start_executor) + validate_workflow_graph(self._edge_groups, self._executors, self._start_executor) context = InProcRunnerContext(self._checkpoint_storage) - return Workflow(self._edge_groups, self._start_executor, context, self._max_iterations) + return Workflow(self._edge_groups, self._executors, self._start_executor, context, self._max_iterations) # endregion diff --git a/python/packages/workflow/tests/test_edge.py b/python/packages/workflow/tests/test_edge.py index ddef2cd649..3e452e95ab 100644 --- a/python/packages/workflow/tests/test_edge.py +++ b/python/packages/workflow/tests/test_edge.py @@ -8,14 +8,17 @@ import pytest from agent_framework.workflow import Executor, WorkflowContext, handler from agent_framework_workflow._edge import ( - Case, - Default, Edge, FanInEdgeGroup, FanOutEdgeGroup, SingleEdgeGroup, SwitchCaseEdgeGroup, + SwitchCaseEdgeGroupCase, + SwitchCaseEdgeGroupDefault, ) +from agent_framework_workflow._edge_runner import create_edge_runner +from agent_framework_workflow._runner_context import InProcRunnerContext, Message +from agent_framework_workflow._shared_state import SharedState @dataclass @@ -35,28 +38,40 @@ class MockMessageSecondary: class MockExecutor(Executor): """A mock executor for testing purposes.""" + call_count: int = 0 + last_message: Any = None + @handler async def mock_handler(self, message: MockMessage, ctx: WorkflowContext) -> None: """A mock handler that does nothing.""" - pass + self.call_count += 1 + self.last_message = message class MockExecutorSecondary(Executor): """A secondary mock executor for testing purposes.""" + call_count: int = 0 + last_message: Any = None + @handler async def mock_handler_secondary(self, message: MockMessageSecondary, ctx: WorkflowContext) -> None: """A secondary mock handler that does nothing.""" - pass + self.call_count += 1 + self.last_message = message class MockAggregator(Executor): """A mock aggregator for testing purposes.""" + call_count: int = 0 + last_message: Any = None + @handler async def mock_aggregator_handler(self, message: list[MockMessage], ctx: WorkflowContext) -> None: """A mock aggregator handler that does nothing.""" - pass + self.call_count += 1 + self.last_message = message # region Edge @@ -67,7 +82,7 @@ def test_create_edge(): source = MockExecutor(id="source_executor") target = MockExecutor(id="target_executor") - edge = Edge(source=source, target=target) + edge = Edge(source_id=source.id, target_id=target.id) assert edge.source_id == "source_executor" assert edge.target_id == "target_executor" @@ -79,9 +94,9 @@ def test_edge_can_handle(): source = MockExecutor(id="source_executor") target = MockExecutor(id="target_executor") - edge = Edge(source=source, target=target) + edge = Edge(source_id=source.id, target_id=target.id) - assert edge.can_handle(MockMessage(data="test")) + assert edge.should_route(MockMessage(data="test")) # endregion Edge @@ -94,10 +109,10 @@ def test_single_edge_group(): source = MockExecutor(id="source_executor") target = MockExecutor(id="target_executor") - edge_group = SingleEdgeGroup(source=source, target=target) + edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id) - assert edge_group.source_executors == [source] - assert edge_group.target_executors == [target] + assert edge_group.source_executor_ids == [source.id] + assert edge_group.target_executor_ids == [target.id] assert edge_group.edges[0].source_id == "source_executor" assert edge_group.edges[0].target_id == "target_executor" @@ -107,92 +122,88 @@ def test_single_edge_group_with_condition(): source = MockExecutor(id="source_executor") target = MockExecutor(id="target_executor") - edge_group = SingleEdgeGroup(source=source, target=target, condition=lambda x: x.data == "test") + edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id, condition=lambda x: x.data == "test") - assert edge_group.source_executors == [source] - assert edge_group.target_executors == [target] + assert edge_group.source_executor_ids == [source.id] + assert edge_group.target_executor_ids == [target.id] assert edge_group.edges[0].source_id == "source_executor" assert edge_group.edges[0].target_id == "target_executor" assert edge_group.edges[0]._condition is not None # type: ignore -async def test_single_edge_group_send_message(): - """Test sending a message through a single edge group.""" +async def test_single_edge_group_send_message() -> None: + """Test sending a message through a single edge runner.""" source = MockExecutor(id="source_executor") target = MockExecutor(id="target_executor") - edge_group = SingleEdgeGroup(source=source, target=target) - - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source.id: source, target.id: target} + edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id) + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() data = MockMessage(data="test") message = Message(data=data, source_id=source.id) - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is True -async def test_single_edge_group_send_message_with_target(): - """Test sending a message through a single edge group.""" +async def test_single_edge_group_send_message_with_target() -> None: + """Test sending a message through a single edge runner.""" source = MockExecutor(id="source_executor") target = MockExecutor(id="target_executor") - edge_group = SingleEdgeGroup(source=source, target=target) - - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source.id: source, target.id: target} + edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id) + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() data = MockMessage(data="test") message = Message(data=data, source_id=source.id, target_id=target.id) - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is True -async def test_single_edge_group_send_message_with_invalid_target(): - """Test sending a message through a single edge group.""" +async def test_single_edge_group_send_message_with_invalid_target() -> None: + """Test sending a message through a single edge runner.""" source = MockExecutor(id="source_executor") target = MockExecutor(id="target_executor") - edge_group = SingleEdgeGroup(source=source, target=target) - - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source.id: source, target.id: target} + edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id) + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() data = MockMessage(data="test") message = Message(data=data, source_id=source.id, target_id="invalid_target") - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is False -async def test_single_edge_group_send_message_with_invalid_data(): - """Test sending a message through a single edge group.""" +async def test_single_edge_group_send_message_with_invalid_data() -> None: + """Test sending a message through a single edge runner with invalid data.""" source = MockExecutor(id="source_executor") target = MockExecutor(id="target_executor") - edge_group = SingleEdgeGroup(source=source, target=target) - - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source.id: source, target.id: target} + edge_group = SingleEdgeGroup(source_id=source.id, target_id=target.id) + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() data = "invalid_data" message = Message(data=data, source_id=source.id) - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is False @@ -208,10 +219,10 @@ def test_source_edge_group(): target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") - edge_group = FanOutEdgeGroup(source=source, targets=[target1, target2]) + edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id]) - assert edge_group.source_executors == [source] - assert edge_group.target_executors == [target1, target2] + assert edge_group.source_executor_ids == [source.id] + assert edge_group.target_executor_ids == [target1.id, target2.id] assert len(edge_group.edges) == 2 assert edge_group.edges[0].source_id == "source_executor" assert edge_group.edges[0].target_id == "target_executor_1" @@ -219,128 +230,122 @@ def test_source_edge_group(): assert edge_group.edges[1].target_id == "target_executor_2" -def test_source_edge_group_invalid_number_of_targets(): +def test_source_edge_group_invalid_number_of_targets() -> None: """Test creating a fan-out group with an invalid number of targets.""" source = MockExecutor(id="source_executor") target = MockExecutor(id="target_executor") with pytest.raises(ValueError, match="FanOutEdgeGroup must contain at least two targets"): - FanOutEdgeGroup(source=source, targets=[target]) + FanOutEdgeGroup(source_id=source.id, target_ids=[target.id]) -async def test_source_edge_group_send_message(): - """Test sending a message through a fan-out group.""" +async def test_source_edge_group_send_message() -> None: + """Test sending a message through a fan-out edge runner.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") - edge_group = FanOutEdgeGroup(source=source, targets=[target1, target2]) - - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id]) + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() data = MockMessage(data="test") message = Message(data=data, source_id=source.id) - with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send: - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) - assert success is True - assert mock_send.call_count == 2 + assert success is True + assert target1.call_count == 1 + assert target2.call_count == 1 -async def test_source_edge_group_send_message_with_target(): +async def test_source_edge_group_send_message_with_target() -> None: """Test sending a message through a fan-out group with a target.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") - edge_group = FanOutEdgeGroup(source=source, targets=[target1, target2]) - - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id]) + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() data = MockMessage(data="test") message = Message(data=data, source_id=source.id, target_id=target1.id) - with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send: - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) - assert success is True - assert mock_send.call_count == 1 - assert mock_send.call_args[0][0].target_id == target1.id + assert success is True + assert target1.call_count == 1 + assert target2.call_count == 0 # target2 should not be called since message targets target1 -async def test_source_edge_group_send_message_with_invalid_target(): +async def test_source_edge_group_send_message_with_invalid_target() -> None: """Test sending a message through a fan-out group with an invalid target.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") - edge_group = FanOutEdgeGroup(source=source, targets=[target1, target2]) - - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id]) + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() data = MockMessage(data="test") message = Message(data=data, source_id=source.id, target_id="invalid_target") - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is False -async def test_source_edge_group_send_message_with_invalid_data(): +async def test_source_edge_group_send_message_with_invalid_data() -> None: """Test sending a message through a fan-out group with invalid data.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") - edge_group = FanOutEdgeGroup(source=source, targets=[target1, target2]) - - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id]) + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() data = "invalid_data" message = Message(data=data, source_id=source.id) - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is False -async def test_source_edge_group_send_message_only_one_successful_send(): +async def test_source_edge_group_send_message_only_one_successful_send() -> None: """Test sending a message through a fan-out group where only one edge can handle the message.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutorSecondary(id="target_executor_2") - edge_group = FanOutEdgeGroup(source=source, targets=[target1, target2]) - - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + edge_group = FanOutEdgeGroup(source_id=source.id, target_ids=[target1.id, target2.id]) + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() data = MockMessage(data="test") message = Message(data=data, source_id=source.id) - with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send: - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) - assert success is True - assert mock_send.call_count == 1 + assert success is True + assert target1.call_count == 1 # target1 can handle MockMessage + assert target2.call_count == 0 # target2 (MockExecutorSecondary) cannot handle MockMessage def test_source_edge_group_with_selection_func(): @@ -350,13 +355,13 @@ def test_source_edge_group_with_selection_func(): target2 = MockExecutor(id="target_executor_2") edge_group = FanOutEdgeGroup( - source=source, - targets=[target1, target2], + source_id=source.id, + target_ids=[target1.id, target2.id], selection_func=lambda data, target_ids: [target1.id], ) - assert edge_group.source_executors == [source] - assert edge_group.target_executors == [target1, target2] + assert edge_group.source_executor_ids == [source.id] + assert edge_group.target_executor_ids == [target1.id, target2.id] assert len(edge_group.edges) == 2 assert edge_group.edges[0].source_id == "source_executor" assert edge_group.edges[0].target_id == "target_executor_1" @@ -364,20 +369,20 @@ def test_source_edge_group_with_selection_func(): assert edge_group.edges[1].target_id == "target_executor_2" -async def test_source_edge_group_with_selection_func_send_message(): +async def test_source_edge_group_with_selection_func_send_message() -> None: """Test sending a message through a fan-out group with a selection function.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") edge_group = FanOutEdgeGroup( - source=source, - targets=[target1, target2], + source_id=source.id, + target_ids=[target1.id, target2.id], selection_func=lambda data, target_ids: [target1.id, target2.id], ) - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() @@ -385,28 +390,28 @@ async def test_source_edge_group_with_selection_func_send_message(): data = MockMessage(data="test") message = Message(data=data, source_id=source.id) - with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send: - success = await edge_group.send_message(message, shared_state, ctx) + with patch("agent_framework_workflow._edge_runner.EdgeRunner._execute_on_target") as mock_send: + success = await edge_runner.send_message(message, shared_state, ctx) assert success is True + assert mock_send.call_count == 2 -async def test_source_edge_group_with_selection_func_send_message_with_invalid_selection_result(): +async def test_source_edge_group_with_selection_func_send_message_with_invalid_selection_result() -> None: """Test sending a message through a fan-out group with a selection func with an invalid selection result.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") edge_group = FanOutEdgeGroup( - source=source, - targets=[target1, target2], + source_id=source.id, + target_ids=[target1.id, target2.id], selection_func=lambda data, target_ids: [target1.id, "invalid_target"], ) - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState - + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() @@ -414,23 +419,23 @@ async def test_source_edge_group_with_selection_func_send_message_with_invalid_s message = Message(data=data, source_id=source.id) with pytest.raises(RuntimeError): - await edge_group.send_message(message, shared_state, ctx) + await edge_runner.send_message(message, shared_state, ctx) -async def test_source_edge_group_with_selection_func_send_message_with_target(): +async def test_source_edge_group_with_selection_func_send_message_with_target() -> None: """Test sending a message through a fan-out group with a selection func with a target.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") edge_group = FanOutEdgeGroup( - source=source, - targets=[target1, target2], + source_id=source.id, + target_ids=[target1.id, target2.id], selection_func=lambda data, target_ids: [target1.id, target2.id], ) - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() @@ -438,28 +443,28 @@ async def test_source_edge_group_with_selection_func_send_message_with_target(): data = MockMessage(data="test") message = Message(data=data, source_id=source.id, target_id=target1.id) - with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send: - success = await edge_group.send_message(message, shared_state, ctx) + with patch("agent_framework_workflow._edge_runner.EdgeRunner._execute_on_target") as mock_send: + success = await edge_runner.send_message(message, shared_state, ctx) assert success is True assert mock_send.call_count == 1 - assert mock_send.call_args[0][0].target_id == target1.id + assert mock_send.call_args[0][0] == target1.id -async def test_source_edge_group_with_selection_func_send_message_with_target_not_in_selection(): +async def test_source_edge_group_with_selection_func_send_message_with_target_not_in_selection() -> None: """Test sending a message through a fan-out group with a selection func with a target not in the selection.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") edge_group = FanOutEdgeGroup( - source=source, - targets=[target1, target2], + source_id=source.id, + target_ids=[target1.id, target2.id], selection_func=lambda data, target_ids: [target1.id], # Only target1 will receive the message ) - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() @@ -467,22 +472,24 @@ async def test_source_edge_group_with_selection_func_send_message_with_target_no data = MockMessage(data="test") message = Message(data=data, source_id=source.id, target_id=target2.id) - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is False -async def test_source_edge_group_with_selection_func_send_message_with_invalid_data(): +async def test_source_edge_group_with_selection_func_send_message_with_invalid_data() -> None: """Test sending a message through a fan-out group with a selection func with invalid data.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") edge_group = FanOutEdgeGroup( - source=source, targets=[target1, target2], selection_func=lambda data, target_ids: [target1.id, target2.id] + source_id=source.id, + target_ids=[target1.id, target2.id], + selection_func=lambda data, target_ids: [target1.id, target2.id], ) - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() @@ -490,22 +497,24 @@ async def test_source_edge_group_with_selection_func_send_message_with_invalid_d data = "invalid_data" message = Message(data=data, source_id=source.id) - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is False -async def test_source_edge_group_with_selection_func_send_message_with_target_invalid_data(): +async def test_source_edge_group_with_selection_func_send_message_with_target_invalid_data() -> None: """Test sending a message through a fan-out group with a selection func with a target and invalid data.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") edge_group = FanOutEdgeGroup( - source=source, targets=[target1, target2], selection_func=lambda data, target_ids: [target1.id, target2.id] + source_id=source.id, + target_ids=[target1.id, target2.id], + selection_func=lambda data, target_ids: [target1.id, target2.id], ) - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() @@ -513,7 +522,7 @@ async def test_source_edge_group_with_selection_func_send_message_with_target_in data = "invalid_data" message = Message(data=data, source_id=source.id, target_id=target1.id) - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is False @@ -528,10 +537,10 @@ def test_target_edge_group(): source2 = MockExecutor(id="source_executor_2") target = MockAggregator(id="target_executor") - edge_group = FanInEdgeGroup(sources=[source1, source2], target=target) + edge_group = FanInEdgeGroup(source_ids=[source1.id, source2.id], target_id=target.id) - assert edge_group.source_executors == [source1, source2] - assert edge_group.target_executors == [target] + assert edge_group.source_executor_ids == [source1.id, source2.id] + assert edge_group.target_executor_ids == [target.id] assert len(edge_group.edges) == 2 assert edge_group.edges[0].source_id == "source_executor_1" assert edge_group.edges[0].target_id == "target_executor" @@ -545,27 +554,27 @@ def test_target_edge_group_invalid_number_of_sources(): target = MockAggregator(id="target_executor") with pytest.raises(ValueError, match="FanInEdgeGroup must contain at least two sources"): - FanInEdgeGroup(sources=[source], target=target) + FanInEdgeGroup(source_ids=[source.id], target_id=target.id) -async def test_target_edge_group_send_message_buffer(): +async def test_target_edge_group_send_message_buffer() -> None: """Test sending a message through a fan-in edge group with buffering.""" source1 = MockExecutor(id="source_executor_1") source2 = MockExecutor(id="source_executor_2") target = MockAggregator(id="target_executor") - edge_group = FanInEdgeGroup(sources=[source1, source2], target=target) + edge_group = FanInEdgeGroup(source_ids=[source1.id, source2.id], target_id=target.id) - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source1.id: source1, source2.id: source2, target.id: target} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() data = MockMessage(data="test") - with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send: - success = await edge_group.send_message( + with patch("agent_framework_workflow._edge_runner.EdgeRunner._execute_on_target") as mock_send: + success = await edge_runner.send_message( Message(data=data, source_id=source1.id), shared_state, ctx, @@ -573,9 +582,9 @@ async def test_target_edge_group_send_message_buffer(): assert success is True assert mock_send.call_count == 0 # The message should be buffered and wait for the second source - assert len(edge_group._buffer[source1.id]) == 1 # type: ignore + assert len(edge_runner._buffer[source1.id]) == 1 # type: ignore - success = await edge_group.send_message( + success = await edge_runner.send_message( Message(data=data, source_id=source2.id), shared_state, ctx, @@ -584,19 +593,19 @@ async def test_target_edge_group_send_message_buffer(): assert mock_send.call_count == 1 # The message should be sent now that both sources have sent their messages # Buffer should be cleared after sending - assert not edge_group._buffer # type: ignore + assert not edge_runner._buffer # type: ignore -async def test_target_edge_group_send_message_with_invalid_target(): +async def test_target_edge_group_send_message_with_invalid_target() -> None: """Test sending a message through a fan-in edge group with an invalid target.""" source1 = MockExecutor(id="source_executor_1") source2 = MockExecutor(id="source_executor_2") target = MockAggregator(id="target_executor") - edge_group = FanInEdgeGroup(sources=[source1, source2], target=target) + edge_group = FanInEdgeGroup(source_ids=[source1.id, source2.id], target_id=target.id) - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source1.id: source1, source2.id: source2, target.id: target} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() @@ -604,20 +613,20 @@ async def test_target_edge_group_send_message_with_invalid_target(): data = MockMessage(data="test") message = Message(data=data, source_id=source1.id, target_id="invalid_target") - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is False -async def test_target_edge_group_send_message_with_invalid_data(): +async def test_target_edge_group_send_message_with_invalid_data() -> None: """Test sending a message through a fan-in edge group with invalid data.""" source1 = MockExecutor(id="source_executor_1") source2 = MockExecutor(id="source_executor_2") target = MockAggregator(id="target_executor") - edge_group = FanInEdgeGroup(sources=[source1, source2], target=target) + edge_group = FanInEdgeGroup(source_ids=[source1.id, source2.id], target_id=target.id) - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source1.id: source1, source2.id: source2, target.id: target} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() @@ -625,7 +634,7 @@ async def test_target_edge_group_send_message_with_invalid_data(): data = "invalid_data" message = Message(data=data, source_id=source1.id) - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is False @@ -634,22 +643,22 @@ async def test_target_edge_group_send_message_with_invalid_data(): # region SwitchCaseEdgeGroup -def test_switch_case_edge_group(): +def test_switch_case_edge_group() -> None: """Test creating a switch case edge group.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") edge_group = SwitchCaseEdgeGroup( - source=source, + source_id=source.id, cases=[ - Case(condition=lambda x: x.data < 0, target=target1), - Default(target=target2), + SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target1.id), + SwitchCaseEdgeGroupDefault(target_id=target2.id), ], ) - assert edge_group.source_executors == [source] - assert edge_group.target_executors == [target1, target2] + assert edge_group.source_executor_ids == [source.id] + assert edge_group.target_executor_ids == [target1.id, target2.id] assert len(edge_group.edges) == 2 assert edge_group.edges[0].source_id == "source_executor" assert edge_group.edges[0].target_id == "target_executor_1" @@ -670,18 +679,18 @@ def test_switch_case_edge_group_invalid_number_of_cases(): ValueError, match=r"SwitchCaseEdgeGroup must contain at least two cases \(including the default case\)." ): SwitchCaseEdgeGroup( - source=source, + source_id=source.id, cases=[ - Case(condition=lambda x: x.data < 0, target=target), + SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target.id), ], ) with pytest.raises(ValueError, match="SwitchCaseEdgeGroup must contain exactly one default case."): SwitchCaseEdgeGroup( - source=source, + source_id=source.id, cases=[ - Case(condition=lambda x: x.data < 0, target=target), - Case(condition=lambda x: x.data >= 0, target=target), + SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target.id), + SwitchCaseEdgeGroupCase(condition=lambda x: x.data >= 0, target_id=target.id), ], ) @@ -694,31 +703,30 @@ def test_switch_case_edge_group_invalid_number_of_default_cases(): with pytest.raises(ValueError, match="SwitchCaseEdgeGroup must contain exactly one default case."): SwitchCaseEdgeGroup( - source=source, + source_id=source.id, cases=[ - Case(condition=lambda x: x.data < 0, target=target1), - Default(target=target2), - Default(target=target2), + SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target1.id), + SwitchCaseEdgeGroupDefault(target_id=target2.id), + SwitchCaseEdgeGroupDefault(target_id=target2.id), ], ) -async def test_switch_case_edge_group_send_message(): +async def test_switch_case_edge_group_send_message() -> None: """Test sending a message through a switch case edge group.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") edge_group = SwitchCaseEdgeGroup( - source=source, + source_id=source.id, cases=[ - Case(condition=lambda x: x.data < 0, target=target1), - Default(target=target2), + SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target1.id), + SwitchCaseEdgeGroupDefault(target_id=target2.id), ], ) - - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() @@ -726,8 +734,8 @@ async def test_switch_case_edge_group_send_message(): data = MockMessage(data=-1) message = Message(data=data, source_id=source.id) - with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send: - success = await edge_group.send_message(message, shared_state, ctx) + with patch("agent_framework_workflow._edge_runner.EdgeRunner._execute_on_target") as mock_send: + success = await edge_runner.send_message(message, shared_state, ctx) assert success is True assert mock_send.call_count == 1 @@ -735,29 +743,29 @@ async def test_switch_case_edge_group_send_message(): # Default condition should data = MockMessage(data=1) message = Message(data=data, source_id=source.id) - with patch("agent_framework_workflow._edge.Edge.send_message") as mock_send: - success = await edge_group.send_message(message, shared_state, ctx) + with patch("agent_framework_workflow._edge_runner.EdgeRunner._execute_on_target") as mock_send: + success = await edge_runner.send_message(message, shared_state, ctx) assert success is True assert mock_send.call_count == 1 -async def test_switch_case_edge_group_send_message_with_invalid_target(): +async def test_switch_case_edge_group_send_message_with_invalid_target() -> None: """Test sending a message through a switch case edge group with an invalid target.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") edge_group = SwitchCaseEdgeGroup( - source=source, + source_id=source.id, cases=[ - Case(condition=lambda x: x.data < 0, target=target1), - Default(target=target2), + SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target1.id), + SwitchCaseEdgeGroupDefault(target_id=target2.id), ], ) - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() @@ -765,26 +773,26 @@ async def test_switch_case_edge_group_send_message_with_invalid_target(): data = MockMessage(data=-1) message = Message(data=data, source_id=source.id, target_id="invalid_target") - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is False -async def test_switch_case_edge_group_send_message_with_valid_target(): +async def test_switch_case_edge_group_send_message_with_valid_target() -> None: """Test sending a message through a switch case edge group with a target.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") edge_group = SwitchCaseEdgeGroup( - source=source, + source_id=source.id, cases=[ - Case(condition=lambda x: x.data < 0, target=target1), - Default(target=target2), + SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target1.id), + SwitchCaseEdgeGroupDefault(target_id=target2.id), ], ) - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() @@ -792,31 +800,31 @@ async def test_switch_case_edge_group_send_message_with_valid_target(): data = MockMessage(data=1) # Condition will fail message = Message(data=data, source_id=source.id, target_id=target1.id) - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is False data = MockMessage(data=-1) # Condition will pass message = Message(data=data, source_id=source.id, target_id=target1.id) - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is True -async def test_switch_case_edge_group_send_message_with_invalid_data(): +async def test_switch_case_edge_group_send_message_with_invalid_data() -> None: """Test sending a message through a switch case edge group with invalid data.""" source = MockExecutor(id="source_executor") target1 = MockExecutor(id="target_executor_1") target2 = MockExecutor(id="target_executor_2") edge_group = SwitchCaseEdgeGroup( - source=source, + source_id=source.id, cases=[ - Case(condition=lambda x: x.data < 0, target=target1), - Default(target=target2), + SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target1.id), + SwitchCaseEdgeGroupDefault(target_id=target2.id), ], ) - from agent_framework_workflow._runner_context import InProcRunnerContext, Message - from agent_framework_workflow._shared_state import SharedState + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) shared_state = SharedState() ctx = InProcRunnerContext() @@ -824,7 +832,7 @@ async def test_switch_case_edge_group_send_message_with_invalid_data(): data = "invalid_data" message = Message(data=data, source_id=source.id) - success = await edge_group.send_message(message, shared_state, ctx) + success = await edge_runner.send_message(message, shared_state, ctx) assert success is False diff --git a/python/packages/workflow/tests/test_runner.py b/python/packages/workflow/tests/test_runner.py index 12745c8f16..0666f1b66c 100644 --- a/python/packages/workflow/tests/test_runner.py +++ b/python/packages/workflow/tests/test_runner.py @@ -37,11 +37,13 @@ def test_create_runner(): # Create a loop edge_groups = [ - SingleEdgeGroup(executor_a, executor_b), - SingleEdgeGroup(executor_b, executor_a), + SingleEdgeGroup(executor_a.id, executor_b.id), + SingleEdgeGroup(executor_b.id, executor_a.id), ] - runner = Runner(edge_groups, shared_state=SharedState(), ctx=InProcRunnerContext()) + executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b} + + runner = Runner(edge_groups, executors, shared_state=SharedState(), ctx=InProcRunnerContext()) assert runner.context is not None and isinstance(runner.context, RunnerContext) @@ -53,14 +55,15 @@ async def test_runner_run_until_convergence(): # Create a loop edges = [ - SingleEdgeGroup(executor_a, executor_b), - SingleEdgeGroup(executor_b, executor_a), + SingleEdgeGroup(executor_a.id, executor_b.id), + SingleEdgeGroup(executor_b.id, executor_a.id), ] + executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b} shared_state = SharedState() ctx = InProcRunnerContext() - runner = Runner(edges, shared_state, ctx) + runner = Runner(edges, executors, shared_state, ctx) result: int | None = None await executor_a.execute( @@ -87,14 +90,15 @@ async def test_runner_run_until_convergence_not_completed(): # Create a loop edges = [ - SingleEdgeGroup(executor_a, executor_b), - SingleEdgeGroup(executor_b, executor_a), + SingleEdgeGroup(executor_a.id, executor_b.id), + SingleEdgeGroup(executor_b.id, executor_a.id), ] + executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b} shared_state = SharedState() ctx = InProcRunnerContext() - runner = Runner(edges, shared_state, ctx, max_iterations=5) + runner = Runner(edges, executors, shared_state, ctx, max_iterations=5) await executor_a.execute( MockMessage(data=0), @@ -117,14 +121,15 @@ async def test_runner_already_running(): # Create a loop edges = [ - SingleEdgeGroup(executor_a, executor_b), - SingleEdgeGroup(executor_b, executor_a), + SingleEdgeGroup(executor_a.id, executor_b.id), + SingleEdgeGroup(executor_b.id, executor_a.id), ] + executors: dict[str, Executor] = {executor_a.id: executor_a, executor_b.id: executor_b} shared_state = SharedState() ctx = InProcRunnerContext() - runner = Runner(edges, shared_state, ctx) + runner = Runner(edges, executors, shared_state, ctx) await executor_a.execute( MockMessage(data=0), diff --git a/python/packages/workflow/tests/test_serialization.py b/python/packages/workflow/tests/test_serialization.py new file mode 100644 index 0000000000..02c405662d --- /dev/null +++ b/python/packages/workflow/tests/test_serialization.py @@ -0,0 +1,636 @@ +# Copyright (c) Microsoft. All rights reserved. + +import json +from typing import Any + +import pytest +from agent_framework.workflow import Executor, WorkflowBuilder, WorkflowContext, handler + +from agent_framework_workflow._edge import ( + Edge, + FanInEdgeGroup, + FanOutEdgeGroup, + SingleEdgeGroup, + SwitchCaseEdgeGroup, + SwitchCaseEdgeGroupCase, + SwitchCaseEdgeGroupDefault, +) + + +class SampleExecutor(Executor): + """Sample executor for serialization testing.""" + + @handler + async def handle_str(self, message: str, ctx: WorkflowContext[str]) -> None: + """Handle string messages.""" + await ctx.send_message(f"Processed: {message}") + + +class SampleAggregator(Executor): + """Sample aggregator executor that can handle lists of messages.""" + + @handler + async def handle_str_list(self, messages: list[str], ctx: WorkflowContext[str]) -> None: + """Handle list of string messages for fan-in aggregation.""" + combined = " | ".join(messages) + await ctx.send_message(f"Aggregated: {combined}") + + +class TestSerializationWorkflowClasses: + """Test serialization of workflow classes.""" + + def test_executor_serialization(self) -> None: + """Test that Executor can be serialized and has correct fields, including type.""" + executor = SampleExecutor(id="test-executor") + + # Test model_dump + data = executor.model_dump() + assert data["id"] == "test-executor" + + # Test type field + assert "type" in data, "Executor should have 'type' field" + assert data["type"] == "SampleExecutor", f"Expected type 'SampleExecutor', got {data['type']}" + + # Test model_dump_json + json_str = executor.model_dump_json() + parsed = json.loads(json_str) + assert parsed["id"] == "test-executor" + + # Test type field in JSON + assert "type" in parsed, "JSON should have 'type' field" + assert parsed["type"] == "SampleExecutor", "JSON should preserve type field" + + def test_edge_serialization(self) -> None: + """Test that Edge can be serialized and has correct fields.""" + # Test edge without condition + edge = Edge(source_id="source", target_id="target") + + # Test model_dump + data = edge.model_dump() + assert data["source_id"] == "source" + assert data["target_id"] == "target" + assert "condition_name" not in data or data["condition_name"] is None + + # Test model_dump_json + json_str = edge.model_dump_json() + parsed = json.loads(json_str) + assert parsed["source_id"] == "source" + assert parsed["target_id"] == "target" + assert "condition_name" not in parsed or parsed["condition_name"] is None + + def test_edge_serialization_with_named_condition(self) -> None: + """Test that Edge with named function condition serializes condition_name correctly.""" + + def is_positive(x: int) -> bool: + return x > 0 + + edge = Edge(source_id="source", target_id="target", condition=is_positive) + + # Test model_dump + data = edge.model_dump() + assert data["source_id"] == "source" + assert data["target_id"] == "target" + assert data["condition_name"] == "is_positive" + + # Test model_dump_json + json_str = edge.model_dump_json() + parsed = json.loads(json_str) + assert parsed["source_id"] == "source" + assert parsed["target_id"] == "target" + assert parsed["condition_name"] == "is_positive" + + def test_edge_serialization_with_lambda_condition(self) -> None: + """Test that Edge with lambda condition serializes condition_name as ''.""" + edge = Edge(source_id="source", target_id="target", condition=lambda x: x > 0) + + # Test model_dump + data = edge.model_dump() + assert data["source_id"] == "source" + assert data["target_id"] == "target" + assert data["condition_name"] == "" + + # Test model_dump_json + json_str = edge.model_dump_json() + parsed = json.loads(json_str) + assert parsed["source_id"] == "source" + assert parsed["target_id"] == "target" + assert parsed["condition_name"] == "" + + def test_single_edge_group_serialization(self) -> None: + """Test that SingleEdgeGroup can be serialized and has correct fields, including edges and type.""" + edge_group = SingleEdgeGroup(source_id="source", target_id="target") + + # Test model_dump + data = edge_group.model_dump() + assert "id" in data + assert data["id"].startswith("SingleEdgeGroup/") + + # Test type field + assert "type" in data, "SingleEdgeGroup should have 'type' field" + assert data["type"] == "SingleEdgeGroup", f"Expected type 'SingleEdgeGroup', got {data['type']}" + + # Verify edges field is present and contains the edge + assert "edges" in data, "SingleEdgeGroup should have 'edges' field" + assert len(data["edges"]) == 1, "SingleEdgeGroup should have exactly one edge" + edge = data["edges"][0] + assert "source_id" in edge, "Edge should have source_id" + assert "target_id" in edge, "Edge should have target_id" + assert edge["source_id"] == "source", f"Expected source_id 'source', got {edge['source_id']}" + assert edge["target_id"] == "target", f"Expected target_id 'target', got {edge['target_id']}" + + # Test model_dump_json + json_str = edge_group.model_dump_json() + parsed = json.loads(json_str) + assert "id" in parsed + assert parsed["id"].startswith("SingleEdgeGroup/") + + # Test type field in JSON + assert "type" in parsed, "JSON should have 'type' field" + assert parsed["type"] == "SingleEdgeGroup", "JSON should preserve type field" + + # Verify edges are preserved in JSON + assert "edges" in parsed, "JSON should have 'edges' field" + assert len(parsed["edges"]) == 1, "JSON should have exactly one edge" + json_edge = parsed["edges"][0] + assert json_edge["source_id"] == "source", "JSON should preserve edge source_id" + assert json_edge["target_id"] == "target", "JSON should preserve edge target_id" + + def test_fan_out_edge_group_serialization(self) -> None: + """Test that FanOutEdgeGroup can be serialized and has correct fields, including edges and type.""" + edge_group = FanOutEdgeGroup(source_id="source", target_ids=["target1", "target2"]) + + # Test model_dump + data = edge_group.model_dump() + assert "id" in data + assert data["id"].startswith("FanOutEdgeGroup/") + + # Test type field + assert "type" in data, "FanOutEdgeGroup should have 'type' field" + assert data["type"] == "FanOutEdgeGroup", f"Expected type 'FanOutEdgeGroup', got {data['type']}" + + # Test selection_func_name field (should be None when no selection function is provided) + assert "selection_func_name" in data, "FanOutEdgeGroup should have 'selection_func_name' field" + assert data["selection_func_name"] is None, ( + "selection_func_name should be None when no selection function is provided" + ) + + # Verify edges field is present and contains the correct edges + assert "edges" in data, "FanOutEdgeGroup should have 'edges' field" + assert len(data["edges"]) == 2, "FanOutEdgeGroup should have exactly two edges" + + edges = data["edges"] + sources = [edge["source_id"] for edge in edges] + targets = [edge["target_id"] for edge in edges] + + assert all(source == "source" for source in sources), f"All edges should have source 'source', got {sources}" + assert set(targets) == {"target1", "target2"}, f"Expected targets {{'target1', 'target2'}}, got {set(targets)}" + + # Test model_dump_json + json_str = edge_group.model_dump_json() + parsed = json.loads(json_str) + assert "id" in parsed + assert parsed["id"].startswith("FanOutEdgeGroup/") + + # Test type field in JSON + assert "type" in parsed, "JSON should have 'type' field" + assert parsed["type"] == "FanOutEdgeGroup", "JSON should preserve type field" + + # Test selection_func_name field in JSON + assert "selection_func_name" in parsed, "JSON should have 'selection_func_name' field" + assert parsed["selection_func_name"] is None, ( + "JSON selection_func_name should be None when no selection function is provided" + ) + + # Verify edges are preserved in JSON + assert "edges" in parsed, "JSON should have 'edges' field" + assert len(parsed["edges"]) == 2, "JSON should have exactly two edges" + json_edges = parsed["edges"] + json_sources = [edge["source_id"] for edge in json_edges] + json_targets = [edge["target_id"] for edge in json_edges] + + assert all(source == "source" for source in json_sources), "JSON should preserve edge sources" + assert set(json_targets) == {"target1", "target2"}, "JSON should preserve edge targets" + + def test_fan_out_edge_group_serialization_with_selection_func(self) -> None: + """Test that FanOutEdgeGroup with named selection function serializes selection_func_name correctly.""" + + def custom_selector(data: Any, targets: list[str]) -> list[str]: + """Custom selection function for testing.""" + return targets[:1] # Select only the first target + + edge_group = FanOutEdgeGroup( + source_id="source", target_ids=["target1", "target2"], selection_func=custom_selector + ) + + # Test model_dump + data = edge_group.model_dump() + assert "selection_func_name" in data, "FanOutEdgeGroup should have 'selection_func_name' field" + assert data["selection_func_name"] == "custom_selector", ( + f"Expected selection_func_name 'custom_selector', got {data['selection_func_name']}" + ) + + # Test model_dump_json + json_str = edge_group.model_dump_json() + parsed = json.loads(json_str) + assert "selection_func_name" in parsed, "JSON should have 'selection_func_name' field" + assert parsed["selection_func_name"] == "custom_selector", "JSON should preserve selection_func_name" + + def test_fan_out_edge_group_serialization_with_lambda_selection_func(self) -> None: + """Test that FanOutEdgeGroup with lambda selection function serializes selection_func_name as ''.""" + edge_group = FanOutEdgeGroup( + source_id="source", target_ids=["target1", "target2"], selection_func=lambda data, targets: targets[:1] + ) + + # Test model_dump + data = edge_group.model_dump() + assert "selection_func_name" in data, "FanOutEdgeGroup should have 'selection_func_name' field" + assert data["selection_func_name"] == "", ( + f"Expected selection_func_name '', got {data['selection_func_name']}" + ) + + # Test model_dump_json + json_str = edge_group.model_dump_json() + parsed = json.loads(json_str) + assert "selection_func_name" in parsed, "JSON should have 'selection_func_name' field" + assert parsed["selection_func_name"] == "", "JSON should preserve selection_func_name as ''" + + def test_fan_in_edge_group_serialization(self) -> None: + """Test that FanInEdgeGroup can be serialized and has correct fields, including edges and type.""" + edge_group = FanInEdgeGroup(source_ids=["source1", "source2"], target_id="target") + + # Test model_dump + data = edge_group.model_dump() + assert "id" in data + assert data["id"].startswith("FanInEdgeGroup/") + + # Test type field + assert "type" in data, "FanInEdgeGroup should have 'type' field" + assert data["type"] == "FanInEdgeGroup", f"Expected type 'FanInEdgeGroup', got {data['type']}" + + # Verify edges field is present and contains the correct edges + assert "edges" in data, "FanInEdgeGroup should have 'edges' field" + assert len(data["edges"]) == 2, "FanInEdgeGroup should have exactly two edges" + + edges = data["edges"] + sources = [edge["source_id"] for edge in edges] + targets = [edge["target_id"] for edge in edges] + + assert set(sources) == {"source1", "source2"}, f"Expected sources {{'source1', 'source2'}}, got {set(sources)}" + assert all(target == "target" for target in targets), f"All edges should have target 'target', got {targets}" + + # Test model_dump_json + json_str = edge_group.model_dump_json() + parsed = json.loads(json_str) + assert "id" in parsed + assert parsed["id"].startswith("FanInEdgeGroup/") + + # Test type field in JSON + assert "type" in parsed, "JSON should have 'type' field" + assert parsed["type"] == "FanInEdgeGroup", "JSON should preserve type field" + + # Verify edges are preserved in JSON + assert "edges" in parsed, "JSON should have 'edges' field" + assert len(parsed["edges"]) == 2, "JSON should have exactly two edges" + json_edges = parsed["edges"] + json_sources = [edge["source_id"] for edge in json_edges] + json_targets = [edge["target_id"] for edge in json_edges] + + assert set(json_sources) == {"source1", "source2"}, "JSON should preserve edge sources" + assert all(target == "target" for target in json_targets), "JSON should preserve edge targets" + + def test_switch_case_edge_group_serialization(self) -> None: + """Test that SwitchCaseEdgeGroup can be serialized and has correct fields, including edges and type.""" + cases = [ + SwitchCaseEdgeGroupCase(condition=lambda x: x > 0, target_id="positive"), + SwitchCaseEdgeGroupDefault(target_id="default"), + ] + edge_group = SwitchCaseEdgeGroup(source_id="source", cases=cases) + + # Test model_dump + data = edge_group.model_dump() + assert "id" in data + assert data["id"].startswith("SwitchCaseEdgeGroup/") + + # Test type field + assert "type" in data, "SwitchCaseEdgeGroup should have 'type' field" + assert data["type"] == "SwitchCaseEdgeGroup", f"Expected type 'SwitchCaseEdgeGroup', got {data['type']}" + + # Test cases field + assert "cases" in data, "SwitchCaseEdgeGroup should have 'cases' field" + assert len(data["cases"]) == 2, "SwitchCaseEdgeGroup should have exactly two cases" + + cases_data = data["cases"] + # Check first case (SwitchCaseEdgeGroupCase) + case_obj = cases_data[0] + assert "target_id" in case_obj, "SwitchCaseEdgeGroupCase should have 'target_id' field" + assert "condition_name" in case_obj, "SwitchCaseEdgeGroupCase should have 'condition_name' field" + assert "type" in case_obj, "SwitchCaseEdgeGroupCase should have 'type' field" + assert case_obj["target_id"] == "positive", f"Expected target_id 'positive', got {case_obj['target_id']}" + assert case_obj["condition_name"] == "", ( + f"Expected condition_name '', got {case_obj['condition_name']}" + ) + assert case_obj["type"] == "Case", f"Expected type 'Case', got {case_obj['type']}" + + # Check default case (SwitchCaseEdgeGroupDefault) + default_obj = cases_data[1] + assert "target_id" in default_obj, "SwitchCaseEdgeGroupDefault should have 'target_id' field" + assert "type" in default_obj, "SwitchCaseEdgeGroupDefault should have 'type' field" + assert default_obj["target_id"] == "default", f"Expected target_id 'default', got {default_obj['target_id']}" + assert default_obj["type"] == "Default", f"Expected type 'Default', got {default_obj['type']}" + + # Verify edges field is present and contains the correct edges + assert "edges" in data, "SwitchCaseEdgeGroup should have 'edges' field" + assert len(data["edges"]) == 2, "SwitchCaseEdgeGroup should have exactly two edges" + + edges = data["edges"] + sources = [edge["source_id"] for edge in edges] + targets = [edge["target_id"] for edge in edges] + + assert all(source == "source" for source in sources), f"All edges should have source 'source', got {sources}" + assert set(targets) == {"positive", "default"}, ( + f"Expected targets {{'positive', 'default'}}, got {set(targets)}" + ) + + # Check condition_name field in edges - SwitchCaseEdgeGroup edges don't have conditions + # because the conditional logic is implemented in the selection_func at the group level + condition_names = [edge.get("condition_name") for edge in edges] + assert all(name is None for name in condition_names), ( + "SwitchCaseEdgeGroup edges should not have condition_name since conditions are handled at group level" + ) + + # Test model_dump_json + json_str = edge_group.model_dump_json() + parsed = json.loads(json_str) + assert "id" in parsed + assert parsed["id"].startswith("SwitchCaseEdgeGroup/") + + # Test type field in JSON + assert "type" in parsed, "JSON should have 'type' field" + assert parsed["type"] == "SwitchCaseEdgeGroup", "JSON should preserve type field" + + # Test cases field in JSON + assert "cases" in parsed, "JSON should have 'cases' field" + assert len(parsed["cases"]) == 2, "JSON should have exactly two cases" + + json_cases = parsed["cases"] + json_case_obj = json_cases[0] + assert json_case_obj["target_id"] == "positive", "JSON should preserve case target_id" + assert json_case_obj["condition_name"] == "", "JSON should preserve case condition_name" + assert json_case_obj["type"] == "Case", "JSON should preserve case type" + + json_default_obj = json_cases[1] + assert json_default_obj["target_id"] == "default", "JSON should preserve default target_id" + assert json_default_obj["type"] == "Default", "JSON should preserve default type" + + # Verify edges are preserved in JSON + assert "edges" in parsed, "JSON should have 'edges' field" + assert len(parsed["edges"]) == 2, "JSON should have exactly two edges" + json_edges = parsed["edges"] + json_sources = [edge["source_id"] for edge in json_edges] + json_targets = [edge["target_id"] for edge in json_edges] + + assert all(source == "source" for source in json_sources), "JSON should preserve edge sources" + assert set(json_targets) == {"positive", "default"}, "JSON should preserve edge targets" + + # Check condition_name field in JSON edges - should be None for SwitchCaseEdgeGroup + json_condition_names = [edge.get("condition_name") for edge in json_edges] + assert all(name is None for name in json_condition_names), ( + "JSON SwitchCaseEdgeGroup edges should not have condition_name" + ) + + def test_switch_case_edge_group_serialization_with_named_condition(self) -> None: + """Test that SwitchCaseEdgeGroup with named condition function serializes condition_name correctly.""" + + def is_positive(x: int) -> bool: + return x > 0 + + cases = [ + SwitchCaseEdgeGroupCase(condition=is_positive, target_id="positive"), + SwitchCaseEdgeGroupDefault(target_id="default"), + ] + edge_group = SwitchCaseEdgeGroup(source_id="source", cases=cases) + + # Test model_dump + data = edge_group.model_dump() + assert "cases" in data, "SwitchCaseEdgeGroup should have 'cases' field" + + cases_data = data["cases"] + case_obj = cases_data[0] + assert case_obj["condition_name"] == "is_positive", ( + f"Expected condition_name 'is_positive', got {case_obj['condition_name']}" + ) + + # Test model_dump_json + json_str = edge_group.model_dump_json() + parsed = json.loads(json_str) + json_cases = parsed["cases"] + json_case_obj = json_cases[0] + assert json_case_obj["condition_name"] == "is_positive", "JSON should preserve named condition_name" + + def test_workflow_serialization(self) -> None: + """Test that Workflow can be serialized and has correct fields, including edges.""" + executor1 = SampleExecutor(id="executor1") + executor2 = SampleExecutor(id="executor2") + + workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build() + + # Test model_dump + data = workflow.model_dump() + assert "edge_groups" in data + assert "executors" in data + assert "start_executor_id" in data + assert "max_iterations" in data + assert "workflow_id" in data + + assert data["start_executor_id"] == "executor1" + assert "executor1" in data["executors"] + assert "executor2" in data["executors"] + + # Verify edge groups contain edges + edge_groups = data["edge_groups"] + assert len(edge_groups) == 1, "Should have exactly one edge group" + edge_group = edge_groups[0] + assert "edges" in edge_group, "Edge group should contain 'edges' field" + assert len(edge_group["edges"]) == 1, "Should have exactly one edge" + + edge = edge_group["edges"][0] + assert "source_id" in edge, "Edge should have source_id" + assert "target_id" in edge, "Edge should have target_id" + assert edge["source_id"] == "executor1", f"Expected source_id 'executor1', got {edge['source_id']}" + assert edge["target_id"] == "executor2", f"Expected target_id 'executor2', got {edge['target_id']}" + + # Test model_dump_json + json_str = workflow.model_dump_json() + parsed = json.loads(json_str) + assert parsed["start_executor_id"] == "executor1" + assert "executor1" in parsed["executors"] + assert "executor2" in parsed["executors"] + + # Verify edges are preserved in JSON serialization + json_edge_groups = parsed["edge_groups"] + assert len(json_edge_groups) == 1, "JSON should have exactly one edge group" + json_edge_group = json_edge_groups[0] + assert "edges" in json_edge_group, "JSON edge group should contain 'edges' field" + json_edge = json_edge_group["edges"][0] + assert json_edge["source_id"] == "executor1", "JSON should preserve edge source_id" + assert json_edge["target_id"] == "executor2", "JSON should preserve edge target_id" + + def test_workflow_serialization_excludes_non_serializable_fields(self) -> None: + """Test that non-serializable fields are excluded from serialization.""" + executor1 = SampleExecutor(id="executor1") + executor2 = SampleExecutor(id="executor2") + + workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build() + + # Test model_dump - should not include private runtime objects + data = workflow.model_dump() + + # These private runtime fields should not be in the serialized data + assert "_runner_context" not in data + assert "_shared_state" not in data + assert "_runner" not in data + + def test_executor_field_validation(self) -> None: + """Test that Executor field validation works correctly.""" + # Valid executor + executor = SampleExecutor(id="valid-id") + assert executor.id == "valid-id" + + # Test validation failure for empty id - pydantic automatically validates min_length=1 + from pydantic import ValidationError + + with pytest.raises(ValidationError): + SampleExecutor(id="") + + def test_edge_field_validation(self) -> None: + """Test that Edge field validation works correctly.""" + # Valid edge + edge = Edge(source_id="source", target_id="target") + assert edge.source_id == "source" + assert edge.target_id == "target" + + # Test validation failure for empty source_id + from pydantic import ValidationError + + with pytest.raises(ValidationError): + Edge(source_id="", target_id="target") + + # Test validation failure for empty target_id + with pytest.raises(ValidationError): + Edge(source_id="source", target_id="") + + +def test_comprehensive_edge_groups_workflow_serialization() -> None: + """Test serialization of a workflow that uses all edge group types: SwitchCase, FanOut, and FanIn.""" + from agent_framework_workflow._edge import Case, Default + + # Create executors for a comprehensive workflow + router = SampleExecutor(id="router") + processor_a = SampleExecutor(id="proc_a") + processor_b = SampleExecutor(id="proc_b") + fanout_hub = SampleExecutor(id="fanout_hub") + parallel_1 = SampleExecutor(id="parallel_1") + parallel_2 = SampleExecutor(id="parallel_2") + aggregator = SampleAggregator(id="aggregator") + + # Build workflow with all three edge group types + workflow = ( + WorkflowBuilder() + .set_start_executor(router) + # 1. SwitchCaseEdgeGroup: Conditional routing + .add_switch_case_edge_group( + router, + [ + Case(condition=lambda msg: len(str(msg)) < 10, target=processor_a), + Default(target=processor_b), + ], + ) + # 2. Direct edges + .add_edge(processor_a, fanout_hub) + .add_edge(processor_b, fanout_hub) + # 3. FanOutEdgeGroup: One-to-many distribution + .add_fan_out_edges(fanout_hub, [parallel_1, parallel_2]) + # 4. FanInEdgeGroup: Many-to-one aggregation + .add_fan_in_edges([parallel_1, parallel_2], aggregator) + .build() + ) + + # Test workflow serialization + data = workflow.model_dump() + + # Verify basic workflow structure + assert "edge_groups" in data + assert "executors" in data + assert "start_executor_id" in data + assert data["start_executor_id"] == "router" + + # Verify all executors are present + expected_executors = {"router", "proc_a", "proc_b", "fanout_hub", "parallel_1", "parallel_2", "aggregator"} + assert set(data["executors"].keys()) == expected_executors + + # Verify edge groups contain all three types + edge_groups = data["edge_groups"] + edge_group_types = [eg.get("id", "").split("/")[0] for eg in edge_groups] + + # Should have: SwitchCaseEdgeGroup, SingleEdgeGroup (x2), FanOutEdgeGroup, FanInEdgeGroup + assert "SwitchCaseEdgeGroup" in edge_group_types, f"Expected SwitchCaseEdgeGroup in {edge_group_types}" + assert "FanOutEdgeGroup" in edge_group_types, f"Expected FanOutEdgeGroup in {edge_group_types}" + assert "FanInEdgeGroup" in edge_group_types, f"Expected FanInEdgeGroup in {edge_group_types}" + assert "SingleEdgeGroup" in edge_group_types, f"Expected SingleEdgeGroup in {edge_group_types}" + + # Test JSON serialization + json_str = workflow.model_dump_json() + parsed = json.loads(json_str) + + # Verify JSON structure matches model_dump + assert parsed["start_executor_id"] == "router" + assert set(parsed["executors"].keys()) == expected_executors + assert len(parsed["edge_groups"]) == len(edge_groups) + + # Verify that serialization excludes non-serializable fields + assert "_runner_context" not in data + assert "_shared_state" not in data + assert "_runner" not in data + + # Test that we can identify each edge group type by examining their structure + switch_case_groups = [eg for eg in edge_groups if eg.get("id", "").startswith("SwitchCaseEdgeGroup/")] + fan_out_groups = [eg for eg in edge_groups if eg.get("id", "").startswith("FanOutEdgeGroup/")] + fan_in_groups = [eg for eg in edge_groups if eg.get("id", "").startswith("FanInEdgeGroup/")] + single_groups = [eg for eg in edge_groups if eg.get("id", "").startswith("SingleEdgeGroup/")] + + assert len(switch_case_groups) == 1, f"Expected 1 SwitchCaseEdgeGroup, got {len(switch_case_groups)}" + assert len(fan_out_groups) == 1, f"Expected 1 FanOutEdgeGroup, got {len(fan_out_groups)}" + assert len(fan_in_groups) == 1, f"Expected 1 FanInEdgeGroup, got {len(fan_in_groups)}" + assert len(single_groups) == 2, f"Expected 2 SingleEdgeGroups, got {len(single_groups)}" + + # The key validation is that all edge group types are present and serializable + # Individual edge group fields may vary based on implementation, + # but each should have at least an 'id' field that identifies its type and 'edges' field + for group_type, groups in [ + ("SwitchCaseEdgeGroup", switch_case_groups), + ("FanOutEdgeGroup", fan_out_groups), + ("FanInEdgeGroup", fan_in_groups), + ("SingleEdgeGroup", single_groups), + ]: + for group in groups: + assert "id" in group, f"{group_type} should have 'id' field" + assert group["id"].startswith(f"{group_type}/"), f"{group_type} id should start with '{group_type}/'" + assert "edges" in group, f"{group_type} should have 'edges' field" + assert isinstance(group["edges"], list), f"{group_type} 'edges' should be a list" + assert len(group["edges"]) > 0, f"{group_type} should have at least one edge" + + # Verify each edge has required fields + for edge in group["edges"]: + assert "source_id" in edge, f"{group_type} edge should have 'source_id'" + assert "target_id" in edge, f"{group_type} edge should have 'target_id'" + assert isinstance(edge["source_id"], str), f"{group_type} edge source_id should be string" + assert isinstance(edge["target_id"], str), f"{group_type} edge target_id should be string" + assert len(edge["source_id"]) > 0, f"{group_type} edge source_id should not be empty" + assert len(edge["target_id"]) > 0, f"{group_type} edge target_id should not be empty" + + # Verify specific edge group edge counts + assert len(switch_case_groups[0]["edges"]) == 2, "SwitchCaseEdgeGroup should have 2 edges (proc_a and proc_b)" + assert len(fan_out_groups[0]["edges"]) == 2, "FanOutEdgeGroup should have 2 edges (parallel_1 and parallel_2)" + assert len(fan_in_groups[0]["edges"]) == 2, "FanInEdgeGroup should have 2 edges (from parallel_1 and parallel_2)" + for single_group in single_groups: + assert len(single_group["edges"]) == 1, "Each SingleEdgeGroup should have exactly 1 edge" diff --git a/python/packages/workflow/tests/test_validation.py b/python/packages/workflow/tests/test_validation.py index 49f694ebc2..dcfdf20705 100644 --- a/python/packages/workflow/tests/test_validation.py +++ b/python/packages/workflow/tests/test_validation.py @@ -161,12 +161,14 @@ def test_graph_connectivity_isolated_executors(): # Create edges that include an isolated executor (self-loop that's not connected to main graph) edge_groups = [ - SingleEdgeGroup(executor1, executor2), - SingleEdgeGroup(executor3, executor3), + SingleEdgeGroup(executor1.id, executor2.id), + SingleEdgeGroup(executor3.id, executor3.id), ] # Self-loop to include in graph + executors: dict[str, Executor] = {executor1.id: executor1, executor2.id: executor2, executor3.id: executor3} + with pytest.raises(GraphConnectivityError) as exc_info: - validate_workflow_graph(edge_groups, executor1) + validate_workflow_graph(edge_groups, executors, executor1) assert "unreachable" in str(exc_info.value).lower() assert "executor3" in str(exc_info.value) @@ -243,15 +245,16 @@ def test_type_compatibility_inheritance(): def test_direct_validation_function(): executor1 = StringExecutor(id="executor1") executor2 = StringExecutor(id="executor2") - edge_groups = [SingleEdgeGroup(executor1, executor2)] + edge_groups = [SingleEdgeGroup(executor1.id, executor2.id)] + executors: dict[str, Executor] = {executor1.id: executor1, executor2.id: executor2} # This should not raise any exceptions - validate_workflow_graph(edge_groups, executor1) + validate_workflow_graph(edge_groups, executors, executor1) # Test with invalid start executor executor3 = StringExecutor(id="executor3") with pytest.raises(GraphConnectivityError): - validate_workflow_graph(edge_groups, executor3) + validate_workflow_graph(edge_groups, executors, executor3) def test_fan_out_validation(): diff --git a/python/packages/workflow/tests/test_workflow.py b/python/packages/workflow/tests/test_workflow.py index 52d677037c..7f6ed73da0 100644 --- a/python/packages/workflow/tests/test_workflow.py +++ b/python/packages/workflow/tests/test_workflow.py @@ -22,33 +22,31 @@ from agent_framework_workflow import Message @dataclass -class MockMessage: +class NumberMessage: """A mock message for testing purposes.""" data: int -class MockExecutor(Executor): - """A mock executor for testing purposes.""" +class IncrementExecutor(Executor): + """An executor that increments message data by a specified amount for testing purposes.""" - def __init__(self, id: str, limit: int = 10): - """Initialize the mock executor with a limit.""" - super().__init__(id=id) - self.limit = limit + limit: int = 10 + increment: int = 1 @handler - async def mock_handler(self, message: MockMessage, ctx: WorkflowContext[MockMessage]) -> None: + async def mock_handler(self, message: NumberMessage, ctx: WorkflowContext[NumberMessage]) -> None: if message.data < self.limit: - await ctx.send_message(MockMessage(data=message.data + 1)) + await ctx.send_message(NumberMessage(data=message.data + self.increment)) else: await ctx.add_event(WorkflowCompletedEvent(data=message.data)) -class MockAggregator(Executor): +class AggregatorExecutor(Executor): """A mock executor that aggregates results from multiple executors.""" @handler - async def mock_handler(self, messages: list[MockMessage], ctx: WorkflowContext[Any]) -> None: + async def mock_handler(self, messages: list[NumberMessage], ctx: WorkflowContext[Any]) -> None: # This mock simply returns the data incremented by 1 await ctx.add_event(WorkflowCompletedEvent(data=sum(msg.data for msg in messages))) @@ -64,25 +62,25 @@ class MockExecutorRequestApproval(Executor): """A mock executor that simulates a request for approval.""" @handler - async def mock_handler_a(self, message: MockMessage, ctx: WorkflowContext[RequestInfoMessage]) -> None: + async def mock_handler_a(self, message: NumberMessage, ctx: WorkflowContext[RequestInfoMessage]) -> None: """A mock handler that requests approval.""" await ctx.set_shared_state(self.id, message.data) await ctx.send_message(RequestInfoMessage()) @handler - async def mock_handler_b(self, message: ApprovalMessage, ctx: WorkflowContext[MockMessage]) -> None: + async def mock_handler_b(self, message: ApprovalMessage, ctx: WorkflowContext[NumberMessage]) -> None: """A mock handler that processes the approval response.""" data = await ctx.get_shared_state(self.id) if message.approved: await ctx.add_event(WorkflowCompletedEvent(data=data)) else: - await ctx.send_message(MockMessage(data=data)) + await ctx.send_message(NumberMessage(data=data)) async def test_workflow_run_streaming(): """Test the workflow run stream.""" - executor_a = MockExecutor(id="executor_a") - executor_b = MockExecutor(id="executor_b") + executor_a = IncrementExecutor(id="executor_a") + executor_b = IncrementExecutor(id="executor_b") workflow = ( WorkflowBuilder() @@ -93,7 +91,7 @@ async def test_workflow_run_streaming(): ) result: int | None = None - async for event in workflow.run_streaming(MockMessage(data=0)): + async for event in workflow.run_streaming(NumberMessage(data=0)): assert isinstance(event, WorkflowEvent) if isinstance(event, WorkflowCompletedEvent): result = event.data @@ -103,8 +101,8 @@ async def test_workflow_run_streaming(): async def test_workflow_run_stream_not_completed(): """Test the workflow run stream.""" - executor_a = MockExecutor(id="executor_a") - executor_b = MockExecutor(id="executor_b") + executor_a = IncrementExecutor(id="executor_a") + executor_b = IncrementExecutor(id="executor_b") workflow = ( WorkflowBuilder() @@ -116,14 +114,14 @@ async def test_workflow_run_stream_not_completed(): ) with pytest.raises(RuntimeError): - async for _ in workflow.run_streaming(MockMessage(data=0)): + async for _ in workflow.run_streaming(NumberMessage(data=0)): pass async def test_workflow_run(): """Test the workflow run.""" - executor_a = MockExecutor(id="executor_a") - executor_b = MockExecutor(id="executor_b") + executor_a = IncrementExecutor(id="executor_a") + executor_b = IncrementExecutor(id="executor_b") workflow = ( WorkflowBuilder() @@ -133,7 +131,7 @@ async def test_workflow_run(): .build() ) - events = await workflow.run(MockMessage(data=0)) + events = await workflow.run(NumberMessage(data=0)) completed_event = events.get_completed_event() assert isinstance(completed_event, WorkflowCompletedEvent) assert completed_event.data == 10 @@ -141,8 +139,8 @@ async def test_workflow_run(): async def test_workflow_run_not_completed(): """Test the workflow run.""" - executor_a = MockExecutor(id="executor_a") - executor_b = MockExecutor(id="executor_b") + executor_a = IncrementExecutor(id="executor_a") + executor_b = IncrementExecutor(id="executor_b") workflow = ( WorkflowBuilder() @@ -154,12 +152,12 @@ async def test_workflow_run_not_completed(): ) with pytest.raises(RuntimeError): - await workflow.run(MockMessage(data=0)) + await workflow.run(NumberMessage(data=0)) async def test_workflow_send_responses_streaming(): """Test the workflow run with approval.""" - executor_a = MockExecutor(id="executor_a") + executor_a = IncrementExecutor(id="executor_a") executor_b = MockExecutorRequestApproval(id="executor_b") request_info_executor = RequestInfoExecutor() @@ -174,7 +172,7 @@ async def test_workflow_send_responses_streaming(): ) request_info_event: RequestInfoEvent | None = None - async for event in workflow.run_streaming(MockMessage(data=0)): + async for event in workflow.run_streaming(NumberMessage(data=0)): if isinstance(event, RequestInfoEvent): request_info_event = event @@ -191,7 +189,7 @@ async def test_workflow_send_responses_streaming(): async def test_workflow_send_responses(): """Test the workflow run with approval.""" - executor_a = MockExecutor(id="executor_a") + executor_a = IncrementExecutor(id="executor_a") executor_b = MockExecutorRequestApproval(id="executor_b") request_info_executor = RequestInfoExecutor() @@ -205,7 +203,7 @@ async def test_workflow_send_responses(): .build() ) - events = await workflow.run(MockMessage(data=0)) + events = await workflow.run(NumberMessage(data=0)) request_info_events = events.get_request_info_events() assert len(request_info_events) == 1 @@ -219,15 +217,15 @@ async def test_workflow_send_responses(): async def test_fan_out(): """Test a fan-out workflow.""" - executor_a = MockExecutor(id="executor_a") - executor_b = MockExecutor(id="executor_b", limit=1) - executor_c = MockExecutor(id="executor_c", limit=2) # This executor will not complete the workflow + executor_a = IncrementExecutor(id="executor_a") + executor_b = IncrementExecutor(id="executor_b", limit=1) + executor_c = IncrementExecutor(id="executor_c", limit=2) # This executor will not complete the workflow workflow = ( WorkflowBuilder().set_start_executor(executor_a).add_fan_out_edges(executor_a, [executor_b, executor_c]).build() ) - events = await workflow.run(MockMessage(data=0)) + events = await workflow.run(NumberMessage(data=0)) # Each executor will emit two events: ExecutorInvokeEvent and ExecutorCompletedEvent # executor_b will also emit a WorkflowCompletedEvent @@ -239,15 +237,15 @@ async def test_fan_out(): async def test_fan_out_multiple_completed_events(): """Test a fan-out workflow with multiple completed events.""" - executor_a = MockExecutor(id="executor_a") - executor_b = MockExecutor(id="executor_b", limit=1) - executor_c = MockExecutor(id="executor_c", limit=1) + executor_a = IncrementExecutor(id="executor_a") + executor_b = IncrementExecutor(id="executor_b", limit=1) + executor_c = IncrementExecutor(id="executor_c", limit=1) workflow = ( WorkflowBuilder().set_start_executor(executor_a).add_fan_out_edges(executor_a, [executor_b, executor_c]).build() ) - events = await workflow.run(MockMessage(data=0)) + events = await workflow.run(NumberMessage(data=0)) # Each executor will emit two events: ExecutorInvokeEvent and ExecutorCompletedEvent # executor_a and executor_b will also emit a WorkflowCompletedEvent @@ -259,10 +257,10 @@ async def test_fan_out_multiple_completed_events(): async def test_fan_in(): """Test a fan-in workflow.""" - executor_a = MockExecutor(id="executor_a") - executor_b = MockExecutor(id="executor_b") - executor_c = MockExecutor(id="executor_c") - aggregator = MockAggregator(id="aggregator") + executor_a = IncrementExecutor(id="executor_a") + executor_b = IncrementExecutor(id="executor_b") + executor_c = IncrementExecutor(id="executor_c") + aggregator = AggregatorExecutor(id="aggregator") workflow = ( WorkflowBuilder() @@ -272,7 +270,7 @@ async def test_fan_in(): .build() ) - events = await workflow.run(MockMessage(data=0)) + events = await workflow.run(NumberMessage(data=0)) # Each executor will emit two events: ExecutorInvokeEvent and ExecutorCompletedEvent # aggregator will also emit a WorkflowCompletedEvent @@ -289,7 +287,7 @@ def simple_executor() -> Executor: async def handle_message(self, message: Message, context: WorkflowContext[None]) -> None: pass - return SimpleExecutor("test_executor") + return SimpleExecutor(id="test_executor") async def test_workflow_with_checkpointing_enabled(simple_executor: Executor): @@ -521,7 +519,7 @@ async def test_workflow_multiple_runs_no_state_collision(): storage = FileCheckpointStorage(temp_dir) # Create executor that tracks state in shared state - state_executor = StateTrackingExecutor("state_executor") + state_executor = StateTrackingExecutor(id="state_executor") # Build workflow with checkpointing workflow = ( @@ -555,3 +553,123 @@ async def test_workflow_multiple_runs_no_state_collision(): assert completed1.data != completed2.data assert completed2.data != completed3.data assert completed1.data != completed3.data + + +async def test_comprehensive_edge_groups_workflow(): + """Test a workflow that uses SwitchCaseEdgeGroup, FanOutEdgeGroup, and FanInEdgeGroup.""" + from agent_framework_workflow._edge import Case, Default + + # Create 6 executors for different roles with different increment values + router = IncrementExecutor(id="router", limit=1000, increment=1) # Increment by 1 + processor_a = IncrementExecutor(id="proc_a", limit=1000, increment=1) # Increment by 1 + processor_b = IncrementExecutor(id="proc_b", limit=1000, increment=2) # Increment by 2 (different from proc_a) + fanout_hub = IncrementExecutor(id="fanout_hub", limit=1000, increment=1) # Increment by 1 + parallel_1 = IncrementExecutor(id="parallel_1", limit=1000, increment=3) # Increment by 3 + parallel_2 = IncrementExecutor( + id="parallel_2", limit=1000, increment=5 + ) # Increment by 5 (different from parallel_1) + aggregator = AggregatorExecutor(id="aggregator") # Combines results from parallel processors + + # Build workflow with different edge group types: + # 1. SwitchCase: router -> (proc_a if data < 5, else proc_b) + # 2. Direct edge: proc_a -> fanout_hub, proc_b -> fanout_hub + # 3. FanOut: fanout_hub -> [parallel_1, parallel_2] + # 4. FanIn: [parallel_1, parallel_2] -> aggregator + workflow = ( + WorkflowBuilder() + .set_start_executor(router) + # Switch-case routing based on message data + .add_switch_case_edge_group( + router, + [ + Case(condition=lambda msg: msg.data < 5, target=processor_a), + Default(target=processor_b), + ], + ) + # Both processors send to fanout hub + .add_edge(processor_a, fanout_hub) + .add_edge(processor_b, fanout_hub) + # Fan out to parallel processors + .add_fan_out_edges(fanout_hub, [parallel_1, parallel_2]) + # Fan in to aggregator + .add_fan_in_edges([parallel_1, parallel_2], aggregator) + .build() + ) + + # Test with small number (should go through processor_a) + # router(2->3) -> switch routes to proc_a -> proc_a(3->4) -> fanout_hub(4->5) + # -> [parallel_1(5->8), parallel_2(5->10)] -> aggregator(8+10=18) + events_small = await workflow.run(NumberMessage(data=2)) + completed_small = events_small.get_completed_event() + assert completed_small is not None + assert completed_small.data == 18 # Exact expected result: 8+10 from parallel processors + + # Test with large number (should go through processor_b) + # router(8->9) -> switch routes to proc_b -> proc_b(9->11) -> fanout_hub(11->12) + # -> [parallel_1(12->15), parallel_2(12->17)] -> aggregator(15+17=32) + events_large = await workflow.run(NumberMessage(data=8)) + completed_large = events_large.get_completed_event() + assert completed_large is not None + assert completed_large.data == 32 # Exact expected result: 15+17 from parallel processors + + # The key verification is that we successfully executed a workflow using all three edge group types + # and that both switch-case paths work (small vs large numbers) + + # Verify we had multiple events indicating complex execution path + assert len(events_small) >= 6 # Should have multiple executors involved + assert len(events_large) >= 6 + + # Verify different paths were taken by checking exact results + assert completed_small.data == 18, f"Small number path should result in 18, got {completed_small.data}" + assert completed_large.data == 32, f"Large number path should result in 32, got {completed_large.data}" + assert completed_small.data != completed_large.data, "Different paths should produce different results" + + # Both tests should complete successfully, proving all edge group types work + + # Additional verification: check that the workflow contains the expected edge group types + edge_groups = workflow.edge_groups + has_switch_case = any(edge_group.__class__.__name__ == "SwitchCaseEdgeGroup" for edge_group in edge_groups) + has_fan_out = any(edge_group.__class__.__name__ == "FanOutEdgeGroup" for edge_group in edge_groups) + has_fan_in = any(edge_group.__class__.__name__ == "FanInEdgeGroup" for edge_group in edge_groups) + + assert has_switch_case, "Workflow should contain SwitchCaseEdgeGroup" + assert has_fan_out, "Workflow should contain FanOutEdgeGroup" + assert has_fan_in, "Workflow should contain FanInEdgeGroup" + + +async def test_workflow_with_simple_cycle_and_exit_condition(): + """Test a simpler workflow with a cycle that has a clear exit condition.""" + + # Create a simple cycle: A -> B -> A, with A having an exit condition + executor_a = IncrementExecutor(id="exec_a", limit=6, increment=2) # Exit when data >= 6 + executor_b = IncrementExecutor(id="exec_b", limit=1000, increment=1) # Never exit, just increment + + # Simple cycle: A -> B -> A, A exits when limit reached + workflow = ( + WorkflowBuilder() + .set_start_executor(executor_a) + .add_edge(executor_a, executor_b) # A -> B + .add_edge(executor_b, executor_a) # B -> A (creates cycle) + .build() + ) + + # Test the cycle + # Expected: exec_a(2->4) -> exec_b(4->5) -> exec_a(5->7, completes because 7 >= 6) + events = await workflow.run(NumberMessage(data=2)) + completed_event = events.get_completed_event() + assert completed_event is not None + assert ( + completed_event.data is not None and completed_event.data >= 6 + ) # Should complete when executor_a reaches its limit + + # Verify cycling occurred (should have events from both executors) + # Check for ExecutorInvokeEvent and ExecutorCompletedEvent types that have executor_id + from agent_framework_workflow._events import ExecutorCompletedEvent, ExecutorInvokeEvent + + executor_events = [e for e in events if isinstance(e, (ExecutorInvokeEvent, ExecutorCompletedEvent))] + executor_ids = {e.executor_id for e in executor_events} + assert "exec_a" in executor_ids, "Should have events from executor A" + assert "exec_b" in executor_ids, "Should have events from executor B" + + # Should have multiple events due to cycling + assert len(events) >= 4, f"Expected at least 4 events due to cycling, got {len(events)}" diff --git a/python/packages/workflow/tests/test_workflow_builder.py b/python/packages/workflow/tests/test_workflow_builder.py index c8ba88a962..469cd5a3c5 100644 --- a/python/packages/workflow/tests/test_workflow_builder.py +++ b/python/packages/workflow/tests/test_workflow_builder.py @@ -61,5 +61,5 @@ def test_workflow_builder_fluent_api(): ) assert len(workflow.edge_groups) == 4 - assert workflow.start_executor.id == executor_a.id + assert workflow.start_executor_id == executor_a.id assert len(workflow.executors) == 6