Skip to content

fix(pydantic): allow union of BaseModel types as stream_type#754

Open
Ghraven wants to merge 2 commits into
apache:mainfrom
Ghraven:fix/stream-type-union-basemodel
Open

fix(pydantic): allow union of BaseModel types as stream_type#754
Ghraven wants to merge 2 commits into
apache:mainfrom
Ghraven:fix/stream-type-union-basemodel

Conversation

@Ghraven
Copy link
Copy Markdown

@Ghraven Ghraven commented May 3, 2026

Fixes #607

What

Broadens the stream_type parameter type annotation in streaming_action.pydantic() and pydantic_streaming_action() to accept a union of pydantic BaseModel types in addition to a single Type[BaseModel] or Type[dict].

Why

Users who pass a union of models (e.g. MyModel1 | MyModel2) hit a static type error at the call site because the parameter was typed as Union[Type[BaseModel], Type[dict]], which only accepts a single concrete model type. The runtime validation logic in _validate_and_extract_signature_types_streaming() already handles union types correctly — this was a type annotation gap only.

Changes

  • burr/core/action.py: stream_type parameter in streaming_action.pydantic() widened to accept object (covers union types, which are not expressible as Type[X]).
  • burr/integrations/pydantic.py: Same change to PartialType alias and the validator function signature.
  • Docstring updated to mention union support.

Testing

No behaviour change — this is a type annotation fix only. Existing tests continue to pass. Users can now write:

union = MyModel1 | MyModel2

@streaming_action.pydantic(
    reads=[...],
    writes=[...],
    state_input_type=...,
    state_output_type=...,
    stream_type=union,
)

without a type error.

Ghraven added 2 commits May 3, 2026 16:56
Fixes apache#607

Broadens the `stream_type` parameter in `streaming_action.pydantic()` and
`pydantic_streaming_action()` to accept a union of pydantic BaseModel types
(e.g. `ModelA | ModelB`) in addition to a single `Type[BaseModel]` or
`Type[dict]`. The internal validation logic already handles this at runtime;
this change removes the overly strict static type annotation that rejected
union types at the call site.
Fixes apache#607

Broadens the `stream_type` parameter in `streaming_action.pydantic()` and
`pydantic_streaming_action()` to accept a union of pydantic BaseModel types
(e.g. `ModelA | ModelB`) in addition to a single `Type[BaseModel]` or
`Type[dict]`. The internal validation logic already handles this at runtime;
this change removes the overly strict static type annotation that rejected
union types at the call site.
@github-actions github-actions Bot added area/core Application, State, Graph, Actions area/integrations External integrations (LLMs, frameworks) area/typing Mypy, type hints, pydantic labels May 3, 2026
@Ghraven
Copy link
Copy Markdown
Author

Ghraven commented May 17, 2026

Hi! Just checking in on this one. If the union type approach doesn't fit how you'd like to handle this, or if you'd prefer a different fix, I'm happy to adjust. Thanks so much for your work on Burr!

@elijahbenizzy
Copy link
Copy Markdown
Contributor

I don't think loosening it to accept everything is correct? See the comment in the original ticket I just added

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area/core Application, State, Graph, Actions area/integrations External integrations (LLMs, frameworks) area/typing Mypy, type hints, pydantic

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Streaming Event type, type hint, should support union type

2 participants