deepset-ai--haystack/test/dataclasses/test_streaming_chunk.py

# SPDX-FileCopyrightText: 2022-present deepset GmbH <info@deepset.ai>
#
# SPDX-License-Identifier: Apache-2.0

import logging
import warnings

import pytest

from haystack import Pipeline, component
from haystack.dataclasses import (
    ComponentInfo,
    ReasoningContent,
    StreamingChunk,
    ToolCall,
    ToolCallDelta,
    ToolCallResult,
)
from haystack.dataclasses.streaming_chunk import FinishReason, _invoke_streaming_callback, select_streaming_callback


@component
class ExampleComponent:
    def __init__(self):
        self.name = "test_component"

    def run(self) -> dict[str, str]:
        return {"test": "Test content"}


class TestStreamingChunk:
    def test_create_chunk_with_content_and_metadata(self):
        chunk = StreamingChunk(content="Test content", meta={"key": "value"})

        assert chunk.content == "Test content"
        assert chunk.meta == {"key": "value"}

    def test_create_chunk_with_only_content(self):
        chunk = StreamingChunk(content="Test content")

        assert chunk.content == "Test content"
        assert chunk.meta == {}

    def test_access_content(self):
        chunk = StreamingChunk(content="Test content", meta={"key": "value"})
        assert chunk.content == "Test content"

    def test_create_chunk_with_empty_content(self):
        chunk = StreamingChunk(content="")
        assert chunk.content == ""
        assert chunk.meta == {}

    def test_create_chunk_with_all_fields(self):
        component_info = ComponentInfo(type="test.component", name="test_component")
        chunk = StreamingChunk(content="Test content", meta={"key": "value"}, component_info=component_info)

        assert chunk.content == "Test content"
        assert chunk.meta == {"key": "value"}
        assert chunk.component_info == component_info

    def test_create_chunk_with_content_and_tool_call(self):
        with pytest.raises(ValueError):
            # Can't have content + tool_call at the same time
            StreamingChunk(
                content="Test content",
                meta={"key": "value"},
                tool_calls=[ToolCallDelta(id="123", tool_name="test_tool", arguments='{"arg1": "value1"}', index=0)],
            )

    def test_create_chunk_with_content_and_tool_call_result(self):
        with pytest.raises(ValueError):
            # Can't have content + tool_call_result at the same time
            StreamingChunk(
                content="Test content",
                meta={"key": "value"},
                tool_call_result=ToolCallResult(
                    result="output",
                    origin=ToolCall(id="123", tool_name="test_tool", arguments={"arg1": "value1"}),
                    error=False,
                ),
            )

    def test_create_chunk_with_content_and_reasoning(self):
        with pytest.raises(ValueError, match="Only one of `content`, `tool_call`, `tool_call_result`"):
            StreamingChunk(
                content="Test content", meta={"key": "value"}, reasoning=ReasoningContent(reasoning_text="thinking")
            )

    def test_reasoning_and_no_index(self):
        with pytest.raises(
            ValueError, match="If `tool_call`, `tool_call_result` or `reasoning` is set, `index` must also be set."
        ):
            StreamingChunk(content="", meta={"key": "value"}, reasoning=ReasoningContent(reasoning_text="thinking"))

    def test_component_info_from_component(self):
        component_info = ComponentInfo.from_component(ExampleComponent())
        assert component_info.type == "test_streaming_chunk.ExampleComponent"

    def test_component_info_from_component_with_name_from_pipeline(self):
        pipeline = Pipeline()
        comp = ExampleComponent()
        pipeline.add_component("pipeline_component", comp)
        component_info = ComponentInfo.from_component(comp)
        assert component_info.type == "test_streaming_chunk.ExampleComponent"
        assert component_info.name == "pipeline_component"

    def test_tool_call_delta(self):
        tool_call = ToolCallDelta(id="123", tool_name="test_tool", arguments='{"arg1": "value1"}', index=0)
        assert tool_call.id == "123"
        assert tool_call.tool_name == "test_tool"
        assert tool_call.arguments == '{"arg1": "value1"}'
        assert tool_call.index == 0

    def test_create_chunk_with_finish_reason(self):
        """Test creating a chunk with the new finish_reason field."""
        chunk = StreamingChunk(content="Test content", finish_reason="stop")

        assert chunk.content == "Test content"
        assert chunk.finish_reason == "stop"
        assert chunk.meta == {}

    def test_create_chunk_with_finish_reason_and_meta(self):
        """Test creating a chunk with both finish_reason field and meta."""
        chunk = StreamingChunk(
            content="Test content", finish_reason="stop", meta={"model": "gpt-4", "usage": {"tokens": 10}}
        )

        assert chunk.content == "Test content"
        assert chunk.finish_reason == "stop"
        assert chunk.meta["model"] == "gpt-4"
        assert chunk.meta["usage"]["tokens"] == 10

    def test_finish_reason_standard_values(self):
        """Test all standard finish_reason values including the new Haystack-specific ones."""
        standard_values: list[FinishReason] = ["stop", "length", "tool_calls", "content_filter", "tool_call_results"]

        for value in standard_values:
            chunk = StreamingChunk(content="Test content", finish_reason=value)
            assert chunk.finish_reason == value

    def test_finish_reason_tool_call_results(self):
        """Test specifically the new tool_call_results finish reason."""
        chunk = StreamingChunk(
            content="", finish_reason="tool_call_results", meta={"finish_reason": "tool_call_results"}
        )

        assert chunk.finish_reason == "tool_call_results"
        assert chunk.meta["finish_reason"] == "tool_call_results"
        assert chunk.content == ""

    def test_to_dict_tool_call_result(self):
        """Test the to_dict method for StreamingChunk with tool_call_result."""
        component_info = ComponentInfo.from_component(ExampleComponent())
        tool_call_result = ToolCallResult(
            result="output", origin=ToolCall(id="123", tool_name="test_tool", arguments={"arg1": "value1"}), error=False
        )

        chunk = StreamingChunk(
            content="",
            meta={"key": "value"},
            index=0,
            component_info=component_info,
            tool_call_result=tool_call_result,
            finish_reason="tool_call_results",
        )

        d = chunk.to_dict()

        assert d["content"] == ""
        assert d["meta"] == {"key": "value"}
        assert d["index"] == 0
        assert d["component_info"]["type"] == "test_streaming_chunk.ExampleComponent"
        assert d["tool_call_result"]["result"] == "output"
        assert d["tool_call_result"]["error"] is False
        assert d["tool_call_result"]["origin"]["id"] == "123"
        assert d["tool_call_result"]["origin"]["arguments"]["arg1"] == "value1"
        assert d["finish_reason"] == "tool_call_results"
        assert d["reasoning"] is None

    def test_to_dict_tool_calls(self):
        """Test the to_dict method for StreamingChunk with tool_calls."""
        component_info = ComponentInfo.from_component(ExampleComponent())
        tool_calls = [
            ToolCallDelta(id="123", tool_name="test_tool", arguments='{"arg1": "value1"}', index=0),
            ToolCallDelta(id="456", tool_name="another_tool", arguments='{"arg2": "value2"}', index=1),
        ]

        chunk = StreamingChunk(
            content="",
            meta={"key": "value"},
            index=0,
            component_info=component_info,
            tool_calls=tool_calls,
            finish_reason="tool_calls",
        )

        d = chunk.to_dict()

        assert d["content"] == ""
        assert d["meta"] == {"key": "value"}
        assert d["index"] == 0
        assert d["component_info"]["type"] == "test_streaming_chunk.ExampleComponent"
        assert len(d["tool_calls"]) == 2
        assert d["tool_calls"][0]["id"] == "123"
        assert d["tool_calls"][0]["index"] == 0
        assert d["tool_calls"][1]["id"] == "456"
        assert d["tool_calls"][1]["index"] == 1
        assert d["finish_reason"] == "tool_calls"
        assert d["reasoning"] is None

    def test_to_dict_reasoning(self):
        """Test the to_dict method for StreamingChunk with reasoning."""
        component_info = ComponentInfo.from_component(ExampleComponent())
        reasoning = ReasoningContent(reasoning_text="thinking", extra={"step": 1})

        chunk = StreamingChunk(
            content="",
            meta={"key": "value"},
            index=0,
            component_info=component_info,
            reasoning=reasoning,
            finish_reason="stop",
        )

        d = chunk.to_dict()

        assert d["content"] == ""
        assert d["meta"] == {"key": "value"}
        assert d["index"] == 0
        assert d["component_info"]["type"] == "test_streaming_chunk.ExampleComponent"
        assert d["reasoning"]["reasoning_text"] == "thinking"
        assert d["reasoning"]["extra"]["step"] == 1
        assert d["finish_reason"] == "stop"
        assert d["tool_calls"] is None
        assert d["tool_call_result"] is None

    def test_from_dict_tool_call_result(self):
        """Test the from_dict method for StreamingChunk with tool_call_result."""
        component_info = {"type": "test_streaming_chunk.ExampleComponent", "name": "test_component"}
        tool_call_result = {
            "result": "output",
            "origin": {"id": "123", "tool_name": "test_tool", "arguments": {"arg1": "value1"}},
            "error": False,
        }

        data = {
            "content": "",
            "meta": {"key": "value"},
            "index": 0,
            "component_info": component_info,
            "tool_call_result": tool_call_result,
            "finish_reason": "tool_call_results",
        }

        chunk = StreamingChunk.from_dict(data)

        assert chunk.content == ""
        assert chunk.meta == {"key": "value"}
        assert chunk.index == 0
        assert chunk.component_info is not None
        assert chunk.component_info.type == "test_streaming_chunk.ExampleComponent"
        assert chunk.component_info.name == "test_component"
        assert chunk.tool_call_result is not None
        assert chunk.tool_call_result.result == "output"
        assert chunk.tool_call_result.error is False
        assert chunk.tool_call_result.origin.id == "123"
        assert chunk.reasoning is None

    def test_from_dict_tool_calls(self):
        """Test the from_dict method for StreamingChunk with tool_calls."""
        component_info = {"type": "test_streaming_chunk.ExampleComponent", "name": "test_component"}
        tool_calls = [{"id": "123", "tool_name": "test_tool", "arguments": '{"arg1": "value1"}', "index": 0}]

        data = {
            "content": "",
            "meta": {"key": "value"},
            "index": 0,
            "component_info": component_info,
            "tool_calls": tool_calls,
            "finish_reason": "tool_calls",
        }

        chunk = StreamingChunk.from_dict(data)

        assert chunk.content == ""
        assert chunk.meta == {"key": "value"}
        assert chunk.index == 0
        assert chunk.component_info is not None
        assert chunk.component_info.type == "test_streaming_chunk.ExampleComponent"
        assert chunk.component_info.name == "test_component"
        assert chunk.tool_calls is not None
        assert chunk.tool_calls[0].tool_name == "test_tool"
        assert chunk.tool_calls[0].index == 0
        assert chunk.finish_reason == "tool_calls"
        assert chunk.reasoning is None

    def test_from_dict_reasoning(self):
        """Test the from_dict method for StreamingChunk with reasoning."""
        component_info = {"type": "test_streaming_chunk.ExampleComponent", "name": "test_component"}
        reasoning = {"reasoning_text": "thinking", "extra": {"step": 1}}

        data = {
            "content": "",
            "meta": {"key": "value"},
            "index": 0,
            "component_info": component_info,
            "reasoning": reasoning,
            "finish_reason": "stop",
        }

        chunk = StreamingChunk.from_dict(data)

        assert chunk.content == ""
        assert chunk.meta == {"key": "value"}
        assert chunk.index == 0
        assert chunk.component_info is not None
        assert chunk.component_info.type == "test_streaming_chunk.ExampleComponent"
        assert chunk.component_info.name == "test_component"
        assert chunk.reasoning is not None
        assert chunk.reasoning.reasoning_text == "thinking"
        assert chunk.reasoning.extra["step"] == 1
        assert chunk.finish_reason == "stop"
        assert chunk.tool_calls is None
        assert chunk.tool_call_result is None

    def test_tool_call_delta_no_warning_on_init(self):
        with warnings.catch_warnings():
            warnings.simplefilter("error", Warning)
            ToolCallDelta(index=0, tool_name="t")

    def test_tool_call_delta_warn_on_inplace_mutation(self):
        tcd = ToolCallDelta(index=0, tool_name="t")
        with pytest.warns(Warning, match="dataclasses.replace"):
            tcd.tool_name = "other"

    def test_component_info_no_warning_on_init(self):
        with warnings.catch_warnings():
            warnings.simplefilter("error", Warning)
            ComponentInfo(type="test.component", name="my_component")

    def test_component_info_warn_on_inplace_mutation(self):
        ci = ComponentInfo(type="test.component", name="my_component")
        with pytest.warns(Warning, match="dataclasses.replace"):
            ci.name = "other"

    def test_streaming_chunk_no_warning_on_init(self):
        with warnings.catch_warnings():
            warnings.simplefilter("error", Warning)
            StreamingChunk(content="test")

    def test_streaming_chunk_warn_on_inplace_mutation(self):
        chunk = StreamingChunk(content="test")
        with pytest.warns(Warning, match="dataclasses.replace"):
            chunk.content = "other"


