初始化提交

backend/app/ai_provider/llm/stream.py

@@ -0,0 +1,115 @@
+from __future__ import annotations
+
+import json
+from collections import defaultdict
+from typing import Any
+
+from app.ai_provider.types.chat import ChatChunk, ChatResult, TokenUsage
+
+
+class ToolCallAccumulator:
+    """将流式 tool call 分片累积为完整的 tool call。
+
+    处理流式场景中常见的多分片 tool call 参数，
+    按索引拼接部分 JSON 字符串。
+    """
+
+    def __init__(self) -> None:
+        self.content: str = ""
+        self.reasoning: str = ""
+        self._tool_call_slots: dict[int, dict[str, Any]] = defaultdict(
+            lambda: {"index": 0, "id": "", "name": "", "args": ""}
+        )
+        self.usage: TokenUsage | None = None
+
+    def add_chunk(self, chunk: ChatChunk) -> None:
+        """处理一个 ChatChunk，累积 content/reasoning/tool-calls/usage。"""
+        if chunk.delta_content:
+            self.content += chunk.delta_content
+        if chunk.delta_reasoning:
+            self.reasoning += chunk.delta_reasoning
+        if chunk.usage is not None:
+            self.usage = chunk.usage
+        for tc in chunk.tool_call_chunks:
+            idx = tc.get("index", 0)
+            if idx not in self._tool_call_slots:
+                self._tool_call_slots[idx] = {"index": idx, "id": "", "name": "", "args": ""}
+            slot = self._tool_call_slots[idx]
+            # The first chunk for a given index carries id/name; subsequent
+            # arg-only chunks have them as None and should NOT overwrite.
+            if tc.get("id"):
+                slot["id"] = tc["id"]
+            if tc.get("name"):
+                slot["name"] = tc["name"]
+            slot["args"] += tc.get("args", "")
+
+    def drain(self) -> list[dict[str, Any]]:
+        """返回所有已累积的 tool call 并重置内部槽位。"""
+        result: list[dict[str, Any]] = []
+        for idx in sorted(self._tool_call_slots.keys()):
+            slot = self._tool_call_slots[idx]
+            result.append(
+                {
+                    "index": slot["index"],
+                    "id": slot["id"] or f"call_{idx}",
+                    "name": slot["name"],
+                    "args": slot["args"],
+                }
+            )
+        self._tool_call_slots.clear()
+        return result
+
+
+class StreamHandler:
+    """收集流式分片并生成最终的 ChatResult。"""
+
+    def __init__(self) -> None:
+        self._accumulator = ToolCallAccumulator()
+        self._chunks: list[ChatChunk] = []
+
+    def add_chunk(self, chunk: ChatChunk) -> None:
+        """收集一个流式分片。"""
+        self._chunks.append(chunk)
+        self._accumulator.add_chunk(chunk)
+
+    @property
+    def chunks(self) -> list[ChatChunk]:
+        """返回所有已收集分片的副本。"""
+        return list(self._chunks)
+
+    def to_result(self, model_name: str = "") -> ChatResult:
+        """从所有累积分片构建最终的 ChatResult。"""
+        raw_tool_calls = self._accumulator.drain()
+        parsed_calls: list[dict[str, Any]] = []
+        for tc in raw_tool_calls:
+            args_raw = tc.get("args", "{}")
+            try:
+                parsed_args = json.loads(args_raw)
+            except json.JSONDecodeError:
+                parsed_args = {"_raw": args_raw}
+            parsed_calls.append(
+                {
+                    "name": tc["name"],
+                    "args": parsed_args,
+                    "id": tc["id"],
+                }
+            )
+        return ChatResult(
+            content=self._accumulator.content,
+            reasoning=self._accumulator.reasoning,
+            tool_calls=parsed_calls,
+            usage=self._accumulator.usage,
+            model_name=model_name,
+        )
+
+
+class StreamFactory:
+    """流处理器和累积器的工厂。"""
+
+    @staticmethod
+    def create_handler() -> StreamHandler:
+        return StreamHandler()
+
+    @staticmethod
+    def create_accumulator() -> ToolCallAccumulator:
+        return ToolCallAccumulator()