hunjing/backend/tests/ai_provider/llm/test_stream.py

from __future__ import annotations

from app.ai_provider.llm.stream import StreamFactory, StreamHandler, ToolCallAccumulator
from app.ai_provider.types.chat import ChatChunk, TokenUsage


class TestToolCallAccumulator:
    def test_accumulate_simple(self):
        acc = ToolCallAccumulator()
        acc.add_chunk(ChatChunk(delta_content="Hello"))
        assert acc.content == "Hello"

    def test_accumulate_reasoning(self):
        acc = ToolCallAccumulator()
        acc.add_chunk(ChatChunk(delta_reasoning="I think"))
        acc.add_chunk(ChatChunk(delta_reasoning=" therefore"))
        assert acc.reasoning == "I think therefore"

    def test_accumulate_tool_call(self):
        acc = ToolCallAccumulator()
        acc.add_chunk(
            ChatChunk(
                tool_call_chunks=[
                    {
                        "index": 0,
                        "id": "call_1",
                        "name": "read_file",
                        "args": '{"path": "test.txt"}',
                    },
                ]
            )
        )
        result = acc.drain()
        assert len(result) == 1
        assert result[0]["name"] == "read_file"

    def test_accumulate_tool_call_args(self):
        acc = ToolCallAccumulator()
        acc.add_chunk(ChatChunk(tool_call_chunks=[{"index": 0, "args": '{"path": "'}]))
        acc.add_chunk(ChatChunk(tool_call_chunks=[{"index": 0, "args": 'test.txt"}'}]))
        result = acc.drain()
        assert result[0]["args"] == '{"path": "test.txt"}'

    def test_accumulate_usage(self):
        acc = ToolCallAccumulator()
        acc.add_chunk(
            ChatChunk(usage=TokenUsage(prompt_tokens=10, completion_tokens=5, total_tokens=15))
        )
        assert acc.usage is not None
        assert acc.usage.prompt_tokens == 10

    def test_empty_drain(self):
        acc = ToolCallAccumulator()
        result = acc.drain()
        assert result == []

    def test_accumulate_multiple_tool_calls(self):
        acc = ToolCallAccumulator()
        acc.add_chunk(
            ChatChunk(
                tool_call_chunks=[
                    {"index": 0, "id": "call_1", "name": "read_file", "args": '{"path": "'},
                    {"index": 1, "id": "call_2", "name": "list_dir", "args": '{"path": "'},
                ]
            )
        )
        acc.add_chunk(
            ChatChunk(
                tool_call_chunks=[
                    {"index": 0, "args": 'x.txt"}'},
                    {"index": 1, "args": '/tmp"}'},
                ]
            )
        )
        result = acc.drain()
        assert len(result) == 2
        assert result[0]["name"] == "read_file"
        assert result[0]["args"] == '{"path": "x.txt"}'
        assert result[1]["name"] == "list_dir"
        assert result[1]["args"] == '{"path": "/tmp"}'


class TestStreamHandler:
    def test_empty_result(self):
        handler = StreamHandler()
        result = handler.to_result(model_name="test-model")
        assert result.content == ""
        assert result.reasoning == ""
        assert result.tool_calls == []
        assert result.usage is None
        assert result.model_name == "test-model"

    def test_stream_to_result(self):
        handler = StreamHandler()
        handler.add_chunk(ChatChunk(delta_content="Hello world"))
        handler.add_chunk(ChatChunk(delta_reasoning="Thinking..."))
        handler.add_chunk(
            ChatChunk(usage=TokenUsage(prompt_tokens=5, completion_tokens=3, total_tokens=8))
        )
        result = handler.to_result(model_name="test-model")
        assert result.content == "Hello world"
        assert result.reasoning == "Thinking..."
        assert result.usage is not None
        assert result.usage.total_tokens == 8
        assert result.model_name == "test-model"

    def test_stream_with_tool_calls(self):
        handler = StreamHandler()
        handler.add_chunk(
            ChatChunk(
                tool_call_chunks=[
                    {"index": 0, "id": "call_1", "name": "read_file", "args": '{"path": "'},
                ]
            )
        )
        handler.add_chunk(
            ChatChunk(
                tool_call_chunks=[
                    {"index": 0, "args": 'test.txt"}'},
                ]
            )
        )
        result = handler.to_result(model_name="test-model")
        assert len(result.tool_calls) == 1
        assert result.tool_calls[0]["name"] == "read_file"
        assert result.tool_calls[0]["args"] == {"path": "test.txt"}

    def test_chunks_property(self):
        handler = StreamHandler()
        handler.add_chunk(ChatChunk(delta_content="Hello"))
        handler.add_chunk(ChatChunk(delta_content=" world"))
        assert len(handler.chunks) == 2

    def test_stream_content_and_tool_calls_together(self):
        handler = StreamHandler()
        handler.add_chunk(ChatChunk(delta_content="I will read the file."))
        handler.add_chunk(
            ChatChunk(
                tool_call_chunks=[
                    {"index": 0, "id": "call_1", "name": "read_file", "args": "{}"},
                ]
            )
        )
        result = handler.to_result()
        assert result.content == "I will read the file."
        assert len(result.tool_calls) == 1

    def test_non_json_args_fallback(self):
        handler = StreamHandler()
        handler.add_chunk(
            ChatChunk(
                tool_call_chunks=[
                    {"index": 0, "id": "call_1", "name": "read_file", "args": "not valid json"},
                ]
            )
        )
        result = handler.to_result()
        assert result.tool_calls[0]["args"] == {"_raw": "not valid json"}


class TestStreamFactory:
    def test_create_handler(self):
        handler = StreamFactory.create_handler()
        assert isinstance(handler, StreamHandler)

    def test_create_accumulator(self):
        acc = StreamFactory.create_accumulator()
        assert isinstance(acc, ToolCallAccumulator)