agent/agent-framework
2
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

[BREAKING] Python: Move single-config fluent methods to constructor parameters (#3693)

Browse Source 
* Move single-config fluent methods to constructor parameters

* Updates

* Adjust magentic and group chat
This commit is contained in:
Evan Mattson
2026-02-07 15:01:52 +09:00
committed by GitHub
Unverified
parent 5d355ac507
commit 74ac470a56
132 changed files with 1341 additions and 2386 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/packages/core/agent_framework/_workflows/_workflow.py
+2 -2
View File
@@ -152,7 +152,7 @@ class Workflow(DictConvertible):
Checkpointing can be configured at build time or runtime:
Build-time (via WorkflowBuilder):
workflow = WorkflowBuilder().with_checkpointing(storage).build()
workflow = WorkflowBuilder(checkpoint_storage=storage).build()
Runtime (via run parameters):
result = await workflow.run(message, checkpoint_storage=runtime_storage)
@@ -428,7 +428,7 @@ class Workflow(DictConvertible):
if not has_checkpointing and checkpoint_storage is None:
raise ValueError(
"Cannot restore from checkpoint: either provide checkpoint_storage parameter "
"or build workflow with WorkflowBuilder.with_checkpointing(checkpoint_storage)."
"or build workflow with WorkflowBuilder(checkpoint_storage=checkpoint_storage)."
)
await self._runner.restore_from_checkpoint(checkpoint_id, checkpoint_storage)
python/packages/core/agent_framework/_workflows/_workflow_builder.py
+33 -179
View File
@@ -138,11 +138,10 @@ class WorkflowBuilder:
# Build a workflow
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor="UpperCase")
.register_executor(lambda: UpperCaseExecutor(id="upper"), name="UpperCase")
.register_executor(lambda: ReverseExecutor(id="reverse"), name="Reverse")
.add_edge("UpperCase", "Reverse")
.set_start_executor("UpperCase")
.build()
)
@@ -156,23 +155,32 @@ class WorkflowBuilder:
max_iterations: int = DEFAULT_MAX_ITERATIONS,
name: str | None = None,
description: str | None = None,
*,
start_executor: Executor | SupportsAgentRun | str,
checkpoint_storage: CheckpointStorage | None = None,
output_executors: list[Executor | SupportsAgentRun | str] | None = None,
):
"""Initialize the WorkflowBuilder with an empty list of edges and no starting executor.
"""Initialize the WorkflowBuilder.
Args:
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.
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._checkpoint_storage: CheckpointStorage | None = None
self._checkpoint_storage: CheckpointStorage | None = checkpoint_storage
self._max_iterations: int = max_iterations
self._name: str | None = name
self._description: str | None = description
# Maps underlying SupportsAgentRun object id -> wrapped Executor so we reuse the same wrapper
# across set_start_executor / add_edge calls. This avoids multiple AgentExecutor instances
# across start_executor / add_edge calls. This avoids multiple AgentExecutor instances
# being created for the same agent.
self._agent_wrappers: dict[str, Executor] = {}
@@ -187,7 +195,10 @@ class WorkflowBuilder:
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] = []
self._output_executors: list[Executor | SupportsAgentRun | str] = output_executors if output_executors else []
# Set the start executor
self._set_start_executor(start_executor)
# Agents auto-wrapped by builder now always stream incremental updates.
@@ -279,10 +290,9 @@ class WorkflowBuilder:
# Build a workflow
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor="UpperCase")
.register_executor(lambda: UpperCaseExecutor(id="upper"), name="UpperCase")
.register_executor(lambda: ReverseExecutor(id="reverse"), name="Reverse")
.set_start_executor("UpperCase")
.add_edge("UpperCase", "Reverse")
.build()
)
@@ -302,9 +312,8 @@ class WorkflowBuilder:
# Register the same executor factory under multiple names
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor="ExecutorA")
.register_executor(lambda: LoggerExecutor(id="logger"), name=["ExecutorA", "ExecutorB"])
.set_start_executor("ExecutorA")
.add_edge("ExecutorA", "ExecutorB")
.build()
"""
@@ -347,7 +356,7 @@ class WorkflowBuilder:
# Build a workflow
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor="SomeOtherExecutor")
.register_executor(lambda: ..., name="SomeOtherExecutor")
.register_agent(
lambda: AnthropicAgent(name="writer", model="claude-3-5-sonnet-20241022"),
@@ -355,7 +364,6 @@ class WorkflowBuilder:
output_response=True,
)
.add_edge("SomeOtherExecutor", "WriterAgent")
.set_start_executor("SomeOtherExecutor")
.build()
)
"""
@@ -420,20 +428,18 @@ class WorkflowBuilder:
# Connect executors with an edge
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor="ProcessorA")
.register_executor(lambda: ProcessorA(id="a"), name="ProcessorA")
.register_executor(lambda: ProcessorB(id="b"), name="ProcessorB")
.add_edge("ProcessorA", "ProcessorB")
.set_start_executor("ProcessorA")
.build()
)
workflow = (
WorkflowBuilder()
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)
.set_start_executor("ProcessorA")
.build()
)
"""
@@ -507,12 +513,11 @@ class WorkflowBuilder:
# Broadcast to multiple validators
workflow = (
WorkflowBuilder()
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"])
.set_start_executor("DataSource")
.build()
)
"""
@@ -600,7 +605,7 @@ class WorkflowBuilder:
# Route based on score value
workflow = (
WorkflowBuilder()
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")
@@ -611,7 +616,6 @@ class WorkflowBuilder:
Default(target="LowScoreHandler"),
],
)
.set_start_executor("Evaluator")
.build()
)
"""
@@ -714,7 +718,7 @@ class WorkflowBuilder:
workflow = (
WorkflowBuilder()
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")
@@ -723,7 +727,6 @@ class WorkflowBuilder:
["WorkerA", "WorkerB"],
selection_func=select_workers,
)
.set_start_executor("TaskDispatcher")
.build()
)
"""
@@ -803,12 +806,11 @@ class WorkflowBuilder:
# Collect results from multiple producers
workflow = (
WorkflowBuilder()
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")
.set_start_executor("Producer1")
.build()
)
"""
@@ -880,12 +882,11 @@ class WorkflowBuilder:
# Chain executors in sequence
workflow = (
WorkflowBuilder()
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"])
.set_start_executor("step1")
.build()
)
"""
@@ -911,46 +912,12 @@ class WorkflowBuilder:
self.add_edge(wrapped[i], wrapped[i + 1])
return self
def set_start_executor(self, executor: Executor | SupportsAgentRun | str) -> Self:
"""Set the starting executor for the workflow.
The start executor is the entry point for the workflow. When the workflow is executed,
the initial message will be sent to this executor.
def _set_start_executor(self, executor: Executor | SupportsAgentRun | str) -> 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.
Returns:
Self: The WorkflowBuilder instance for method chaining.
Example:
.. code-block:: python
from typing_extensions import Never
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
class EntryPoint(Executor):
@handler
async def process(self, text: str, ctx: WorkflowContext[str]) -> None:
await ctx.send_message(text.upper())
class Processor(Executor):
@handler
async def process(self, text: str, ctx: WorkflowContext[Never, str]) -> None:
await ctx.yield_output(text)
workflow = (
WorkflowBuilder()
.register_executor(lambda: EntryPoint(id="entry"), name="EntryPoint")
.register_executor(lambda: Processor(id="proc"), name="Processor")
.add_edge("EntryPoint", "Processor")
.set_start_executor("EntryPoint")
.build()
)
"""
if self._start_executor is not None:
start_id = self._start_executor if isinstance(self._start_executor, str) else self._start_executor.id
@@ -966,123 +933,9 @@ class WorkflowBuilder:
existing = self._executors.get(wrapped.id)
if existing is not wrapped:
self._add_executor(wrapped)
return self
def set_max_iterations(self, max_iterations: int) -> Self:
"""Set the maximum number of iterations for the workflow.
When a workflow contains cycles, this limit prevents infinite loops by capping
the total number of executor invocations. The default is 100 iterations.
Args:
max_iterations: The maximum number of iterations the workflow will run for convergence.
Returns:
Self: The WorkflowBuilder instance for method chaining.
Example:
.. code-block:: python
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
class StepA(Executor):
@handler
async def process(self, count: int, ctx: WorkflowContext[int]) -> None:
if count < 10:
await ctx.send_message(count + 1)
class StepB(Executor):
@handler
async def process(self, count: int, ctx: WorkflowContext[int]) -> None:
await ctx.send_message(count)
# Set a custom iteration limit for workflow with cycles
workflow = (
WorkflowBuilder()
.set_max_iterations(500)
.register_executor(lambda: StepA(id="step_a"), name="StepA")
.register_executor(lambda: StepB(id="step_b"), name="StepB")
.add_edge("StepA", "StepB")
.add_edge("StepB", "StepA") # Cycle
.set_start_executor("StepA")
.build()
)
"""
self._max_iterations = max_iterations
return self
# Removed explicit set_agent_streaming() API; agents always stream updates.
def with_checkpointing(self, checkpoint_storage: CheckpointStorage) -> Self:
"""Enable checkpointing with the specified storage.
Checkpointing allows workflows to save their state periodically, enabling
pause/resume functionality and recovery from failures. The checkpoint storage
implementation determines where checkpoints are persisted.
Args:
checkpoint_storage: The checkpoint storage implementation to use.
Returns:
Self: The WorkflowBuilder instance for method chaining.
Example:
.. code-block:: python
from typing_extensions import Never
from agent_framework import Executor, WorkflowBuilder, WorkflowContext, handler
from agent_framework import FileCheckpointStorage
class ProcessorA(Executor):
@handler
async def process(self, text: str, ctx: WorkflowContext[str]) -> None:
await ctx.send_message(text.upper())
class ProcessorB(Executor):
@handler
async def process(self, text: str, ctx: WorkflowContext[Never, str]) -> None:
await ctx.yield_output(text)
# Enable checkpointing with file-based storage
storage = FileCheckpointStorage("./checkpoints")
workflow = (
WorkflowBuilder()
.register_executor(lambda: ProcessorA(id="proc_a"), name="ProcessorA")
.register_executor(lambda: ProcessorB(id="proc_b"), name="ProcessorB")
.add_edge("ProcessorA", "ProcessorB")
.set_start_executor("ProcessorA")
.with_checkpointing(storage)
.build()
)
# Run with checkpoint saving
events = await workflow.run("input")
"""
self._checkpoint_storage = checkpoint_storage
return self
def with_output_from(self, executors: list[Executor | SupportsAgentRun | str]) -> Self:
"""Specify which executors' outputs should be collected as workflow outputs.
By default, outputs from all executors are collected. This method allows
filtering to only include outputs from specified executors.
Args:
executors: A list of executors or registered names of the executor factories
whose outputs should be collected.
Returns:
Self: The WorkflowBuilder instance for method chaining.
"""
self._output_executors = list(executors)
return self
def _resolve_edge_registry(self) -> tuple[Executor, dict[str, Executor], list[EdgeGroup]]:
"""Resolve deferred edge registrations into executors and edge groups.
@@ -1097,7 +950,9 @@ class WorkflowBuilder:
as they are already part of the workflow builder's internal state.
"""
if not self._start_executor:
raise ValueError("Starting executor must be set using set_start_executor before building the workflow.")
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):
@@ -1200,9 +1055,8 @@ class WorkflowBuilder:
# Build and execute a workflow
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor="MyExecutor")
.register_executor(lambda: MyExecutor(id="executor"), name="MyExecutor")
.set_start_executor("MyExecutor")
.build()
)
python/packages/core/tests/workflow/test_agent_executor.py
+2 -2
View File
@@ -70,7 +70,7 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
executor = AgentExecutor(initial_agent, agent_thread=initial_thread)
# Build workflow with checkpointing enabled
wf = SequentialBuilder().participants([executor]).with_checkpointing(storage).build()
wf = SequentialBuilder(participants=[executor], checkpoint_storage=storage).build()
# Run the workflow with a user message
first_run_output: AgentExecutorResponse | None = None
@@ -124,7 +124,7 @@ async def test_agent_executor_checkpoint_stores_and_restores_state() -> None:
assert restored_agent.call_count == 0
# Build new workflow with the restored executor
wf_resume = SequentialBuilder().participants([restored_executor]).with_checkpointing(storage).build()
wf_resume = SequentialBuilder(participants=[restored_executor], checkpoint_storage=storage).build()
# Resume from checkpoint
resumed_output: AgentExecutorResponse | None = None
python/packages/core/tests/workflow/test_agent_executor_tool_calls.py
+5 -13
View File
@@ -96,7 +96,7 @@ async def test_agent_executor_emits_tool_calls_in_streaming_mode() -> None:
agent = _ToolCallingAgent(id="tool_agent", name="ToolAgent")
agent_exec = AgentExecutor(agent, id="tool_exec")
workflow = WorkflowBuilder().set_start_executor(agent_exec).build()
workflow = WorkflowBuilder(start_executor=agent_exec).build()
# Act: run in streaming mode
events: list[WorkflowEvent[AgentResponseUpdate]] = []
@@ -249,11 +249,7 @@ async def test_agent_executor_tool_call_with_approval() -> None:
)
workflow = (
WorkflowBuilder()
.set_start_executor(agent)
.add_edge(agent, test_executor)
.with_output_from([test_executor])
.build()
WorkflowBuilder(start_executor=agent, output_executors=[test_executor]).add_edge(agent, test_executor).build()
)
# Act
@@ -286,7 +282,7 @@ async def test_agent_executor_tool_call_with_approval_streaming() -> None:
tools=[mock_tool_requiring_approval],
)
workflow = WorkflowBuilder().set_start_executor(agent).add_edge(agent, test_executor).build()
workflow = WorkflowBuilder(start_executor=agent).add_edge(agent, test_executor).build()
# Act
request_info_events: list[WorkflowEvent] = []
@@ -324,11 +320,7 @@ async def test_agent_executor_parallel_tool_call_with_approval() -> None:
)
workflow = (
WorkflowBuilder()
.set_start_executor(agent)
.add_edge(agent, test_executor)
.with_output_from([test_executor])
.build()
WorkflowBuilder(start_executor=agent, output_executors=[test_executor]).add_edge(agent, test_executor).build()
)
# Act
@@ -363,7 +355,7 @@ async def test_agent_executor_parallel_tool_call_with_approval_streaming() -> No
tools=[mock_tool_requiring_approval],
)
workflow = WorkflowBuilder().set_start_executor(agent).add_edge(agent, test_executor).build()
workflow = WorkflowBuilder(start_executor=agent).add_edge(agent, test_executor).build()
# Act
request_info_events: list[WorkflowEvent] = []
python/packages/core/tests/workflow/test_checkpoint_validation.py
+3 -2
View File
@@ -30,8 +30,9 @@ def build_workflow(storage: InMemoryCheckpointStorage, finish_id: str = "finish"
start = StartExecutor(id="start")
finish = FinishExecutor(id=finish_id)
builder = WorkflowBuilder(max_iterations=3).set_start_executor(start).add_edge(start, finish)
builder = builder.with_checkpointing(checkpoint_storage=storage)
builder = WorkflowBuilder(max_iterations=3, start_executor=start, checkpoint_storage=storage).add_edge(
start, finish
)
return builder.build()
python/packages/core/tests/workflow/test_executor.py
+5 -5
View File
@@ -153,7 +153,7 @@ async def test_executor_invoked_event_contains_input_data():
upper = UpperCaseExecutor(id="upper")
collector = CollectorExecutor(id="collector")
workflow = WorkflowBuilder().add_edge(upper, collector).set_start_executor(upper).build()
workflow = WorkflowBuilder(start_executor=upper).add_edge(upper, collector).build()
events = await workflow.run("hello world")
invoked_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_invoked"]
@@ -190,7 +190,7 @@ async def test_executor_completed_event_contains_sent_messages():
sender = MultiSenderExecutor(id="sender")
collector = CollectorExecutor(id="collector")
workflow = WorkflowBuilder().add_edge(sender, collector).set_start_executor(sender).build()
workflow = WorkflowBuilder(start_executor=sender).add_edge(sender, collector).build()
events = await workflow.run("hello")
completed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
@@ -217,7 +217,7 @@ async def test_executor_completed_event_includes_yielded_outputs():
await ctx.yield_output(text.upper())
executor = YieldOnlyExecutor(id="yielder")
workflow = WorkflowBuilder().set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).build()
events = await workflow.run("test")
completed_events = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "executor_completed"]
@@ -260,7 +260,7 @@ async def test_executor_events_with_complex_message_types():
processor = ProcessorExecutor(id="processor")
collector = CollectorExecutor(id="collector")
workflow = WorkflowBuilder().add_edge(processor, collector).set_start_executor(processor).build()
workflow = WorkflowBuilder(start_executor=processor).add_edge(processor, collector).build()
input_request = Request(query="hello", limit=3)
events = await workflow.run(input_request)
@@ -539,7 +539,7 @@ async def test_executor_invoked_event_data_not_mutated_by_handler():
# Verify mutation happened
assert len(messages) == original_len + 1
workflow = WorkflowBuilder().set_start_executor(mutator).build()
workflow = WorkflowBuilder(start_executor=mutator).build()
# Run with a single user message
input_messages = [ChatMessage(role="user", text="hello")]
python/packages/core/tests/workflow/test_full_conversation.py
+2 -8
View File
@@ -76,13 +76,7 @@ async def test_agent_executor_populates_full_conversation_non_streaming() -> Non
agent_exec = AgentExecutor(agent, id="agent1-exec")
capturer = _CaptureFullConversation(id="capture")
wf = (
WorkflowBuilder()
.set_start_executor(agent_exec)
.add_edge(agent_exec, capturer)
.with_output_from([capturer])
.build()
)
wf = WorkflowBuilder(start_executor=agent_exec, output_executors=[capturer]).add_edge(agent_exec, capturer).build()
# Act: use run() to test non-streaming mode
result = await wf.run("hello world")
@@ -144,7 +138,7 @@ async def test_sequential_adapter_uses_full_conversation() -> None:
a1 = _CaptureAgent(id="agent1", name="A1", reply_text="A1 reply")
a2 = _CaptureAgent(id="agent2", name="A2", reply_text="A2 reply")
wf = SequentialBuilder().participants([a1, a2]).build()
wf = SequentialBuilder(participants=[a1, a2]).build()
# Act
async for ev in wf.run("hello seq", stream=True):
python/packages/core/tests/workflow/test_function_executor.py
+2 -2
View File
@@ -236,7 +236,7 @@ class TestFunctionExecutor:
assert reverse_spec["output_types"] == [Any] # First parameter is Any
assert reverse_spec["workflow_output_types"] == [str] # Second parameter is str
workflow = WorkflowBuilder().add_edge(to_upper, reverse_text).set_start_executor(to_upper).build()
workflow = WorkflowBuilder(start_executor=to_upper).add_edge(to_upper, reverse_text).build()
# Run workflow
events = await workflow.run("hello world")
@@ -345,7 +345,7 @@ class TestFunctionExecutor:
# Since single-parameter functions can't send messages,
# they're typically used as terminal nodes or for side effects
WorkflowBuilder().set_start_executor(double_value).build()
WorkflowBuilder(start_executor=double_value).build()
# For testing purposes, we can check that the handler is registered correctly
assert double_value.can_handle(Message(data=5, source_id="mock"))
python/packages/core/tests/workflow/test_request_info_and_response.py
+8 -8
View File
@@ -178,7 +178,7 @@ class TestRequestInfoAndResponse:
async def test_approval_workflow(self):
"""Test end-to-end workflow with approval request."""
executor = ApprovalRequiredExecutor(id="approval_executor")
workflow = WorkflowBuilder().set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).build()
# First run the workflow until it emits a request
request_info_event: WorkflowEvent | None = None
@@ -203,7 +203,7 @@ class TestRequestInfoAndResponse:
async def test_calculation_workflow(self):
"""Test end-to-end workflow with calculation request."""
executor = CalculationExecutor(id="calc_executor")
workflow = WorkflowBuilder().set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).build()
# First run the workflow until it emits a calculation request
request_info_event: WorkflowEvent | None = None
@@ -230,7 +230,7 @@ class TestRequestInfoAndResponse:
async def test_multiple_requests_workflow(self):
"""Test workflow with multiple concurrent requests."""
executor = MultiRequestExecutor(id="multi_executor")
workflow = WorkflowBuilder().set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).build()
# Collect all request events by running the full stream
request_events: list[WorkflowEvent] = []
@@ -264,7 +264,7 @@ class TestRequestInfoAndResponse:
async def test_denied_approval_workflow(self):
"""Test workflow when approval is denied."""
executor = ApprovalRequiredExecutor(id="approval_executor")
workflow = WorkflowBuilder().set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).build()
# First run the workflow until it emits a request
request_info_event: WorkflowEvent | None = None
@@ -287,7 +287,7 @@ class TestRequestInfoAndResponse:
async def test_workflow_state_with_pending_requests(self):
"""Test workflow state when waiting for responses."""
executor = ApprovalRequiredExecutor(id="approval_executor")
workflow = WorkflowBuilder().set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).build()
# Run workflow until idle with pending requests
request_info_event: WorkflowEvent | None = None
@@ -312,7 +312,7 @@ class TestRequestInfoAndResponse:
async def test_invalid_calculation_input(self):
"""Test workflow handling of invalid calculation input."""
executor = CalculationExecutor(id="calc_executor")
workflow = WorkflowBuilder().set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).build()
# Send invalid input (no numbers)
completed = False
@@ -334,7 +334,7 @@ class TestRequestInfoAndResponse:
# Create workflow with checkpointing enabled
executor = ApprovalRequiredExecutor(id="approval_executor")
workflow = WorkflowBuilder().set_start_executor(executor).with_checkpointing(storage).build()
workflow = WorkflowBuilder(start_executor=executor, checkpoint_storage=storage).build()
# Step 1: Run workflow to completion to ensure checkpoints are created
request_info_event: WorkflowEvent | None = None
@@ -372,7 +372,7 @@ class TestRequestInfoAndResponse:
# Step 4: Create a fresh workflow and restore from checkpoint
new_executor = ApprovalRequiredExecutor(id="approval_executor")
restored_workflow = WorkflowBuilder().set_start_executor(new_executor).with_checkpointing(storage).build()
restored_workflow = WorkflowBuilder(start_executor=new_executor, checkpoint_storage=storage).build()
# Step 5: Resume from checkpoint and verify the request can be continued
completed = False
python/packages/core/tests/workflow/test_serialization.py
+13 -18
View File
@@ -413,16 +413,14 @@ class TestSerializationWorkflowClasses:
"""
# Create innermost workflow
inner_executor = SampleExecutor(id="inner-exec")
inner_workflow = WorkflowBuilder().set_start_executor(inner_executor).set_max_iterations(10).build()
inner_workflow = WorkflowBuilder(max_iterations=10, start_executor=inner_executor).build()
# Create middle workflow with WorkflowExecutor
inner_workflow_executor = WorkflowExecutor(workflow=inner_workflow, id="inner-workflow-exec")
middle_executor = SampleExecutor(id="middle-exec")
middle_workflow = (
WorkflowBuilder()
.set_start_executor(middle_executor)
WorkflowBuilder(max_iterations=20, start_executor=middle_executor)
.add_edge(middle_executor, inner_workflow_executor)
.set_max_iterations(20)
.build()
)
@@ -430,10 +428,8 @@ class TestSerializationWorkflowClasses:
middle_workflow_executor = WorkflowExecutor(workflow=middle_workflow, id="middle-workflow-exec")
outer_executor = SampleExecutor(id="outer-exec")
outer_workflow = (
WorkflowBuilder()
.set_start_executor(outer_executor)
WorkflowBuilder(max_iterations=30, start_executor=outer_executor)
.add_edge(outer_executor, middle_workflow_executor)
.set_max_iterations(30)
.build()
)
@@ -543,7 +539,7 @@ class TestSerializationWorkflowClasses:
executor1 = SampleExecutor(id="executor1")
executor2 = SampleExecutor(id="executor2")
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
# Test model_dump
data = workflow.to_dict()
@@ -616,7 +612,7 @@ class TestSerializationWorkflowClasses:
executor1 = SampleExecutor(id="executor1")
executor2 = SampleExecutor(id="executor2")
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
# Test model_dump - should not include private runtime objects
data = workflow.to_dict()
@@ -629,11 +625,11 @@ class TestSerializationWorkflowClasses:
def test_workflow_name_description_serialization(self) -> None:
"""Test that workflow name and description are serialized correctly."""
# Test 1: With name and description
workflow1 = (
WorkflowBuilder(name="Test Pipeline", description="Test workflow description")
.set_start_executor(SampleExecutor(id="e1"))
.build()
)
workflow1 = WorkflowBuilder(
name="Test Pipeline",
description="Test workflow description",
start_executor=SampleExecutor(id="e1"),
).build()
assert workflow1.name == "Test Pipeline"
assert workflow1.description == "Test workflow description"
@@ -649,7 +645,7 @@ class TestSerializationWorkflowClasses:
assert parsed1["description"] == "Test workflow description"
# Test 2: Without name and description (defaults)
workflow2 = WorkflowBuilder().set_start_executor(SampleExecutor(id="e2")).build()
workflow2 = WorkflowBuilder(start_executor=SampleExecutor(id="e2")).build()
assert workflow2.name is None
assert workflow2.description is None
@@ -659,7 +655,7 @@ class TestSerializationWorkflowClasses:
assert "description" not in data2
# Test 3: With only name (no description)
workflow3 = WorkflowBuilder(name="Named Only").set_start_executor(SampleExecutor(id="e3")).build()
workflow3 = WorkflowBuilder(name="Named Only", start_executor=SampleExecutor(id="e3")).build()
assert workflow3.name == "Named Only"
assert workflow3.description is None
@@ -706,8 +702,7 @@ def test_comprehensive_edge_groups_workflow_serialization() -> None:
# Build workflow with all three edge group types
workflow = (
WorkflowBuilder()
.set_start_executor(router)
WorkflowBuilder(start_executor=router)
# 1. SwitchCaseEdgeGroup: Conditional routing
.add_switch_case_edge_group(
router,
python/packages/core/tests/workflow/test_sub_workflow.py
+7 -14
View File
@@ -167,8 +167,7 @@ def create_email_validation_workflow() -> Workflow:
email_domain_validator = EmailDomainValidator()
return (
WorkflowBuilder()
.set_start_executor(email_format_validator)
WorkflowBuilder(start_executor=email_format_validator)
.add_edge(email_format_validator, email_domain_validator)
.build()
)
@@ -184,8 +183,7 @@ async def test_basic_sub_workflow() -> None:
workflow_executor = WorkflowExecutor(validation_workflow, "email_validation_workflow")
main_workflow = (
WorkflowBuilder()
.set_start_executor(parent)
WorkflowBuilder(start_executor=parent)
.add_edge(parent, workflow_executor)
.add_edge(workflow_executor, parent)
.build()
@@ -223,8 +221,7 @@ async def test_sub_workflow_with_interception():
workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow")
main_workflow = (
WorkflowBuilder()
.set_start_executor(parent)
WorkflowBuilder(start_executor=parent)
.add_edge(parent, workflow_executor)
.add_edge(workflow_executor, parent)
.build()
@@ -340,8 +337,7 @@ async def test_workflow_scoped_interception() -> None:
executor_b = WorkflowExecutor(workflow_b, "workflow_b")
main_workflow = (
WorkflowBuilder()
.set_start_executor(parent)
WorkflowBuilder(start_executor=parent)
.add_edge(parent, executor_a)
.add_edge(parent, executor_b)
.add_edge(executor_a, parent)
@@ -422,8 +418,7 @@ async def test_concurrent_sub_workflow_execution() -> None:
workflow_executor = WorkflowExecutor(validation_workflow, "email_workflow")
main_workflow = (
WorkflowBuilder()
.set_start_executor(processor)
WorkflowBuilder(start_executor=processor)
.add_edge(processor, workflow_executor)
.add_edge(workflow_executor, processor)
.build()
@@ -564,16 +559,14 @@ class CheckpointTestCoordinator(Executor):
def _build_checkpoint_test_workflow(storage: InMemoryCheckpointStorage) -> Workflow:
"""Build the main workflow with checkpointing for testing."""
two_step_executor = TwoStepSubWorkflowExecutor()
sub_workflow = WorkflowBuilder().set_start_executor(two_step_executor).build()
sub_workflow = WorkflowBuilder(start_executor=two_step_executor).build()
sub_workflow_executor = WorkflowExecutor(sub_workflow, id="sub_workflow_executor")
coordinator = CheckpointTestCoordinator()
return (
WorkflowBuilder()
.set_start_executor(coordinator)
WorkflowBuilder(start_executor=coordinator, checkpoint_storage=storage)
.add_edge(coordinator, sub_workflow_executor)
.add_edge(sub_workflow_executor, coordinator)
.with_checkpointing(storage)
.build()
)
python/packages/core/tests/workflow/test_validation.py
+33 -67
View File
@@ -69,9 +69,8 @@ def test_valid_workflow_passes_validation():
# Create a valid workflow
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor=executor1)
.add_edge(executor1, executor2)
.set_start_executor(executor1)
.build() # This should not raise any exceptions
)
@@ -83,7 +82,7 @@ def test_duplicate_executor_ids_fail_validation():
executor2 = IntExecutor(id="dup")
with pytest.raises(ValueError) as exc_info:
(WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build())
(WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build())
assert str(exc_info.value) == "Duplicate executor ID 'dup' detected in workflow."
@@ -93,9 +92,7 @@ def test_edge_duplication_validation_fails():
executor2 = StringExecutor(id="executor2")
with pytest.raises(EdgeDuplicationError) as exc_info:
WorkflowBuilder().add_edge(executor1, executor2).add_edge(executor1, executor2).set_start_executor(
executor1
).build()
WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).add_edge(executor1, executor2).build()
assert "executor1->executor2" in str(exc_info.value)
assert exc_info.value.validation_type == ValidationTypeEnum.EDGE_DUPLICATION
@@ -106,7 +103,7 @@ def test_type_compatibility_validation_fails():
int_executor = IntExecutor(id="int_executor")
with pytest.raises(TypeCompatibilityError) as exc_info:
WorkflowBuilder().add_edge(string_executor, int_executor).set_start_executor(string_executor).build()
WorkflowBuilder(start_executor=string_executor).add_edge(string_executor, int_executor).build()
error = exc_info.value
assert error.source_executor_id == "string_executor"
@@ -119,7 +116,7 @@ def test_type_compatibility_with_any_type_passes():
any_executor = AnyExecutor(id="any_executor")
# This should not raise an exception
workflow = WorkflowBuilder().add_edge(string_executor, any_executor).set_start_executor(string_executor).build()
workflow = WorkflowBuilder(start_executor=string_executor).add_edge(string_executor, any_executor).build()
assert workflow is not None
@@ -129,9 +126,7 @@ def test_type_compatibility_with_no_output_types():
string_executor = StringExecutor(id="string_executor")
# This should pass validation since no output types are specified
workflow = (
WorkflowBuilder().add_edge(no_output_executor, string_executor).set_start_executor(no_output_executor).build()
)
workflow = WorkflowBuilder(start_executor=no_output_executor).add_edge(no_output_executor, string_executor).build()
assert workflow is not None
@@ -141,9 +136,7 @@ def test_multi_type_executor_compatibility():
multi_type_executor = MultiTypeExecutor(id="multi_type")
# String executor outputs strings, multi-type can handle strings
workflow = (
WorkflowBuilder().add_edge(string_executor, multi_type_executor).set_start_executor(string_executor).build()
)
workflow = WorkflowBuilder(start_executor=string_executor).add_edge(string_executor, multi_type_executor).build()
assert workflow is not None
@@ -154,9 +147,7 @@ def test_graph_connectivity_unreachable_executors():
executor3 = StringExecutor(id="executor3") # This will be unreachable
with pytest.raises(GraphConnectivityError) as exc_info:
WorkflowBuilder().add_edge(executor1, executor2).add_edge(executor3, executor2).set_start_executor(
executor1
).build()
WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).add_edge(executor3, executor2).build()
assert "unreachable" in str(exc_info.value).lower()
assert "executor3" in str(exc_info.value)
@@ -189,19 +180,14 @@ def test_disconnected_start_executor_not_in_graph():
executor3 = StringExecutor(id="executor3") # Not in graph
with pytest.raises(GraphConnectivityError) as exc_info:
WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor3).build()
WorkflowBuilder(start_executor=executor3).add_edge(executor1, executor2).build()
assert "The following executors are unreachable from the start executor 'executor3'" in str(exc_info.value)
def test_missing_start_executor():
executor1 = StringExecutor(id="executor1")
executor2 = StringExecutor(id="executor2")
with pytest.raises(ValueError) as exc_info:
WorkflowBuilder().add_edge(executor1, executor2).build()
assert "Starting executor must be set" in str(exc_info.value)
with pytest.raises(TypeError):
WorkflowBuilder() # type: ignore[call-arg]
def test_workflow_validation_error_base_class():
@@ -219,12 +205,11 @@ def test_complex_workflow_validation():
executor4 = AnyExecutor(id="executor4")
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor=executor1)
.add_edge(executor1, executor2) # str -> MultiType (compatible)
.add_edge(executor2, executor3) # MultiType -> str (compatible)
.add_edge(executor2, executor4) # MultiType -> Any (compatible)
.add_edge(executor3, executor4) # str -> Any (compatible)
.set_start_executor(executor1)
.build()
)
@@ -246,7 +231,7 @@ def test_type_compatibility_inheritance():
derived_executor = DerivedExecutor(id="derived")
# This should pass since both handle str
workflow = WorkflowBuilder().add_edge(base_executor, derived_executor).set_start_executor(base_executor).build()
workflow = WorkflowBuilder(start_executor=base_executor).add_edge(base_executor, derived_executor).build()
assert workflow is not None
@@ -271,7 +256,7 @@ def test_fan_out_validation():
target1 = StringExecutor(id="target1")
target2 = AnyExecutor(id="target2")
workflow = WorkflowBuilder().add_fan_out_edges(source, [target1, target2]).set_start_executor(source).build()
workflow = WorkflowBuilder(start_executor=source).add_fan_out_edges(source, [target1, target2]).build()
assert workflow is not None
@@ -284,11 +269,10 @@ def test_fan_in_validation():
# Create a proper fan-in by having a start executor that connects to both sources
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor=start_executor)
.add_edge(start_executor, source1) # Start connects to source1
.add_edge(start_executor, source2) # Start connects to source2
.add_fan_in_edges([source1, source2], target) # Both sources fan-in to target
.set_start_executor(start_executor)
.build()
)
@@ -300,7 +284,7 @@ def test_chain_validation():
executor2 = StringExecutor(id="executor2")
executor3 = AnyExecutor(id="executor3")
workflow = WorkflowBuilder().add_chain([executor1, executor2, executor3]).set_start_executor(executor1).build()
workflow = WorkflowBuilder(start_executor=executor1).add_chain([executor1, executor2, executor3]).build()
assert workflow is not None
@@ -313,9 +297,7 @@ def test_logging_for_missing_output_types(caplog: Any) -> None:
string_executor = StringExecutor(id="string_executor")
# This should trigger a warning log
workflow = (
WorkflowBuilder().add_edge(no_output_executor, string_executor).set_start_executor(no_output_executor).build()
)
workflow = WorkflowBuilder(start_executor=no_output_executor).add_edge(no_output_executor, string_executor).build()
assert workflow is not None
assert "has no output type annotations" in caplog.text
@@ -338,9 +320,7 @@ def test_logging_for_missing_input_types(caplog: Any) -> None:
no_input_executor = NoInputTypesExecutor(id="no_input")
# This should pass since NoInputTypesExecutor has no proper input types
workflow = (
WorkflowBuilder().add_edge(string_executor, no_input_executor).set_start_executor(string_executor).build()
)
workflow = WorkflowBuilder(start_executor=string_executor).add_edge(string_executor, no_input_executor).build()
assert workflow is not None
@@ -351,7 +331,7 @@ def test_self_loop_detection_warning(caplog: Any) -> None:
executor = StringExecutor(id="self_loop_executor")
# Create a self-loop
workflow = WorkflowBuilder().add_edge(executor, executor).set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).add_edge(executor, executor).build()
assert workflow is not None
assert "Self-loop detected" in caplog.text
@@ -365,7 +345,7 @@ def test_handler_validation_basic(caplog: Any) -> None:
start_executor = StringExecutor(id="start")
target_executor = StringExecutor(id="target")
workflow = WorkflowBuilder().add_edge(start_executor, target_executor).set_start_executor(start_executor).build()
workflow = WorkflowBuilder(start_executor=start_executor).add_edge(start_executor, target_executor).build()
assert workflow is not None
# Just ensure the validation runs without errors
@@ -377,7 +357,7 @@ def test_dead_end_detection(caplog: Any) -> None:
executor1 = StringExecutor(id="executor1")
executor2 = StringExecutor(id="executor2") # This will be a dead end
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
assert workflow is not None
assert "Dead-end executors detected" in caplog.text
@@ -391,7 +371,7 @@ def test_successful_type_compatibility_logging(caplog: Any) -> None:
executor1 = StringExecutor(id="executor1")
executor2 = StringExecutor(id="executor2")
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
assert workflow is not None
assert "Type compatibility validated for edge" in caplog.text
@@ -406,11 +386,7 @@ def test_multiple_dead_ends_detection(caplog: Any) -> None:
executor3 = StringExecutor(id="executor3") # Dead end
workflow = (
WorkflowBuilder()
.add_edge(executor1, executor2)
.add_edge(executor1, executor3)
.set_start_executor(executor1)
.build()
WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).add_edge(executor1, executor3).build()
)
assert workflow is not None
@@ -426,7 +402,7 @@ def test_single_executor_workflow(caplog: Any) -> None:
executor2 = StringExecutor(id="executor2")
# Create a simple two-executor workflow to avoid graph validation issues
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
assert workflow is not None
# Should detect executor2 as dead end
@@ -438,7 +414,7 @@ def test_enhanced_type_compatibility_error_details():
int_executor = IntExecutor(id="int_executor")
with pytest.raises(TypeCompatibilityError) as exc_info:
WorkflowBuilder().add_edge(string_executor, int_executor).set_start_executor(string_executor).build()
WorkflowBuilder(start_executor=string_executor).add_edge(string_executor, int_executor).build()
error = exc_info.value
# Verify enhanced error contains detailed type information
@@ -463,7 +439,7 @@ def test_union_type_compatibility_validation() -> None:
union_input = UnionInputExecutor(id="union_input")
# This should pass validation due to type compatibility (str)
workflow = WorkflowBuilder().add_edge(union_output, union_input).set_start_executor(union_output).build()
workflow = WorkflowBuilder(start_executor=union_output).add_edge(union_output, union_input).build()
assert workflow is not None
@@ -483,7 +459,7 @@ def test_generic_type_compatibility() -> None:
list_input = ListInputExecutor(id="list_input")
# This should pass validation for generic type compatibility
workflow = WorkflowBuilder().add_edge(list_output, list_input).set_start_executor(list_output).build()
workflow = WorkflowBuilder(start_executor=list_output).add_edge(list_output, list_input).build()
assert workflow is not None
@@ -539,7 +515,7 @@ def test_handler_ctx_none_is_allowed() -> None:
none_exec = NoneExecutor(id="n")
# Should build successfully
wf = WorkflowBuilder().add_edge(start, none_exec).set_start_executor(start).build()
wf = WorkflowBuilder(start_executor=start).add_edge(start, none_exec).build()
assert wf is not None
@@ -555,7 +531,7 @@ def test_handler_ctx_any_is_allowed_but_skips_type_checks(caplog: Any) -> None:
any_out = AnyOutExecutor(id="a")
# Builds; later edges from this executor will skip type compatibility when outputs are unspecified
wf = WorkflowBuilder().add_edge(start, any_out).set_start_executor(start).build()
wf = WorkflowBuilder(start_executor=start).add_edge(start, any_out).build()
assert wf is not None
@@ -575,11 +551,7 @@ def test_output_validation_with_valid_output_executors():
# Build workflow with valid output executors
workflow = (
WorkflowBuilder()
.add_edge(executor1, executor2)
.set_start_executor(executor1)
.with_output_from([executor2])
.build()
WorkflowBuilder(start_executor=executor1, output_executors=[executor2]).add_edge(executor1, executor2).build()
)
assert workflow is not None
@@ -593,11 +565,9 @@ def test_output_validation_with_multiple_valid_output_executors():
executor3 = OutputExecutor(id="executor3")
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor=executor1, output_executors=[executor1, executor3])
.add_edge(executor1, executor2)
.add_edge(executor2, executor3)
.set_start_executor(executor1)
.with_output_from([executor1, executor3])
.build()
)
@@ -628,10 +598,8 @@ def test_output_validation_fails_for_executor_without_output_types():
with pytest.raises(WorkflowValidationError) as exc_info:
(
WorkflowBuilder()
WorkflowBuilder(start_executor=executor1, output_executors=[no_output_executor])
.add_edge(executor1, no_output_executor)
.set_start_executor(executor1)
.with_output_from([no_output_executor])
.build()
)
@@ -645,9 +613,7 @@ def test_output_validation_empty_list_passes():
executor1 = OutputExecutor(id="executor1")
executor2 = OutputExecutor(id="executor2")
workflow = (
WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).with_output_from([]).build()
)
workflow = WorkflowBuilder(start_executor=executor1, output_executors=[]).add_edge(executor1, executor2).build()
assert workflow is not None
# All executors are outputs
python/packages/core/tests/workflow/test_viz.py
+17 -41
View File
@@ -31,7 +31,7 @@ def basic_sub_workflow():
sub_exec1 = MockExecutor(id="sub_exec1")
sub_exec2 = MockExecutor(id="sub_exec2")
sub_workflow = WorkflowBuilder().add_edge(sub_exec1, sub_exec2).set_start_executor(sub_exec1).build()
sub_workflow = WorkflowBuilder(start_executor=sub_exec1).add_edge(sub_exec1, sub_exec2).build()
# Create a workflow executor that wraps the sub-workflow
workflow_executor = WorkflowExecutor(sub_workflow, id="workflow_executor_1")
@@ -41,10 +41,9 @@ def basic_sub_workflow():
final_exec = MockExecutor(id="final_executor")
main_workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor=main_exec)
.add_edge(main_exec, workflow_executor)
.add_edge(workflow_executor, final_exec)
.set_start_executor(main_exec)
.build()
)
@@ -65,7 +64,7 @@ def test_workflow_viz_to_digraph():
executor1 = MockExecutor(id="executor1")
executor2 = MockExecutor(id="executor2")
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
viz = WorkflowViz(workflow)
dot_content = viz.to_digraph()
@@ -84,7 +83,7 @@ def test_workflow_viz_export_dot():
executor1 = MockExecutor(id="executor1")
executor2 = MockExecutor(id="executor2")
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
viz = WorkflowViz(workflow)
@@ -104,7 +103,7 @@ def test_workflow_viz_export_dot_with_filename(tmp_path):
executor1 = MockExecutor(id="executor1")
executor2 = MockExecutor(id="executor2")
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
viz = WorkflowViz(workflow)
@@ -128,12 +127,11 @@ def test_workflow_viz_complex_workflow():
executor4 = MockExecutor(id="end")
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor=executor1)
.add_edge(executor1, executor2)
.add_edge(executor1, executor3)
.add_edge(executor2, executor4)
.add_edge(executor3, executor4)
.set_start_executor(executor1)
.build()
)
@@ -162,7 +160,7 @@ def test_workflow_viz_export_svg():
executor1 = MockExecutor(id="executor1")
executor2 = MockExecutor(id="executor2")
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
viz = WorkflowViz(workflow)
@@ -178,7 +176,7 @@ def test_workflow_viz_unsupported_format():
executor1 = MockExecutor(id="executor1")
executor2 = MockExecutor(id="executor2")
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
viz = WorkflowViz(workflow)
@@ -196,7 +194,7 @@ def test_workflow_viz_graphviz_binary_not_found():
executor1 = MockExecutor(id="executor1")
executor2 = MockExecutor(id="executor2")
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
viz = WorkflowViz(workflow)
# Mock graphviz.Source.render to raise ExecutableNotFound
@@ -224,13 +222,7 @@ def test_workflow_viz_conditional_edge():
def only_if_foo(msg: str) -> bool: # pragma: no cover - simple predicate
return msg == "foo"
wf = (
WorkflowBuilder()
.add_edge(start, mid, condition=only_if_foo)
.add_edge(mid, end)
.set_start_executor(start)
.build()
)
wf = WorkflowBuilder(start_executor=start).add_edge(start, mid, condition=only_if_foo).add_edge(mid, end).build()
dot = WorkflowViz(wf).to_digraph()
@@ -249,13 +241,7 @@ def test_workflow_viz_fan_in_edge_group():
t = ListStrTargetExecutor(id="t")
# Build a connected workflow: start fans out to s1 and s2, which then fan-in to t
wf = (
WorkflowBuilder()
.add_fan_out_edges(start, [s1, s2])
.add_fan_in_edges([s1, s2], t)
.set_start_executor(start)
.build()
)
wf = WorkflowBuilder(start_executor=start).add_fan_out_edges(start, [s1, s2]).add_fan_in_edges([s1, s2], t).build()
dot = WorkflowViz(wf).to_digraph()
@@ -287,7 +273,7 @@ def test_workflow_viz_to_mermaid_basic():
executor1 = MockExecutor(id="executor1")
executor2 = MockExecutor(id="executor2")
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
mermaid = WorkflowViz(workflow).to_mermaid()
# Start node and normal node
@@ -305,7 +291,7 @@ def test_workflow_viz_mermaid_conditional_edge():
def only_if_foo(msg: str) -> bool: # pragma: no cover - simple predicate
return msg == "foo"
wf = WorkflowBuilder().add_edge(start, mid, condition=only_if_foo).set_start_executor(start).build()
wf = WorkflowBuilder(start_executor=start).add_edge(start, mid, condition=only_if_foo).build()
mermaid = WorkflowViz(wf).to_mermaid()
assert "start -. conditional .-> mid" in mermaid
@@ -318,13 +304,7 @@ def test_workflow_viz_mermaid_fan_in_edge_group():
s2 = MockExecutor(id="s2")
t = ListStrTargetExecutor(id="t")
wf = (
WorkflowBuilder()
.add_fan_out_edges(start, [s1, s2])
.add_fan_in_edges([s1, s2], t)
.set_start_executor(start)
.build()
)
wf = WorkflowBuilder(start_executor=start).add_fan_out_edges(start, [s1, s2]).add_fan_in_edges([s1, s2], t).build()
mermaid = WorkflowViz(wf).to_mermaid()
lines = [line.strip() for line in mermaid.splitlines()]
@@ -398,23 +378,19 @@ def test_workflow_viz_nested_sub_workflows():
"""Test visualization of deeply nested sub-workflows."""
# Create innermost sub-workflow
inner_exec = MockExecutor(id="inner_exec")
inner_workflow = WorkflowBuilder().set_start_executor(inner_exec).build()
inner_workflow = WorkflowBuilder(start_executor=inner_exec).build()
# Create middle sub-workflow that contains the inner one
inner_workflow_executor = WorkflowExecutor(inner_workflow, id="inner_wf_exec")
middle_exec = MockExecutor(id="middle_exec")
middle_workflow = (
WorkflowBuilder().add_edge(middle_exec, inner_workflow_executor).set_start_executor(middle_exec).build()
)
middle_workflow = WorkflowBuilder(start_executor=middle_exec).add_edge(middle_exec, inner_workflow_executor).build()
# Create outer workflow
middle_workflow_executor = WorkflowExecutor(middle_workflow, id="middle_wf_exec")
outer_exec = MockExecutor(id="outer_exec")
outer_workflow = (
WorkflowBuilder().add_edge(outer_exec, middle_workflow_executor).set_start_executor(outer_exec).build()
)
outer_workflow = WorkflowBuilder(start_executor=outer_exec).add_edge(outer_exec, middle_workflow_executor).build()
viz = WorkflowViz(outer_workflow)
dot_content = viz.to_digraph()
python/packages/core/tests/workflow/test_workflow.py
+34 -67
View File
@@ -110,8 +110,7 @@ async def test_workflow_run_streaming() -> None:
executor_b = IncrementExecutor(id="executor_b")
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
WorkflowBuilder(start_executor=executor_a)
.add_edge(executor_a, executor_b)
.add_edge(executor_b, executor_a)
.build()
@@ -132,11 +131,9 @@ async def test_workflow_run_stream_not_completed():
executor_b = IncrementExecutor(id="executor_b")
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
WorkflowBuilder(max_iterations=5, start_executor=executor_a)
.add_edge(executor_a, executor_b)
.add_edge(executor_b, executor_a)
.set_max_iterations(5)
.build()
)
@@ -151,8 +148,7 @@ async def test_workflow_run():
executor_b = IncrementExecutor(id="executor_b")
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
WorkflowBuilder(start_executor=executor_a)
.add_edge(executor_a, executor_b)
.add_edge(executor_b, executor_a)
.build()
@@ -170,11 +166,9 @@ async def test_workflow_run_not_completed():
executor_b = IncrementExecutor(id="executor_b")
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
WorkflowBuilder(max_iterations=5, start_executor=executor_a)
.add_edge(executor_a, executor_b)
.add_edge(executor_b, executor_a)
.set_max_iterations(5)
.build()
)
@@ -189,7 +183,7 @@ async def test_fan_out():
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()
WorkflowBuilder(start_executor=executor_a).add_fan_out_edges(executor_a, [executor_b, executor_c]).build()
)
events = await workflow.run(NumberMessage(data=0))
@@ -214,7 +208,7 @@ async def test_fan_out_multiple_completed_events():
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()
WorkflowBuilder(start_executor=executor_a).add_fan_out_edges(executor_a, [executor_b, executor_c]).build()
)
events = await workflow.run(NumberMessage(data=0))
@@ -239,8 +233,7 @@ async def test_fan_in():
aggregator = AggregatorExecutor(id="aggregator")
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
WorkflowBuilder(start_executor=executor_a)
.add_fan_out_edges(executor_a, [executor_b, executor_c])
.add_fan_in_edges([executor_b, executor_c], aggregator)
.build()
@@ -276,10 +269,8 @@ async def test_workflow_with_checkpointing_enabled(simple_executor: Executor):
# Build workflow with checkpointing - should not raise any errors
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor=simple_executor, checkpoint_storage=storage)
.add_edge(simple_executor, simple_executor) # Self-loop to satisfy graph requirements
.set_start_executor(simple_executor)
.with_checkpointing(storage)
.build()
)
@@ -295,9 +286,8 @@ async def test_workflow_checkpointing_not_enabled_for_external_restore(
"""Test that external checkpoint restoration fails when workflow doesn't support checkpointing."""
# Build workflow WITHOUT checkpointing
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor=simple_executor)
.add_edge(simple_executor, simple_executor) # Self-loop to satisfy graph requirements
.set_start_executor(simple_executor)
.build()
)
@@ -315,9 +305,8 @@ async def test_workflow_run_stream_from_checkpoint_no_checkpointing_enabled(
):
# Build workflow WITHOUT checkpointing
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor=simple_executor)
.add_edge(simple_executor, simple_executor) # Self-loop to satisfy graph requirements
.set_start_executor(simple_executor)
.build()
)
@@ -340,10 +329,8 @@ async def test_workflow_run_stream_from_checkpoint_invalid_checkpoint(
# Build workflow with checkpointing
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor=simple_executor, checkpoint_storage=storage)
.add_edge(simple_executor, simple_executor) # Self-loop to satisfy graph requirements
.set_start_executor(simple_executor)
.with_checkpointing(storage)
.build()
)
@@ -376,7 +363,7 @@ async def test_workflow_run_stream_from_checkpoint_with_external_storage(
# Create a workflow WITHOUT checkpointing
workflow_without_checkpointing = (
WorkflowBuilder().add_edge(simple_executor, simple_executor).set_start_executor(simple_executor).build()
WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
)
# Resume from checkpoint using external storage parameter
@@ -411,10 +398,8 @@ async def test_workflow_run_from_checkpoint_non_streaming(simple_executor: Execu
# Build workflow with checkpointing
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor=simple_executor, checkpoint_storage=storage)
.add_edge(simple_executor, simple_executor)
.set_start_executor(simple_executor)
.with_checkpointing(storage)
.build()
)
@@ -452,10 +437,8 @@ async def test_workflow_run_stream_from_checkpoint_with_responses(
# Build workflow with checkpointing
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor=simple_executor, checkpoint_storage=storage)
.add_edge(simple_executor, simple_executor)
.set_start_executor(simple_executor)
.with_checkpointing(storage)
.build()
)
@@ -512,10 +495,8 @@ async def test_workflow_multiple_runs_no_state_collision():
# Build workflow with checkpointing
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor=state_executor, checkpoint_storage=storage)
.add_edge(state_executor, state_executor) # Self-loop to satisfy graph requirements
.set_start_executor(state_executor)
.with_checkpointing(storage)
.build()
)
@@ -552,9 +533,7 @@ async def test_workflow_checkpoint_runtime_only_configuration(
storage = FileCheckpointStorage(temp_dir)
# Build workflow WITHOUT checkpointing at build time
workflow = (
WorkflowBuilder().add_edge(simple_executor, simple_executor).set_start_executor(simple_executor).build()
)
workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
# Run with runtime checkpoint storage - should create checkpoints
test_message = Message(data="runtime checkpoint test", source_id="test", target_id=None)
@@ -575,7 +554,7 @@ async def test_workflow_checkpoint_runtime_only_configuration(
# Create new workflow instance (still without build-time checkpointing)
workflow_resume = (
WorkflowBuilder().add_edge(simple_executor, simple_executor).set_start_executor(simple_executor).build()
WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
)
# Resume from checkpoint using runtime checkpoint storage
@@ -602,10 +581,8 @@ async def test_workflow_checkpoint_runtime_overrides_buildtime(
# Build workflow with build-time checkpointing
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor=simple_executor, checkpoint_storage=buildtime_storage)
.add_edge(simple_executor, simple_executor)
.set_start_executor(simple_executor)
.with_checkpointing(buildtime_storage)
.build()
)
@@ -643,8 +620,7 @@ async def test_comprehensive_edge_groups_workflow():
# 3. FanOut: fanout_hub -> [parallel_1, parallel_2]
# 4. FanIn: [parallel_1, parallel_2] -> aggregator
workflow = (
WorkflowBuilder()
.set_start_executor(router)
WorkflowBuilder(start_executor=router)
# Switch-case routing based on message data
.add_switch_case_edge_group(
router,
@@ -713,8 +689,7 @@ async def test_workflow_with_simple_cycle_and_exit_condition():
# Simple cycle: A -> B -> A, A exits when limit reached
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
WorkflowBuilder(start_executor=executor_a)
.add_edge(executor_a, executor_b) # A -> B
.add_edge(executor_b, executor_a) # B -> A (creates cycle)
.build()
@@ -746,7 +721,7 @@ async def test_workflow_concurrent_execution_prevention():
"""Test that concurrent workflow executions are prevented."""
# Create a simple workflow that takes some time to execute
executor = IncrementExecutor(id="slow_executor", limit=3, increment=1)
workflow = WorkflowBuilder().set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).build()
# Create a task that will run the workflow
async def run_workflow():
@@ -778,7 +753,7 @@ async def test_workflow_concurrent_execution_prevention_streaming():
"""Test that concurrent workflow streaming executions are prevented."""
# Create a simple workflow
executor = IncrementExecutor(id="slow_executor", limit=3, increment=1)
workflow = WorkflowBuilder().set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).build()
# Create an async generator that will consume the stream slowly
async def consume_stream_slowly():
@@ -814,7 +789,7 @@ async def test_workflow_concurrent_execution_prevention_mixed_methods():
"""Test that concurrent executions are prevented across different execution methods."""
# Create a simple workflow
executor = IncrementExecutor(id="slow_executor", limit=3, increment=1)
workflow = WorkflowBuilder().set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).build()
# Start a streaming execution
async def consume_stream():
@@ -884,7 +859,7 @@ async def test_agent_streaming_vs_non_streaming() -> None:
agent = _StreamingTestAgent(id="test_agent", name="TestAgent", reply_text="Hello World")
agent_exec = AgentExecutor(agent, id="agent_exec")
workflow = WorkflowBuilder().set_start_executor(agent_exec).build()
workflow = WorkflowBuilder(start_executor=agent_exec).build()
# Test non-streaming mode with run()
result = await workflow.run("test message")
@@ -934,7 +909,7 @@ async def test_agent_streaming_vs_non_streaming() -> None:
async def test_workflow_run_parameter_validation(simple_executor: Executor) -> None:
"""Test that stream properly validate parameter combinations."""
workflow = WorkflowBuilder().add_edge(simple_executor, simple_executor).set_start_executor(simple_executor).build()
workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
test_message = Message(data="test", source_id="test", target_id=None)
@@ -965,7 +940,7 @@ async def test_workflow_run_stream_parameter_validation(
simple_executor: Executor,
) -> None:
"""Test stream=True specific parameter validation scenarios."""
workflow = WorkflowBuilder().add_edge(simple_executor, simple_executor).set_start_executor(simple_executor).build()
workflow = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, simple_executor).build()
test_message = Message(data="test", source_id="test", target_id=None)
@@ -1014,7 +989,7 @@ async def test_output_executors_empty_yields_all_outputs() -> None:
executor_b = OutputProducerExecutor(id="executor_b", output_value=20)
# Build workflow with a -> b
workflow = WorkflowBuilder().set_start_executor(executor_a).add_edge(executor_a, executor_b).build()
workflow = WorkflowBuilder(start_executor=executor_a).add_edge(executor_a, executor_b).build()
result = await workflow.run(NumberMessage(data=0))
outputs = result.get_outputs()
@@ -1037,10 +1012,8 @@ async def test_output_executors_filters_outputs_non_streaming() -> None:
# Build workflow with a -> b
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
WorkflowBuilder(start_executor=executor_a, output_executors=[executor_b])
.add_edge(executor_a, executor_b)
.with_output_from([executor_b])
.build()
)
@@ -1064,10 +1037,8 @@ async def test_output_executors_filters_outputs_streaming() -> None:
# Build workflow with a -> b
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
WorkflowBuilder(start_executor=executor_a, output_executors=[executor_a])
.add_edge(executor_a, executor_b)
.with_output_from([executor_a])
.build()
)
@@ -1092,11 +1063,9 @@ async def test_output_executors_with_multiple_specified_executors() -> None:
# Build workflow with a -> b -> c
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
WorkflowBuilder(start_executor=executor_a, output_executors=[executor_a, executor_c])
.add_edge(executor_a, executor_b)
.add_edge(executor_b, executor_c)
.with_output_from([executor_a, executor_c])
.build()
)
@@ -1114,7 +1083,7 @@ async def test_output_executors_with_nonexistent_executor_id() -> None:
"""Test that specifying a non-existent executor ID doesn't break the workflow."""
executor_a = OutputProducerExecutor(id="executor_a", output_value=42)
workflow = WorkflowBuilder().set_start_executor(executor_a).build()
workflow = WorkflowBuilder(start_executor=executor_a).build()
# Set output_executors to an ID that doesn't exist
workflow._output_executors = ["nonexistent_executor"] # type: ignore
@@ -1157,11 +1126,9 @@ async def test_output_executors_filtering_with_fan_in() -> None:
# Build fan-in workflow: start -> [a, b] -> aggregator
workflow = (
WorkflowBuilder()
.set_start_executor(executor_start)
WorkflowBuilder(start_executor=executor_start, output_executors=[aggregator])
.add_fan_out_edges(executor_start, [executor_a, executor_b])
.add_fan_in_edges([executor_a, executor_b], aggregator)
.with_output_from([aggregator])
.build()
)
@@ -1178,7 +1145,7 @@ async def test_output_executors_filtering_with_run_responses() -> None:
"""Test output filtering works correctly with run(responses=...) method."""
executor = MockExecutorRequestApproval(id="approval_executor")
workflow = WorkflowBuilder().set_start_executor(executor).with_output_from([executor]).build()
workflow = WorkflowBuilder(start_executor=executor, output_executors=[executor]).build()
# Run workflow which will request approval
result = await workflow.run(NumberMessage(data=42))
@@ -1201,7 +1168,7 @@ async def test_output_executors_filtering_with_run_responses_streaming() -> None
"""Test output filtering works correctly with run(responses=..., stream=True) method."""
executor = MockExecutorRequestApproval(id="approval_executor")
workflow = WorkflowBuilder().set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).build()
# Run workflow which will request approval
events_list: list[WorkflowEvent] = []
python/packages/core/tests/workflow/test_workflow_agent.py
+24 -31
View File
@@ -150,7 +150,7 @@ class TestWorkflowAgent:
executor1 = SimpleExecutor(id="executor1", response_text="Step1", streaming=False)
executor2 = SimpleExecutor(id="executor2", response_text="Step2", streaming=False)
workflow = WorkflowBuilder().set_start_executor(executor1).add_edge(executor1, executor2).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
agent = WorkflowAgent(workflow=workflow, name="Test Agent")
@@ -194,7 +194,7 @@ class TestWorkflowAgent:
executor2 = SimpleExecutor(id="stream2", response_text="Streaming2")
# Create workflow with just one executor
workflow = WorkflowBuilder().set_start_executor(executor1).add_edge(executor1, executor2).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
agent = WorkflowAgent(workflow=workflow, name="Streaming Test Agent")
@@ -224,7 +224,7 @@ class TestWorkflowAgent:
requesting_executor = RequestingExecutor(id="requester", streaming=False)
workflow = (
WorkflowBuilder().set_start_executor(simple_executor).add_edge(simple_executor, requesting_executor).build()
WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, requesting_executor).build()
)
agent = WorkflowAgent(workflow=workflow, name="Request Test Agent")
@@ -296,7 +296,7 @@ class TestWorkflowAgent:
"""Test that Workflow.as_agent() creates a properly configured WorkflowAgent."""
# Create a simple workflow
executor = SimpleExecutor(id="executor1", response_text="Response")
workflow = WorkflowBuilder().set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).build()
# Test as_agent with a name
agent = workflow.as_agent(name="TestAgent")
@@ -322,7 +322,7 @@ class TestWorkflowAgent:
# Create a simple workflow
executor = _Executor(id="test")
workflow = WorkflowBuilder().set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).build()
# Try to create an agent with unsupported input types
with pytest.raises(ValueError, match="Workflow's start executor cannot handle list\\[ChatMessage\\]"):
@@ -341,7 +341,7 @@ class TestWorkflowAgent:
input_text = messages[0].text if messages else "no input"
await ctx.yield_output(f"processed: {input_text}")
workflow = WorkflowBuilder().set_start_executor(yielding_executor).build()
workflow = WorkflowBuilder(start_executor=yielding_executor).build()
# Run directly - should return output event (type='output') in result
direct_result = await workflow.run([ChatMessage(role="user", text="hello")])
@@ -365,7 +365,7 @@ class TestWorkflowAgent:
await ctx.yield_output("first output")
await ctx.yield_output("second output")
workflow = WorkflowBuilder().set_start_executor(yielding_executor).build()
workflow = WorkflowBuilder(start_executor=yielding_executor).build()
agent = workflow.as_agent("test-agent")
updates: list[AgentResponseUpdate] = []
@@ -387,7 +387,7 @@ class TestWorkflowAgent:
await ctx.yield_output(Content.from_data(data=b"binary data", media_type="application/octet-stream"))
await ctx.yield_output(Content.from_uri(uri="https://example.com/image.png", media_type="image/png"))
workflow = WorkflowBuilder().set_start_executor(content_yielding_executor).build()
workflow = WorkflowBuilder(start_executor=content_yielding_executor).build()
agent = workflow.as_agent("content-test-agent")
result = await agent.run("test")
@@ -417,7 +417,7 @@ class TestWorkflowAgent:
)
await ctx.yield_output(msg)
workflow = WorkflowBuilder().set_start_executor(chat_message_executor).build()
workflow = WorkflowBuilder(start_executor=chat_message_executor).build()
agent = workflow.as_agent("chat-msg-agent")
result = await agent.run("test")
@@ -448,7 +448,7 @@ class TestWorkflowAgent:
custom = CustomData(42)
await ctx.yield_output(custom)
workflow = WorkflowBuilder().set_start_executor(raw_yielding_executor).build()
workflow = WorkflowBuilder(start_executor=raw_yielding_executor).build()
agent = workflow.as_agent("raw-test-agent")
updates: list[AgentResponseUpdate] = []
@@ -490,7 +490,7 @@ class TestWorkflowAgent:
]
await ctx.yield_output(msg_list)
workflow = WorkflowBuilder().set_start_executor(list_yielding_executor).build()
workflow = WorkflowBuilder(start_executor=list_yielding_executor).build()
agent = workflow.as_agent("list-msg-agent")
# Verify streaming returns the update with all 4 contents before coalescing
@@ -521,7 +521,7 @@ class TestWorkflowAgent:
"""
# Create an executor that captures all received messages
capturing_executor = ConversationHistoryCapturingExecutor(id="capturing", streaming=False)
workflow = WorkflowBuilder().set_start_executor(capturing_executor).build()
workflow = WorkflowBuilder(start_executor=capturing_executor).build()
agent = WorkflowAgent(workflow=workflow, name="Thread History Test Agent")
# Create a thread with existing conversation history
@@ -551,7 +551,7 @@ class TestWorkflowAgent:
"""
# Create an executor that captures all received messages
capturing_executor = ConversationHistoryCapturingExecutor(id="capturing_stream")
workflow = WorkflowBuilder().set_start_executor(capturing_executor).build()
workflow = WorkflowBuilder(start_executor=capturing_executor).build()
agent = WorkflowAgent(workflow=workflow, name="Thread Stream Test Agent")
# Create a thread with existing conversation history
@@ -579,7 +579,7 @@ class TestWorkflowAgent:
async def test_empty_thread_works_correctly(self) -> None:
"""Test that an empty thread (no message store) works correctly."""
capturing_executor = ConversationHistoryCapturingExecutor(id="empty_thread_test")
workflow = WorkflowBuilder().set_start_executor(capturing_executor).build()
workflow = WorkflowBuilder(start_executor=capturing_executor).build()
agent = WorkflowAgent(workflow=workflow, name="Empty Thread Test Agent")
# Create an empty thread
@@ -597,7 +597,7 @@ class TestWorkflowAgent:
from agent_framework import InMemoryCheckpointStorage
capturing_executor = ConversationHistoryCapturingExecutor(id="checkpoint_test")
workflow = WorkflowBuilder().set_start_executor(capturing_executor).build()
workflow = WorkflowBuilder(start_executor=capturing_executor).build()
agent = WorkflowAgent(workflow=workflow, name="Checkpoint Test Agent")
# Create checkpoint storage
@@ -675,17 +675,11 @@ class TestWorkflowAgent:
await ctx.send_message(AgentExecutorRequest(messages=messages, should_respond=True))
# Build workflow: start -> agent1 (no output) -> agent2 (output_response=True)
workflow = (
WorkflowBuilder()
.register_executor(lambda: start_executor, "start")
.register_agent(lambda: MockAgent("agent1", "Agent1 output - should NOT appear"), "agent1")
.register_agent(lambda: MockAgent("agent2", "Agent2 output - SHOULD appear"), "agent2")
.set_start_executor("start")
.add_edge("start", "agent1")
.add_edge("agent1", "agent2")
.with_output_from(["start", "agent2"])
.build()
)
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()
agent = WorkflowAgent(workflow=workflow, name="Test Agent")
result = await agent.run("Test input")
@@ -765,10 +759,9 @@ class TestWorkflowAgent:
# Build workflow with single agent
workflow = (
WorkflowBuilder()
WorkflowBuilder(start_executor="start")
.register_executor(lambda: start_executor, "start")
.register_agent(lambda: MockAgent("agent", "Unique response text"), "agent")
.set_start_executor("start")
.add_edge("start", "agent")
.build()
)
@@ -794,7 +787,7 @@ class TestWorkflowAgentAuthorName:
"""
# Create workflow with executor that emits AgentResponseUpdate without author_name
executor1 = SimpleExecutor(id="my_executor_id", response_text="Response", streaming=True)
workflow = WorkflowBuilder().set_start_executor(executor1).build()
workflow = WorkflowBuilder(start_executor=executor1).build()
agent = WorkflowAgent(workflow=workflow, name="Test Agent")
# Collect streaming updates
@@ -830,7 +823,7 @@ class TestWorkflowAgentAuthorName:
await ctx.yield_output(update)
executor = AuthorNameExecutor(id="executor_id")
workflow = WorkflowBuilder().set_start_executor(executor).build()
workflow = WorkflowBuilder(start_executor=executor).build()
agent = WorkflowAgent(workflow=workflow, name="Test Agent")
# Collect streaming updates
@@ -848,7 +841,7 @@ class TestWorkflowAgentAuthorName:
executor1 = SimpleExecutor(id="first_executor", response_text="First")
executor2 = SimpleExecutor(id="second_executor", response_text="Second")
workflow = WorkflowBuilder().set_start_executor(executor1).add_edge(executor1, executor2).build()
workflow = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
agent = WorkflowAgent(workflow=workflow, name="Multi-Executor Agent")
# Collect streaming updates
python/packages/core/tests/workflow/test_workflow_builder.py
+62 -105
View File
@@ -45,7 +45,7 @@ def test_builder_accepts_agents_directly():
agent1 = DummyAgent(id="agent1", name="writer")
agent2 = DummyAgent(id="agent2", name="reviewer")
wf = WorkflowBuilder().set_start_executor(agent1).add_edge(agent1, agent2).build()
wf = WorkflowBuilder(start_executor=agent1).add_edge(agent1, agent2).build()
# Confirm auto-wrapped executors use agent names as IDs
assert wf.start_executor_id == "writer"
@@ -79,10 +79,8 @@ class MockAggregator(Executor):
def test_workflow_builder_without_start_executor_throws():
"""Test creating a workflow builder without a start executor."""
builder = WorkflowBuilder()
with pytest.raises(ValueError):
builder.build()
with pytest.raises(TypeError):
WorkflowBuilder() # type: ignore[call-arg]
def test_workflow_builder_fluent_api():
@@ -95,13 +93,11 @@ def test_workflow_builder_fluent_api():
executor_f = MockExecutor(id="executor_f")
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
WorkflowBuilder(max_iterations=5, start_executor=executor_a)
.add_edge(executor_a, executor_b)
.add_fan_out_edges(executor_b, [executor_c, executor_d])
.add_fan_in_edges([executor_c, executor_d], executor_e)
.add_chain([executor_e, executor_f])
.set_max_iterations(5)
.build()
)
@@ -115,9 +111,8 @@ def test_add_agent_reuses_same_wrapper():
reuse_agent = DummyAgent(id="agent_reuse", name="reuse_agent")
agent_a = DummyAgent(id="agent_a", name="agent_a")
builder = WorkflowBuilder()
builder = WorkflowBuilder(start_executor=reuse_agent)
# Use the same agent instance in add_edge - should reuse the same wrapper
builder.set_start_executor(reuse_agent)
builder.add_edge(reuse_agent, agent_a)
builder.add_edge(agent_a, reuse_agent)
@@ -133,10 +128,10 @@ def test_add_agent_duplicate_id_raises_error():
"""Test that adding agents with duplicate IDs raises an error."""
agent1 = DummyAgent(id="agent1", name="first")
agent2 = DummyAgent(id="agent2", name="first") # Same name as agent1
builder = WorkflowBuilder()
builder = WorkflowBuilder(start_executor=agent1)
with pytest.raises(ValueError, match="Duplicate executor ID"):
builder.set_start_executor(agent1).add_edge(agent1, agent2).build()
builder.add_edge(agent1, agent2).build()
# Tests for new executor registration patterns
@@ -144,7 +139,7 @@ def test_add_agent_duplicate_id_raises_error():
def test_register_executor_basic():
"""Test basic executor registration with lazy initialization."""
builder = WorkflowBuilder()
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")
@@ -153,14 +148,14 @@ def test_register_executor_basic():
assert result is builder
# Build workflow and verify executor is instantiated
workflow = builder.set_start_executor("TestExecutor").build()
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()
builder = WorkflowBuilder(start_executor="ExecutorA")
# Register multiple executors - IDs must match registered names
builder.register_executor(lambda: MockExecutor(id="ExecutorA"), name="ExecutorA")
@@ -168,13 +163,7 @@ def test_register_multiple_executors():
builder.register_executor(lambda: MockExecutor(id="ExecutorC"), name="ExecutorC")
# Build workflow with edges using registered names
workflow = (
builder
.set_start_executor("ExecutorA")
.add_edge("ExecutorA", "ExecutorB")
.add_edge("ExecutorB", "ExecutorC")
.build()
)
workflow = builder.add_edge("ExecutorA", "ExecutorB").add_edge("ExecutorB", "ExecutorC").build()
# Verify all executors are present
assert "ExecutorA" in workflow.executors
@@ -185,7 +174,7 @@ def test_register_multiple_executors():
def test_register_with_multiple_names():
"""Test registering the same factory function under multiple names."""
builder = WorkflowBuilder()
builder = WorkflowBuilder(start_executor="ExecutorA")
# Register same executor factory under multiple names
# Note: Each call creates a new instance, so IDs won't conflict
@@ -198,7 +187,7 @@ def test_register_with_multiple_names():
builder.register_executor(make_executor, name=["ExecutorA", "ExecutorB"])
# Set up workflow
workflow = builder.set_start_executor("ExecutorA").add_edge("ExecutorA", "ExecutorB").build()
workflow = builder.add_edge("ExecutorA", "ExecutorB").build()
# Verify both executors are present
assert "ExecutorA" in workflow.executors
@@ -208,7 +197,7 @@ def test_register_with_multiple_names():
def test_register_duplicate_name_raises_error():
"""Test that registering duplicate names raises an error."""
builder = WorkflowBuilder()
builder = WorkflowBuilder(start_executor="MyExecutor")
# Register first executor
builder.register_executor(lambda: MockExecutor(id="executor_1"), name="MyExecutor")
@@ -220,12 +209,11 @@ def test_register_duplicate_name_raises_error():
def test_register_duplicate_id_raises_error():
"""Test that registering duplicate id raises an error."""
builder = WorkflowBuilder()
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")
builder.set_start_executor("MyExecutor1")
# Registering second executor with same ID should raise ValueError
with pytest.raises(ValueError, match="Executor with ID 'executor' has already been registered."):
@@ -234,7 +222,7 @@ def test_register_duplicate_id_raises_error():
def test_register_agent_basic():
"""Test basic agent registration with lazy initialization."""
builder = WorkflowBuilder()
builder = WorkflowBuilder(start_executor="TestAgent")
# Register an agent factory
result = builder.register_agent(lambda: DummyAgent(id="agent_test", name="test_agent"), name="TestAgent")
@@ -243,14 +231,14 @@ def test_register_agent_basic():
assert result is builder
# Build workflow and verify agent is wrapped in AgentExecutor
workflow = builder.set_start_executor("TestAgent").build()
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()
builder = WorkflowBuilder(start_executor="ThreadedAgent")
custom_thread = AgentThread()
# Register agent with custom thread
@@ -261,7 +249,7 @@ def test_register_agent_with_thread():
)
# Build workflow and verify agent executor configuration
workflow = builder.set_start_executor("ThreadedAgent").build()
workflow = builder.build()
executor = workflow.executors["threaded_agent"]
assert isinstance(executor, AgentExecutor)
@@ -271,7 +259,7 @@ def test_register_agent_with_thread():
def test_register_agent_duplicate_name_raises_error():
"""Test that registering agents with duplicate names raises an error."""
builder = WorkflowBuilder()
builder = WorkflowBuilder(start_executor="MyAgent")
# Register first agent
builder.register_agent(lambda: DummyAgent(id="agent1", name="first"), name="MyAgent")
@@ -283,14 +271,14 @@ def test_register_agent_duplicate_name_raises_error():
def test_register_and_add_edge_with_strings():
"""Test that registered executors can be connected using string names."""
builder = WorkflowBuilder()
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.set_start_executor("Source").add_edge("Source", "Target").build()
workflow = builder.add_edge("Source", "Target").build()
# Verify edge is created correctly
assert workflow.start_executor_id == "source"
@@ -300,14 +288,14 @@ def test_register_and_add_edge_with_strings():
def test_register_agent_and_add_edge_with_strings():
"""Test that registered agents can be connected using string names."""
builder = WorkflowBuilder()
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.set_start_executor("Writer").add_edge("Writer", "Reviewer").build()
workflow = builder.add_edge("Writer", "Reviewer").build()
# Verify edge is created correctly
assert workflow.start_executor_id == "writer"
@@ -318,7 +306,7 @@ def test_register_agent_and_add_edge_with_strings():
def test_register_with_fan_out_edges():
"""Test using registered names with fan-out edge groups."""
builder = WorkflowBuilder()
builder = WorkflowBuilder(start_executor="Source")
# Register executors - IDs must match registered names
builder.register_executor(lambda: MockExecutor(id="Source"), name="Source")
@@ -326,7 +314,7 @@ def test_register_with_fan_out_edges():
builder.register_executor(lambda: MockExecutor(id="Target2"), name="Target2")
# Add fan-out edges using registered names
workflow = builder.set_start_executor("Source").add_fan_out_edges("Source", ["Target1", "Target2"]).build()
workflow = builder.add_fan_out_edges("Source", ["Target1", "Target2"]).build()
# Verify all executors are present
assert "Source" in workflow.executors
@@ -336,7 +324,7 @@ def test_register_with_fan_out_edges():
def test_register_with_fan_in_edges():
"""Test using registered names with fan-in edge groups."""
builder = WorkflowBuilder()
builder = WorkflowBuilder(start_executor="Source1")
# Register executors - IDs must match registered names
builder.register_executor(lambda: MockExecutor(id="Source1"), name="Source1")
@@ -345,13 +333,7 @@ def test_register_with_fan_in_edges():
# Add fan-in edges using registered names
# Both Source1 and Source2 need to be reachable, so connect Source1 to Source2
workflow = (
builder
.set_start_executor("Source1")
.add_edge("Source1", "Source2")
.add_fan_in_edges(["Source1", "Source2"], "Aggregator")
.build()
)
workflow = builder.add_edge("Source1", "Source2").add_fan_in_edges(["Source1", "Source2"], "Aggregator").build()
# Verify all executors are present
assert "Source1" in workflow.executors
@@ -361,7 +343,7 @@ def test_register_with_fan_in_edges():
def test_register_with_chain():
"""Test using registered names with add_chain."""
builder = WorkflowBuilder()
builder = WorkflowBuilder(start_executor="Step1")
# Register executors - IDs must match registered names
builder.register_executor(lambda: MockExecutor(id="Step1"), name="Step1")
@@ -369,7 +351,7 @@ def test_register_with_chain():
builder.register_executor(lambda: MockExecutor(id="Step3"), name="Step3")
# Add chain using registered names
workflow = builder.add_chain(["Step1", "Step2", "Step3"]).set_start_executor("Step1").build()
workflow = builder.add_chain(["Step1", "Step2", "Step3"]).build()
# Verify all executors are present
assert "Step1" in workflow.executors
@@ -387,15 +369,12 @@ def test_register_factory_called_only_once():
call_count += 1
return MockExecutor(id="Test")
builder = WorkflowBuilder()
builder = WorkflowBuilder(start_executor="Test")
builder.register_executor(factory, name="Test")
# Factory should not be called yet
assert call_count == 0
# Add edge without building
builder.set_start_executor("Test")
# Factory should still not be called
assert call_count == 0
@@ -409,7 +388,7 @@ def test_register_factory_called_only_once():
def test_mixing_eager_and_lazy_initialization_error():
"""Test that mixing eager executor instances with lazy string names raises appropriate error."""
builder = WorkflowBuilder()
builder = WorkflowBuilder(start_executor="Lazy")
# Create an eager executor instance
eager_executor = MockExecutor(id="eager")
@@ -430,7 +409,7 @@ def test_mixing_eager_and_lazy_initialization_error():
def test_register_with_condition():
"""Test adding edges with conditions using registered names."""
builder = WorkflowBuilder()
builder = WorkflowBuilder(start_executor="Source")
def condition_func(msg: MockMessage) -> bool:
return msg.data > 0
@@ -440,7 +419,7 @@ def test_register_with_condition():
builder.register_executor(lambda: MockExecutor(id="Target"), name="Target")
# Add edge with condition
workflow = builder.set_start_executor("Source").add_edge("Source", "Target", condition=condition_func).build()
workflow = builder.add_edge("Source", "Target", condition=condition_func).build()
# Verify workflow is built correctly
assert "Source" in workflow.executors
@@ -457,14 +436,14 @@ def test_register_agent_creates_unique_instances():
return agent
# Build first workflow
builder1 = WorkflowBuilder()
builder1 = WorkflowBuilder(start_executor="Agent")
builder1.register_agent(agent_factory, name="Agent")
_ = builder1.set_start_executor("Agent").build()
_ = builder1.build()
# Build second workflow
builder2 = WorkflowBuilder()
builder2 = WorkflowBuilder(start_executor="Agent")
builder2.register_agent(agent_factory, name="Agent")
_ = builder2.set_start_executor("Agent").build()
_ = builder2.build()
# Verify that two different agent instances were created
assert len(instance_ids) == 2
@@ -477,11 +456,10 @@ def test_register_agent_creates_unique_instances():
def test_with_output_from_returns_builder():
"""Test that with_output_from returns the builder for method chaining."""
executor_a = MockExecutor(id="executor_a")
builder = WorkflowBuilder()
builder = WorkflowBuilder(output_executors=[executor_a], start_executor=executor_a)
result = builder.with_output_from([executor_a])
assert result is builder
# Verify builder was created with output_executors
assert builder._output_executors == [executor_a]
def test_with_output_from_with_executor_instances():
@@ -490,10 +468,8 @@ def test_with_output_from_with_executor_instances():
executor_b = MockExecutor(id="executor_b")
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
WorkflowBuilder(start_executor=executor_a, output_executors=[executor_b])
.add_edge(executor_a, executor_b)
.with_output_from([executor_b])
.build()
)
@@ -506,9 +482,7 @@ def test_with_output_from_with_agent_instances():
agent_a = DummyAgent(id="agent_a", name="writer")
agent_b = DummyAgent(id="agent_b", name="reviewer")
workflow = (
WorkflowBuilder().set_start_executor(agent_a).add_edge(agent_a, agent_b).with_output_from([agent_b]).build()
)
workflow = WorkflowBuilder(start_executor=agent_a, output_executors=[agent_b]).add_edge(agent_a, agent_b).build()
# Verify that the workflow was built with the agent's name as output executor
assert workflow._output_executors == ["reviewer"] # type: ignore
@@ -516,15 +490,10 @@ def test_with_output_from_with_agent_instances():
def test_with_output_from_with_registered_names():
"""Test with_output_from with registered factory names (strings)."""
workflow = (
WorkflowBuilder()
.register_executor(lambda: MockExecutor(id="ExecutorA"), name="ExecutorAFactory")
.register_executor(lambda: MockExecutor(id="ExecutorB"), name="ExecutorBFactory")
.set_start_executor("ExecutorAFactory")
.add_edge("ExecutorAFactory", "ExecutorBFactory")
.with_output_from(["ExecutorBFactory"])
.build()
)
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()
# Verify that the workflow was built with the correct output executors
assert workflow._output_executors == ["ExecutorB"] # type: ignore
@@ -537,11 +506,9 @@ def test_with_output_from_with_multiple_executors():
executor_c = MockExecutor(id="executor_c")
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
WorkflowBuilder(start_executor=executor_a, output_executors=[executor_a, executor_c])
.add_edge(executor_a, executor_b)
.add_edge(executor_b, executor_c)
.with_output_from([executor_a, executor_c])
.build()
)
@@ -549,51 +516,41 @@ def test_with_output_from_with_multiple_executors():
assert set(workflow._output_executors) == {"executor_a", "executor_c"} # type: ignore
def test_with_output_from_can_be_called_multiple_times():
"""Test that calling with_output_from multiple times overwrites the previous setting."""
def test_with_output_from_can_be_set_to_different_value():
"""Test that output_executors can be set at construction time."""
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b")
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
WorkflowBuilder(start_executor=executor_a, output_executors=[executor_b])
.add_edge(executor_a, executor_b)
.with_output_from([executor_a])
.with_output_from([executor_b]) # This should overwrite the previous setting
.build()
)
# Verify that only the last setting is applied
# Verify that the setting is applied
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."""
workflow = (
WorkflowBuilder()
.register_agent(lambda: DummyAgent(id="agent1", name="writer"), name="WriterAgent")
.register_agent(lambda: DummyAgent(id="agent2", name="reviewer"), name="ReviewerAgent")
.set_start_executor("WriterAgent")
.add_edge("WriterAgent", "ReviewerAgent")
.with_output_from(["ReviewerAgent"])
.build()
)
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()
# Verify that the workflow was built with the agent's resolved name
assert workflow._output_executors == ["reviewer"] # type: ignore
def test_with_output_from_in_fluent_chain():
"""Test that with_output_from works correctly in a fluent builder chain."""
def test_with_output_from_in_constructor():
"""Test that output_executors works correctly when set in the constructor."""
executor_a = MockExecutor(id="executor_a")
executor_b = MockExecutor(id="executor_b")
executor_c = MockExecutor(id="executor_c")
# Build workflow with with_output_from in the middle of the chain
# Build workflow with output_executors in the constructor
workflow = (
WorkflowBuilder()
.set_start_executor(executor_a)
.with_output_from([executor_c]) # Set early in the chain
WorkflowBuilder(start_executor=executor_a, output_executors=[executor_c])
.add_edge(executor_a, executor_b)
.add_edge(executor_b, executor_c)
.build()
@@ -607,13 +564,13 @@ def test_with_output_from_with_invalid_executor_raises_validation_error():
"""Test that with_output_from with an invalid executor raises an error."""
executor_a = MockExecutor(id="executor_a")
builder = WorkflowBuilder().set_start_executor(executor_a)
builder = WorkflowBuilder(start_executor=executor_a, output_executors=[MockExecutor(id="executor_b")])
# Attempting to set output from an executor not in the workflow should raise an error
with pytest.raises(
WorkflowValidationError, match="Output executor 'executor_b' is not present in the workflow graph"
):
builder.with_output_from([MockExecutor(id="executor_b")]).build()
builder.build()
# endregion
python/packages/core/tests/workflow/test_workflow_context.py
+6 -6
View File
@@ -93,7 +93,7 @@ async def test_workflow_context_type_annotations_no_parameter() -> None:
async def func1(text: str, ctx: WorkflowContext) -> None:
await ctx.add_event(_TestEvent())
wf = WorkflowBuilder().set_start_executor(func1).build()
wf = WorkflowBuilder(start_executor=func1).build()
events = await wf.run("hello")
test_events = [e for e in events if isinstance(e, _TestEvent)]
assert len(test_events) == 1
@@ -110,7 +110,7 @@ async def test_workflow_context_type_annotations_no_parameter() -> None:
assert executor1.output_types == []
assert executor1.workflow_output_types == []
wf2 = WorkflowBuilder().set_start_executor(executor1).build()
wf2 = WorkflowBuilder(start_executor=executor1).build()
events2 = await wf2.run("hello")
test_events2 = [e for e in events2 if isinstance(e, _TestEvent)]
assert len(test_events2) == 1
@@ -126,7 +126,7 @@ async def test_workflow_context_type_annotations_message_type_parameter() -> Non
async def func2(text: str, ctx: WorkflowContext) -> None:
await ctx.add_event(_TestEvent(data=text))
wf = WorkflowBuilder().add_edge(func1, func2).set_start_executor(func1).build()
wf = WorkflowBuilder(start_executor=func1).add_edge(func1, func2).build()
events = await wf.run("hello")
test_events = [e for e in events if isinstance(e, _TestEvent)]
assert len(test_events) == 1
@@ -153,7 +153,7 @@ async def test_workflow_context_type_annotations_message_type_parameter() -> Non
assert executor2.output_types == []
assert executor2.workflow_output_types == []
wf2 = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
wf2 = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
events2 = await wf2.run("hello")
test_events2 = [e for e in events2 if isinstance(e, _TestEvent)]
assert len(test_events2) == 1
@@ -171,7 +171,7 @@ async def test_workflow_context_type_annotations_message_and_output_type_paramet
await ctx.add_event(_TestEvent(data=text))
await ctx.yield_output(text)
wf = WorkflowBuilder().add_edge(func1, func2).set_start_executor(func1).build()
wf = WorkflowBuilder(start_executor=func1).add_edge(func1, func2).build()
events = await wf.run("hello")
outputs = events.get_outputs()
assert len(outputs) == 1
@@ -199,7 +199,7 @@ async def test_workflow_context_type_annotations_message_and_output_type_paramet
assert executor2.output_types == []
assert executor2.workflow_output_types == [str]
wf2 = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
wf2 = WorkflowBuilder(start_executor=executor1).add_edge(executor1, executor2).build()
events2 = await wf2.run("hello")
outputs2 = events2.get_outputs()
assert len(outputs2) == 1
python/packages/core/tests/workflow/test_workflow_kwargs.py
+30 -33
View File
@@ -78,7 +78,7 @@ class _KwargsCapturingAgent(BaseAgent):
async def test_sequential_kwargs_flow_to_agent() -> None:
"""Test that kwargs passed to SequentialBuilder workflow flow through to agent."""
agent = _KwargsCapturingAgent(name="seq_agent")
workflow = SequentialBuilder().participants([agent]).build()
workflow = SequentialBuilder(participants=[agent]).build()
custom_data = {"endpoint": "https://api.example.com", "version": "v1"}
user_token = {"user_name": "alice", "access_level": "admin"}
@@ -105,7 +105,7 @@ async def test_sequential_kwargs_flow_to_multiple_agents() -> None:
"""Test that kwargs flow to all agents in a sequential workflow."""
agent1 = _KwargsCapturingAgent(name="agent1")
agent2 = _KwargsCapturingAgent(name="agent2")
workflow = SequentialBuilder().participants([agent1, agent2]).build()
workflow = SequentialBuilder(participants=[agent1, agent2]).build()
custom_data = {"key": "value"}
@@ -123,7 +123,7 @@ async def test_sequential_kwargs_flow_to_multiple_agents() -> None:
async def test_sequential_run_kwargs_flow() -> None:
"""Test that kwargs flow through workflow.run() (non-streaming)."""
agent = _KwargsCapturingAgent(name="run_agent")
workflow = SequentialBuilder().participants([agent]).build()
workflow = SequentialBuilder(participants=[agent]).build()
_ = await workflow.run("test message", custom_data={"test": True})
@@ -141,7 +141,7 @@ async def test_concurrent_kwargs_flow_to_agents() -> None:
"""Test that kwargs flow to all agents in a concurrent workflow."""
agent1 = _KwargsCapturingAgent(name="concurrent1")
agent2 = _KwargsCapturingAgent(name="concurrent2")
workflow = ConcurrentBuilder().participants([agent1, agent2]).build()
workflow = ConcurrentBuilder(participants=[agent1, agent2]).build()
custom_data = {"batch_id": "123"}
user_token = {"user_name": "bob"}
@@ -188,13 +188,11 @@ async def test_groupchat_kwargs_flow_to_agents() -> None:
names = list(state.participants.keys())
return names[(turn_count - 1) % len(names)]
workflow = (
GroupChatBuilder()
.participants([agent1, agent2])
.with_orchestrator(selection_func=simple_selector)
.with_max_rounds(2) # Limit rounds to prevent infinite loop
.build()
)
workflow = GroupChatBuilder(
participants=[agent1, agent2],
max_rounds=2, # Limit rounds to prevent infinite loop
selection_func=simple_selector,
).build()
custom_data = {"session_id": "group123"}
@@ -230,7 +228,7 @@ async def test_kwargs_stored_in_state() -> None:
await ctx.send_message(msgs)
inspector = _StateInspector(id="inspector")
workflow = SequentialBuilder().participants([inspector]).build()
workflow = SequentialBuilder(participants=[inspector]).build()
async for event in workflow.run("test", my_kwarg="my_value", another=123, stream=True):
if event.type == "status" and event.state == WorkflowRunState.IDLE:
@@ -255,7 +253,7 @@ async def test_empty_kwargs_stored_as_empty_dict() -> None:
await ctx.send_message(msgs)
checker = _StateChecker(id="checker")
workflow = SequentialBuilder().participants([checker]).build()
workflow = SequentialBuilder(participants=[checker]).build()
# Run without any kwargs
async for event in workflow.run("test", stream=True):
@@ -275,7 +273,7 @@ async def test_empty_kwargs_stored_as_empty_dict() -> None:
async def test_kwargs_with_none_values() -> None:
"""Test that kwargs with None values are passed through correctly."""
agent = _KwargsCapturingAgent(name="none_test")
workflow = SequentialBuilder().participants([agent]).build()
workflow = SequentialBuilder(participants=[agent]).build()
async for event in workflow.run("test", optional_param=None, other_param="value", stream=True):
if event.type == "status" and event.state == WorkflowRunState.IDLE:
@@ -291,7 +289,7 @@ async def test_kwargs_with_none_values() -> None:
async def test_kwargs_with_complex_nested_data() -> None:
"""Test that complex nested data structures flow through correctly."""
agent = _KwargsCapturingAgent(name="nested_test")
workflow = SequentialBuilder().participants([agent]).build()
workflow = SequentialBuilder(participants=[agent]).build()
complex_data = {
"level1": {
@@ -318,8 +316,8 @@ async def test_kwargs_preserved_across_workflow_reruns() -> None:
agent = _KwargsCapturingAgent(name="rerun_test")
# Build separate workflows for each run to avoid "already running" error
workflow1 = SequentialBuilder().participants([agent]).build()
workflow2 = SequentialBuilder().participants([agent]).build()
workflow1 = SequentialBuilder(participants=[agent]).build()
workflow2 = SequentialBuilder(participants=[agent]).build()
# First run
async for event in workflow1.run("run1", run_id="first", stream=True):
@@ -349,11 +347,10 @@ async def test_handoff_kwargs_flow_to_agents() -> None:
agent2 = _KwargsCapturingAgent(name="specialist")
workflow = (
HandoffBuilder()
HandoffBuilder(termination_condition=lambda conv: len(conv) >= 4)
.participants([agent1, agent2])
.with_start_agent(agent1)
.with_autonomous_mode()
.with_termination_condition(lambda conv: len(conv) >= 4)
.build()
)
@@ -413,7 +410,7 @@ async def test_magentic_kwargs_flow_to_agents() -> None:
agent = _KwargsCapturingAgent(name="agent1")
manager = _MockManager()
workflow = MagenticBuilder().participants([agent]).with_manager(manager=manager).build()
workflow = MagenticBuilder(participants=[agent], manager=manager).build()
custom_data = {"session_id": "magentic123"}
@@ -463,7 +460,7 @@ async def test_magentic_kwargs_stored_in_state() -> None:
agent = _KwargsCapturingAgent(name="agent1")
manager = _MockManager()
magentic_workflow = MagenticBuilder().participants([agent]).with_manager(manager=manager).build()
magentic_workflow = MagenticBuilder(participants=[agent], manager=manager).build()
# Use MagenticWorkflow.run() which goes through the kwargs attachment path
custom_data = {"magentic_key": "magentic_value"}
@@ -485,7 +482,7 @@ async def test_magentic_kwargs_stored_in_state() -> None:
async def test_workflow_as_agent_run_propagates_kwargs_to_underlying_agent() -> None:
"""Test that kwargs passed to workflow_agent.run() flow through to the underlying agents."""
agent = _KwargsCapturingAgent(name="inner_agent")
workflow = SequentialBuilder().participants([agent]).build()
workflow = SequentialBuilder(participants=[agent]).build()
workflow_agent = workflow.as_agent(name="TestWorkflowAgent")
custom_data = {"endpoint": "https://api.example.com", "version": "v1"}
@@ -509,7 +506,7 @@ async def test_workflow_as_agent_run_propagates_kwargs_to_underlying_agent() ->
async def test_workflow_as_agent_run_stream_propagates_kwargs_to_underlying_agent() -> None:
"""Test that kwargs passed to workflow_agent.run() flow through to the underlying agents."""
agent = _KwargsCapturingAgent(name="inner_agent")
workflow = SequentialBuilder().participants([agent]).build()
workflow = SequentialBuilder(participants=[agent]).build()
workflow_agent = workflow.as_agent(name="TestWorkflowAgent")
custom_data = {"session_id": "xyz123"}
@@ -536,7 +533,7 @@ async def test_workflow_as_agent_propagates_kwargs_to_multiple_agents() -> None:
"""Test that kwargs flow to all agents when using workflow.as_agent()."""
agent1 = _KwargsCapturingAgent(name="agent1")
agent2 = _KwargsCapturingAgent(name="agent2")
workflow = SequentialBuilder().participants([agent1, agent2]).build()
workflow = SequentialBuilder(participants=[agent1, agent2]).build()
workflow_agent = workflow.as_agent(name="MultiAgentWorkflow")
custom_data = {"batch_id": "batch-001"}
@@ -553,7 +550,7 @@ async def test_workflow_as_agent_propagates_kwargs_to_multiple_agents() -> None:
async def test_workflow_as_agent_kwargs_with_none_values() -> None:
"""Test that kwargs with None values are passed through correctly via as_agent()."""
agent = _KwargsCapturingAgent(name="none_test_agent")
workflow = SequentialBuilder().participants([agent]).build()
workflow = SequentialBuilder(participants=[agent]).build()
workflow_agent = workflow.as_agent(name="NoneTestWorkflow")
_ = await workflow_agent.run("test", optional_param=None, other_param="value")
@@ -568,7 +565,7 @@ async def test_workflow_as_agent_kwargs_with_none_values() -> None:
async def test_workflow_as_agent_kwargs_with_complex_nested_data() -> None:
"""Test that complex nested data structures flow through correctly via as_agent()."""
agent = _KwargsCapturingAgent(name="nested_agent")
workflow = SequentialBuilder().participants([agent]).build()
workflow = SequentialBuilder(participants=[agent]).build()
workflow_agent = workflow.as_agent(name="NestedDataWorkflow")
complex_data = {
@@ -606,13 +603,13 @@ async def test_subworkflow_kwargs_propagation() -> None:
inner_agent = _KwargsCapturingAgent(name="inner_agent")
# Build the inner (sub) workflow with the agent
inner_workflow = SequentialBuilder().participants([inner_agent]).build()
inner_workflow = SequentialBuilder(participants=[inner_agent]).build()
# Wrap the inner workflow in a WorkflowExecutor so it can be used as a subworkflow
subworkflow_executor = WorkflowExecutor(workflow=inner_workflow, id="subworkflow_executor")
# Build the outer (parent) workflow containing the subworkflow
outer_workflow = SequentialBuilder().participants([subworkflow_executor]).build()
outer_workflow = SequentialBuilder(participants=[subworkflow_executor]).build()
# Define kwargs that should propagate to subworkflow
custom_data = {"api_key": "secret123", "endpoint": "https://api.example.com"}
@@ -670,13 +667,13 @@ async def test_subworkflow_kwargs_accessible_via_state() -> None:
# Build inner workflow with State reader
state_reader = _StateReader(id="state_reader")
inner_workflow = SequentialBuilder().participants([state_reader]).build()
inner_workflow = SequentialBuilder(participants=[state_reader]).build()
# Wrap as subworkflow
subworkflow_executor = WorkflowExecutor(workflow=inner_workflow, id="subworkflow")
# Build outer workflow
outer_workflow = SequentialBuilder().participants([subworkflow_executor]).build()
outer_workflow = SequentialBuilder(participants=[subworkflow_executor]).build()
# Run with kwargs
async for event in outer_workflow.run(
@@ -715,15 +712,15 @@ async def test_nested_subworkflow_kwargs_propagation() -> None:
inner_agent = _KwargsCapturingAgent(name="deeply_nested_agent")
# Build inner workflow
inner_workflow = SequentialBuilder().participants([inner_agent]).build()
inner_workflow = SequentialBuilder(participants=[inner_agent]).build()
inner_executor = WorkflowExecutor(workflow=inner_workflow, id="inner_executor")
# Build middle workflow containing inner
middle_workflow = SequentialBuilder().participants([inner_executor]).build()
middle_workflow = SequentialBuilder(participants=[inner_executor]).build()
middle_executor = WorkflowExecutor(workflow=middle_workflow, id="middle_executor")
# Build outer workflow containing middle
outer_workflow = SequentialBuilder().participants([middle_executor]).build()
outer_workflow = SequentialBuilder(participants=[middle_executor]).build()
# Run with kwargs
async for event in outer_workflow.run(
python/packages/core/tests/workflow/test_workflow_observability.py
+11 -12
View File
@@ -268,8 +268,7 @@ async def test_end_to_end_workflow_tracing(span_exporter: InMemorySpanExporter)
# Create workflow with fan-in: executor1 -> [executor2, executor3] -> aggregator
workflow = (
WorkflowBuilder()
.set_start_executor(executor1)
WorkflowBuilder(start_executor=executor1)
.add_fan_out_edges(executor1, [executor2, executor3])
.add_fan_in_edges([executor2, executor3], aggregator)
.build()
@@ -297,11 +296,11 @@ async def test_end_to_end_workflow_tracing(span_exporter: InMemorySpanExporter)
span_exporter.clear()
# Test workflow with name and description - verify OTEL attributes
(
WorkflowBuilder(name="Test Pipeline", description="Test workflow description")
.set_start_executor(MockExecutor("start"))
.build()
)
WorkflowBuilder(
name="Test Pipeline",
description="Test workflow description",
start_executor=MockExecutor("start"),
).build()
build_spans_with_metadata = [s for s in span_exporter.get_finished_spans() if s.name == "workflow.build"]
assert len(build_spans_with_metadata) == 1
@@ -412,7 +411,7 @@ async def test_workflow_error_handling_in_tracing(span_exporter: InMemorySpanExp
raise ValueError("Test error")
failing_executor = FailingExecutor()
workflow = WorkflowBuilder().set_start_executor(failing_executor).build()
workflow = WorkflowBuilder(start_executor=failing_executor).build()
# Run workflow and expect error
with pytest.raises(ValueError, match="Test error"):
@@ -475,10 +474,10 @@ 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 not setting start executor
builder = WorkflowBuilder()
# Test validation error by referencing a non-existent start executor
builder = WorkflowBuilder(start_executor="NonExistent")
with pytest.raises(ValueError, match="Starting executor must be set"):
with pytest.raises(ValueError):
builder.build()
spans = span_exporter.get_finished_spans()
@@ -501,5 +500,5 @@ async def test_workflow_build_error_tracing(span_exporter: InMemorySpanExporter)
error_event = error_events[0]
assert error_event.attributes is not None
assert "Starting executor must be set" in str(error_event.attributes.get("build.error.message"))
assert "starting executor" in str(error_event.attributes.get("build.error.message")).lower()
assert error_event.attributes.get("build.error.type") == "ValueError"
python/packages/core/tests/workflow/test_workflow_states.py
+9 -9
View File
@@ -28,7 +28,7 @@ class FailingExecutor(Executor):
async def test_executor_failed_and_workflow_failed_events_streaming():
failing = FailingExecutor(id="f")
wf: Workflow = WorkflowBuilder().set_start_executor(failing).build()
wf: Workflow = WorkflowBuilder(start_executor=failing).build()
events: list[object] = []
with pytest.raises(RuntimeError, match="boom"):
@@ -86,7 +86,7 @@ async def test_executor_failed_event_from_second_executor_in_chain():
"""Test that executor_failed event is emitted when a non-start executor fails."""
passthrough = PassthroughExecutor(id="passthrough")
failing = FailingExecutor(id="failing")
wf: Workflow = WorkflowBuilder().set_start_executor(passthrough).add_edge(passthrough, failing).build()
wf: Workflow = WorkflowBuilder(start_executor=passthrough).add_edge(passthrough, failing).build()
events: list[object] = []
with pytest.raises(RuntimeError, match="boom"):
@@ -131,7 +131,7 @@ class Requester(Executor):
async def test_idle_with_pending_requests_status_streaming():
simple_executor = SimpleExecutor(id="simple")
requester = Requester(id="req")
wf = WorkflowBuilder().set_start_executor(simple_executor).add_edge(simple_executor, requester).build()
wf = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, requester).build()
events = [ev async for ev in wf.run("start", stream=True)] # Consume stream fully
@@ -153,7 +153,7 @@ class Completer(Executor):
async def test_completed_status_streaming():
c = Completer(id="c")
wf = WorkflowBuilder().set_start_executor(c).build()
wf = WorkflowBuilder(start_executor=c).build()
events = [ev async for ev in wf.run("ok", stream=True)] # no raise
# Last status should be IDLE
status = [e for e in events if isinstance(e, WorkflowEvent) and e.type == "status"]
@@ -163,7 +163,7 @@ async def test_completed_status_streaming():
async def test_started_and_completed_event_origins():
c = Completer(id="c-origin")
wf = WorkflowBuilder().set_start_executor(c).build()
wf = WorkflowBuilder(start_executor=c).build()
events = [ev async for ev in wf.run("payload", stream=True)]
started = next(e for e in events if isinstance(e, WorkflowEvent) and e.type == "started")
@@ -181,21 +181,21 @@ async def test_started_and_completed_event_origins():
async def test_non_streaming_final_state_helpers():
# Completed case
c = Completer(id="c")
wf1 = WorkflowBuilder().set_start_executor(c).build()
wf1 = WorkflowBuilder(start_executor=c).build()
result1: WorkflowRunResult = await wf1.run("done")
assert result1.get_final_state() == WorkflowRunState.IDLE
# Idle-with-pending-request case
simple_executor = SimpleExecutor(id="simple")
requester = Requester(id="req")
wf2 = WorkflowBuilder().set_start_executor(simple_executor).add_edge(simple_executor, requester).build()
wf2 = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, requester).build()
result2: WorkflowRunResult = await wf2.run("start")
assert result2.get_final_state() == WorkflowRunState.IDLE_WITH_PENDING_REQUESTS
async def test_run_includes_status_events_completed():
c = Completer(id="c2")
wf = WorkflowBuilder().set_start_executor(c).build()
wf = WorkflowBuilder(start_executor=c).build()
result: WorkflowRunResult = await wf.run("ok")
timeline = result.status_timeline()
assert timeline, "Expected status timeline in non-streaming run() results"
@@ -205,7 +205,7 @@ async def test_run_includes_status_events_completed():
async def test_run_includes_status_events_idle_with_requests():
simple_executor = SimpleExecutor(id="simple")
requester = Requester(id="req2")
wf = WorkflowBuilder().set_start_executor(simple_executor).add_edge(simple_executor, requester).build()
wf = WorkflowBuilder(start_executor=simple_executor).add_edge(simple_executor, requester).build()
result: WorkflowRunResult = await wf.run("start")
timeline = result.status_timeline()
assert timeline, "Expected status timeline in non-streaming run() results"