def sync_callback(chunk):
    pass


async def async_callback(chunk):
    pass


class TestSelectStreamingCallback:
    def test_runtime_precedes_init(self):
        selected = select_streaming_callback(sync_callback, async_callback, requires_async=True)
        assert selected is async_callback

    def test_runtime_only(self):
        selected = select_streaming_callback(None, sync_callback, requires_async=False)
        assert selected is sync_callback

    def test_init_fallback_when_no_runtime(self):
        selected = select_streaming_callback(async_callback, None, requires_async=True)
        assert selected is async_callback

    def test_both_none(self):
        assert select_streaming_callback(None, None, requires_async=False) is None
        assert select_streaming_callback(None, None, requires_async=True) is None

    def test_async_callback_in_sync_context_raises(self):
        with pytest.raises(ValueError, match="cannot be a coroutine"):
            select_streaming_callback(None, async_callback, requires_async=False)
        with pytest.raises(ValueError, match="cannot be a coroutine"):
            select_streaming_callback(async_callback, None, requires_async=False)

    def test_sync_callback_in_async_context_warns_but_returns(self, caplog):
        with caplog.at_level(logging.WARNING):
            selected = select_streaming_callback(None, sync_callback, requires_async=True)
        assert selected is sync_callback
        assert "sync streaming callback" in caplog.text
        assert "runtime" in caplog.text

        caplog.clear()
        with caplog.at_level(logging.WARNING):
            selected = select_streaming_callback(sync_callback, None, requires_async=True)
        assert selected is sync_callback
        assert "sync streaming callback" in caplog.text
        assert "initialization" in caplog.text


class TestInvokeStreamingCallback:
    async def test_invokes_sync_callback(self):
        captured: list[StreamingChunk] = []

        def callback(chunk):
            captured.append(chunk)

        chunk = StreamingChunk(content="hello")
        await _invoke_streaming_callback(callback, chunk)
        assert captured == [chunk]

    async def test_invokes_async_callback(self):
        captured: list[StreamingChunk] = []

        async def callback(chunk):
            captured.append(chunk)

        chunk = StreamingChunk(content="world")
        await _invoke_streaming_callback(callback, chunk)
        assert captured == [chunk]