diff --git a/burr/core/action.py b/burr/core/action.py index 08771d411..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]], + stream_type: Union[Type["BaseModel"], Type[dict], "types.UnionType"], 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) + diff --git a/burr/integrations/pydantic.py b/burr/integrations/pydantic.py index 300cbb6a0..b667df72d 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] @@ -288,9 +292,35 @@ 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[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[ @@ -299,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__}") diff --git a/tests/integrations/test_burr_pydantic.py b/tests/integrations/test_burr_pydantic.py index 567f8be65..87b28a8b8 100644 --- a/tests/integrations/test_burr_pydantic.py +++ b/tests/integrations/test_burr_pydantic.py @@ -779,3 +779,246 @@ 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 + + +@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