mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
[BREAKING] Python: Remove workflow register factory methods. Update tests and samples (#3781)
* Remove workflow register factory methods. Update tests and samples * Address Copilot feedback
This commit is contained in:
Evan Mattson
committed by
GitHub
Unverified
GitHub
parent
f407f726a7
commit
a4c9e43afb
@@ -3,11 +3,9 @@
|
||||
import logging
|
||||
import sys
|
||||
from collections.abc import Callable, Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from .._agents import SupportsAgentRun
|
||||
from .._threads import AgentThread
|
||||
from ..observability import OtelAttr, capture_exception, create_workflow_span
|
||||
from ._agent_executor import AgentExecutor
|
||||
from ._agent_utils import resolve_agent_id
|
||||
@@ -40,76 +38,6 @@ else:
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class _EdgeRegistration:
|
||||
"""A data class representing an edge registration in the workflow builder.
|
||||
|
||||
Args:
|
||||
source: The registered source name.
|
||||
target: The registered target name.
|
||||
condition: An optional condition function `(data) -> bool | Awaitable[bool]`.
|
||||
"""
|
||||
|
||||
source: str
|
||||
target: str
|
||||
condition: EdgeCondition | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FanOutEdgeRegistration:
|
||||
"""A data class representing a fan-out edge registration in the workflow builder.
|
||||
|
||||
Args:
|
||||
source: The registered source name.
|
||||
targets: A list of registered target names.
|
||||
"""
|
||||
|
||||
source: str
|
||||
targets: list[str]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _FanInEdgeRegistration:
|
||||
"""A data class representing a fan-in edge registration in the workflow builder.
|
||||
|
||||
Args:
|
||||
sources: A list of registered source names.
|
||||
target: The registered target name.
|
||||
"""
|
||||
|
||||
sources: list[str]
|
||||
target: str
|
||||
|
||||
|
||||
@dataclass
|
||||
class _SwitchCaseEdgeGroupRegistration:
|
||||
"""A data class representing a switch-case edge group registration in the workflow builder.
|
||||
|
||||
Args:
|
||||
source: The registered source name.
|
||||
cases: A list of case objects that determine the target executor for each message.
|
||||
"""
|
||||
|
||||
source: str
|
||||
cases: list[Case | Default]
|
||||
|
||||
|
||||
@dataclass
|
||||
class _MultiSelectionEdgeGroupRegistration:
|
||||
"""A data class representing a multi-selection edge group registration in the workflow builder.
|
||||
|
||||
Args:
|
||||
source: The registered source name.
|
||||
targets: A list of registered target names.
|
||||
selection_func: A function that selects target executors for messages.
|
||||
Takes (message, list[registered target names]) and returns list[registered target names].
|
||||
"""
|
||||
|
||||
source: str
|
||||
targets: list[str]
|
||||
selection_func: Callable[[Any, list[str]], list[str]]
|
||||
|
||||
|
||||
class WorkflowBuilder:
|
||||
"""A builder class for constructing workflows.
|
||||
|
||||
@@ -136,14 +64,10 @@ class WorkflowBuilder:
|
||||
await ctx.yield_output(text[::-1])
|
||||
|
||||
|
||||
# Build a workflow
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor="UpperCase")
|
||||
.register_executor(lambda: UpperCaseExecutor(id="upper"), name="UpperCase")
|
||||
.register_executor(lambda: ReverseExecutor(id="reverse"), name="Reverse")
|
||||
.add_edge("UpperCase", "Reverse")
|
||||
.build()
|
||||
)
|
||||
upper = UpperCaseExecutor(id="upper")
|
||||
reverse = ReverseExecutor(id="reverse")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=upper).add_edge(upper, reverse).build()
|
||||
|
||||
# Run the workflow
|
||||
events = await workflow.run("hello")
|
||||
@@ -156,9 +80,9 @@ class WorkflowBuilder:
|
||||
name: str | None = None,
|
||||
description: str | None = None,
|
||||
*,
|
||||
start_executor: Executor | SupportsAgentRun | str,
|
||||
start_executor: Executor | SupportsAgentRun,
|
||||
checkpoint_storage: CheckpointStorage | None = None,
|
||||
output_executors: list[Executor | SupportsAgentRun | str] | None = None,
|
||||
output_executors: list[Executor | SupportsAgentRun] | None = None,
|
||||
):
|
||||
"""Initialize the WorkflowBuilder.
|
||||
|
||||
@@ -166,15 +90,15 @@ class WorkflowBuilder:
|
||||
max_iterations: Maximum number of iterations for workflow convergence. Default is 100.
|
||||
name: Optional human-readable name for the workflow.
|
||||
description: Optional description of what the workflow does.
|
||||
start_executor: The starting executor for the workflow. Can be an Executor instance,
|
||||
SupportsAgentRun instance, or the name of a registered executor factory.
|
||||
start_executor: The starting executor for the workflow. Can be an Executor instance
|
||||
or SupportsAgentRun instance.
|
||||
checkpoint_storage: Optional checkpoint storage for enabling workflow state persistence.
|
||||
output_executors: Optional list of executors whose outputs should be collected.
|
||||
If not provided, outputs from all executors are collected.
|
||||
"""
|
||||
self._edge_groups: list[EdgeGroup] = []
|
||||
self._executors: dict[str, Executor] = {}
|
||||
self._start_executor: Executor | str | None = None
|
||||
self._start_executor: Executor | None = None
|
||||
self._checkpoint_storage: CheckpointStorage | None = checkpoint_storage
|
||||
self._max_iterations: int = max_iterations
|
||||
self._name: str | None = name
|
||||
@@ -184,18 +108,8 @@ class WorkflowBuilder:
|
||||
# being created for the same agent.
|
||||
self._agent_wrappers: dict[str, Executor] = {}
|
||||
|
||||
# Registrations for lazy initialization of executors
|
||||
self._edge_registry: list[
|
||||
_EdgeRegistration
|
||||
| _FanOutEdgeRegistration
|
||||
| _SwitchCaseEdgeGroupRegistration
|
||||
| _MultiSelectionEdgeGroupRegistration
|
||||
| _FanInEdgeRegistration
|
||||
] = []
|
||||
self._executor_registry: dict[str, Callable[[], Executor]] = {}
|
||||
|
||||
# Output executors filter; if set, only outputs from these executors are yielded
|
||||
self._output_executors: list[Executor | SupportsAgentRun | str] = output_executors if output_executors else []
|
||||
self._output_executors: list[Executor | SupportsAgentRun] = output_executors if output_executors else []
|
||||
|
||||
# Set the start executor
|
||||
self._set_start_executor(start_executor)
|
||||
@@ -258,133 +172,10 @@ class WorkflowBuilder:
|
||||
f"WorkflowBuilder expected an Executor or SupportsAgentRun instance; got {type(candidate).__name__}."
|
||||
)
|
||||
|
||||
def register_executor(self, factory_func: Callable[[], Executor], name: str | list[str]) -> Self:
|
||||
"""Register an executor factory function for lazy initialization.
|
||||
|
||||
This method allows you to register a factory function that creates an executor.
|
||||
The executor will be instantiated only when the workflow is built, enabling
|
||||
deferred initialization and potentially reducing startup time.
|
||||
|
||||
Args:
|
||||
factory_func: A callable that returns an Executor instance when called.
|
||||
name: The name(s) of the registered executor factory. This doesn't have to match
|
||||
the executor's ID, but it must be unique within the workflow.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
from typing_extensions import Never
|
||||
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
|
||||
|
||||
|
||||
class UpperCaseExecutor(Executor):
|
||||
@handler
|
||||
async def process(self, text: str, ctx: WorkflowContext[str]) -> None:
|
||||
await ctx.send_message(text.upper())
|
||||
|
||||
|
||||
class ReverseExecutor(Executor):
|
||||
@handler
|
||||
async def process(self, text: str, ctx: WorkflowContext[Never, str]) -> None:
|
||||
await ctx.yield_output(text[::-1])
|
||||
|
||||
|
||||
# Build a workflow
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor="UpperCase")
|
||||
.register_executor(lambda: UpperCaseExecutor(id="upper"), name="UpperCase")
|
||||
.register_executor(lambda: ReverseExecutor(id="reverse"), name="Reverse")
|
||||
.add_edge("UpperCase", "Reverse")
|
||||
.build()
|
||||
)
|
||||
|
||||
If multiple names are provided, the same factory function will be registered under each name.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import WorkflowBuilder, Executor, WorkflowContext, handler
|
||||
|
||||
|
||||
class LoggerExecutor(Executor):
|
||||
@handler
|
||||
async def log(self, message: str, ctx: WorkflowContext) -> None:
|
||||
print(f"Log: {message}")
|
||||
|
||||
|
||||
# Register the same executor factory under multiple names
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor="ExecutorA")
|
||||
.register_executor(lambda: LoggerExecutor(id="logger"), name=["ExecutorA", "ExecutorB"])
|
||||
.add_edge("ExecutorA", "ExecutorB")
|
||||
.build()
|
||||
"""
|
||||
names = [name] if isinstance(name, str) else name
|
||||
|
||||
for n in names:
|
||||
if n in self._executor_registry:
|
||||
raise ValueError(f"An executor factory with the name '{n}' is already registered.")
|
||||
|
||||
for n in names:
|
||||
self._executor_registry[n] = factory_func
|
||||
|
||||
return self
|
||||
|
||||
def register_agent(
|
||||
self,
|
||||
factory_func: Callable[[], SupportsAgentRun],
|
||||
name: str,
|
||||
agent_thread: AgentThread | None = None,
|
||||
) -> Self:
|
||||
"""Register an agent factory function for lazy initialization.
|
||||
|
||||
This method allows you to register a factory function that creates an agent.
|
||||
The agent will be instantiated and wrapped in an AgentExecutor only when the workflow is built,
|
||||
enabling deferred initialization and potentially reducing startup time.
|
||||
|
||||
Args:
|
||||
factory_func: A callable that returns an SupportsAgentRun instance when called.
|
||||
name: The name of the registered agent factory. This doesn't have to match
|
||||
the agent's internal name. But it must be unique within the workflow.
|
||||
agent_thread: The thread to use for running the agent. If None, a new thread will be created when
|
||||
the agent is instantiated.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
from agent_framework import WorkflowBuilder
|
||||
from agent_framework_anthropic import AnthropicAgent
|
||||
|
||||
|
||||
# Build a workflow
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor="SomeOtherExecutor")
|
||||
.register_executor(lambda: ..., name="SomeOtherExecutor")
|
||||
.register_agent(
|
||||
lambda: AnthropicAgent(name="writer", model="claude-3-5-sonnet-20241022"),
|
||||
name="WriterAgent",
|
||||
output_response=True,
|
||||
)
|
||||
.add_edge("SomeOtherExecutor", "WriterAgent")
|
||||
.build()
|
||||
)
|
||||
"""
|
||||
if name in self._executor_registry:
|
||||
raise ValueError(f"An agent factory with the name '{name}' is already registered.")
|
||||
|
||||
def wrapped_factory() -> AgentExecutor:
|
||||
agent = factory_func()
|
||||
return AgentExecutor(
|
||||
agent,
|
||||
agent_thread=agent_thread,
|
||||
)
|
||||
|
||||
self._executor_registry[name] = wrapped_factory
|
||||
|
||||
return self
|
||||
|
||||
def add_edge(
|
||||
self,
|
||||
source: Executor | SupportsAgentRun | str,
|
||||
target: Executor | SupportsAgentRun | str,
|
||||
source: Executor | SupportsAgentRun,
|
||||
target: Executor | SupportsAgentRun,
|
||||
condition: EdgeCondition | None = None,
|
||||
) -> Self:
|
||||
"""Add a directed edge between two executors.
|
||||
@@ -393,17 +184,12 @@ class WorkflowBuilder:
|
||||
Messages sent by the source executor will be routed to the target executor.
|
||||
|
||||
Args:
|
||||
source: The source executor or registered name of the source factory for the edge.
|
||||
target: The target executor or registered name of the target factory for the edge.
|
||||
source: The source executor or agent for the edge.
|
||||
target: The target executor or agent for the edge.
|
||||
condition: An optional condition function `(data) -> bool | Awaitable[bool]`
|
||||
that determines whether the edge should be traversed.
|
||||
Example: `lambda data: data["ready"]`.
|
||||
|
||||
Note: If instances are provided for both source and target, they will be shared across
|
||||
all workflow instances created from the built Workflow. To avoid this, consider
|
||||
registering the executors and agents using `register_executor` and `register_agent`
|
||||
and referencing them by factory name for lazy initialization instead.
|
||||
|
||||
Returns:
|
||||
Self: The WorkflowBuilder instance for method chaining.
|
||||
|
||||
@@ -426,39 +212,13 @@ class WorkflowBuilder:
|
||||
await ctx.yield_output(f"Processed {count} characters")
|
||||
|
||||
|
||||
# Connect executors with an edge
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor="ProcessorA")
|
||||
.register_executor(lambda: ProcessorA(id="a"), name="ProcessorA")
|
||||
.register_executor(lambda: ProcessorB(id="b"), name="ProcessorB")
|
||||
.add_edge("ProcessorA", "ProcessorB")
|
||||
.build()
|
||||
)
|
||||
a = ProcessorA(id="a")
|
||||
b = ProcessorB(id="b")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor="ProcessorA")
|
||||
.register_executor(lambda: ProcessorA(id="a"), name="ProcessorA")
|
||||
.register_executor(lambda: ProcessorB(id="b"), name="ProcessorB")
|
||||
.add_edge("ProcessorA", "ProcessorB", condition=only_large_numbers)
|
||||
.build()
|
||||
)
|
||||
workflow = WorkflowBuilder(start_executor=a).add_edge(a, b).build()
|
||||
"""
|
||||
if (isinstance(source, str) and not isinstance(target, str)) or (
|
||||
not isinstance(source, str) and isinstance(target, str)
|
||||
):
|
||||
raise ValueError(
|
||||
"Both source and target must be either registered factory names (str) or "
|
||||
"Executor/SupportsAgentRun instances."
|
||||
)
|
||||
|
||||
if isinstance(source, str) and isinstance(target, str):
|
||||
# Both are names; defer resolution to build time
|
||||
self._edge_registry.append(_EdgeRegistration(source=source, target=target, condition=condition))
|
||||
return self
|
||||
|
||||
# Both are Executor/SupportsAgentRun instances; wrap and add now
|
||||
source_exec = self._maybe_wrap_agent(source) # type: ignore[arg-type]
|
||||
target_exec = self._maybe_wrap_agent(target) # type: ignore[arg-type]
|
||||
source_exec = self._maybe_wrap_agent(source)
|
||||
target_exec = self._maybe_wrap_agent(target)
|
||||
source_id = self._add_executor(source_exec)
|
||||
target_id = self._add_executor(target_exec)
|
||||
self._edge_groups.append(SingleEdgeGroup(source_id, target_id, condition))
|
||||
@@ -466,8 +226,8 @@ class WorkflowBuilder:
|
||||
|
||||
def add_fan_out_edges(
|
||||
self,
|
||||
source: Executor | SupportsAgentRun | str,
|
||||
targets: Sequence[Executor | SupportsAgentRun | str],
|
||||
source: Executor | SupportsAgentRun,
|
||||
targets: Sequence[Executor | SupportsAgentRun],
|
||||
) -> Self:
|
||||
"""Add multiple edges to the workflow where messages from the source will be sent to all targets.
|
||||
|
||||
@@ -475,17 +235,12 @@ class WorkflowBuilder:
|
||||
Messages from the source will be broadcast to all target executors concurrently.
|
||||
|
||||
Args:
|
||||
source: The source executor or registered name of the source factory for the edges.
|
||||
targets: A list of target executors or registered names of the target factories for the edges.
|
||||
source: The source executor or agent for the edges.
|
||||
targets: A list of target executors or agents for the edges.
|
||||
|
||||
Returns:
|
||||
Self: The WorkflowBuilder instance for method chaining.
|
||||
|
||||
Note: If instances are provided for source and targets, they will be shared across
|
||||
all workflow instances created from the built Workflow. To avoid this, consider
|
||||
registering the executors and agents using `register_executor` and `register_agent`
|
||||
and referencing them by factory name for lazy initialization instead.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
@@ -511,32 +266,14 @@ class WorkflowBuilder:
|
||||
print(f"ValidatorB: {data}")
|
||||
|
||||
|
||||
# Broadcast to multiple validators
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor="DataSource")
|
||||
.register_executor(lambda: DataSource(id="source"), name="DataSource")
|
||||
.register_executor(lambda: ValidatorA(id="val_a"), name="ValidatorA")
|
||||
.register_executor(lambda: ValidatorB(id="val_b"), name="ValidatorB")
|
||||
.add_fan_out_edges("DataSource", ["ValidatorA", "ValidatorB"])
|
||||
.build()
|
||||
)
|
||||
source = DataSource(id="source")
|
||||
val_a = ValidatorA(id="val_a")
|
||||
val_b = ValidatorB(id="val_b")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=source).add_fan_out_edges(source, [val_a, val_b]).build()
|
||||
"""
|
||||
if (isinstance(source, str) and not all(isinstance(t, str) for t in targets)) or (
|
||||
not isinstance(source, str) and any(isinstance(t, str) for t in targets)
|
||||
):
|
||||
raise ValueError(
|
||||
"Both source and targets must be either registered factory names (str) or "
|
||||
"Executor/SupportsAgentRun instances."
|
||||
)
|
||||
|
||||
if isinstance(source, str) and all(isinstance(t, str) for t in targets):
|
||||
# Both are names; defer resolution to build time
|
||||
self._edge_registry.append(_FanOutEdgeRegistration(source=source, targets=list(targets))) # type: ignore
|
||||
return self
|
||||
|
||||
# Both are Executor/SupportsAgentRun instances; wrap and add now
|
||||
source_exec = self._maybe_wrap_agent(source) # type: ignore[arg-type]
|
||||
target_execs = [self._maybe_wrap_agent(t) for t in targets] # type: ignore[arg-type]
|
||||
source_exec = self._maybe_wrap_agent(source)
|
||||
target_execs = [self._maybe_wrap_agent(t) for t in targets]
|
||||
source_id = self._add_executor(source_exec)
|
||||
target_ids = [self._add_executor(t) for t in target_execs]
|
||||
self._edge_groups.append(FanOutEdgeGroup(source_id, target_ids)) # type: ignore[call-arg]
|
||||
@@ -545,7 +282,7 @@ class WorkflowBuilder:
|
||||
|
||||
def add_switch_case_edge_group(
|
||||
self,
|
||||
source: Executor | SupportsAgentRun | str,
|
||||
source: Executor | SupportsAgentRun,
|
||||
cases: Sequence[Case | Default],
|
||||
) -> Self:
|
||||
"""Add an edge group that represents a switch-case statement.
|
||||
@@ -562,17 +299,12 @@ class WorkflowBuilder:
|
||||
(i.e., no condition matched).
|
||||
|
||||
Args:
|
||||
source: The source executor or registered name of the source factory for the edge group.
|
||||
source: The source executor or agent for the edge group.
|
||||
cases: A list of case objects that determine the target executor for each message.
|
||||
|
||||
Returns:
|
||||
Self: The WorkflowBuilder instance for method chaining.
|
||||
|
||||
Note: If instances are provided for source and case targets, they will be shared across
|
||||
all workflow instances created from the built Workflow. To avoid this, consider
|
||||
registering the executors and agents using `register_executor` and `register_agent`
|
||||
and referencing them by factory name for lazy initialization instead.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
@@ -603,37 +335,23 @@ class WorkflowBuilder:
|
||||
print(f"Low score: {result.score}")
|
||||
|
||||
|
||||
# Route based on score value
|
||||
evaluator = Evaluator(id="eval")
|
||||
high = HighScoreHandler(id="high")
|
||||
low = LowScoreHandler(id="low")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor="Evaluator")
|
||||
.register_executor(lambda: Evaluator(id="eval"), name="Evaluator")
|
||||
.register_executor(lambda: HighScoreHandler(id="high"), name="HighScoreHandler")
|
||||
.register_executor(lambda: LowScoreHandler(id="low"), name="LowScoreHandler")
|
||||
WorkflowBuilder(start_executor=evaluator)
|
||||
.add_switch_case_edge_group(
|
||||
"Evaluator",
|
||||
evaluator,
|
||||
[
|
||||
Case(condition=lambda r: r.score > 10, target="HighScoreHandler"),
|
||||
Default(target="LowScoreHandler"),
|
||||
Case(condition=lambda r: r.score > 10, target=high),
|
||||
Default(target=low),
|
||||
],
|
||||
)
|
||||
.build()
|
||||
)
|
||||
"""
|
||||
if (isinstance(source, str) and not all(isinstance(case.target, str) for case in cases)) or (
|
||||
not isinstance(source, str) and any(isinstance(case.target, str) for case in cases)
|
||||
):
|
||||
raise ValueError(
|
||||
"Both source and case targets must be either registered factory names (str) "
|
||||
"or Executor/SupportsAgentRun instances."
|
||||
)
|
||||
|
||||
if isinstance(source, str) and all(isinstance(case.target, str) for case in cases):
|
||||
# Source is a name; defer resolution to build time
|
||||
self._edge_registry.append(_SwitchCaseEdgeGroupRegistration(source=source, cases=list(cases))) # type: ignore
|
||||
return self
|
||||
|
||||
# Source is an Executor/SupportsAgentRun instance; wrap and add now
|
||||
source_exec = self._maybe_wrap_agent(source) # type: ignore[arg-type]
|
||||
source_exec = self._maybe_wrap_agent(source)
|
||||
source_id = self._add_executor(source_exec)
|
||||
# Convert case data types to internal types that only uses target_id.
|
||||
internal_cases: list[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault] = []
|
||||
@@ -651,8 +369,8 @@ class WorkflowBuilder:
|
||||
|
||||
def add_multi_selection_edge_group(
|
||||
self,
|
||||
source: Executor | SupportsAgentRun | str,
|
||||
targets: Sequence[Executor | SupportsAgentRun | str],
|
||||
source: Executor | SupportsAgentRun,
|
||||
targets: Sequence[Executor | SupportsAgentRun],
|
||||
selection_func: Callable[[Any, list[str]], list[str]],
|
||||
) -> Self:
|
||||
"""Add an edge group that represents a multi-selection execution model.
|
||||
@@ -665,19 +383,14 @@ class WorkflowBuilder:
|
||||
and return a list of executor IDs indicating which target executors should receive the message.
|
||||
|
||||
Args:
|
||||
source: The source executor or registered name of the source factory for the edge group.
|
||||
targets: A list of target executors or registered names of the target factories for the edges.
|
||||
source: The source executor or agent for the edge group.
|
||||
targets: A list of target executors or agents for the edges.
|
||||
selection_func: A function that selects target executors for messages.
|
||||
Takes (message, list[executor_id]) and returns list[executor_id].
|
||||
|
||||
Returns:
|
||||
Self: The WorkflowBuilder instance for method chaining.
|
||||
|
||||
Note: If instances are provided for source and targets, they will be shared across
|
||||
all workflow instances created from the built Workflow. To avoid this, consider
|
||||
registering the executors and agents using `register_executor` and `register_agent`
|
||||
and referencing them by factory name for lazy initialization instead.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
@@ -710,6 +423,11 @@ class WorkflowBuilder:
|
||||
print(f"WorkerB processing: {task.data}")
|
||||
|
||||
|
||||
dispatcher = TaskDispatcher(id="dispatcher")
|
||||
worker_a = WorkerA(id="worker_a")
|
||||
worker_b = WorkerB(id="worker_b")
|
||||
|
||||
|
||||
# Select workers based on task priority
|
||||
def select_workers(task: Task, available: list[str]) -> list[str]:
|
||||
if task.priority == "high":
|
||||
@@ -718,40 +436,17 @@ class WorkflowBuilder:
|
||||
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor="TaskDispatcher")
|
||||
.register_executor(lambda: TaskDispatcher(id="dispatcher"), name="TaskDispatcher")
|
||||
.register_executor(lambda: WorkerA(id="worker_a"), name="WorkerA")
|
||||
.register_executor(lambda: WorkerB(id="worker_b"), name="WorkerB")
|
||||
WorkflowBuilder(start_executor=dispatcher)
|
||||
.add_multi_selection_edge_group(
|
||||
"TaskDispatcher",
|
||||
["WorkerA", "WorkerB"],
|
||||
dispatcher,
|
||||
[worker_a, worker_b],
|
||||
selection_func=select_workers,
|
||||
)
|
||||
.build()
|
||||
)
|
||||
"""
|
||||
if (isinstance(source, str) and not all(isinstance(t, str) for t in targets)) or (
|
||||
not isinstance(source, str) and any(isinstance(t, str) for t in targets)
|
||||
):
|
||||
raise ValueError(
|
||||
"Both source and targets must be either registered factory names (str) or "
|
||||
"Executor/SupportsAgentRun instances."
|
||||
)
|
||||
|
||||
if isinstance(source, str) and all(isinstance(t, str) for t in targets):
|
||||
# Both are names; defer resolution to build time
|
||||
self._edge_registry.append(
|
||||
_MultiSelectionEdgeGroupRegistration(
|
||||
source=source,
|
||||
targets=list(targets), # type: ignore
|
||||
selection_func=selection_func,
|
||||
)
|
||||
)
|
||||
return self
|
||||
|
||||
# Both are Executor/SupportsAgentRun instances; wrap and add now
|
||||
source_exec = self._maybe_wrap_agent(source) # type: ignore
|
||||
target_execs = [self._maybe_wrap_agent(t) for t in targets] # type: ignore
|
||||
source_exec = self._maybe_wrap_agent(source)
|
||||
target_execs = [self._maybe_wrap_agent(t) for t in targets]
|
||||
source_id = self._add_executor(source_exec)
|
||||
target_ids = [self._add_executor(t) for t in target_execs]
|
||||
self._edge_groups.append(FanOutEdgeGroup(source_id, target_ids, selection_func)) # type: ignore[call-arg]
|
||||
@@ -760,8 +455,8 @@ class WorkflowBuilder:
|
||||
|
||||
def add_fan_in_edges(
|
||||
self,
|
||||
sources: Sequence[Executor | SupportsAgentRun | str],
|
||||
target: Executor | SupportsAgentRun | str,
|
||||
sources: Sequence[Executor | SupportsAgentRun],
|
||||
target: Executor | SupportsAgentRun,
|
||||
) -> Self:
|
||||
"""Add multiple edges from sources to a single target executor.
|
||||
|
||||
@@ -773,17 +468,12 @@ class WorkflowBuilder:
|
||||
types of the source executors.
|
||||
|
||||
Args:
|
||||
sources: A list of source executors or registered names of the source factories for the edges.
|
||||
target: The target executor or registered name of the target factory for the edges.
|
||||
sources: A list of source executors or agents for the edges.
|
||||
target: The target executor or agent for the edges.
|
||||
|
||||
Returns:
|
||||
Self: The WorkflowBuilder instance for method chaining.
|
||||
|
||||
Note: If instances are provided for sources and target, they will be shared across
|
||||
all workflow instances created from the built Workflow. To avoid this, consider
|
||||
registering the executors and agents using `register_executor` and `register_agent`
|
||||
and referencing them by factory name for lazy initialization instead.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
@@ -804,39 +494,21 @@ class WorkflowBuilder:
|
||||
await ctx.yield_output(f"Combined: {combined}")
|
||||
|
||||
|
||||
# Collect results from multiple producers
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor="Producer1")
|
||||
.register_executor(lambda: Producer(id="prod_1"), name="Producer1")
|
||||
.register_executor(lambda: Producer(id="prod_2"), name="Producer2")
|
||||
.register_executor(lambda: Aggregator(id="agg"), name="Aggregator")
|
||||
.add_fan_in_edges(["Producer1", "Producer2"], "Aggregator")
|
||||
.build()
|
||||
)
|
||||
prod_1 = Producer(id="prod_1")
|
||||
prod_2 = Producer(id="prod_2")
|
||||
agg = Aggregator(id="agg")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=prod_1).add_fan_in_edges([prod_1, prod_2], agg).build()
|
||||
"""
|
||||
if (all(isinstance(s, str) for s in sources) and not isinstance(target, str)) or (
|
||||
not all(isinstance(s, str) for s in sources) and isinstance(target, str)
|
||||
):
|
||||
raise ValueError(
|
||||
"Both sources and target must be either registered factory names (str) or "
|
||||
"Executor/SupportsAgentRun instances."
|
||||
)
|
||||
|
||||
if all(isinstance(s, str) for s in sources) and isinstance(target, str):
|
||||
# Both are names; defer resolution to build time
|
||||
self._edge_registry.append(_FanInEdgeRegistration(sources=list(sources), target=target)) # type: ignore
|
||||
return self
|
||||
|
||||
# Both are Executor/SupportsAgentRun instances; wrap and add now
|
||||
source_execs = [self._maybe_wrap_agent(s) for s in sources] # type: ignore
|
||||
target_exec = self._maybe_wrap_agent(target) # type: ignore
|
||||
source_execs = [self._maybe_wrap_agent(s) for s in sources]
|
||||
target_exec = self._maybe_wrap_agent(target)
|
||||
source_ids = [self._add_executor(s) for s in source_execs]
|
||||
target_id = self._add_executor(target_exec)
|
||||
self._edge_groups.append(FanInEdgeGroup(source_ids, target_id)) # type: ignore[call-arg]
|
||||
|
||||
return self
|
||||
|
||||
def add_chain(self, executors: Sequence[Executor | SupportsAgentRun | str]) -> Self:
|
||||
def add_chain(self, executors: Sequence[Executor | SupportsAgentRun]) -> Self:
|
||||
"""Add a chain of executors to the workflow.
|
||||
|
||||
The output of each executor in the chain will be sent to the next executor in the chain.
|
||||
@@ -845,16 +517,11 @@ class WorkflowBuilder:
|
||||
Cycles in the chain are not allowed, meaning an executor cannot appear more than once in the chain.
|
||||
|
||||
Args:
|
||||
executors: A list of executors or registered names of the executor factories to chain together.
|
||||
executors: A list of executors or agents to chain together.
|
||||
|
||||
Returns:
|
||||
Self: The WorkflowBuilder instance for method chaining.
|
||||
|
||||
Note: If executor instances are provided, they will be shared across all workflow instances created
|
||||
from the built Workflow. To avoid this, consider registering the executors and agents using
|
||||
`register_executor` and `register_agent` and referencing them by factory name for lazy
|
||||
initialization instead.
|
||||
|
||||
Example:
|
||||
.. code-block:: python
|
||||
|
||||
@@ -880,148 +547,37 @@ class WorkflowBuilder:
|
||||
await ctx.yield_output(f"Final: {text}")
|
||||
|
||||
|
||||
# Chain executors in sequence
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor="step1")
|
||||
.register_executor(lambda: Step1(id="step1"), name="step1")
|
||||
.register_executor(lambda: Step2(id="step2"), name="step2")
|
||||
.register_executor(lambda: Step3(id="step3"), name="step3")
|
||||
.add_chain(["step1", "step2", "step3"])
|
||||
.build()
|
||||
)
|
||||
step1 = Step1(id="step1")
|
||||
step2 = Step2(id="step2")
|
||||
step3 = Step3(id="step3")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=step1).add_chain([step1, step2, step3]).build()
|
||||
"""
|
||||
if len(executors) < 2:
|
||||
raise ValueError("At least two executors are required to form a chain.")
|
||||
|
||||
if not all(isinstance(e, str) for e in executors) and any(isinstance(e, str) for e in executors):
|
||||
raise ValueError(
|
||||
"All executors in the chain must be either registered factory names (str) "
|
||||
"or Executor/SupportsAgentRun instances."
|
||||
)
|
||||
|
||||
if all(isinstance(e, str) for e in executors):
|
||||
# All are names; defer resolution to build time
|
||||
for i in range(len(executors) - 1):
|
||||
self.add_edge(executors[i], executors[i + 1])
|
||||
return self
|
||||
|
||||
# All are Executor/SupportsAgentRun instances; wrap and add now
|
||||
# Wrap each candidate first to ensure stable IDs before adding edges
|
||||
wrapped: list[Executor] = [self._maybe_wrap_agent(e) for e in executors] # type: ignore[arg-type]
|
||||
wrapped: list[Executor] = [self._maybe_wrap_agent(e) for e in executors]
|
||||
for i in range(len(wrapped) - 1):
|
||||
self.add_edge(wrapped[i], wrapped[i + 1])
|
||||
return self
|
||||
|
||||
def _set_start_executor(self, executor: Executor | SupportsAgentRun | str) -> None:
|
||||
def _set_start_executor(self, executor: Executor | SupportsAgentRun) -> None:
|
||||
"""Set the starting executor for the workflow (internal method).
|
||||
|
||||
Args:
|
||||
executor: The starting executor, which can be an Executor instance, SupportsAgentRun instance,
|
||||
or the name of a registered executor factory.
|
||||
executor: The starting executor, which can be an Executor instance or SupportsAgentRun instance.
|
||||
"""
|
||||
if self._start_executor is not None:
|
||||
start_id = self._start_executor if isinstance(self._start_executor, str) else self._start_executor.id
|
||||
logger.warning(f"Overwriting existing start executor: {start_id} for the workflow.")
|
||||
logger.warning(f"Overwriting existing start executor: {self._start_executor.id} for the workflow.")
|
||||
|
||||
if isinstance(executor, str):
|
||||
self._start_executor = executor
|
||||
else:
|
||||
wrapped = self._maybe_wrap_agent(executor) # type: ignore[arg-type]
|
||||
self._start_executor = wrapped
|
||||
# Ensure the start executor is present in the executor map so validation succeeds
|
||||
# even if no edges are added yet, or before edges wrap the same agent again.
|
||||
existing = self._executors.get(wrapped.id)
|
||||
if existing is not wrapped:
|
||||
self._add_executor(wrapped)
|
||||
|
||||
# Removed explicit set_agent_streaming() API; agents always stream updates.
|
||||
|
||||
def _resolve_edge_registry(self) -> tuple[Executor, dict[str, Executor], list[EdgeGroup]]:
|
||||
"""Resolve deferred edge registrations into executors and edge groups.
|
||||
|
||||
Returns:
|
||||
tuple: A tuple containing:
|
||||
- The starting Executor instance.
|
||||
- A dictionary mapping registered factory names to resolved Executor instances.
|
||||
- A list of EdgeGroup instances representing the workflow edges composed of resolved executors.
|
||||
|
||||
Notes:
|
||||
Non-factory executors (i.e., those added directly) are not included in the returned list,
|
||||
as they are already part of the workflow builder's internal state.
|
||||
"""
|
||||
if not self._start_executor:
|
||||
raise ValueError(
|
||||
"Starting executor must be set via the start_executor constructor parameter before building."
|
||||
)
|
||||
|
||||
start_executor: Executor | None = None
|
||||
if isinstance(self._start_executor, Executor):
|
||||
start_executor = self._start_executor
|
||||
|
||||
# Maps registered factory names to created executor instances for edge resolution
|
||||
factory_name_to_instance: dict[str, Executor] = {}
|
||||
# Maps executor IDs to created executor instances to prevent duplicates
|
||||
executor_id_to_instance: dict[str, Executor] = {}
|
||||
deferred_edge_groups: list[EdgeGroup] = []
|
||||
for name, exec_factory in self._executor_registry.items():
|
||||
instance = exec_factory()
|
||||
if instance.id in executor_id_to_instance:
|
||||
raise ValueError(f"Executor with ID '{instance.id}' has already been registered.")
|
||||
if instance.id in self._executors:
|
||||
raise ValueError(f"Executor ID collision: An executor with ID '{instance.id}' already exists.")
|
||||
executor_id_to_instance[instance.id] = instance
|
||||
|
||||
if isinstance(self._start_executor, str) and name == self._start_executor:
|
||||
start_executor = instance
|
||||
|
||||
# All executors will get their own internal edge group for receiving system messages
|
||||
deferred_edge_groups.append(InternalEdgeGroup(instance.id)) # type: ignore[call-arg]
|
||||
factory_name_to_instance[name] = instance
|
||||
|
||||
def _get_executor(name: str) -> Executor:
|
||||
"""Helper to get executor by the registered name. Raises if not found."""
|
||||
if name not in factory_name_to_instance:
|
||||
raise ValueError(f"Factory '{name}' has not been registered.")
|
||||
return factory_name_to_instance[name]
|
||||
|
||||
for registration in self._edge_registry:
|
||||
match registration:
|
||||
case _EdgeRegistration(source, target, condition):
|
||||
source_exec: Executor = _get_executor(source)
|
||||
target_exec: Executor = _get_executor(target)
|
||||
deferred_edge_groups.append(SingleEdgeGroup(source_exec.id, target_exec.id, condition)) # type: ignore[call-arg]
|
||||
case _FanOutEdgeRegistration(source, targets):
|
||||
source_exec = _get_executor(source)
|
||||
target_execs = [_get_executor(t) for t in targets]
|
||||
deferred_edge_groups.append(FanOutEdgeGroup(source_exec.id, [t.id for t in target_execs])) # type: ignore[call-arg]
|
||||
case _SwitchCaseEdgeGroupRegistration(source, cases):
|
||||
source_exec = _get_executor(source)
|
||||
cases_converted: list[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault] = []
|
||||
for case in cases:
|
||||
if not isinstance(case.target, str):
|
||||
raise ValueError("Switch case target must be a registered factory name (str) if deferred.")
|
||||
target_exec = _get_executor(case.target)
|
||||
if isinstance(case, Default):
|
||||
cases_converted.append(SwitchCaseEdgeGroupDefault(target_id=target_exec.id))
|
||||
else:
|
||||
cases_converted.append(
|
||||
SwitchCaseEdgeGroupCase(condition=case.condition, target_id=target_exec.id)
|
||||
)
|
||||
deferred_edge_groups.append(SwitchCaseEdgeGroup(source_exec.id, cases_converted)) # type: ignore[call-arg]
|
||||
case _MultiSelectionEdgeGroupRegistration(source, targets, selection_func):
|
||||
source_exec = _get_executor(source)
|
||||
target_execs = [_get_executor(t) for t in targets]
|
||||
deferred_edge_groups.append(
|
||||
FanOutEdgeGroup(source_exec.id, [t.id for t in target_execs], selection_func) # type: ignore[call-arg]
|
||||
)
|
||||
case _FanInEdgeRegistration(sources, target):
|
||||
source_execs = [_get_executor(s) for s in sources]
|
||||
target_exec = _get_executor(target)
|
||||
deferred_edge_groups.append(FanInEdgeGroup([s.id for s in source_execs], target_exec.id)) # type: ignore[call-arg]
|
||||
if start_executor is None:
|
||||
raise ValueError("Failed to resolve starting executor from registered factories.")
|
||||
|
||||
return (start_executor, factory_name_to_instance, deferred_edge_groups)
|
||||
wrapped = self._maybe_wrap_agent(executor)
|
||||
self._start_executor = wrapped
|
||||
# Ensure the start executor is present in the executor map so validation succeeds
|
||||
# even if no edges are added yet, or before edges wrap the same agent again.
|
||||
existing = self._executors.get(wrapped.id)
|
||||
if existing is not wrapped:
|
||||
self._add_executor(wrapped)
|
||||
|
||||
def build(self) -> Workflow:
|
||||
"""Build and return the constructed workflow.
|
||||
@@ -1053,12 +609,9 @@ class WorkflowBuilder:
|
||||
await ctx.yield_output(text.upper())
|
||||
|
||||
|
||||
# Build and execute a workflow
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor="MyExecutor")
|
||||
.register_executor(lambda: MyExecutor(id="executor"), name="MyExecutor")
|
||||
.build()
|
||||
)
|
||||
executor = MyExecutor(id="executor")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=executor).build()
|
||||
|
||||
# The workflow is now immutable and ready to run
|
||||
events = await workflow.run("hello")
|
||||
@@ -1074,23 +627,17 @@ class WorkflowBuilder:
|
||||
# Add workflow build started event
|
||||
span.add_event(OtelAttr.BUILD_STARTED)
|
||||
|
||||
# Resolve lazy edge registrations
|
||||
start_executor, deferred_executors, deferred_edge_groups = self._resolve_edge_registry()
|
||||
executors = self._executors | {exe.id: exe for exe in deferred_executors.values()}
|
||||
edge_groups = self._edge_groups + deferred_edge_groups
|
||||
output_executors = (
|
||||
[
|
||||
deferred_executors[factory_name].id
|
||||
for factory_name in self._output_executors
|
||||
if isinstance(factory_name, str)
|
||||
]
|
||||
+ [ex.id for ex in self._output_executors if isinstance(ex, Executor)]
|
||||
+ [
|
||||
resolve_agent_id(agent)
|
||||
for agent in self._output_executors
|
||||
if isinstance(agent, SupportsAgentRun)
|
||||
]
|
||||
)
|
||||
if not self._start_executor:
|
||||
raise ValueError(
|
||||
"Starting executor must be set via the start_executor constructor parameter before building."
|
||||
)
|
||||
|
||||
start_executor = self._start_executor
|
||||
executors = self._executors
|
||||
edge_groups = self._edge_groups
|
||||
output_executors = [ex.id for ex in self._output_executors if isinstance(ex, Executor)] + [
|
||||
resolve_agent_id(agent) for agent in self._output_executors if isinstance(agent, SupportsAgentRun)
|
||||
]
|
||||
|
||||
# Perform validation before creating the workflow
|
||||
validate_workflow_graph(
|
||||
|
||||
@@ -670,16 +670,20 @@ class TestWorkflowAgent:
|
||||
return ResponseStream(_iter(), finalizer=AgentResponse.from_updates)
|
||||
|
||||
@executor
|
||||
async def start_executor(messages: list[ChatMessage], ctx: WorkflowContext[AgentExecutorRequest, str]) -> None:
|
||||
async def start_exec(messages: list[ChatMessage], ctx: WorkflowContext[AgentExecutorRequest, str]) -> None:
|
||||
await ctx.yield_output("Start output")
|
||||
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
|
||||
|
||||
# Build workflow: start -> agent1 (no output) -> agent2 (output_response=True)
|
||||
builder = WorkflowBuilder(start_executor="start", output_executors=["start", "agent2"])
|
||||
builder.register_executor(lambda: start_executor, "start")
|
||||
builder.register_agent(lambda: MockAgent("agent1", "Agent1 output - should NOT appear"), "agent1")
|
||||
builder.register_agent(lambda: MockAgent("agent2", "Agent2 output - SHOULD appear"), "agent2")
|
||||
workflow = builder.add_edge("start", "agent1").add_edge("agent1", "agent2").build()
|
||||
agent1 = MockAgent("agent1", "Agent1 output - should NOT appear")
|
||||
agent2 = MockAgent("agent2", "Agent2 output - SHOULD appear")
|
||||
|
||||
# Build workflow: start -> agent1 (no output) -> agent2 (output visible)
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor=start_exec, output_executors=[start_exec, agent2])
|
||||
.add_edge(start_exec, agent1)
|
||||
.add_edge(agent1, agent2)
|
||||
.build()
|
||||
)
|
||||
|
||||
agent = WorkflowAgent(workflow=workflow, name="Test Agent")
|
||||
result = await agent.run("Test input")
|
||||
@@ -754,17 +758,13 @@ class TestWorkflowAgent:
|
||||
return ResponseStream(_iter(), finalizer=AgentResponse.from_updates)
|
||||
|
||||
@executor
|
||||
async def start_executor(messages: list[ChatMessage], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
|
||||
async def start_exec(messages: list[ChatMessage], ctx: WorkflowContext[AgentExecutorRequest]) -> None:
|
||||
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
|
||||
|
||||
mock_agent = MockAgent("agent", "Unique response text")
|
||||
|
||||
# Build workflow with single agent
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor="start")
|
||||
.register_executor(lambda: start_executor, "start")
|
||||
.register_agent(lambda: MockAgent("agent", "Unique response text"), "agent")
|
||||
.add_edge("start", "agent")
|
||||
.build()
|
||||
)
|
||||
workflow = WorkflowBuilder(start_executor=start_exec).add_edge(start_exec, mock_agent).build()
|
||||
|
||||
agent = WorkflowAgent(workflow=workflow, name="Test Agent")
|
||||
result = await agent.run("Test input")
|
||||
|
||||
@@ -134,322 +134,65 @@ def test_add_agent_duplicate_id_raises_error():
|
||||
builder.add_edge(agent1, agent2).build()
|
||||
|
||||
|
||||
# Tests for new executor registration patterns
|
||||
def test_fan_out_edges_with_direct_instances():
|
||||
"""Test fan-out edges with direct executor instances."""
|
||||
source = MockExecutor(id="Source")
|
||||
target1 = MockExecutor(id="Target1")
|
||||
target2 = MockExecutor(id="Target2")
|
||||
|
||||
workflow = WorkflowBuilder(start_executor=source).add_fan_out_edges(source, [target1, target2]).build()
|
||||
|
||||
def test_register_executor_basic():
|
||||
"""Test basic executor registration with lazy initialization."""
|
||||
builder = WorkflowBuilder(start_executor="TestExecutor")
|
||||
|
||||
# Register an executor factory - ID must match the registered name
|
||||
result = builder.register_executor(lambda: MockExecutor(id="TestExecutor"), name="TestExecutor")
|
||||
|
||||
# Verify that register returns the builder for chaining
|
||||
assert result is builder
|
||||
|
||||
# Build workflow and verify executor is instantiated
|
||||
workflow = builder.build()
|
||||
assert "TestExecutor" in workflow.executors
|
||||
assert isinstance(workflow.executors["TestExecutor"], MockExecutor)
|
||||
|
||||
|
||||
def test_register_multiple_executors():
|
||||
"""Test registering multiple executors and connecting them with edges."""
|
||||
builder = WorkflowBuilder(start_executor="ExecutorA")
|
||||
|
||||
# Register multiple executors - IDs must match registered names
|
||||
builder.register_executor(lambda: MockExecutor(id="ExecutorA"), name="ExecutorA")
|
||||
builder.register_executor(lambda: MockExecutor(id="ExecutorB"), name="ExecutorB")
|
||||
builder.register_executor(lambda: MockExecutor(id="ExecutorC"), name="ExecutorC")
|
||||
|
||||
# Build workflow with edges using registered names
|
||||
workflow = builder.add_edge("ExecutorA", "ExecutorB").add_edge("ExecutorB", "ExecutorC").build()
|
||||
|
||||
# Verify all executors are present
|
||||
assert "ExecutorA" in workflow.executors
|
||||
assert "ExecutorB" in workflow.executors
|
||||
assert "ExecutorC" in workflow.executors
|
||||
assert workflow.start_executor_id == "ExecutorA"
|
||||
|
||||
|
||||
def test_register_with_multiple_names():
|
||||
"""Test registering the same factory function under multiple names."""
|
||||
builder = WorkflowBuilder(start_executor="ExecutorA")
|
||||
|
||||
# Register same executor factory under multiple names
|
||||
# Note: Each call creates a new instance, so IDs won't conflict
|
||||
counter = {"val": 0}
|
||||
|
||||
def make_executor():
|
||||
counter["val"] += 1
|
||||
return MockExecutor(id="ExecutorA" if counter["val"] == 1 else "ExecutorB")
|
||||
|
||||
builder.register_executor(make_executor, name=["ExecutorA", "ExecutorB"])
|
||||
|
||||
# Set up workflow
|
||||
workflow = builder.add_edge("ExecutorA", "ExecutorB").build()
|
||||
|
||||
# Verify both executors are present
|
||||
assert "ExecutorA" in workflow.executors
|
||||
assert "ExecutorB" in workflow.executors
|
||||
assert workflow.start_executor_id == "ExecutorA"
|
||||
|
||||
|
||||
def test_register_duplicate_name_raises_error():
|
||||
"""Test that registering duplicate names raises an error."""
|
||||
builder = WorkflowBuilder(start_executor="MyExecutor")
|
||||
|
||||
# Register first executor
|
||||
builder.register_executor(lambda: MockExecutor(id="executor_1"), name="MyExecutor")
|
||||
|
||||
# Registering second executor with same name should raise ValueError
|
||||
with pytest.raises(ValueError, match="already registered"):
|
||||
builder.register_executor(lambda: MockExecutor(id="executor_2"), name="MyExecutor")
|
||||
|
||||
|
||||
def test_register_duplicate_id_raises_error():
|
||||
"""Test that registering duplicate id raises an error."""
|
||||
builder = WorkflowBuilder(start_executor="MyExecutor1")
|
||||
|
||||
# Register first executor
|
||||
builder.register_executor(lambda: MockExecutor(id="executor"), name="MyExecutor1")
|
||||
builder.register_executor(lambda: MockExecutor(id="executor"), name="MyExecutor2")
|
||||
|
||||
# Registering second executor with same ID should raise ValueError
|
||||
with pytest.raises(ValueError, match="Executor with ID 'executor' has already been registered."):
|
||||
builder.build()
|
||||
|
||||
|
||||
def test_register_agent_basic():
|
||||
"""Test basic agent registration with lazy initialization."""
|
||||
builder = WorkflowBuilder(start_executor="TestAgent")
|
||||
|
||||
# Register an agent factory
|
||||
result = builder.register_agent(lambda: DummyAgent(id="agent_test", name="test_agent"), name="TestAgent")
|
||||
|
||||
# Verify that register_agent returns the builder for chaining
|
||||
assert result is builder
|
||||
|
||||
# Build workflow and verify agent is wrapped in AgentExecutor
|
||||
workflow = builder.build()
|
||||
assert "test_agent" in workflow.executors
|
||||
assert isinstance(workflow.executors["test_agent"], AgentExecutor)
|
||||
|
||||
|
||||
def test_register_agent_with_thread():
|
||||
"""Test registering an agent with a custom thread."""
|
||||
builder = WorkflowBuilder(start_executor="ThreadedAgent")
|
||||
custom_thread = AgentThread()
|
||||
|
||||
# Register agent with custom thread
|
||||
builder.register_agent(
|
||||
lambda: DummyAgent(id="agent_with_thread", name="threaded_agent"),
|
||||
name="ThreadedAgent",
|
||||
agent_thread=custom_thread,
|
||||
)
|
||||
|
||||
# Build workflow and verify agent executor configuration
|
||||
workflow = builder.build()
|
||||
executor = workflow.executors["threaded_agent"]
|
||||
|
||||
assert isinstance(executor, AgentExecutor)
|
||||
assert executor.id == "threaded_agent"
|
||||
assert executor._agent_thread is custom_thread # type: ignore
|
||||
|
||||
|
||||
def test_register_agent_duplicate_name_raises_error():
|
||||
"""Test that registering agents with duplicate names raises an error."""
|
||||
builder = WorkflowBuilder(start_executor="MyAgent")
|
||||
|
||||
# Register first agent
|
||||
builder.register_agent(lambda: DummyAgent(id="agent1", name="first"), name="MyAgent")
|
||||
|
||||
# Registering second agent with same name should raise ValueError
|
||||
with pytest.raises(ValueError, match="already registered"):
|
||||
builder.register_agent(lambda: DummyAgent(id="agent2", name="second"), name="MyAgent")
|
||||
|
||||
|
||||
def test_register_and_add_edge_with_strings():
|
||||
"""Test that registered executors can be connected using string names."""
|
||||
builder = WorkflowBuilder(start_executor="Source")
|
||||
|
||||
# Register executors
|
||||
builder.register_executor(lambda: MockExecutor(id="source"), name="Source")
|
||||
builder.register_executor(lambda: MockExecutor(id="target"), name="Target")
|
||||
|
||||
# Add edge using string names
|
||||
workflow = builder.add_edge("Source", "Target").build()
|
||||
|
||||
# Verify edge is created correctly
|
||||
assert workflow.start_executor_id == "source"
|
||||
assert "source" in workflow.executors
|
||||
assert "target" in workflow.executors
|
||||
|
||||
|
||||
def test_register_agent_and_add_edge_with_strings():
|
||||
"""Test that registered agents can be connected using string names."""
|
||||
builder = WorkflowBuilder(start_executor="Writer")
|
||||
|
||||
# Register agents
|
||||
builder.register_agent(lambda: DummyAgent(id="writer_id", name="writer"), name="Writer")
|
||||
builder.register_agent(lambda: DummyAgent(id="reviewer_id", name="reviewer"), name="Reviewer")
|
||||
|
||||
# Add edge using string names
|
||||
workflow = builder.add_edge("Writer", "Reviewer").build()
|
||||
|
||||
# Verify edge is created correctly
|
||||
assert workflow.start_executor_id == "writer"
|
||||
assert "writer" in workflow.executors
|
||||
assert "reviewer" in workflow.executors
|
||||
assert all(isinstance(e, AgentExecutor) for e in workflow.executors.values())
|
||||
|
||||
|
||||
def test_register_with_fan_out_edges():
|
||||
"""Test using registered names with fan-out edge groups."""
|
||||
builder = WorkflowBuilder(start_executor="Source")
|
||||
|
||||
# Register executors - IDs must match registered names
|
||||
builder.register_executor(lambda: MockExecutor(id="Source"), name="Source")
|
||||
builder.register_executor(lambda: MockExecutor(id="Target1"), name="Target1")
|
||||
builder.register_executor(lambda: MockExecutor(id="Target2"), name="Target2")
|
||||
|
||||
# Add fan-out edges using registered names
|
||||
workflow = builder.add_fan_out_edges("Source", ["Target1", "Target2"]).build()
|
||||
|
||||
# Verify all executors are present
|
||||
assert "Source" in workflow.executors
|
||||
assert "Target1" in workflow.executors
|
||||
assert "Target2" in workflow.executors
|
||||
|
||||
|
||||
def test_register_with_fan_in_edges():
|
||||
"""Test using registered names with fan-in edge groups."""
|
||||
builder = WorkflowBuilder(start_executor="Source1")
|
||||
def test_fan_in_edges_with_direct_instances():
|
||||
"""Test fan-in edges with direct executor instances."""
|
||||
source1 = MockExecutor(id="Source1")
|
||||
source2 = MockExecutor(id="Source2")
|
||||
aggregator = MockAggregator(id="Aggregator")
|
||||
|
||||
# Register executors - IDs must match registered names
|
||||
builder.register_executor(lambda: MockExecutor(id="Source1"), name="Source1")
|
||||
builder.register_executor(lambda: MockExecutor(id="Source2"), name="Source2")
|
||||
builder.register_executor(lambda: MockAggregator(id="Aggregator"), name="Aggregator")
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor=source1)
|
||||
.add_edge(source1, source2)
|
||||
.add_fan_in_edges([source1, source2], aggregator)
|
||||
.build()
|
||||
)
|
||||
|
||||
# Add fan-in edges using registered names
|
||||
# Both Source1 and Source2 need to be reachable, so connect Source1 to Source2
|
||||
workflow = builder.add_edge("Source1", "Source2").add_fan_in_edges(["Source1", "Source2"], "Aggregator").build()
|
||||
|
||||
# Verify all executors are present
|
||||
assert "Source1" in workflow.executors
|
||||
assert "Source2" in workflow.executors
|
||||
assert "Aggregator" in workflow.executors
|
||||
|
||||
|
||||
def test_register_with_chain():
|
||||
"""Test using registered names with add_chain."""
|
||||
builder = WorkflowBuilder(start_executor="Step1")
|
||||
def test_chain_with_direct_instances():
|
||||
"""Test add_chain with direct executor instances."""
|
||||
step1 = MockExecutor(id="Step1")
|
||||
step2 = MockExecutor(id="Step2")
|
||||
step3 = MockExecutor(id="Step3")
|
||||
|
||||
# Register executors - IDs must match registered names
|
||||
builder.register_executor(lambda: MockExecutor(id="Step1"), name="Step1")
|
||||
builder.register_executor(lambda: MockExecutor(id="Step2"), name="Step2")
|
||||
builder.register_executor(lambda: MockExecutor(id="Step3"), name="Step3")
|
||||
workflow = WorkflowBuilder(start_executor=step1).add_chain([step1, step2, step3]).build()
|
||||
|
||||
# Add chain using registered names
|
||||
workflow = builder.add_chain(["Step1", "Step2", "Step3"]).build()
|
||||
|
||||
# Verify all executors are present
|
||||
assert "Step1" in workflow.executors
|
||||
assert "Step2" in workflow.executors
|
||||
assert "Step3" in workflow.executors
|
||||
assert workflow.start_executor_id == "Step1"
|
||||
|
||||
|
||||
def test_register_factory_called_only_once():
|
||||
"""Test that registered factory functions are called only during build."""
|
||||
call_count = 0
|
||||
|
||||
def factory():
|
||||
nonlocal call_count
|
||||
call_count += 1
|
||||
return MockExecutor(id="Test")
|
||||
|
||||
builder = WorkflowBuilder(start_executor="Test")
|
||||
builder.register_executor(factory, name="Test")
|
||||
|
||||
# Factory should not be called yet
|
||||
assert call_count == 0
|
||||
|
||||
# Factory should still not be called
|
||||
assert call_count == 0
|
||||
|
||||
# Build workflow
|
||||
workflow = builder.build()
|
||||
|
||||
# Factory should now be called exactly once
|
||||
assert call_count == 1
|
||||
assert "Test" in workflow.executors
|
||||
|
||||
|
||||
def test_mixing_eager_and_lazy_initialization_error():
|
||||
"""Test that mixing eager executor instances with lazy string names raises appropriate error."""
|
||||
builder = WorkflowBuilder(start_executor="Lazy")
|
||||
|
||||
# Create an eager executor instance
|
||||
eager_executor = MockExecutor(id="eager")
|
||||
|
||||
# Register a lazy executor
|
||||
builder.register_executor(lambda: MockExecutor(id="Lazy"), name="Lazy")
|
||||
|
||||
# Mixing eager and lazy should raise an error during add_edge
|
||||
with pytest.raises(
|
||||
ValueError,
|
||||
match=(
|
||||
r"Both source and target must be either registered factory names \(str\) "
|
||||
r"or Executor/SupportsAgentRun instances\."
|
||||
),
|
||||
):
|
||||
builder.add_edge(eager_executor, "Lazy")
|
||||
|
||||
|
||||
def test_register_with_condition():
|
||||
"""Test adding edges with conditions using registered names."""
|
||||
builder = WorkflowBuilder(start_executor="Source")
|
||||
def test_add_edge_with_condition():
|
||||
"""Test adding edges with conditions using direct executor instances."""
|
||||
source = MockExecutor(id="Source")
|
||||
target = MockExecutor(id="Target")
|
||||
|
||||
def condition_func(msg: MockMessage) -> bool:
|
||||
return msg.data > 0
|
||||
|
||||
# Register executors - IDs must match registered names
|
||||
builder.register_executor(lambda: MockExecutor(id="Source"), name="Source")
|
||||
builder.register_executor(lambda: MockExecutor(id="Target"), name="Target")
|
||||
workflow = WorkflowBuilder(start_executor=source).add_edge(source, target, condition=condition_func).build()
|
||||
|
||||
# Add edge with condition
|
||||
workflow = builder.add_edge("Source", "Target", condition=condition_func).build()
|
||||
|
||||
# Verify workflow is built correctly
|
||||
assert "Source" in workflow.executors
|
||||
assert "Target" in workflow.executors
|
||||
|
||||
|
||||
def test_register_agent_creates_unique_instances():
|
||||
"""Test that registered agent factories create new instances on each build."""
|
||||
instance_ids: list[int] = []
|
||||
|
||||
def agent_factory() -> DummyAgent:
|
||||
agent = DummyAgent(id=f"agent_{len(instance_ids)}", name="test")
|
||||
instance_ids.append(id(agent))
|
||||
return agent
|
||||
|
||||
# Build first workflow
|
||||
builder1 = WorkflowBuilder(start_executor="Agent")
|
||||
builder1.register_agent(agent_factory, name="Agent")
|
||||
_ = builder1.build()
|
||||
|
||||
# Build second workflow
|
||||
builder2 = WorkflowBuilder(start_executor="Agent")
|
||||
builder2.register_agent(agent_factory, name="Agent")
|
||||
_ = builder2.build()
|
||||
|
||||
# Verify that two different agent instances were created
|
||||
assert len(instance_ids) == 2
|
||||
assert instance_ids[0] != instance_ids[1]
|
||||
|
||||
|
||||
# region with_output_from tests
|
||||
|
||||
|
||||
@@ -488,14 +231,17 @@ def test_with_output_from_with_agent_instances():
|
||||
assert workflow._output_executors == ["reviewer"] # type: ignore
|
||||
|
||||
|
||||
def test_with_output_from_with_registered_names():
|
||||
"""Test with_output_from with registered factory names (strings)."""
|
||||
builder = WorkflowBuilder(start_executor="ExecutorAFactory", output_executors=["ExecutorBFactory"])
|
||||
builder.register_executor(lambda: MockExecutor(id="ExecutorA"), name="ExecutorAFactory")
|
||||
builder.register_executor(lambda: MockExecutor(id="ExecutorB"), name="ExecutorBFactory")
|
||||
workflow = builder.add_edge("ExecutorAFactory", "ExecutorBFactory").build()
|
||||
def test_with_output_from_with_executor_instances_by_id():
|
||||
"""Test with_output_from with direct executor instances resolves to executor IDs."""
|
||||
executor_a = MockExecutor(id="ExecutorA")
|
||||
executor_b = MockExecutor(id="ExecutorB")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor=executor_a, output_executors=[executor_b])
|
||||
.add_edge(executor_a, executor_b)
|
||||
.build()
|
||||
)
|
||||
|
||||
# Verify that the workflow was built with the correct output executors
|
||||
assert workflow._output_executors == ["ExecutorB"] # type: ignore
|
||||
|
||||
|
||||
@@ -531,14 +277,17 @@ def test_with_output_from_can_be_set_to_different_value():
|
||||
assert workflow._output_executors == ["executor_b"] # type: ignore
|
||||
|
||||
|
||||
def test_with_output_from_with_registered_agents():
|
||||
"""Test with_output_from with registered agent factory names."""
|
||||
builder = WorkflowBuilder(start_executor="WriterAgent", output_executors=["ReviewerAgent"])
|
||||
builder.register_agent(lambda: DummyAgent(id="agent1", name="writer"), name="WriterAgent")
|
||||
builder.register_agent(lambda: DummyAgent(id="agent2", name="reviewer"), name="ReviewerAgent")
|
||||
workflow = builder.add_edge("WriterAgent", "ReviewerAgent").build()
|
||||
def test_with_output_from_with_agent_instances_resolves_name():
|
||||
"""Test with_output_from with agent instances resolves to agent names."""
|
||||
agent_writer = DummyAgent(id="agent1", name="writer")
|
||||
agent_reviewer = DummyAgent(id="agent2", name="reviewer")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder(start_executor=agent_writer, output_executors=[agent_reviewer])
|
||||
.add_edge(agent_writer, agent_reviewer)
|
||||
.build()
|
||||
)
|
||||
|
||||
# Verify that the workflow was built with the agent's resolved name
|
||||
assert workflow._output_executors == ["reviewer"] # type: ignore
|
||||
|
||||
|
||||
|
||||
@@ -474,8 +474,9 @@ async def test_message_trace_context_serialization(span_exporter: InMemorySpanEx
|
||||
async def test_workflow_build_error_tracing(span_exporter: InMemorySpanExporter) -> None:
|
||||
"""Test that build errors are properly recorded in build spans."""
|
||||
|
||||
# Test validation error by referencing a non-existent start executor
|
||||
builder = WorkflowBuilder(start_executor="NonExistent")
|
||||
# Create a valid builder, then clear the start executor to trigger a build-time ValueError
|
||||
builder = WorkflowBuilder(start_executor=MockExecutor(id="mock"))
|
||||
builder._start_executor = None # type: ignore[assignment]
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
builder.build()
|
||||
|
||||
Reference in New Issue
Block a user