mirror of
https://github.com/microsoft/agent-framework.git
synced 2026-06-16 21:04:09 +08:00
Python: Define Workflow and Executor APIs (#272)
* Workflow init commit * Add samples and clean up * ExecutionContext -> WorkflowContext * Address comments 1 * Fix mypy * flatting folder structure, and rename contexts * Remove add_loop * Add map reduce sample, remove Activation conditions * Add AgentExecutor and allow multiple handlers per executor * Minor improvement * Add RequestInfoExecutor * Add unit tests part 1 * Address comments 2 * Pre-commit update * Add run method and more unit tests * Add xml docs * run_stream -> run_streaming * message_handler -> handler --------- Co-authored-by: Chris <66376200+crickman@users.noreply.github.com> Co-authored-by: Evan Mattson <evan.mattson@microsoft.com>
This commit is contained in:
Tao Chen
committed by
GitHub
Unverified
GitHub
parent
0d4d7abde1
commit
c8694a8c76
@@ -0,0 +1,47 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
from agent_framework.workflow import Executor, WorkflowContext, handler
|
||||
|
||||
from agent_framework_workflow._edge import Edge
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockMessage:
|
||||
"""A mock message for testing purposes."""
|
||||
|
||||
data: Any
|
||||
|
||||
|
||||
class MockExecutor(Executor):
|
||||
"""A mock executor for testing purposes."""
|
||||
|
||||
@handler
|
||||
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext) -> None:
|
||||
"""A mock handler that does nothing."""
|
||||
pass
|
||||
|
||||
|
||||
def test_create_edge():
|
||||
"""Test creating an edge with a source and target executor."""
|
||||
source = MockExecutor(id="source_executor")
|
||||
target = MockExecutor(id="target_executor")
|
||||
|
||||
edge = Edge(source=source, target=target)
|
||||
|
||||
assert edge.source_id == "source_executor"
|
||||
assert edge.target_id == "target_executor"
|
||||
assert edge.id == f"{edge.source_id}{Edge.ID_SEPARATOR}{edge.target_id}"
|
||||
assert (edge.source_id, edge.target_id) == Edge.source_and_target_from_id(edge.id)
|
||||
|
||||
|
||||
def test_edge_can_handle():
|
||||
"""Test creating an edge with a source and target executor."""
|
||||
source = MockExecutor(id="source_executor")
|
||||
target = MockExecutor(id="target_executor")
|
||||
|
||||
edge = Edge(source=source, target=target)
|
||||
|
||||
assert edge.can_handle(MockMessage(data="test"))
|
||||
@@ -0,0 +1,102 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import pytest
|
||||
from agent_framework.workflow import Executor, WorkflowContext, handler
|
||||
|
||||
|
||||
def test_executor_without_handlers():
|
||||
"""Test that an executor without handlers raises an error when trying to run."""
|
||||
|
||||
class MockExecutorWithoutHandlers(Executor):
|
||||
"""A mock executor that does not implement any handlers."""
|
||||
|
||||
pass
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
MockExecutorWithoutHandlers()
|
||||
|
||||
|
||||
def test_executor_handler_without_annotations():
|
||||
"""Test that an executor with one handler without annotations raises an error when trying to run."""
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
|
||||
class MockExecutorWithOneHandlerWithoutAnnotations(Executor): # type: ignore
|
||||
"""A mock executor with one handler that does not implement any annotations."""
|
||||
|
||||
@handler
|
||||
async def handle(self, message, ctx) -> None: # type: ignore
|
||||
"""A mock handler that does not implement any annotations."""
|
||||
pass
|
||||
|
||||
|
||||
def test_executor_invalid_handler_signature():
|
||||
"""Test that an executor with an invalid handler signature raises an error when trying to run."""
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
|
||||
class MockExecutorWithInvalidHandlerSignature(Executor): # type: ignore
|
||||
"""A mock executor with an invalid handler signature."""
|
||||
|
||||
@handler # type: ignore
|
||||
async def handle(self, message, other, ctx) -> None: # type: ignore
|
||||
"""A mock handler with an invalid signature."""
|
||||
pass
|
||||
|
||||
|
||||
def test_executor_with_valid_handlers():
|
||||
"""Test that an executor with valid handlers can be instantiated and run."""
|
||||
|
||||
class MockExecutorWithValidHandlers(Executor): # type: ignore
|
||||
"""A mock executor with valid handlers."""
|
||||
|
||||
@handler
|
||||
async def handle_text(self, text: str, ctx: WorkflowContext) -> None: # type: ignore
|
||||
"""A mock handler with a valid signature."""
|
||||
pass
|
||||
|
||||
@handler
|
||||
async def handle_number(self, number: int, ctx: WorkflowContext) -> None: # type: ignore
|
||||
"""Another mock handler with a valid signature."""
|
||||
pass
|
||||
|
||||
executor = MockExecutorWithValidHandlers()
|
||||
assert executor.id is not None
|
||||
assert len(executor._handlers) == 2 # type: ignore
|
||||
assert executor.can_handle("text") is True
|
||||
assert executor.can_handle(42) is True
|
||||
assert executor.can_handle(3.14) is False
|
||||
|
||||
|
||||
def test_executor_handlers_with_output_types():
|
||||
"""Test that an executor with handlers that specify output types can be instantiated and run."""
|
||||
|
||||
class MockExecutorWithOutputTypes(Executor): # type: ignore
|
||||
"""A mock executor with handlers that specify output types."""
|
||||
|
||||
@handler(output_types=[str])
|
||||
async def handle_string(self, text: str, ctx: WorkflowContext) -> None: # type: ignore
|
||||
"""A mock handler that outputs a string."""
|
||||
pass
|
||||
|
||||
@handler(output_types=[int])
|
||||
async def handle_integer(self, number: int, ctx: WorkflowContext) -> None: # type: ignore
|
||||
"""A mock handler that outputs an integer."""
|
||||
pass
|
||||
|
||||
executor = MockExecutorWithOutputTypes()
|
||||
assert len(executor._handlers) == 2 # type: ignore
|
||||
|
||||
string_handler = executor._handlers[str] # type: ignore
|
||||
assert string_handler is not None
|
||||
assert string_handler._handler_spec is not None # type: ignore
|
||||
assert string_handler._handler_spec["name"] == "handle_string" # type: ignore
|
||||
assert string_handler._handler_spec["message_type"] is str # type: ignore
|
||||
assert string_handler._handler_spec["output_types"] == [str] # type: ignore
|
||||
|
||||
int_handler = executor._handlers[int] # type: ignore
|
||||
assert int_handler is not None
|
||||
assert int_handler._handler_spec is not None # type: ignore
|
||||
assert int_handler._handler_spec["name"] == "handle_integer" # type: ignore
|
||||
assert int_handler._handler_spec["message_type"] is int # type: ignore
|
||||
assert int_handler._handler_spec["output_types"] == [int] # type: ignore
|
||||
@@ -0,0 +1,145 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import asyncio
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
from agent_framework.workflow import Executor, WorkflowCompletedEvent, WorkflowContext, WorkflowEvent, handler
|
||||
|
||||
from agent_framework_workflow._edge import Edge
|
||||
from agent_framework_workflow._runner import Runner
|
||||
from agent_framework_workflow._runner_context import InProcRunnerContext, RunnerContext
|
||||
from agent_framework_workflow._shared_state import SharedState
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockMessage:
|
||||
"""A mock message for testing purposes."""
|
||||
|
||||
data: int
|
||||
|
||||
|
||||
class MockExecutor(Executor):
|
||||
"""A mock executor for testing purposes."""
|
||||
|
||||
@handler(output_types=[MockMessage])
|
||||
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext) -> None:
|
||||
if message.data < 10:
|
||||
await ctx.send_message(MockMessage(data=message.data + 1))
|
||||
else:
|
||||
await ctx.add_event(WorkflowCompletedEvent(data=message.data))
|
||||
|
||||
|
||||
def test_create_runner():
|
||||
"""Test creating a runner with edges and shared state."""
|
||||
executor_a = MockExecutor(id="executor_a")
|
||||
executor_b = MockExecutor(id="executor_b")
|
||||
|
||||
# Create a loop
|
||||
edges = [
|
||||
Edge(source=executor_a, target=executor_b),
|
||||
Edge(source=executor_b, target=executor_a),
|
||||
]
|
||||
|
||||
runner = Runner(edges, shared_state=SharedState(), ctx=InProcRunnerContext())
|
||||
|
||||
assert runner.context is not None and isinstance(runner.context, RunnerContext)
|
||||
|
||||
|
||||
async def test_runner_run_until_convergence():
|
||||
"""Test running the runner with a simple workflow."""
|
||||
executor_a = MockExecutor(id="executor_a")
|
||||
executor_b = MockExecutor(id="executor_b")
|
||||
|
||||
# Create a loop
|
||||
edges = [
|
||||
Edge(source=executor_a, target=executor_b),
|
||||
Edge(source=executor_b, target=executor_a),
|
||||
]
|
||||
|
||||
shared_state = SharedState()
|
||||
ctx = InProcRunnerContext()
|
||||
|
||||
runner = Runner(edges, shared_state, ctx)
|
||||
|
||||
result: int | None = None
|
||||
await executor_a.execute(
|
||||
MockMessage(data=0),
|
||||
WorkflowContext(
|
||||
executor_id=executor_a.id,
|
||||
source_executor_ids=["START"],
|
||||
shared_state=shared_state,
|
||||
runner_context=ctx,
|
||||
),
|
||||
)
|
||||
async for event in runner.run_until_convergence():
|
||||
assert isinstance(event, WorkflowEvent)
|
||||
if isinstance(event, WorkflowCompletedEvent):
|
||||
result = event.data
|
||||
|
||||
assert result is not None and result == 10
|
||||
|
||||
|
||||
async def test_runner_run_until_convergence_not_completed():
|
||||
"""Test running the runner with a simple workflow."""
|
||||
executor_a = MockExecutor(id="executor_a")
|
||||
executor_b = MockExecutor(id="executor_b")
|
||||
|
||||
# Create a loop
|
||||
edges = [
|
||||
Edge(source=executor_a, target=executor_b),
|
||||
Edge(source=executor_b, target=executor_a),
|
||||
]
|
||||
|
||||
shared_state = SharedState()
|
||||
ctx = InProcRunnerContext()
|
||||
|
||||
runner = Runner(edges, shared_state, ctx, max_iterations=5)
|
||||
|
||||
await executor_a.execute(
|
||||
MockMessage(data=0),
|
||||
WorkflowContext(
|
||||
executor_id=executor_a.id,
|
||||
source_executor_ids=["START"],
|
||||
shared_state=shared_state,
|
||||
runner_context=ctx,
|
||||
),
|
||||
)
|
||||
with pytest.raises(RuntimeError, match="Runner did not converge after 5 iterations."):
|
||||
async for event in runner.run_until_convergence():
|
||||
assert not isinstance(event, WorkflowCompletedEvent)
|
||||
|
||||
|
||||
async def test_runner_already_running():
|
||||
"""Test that running the runner while it is already running raises an error."""
|
||||
executor_a = MockExecutor(id="executor_a")
|
||||
executor_b = MockExecutor(id="executor_b")
|
||||
|
||||
# Create a loop
|
||||
edges = [
|
||||
Edge(source=executor_a, target=executor_b),
|
||||
Edge(source=executor_b, target=executor_a),
|
||||
]
|
||||
|
||||
shared_state = SharedState()
|
||||
ctx = InProcRunnerContext()
|
||||
|
||||
runner = Runner(edges, shared_state, ctx)
|
||||
|
||||
await executor_a.execute(
|
||||
MockMessage(data=0),
|
||||
WorkflowContext(
|
||||
executor_id=executor_a.id,
|
||||
source_executor_ids=["START"],
|
||||
shared_state=shared_state,
|
||||
runner_context=ctx,
|
||||
),
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError, match="Runner is already running."):
|
||||
|
||||
async def _run():
|
||||
async for _ in runner.run_until_convergence():
|
||||
pass
|
||||
|
||||
await asyncio.gather(_run(), _run())
|
||||
@@ -0,0 +1,555 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
|
||||
from agent_framework_workflow import (
|
||||
EdgeDuplicationError,
|
||||
Executor,
|
||||
GraphConnectivityError,
|
||||
TypeCompatibilityError,
|
||||
ValidationTypeEnum,
|
||||
WorkflowBuilder,
|
||||
WorkflowContext,
|
||||
WorkflowValidationError,
|
||||
handler,
|
||||
validate_workflow_graph,
|
||||
)
|
||||
from agent_framework_workflow._edge import Edge
|
||||
|
||||
|
||||
class StringExecutor(Executor):
|
||||
@handler(output_types=[str])
|
||||
async def handle_string(self, message: str, ctx: WorkflowContext) -> None:
|
||||
await ctx.send_message(message.upper())
|
||||
|
||||
|
||||
class StringAggregator(Executor):
|
||||
"""A mock executor that aggregates results from multiple executors."""
|
||||
|
||||
@handler(output_types=[str])
|
||||
async def mock_handler(self, messages: list[str], ctx: WorkflowContext) -> None:
|
||||
# This mock simply returns the data incremented by 1
|
||||
await ctx.send_message("Aggregated: " + ", ".join(messages))
|
||||
|
||||
|
||||
class IntExecutor(Executor):
|
||||
@handler(output_types=[int])
|
||||
async def handle_int(self, message: int, ctx: WorkflowContext) -> None:
|
||||
await ctx.send_message(message * 2)
|
||||
|
||||
|
||||
class AnyExecutor(Executor):
|
||||
@handler
|
||||
async def handle_any(self, message: Any, ctx: WorkflowContext) -> None:
|
||||
await ctx.send_message(f"Processed: {message}")
|
||||
|
||||
|
||||
class NoOutputTypesExecutor(Executor):
|
||||
@handler
|
||||
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
|
||||
await ctx.send_message("processed")
|
||||
|
||||
|
||||
class MultiTypeExecutor(Executor):
|
||||
@handler(output_types=[str])
|
||||
async def handle_string(self, message: str, ctx: WorkflowContext) -> None:
|
||||
await ctx.send_message(f"String: {message}")
|
||||
|
||||
@handler(output_types=[int])
|
||||
async def handle_int(self, message: int, ctx: WorkflowContext) -> None:
|
||||
await ctx.send_message(f"Int: {message}")
|
||||
|
||||
|
||||
def test_valid_workflow_passes_validation():
|
||||
executor1 = StringExecutor(id="string_executor")
|
||||
executor2 = StringExecutor(id="string_executor_2")
|
||||
|
||||
# Create a valid workflow
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.add_edge(executor1, executor2)
|
||||
.set_start_executor(executor1)
|
||||
.build() # This should not raise any exceptions
|
||||
)
|
||||
|
||||
assert workflow is not None
|
||||
|
||||
|
||||
def test_edge_duplication_validation_fails():
|
||||
executor1 = StringExecutor(id="executor1")
|
||||
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()
|
||||
|
||||
assert "executor1->executor2" in str(exc_info.value)
|
||||
assert exc_info.value.validation_type == ValidationTypeEnum.EDGE_DUPLICATION
|
||||
|
||||
|
||||
def test_type_compatibility_validation_fails():
|
||||
string_executor = StringExecutor(id="string_executor")
|
||||
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()
|
||||
|
||||
error = exc_info.value
|
||||
assert error.source_executor_id == "string_executor"
|
||||
assert error.target_executor_id == "int_executor"
|
||||
assert error.validation_type == ValidationTypeEnum.TYPE_COMPATIBILITY
|
||||
|
||||
|
||||
def test_type_compatibility_with_any_type_passes():
|
||||
string_executor = StringExecutor(id="string_executor")
|
||||
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()
|
||||
|
||||
assert workflow is not None
|
||||
|
||||
|
||||
def test_type_compatibility_with_no_output_types():
|
||||
no_output_executor = NoOutputTypesExecutor(id="no_output")
|
||||
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()
|
||||
)
|
||||
|
||||
assert workflow is not None
|
||||
|
||||
|
||||
def test_multi_type_executor_compatibility():
|
||||
string_executor = StringExecutor(id="string_executor")
|
||||
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()
|
||||
)
|
||||
|
||||
assert workflow is not None
|
||||
|
||||
|
||||
def test_graph_connectivity_unreachable_executors():
|
||||
executor1 = StringExecutor(id="executor1")
|
||||
executor2 = StringExecutor(id="executor2")
|
||||
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()
|
||||
|
||||
assert "unreachable" in str(exc_info.value).lower()
|
||||
assert "executor3" in str(exc_info.value)
|
||||
assert exc_info.value.validation_type == ValidationTypeEnum.GRAPH_CONNECTIVITY
|
||||
|
||||
|
||||
def test_graph_connectivity_isolated_executors():
|
||||
executor1 = StringExecutor(id="executor1")
|
||||
executor2 = StringExecutor(id="executor2")
|
||||
executor3 = StringExecutor(id="executor3") # This will be isolated
|
||||
|
||||
# Create edges that include an isolated executor (self-loop that's not connected to main graph)
|
||||
edges = [Edge(executor1, executor2), Edge(executor3, executor3)] # Self-loop to include in graph
|
||||
|
||||
with pytest.raises(GraphConnectivityError) as exc_info:
|
||||
validate_workflow_graph(edges, executor1)
|
||||
|
||||
assert "unreachable" in str(exc_info.value).lower()
|
||||
assert "executor3" in str(exc_info.value)
|
||||
|
||||
|
||||
def test_start_executor_not_in_graph():
|
||||
executor1 = StringExecutor(id="executor1")
|
||||
executor2 = StringExecutor(id="executor2")
|
||||
executor3 = StringExecutor(id="executor3") # Not in graph
|
||||
|
||||
with pytest.raises(GraphConnectivityError) as exc_info:
|
||||
WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor3).build()
|
||||
|
||||
assert "not present in the workflow graph" 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)
|
||||
|
||||
|
||||
def test_workflow_validation_error_base_class():
|
||||
error = WorkflowValidationError("Test message", ValidationTypeEnum.EDGE_DUPLICATION)
|
||||
assert str(error) == "[EDGE_DUPLICATION] Test message"
|
||||
assert error.message == "Test message"
|
||||
assert error.validation_type == ValidationTypeEnum.EDGE_DUPLICATION
|
||||
|
||||
|
||||
def test_complex_workflow_validation():
|
||||
# Create a workflow with multiple paths
|
||||
executor1 = StringExecutor(id="executor1")
|
||||
executor2 = MultiTypeExecutor(id="executor2")
|
||||
executor3 = StringExecutor(id="executor3")
|
||||
executor4 = AnyExecutor(id="executor4")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.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()
|
||||
)
|
||||
|
||||
assert workflow is not None
|
||||
|
||||
|
||||
def test_type_compatibility_inheritance():
|
||||
class BaseExecutor(Executor):
|
||||
@handler(output_types=[str])
|
||||
async def handle_base(self, message: str, ctx: WorkflowContext) -> None:
|
||||
await ctx.send_message("base")
|
||||
|
||||
class DerivedExecutor(Executor):
|
||||
@handler(output_types=[str])
|
||||
async def handle_derived(self, message: str, ctx: WorkflowContext) -> None:
|
||||
await ctx.send_message("derived")
|
||||
|
||||
base_executor = BaseExecutor(id="base")
|
||||
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()
|
||||
|
||||
assert workflow is not None
|
||||
|
||||
|
||||
def test_direct_validation_function():
|
||||
executor1 = StringExecutor(id="executor1")
|
||||
executor2 = StringExecutor(id="executor2")
|
||||
edges = [Edge(executor1, executor2)]
|
||||
|
||||
# This should not raise any exceptions
|
||||
validate_workflow_graph(edges, executor1)
|
||||
|
||||
# Test with invalid start executor
|
||||
executor3 = StringExecutor(id="executor3")
|
||||
with pytest.raises(GraphConnectivityError):
|
||||
validate_workflow_graph(edges, executor3)
|
||||
|
||||
|
||||
def test_fan_out_validation():
|
||||
source = StringExecutor(id="source")
|
||||
target1 = StringExecutor(id="target1")
|
||||
target2 = AnyExecutor(id="target2")
|
||||
|
||||
workflow = WorkflowBuilder().add_fan_out_edges(source, [target1, target2]).set_start_executor(source).build()
|
||||
|
||||
assert workflow is not None
|
||||
|
||||
|
||||
def test_fan_in_validation():
|
||||
start_executor = StringExecutor(id="start")
|
||||
source1 = StringExecutor(id="source1")
|
||||
source2 = StringExecutor(id="source2")
|
||||
target = StringAggregator(id="target")
|
||||
|
||||
# Create a proper fan-in by having a start executor that connects to both sources
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.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()
|
||||
)
|
||||
|
||||
assert workflow is not None
|
||||
|
||||
|
||||
def test_chain_validation():
|
||||
executor1 = StringExecutor(id="executor1")
|
||||
executor2 = StringExecutor(id="executor2")
|
||||
executor3 = AnyExecutor(id="executor3")
|
||||
|
||||
workflow = WorkflowBuilder().add_chain([executor1, executor2, executor3]).set_start_executor(executor1).build()
|
||||
|
||||
assert workflow is not None
|
||||
|
||||
|
||||
def test_logging_for_missing_output_types(caplog: Any) -> None:
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
# Create executor without output types
|
||||
no_output_executor = NoOutputTypesExecutor(id="no_output")
|
||||
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()
|
||||
)
|
||||
|
||||
assert workflow is not None
|
||||
assert "has no output type annotations" in caplog.text
|
||||
assert "Consider adding output_types to @handler decorators" in caplog.text
|
||||
|
||||
|
||||
def test_logging_for_missing_input_types(caplog: Any) -> None:
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
class NoInputTypesExecutor(Executor):
|
||||
# Handler without type annotation for input parameter
|
||||
async def handle_message(self, message: Any, ctx: WorkflowContext) -> None:
|
||||
await ctx.send_message("processed")
|
||||
|
||||
def _discover_handlers(self) -> None:
|
||||
# Override to manually register handler without type info
|
||||
self._handlers[str] = self.handle_message
|
||||
|
||||
string_executor = StringExecutor(id="string_executor")
|
||||
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()
|
||||
)
|
||||
|
||||
assert workflow is not None
|
||||
|
||||
|
||||
def test_self_loop_detection_warning(caplog: Any) -> None:
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
executor = StringExecutor(id="self_loop_executor")
|
||||
|
||||
# Create a self-loop
|
||||
workflow = WorkflowBuilder().add_edge(executor, executor).set_start_executor(executor).build()
|
||||
|
||||
assert workflow is not None
|
||||
assert "Self-loop detected" in caplog.text
|
||||
assert "may cause infinite recursion" in caplog.text
|
||||
|
||||
|
||||
def test_handler_validation_basic(caplog: Any) -> None:
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
# Test basic handler validation - ensure the validation code runs without errors
|
||||
start_executor = StringExecutor(id="start")
|
||||
target_executor = StringExecutor(id="target")
|
||||
|
||||
workflow = WorkflowBuilder().add_edge(start_executor, target_executor).set_start_executor(start_executor).build()
|
||||
|
||||
assert workflow is not None
|
||||
# Just ensure the validation runs without errors
|
||||
|
||||
|
||||
def test_dead_end_detection(caplog: Any) -> None:
|
||||
caplog.set_level(logging.INFO)
|
||||
|
||||
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()
|
||||
|
||||
assert workflow is not None
|
||||
assert "Dead-end executors detected" in caplog.text
|
||||
assert "executor2" in caplog.text
|
||||
assert "Verify these are intended as final nodes" in caplog.text
|
||||
|
||||
|
||||
def test_cycle_detection_warning(caplog: Any) -> None:
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
executor1 = StringExecutor(id="executor1")
|
||||
executor2 = StringExecutor(id="executor2")
|
||||
executor3 = StringExecutor(id="executor3")
|
||||
|
||||
# Create a cycle: executor1 -> executor2 -> executor3 -> executor1
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.add_edge(executor1, executor2)
|
||||
.add_edge(executor2, executor3)
|
||||
.add_edge(executor3, executor1)
|
||||
.set_start_executor(executor1)
|
||||
.build()
|
||||
)
|
||||
|
||||
assert workflow is not None
|
||||
assert "Cycle detected in the workflow graph" in caplog.text
|
||||
assert "Ensure proper termination conditions exist" in caplog.text
|
||||
|
||||
|
||||
def test_successful_type_compatibility_logging(caplog: Any) -> None:
|
||||
caplog.set_level(logging.DEBUG)
|
||||
|
||||
executor1 = StringExecutor(id="executor1")
|
||||
executor2 = StringExecutor(id="executor2")
|
||||
|
||||
workflow = WorkflowBuilder().add_edge(executor1, executor2).set_start_executor(executor1).build()
|
||||
|
||||
assert workflow is not None
|
||||
assert "Type compatibility validated for edge" in caplog.text
|
||||
assert "Compatible type pairs" in caplog.text
|
||||
|
||||
|
||||
def test_complex_cycle_detection(caplog: Any) -> None:
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
# Create a more complex graph with multiple cycles
|
||||
executor1 = StringExecutor(id="executor1")
|
||||
executor2 = StringExecutor(id="executor2")
|
||||
executor3 = StringExecutor(id="executor3")
|
||||
executor4 = StringExecutor(id="executor4")
|
||||
|
||||
# Create multiple paths and cycles
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.add_edge(executor1, executor2)
|
||||
.add_edge(executor2, executor3)
|
||||
.add_edge(executor3, executor4)
|
||||
.add_edge(executor4, executor2) # Creates cycle: executor2 -> executor3 -> executor4 -> executor2
|
||||
.set_start_executor(executor1)
|
||||
.build()
|
||||
)
|
||||
|
||||
assert workflow is not None
|
||||
assert "Cycle detected in the workflow graph" in caplog.text
|
||||
|
||||
|
||||
def test_no_cycles_in_simple_chain(caplog: Any) -> None:
|
||||
caplog.set_level(logging.WARNING)
|
||||
|
||||
executor1 = StringExecutor(id="executor1")
|
||||
executor2 = StringExecutor(id="executor2")
|
||||
executor3 = StringExecutor(id="executor3")
|
||||
|
||||
# Simple chain without cycles
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.add_edge(executor1, executor2)
|
||||
.add_edge(executor2, executor3)
|
||||
.set_start_executor(executor1)
|
||||
.build()
|
||||
)
|
||||
|
||||
assert workflow is not None
|
||||
# Should not log cycle detection
|
||||
assert "Cycle detected" not in caplog.text
|
||||
|
||||
|
||||
def test_multiple_dead_ends_detection(caplog: Any) -> None:
|
||||
caplog.set_level(logging.INFO)
|
||||
|
||||
executor1 = StringExecutor(id="executor1")
|
||||
executor2 = StringExecutor(id="executor2") # Dead end
|
||||
executor3 = StringExecutor(id="executor3") # Dead end
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.add_edge(executor1, executor2)
|
||||
.add_edge(executor1, executor3)
|
||||
.set_start_executor(executor1)
|
||||
.build()
|
||||
)
|
||||
|
||||
assert workflow is not None
|
||||
assert "Dead-end executors detected" in caplog.text
|
||||
assert "executor2" in caplog.text and "executor3" in caplog.text
|
||||
|
||||
|
||||
def test_single_executor_workflow(caplog: Any) -> None:
|
||||
caplog.set_level(logging.INFO)
|
||||
|
||||
# Test workflow with minimal structure
|
||||
executor1 = StringExecutor(id="executor1")
|
||||
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()
|
||||
|
||||
assert workflow is not None
|
||||
# Should detect executor2 as dead end
|
||||
assert "Dead-end executors detected" in caplog.text
|
||||
|
||||
|
||||
def test_enhanced_type_compatibility_error_details():
|
||||
string_executor = StringExecutor(id="string_executor")
|
||||
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()
|
||||
|
||||
error = exc_info.value
|
||||
# Verify enhanced error contains detailed type information
|
||||
assert "Source executor outputs types" in str(error)
|
||||
assert "target executor can only handle types" in str(error)
|
||||
assert error.source_types is not None
|
||||
assert error.target_types is not None
|
||||
|
||||
|
||||
def test_union_type_compatibility_validation() -> None:
|
||||
class UnionOutputExecutor(Executor):
|
||||
@handler(output_types=[str, int])
|
||||
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
|
||||
await ctx.send_message("output")
|
||||
|
||||
class UnionInputExecutor(Executor):
|
||||
@handler(output_types=[str])
|
||||
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
|
||||
await ctx.send_message("processed")
|
||||
|
||||
union_output = UnionOutputExecutor(id="union_output")
|
||||
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()
|
||||
|
||||
assert workflow is not None
|
||||
|
||||
|
||||
def test_generic_type_compatibility() -> None:
|
||||
class ListOutputExecutor(Executor):
|
||||
@handler(output_types=[list[str]])
|
||||
async def handle_message(self, message: str, ctx: WorkflowContext) -> None:
|
||||
await ctx.send_message(["output"])
|
||||
|
||||
class ListInputExecutor(Executor):
|
||||
@handler(output_types=[str])
|
||||
async def handle_message(self, message: list[str], ctx: WorkflowContext) -> None:
|
||||
await ctx.send_message("processed")
|
||||
|
||||
list_output = ListOutputExecutor(id="list_output")
|
||||
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()
|
||||
|
||||
assert workflow is not None
|
||||
|
||||
|
||||
def test_validation_enum_usage() -> None:
|
||||
# Test that all validation types use the enum correctly
|
||||
edge_error = EdgeDuplicationError("test->test")
|
||||
assert edge_error.validation_type == ValidationTypeEnum.EDGE_DUPLICATION
|
||||
|
||||
type_error = TypeCompatibilityError("source", "target", [str], [int])
|
||||
assert type_error.validation_type == ValidationTypeEnum.TYPE_COMPATIBILITY
|
||||
|
||||
graph_error = GraphConnectivityError("test message")
|
||||
assert graph_error.validation_type == ValidationTypeEnum.GRAPH_CONNECTIVITY
|
||||
|
||||
# Test enum string representation
|
||||
assert str(ValidationTypeEnum.EDGE_DUPLICATION) == "ValidationTypeEnum.EDGE_DUPLICATION"
|
||||
assert ValidationTypeEnum.EDGE_DUPLICATION.value == "EDGE_DUPLICATION"
|
||||
@@ -0,0 +1,277 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from dataclasses import dataclass
|
||||
|
||||
import pytest
|
||||
from agent_framework.workflow import (
|
||||
Executor,
|
||||
RequestInfoEvent,
|
||||
RequestInfoExecutor,
|
||||
RequestInfoMessage,
|
||||
WorkflowBuilder,
|
||||
WorkflowCompletedEvent,
|
||||
WorkflowContext,
|
||||
WorkflowEvent,
|
||||
handler,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockMessage:
|
||||
"""A mock message for testing purposes."""
|
||||
|
||||
data: int
|
||||
|
||||
|
||||
class MockExecutor(Executor):
|
||||
"""A mock executor for testing purposes."""
|
||||
|
||||
def __init__(self, id: str, limit: int = 10):
|
||||
"""Initialize the mock executor with a limit."""
|
||||
super().__init__(id=id)
|
||||
self.limit = limit
|
||||
|
||||
@handler(output_types=[MockMessage])
|
||||
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext) -> None:
|
||||
if message.data < self.limit:
|
||||
await ctx.send_message(MockMessage(data=message.data + 1))
|
||||
else:
|
||||
await ctx.add_event(WorkflowCompletedEvent(data=message.data))
|
||||
|
||||
|
||||
class MockAggregator(Executor):
|
||||
"""A mock executor that aggregates results from multiple executors."""
|
||||
|
||||
@handler
|
||||
async def mock_handler(self, messages: list[MockMessage], ctx: WorkflowContext) -> None:
|
||||
# This mock simply returns the data incremented by 1
|
||||
await ctx.add_event(WorkflowCompletedEvent(data=sum(msg.data for msg in messages)))
|
||||
|
||||
|
||||
@dataclass
|
||||
class ApprovalMessage:
|
||||
"""A mock message for approval requests."""
|
||||
|
||||
approved: bool
|
||||
|
||||
|
||||
class MockExecutorRequestApproval(Executor):
|
||||
"""A mock executor that simulates a request for approval."""
|
||||
|
||||
@handler(output_types=[RequestInfoMessage])
|
||||
async def mock_handler_a(self, message: MockMessage, ctx: WorkflowContext) -> None:
|
||||
"""A mock handler that requests approval."""
|
||||
await ctx.set_shared_state(self.id, message.data)
|
||||
await ctx.send_message(RequestInfoMessage())
|
||||
|
||||
@handler(output_types=[MockMessage])
|
||||
async def mock_handler_b(self, message: ApprovalMessage, ctx: WorkflowContext) -> None:
|
||||
"""A mock handler that processes the approval response."""
|
||||
data = await ctx.get_shared_state(self.id)
|
||||
if message.approved:
|
||||
await ctx.add_event(WorkflowCompletedEvent(data=data))
|
||||
else:
|
||||
await ctx.send_message(MockMessage(data=data))
|
||||
|
||||
|
||||
async def test_workflow_run_streaming():
|
||||
"""Test the workflow run stream."""
|
||||
executor_a = MockExecutor(id="executor_a")
|
||||
executor_b = MockExecutor(id="executor_b")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(executor_a)
|
||||
.add_edge(executor_a, executor_b)
|
||||
.add_edge(executor_b, executor_a)
|
||||
.build()
|
||||
)
|
||||
|
||||
result: int | None = None
|
||||
async for event in workflow.run_streaming(MockMessage(data=0)):
|
||||
assert isinstance(event, WorkflowEvent)
|
||||
if isinstance(event, WorkflowCompletedEvent):
|
||||
result = event.data
|
||||
|
||||
assert result is not None and result == 10
|
||||
|
||||
|
||||
async def test_workflow_run_stream_not_completed():
|
||||
"""Test the workflow run stream."""
|
||||
executor_a = MockExecutor(id="executor_a")
|
||||
executor_b = MockExecutor(id="executor_b")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(executor_a)
|
||||
.add_edge(executor_a, executor_b)
|
||||
.add_edge(executor_b, executor_a)
|
||||
.set_max_iterations(5)
|
||||
.build()
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
async for _ in workflow.run_streaming(MockMessage(data=0)):
|
||||
pass
|
||||
|
||||
|
||||
async def test_workflow_run():
|
||||
"""Test the workflow run."""
|
||||
executor_a = MockExecutor(id="executor_a")
|
||||
executor_b = MockExecutor(id="executor_b")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(executor_a)
|
||||
.add_edge(executor_a, executor_b)
|
||||
.add_edge(executor_b, executor_a)
|
||||
.build()
|
||||
)
|
||||
|
||||
events = await workflow.run(MockMessage(data=0))
|
||||
completed_event = events.get_completed_event()
|
||||
assert isinstance(completed_event, WorkflowCompletedEvent)
|
||||
assert completed_event.data == 10
|
||||
|
||||
|
||||
async def test_workflow_run_not_completed():
|
||||
"""Test the workflow run."""
|
||||
executor_a = MockExecutor(id="executor_a")
|
||||
executor_b = MockExecutor(id="executor_b")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(executor_a)
|
||||
.add_edge(executor_a, executor_b)
|
||||
.add_edge(executor_b, executor_a)
|
||||
.set_max_iterations(5)
|
||||
.build()
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError):
|
||||
await workflow.run(MockMessage(data=0))
|
||||
|
||||
|
||||
async def test_workflow_send_responses_streaming():
|
||||
"""Test the workflow run with approval."""
|
||||
executor_a = MockExecutor(id="executor_a")
|
||||
executor_b = MockExecutorRequestApproval(id="executor_b")
|
||||
request_info_executor = RequestInfoExecutor()
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(executor_a)
|
||||
.add_edge(executor_a, executor_b)
|
||||
.add_edge(executor_b, executor_a)
|
||||
.add_edge(executor_b, request_info_executor)
|
||||
.add_edge(request_info_executor, executor_b)
|
||||
.build()
|
||||
)
|
||||
|
||||
request_info_event: RequestInfoEvent | None = None
|
||||
async for event in workflow.run_streaming(MockMessage(data=0)):
|
||||
if isinstance(event, RequestInfoEvent):
|
||||
request_info_event = event
|
||||
|
||||
assert request_info_event is not None
|
||||
result: int | None = None
|
||||
async for event in workflow.send_responses_streaming({
|
||||
request_info_event.request_id: ApprovalMessage(approved=True)
|
||||
}):
|
||||
if isinstance(event, WorkflowCompletedEvent):
|
||||
result = event.data
|
||||
|
||||
assert result is not None and result == 1 # The data should be incremented by 1 from the initial message
|
||||
|
||||
|
||||
async def test_workflow_send_responses():
|
||||
"""Test the workflow run with approval."""
|
||||
executor_a = MockExecutor(id="executor_a")
|
||||
executor_b = MockExecutorRequestApproval(id="executor_b")
|
||||
request_info_executor = RequestInfoExecutor()
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(executor_a)
|
||||
.add_edge(executor_a, executor_b)
|
||||
.add_edge(executor_b, executor_a)
|
||||
.add_edge(executor_b, request_info_executor)
|
||||
.add_edge(request_info_executor, executor_b)
|
||||
.build()
|
||||
)
|
||||
|
||||
events = await workflow.run(MockMessage(data=0))
|
||||
request_info_events = events.get_request_info_events()
|
||||
|
||||
assert len(request_info_events) == 1
|
||||
|
||||
result = await workflow.send_responses({request_info_events[0].request_id: ApprovalMessage(approved=True)})
|
||||
|
||||
completed_event = result.get_completed_event()
|
||||
assert isinstance(completed_event, WorkflowCompletedEvent)
|
||||
assert completed_event.data == 1 # The data should be incremented by 1 from the initial message
|
||||
|
||||
|
||||
async def test_fan_out():
|
||||
"""Test a fan-out workflow."""
|
||||
executor_a = MockExecutor(id="executor_a")
|
||||
executor_b = MockExecutor(id="executor_b", limit=1)
|
||||
executor_c = MockExecutor(id="executor_c", limit=2) # This executor will not complete the workflow
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder().set_start_executor(executor_a).add_fan_out_edges(executor_a, [executor_b, executor_c]).build()
|
||||
)
|
||||
|
||||
events = await workflow.run(MockMessage(data=0))
|
||||
|
||||
# Each executor will emit two events: ExecutorInvokeEvent and ExecutorCompletedEvent
|
||||
# executor_b will also emit a WorkflowCompletedEvent
|
||||
assert len(events) == 7
|
||||
|
||||
completed_event = events.get_completed_event()
|
||||
assert completed_event is not None and completed_event.data == 1
|
||||
|
||||
|
||||
async def test_fan_out_multiple_completed_events():
|
||||
"""Test a fan-out workflow with multiple completed events."""
|
||||
executor_a = MockExecutor(id="executor_a")
|
||||
executor_b = MockExecutor(id="executor_b", limit=1)
|
||||
executor_c = MockExecutor(id="executor_c", limit=1)
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder().set_start_executor(executor_a).add_fan_out_edges(executor_a, [executor_b, executor_c]).build()
|
||||
)
|
||||
|
||||
events = await workflow.run(MockMessage(data=0))
|
||||
|
||||
# Each executor will emit two events: ExecutorInvokeEvent and ExecutorCompletedEvent
|
||||
# executor_a and executor_b will also emit a WorkflowCompletedEvent
|
||||
assert len(events) == 8
|
||||
|
||||
with pytest.raises(ValueError):
|
||||
events.get_completed_event()
|
||||
|
||||
|
||||
async def test_fan_in():
|
||||
"""Test a fan-in workflow."""
|
||||
executor_a = MockExecutor(id="executor_a")
|
||||
executor_b = MockExecutor(id="executor_b")
|
||||
executor_c = MockExecutor(id="executor_c")
|
||||
aggregator = MockAggregator(id="aggregator")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_start_executor(executor_a)
|
||||
.add_fan_out_edges(executor_a, [executor_b, executor_c])
|
||||
.add_fan_in_edges([executor_b, executor_c], aggregator)
|
||||
.build()
|
||||
)
|
||||
|
||||
events = await workflow.run(MockMessage(data=0))
|
||||
|
||||
# Each executor will emit two events: ExecutorInvokeEvent and ExecutorCompletedEvent
|
||||
# aggregator will also emit a WorkflowCompletedEvent
|
||||
assert len(events) == 9
|
||||
|
||||
completed_event = events.get_completed_event()
|
||||
assert completed_event is not None and completed_event.data == 4
|
||||
@@ -0,0 +1,65 @@
|
||||
# Copyright (c) Microsoft. All rights reserved.
|
||||
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import pytest
|
||||
from agent_framework.workflow import Executor, WorkflowBuilder, WorkflowContext, handler
|
||||
|
||||
|
||||
@dataclass
|
||||
class MockMessage:
|
||||
"""A mock message for testing purposes."""
|
||||
|
||||
data: Any
|
||||
|
||||
|
||||
class MockExecutor(Executor):
|
||||
"""A mock executor for testing purposes."""
|
||||
|
||||
@handler(output_types=[MockMessage])
|
||||
async def mock_handler(self, message: MockMessage, ctx: WorkflowContext) -> None:
|
||||
"""A mock handler that does nothing."""
|
||||
pass
|
||||
|
||||
|
||||
class MockAggregator(Executor):
|
||||
"""A mock executor that aggregates results from multiple executors."""
|
||||
|
||||
@handler(output_types=[MockMessage])
|
||||
async def mock_handler(self, messages: list[MockMessage], ctx: WorkflowContext) -> None:
|
||||
# This mock simply returns the data incremented by 1
|
||||
pass
|
||||
|
||||
|
||||
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()
|
||||
|
||||
|
||||
def test_workflow_builder_fluent_api():
|
||||
"""Test the fluent API of the workflow builder."""
|
||||
executor_a = MockExecutor(id="executor_a")
|
||||
executor_b = MockExecutor(id="executor_b")
|
||||
executor_c = MockExecutor(id="executor_c")
|
||||
executor_d = MockExecutor(id="executor_d")
|
||||
executor_e = MockAggregator(id="executor_e")
|
||||
executor_f = MockExecutor(id="executor_f")
|
||||
|
||||
workflow = (
|
||||
WorkflowBuilder()
|
||||
.set_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()
|
||||
)
|
||||
|
||||
assert len(workflow.edges) == 6
|
||||
assert workflow.start_executor.id == executor_a.id
|
||||
assert len(workflow.executors) == 6
|
||||
Reference in New Issue
Block a user