From b28e1db478cee942cb0f53c6ba627573ec0fb0f5 Mon Sep 17 00:00:00 2001 From: Tao Chen Date: Tue, 23 Sep 2025 14:02:00 -0700 Subject: [PATCH] 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 --- .../_workflow/_typing_utils.py | 39 +++-- .../packages/main/tests/workflow/test_edge.py | 94 ++++++++++++ .../main/tests/workflow/test_typing_utils.py | 134 ++++++++++++++++++ .../getting_started/workflow/README.md | 1 + .../aggregate_results_of_different_types.py | 100 +++++++++++++ python/uv.lock | 6 +- 6 files changed, 363 insertions(+), 11 deletions(-) create mode 100644 python/packages/main/tests/workflow/test_typing_utils.py create mode 100644 python/samples/getting_started/workflow/parallelism/aggregate_results_of_different_types.py diff --git a/python/packages/main/agent_framework/_workflow/_typing_utils.py b/python/packages/main/agent_framework/_workflow/_typing_utils.py index 586178d8d6..f085fee5b1 100644 --- a/python/packages/main/agent_framework/_workflow/_typing_utils.py +++ b/python/packages/main/agent_framework/_workflow/_typing_utils.py @@ -1,7 +1,9 @@ # Copyright (c) Microsoft. All rights reserved. import logging +from collections.abc import Mapping from dataclasses import fields, is_dataclass +from types import UnionType from typing import Any, Union, get_args, get_origin logger = logging.getLogger(__name__) @@ -54,7 +56,7 @@ def _coerce_to_type(value: Any, target_type: type) -> Any | None: return None -def is_instance_of(data: Any, target_type: type) -> bool: +def is_instance_of(data: Any, target_type: type | UnionType | Any) -> bool: """Check if the data is an instance of the target type. Args: @@ -77,17 +79,28 @@ def is_instance_of(data: Any, target_type: type) -> bool: # Case 2: target_type is Optional[T] or Union[T1, T2, ...] # Optional[T] is really just as Union[T, None] + if origin is UnionType: + return any(is_instance_of(data, arg) for arg in args) + + # Case 2b: Handle typing.Union (legacy Union syntax) if origin is Union: return any(is_instance_of(data, arg) for arg in args) # Case 3: target_type is a generic type if origin in [list, set]: - return isinstance(data, origin) and all(is_instance_of(item, args[0]) for item in data) # type: ignore + return isinstance(data, origin) and ( + not args or all(any(is_instance_of(item, arg) for arg in args) for item in data) + ) # type: ignore # Case 4: target_type is a tuple if origin is tuple: + if len(args) == 2 and args[1] is Ellipsis: # Tuple[T, ...] case + element_type = args[0] + return isinstance(data, tuple) and all(is_instance_of(item, element_type) for item in data) if len(args) == 1 and args[0] is Ellipsis: # Tuple[...] case return isinstance(data, tuple) + if len(args) == 0: + return isinstance(data, tuple) return ( isinstance(data, tuple) and len(data) == len(args) # type: ignore @@ -96,9 +109,12 @@ def is_instance_of(data: Any, target_type: type) -> bool: # Case 5: target_type is a dict if origin is dict: - return isinstance(data, dict) and all( - is_instance_of(key, args[0]) and is_instance_of(value, args[1]) - for key, value in data.items() # type: ignore + return isinstance(data, dict) and ( + not args + or all( + is_instance_of(key, args[0]) and is_instance_of(value, args[1]) + for key, value in data.items() # type: ignore + ) ) # Case 6: target_type is RequestResponse[T, U] - validate generic parameters @@ -114,10 +130,17 @@ def is_instance_of(data: Any, target_type: type) -> bool: and data.original_request is not None and not is_instance_of(data.original_request, request_type) ): - coerced = _coerce_to_type(data.original_request, request_type) - if coerced is None: + # Checkpoint decoding can leave original_request as a plain mapping. In that + # case we coerce it back into the expected request type so downstream handlers + # and validators still receive a fully typed RequestResponse instance. + original_request = data.original_request + if isinstance(original_request, Mapping): + coerced = _coerce_to_type(dict(original_request), request_type) + if coerced is None or not isinstance(coerced, request_type): + return False + data.original_request = coerced + else: return False - data.original_request = coerced if hasattr(data, "data") and data.data is not None and not is_instance_of(data.data, response_type): return False return True diff --git a/python/packages/main/tests/workflow/test_edge.py b/python/packages/main/tests/workflow/test_edge.py index b0ef12b3cd..b618c03964 100644 --- a/python/packages/main/tests/workflow/test_edge.py +++ b/python/packages/main/tests/workflow/test_edge.py @@ -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 diff --git a/python/packages/main/tests/workflow/test_typing_utils.py b/python/packages/main/tests/workflow/test_typing_utils.py new file mode 100644 index 0000000000..157225096a --- /dev/null +++ b/python/packages/main/tests/workflow/test_typing_utils.py @@ -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 diff --git a/python/samples/getting_started/workflow/README.md b/python/samples/getting_started/workflow/README.md index 8758c53b6d..e34f146643 100644 --- a/python/samples/getting_started/workflow/README.md +++ b/python/samples/getting_started/workflow/README.md @@ -87,6 +87,7 @@ Once comfortable with these, explore the rest of the samples below. | Sample | File | Concepts | |---|---|---| | Concurrent (Fan-out/Fan-in) | [parallelism/fan_out_fan_in_edges.py](./parallelism/fan_out_fan_in_edges.py) | Dispatch to multiple executors and aggregate results | +| Aggregate Results of Different Types | [parallelism/aggregate_results_of_different_types.py](./parallelism/aggregate_results_of_different_types.py) | Handle results of different types from multiple concurrent executors | | Map-Reduce with Visualization | [parallelism/map_reduce_and_visualization.py](./parallelism/map_reduce_and_visualization.py) | Fan-out/fan-in pattern with diagram export | ### state-management diff --git a/python/samples/getting_started/workflow/parallelism/aggregate_results_of_different_types.py b/python/samples/getting_started/workflow/parallelism/aggregate_results_of_different_types.py new file mode 100644 index 0000000000..16628fb54f --- /dev/null +++ b/python/samples/getting_started/workflow/parallelism/aggregate_results_of_different_types.py @@ -0,0 +1,100 @@ +# Copyright (c) Microsoft. All rights reserved. + +import asyncio +import random + +from agent_framework import Executor, WorkflowBuilder, WorkflowCompletedEvent, WorkflowContext, handler + +""" +Sample: Concurrent fan out and fan in with two different tasks that output results of different types. + +Purpose: +Show how to construct a parallel branch pattern in workflows. Demonstrate: +- Fan out by targeting multiple executors from one dispatcher. +- Fan in by collecting a list of results from the executors. +- Simple tracing using AgentRunEvent to observe execution order and progress. + +Prerequisites: +- Familiarity with WorkflowBuilder, executors, edges, events, and streaming runs. +""" + + +class Dispatcher(Executor): + """ + The sole purpose of this decorator is to dispatch the input of the workflow to + other executors. + """ + + @handler + async def handle(self, numbers: list[int], ctx: WorkflowContext[list[int]]): + if not numbers: + raise RuntimeError("Input must be a valid list of integers.") + + await ctx.send_message(numbers) + + +class Average(Executor): + """Calculate the average of a list of integers.""" + + @handler + async def handle(self, numbers: list[int], ctx: WorkflowContext[float]): + average: float = sum(numbers) / len(numbers) + await ctx.send_message(average) + + +class Sum(Executor): + """Calculate the sum of a list of integers.""" + + @handler + async def handle(self, numbers: list[int], ctx: WorkflowContext[int]): + total: int = sum(numbers) + await ctx.send_message(total) + + +class Aggregator(Executor): + """Aggregate the results from the different tasks and emit the `WorkflowCompletedEvent`.""" + + @handler + async def handle(self, results: list[int | float], ctx: WorkflowContext[None]): + """Receive the results from the source executors. + + The framework will automatically collect messages from the source executors + and deliver them as a list. + + Args: + results (list[int | float]): execution results from upstream executors. + The type annotation must be a list of union types that the upstream + executors will produce. + cts (WorkflowContext[None]): A workflow context. + """ + await ctx.add_event(WorkflowCompletedEvent(data=results)) + + +async def main() -> None: + # 1) Create the executors + dispatcher = Dispatcher(id="dispatcher") + average = Average(id="average") + summation = Sum(id="summation") + aggregator = Aggregator(id="aggregator") + + # 2) Build a simple fan out and fan in workflow + workflow = ( + WorkflowBuilder() + .set_start_executor(dispatcher) + .add_fan_out_edges(dispatcher, [average, summation]) + .add_fan_in_edges([average, summation], aggregator) + .build() + ) + + # 3) Run the workflow + completion: WorkflowCompletedEvent | None = None + async for event in workflow.run_stream([random.randint(1, 100) for _ in range(10)]): + if isinstance(event, WorkflowCompletedEvent): + completion = event + + if completion is not None: + print(completion.data) + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/python/uv.lock b/python/uv.lock index 946ba301cc..f7d3ddd404 100644 --- a/python/uv.lock +++ b/python/uv.lock @@ -3578,7 +3578,7 @@ wheels = [ [[package]] name = "posthog" -version = "6.7.5" +version = "6.7.6" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "backoff", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, @@ -3588,9 +3588,9 @@ dependencies = [ { name = "six", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, { name = "typing-extensions", marker = "sys_platform == 'darwin' or sys_platform == 'linux' or sys_platform == 'win32'" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/1f/f3/d9055bcd190980730bdc600318b34290e85fcb0afb1e31f83cdc33f92615/posthog-6.7.5.tar.gz", hash = "sha256:f4f32b4a4b0df531ae8f80f255a33a49e8880c8c1b62712e6b640535e33a905f", size = 118558, upload-time = "2025-09-16T12:40:34.431Z" } +sdist = { url = "https://files.pythonhosted.org/packages/e2/ce/11d6fa30ab517018796e1d675498992da585479e7079770ec8fa99a61561/posthog-6.7.6.tar.gz", hash = "sha256:ee5c5ad04b857d96d9b7a4f715e23916a2f206bfcf25e5a9d328a3d27664b0d3", size = 119129, upload-time = "2025-09-22T18:11:12.365Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/42/b4/b40f8467252b4ff481e54a9767b211b4ff83114e6d0b6f481852d0ef3e46/posthog-6.7.5-py3-none-any.whl", hash = "sha256:95b00f915365939e63fa183635bad1caaf89cf4a24b63c8bb6983f2a22a56cb3", size = 136766, upload-time = "2025-09-16T12:40:32.741Z" }, + { url = "https://files.pythonhosted.org/packages/de/84/586422d8861b5391c8414360b10f603c0b7859bb09ad688e64430ed0df7b/posthog-6.7.6-py3-none-any.whl", hash = "sha256:b09a7e65a042ec416c28874b397d3accae412a80a8b0ef3fa686fbffc99e4d4b", size = 137348, upload-time = "2025-09-22T18:11:10.807Z" }, ] [[package]]