implement union type support for streaming actions

burr/core/action.py

@@ -1507,15 +1507,15 @@ 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], object],
         tags: Optional[List[str]] = None,
     ) -> Callable:
         """Creates a streaming action that uses pydantic models.
 
         :param reads: The fields this consumes from the state.
         :param writes: The fields this writes to the state.
         :param stream_type: The pydantic model or dictionary type that is used to represent the partial results.
-            Use a dict if you want this untyped.
+            Use a dict if you want this untyped. Can also be a union of models (e.g., Model1 | Model2).
         :param state_input_type: The pydantic model type that is used to represent the input state.
         :param state_output_type: The pydantic model type that is used to represent the output state.
         :param tags: Optional list of tags to associate with this action

burr/integrations/pydantic.py

@@ -269,7 +269,15 @@ async def async_action_function(state: State, **kwargs) -> State:
     return decorator
 
 
-PartialType = Union[Type[pydantic.BaseModel], Type[dict]]
+# PartialType represents the type of intermediate results in a streaming action.
+# It can be:
+# - Type[pydantic.BaseModel]: A Pydantic model class
+# - Type[dict]: The dict type
+# - A union of BaseModel types (e.g., Model1 | Model2, or Union[Model1, Model2])
+# We use 'object' to accept union types since they are valid runtime values even if
+# the type system cannot precisely express them. Union types created with | (Python 3.10+)
+# or typing.Union are both supported.
+PartialType = Union[Type[pydantic.BaseModel], Type[dict], object]
 
 PydanticStreamingActionFunctionSync = Callable[
     ..., Generator[Tuple[Union[pydantic.BaseModel, dict], Optional[pydantic.BaseModel]], None, None]
@@ -290,11 +298,11 @@ 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[
-    Type[pydantic.BaseModel], Type[pydantic.BaseModel], Union[Type[dict], Type[pydantic.BaseModel]]
+    Type[pydantic.BaseModel], Type[pydantic.BaseModel], PartialType
 ]:
     if stream_type is None:
         # TODO -- derive from the signature