diff --git a/python/packages/core/agent_framework/_workflows/_edge.py b/python/packages/core/agent_framework/_workflows/_edge.py index 5a65b0563bf..549bd9c7476 100644 --- a/python/packages/core/agent_framework/_workflows/_edge.py +++ b/python/packages/core/agent_framework/_workflows/_edge.py @@ -871,7 +871,8 @@ def selection_func(message: Any, targets: list[str]) -> list[str]: return [case.target_id] raise RuntimeError("No matching case found in SwitchCaseEdgeGroup") - target_ids = [case.target_id for case in cases] + # Several cases may route to the same executor, so keep one edge per target. + target_ids = list(dict.fromkeys(case.target_id for case in cases)) # Call FanOutEdgeGroup constructor directly to avoid type checking issues edges = [Edge(source_id=source_id, target_id=target) for target in target_ids] EdgeGroup.__init__(self, edges, id=id, type=self.__class__.__name__) diff --git a/python/packages/core/tests/workflow/test_edge.py b/python/packages/core/tests/workflow/test_edge.py index 42f8ea7bb3b..2c62e245be1 100644 --- a/python/packages/core/tests/workflow/test_edge.py +++ b/python/packages/core/tests/workflow/test_edge.py @@ -1364,6 +1364,33 @@ def test_switch_case_edge_group() -> None: assert edge_group._selection_func(MockMessage(data=1), [target1.id, target2.id]) == [target2.id] # type: ignore +def test_switch_case_edge_group_cases_share_a_target() -> None: + """Test that cases routing to the same target create one edge per target.""" + source = MockExecutor(id="source_executor") + target1 = MockExecutor(id="target_executor_1") + target2 = MockExecutor(id="target_executor_2") + + edge_group = SwitchCaseEdgeGroup( + source_id=source.id, + cases=[ + SwitchCaseEdgeGroupCase(condition=lambda x: x.data < 0, target_id=target2.id), + SwitchCaseEdgeGroupCase(condition=lambda x: x.data == 0, target_id=target1.id), + SwitchCaseEdgeGroupCase(condition=lambda x: x.data == 1, target_id=target2.id), + SwitchCaseEdgeGroupDefault(target_id=target1.id), + ], + ) + + assert len(edge_group.cases) == 4 + assert [edge.target_id for edge in edge_group.edges] == [target2.id, target1.id] + assert edge_group.target_executor_ids == [target2.id, target1.id] + + targets = [target2.id, target1.id] + assert edge_group._selection_func(MockMessage(data=-1), targets) == [target2.id] # type: ignore + assert edge_group._selection_func(MockMessage(data=0), targets) == [target1.id] # type: ignore + assert edge_group._selection_func(MockMessage(data=1), targets) == [target2.id] # type: ignore + assert edge_group._selection_func(MockMessage(data=2), targets) == [target1.id] # type: ignore + + def test_switch_case_edge_group_invalid_number_of_cases(): """Test creating a switch case edge group with an invalid number of cases.""" source = MockExecutor(id="source_executor") diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 83735c95f14..76a285e6d93 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -1112,6 +1112,37 @@ async def test_comprehensive_edge_groups_workflow(): assert has_fan_in, "Workflow should contain FanInEdgeGroup" +async def test_switch_case_cases_can_share_a_target(): + """Test that two switch-case branches can route to the same executor.""" + from agent_framework import Case, Default + + class TaggingExecutor(Executor): + @handler + async def tag(self, message: NumberMessage, ctx: WorkflowContext[Any, str]) -> None: + await ctx.yield_output(f"{self.id}:{message.data}") + + router = IncrementExecutor(id="router", limit=1000, increment=0) + shared = TaggingExecutor(id="shared") + fallback = TaggingExecutor(id="fallback") + + workflow = ( + WorkflowBuilder(start_executor=router) + .add_switch_case_edge_group( + router, + [ + Case(condition=lambda msg: msg.data == 1, target=shared), + Case(condition=lambda msg: msg.data == 2, target=shared), + Default(target=fallback), + ], + ) + .build() + ) + + assert (await workflow.run(NumberMessage(data=1))).get_outputs() == ["shared:1"] + assert (await workflow.run(NumberMessage(data=2))).get_outputs() == ["shared:2"] + assert (await workflow.run(NumberMessage(data=3))).get_outputs() == ["fallback:3"] + + async def test_workflow_with_simple_cycle_and_exit_condition(): """Test a simpler workflow with a cycle that has a clear exit condition."""