116 lines
3.9 KiB
Python
116 lines
3.9 KiB
Python
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()
|