11# Copyright (c) Microsoft. All rights reserved.
22
3+ import inspect
34import logging
45import uuid
5- from collections .abc import Callable , Sequence
6+ from collections .abc import Awaitable , Callable , Sequence
67from dataclasses import dataclass , field
7- from typing import Any , ClassVar
8+ from typing import Any , ClassVar , TypeAlias , TypeVar
89
910from ._const import INTERNAL_SOURCE_ID
1011from ._executor import Executor
1112from ._model_utils import DictConvertible , encode_value
1213
1314logger = logging .getLogger (__name__ )
1415
16+ # Type alias for edge condition functions.
17+ # Conditions receive the message data and return bool (sync or async).
18+ EdgeCondition : TypeAlias = Callable [[Any ], bool | Awaitable [bool ]]
19+
1520
1621def _extract_function_name (func : Callable [..., Any ]) -> str :
1722 """Map a Python callable to a concise, human-focused identifier.
@@ -71,26 +76,27 @@ class Edge(DictConvertible):
7176 serialising the edge down to primitives we can reconstruct the topology of
7277 a workflow irrespective of the original Python process.
7378
79+ Edge conditions receive the message data and return a boolean (sync or async).
80+
7481 Examples:
7582 .. code-block:: python
7683
77- edge = Edge(source_id="ingest", target_id="score", condition=lambda payload: payload["ready"])
78- assert edge.should_route({"ready": True}) is True
79- assert edge.should_route({"ready": False}) is False
84+ edge = Edge(source_id="ingest", target_id="score", condition=lambda data: data["ready"])
85+ assert await edge.should_route({"ready": True}) is True
8086 """
8187
8288 ID_SEPARATOR : ClassVar [str ] = "->"
8389
8490 source_id : str
8591 target_id : str
8692 condition_name : str | None
87- _condition : Callable [[ Any ], bool ] | None = field (default = None , repr = False , compare = False )
93+ _condition : EdgeCondition | None = field (default = None , repr = False , compare = False )
8894
8995 def __init__ (
9096 self ,
9197 source_id : str ,
9298 target_id : str ,
93- condition : Callable [[ Any ], bool ] | None = None ,
99+ condition : EdgeCondition | None = None ,
94100 * ,
95101 condition_name : str | None = None ,
96102 ) -> None :
@@ -103,9 +109,9 @@ def __init__(
103109 target_id:
104110 Canonical identifier of the downstream executor instance.
105111 condition:
106- Optional predicate that receives the message payload and returns
107- `True` when the edge should be traversed. When omitted, the edge is
108- considered unconditionally active.
112+ Optional predicate that receives the message data and returns
113+ `True` when the edge should be traversed. Can be sync or async.
114+ When omitted, the edge is unconditionally active.
109115 condition_name:
110116 Optional override that pins a human-friendly name for the condition
111117 when the callable cannot be introspected (for example after
@@ -125,7 +131,9 @@ def __init__(
125131 self .source_id = source_id
126132 self .target_id = target_id
127133 self ._condition = condition
128- self .condition_name = _extract_function_name (condition ) if condition is not None else condition_name
134+ self .condition_name = (
135+ _extract_function_name (condition ) if condition is not None and condition_name is None else condition_name
136+ )
129137
130138 @property
131139 def id (self ) -> str :
@@ -144,25 +152,44 @@ def id(self) -> str:
144152 """
145153 return f"{ self .source_id } { self .ID_SEPARATOR } { self .target_id } "
146154
147- def should_route (self , data : Any ) -> bool :
148- """Evaluate the edge predicate against an incoming payload.
155+ @property
156+ def has_condition (self ) -> bool :
157+ """Check if this edge has a condition.
158+
159+ Returns True if the edge was configured with a condition function.
160+ """
161+ return self ._condition is not None
162+
163+ async def should_route (self , data : Any ) -> bool :
164+ """Evaluate the edge predicate against payload.
149165
150166 When the edge was defined without an explicit predicate the method
151167 returns `True`, signalling an unconditional routing rule. Otherwise the
152168 user-supplied callable decides whether the message should proceed along
153169 this edge. Any exception raised by the callable is deliberately allowed
154170 to surface to the caller to avoid masking logic bugs.
155171
172+ The condition receives the message data and may be sync or async.
173+
174+ Args:
175+ data: The message payload
176+
177+ Returns:
178+ True if the edge should be traversed, False otherwise.
179+
156180 Examples:
157181 .. code-block:: python
158182
159- edge = Edge("stage1", "stage2", condition=lambda payload: payload ["score"] > 0.8)
160- assert edge.should_route({"score": 0.9}) is True
161- assert edge.should_route({"score": 0.4}) is False
183+ edge = Edge("stage1", "stage2", condition=lambda data: data ["score"] > 0.8)
184+ assert await edge.should_route({"score": 0.9}) is True
185+ assert await edge.should_route({"score": 0.4}) is False
162186 """
163187 if self ._condition is None :
164188 return True
165- return self ._condition (data )
189+ result = self ._condition (data )
190+ if inspect .isawaitable (result ):
191+ return bool (await result )
192+ return bool (result )
166193
167194 def to_dict (self ) -> dict [str , Any ]:
168195 """Produce a JSON-serialisable view of the edge metadata.
@@ -281,6 +308,8 @@ class EdgeGroup(DictConvertible):
281308
282309 from builtins import type as builtin_type
283310
311+ _T_EdgeGroup = TypeVar ("_T_EdgeGroup" , bound = "EdgeGroup" )
312+
284313 _TYPE_REGISTRY : ClassVar [dict [str , builtin_type ["EdgeGroup" ]]] = {}
285314
286315 def __init__ (
@@ -363,7 +392,7 @@ def to_dict(self) -> dict[str, Any]:
363392 }
364393
365394 @classmethod
366- def register (cls , subclass : builtin_type ["EdgeGroup" ]) -> builtin_type ["EdgeGroup" ]:
395+ def register (cls , subclass : builtin_type [_T_EdgeGroup ]) -> builtin_type [_T_EdgeGroup ]:
367396 """Register a subclass so deserialisation can recover the right type.
368397
369398 Registration is typically performed via the decorator syntax applied to
@@ -443,12 +472,18 @@ def __init__(
443472 self ,
444473 source_id : str ,
445474 target_id : str ,
446- condition : Callable [[ Any ], bool ] | None = None ,
475+ condition : EdgeCondition | None = None ,
447476 * ,
448477 id : str | None = None ,
449478 ) -> None :
450479 """Create a one-to-one edge group between two executors.
451480
481+ Args:
482+ source_id: The source executor ID.
483+ target_id: The target executor ID.
484+ condition: Optional condition function `(data) -> bool | Awaitable[bool]`.
485+ id: Optional explicit ID for the edge group.
486+
452487 Examples:
453488 .. code-block:: python
454489
0 commit comments