agent/agent-framework
1
0
Fork 0
mirror of https://github.com/microsoft/agent-framework.git synced 2026-06-16 21:04:09 +08:00
Code Issues Packages Projects Releases Wiki Activity

Python: Move Workflow, Edge, EdgeGroup and Executor to AFBaseModel (#472)

Browse Source 
* 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
This commit is contained in:
Eric Zhu
2025-08-22 15:02:34 -07:00
committed by GitHub
Unverified
parent 88c51013c6
commit 6e9d35830f
13 changed files with 1588 additions and 633 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/workflow/agent_framework_workflow/_edge.py
+171 -241
View File
@@ -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 == "<lambda>":
return "<lambda>"
return name
# Fallback for other callable objects
return "<callable>"
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)
python/packages/workflow/agent_framework_workflow/_edge_runner.py
+198
View File
@@ -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)}")
python/packages/workflow/agent_framework_workflow/_executor.py
+39 -18
View File
@@ -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 <class_name>/<uuid>.
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.
python/packages/workflow/agent_framework_workflow/_runner.py
+21 -22
View File
@@ -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
python/packages/workflow/agent_framework_workflow/_validation.py
+28 -23
View File
@@ -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)
python/packages/workflow/agent_framework_workflow/_viz.py
+14 -14
View File
@@ -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))
python/packages/workflow/agent_framework_workflow/_workflow.py
+82 -50
View File
@@ -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
python/packages/workflow/tests/test_edge.py
+209 -201
View File
@@ -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
python/packages/workflow/tests/test_runner.py
+17 -12
View File
@@ -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),
python/packages/workflow/tests/test_serialization.py
+636
View File
@@ -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 '<lambda>'."""
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"] == "<lambda>"
# 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"] == "<lambda>"
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 '<lambda>'."""
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"] == "<lambda>", (
f"Expected selection_func_name '<lambda>', 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"] == "<lambda>", "JSON should preserve selection_func_name as '<lambda>'"
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"] == "<lambda>", (
f"Expected condition_name '<lambda>', 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"] == "<lambda>", "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"
python/packages/workflow/tests/test_validation.py
+9 -6
View File
@@ -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():
python/packages/workflow/tests/test_workflow.py
+163 -45
View File
@@ -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)}"
python/packages/workflow/tests/test_workflow_builder.py
+1 -1
View File
@@ -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