初始化提交

backend/tests/ai_provider/llm/test_stream.py

@@ -0,0 +1,168 @@
+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)