Python: Support async conditions in switch-case edge groups

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")
 

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,

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.