chore: import upstream snapshot with attribution

haystack/dataclasses/streaming_chunk.py

@@ -0,0 +1,280 @@
+# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
+#
+# SPDX-License-Identifier: Apache-2.0
+
+import inspect
+from collections.abc import Awaitable, Callable
+from dataclasses import asdict, dataclass, field
+from typing import Any, Literal, overload
+
+from haystack import logging
+from haystack.core.component import Component
+from haystack.dataclasses.chat_message import ReasoningContent, ToolCallResult
+from haystack.utils.dataclasses import _warn_on_inplace_mutation
+
+logger = logging.getLogger(__name__)
+
+# Type alias for standard finish_reason values following OpenAI's convention
+# plus Haystack-specific value ("tool_call_results")
+FinishReason = Literal["stop", "length", "tool_calls", "content_filter", "tool_call_results"]
+
+
+@_warn_on_inplace_mutation
+@dataclass
+class ToolCallDelta:
+    """
+    Represents a Tool call prepared by the model, usually contained in an assistant message.
+
+    :param index: The index of the Tool call in the list of Tool calls.
+    :param tool_name: The name of the Tool to call.
+    :param arguments: Either the full arguments in JSON format or a delta of the arguments.
+    :param id: The ID of the Tool call.
+    :param extra: Dictionary of extra information about the Tool call. Use to store provider-specific
+        information. To avoid serialization issues, values should be JSON serializable.
+    """
+
+    index: int
+    tool_name: str | None = field(default=None)
+    arguments: str | None = field(default=None)
+    id: str | None = field(default=None)
+    extra: dict[str, Any] | None = field(default=None)
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Returns a dictionary representation of the ToolCallDelta.
+
+        :returns: A dictionary with keys 'index', 'tool_name', 'arguments', 'id', and 'extra'.
+        """
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "ToolCallDelta":
+        """
+        Creates a ToolCallDelta from a serialized representation.
+
+        :param data: Dictionary containing ToolCallDelta's attributes.
+        :returns: A ToolCallDelta instance.
+        """
+        return ToolCallDelta(**data)
+
+
+@_warn_on_inplace_mutation
+@dataclass
+class ComponentInfo:
+    """
+    The `ComponentInfo` class encapsulates information about a component.
+
+    :param type: The type of the component.
+    :param name: The name of the component assigned when adding it to a pipeline.
+
+    """
+
+    type: str
+    name: str | None = field(default=None)
+
+    @classmethod
+    def from_component(cls, component: Component) -> "ComponentInfo":
+        """
+        Create a `ComponentInfo` object from a `Component` instance.
+
+        :param component:
+            The `Component` instance.
+        :returns:
+            The `ComponentInfo` object with the type and name of the given component.
+        """
+        component_type = f"{component.__class__.__module__}.{component.__class__.__name__}"
+        component_name = getattr(component, "__component_name__", None)
+        return cls(type=component_type, name=component_name)
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Returns a dictionary representation of ComponentInfo.
+
+        :returns: A dictionary with keys 'type' and 'name'.
+        """
+        return asdict(self)
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "ComponentInfo":
+        """
+        Creates a ComponentInfo from a serialized representation.
+
+        :param data: Dictionary containing ComponentInfo's attributes.
+        :returns: A ComponentInfo instance.
+        """
+        return ComponentInfo(**data)
+
+
+@_warn_on_inplace_mutation
+@dataclass
+class StreamingChunk:
+    """
+    The `StreamingChunk` class encapsulates a segment of streamed content along with associated metadata.
+
+    This structure facilitates the handling and processing of streamed data in a systematic manner.
+
+    :param content: The content of the message chunk as a string.
+    :param meta: A dictionary containing metadata related to the message chunk.
+    :param component_info: A `ComponentInfo` object containing information about the component that generated the chunk,
+        such as the component name and type.
+    :param index: An optional integer index representing which content block this chunk belongs to.
+    :param tool_calls: An optional list of ToolCallDelta object representing a tool call associated with the message
+        chunk.
+    :param tool_call_result: An optional ToolCallResult object representing the result of a tool call.
+    :param start: A boolean indicating whether this chunk marks the start of a content block.
+    :param finish_reason: An optional value indicating the reason the generation finished.
+        Standard values follow OpenAI's convention: "stop", "length", "tool_calls", "content_filter",
+        plus Haystack-specific value "tool_call_results".
+    :param reasoning: An optional ReasoningContent object representing the reasoning content associated
+        with the message chunk.
+    """
+
+    content: str
+    meta: dict[str, Any] = field(default_factory=dict, hash=False)
+    component_info: ComponentInfo | None = field(default=None)
+    index: int | None = field(default=None)
+    tool_calls: list[ToolCallDelta] | None = field(default=None)
+    tool_call_result: ToolCallResult | None = field(default=None)
+    start: bool = field(default=False)
+    finish_reason: FinishReason | None = field(default=None)
+    reasoning: ReasoningContent | None = field(default=None)
+
+    def __post_init__(self) -> None:
+        fields_set = sum(bool(x) for x in (self.content, self.tool_calls, self.tool_call_result, self.reasoning))
+        if fields_set > 1:
+            raise ValueError(
+                "Only one of `content`, `tool_call`, `tool_call_result` or `reasoning` may be set in a StreamingChunk. "
+                f"Got content: '{self.content}', tool_call: '{self.tool_calls}', "
+                f"tool_call_result: '{self.tool_call_result}', reasoning: '{self.reasoning}'."
+            )
+
+        # NOTE: We don't enforce this for self.content otherwise it would be a breaking change
+        if (self.tool_calls or self.tool_call_result or self.reasoning) and self.index is None:
+            raise ValueError("If `tool_call`, `tool_call_result` or `reasoning` is set, `index` must also be set.")
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Returns a dictionary representation of the StreamingChunk.
+
+        :returns: Serialized dictionary representation of the calling object.
+        """
+        return {
+            "content": self.content,
+            "meta": self.meta,
+            "component_info": self.component_info.to_dict() if self.component_info else None,
+            "index": self.index,
+            "tool_calls": [tc.to_dict() for tc in self.tool_calls] if self.tool_calls else None,
+            "tool_call_result": self.tool_call_result.to_dict() if self.tool_call_result else None,
+            "start": self.start,
+            "finish_reason": self.finish_reason,
+            "reasoning": self.reasoning.to_dict() if self.reasoning else None,
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, Any]) -> "StreamingChunk":
+        """
+        Creates a deserialized StreamingChunk instance from a serialized representation.
+
+        :param data: Dictionary containing the StreamingChunk's attributes.
+        :returns: A StreamingChunk instance.
+        """
+        if "content" not in data:
+            raise ValueError("Missing required field `content` in StreamingChunk deserialization.")
+
+        return StreamingChunk(
+            content=data["content"],
+            meta=data.get("meta", {}),
+            component_info=ComponentInfo.from_dict(data["component_info"]) if data.get("component_info") else None,
+            index=data.get("index"),
+            tool_calls=[ToolCallDelta.from_dict(tc) for tc in data["tool_calls"]] if data.get("tool_calls") else None,
+            tool_call_result=ToolCallResult.from_dict(data["tool_call_result"])
+            if data.get("tool_call_result")
+            else None,
+            start=data.get("start", False),
+            finish_reason=data.get("finish_reason"),
+            reasoning=ReasoningContent.from_dict(data["reasoning"]) if data.get("reasoning") else None,
+        )
+
+
+SyncStreamingCallbackT = Callable[[StreamingChunk], None]
+AsyncStreamingCallbackT = Callable[[StreamingChunk], Awaitable[None]]
+
+StreamingCallbackT = SyncStreamingCallbackT | AsyncStreamingCallbackT
+
+
+def _is_callable_async_compatible(func: Callable) -> bool:
+    """
+    Returns if the given callable is usable inside a component's `run_async` method.
+
+    :param func:
+        The callable to check.
+    :returns:
+        True if the callable is compatible, False otherwise.
+    """
+    return inspect.iscoroutinefunction(func)
+
+
+@overload
+def select_streaming_callback(
+    init_callback: StreamingCallbackT | None,
+    runtime_callback: StreamingCallbackT | None,
+    requires_async: Literal[False],
+) -> SyncStreamingCallbackT | None: ...
+@overload
+def select_streaming_callback(
+    init_callback: StreamingCallbackT | None, runtime_callback: StreamingCallbackT | None, requires_async: Literal[True]
+) -> StreamingCallbackT | None: ...
+
+
+def select_streaming_callback(
+    init_callback: StreamingCallbackT | None, runtime_callback: StreamingCallbackT | None, requires_async: bool
+) -> StreamingCallbackT | None:
+    """
+    Picks the correct streaming callback given an optional initial and runtime callback.
+
+    The runtime callback takes precedence over the initial callback.
+
+    In an async context (`requires_async=True`), a sync callback is accepted but emits a warning: it will run inline on
+    the event loop and may block it. In a sync context (`requires_async=False`), an async callback is rejected because
+    there is no way to await it.
+
+    :param init_callback:
+        The initial callback.
+    :param runtime_callback:
+        The runtime callback.
+    :param requires_async:
+        Whether the selected callback will be invoked from an async context.
+    :returns:
+        The selected callback.
+    """
+    if init_callback is not None:
+        if requires_async and not _is_callable_async_compatible(init_callback):
+            logger.warning(
+                "A sync streaming callback was provided at initialization for use in an async context. "
+                "It will run synchronously on the event loop and may block it."
+            )
+        if not requires_async and _is_callable_async_compatible(init_callback):
+            raise ValueError("The init callback cannot be a coroutine.")
+
+    if runtime_callback is not None:
+        if requires_async and not _is_callable_async_compatible(runtime_callback):
+            logger.warning(
+                "A sync streaming callback was provided at runtime for use in an async context. "
+                "It will run synchronously on the event loop and may block it."
+            )
+        if not requires_async and _is_callable_async_compatible(runtime_callback):
+            raise ValueError("The runtime callback cannot be a coroutine.")
+
+    return runtime_callback or init_callback
+
+
+async def _invoke_streaming_callback(callback: StreamingCallbackT, chunk: StreamingChunk) -> None:
+    """
+    Invokes a streaming callback in an async context, handling both sync and async callbacks.
+
+    :param callback: The streaming callback to invoke.
+    :param chunk: The streaming chunk to pass to the callback.
+    """
+    result = callback(chunk)
+    if inspect.isawaitable(result):
+        await result