Python: Allow union types in FanIn edge group (#868)

* Improve type utils * Add sample * Add Union * Add more test cases * Add more test cases * Fix RequestResponse typing to only coerce mapping original_request --------- Co-authored-by: Evan Mattson <evan.mattson@microsoft.com>
2026-06-16 21:04:09 +08:00 · 2025-09-23 14:02:00 -07:00
parent 2133043f11
commit b28e1db478
6 changed files with 363 additions and 11 deletions
@@ -79,6 +79,33 @@ class MockAggregator(Executor):
        self.call_count += 1
        self.last_message = message

+    @handler
+    async def mock_aggregator_handler_secondary(
+        self,
+        message: list[MockMessageSecondary],
+        ctx: WorkflowContext,
+    ) -> None:
+        """A mock aggregator handler that does nothing."""
+        self.call_count += 1
+        self.last_message = message
+
+
+class MockAggregatorSecondary(Executor):
+    """A mock aggregator that has a handler for a union type for testing purposes."""
+
+    call_count: int = 0
+    last_message: Any = None
+
+    @handler
+    async def mock_aggregator_handler_combine(
+        self,
+        message: list[MockMessage | MockMessageSecondary],
+        ctx: WorkflowContext,
+    ) -> None:
+        """A mock aggregator handler that does nothing."""
+        self.call_count += 1
+        self.last_message = message
+

 # region Edge

@@ -1062,6 +1089,73 @@ async def test_fan_in_edge_group_tracing_type_mismatch(span_exporter) -> None:
    assert span.attributes.get("edge_group.delivery_status") == EdgeGroupDeliveryStatus.DROPPED_TYPE_MISMATCH.value


+async def test_fan_in_edge_group_with_multiple_message_types() -> None:
+    source1 = MockExecutor(id="source_executor_1")
+    source2 = MockExecutor(id="source_executor_2")
+    target = MockAggregatorSecondary(id="target_executor")
+
+    edge_group = FanInEdgeGroup(source_ids=[source1.id, source2.id], target_id=target.id)
+
+    executors: dict[str, Executor] = {source1.id: source1, source2.id: source2, target.id: target}
+    edge_runner = create_edge_runner(edge_group, executors)
+
+    shared_state = SharedState()
+    ctx = InProcRunnerContext()
+
+    data = MockMessage(data="test")
+
+    success = await edge_runner.send_message(
+        Message(data=data, source_id=source1.id),
+        shared_state,
+        ctx,
+    )
+    assert success
+
+    data2 = MockMessageSecondary(data="test")
+    success = await edge_runner.send_message(
+        Message(data=data2, source_id=source2.id),
+        shared_state,
+        ctx,
+    )
+    assert success
+
+
+async def test_fan_in_edge_group_with_multiple_message_types_failed() -> None:
+    source1 = MockExecutor(id="source_executor_1")
+    source2 = MockExecutor(id="source_executor_2")
+    target = MockAggregator(id="target_executor")
+
+    edge_group = FanInEdgeGroup(source_ids=[source1.id, source2.id], target_id=target.id)
+
+    executors: dict[str, Executor] = {source1.id: source1, source2.id: source2, target.id: target}
+    edge_runner = create_edge_runner(edge_group, executors)
+
+    shared_state = SharedState()
+    ctx = InProcRunnerContext()
+
+    data = MockMessage(data="test")
+
+    success = await edge_runner.send_message(
+        Message(data=data, source_id=source1.id),
+        shared_state,
+        ctx,
+    )
+    assert success
+
+    with pytest.raises(RuntimeError):
+        # Although `MockAggregator` can handle `list[MockMessage]` and `list[MockMessageSecondary]`
+        # separately (i.e., it has handlers for each type individually), it cannot handle
+        # `list[MockMessage | MockMessageSecondary]` (a list containing a mix of both types).
+        # With the fan-in edge group, the target executor must handle all message types from the
+        # source executors as a union.
+        data2 = MockMessageSecondary(data="test")
+        _ = await edge_runner.send_message(
+            Message(data=data2, source_id=source2.id),
+            shared_state,
+            ctx,
+        )
+
+
 # endregion FanInEdgeGroup

 # region SwitchCaseEdgeGroup
@@ -0,0 +1,134 @@
+# Copyright (c) Microsoft. All rights reserved.
+
+from dataclasses import dataclass
+from typing import Any, Generic, TypeVar, Union
+
+from agent_framework._workflow import RequestInfoMessage, RequestResponse
+from agent_framework._workflow._typing_utils import is_instance_of
+
+
+def test_basic_types() -> None:
+    """Test basic built-in types."""
+    assert is_instance_of(5, int)
+    assert is_instance_of("hello", str)
+    assert is_instance_of(None, type(None))
+
+
+def test_union_types() -> None:
+    """Test union types (|) and optional types."""
+    assert is_instance_of(5, int | str)
+    assert is_instance_of("hello", int | str)
+    assert is_instance_of(5, Union[int, str])
+    assert not is_instance_of(5.0, int | str)
+
+
+def test_list_types() -> None:
+    """Test list types with various element types."""
+    assert is_instance_of([], list)
+    assert is_instance_of([1, 2, 3], list)
+    assert is_instance_of([1, 2, 3], list[int])
+    assert is_instance_of([1, 2, 3], list[int | str])
+    assert is_instance_of([1, "a", 3], list[int | str])
+    assert is_instance_of([1, "a", 3], list[Union[int, str]])
+    assert not is_instance_of([1, 2.0, 3], dict)
+    assert not is_instance_of([1, 2.0, 3], list[int | str])
+
+
+def test_tuple_types() -> None:
+    """Test tuple types with fixed and variable lengths."""
+    assert is_instance_of((1, "a"), tuple)
+    assert is_instance_of((1, "a"), tuple[int, str])
+    assert is_instance_of((1, "a", 3), tuple[int | str, ...])
+    assert is_instance_of((1, 2.0, "a"), tuple[...])  # type: ignore
+    assert not is_instance_of((1, 2.0, 3), tuple[int | str, ...])
+    assert not is_instance_of((1, 2.0, 3), dict)
+
+
+def test_dict_types() -> None:
+    """Test dictionary types with typed keys and values."""
+    assert is_instance_of({"key": "value"}, dict)
+    assert is_instance_of({"key": "value"}, dict[str, str])
+    assert is_instance_of({"key": 5, "another_key": "value"}, dict[str, int | str])
+    assert not is_instance_of({"key": 5, "another_key": 3.0}, dict[str, int | str])
+    assert not is_instance_of({"key": 5, "another_key": 3.0}, list)
+
+
+def test_set_types() -> None:
+    """Test set types with various element types."""
+    assert is_instance_of({1, 2, 3}, set)
+    assert is_instance_of({1, 2, 3}, set[int])
+    assert is_instance_of({1, 2, 3}, set[int | str])
+    assert is_instance_of({1, "a", 3}, set[int | str])
+    assert is_instance_of({1, "a", 3}, set[Union[int, str]])
+    assert is_instance_of(set(), set[int])
+    assert not is_instance_of({1, 2.0, 3}, set[int | str])
+    assert not is_instance_of({1, 2, 3}, list)
+    assert not is_instance_of({1, 2, 3}, dict)
+
+
+def test_any_type() -> None:
+    """Test Any type - should accept all values."""
+    assert is_instance_of(5, Any)
+    assert is_instance_of("hello", Any)
+    assert is_instance_of([1, 2, 3], Any)
+
+
+def test_nested_types() -> None:
+    """Test complex nested type structures."""
+    assert is_instance_of([{"key": [1, 2]}, {"another_key": [3]}], list[dict[str, list[int]]])
+    assert not is_instance_of([{"key": [1, 2]}, {"another_key": [3.0]}], list[dict[str, list[int]]])
+
+
+def test_custom_type() -> None:
+    """Test custom object type checking."""
+
+    @dataclass
+    class CustomClass:
+        value: int
+
+    instance = CustomClass(10)
+    assert is_instance_of(instance, CustomClass)
+    assert not is_instance_of(instance, dict)
+
+
+def test_request_response_type() -> None:
+    """Test RequestResponse generic type checking."""
+
+    request_instance = RequestResponse[RequestInfoMessage, str](
+        is_handled=False,
+        original_request=RequestInfoMessage(),
+    )
+
+    class CustomRequestInfoMessage(RequestInfoMessage):
+        info: str
+
+    assert is_instance_of(request_instance, RequestResponse[RequestInfoMessage, str])
+    assert not is_instance_of(request_instance, RequestResponse[CustomRequestInfoMessage, str])
+
+
+def test_custom_generic_type() -> None:
+    """Test custom generic type checking."""
+
+    T = TypeVar("T")
+    U = TypeVar("U")
+
+    class CustomClass(Generic[T, U]):
+        def __init__(self, request: T, response: U, extra: Any | None = None) -> None:
+            self.request = request
+            self.response = response
+            self.extra = extra
+
+    instance = CustomClass[int, str](request=5, response="response")
+
+    assert is_instance_of(instance, CustomClass[int, str])
+    # Generic parameters are not strictly enforced at runtime
+    assert is_instance_of(instance, CustomClass[str, str])
+
+
+def test_edge_cases() -> None:
+    """Test edge cases and unusual scenarios."""
+    assert is_instance_of([], list[int])  # Empty list should be valid
+    assert is_instance_of((), tuple[int, ...])  # Empty tuple should be valid
+    assert is_instance_of({}, dict[str, int])  # Empty dict should be valid
+    assert is_instance_of(None, int | None)  # Optional type with None
+    assert not is_instance_of(5, str | None)  # Optional type without matching type