From 0804ca3ec8a84dab3a5eed44cd4a3e547371aa33 Mon Sep 17 00:00:00 2001 From: Kourtney Miranda Date: Fri, 3 Jul 2026 14:39:11 -0700 Subject: [PATCH 1/6] Add support for UnionType in PartialType definition --- burr/integrations/pydantic.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/burr/integrations/pydantic.py b/burr/integrations/pydantic.py index 300cbb6a0..6d36476fd 100644 --- a/burr/integrations/pydantic.py +++ b/burr/integrations/pydantic.py @@ -19,6 +19,7 @@ import copy import inspect +import sys import types import typing from typing import ( @@ -269,7 +270,10 @@ async def async_action_function(state: State, **kwargs) -> State: return decorator -PartialType = Union[Type[pydantic.BaseModel], Type[dict]] +if sys.version_info >= (3, 10): + PartialType = Union[Type[pydantic.BaseModel], Type[dict], types.UnionType] # noqa: E501 +else: + PartialType = Union[Type[pydantic.BaseModel], Type[dict]] PydanticStreamingActionFunctionSync = Callable[ ..., Generator[Tuple[Union[pydantic.BaseModel, dict], Optional[pydantic.BaseModel]], None, None] @@ -290,7 +294,7 @@ async def async_action_function(state: State, **kwargs) -> State: def _validate_and_extract_signature_types_streaming( fn: PydanticStreamingActionFunction, - stream_type: Optional[Union[Type[pydantic.BaseModel], Type[dict]]], + stream_type: Optional[PartialType], state_input_type: Optional[Type[pydantic.BaseModel]] = None, state_output_type: Optional[Type[pydantic.BaseModel]] = None, ) -> Tuple[ @@ -421,3 +425,4 @@ def construct_data(self, state: State) -> StateModel: def construct_state(self, data: StateModel) -> State: return State(model_to_dict(data)) + From b9ca82227d9797aafda7eaa00147def99f56103b Mon Sep 17 00:00:00 2001 From: Kourtney Miranda Date: Mon, 6 Jul 2026 13:53:13 -0700 Subject: [PATCH 2/6] Allow Any type for stream_type in action.py --- burr/core/action.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/burr/core/action.py b/burr/core/action.py index 08771d411..545c9fd01 100644 --- a/burr/core/action.py +++ b/burr/core/action.py @@ -1507,7 +1507,7 @@ def pydantic( writes: List[str], state_input_type: Type["BaseModel"], state_output_type: Type["BaseModel"], - stream_type: Union[Type["BaseModel"], Type[dict]], + stream_type: Union[Type["BaseModel"], Type[dict], Any], tags: Optional[List[str]] = None, ) -> Callable: """Creates a streaming action that uses pydantic models. @@ -1607,3 +1607,4 @@ def create_action(action_: Union[Callable, ActionT], name: str) -> ActionT: f"Object {action_} is not a valid action. Have you decorated it with @action or @streaming_action?" ) return action_.with_name(name) + From 154305d16480a2186f84f16d447b76eda5aa195f Mon Sep 17 00:00:00 2001 From: Kourtney Miranda Date: Mon, 6 Jul 2026 14:00:35 -0700 Subject: [PATCH 3/6] Add tests for union stream type in Pydantic Add regression tests for union stream type support in Pydantic decorators. These tests ensure compatibility with Python 3.10's union type syntax. --- tests/integrations/test_burr_pydantic.py | 111 +++++++++++++++++++++++ 1 file changed, 111 insertions(+) diff --git a/tests/integrations/test_burr_pydantic.py b/tests/integrations/test_burr_pydantic.py index 567f8be65..621165fde 100644 --- a/tests/integrations/test_burr_pydantic.py +++ b/tests/integrations/test_burr_pydantic.py @@ -779,3 +779,114 @@ async def final_result_streamed( assert state.data.count == 20 assert isinstance(result, IntermediateModel) assert result.result == 20 + + +# --- Tests for stream_type union support (Issue #607) --- + + +import sys + + +@pytest.mark.skipif( + sys.version_info < (3, 10), + reason="Union type syntax with | requires Python 3.10+", +) +def test_validate_streaming_signature_accepts_union_stream_type(): + """Regression test: _validate_and_extract_signature_types_streaming should accept + a union stream_type (e.g. ModelA | ModelB) on Python 3.10+. + + Before the fix, PartialType did not include types.UnionType, causing a TypeError + when a union type was passed as stream_type. + """ + from burr.integrations.pydantic import _validate_and_extract_signature_types_streaming + + class ModelA(BaseModel): + value: int + + class ModelB(BaseModel): + label: str + + def dummy_fn( + state: AppStateModel, + ) -> Generator[Tuple[ModelA, Optional[AppStateModel]], None, None]: + yield ModelA(value=1), state + + union_type = ModelA | ModelB # This is types.UnionType on Python 3.10+ + + state_in, state_out, resolved_stream_type = _validate_and_extract_signature_types_streaming( + dummy_fn, + stream_type=union_type, + state_input_type=AppStateModel, + state_output_type=AppStateModel, + ) + assert state_in is AppStateModel + assert state_out is AppStateModel + assert resolved_stream_type is union_type + + +@pytest.mark.skipif( + sys.version_info < (3, 10), + reason="Union type syntax with | requires Python 3.10+", +) +def test_pydantic_streaming_action_accepts_union_stream_type(): + """Regression test: @pydantic_streaming_action should accept a union stream_type + on Python 3.10+. + + This is the user-facing decorator that was broken by issue #607. + """ + + class ResultA(BaseModel): + value: int + + class ResultB(BaseModel): + label: str + + @pydantic_streaming_action( + reads=["count", "times_called"], + writes=["count", "times_called"], + stream_type=ResultA | ResultB, + state_input_type=AppStateModel, + state_output_type=AppStateModel, + ) + def act( + state: AppStateModel, + ) -> Generator[Tuple[ResultA, Optional[AppStateModel]], None, None]: + state.count += 1 + yield ResultA(value=state.count), state + + assert hasattr(act, "bind") + action_fn = getattr(act, FunctionBasedAction.ACTION_FUNCTION, None) + assert action_fn is not None + + +@pytest.mark.skipif( + sys.version_info < (3, 10), + reason="Union type syntax with | requires Python 3.10+", +) +def test_streaming_action_pydantic_decorator_accepts_union_stream_type(): + """Regression test: @streaming_action.pydantic should accept a union stream_type + on Python 3.10+, matching pydantic_streaming_action behavior. + """ + + class ResultA(BaseModel): + value: int + + class ResultB(BaseModel): + label: str + + @streaming_action.pydantic( + reads=["count", "times_called"], + writes=["count", "times_called"], + stream_type=ResultA | ResultB, + state_input_type=AppStateModel, + state_output_type=AppStateModel, + ) + def act( + state: AppStateModel, + ) -> Generator[Tuple[ResultA, Optional[AppStateModel]], None, None]: + state.count += 1 + yield ResultA(value=state.count), state + + assert hasattr(act, "bind") + assert getattr(act, FunctionBasedAction.ACTION_FUNCTION, None) is not None + From fc49b0caae280cb31c14d4394d6ab620c40101dc Mon Sep 17 00:00:00 2001 From: Kourtney Miranda Date: Wed, 29 Jul 2026 15:44:33 -0700 Subject: [PATCH 4/6] fix: validate stream_type unions at decoration time (review on #845) Add _is_valid_stream_type to check that stream_type is a BaseModel subclass, dict, or a union (typing.Union[...] or PEP 604 X | Y) of those, recursing into union members so invalid unions are rejected at decoration time. Keep PartialType covering types.UnionType. Addresses review feedback on PR #845. --- burr/integrations/pydantic.py | 33 ++++++++++++++++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/burr/integrations/pydantic.py b/burr/integrations/pydantic.py index 6d36476fd..b667df72d 100644 --- a/burr/integrations/pydantic.py +++ b/burr/integrations/pydantic.py @@ -292,6 +292,32 @@ async def async_action_function(state: State, **kwargs) -> State: ) +def _is_valid_stream_type(stream_type: typing.Any) -> bool: + """Return whether ``stream_type`` is a valid intermediate-result type for a + streaming pydantic action. + + A valid ``stream_type`` is one of: + + - a ``pydantic.BaseModel`` subclass, + - ``dict`` (or a ``dict`` subclass), used when the partial result is untyped, + - a union of the above, written either as ``typing.Union[A, B]`` or with the + PEP 604 ``A | B`` syntax (Python 3.10+). + + Unions are validated recursively so that a union with an invalid member + (e.g. ``ModelA | int``) is rejected at decoration time rather than silently + accepted and surfacing later during streaming. This also answers the + per-member question raised on issue #607: each member is checked + individually rather than assuming they are all the same. + """ + origin = typing.get_origin(stream_type) + union_type = getattr(types, "UnionType", None) + is_union = origin is Union or (union_type is not None and origin is union_type) + if is_union: + args = typing.get_args(stream_type) + return len(args) > 0 and all(_is_valid_stream_type(arg) for arg in args) + return isinstance(stream_type, type) and issubclass(stream_type, (pydantic.BaseModel, dict)) + + def _validate_and_extract_signature_types_streaming( fn: PydanticStreamingActionFunction, stream_type: Optional[PartialType], @@ -303,6 +329,12 @@ def _validate_and_extract_signature_types_streaming( if stream_type is None: # TODO -- derive from the signature raise ValueError(f"stream_type is required for function: {fn.__qualname__}") + if not _is_valid_stream_type(stream_type): + raise ValueError( + f"stream_type for function: {fn.__qualname__} must be a pydantic.BaseModel " + f"subclass, dict, or a union of those types (e.g. ModelA | ModelB). " + f"Got: {stream_type!r}." + ) if state_input_type is None: # TODO -- derive from the signature raise ValueError(f"state_input_type is required for function: {fn.__qualname__}") @@ -425,4 +457,3 @@ def construct_data(self, state: State) -> StateModel: def construct_state(self, data: StateModel) -> State: return State(model_to_dict(data)) - From 2e13114d9d0960e0cd23624735fc16eb71a817ae Mon Sep 17 00:00:00 2001 From: Kourtney Miranda Date: Wed, 29 Jul 2026 15:50:20 -0700 Subject: [PATCH 5/6] fix: make stream_type hint meaningful (types.UnionType instead of Any) Replace the Any in the streaming_action.pydantic stream_type annotation, which collapsed the Union to Any, with a types.UnionType forward reference so the hint stays meaningful while accepting X | Y. Addresses review feedback on PR #845. --- burr/core/action.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/burr/core/action.py b/burr/core/action.py index 545c9fd01..e434995f4 100644 --- a/burr/core/action.py +++ b/burr/core/action.py @@ -1507,7 +1507,7 @@ def pydantic( writes: List[str], state_input_type: Type["BaseModel"], state_output_type: Type["BaseModel"], - stream_type: Union[Type["BaseModel"], Type[dict], Any], + stream_type: Union[Type["BaseModel"], Type[dict], "types.UnionType"], tags: Optional[List[str]] = None, ) -> Callable: """Creates a streaming action that uses pydantic models. From aa5979f9bd07964fe7ad3bee19cdb7fd9265f8ab Mon Sep 17 00:00:00 2001 From: Kourtney Miranda Date: Wed, 29 Jul 2026 15:54:01 -0700 Subject: [PATCH 6/6] test: add guarding tests for stream_type union validation Add tests that accept both typing.Union[...] and PEP 604 X | Y stream types, and tests that assert invalid unions and non-model/non-dict stream types are rejected at decoration time. The reject tests fail on unmodified main, so they actually guard the fix. Addresses review feedback on PR #845. --- tests/integrations/test_burr_pydantic.py | 132 +++++++++++++++++++++++ 1 file changed, 132 insertions(+) diff --git a/tests/integrations/test_burr_pydantic.py b/tests/integrations/test_burr_pydantic.py index 621165fde..87b28a8b8 100644 --- a/tests/integrations/test_burr_pydantic.py +++ b/tests/integrations/test_burr_pydantic.py @@ -890,3 +890,135 @@ def act( assert hasattr(act, "bind") assert getattr(act, FunctionBasedAction.ACTION_FUNCTION, None) is not None + +@pytest.mark.skipif( + sys.version_info < (3, 10), + reason="Union type syntax with | requires Python 3.10+", +) +def test_pydantic_streaming_action_accepts_pep604_union_stream_type(): + """A PEP 604 union (ModelA | ModelB) of BaseModel subclasses is accepted.""" + + class ResultA(BaseModel): + value: int + + class ResultB(BaseModel): + label: str + + @pydantic_streaming_action( + reads=["count", "times_called"], + writes=["count", "times_called"], + stream_type=ResultA | ResultB, + state_input_type=AppStateModel, + state_output_type=AppStateModel, + ) + def act( + state: AppStateModel, + ) -> Generator[Tuple[ResultA, Optional[AppStateModel]], None, None]: + state.count += 1 + yield ResultA(value=state.count), state + + assert getattr(act, FunctionBasedAction.ACTION_FUNCTION, None) is not None + + +def test_pydantic_streaming_action_accepts_typing_union_stream_type(): + """A typing.Union[ModelA, ModelB] stream_type is accepted (works on Python 3.9+).""" + from typing import Union as TypingUnion + + class ResultA(BaseModel): + value: int + + class ResultB(BaseModel): + label: str + + @pydantic_streaming_action( + reads=["count", "times_called"], + writes=["count", "times_called"], + stream_type=TypingUnion[ResultA, ResultB], + state_input_type=AppStateModel, + state_output_type=AppStateModel, + ) + def act( + state: AppStateModel, + ) -> Generator[Tuple[ResultA, Optional[AppStateModel]], None, None]: + state.count += 1 + yield ResultA(value=state.count), state + + assert getattr(act, FunctionBasedAction.ACTION_FUNCTION, None) is not None + + +@pytest.mark.skipif( + sys.version_info < (3, 10), + reason="Union type syntax with | requires Python 3.10+", +) +def test_pydantic_streaming_action_rejects_union_with_invalid_member(): + """A union with a non-model, non-dict member must be rejected at decoration time. + + This guards the fix: on unmodified main no stream_type validation happens, so + this invalid union is silently accepted there and this test fails without the fix. + """ + + class ResultA(BaseModel): + value: int + + with pytest.raises(ValueError): + + @pydantic_streaming_action( + reads=["count", "times_called"], + writes=["count", "times_called"], + stream_type=ResultA | int, + state_input_type=AppStateModel, + state_output_type=AppStateModel, + ) + def act( + state: AppStateModel, + ) -> Generator[Tuple[ResultA, Optional[AppStateModel]], None, None]: + state.count += 1 + yield ResultA(value=state.count), state + + +def test_pydantic_streaming_action_rejects_typing_union_with_invalid_member(): + """typing.Union with an invalid member must also be rejected at decoration time.""" + from typing import Union as TypingUnion + + class ResultA(BaseModel): + value: int + + with pytest.raises(ValueError): + + @pydantic_streaming_action( + reads=["count", "times_called"], + writes=["count", "times_called"], + stream_type=TypingUnion[ResultA, int], + state_input_type=AppStateModel, + state_output_type=AppStateModel, + ) + def act( + state: AppStateModel, + ) -> Generator[Tuple[ResultA, Optional[AppStateModel]], None, None]: + state.count += 1 + yield ResultA(value=state.count), state + + +def test_pydantic_streaming_action_rejects_non_model_stream_type(): + """A scalar / non-model, non-dict stream_type must be rejected at decoration time. + + Guards the fix: on unmodified main this is silently accepted. + """ + + class ResultA(BaseModel): + value: int + + with pytest.raises(ValueError): + + @pydantic_streaming_action( + reads=["count", "times_called"], + writes=["count", "times_called"], + stream_type=int, + state_input_type=AppStateModel, + state_output_type=AppStateModel, + ) + def act( + state: AppStateModel, + ) -> Generator[Tuple[ResultA, Optional[AppStateModel]], None, None]: + state.count += 1 + yield ResultA(value=state.count), state