diff --git a/python/packages/core/agent_framework/_workflows/_edge.py b/python/packages/core/agent_framework/_workflows/_edge.py index 549bd9c7476..44e6a0a98b7 100644 --- a/python/packages/core/agent_framework/_workflows/_edge.py +++ b/python/packages/core/agent_framework/_workflows/_edge.py @@ -264,7 +264,7 @@ def __init__(self) -> None: assert case.target.id == "json" """ - condition: Callable[[Any], bool] + condition: EdgeCondition target: Executor | SupportsAgentRun @@ -508,14 +508,14 @@ class FanOutEdgeGroup(EdgeGroup): """ selection_func_name: str | None - _selection_func: Callable[[Any, list[str]], list[str]] | None + _selection_func: Callable[[Any, list[str]], list[str] | Awaitable[list[str]]] | None _target_ids: list[str] def __init__( self, source_id: str, target_ids: Sequence[str], - selection_func: Callable[[Any, list[str]], list[str]] | None = None, + selection_func: Callable[[Any, list[str]], list[str] | Awaitable[list[str]]] | None = None, *, selection_func_name: str | None = None, id: str | None = None, @@ -579,7 +579,7 @@ def target_ids(self) -> list[str]: return list(self._target_ids) @property - def selection_func(self) -> Callable[[Any, list[str]], list[str]] | None: + def selection_func(self) -> Callable[[Any, list[str]], list[str] | Awaitable[list[str]]] | None: """Expose the runtime callable used to select active fan-out targets. When no selection function was supplied the property returns `None`, @@ -660,11 +660,11 @@ class SwitchCaseEdgeGroupCase(DictConvertible): target_id: str condition_name: str | None type: str - _condition: Callable[[Any], bool] = field(repr=False, compare=False) + _condition: EdgeCondition = field(repr=False, compare=False) def __init__( self, - condition: Callable[[Any], bool] | None, + condition: EdgeCondition | None, target_id: str, *, condition_name: str | None = None, @@ -702,7 +702,7 @@ def __init__( self.condition_name = condition_name @property - def condition(self) -> Callable[[Any], bool]: + def condition(self) -> EdgeCondition: """Return the predicate associated with this case. The placeholder installed during deserialisation raises a @@ -859,15 +859,36 @@ def __init__( "This may result in unexpected behavior." ) - def selection_func(message: Any, targets: list[str]) -> list[str]: - for case in cases: + def selection_func(message: Any, targets: list[str]) -> list[str] | Awaitable[list[str]]: + for i, case in enumerate(cases): if isinstance(case, SwitchCaseEdgeGroupDefault): return [case.target_id] # Errors raised by a case predicate deliberately surface to the caller, matching # `Edge.should_route`. Swallowing them would route the message to the default # branch, turning a broken predicate into silent misrouting rather than a # visible failure. `FanOutEdgeRunner` records the error on the edge-group span. - if case.condition(message): + result = case.condition(message) + if inspect.isawaitable(result): + + async def _async_eval( + pending_result: Any, + pending_target: str, + remaining_cases: list[SwitchCaseEdgeGroupCase | SwitchCaseEdgeGroupDefault], + ) -> list[str]: + if await pending_result: + return [pending_target] + for rem_case in remaining_cases: + if isinstance(rem_case, SwitchCaseEdgeGroupDefault): + return [rem_case.target_id] + rem_res = rem_case.condition(message) + if inspect.isawaitable(rem_res): + rem_res = await rem_res + if rem_res: + return [rem_case.target_id] + raise RuntimeError("No matching case found in SwitchCaseEdgeGroup") + + return _async_eval(result, case.target_id, list(cases[i + 1 :])) + if result: return [case.target_id] raise RuntimeError("No matching case found in SwitchCaseEdgeGroup") diff --git a/python/packages/core/agent_framework/_workflows/_edge_runner.py b/python/packages/core/agent_framework/_workflows/_edge_runner.py index 89d2ecf3a6b..4e9e8c4121a 100644 --- a/python/packages/core/agent_framework/_workflows/_edge_runner.py +++ b/python/packages/core/agent_framework/_workflows/_edge_runner.py @@ -1,11 +1,12 @@ # Copyright (c) Microsoft. All rights reserved. import asyncio +import inspect import json import logging from abc import ABC, abstractmethod from collections import defaultdict -from collections.abc import Callable, Coroutine +from collections.abc import Awaitable, Callable, Coroutine from typing import Any, cast from ..exceptions import WorkflowCheckpointException @@ -233,7 +234,8 @@ def __init__(self, edge_group: FanOutEdgeGroup, executors: dict[str, Executor]) self._target_ids = edge_group.target_executor_ids self._target_map = {edge.target_id: edge for edge in self._edges} self._selection_func = cast( - Callable[[Any, list[str]], list[str]] | None, getattr(edge_group, "selection_func", None) + Callable[[Any, list[str]], list[str] | Awaitable[list[str]]] | None, + getattr(edge_group, "selection_func", None), ) async def send_message( @@ -278,9 +280,12 @@ async def send_message( }) return False - selection_results = ( - self._selection_func(message.data, self._target_ids) if self._selection_func else self._target_ids - ) + if self._selection_func: + selection_results = self._selection_func(message.data, self._target_ids) + if inspect.isawaitable(selection_results): + selection_results = await selection_results + else: + selection_results = self._target_ids if not self._validate_selection_result(selection_results): span.set_attributes({ OtelAttr.EDGE_GROUP_DELIVERED: False, diff --git a/python/packages/core/agent_framework/_workflows/_workflow_builder.py b/python/packages/core/agent_framework/_workflows/_workflow_builder.py index 1aa75a3405d..3eed9ca800a 100644 --- a/python/packages/core/agent_framework/_workflows/_workflow_builder.py +++ b/python/packages/core/agent_framework/_workflows/_workflow_builder.py @@ -3,7 +3,7 @@ import logging import sys import uuid -from collections.abc import Callable, Sequence +from collections.abc import Awaitable, Callable, Sequence from typing import Any, ClassVar, Literal from .._agents import SupportsAgentRun @@ -426,7 +426,7 @@ def add_multi_selection_edge_group( self, source: Executor | SupportsAgentRun, targets: Sequence[Executor | SupportsAgentRun], - selection_func: Callable[[Any, list[str]], list[str]], + selection_func: Callable[[Any, list[str]], list[str] | Awaitable[list[str]]], ) -> Self: """Add an edge group that represents a multi-selection execution model. diff --git a/python/packages/core/tests/workflow/test_edge.py b/python/packages/core/tests/workflow/test_edge.py index e054a73e2a3..08e3b9c2e7a 100644 --- a/python/packages/core/tests/workflow/test_edge.py +++ b/python/packages/core/tests/workflow/test_edge.py @@ -1,5 +1,6 @@ # Copyright (c) Microsoft. All rights reserved. +import asyncio from dataclasses import dataclass from typing import Any from unittest.mock import patch @@ -1703,4 +1704,341 @@ async def test_switch_case_edge_group_send_message_with_invalid_data() -> None: assert success is False +async def test_switch_case_edge_group_async_condition_matches() -> None: + """An async case condition is awaited and correctly routes the message.""" + source = MockExecutor(id="source_executor") + target1 = MockExecutor(id="target_executor_1") + target2 = MockExecutor(id="target_executor_2") + + async def async_is_negative(message: MockMessage) -> bool: + await asyncio.sleep(0) + return message.data < 0 + + edge_group = SwitchCaseEdgeGroup( + source_id=source.id, + cases=[ + SwitchCaseEdgeGroupCase(condition=async_is_negative, target_id=target1.id), + SwitchCaseEdgeGroupDefault(target_id=target2.id), + ], + ) + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) + + state = State() + ctx = InProcRunnerContext() + + message = WorkflowMessage(data=MockMessage(data=-1), source_id=source.id) + with patch("agent_framework._workflows._edge_runner.EdgeRunner._execute_on_target") as mock_send: + success = await edge_runner.send_message(message, state, ctx) + assert success is True + assert mock_send.call_count == 1 + assert mock_send.call_args[0][0] == target1.id + + +async def test_switch_case_edge_group_async_condition_fallthrough_to_default() -> None: + """An async condition evaluating to False falls through to the default branch.""" + source = MockExecutor(id="source_executor") + target1 = MockExecutor(id="target_executor_1") + target2 = MockExecutor(id="target_executor_2") + + async def async_is_negative(message: MockMessage) -> bool: + await asyncio.sleep(0) + return message.data < 0 + + edge_group = SwitchCaseEdgeGroup( + source_id=source.id, + cases=[ + SwitchCaseEdgeGroupCase(condition=async_is_negative, target_id=target1.id), + SwitchCaseEdgeGroupDefault(target_id=target2.id), + ], + ) + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) + + state = State() + ctx = InProcRunnerContext() + + message = WorkflowMessage(data=MockMessage(data=10), source_id=source.id) + with patch("agent_framework._workflows._edge_runner.EdgeRunner._execute_on_target") as mock_send: + success = await edge_runner.send_message(message, state, ctx) + assert success is True + assert mock_send.call_count == 1 + assert mock_send.call_args[0][0] == target2.id + + +async def test_switch_case_edge_group_mixed_sync_and_async_conditions_order() -> None: + """Conditions are evaluated in order even with mixed sync and async conditions.""" + source = MockExecutor(id="source_executor") + target1 = MockExecutor(id="target_executor_1") + target2 = MockExecutor(id="target_executor_2") + target3 = MockExecutor(id="target_executor_3") + + async def async_is_large(message: MockMessage) -> bool: + return message.data > 100 + + def sync_is_positive(message: MockMessage) -> bool: + return message.data > 0 + + edge_group = SwitchCaseEdgeGroup( + source_id=source.id, + cases=[ + SwitchCaseEdgeGroupCase(condition=async_is_large, target_id=target1.id), + SwitchCaseEdgeGroupCase(condition=sync_is_positive, target_id=target2.id), + SwitchCaseEdgeGroupDefault(target_id=target3.id), + ], + ) + executors: dict[str, Executor] = { + source.id: source, + target1.id: target1, + target2.id: target2, + target3.id: target3, + } + edge_runner = create_edge_runner(edge_group, executors) + + state = State() + ctx = InProcRunnerContext() + + with patch("agent_framework._workflows._edge_runner.EdgeRunner._execute_on_target") as mock_send: + await edge_runner.send_message(WorkflowMessage(data=MockMessage(data=150), source_id=source.id), state, ctx) + assert mock_send.call_args[0][0] == target1.id + + with patch("agent_framework._workflows._edge_runner.EdgeRunner._execute_on_target") as mock_send: + await edge_runner.send_message(WorkflowMessage(data=MockMessage(data=50), source_id=source.id), state, ctx) + assert mock_send.call_args[0][0] == target2.id + + with patch("agent_framework._workflows._edge_runner.EdgeRunner._execute_on_target") as mock_send: + await edge_runner.send_message(WorkflowMessage(data=MockMessage(data=-10), source_id=source.id), state, ctx) + assert mock_send.call_args[0][0] == target3.id + + +async def test_switch_case_edge_group_async_condition_error_propagates() -> None: + """An error inside an async condition propagates and does not silently fall through to default.""" + source = MockExecutor(id="source_executor") + target1 = MockExecutor(id="target_executor_1") + target2 = MockExecutor(id="target_executor_2") + + async def failing_async_condition(message: MockMessage) -> bool: + raise ValueError("async failure") + + edge_group = SwitchCaseEdgeGroup( + source_id=source.id, + cases=[ + SwitchCaseEdgeGroupCase(condition=failing_async_condition, target_id=target1.id), + SwitchCaseEdgeGroupDefault(target_id=target2.id), + ], + ) + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) + + message = WorkflowMessage(data=MockMessage(data=1), source_id=source.id) + with patch("agent_framework._workflows._edge_runner.EdgeRunner._execute_on_target") as mock_send: + with pytest.raises(ValueError, match="async failure"): + await edge_runner.send_message(message, State(), InProcRunnerContext()) + assert mock_send.call_count == 0 + + +async def test_fan_out_edge_group_with_async_selection_func() -> None: + """FanOutEdgeRunner awaits an async selection_func on FanOutEdgeGroup.""" + source = MockExecutor(id="source_executor") + target1 = MockExecutor(id="target_executor_1") + target2 = MockExecutor(id="target_executor_2") + + async def async_selector(data: Any, available_targets: list[str]) -> list[str]: + await asyncio.sleep(0) + return [target1.id] + + edge_group = FanOutEdgeGroup( + source_id=source.id, + target_ids=[target1.id, target2.id], + selection_func=async_selector, + ) + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) + + state = State() + ctx = InProcRunnerContext() + + message = WorkflowMessage(data=MockMessage(data=1), source_id=source.id) + with patch("agent_framework._workflows._edge_runner.EdgeRunner._execute_on_target") as mock_send: + success = await edge_runner.send_message(message, state, ctx) + assert success is True + assert mock_send.call_count == 1 + assert mock_send.call_args[0][0] == target1.id + + +async def test_switch_case_edge_group_async_condition_future_and_custom_awaitable() -> None: + """Async condition results from asyncio.Future and custom awaitable objects are properly awaited.""" + source = MockExecutor(id="source_executor") + target1 = MockExecutor(id="target_executor_1") + target2 = MockExecutor(id="target_executor_2") + target3 = MockExecutor(id="target_executor_3") + + class CustomAwaitableBool: + def __init__(self, value: bool) -> None: + self.value = value + + def __await__(self): + async def _coro() -> bool: + await asyncio.sleep(0) + return self.value + + return _coro().__await__() + + def future_condition(message: MockMessage) -> asyncio.Future[bool]: + loop = asyncio.get_running_loop() + fut: asyncio.Future[bool] = loop.create_future() + fut.set_result(message.data == 100) + return fut + + def custom_awaitable_condition(message: MockMessage) -> CustomAwaitableBool: + return CustomAwaitableBool(message.data == 200) + + edge_group = SwitchCaseEdgeGroup( + source_id=source.id, + cases=[ + SwitchCaseEdgeGroupCase(condition=future_condition, target_id=target1.id), + SwitchCaseEdgeGroupCase(condition=custom_awaitable_condition, target_id=target2.id), + SwitchCaseEdgeGroupDefault(target_id=target3.id), + ], + ) + executors: dict[str, Executor] = { + source.id: source, + target1.id: target1, + target2.id: target2, + target3.id: target3, + } + edge_runner = create_edge_runner(edge_group, executors) + state = State() + ctx = InProcRunnerContext() + + # Matches Future condition + with patch("agent_framework._workflows._edge_runner.EdgeRunner._execute_on_target") as mock_send: + await edge_runner.send_message(WorkflowMessage(data=MockMessage(data=100), source_id=source.id), state, ctx) + assert mock_send.call_args[0][0] == target1.id + + # Matches Custom awaitable condition + with patch("agent_framework._workflows._edge_runner.EdgeRunner._execute_on_target") as mock_send: + await edge_runner.send_message(WorkflowMessage(data=MockMessage(data=200), source_id=source.id), state, ctx) + assert mock_send.call_args[0][0] == target2.id + + # Falls through to default when neither matches + with patch("agent_framework._workflows._edge_runner.EdgeRunner._execute_on_target") as mock_send: + await edge_runner.send_message(WorkflowMessage(data=MockMessage(data=0), source_id=source.id), state, ctx) + assert mock_send.call_args[0][0] == target3.id + + +async def test_switch_case_edge_group_sync_match_skips_later_async_predicate() -> None: + """A synchronous match on an earlier case skips evaluating later async predicates.""" + source = MockExecutor(id="source_executor") + target1 = MockExecutor(id="target_executor_1") + target2 = MockExecutor(id="target_executor_2") + target3 = MockExecutor(id="target_executor_3") + + async_called = False + + async def bomb_condition(message: MockMessage) -> bool: + nonlocal async_called + async_called = True + raise RuntimeError("Should never be called") + + edge_group = SwitchCaseEdgeGroup( + source_id=source.id, + cases=[ + SwitchCaseEdgeGroupCase(condition=lambda msg: msg.data == 1, target_id=target1.id), + SwitchCaseEdgeGroupCase(condition=bomb_condition, target_id=target2.id), + SwitchCaseEdgeGroupDefault(target_id=target3.id), + ], + ) + executors: dict[str, Executor] = { + source.id: source, + target1.id: target1, + target2.id: target2, + target3.id: target3, + } + edge_runner = create_edge_runner(edge_group, executors) + + state = State() + ctx = InProcRunnerContext() + + with patch("agent_framework._workflows._edge_runner.EdgeRunner._execute_on_target") as mock_send: + success = await edge_runner.send_message( + WorkflowMessage(data=MockMessage(data=1), source_id=source.id), state, ctx + ) + assert success is True + assert mock_send.call_args[0][0] == target1.id + assert async_called is False + + +async def test_switch_case_edge_group_async_condition_cancellation() -> None: + """Cancellation during async condition evaluation propagates CancelledError and skips next cases.""" + source = MockExecutor(id="source_executor") + target1 = MockExecutor(id="target_executor_1") + target2 = MockExecutor(id="target_executor_2") + target3 = MockExecutor(id="target_executor_3") + + case2_called = False + + async def cancelling_condition(message: MockMessage) -> bool: + task = asyncio.current_task() + if task: + task.cancel() + await asyncio.sleep(0.01) + return False + + def case2_condition(message: MockMessage) -> bool: + nonlocal case2_called + case2_called = True + return True + + edge_group = SwitchCaseEdgeGroup( + source_id=source.id, + cases=[ + SwitchCaseEdgeGroupCase(condition=cancelling_condition, target_id=target1.id), + SwitchCaseEdgeGroupCase(condition=case2_condition, target_id=target2.id), + SwitchCaseEdgeGroupDefault(target_id=target3.id), + ], + ) + executors: dict[str, Executor] = { + source.id: source, + target1.id: target1, + target2.id: target2, + target3.id: target3, + } + edge_runner = create_edge_runner(edge_group, executors) + + with patch("agent_framework._workflows._edge_runner.EdgeRunner._execute_on_target") as mock_send: + with pytest.raises(asyncio.CancelledError): + await edge_runner.send_message( + WorkflowMessage(data=MockMessage(data=1), source_id=source.id), State(), InProcRunnerContext() + ) + assert mock_send.call_count == 0 + assert case2_called is False + + +async def test_fan_out_edge_group_with_async_selection_func_invalid_result_raises() -> None: + """FanOutEdgeRunner raises RuntimeError when an async selection_func returns invalid target IDs.""" + source = MockExecutor(id="source_executor") + target1 = MockExecutor(id="target_executor_1") + target2 = MockExecutor(id="target_executor_2") + + async def async_invalid_selector(data: Any, available_targets: list[str]) -> list[str]: + await asyncio.sleep(0) + return [target1.id, "invalid_target_id"] + + edge_group = FanOutEdgeGroup( + source_id=source.id, + target_ids=[target1.id, target2.id], + selection_func=async_invalid_selector, + ) + executors: dict[str, Executor] = {source.id: source, target1.id: target1, target2.id: target2} + edge_runner = create_edge_runner(edge_group, executors) + + state = State() + ctx = InProcRunnerContext() + + message = WorkflowMessage(data=MockMessage(data=1), source_id=source.id) + with pytest.raises(RuntimeError): + await edge_runner.send_message(message, state, ctx) + + # endregion SwitchCaseEdgeGroup diff --git a/python/packages/core/tests/workflow/test_workflow.py b/python/packages/core/tests/workflow/test_workflow.py index 76a285e6d93..2ecc00c9d35 100644 --- a/python/packages/core/tests/workflow/test_workflow.py +++ b/python/packages/core/tests/workflow/test_workflow.py @@ -19,7 +19,9 @@ AgentRunInputs, AgentSession, BaseAgent, + Case, Content, + Default, Executor, FileCheckpointStorage, InMemoryCheckpointStorage, @@ -1112,6 +1114,45 @@ async def test_comprehensive_edge_groups_workflow(): assert has_fan_in, "Workflow should contain FanInEdgeGroup" +async def test_workflow_switch_case_with_async_condition() -> None: + """Test that a workflow with switch-case edge group correctly awaits async case conditions.""" + + class LeafExecutor(Executor): + @handler + async def handle(self, message: NumberMessage, ctx: WorkflowContext[Any, str]) -> None: + await ctx.yield_output(f"{self.id}:{message.data}") + + router = IncrementExecutor(id="router", limit=1000, increment=1) + large_handler = LeafExecutor(id="large") + small_handler = LeafExecutor(id="small") + + async def async_is_large(msg: NumberMessage) -> bool: + await asyncio.sleep(0) + return msg.data > 5 + + workflow = ( + WorkflowBuilder(start_executor=router) + .add_switch_case_edge_group( + router, + [ + Case(condition=async_is_large, target=large_handler), + Default(target=small_handler), + ], + ) + .build() + ) + + # 1. Input 10 -> router sends 11 -> async_is_large(11) is True -> routes to large_handler + events_large = await workflow.run(NumberMessage(data=10)) + assert events_large.get_final_state() == WorkflowRunState.IDLE + assert events_large.get_outputs() == ["large:11"] + + # 2. Input 2 -> router sends 3 -> async_is_large(3) is False -> routes to default (small_handler) + events_small = await workflow.run(NumberMessage(data=2)) + assert events_small.get_final_state() == WorkflowRunState.IDLE + assert events_small.get_outputs() == ["small:3"] + + 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