"""LLM 抽象层。 为什么需要这层：本机 vLLM 和 DeepSeek 虽然都是 OpenAI 兼容协议， 但有三处实测差异，写死任何一个都会出问题： 1. 思考内容的字段名 本机 vLLM : delta.reasoning DeepSeek : delta.reasoning_content 只认一个会**静默丢弃**思考内容——不报错，只是看不见。 2. 关闭思考的方式 本机 vLLM : chat_template_kwargs={"enable_thinking": false} DeepSeek : thinking={"type": "disabled"} 3. system 消息条数 本机 vLLM : 只允许一条，且必须在开头，多条直接 400 DeepSeek : 宽松 所以统一由本层归一化成单条 system。 以上均来自实测，见 P0-本地LLM验证报告.md / P0-DeepSeek验证报告.md。 """ from __future__ import annotations import asyncio import json import logging import time from abc import ABC, abstractmethod from dataclasses import dataclass, field from enum import Enum from typing import Any, AsyncIterator, Sequence import aiohttp from ..config import LLMConfig log = logging.getLogger(__name__) class ChunkKind(str, Enum): REASONING = "reasoning" CONTENT = "content" TOOL_CALL = "tool_call" DONE = "done" @dataclass class Chunk: kind: ChunkKind text: str = "" tool_index: int = 0 tool_id: str = "" tool_name: str = "" tool_arguments: str = "" finish_reason: str | None = None usage: dict[str, Any] | None = None @dataclass class ToolCall: id: str name: str arguments: str def parsed_arguments(self) -> dict[str, Any]: try: return json.loads(self.arguments or "{}") except json.JSONDecodeError: return {} @dataclass class LLMResult: content: str = "" reasoning: str = "" tool_calls: list[ToolCall] = field(default_factory=list) finish_reason: str | None = None usage: dict[str, Any] | None = None # 计时：首 token 是"思考"还是"回答"要分开，否则会误判延迟。 # 实测开思考时 138ms 就出首块，但那是思考内容，用户要听的回答还没开始。 first_reasoning_ms: float | None = None first_content_ms: float | None = None total_ms: float = 0.0 backend: str = "" @property def has_tool_calls(self) -> bool: return bool(self.tool_calls) class LLMProvider(ABC): """一个 LLM 后端。 子类只需实现 `_build_body`（怎么关思考、怎么放参数）； 消息归一化、SSE 解析、计时、取消都在基类做，保证各后端行为一致。 """ #: 是否支持图片输入。本机 vLLM 实例实测不支持（"At most 0 image(s)"）。 supports_vision: bool = False def __init__(self, cfg: LLMConfig, session: aiohttp.ClientSession | None = None): self.cfg = cfg self.name = cfg.name self._session = session self._owns_session = session is None # ------------------------------------------------------------ 供子类覆写 @abstractmethod def _build_body(self, messages: list[dict], *, tools: Sequence[dict] | None, max_tokens: int) -> dict: """构造请求体。各后端的思考开关差异在这里处理。""" def _headers(self) -> dict[str, str]: return { "Authorization": f"Bearer {self.cfg.api_key}", "Content-Type": "application/json", } @property def endpoint(self) -> str: return self.cfg.base_url.rstrip("/") + "/chat/completions" # ------------------------------------------------------------ 消息归一化 @staticmethod def normalize_messages(messages: Sequence[dict]) -> list[dict]: """把多条 system 合并成一条并置于开头。 本机 vLLM 对多条 system 直接 400（实测 "System message must be at the beginning."）。而记忆注入天然想用第二条 system，所以在这里统一抹平。 顺带这也满足方案 §8 的要求："注入的记忆每轮替换，不是追加"—— 调用方每轮重新构造 system 内容即可。 """ systems, rest = [], [] for m in messages: if m.get("role") == "system": c = m.get("content") if isinstance(c, str) and c.strip(): systems.append(c.strip()) else: rest.append(m) out = [] if systems: out.append({"role": "system", "content": "\n\n".join(systems)}) out.extend(rest) return out # ------------------------------------------------------------ 流式请求 async def _ensure_session(self) -> aiohttp.ClientSession: if self._session is None or self._session.closed: self._session = aiohttp.ClientSession() self._owns_session = True return self._session async def aclose(self) -> None: if self._owns_session and self._session and not self._session.closed: await self._session.close() async def stream( self, messages: Sequence[dict], *, tools: Sequence[dict] | None = None, max_tokens: int | None = None, ) -> AsyncIterator[Chunk]: """流式产出。调用方取消（CancelledError）时连接会随之关闭。 这就是方案 §5 要求的"打断时取消旧回复生成"的底层能力。 """ body = self._build_body( self.normalize_messages(messages), tools=tools, max_tokens=max_tokens or self.cfg.max_tokens, ) session = await self._ensure_session() timeout = aiohttp.ClientTimeout( total=self.cfg.timeout_s, sock_connect=5.0, sock_read=self.cfg.first_token_timeout_s, ) async with session.post(self.endpoint, json=body, headers=self._headers(), timeout=timeout) as resp: if resp.status != 200: detail = (await resp.text())[:400] raise LLMError(f"{self.name} HTTP {resp.status}: {detail}") async for chunk in self._iter_sse(resp): yield chunk @staticmethod async def _iter_sse(resp: aiohttp.ClientResponse) -> AsyncIterator[Chunk]: """解析 SSE。字段名差异在这里吸收。""" # 工具调用的参数是分片到达的，按 index 累积 tool_acc: dict[int, dict] = {} async for raw in resp.content: line = raw.decode("utf-8", "replace").strip() if not line or not line.startswith("data:"): continue payload = line[5:].strip() if payload == "[DONE]": break try: data = json.loads(payload) except json.JSONDecodeError: continue usage = data.get("usage") choices = data.get("choices") or [] if not choices: if usage: yield Chunk(kind=ChunkKind.DONE, usage=usage) continue choice = choices[0] delta = choice.get("delta") or {} finish = choice.get("finish_reason") # 两个字段名都认。少认一个就会静默丢思考内容。 reasoning = delta.get("reasoning_content") or delta.get("reasoning") if reasoning: yield Chunk(kind=ChunkKind.REASONING, text=reasoning) content = delta.get("content") if content: yield Chunk(kind=ChunkKind.CONTENT, text=content) for tc in delta.get("tool_calls") or []: idx = tc.get("index", 0) slot = tool_acc.setdefault(idx, {"id": "", "name": "", "arguments": ""}) if tc.get("id"): slot["id"] = tc["id"] fn = tc.get("function") or {} if fn.get("name"): slot["name"] = fn["name"] if fn.get("arguments"): slot["arguments"] += fn["arguments"] yield Chunk(kind=ChunkKind.TOOL_CALL, tool_index=idx, tool_id=slot["id"], tool_name=slot["name"], tool_arguments=slot["arguments"]) if finish: yield Chunk(kind=ChunkKind.DONE, finish_reason=finish, usage=usage) # 注意：**不能在这里 return**。开了 stream_options.include_usage 后， # usage 是在 finish_reason 之后的**单独一帧**里到达的 # （该帧 choices 为空）。提前返回会把 usage 丢掉。 # 循环由 [DONE] 或流结束自然终止。 async def complete( self, messages: Sequence[dict], *, tools: Sequence[dict] | None = None, max_tokens: int | None = None, ) -> LLMResult: """收完整个流，返回结构化结果。""" t0 = time.perf_counter() res = LLMResult(backend=self.name) parts: list[str] = [] reason: list[str] = [] acc: dict[int, dict] = {} async for ch in self.stream(messages, tools=tools, max_tokens=max_tokens): if ch.kind == ChunkKind.CONTENT: if res.first_content_ms is None: res.first_content_ms = (time.perf_counter() - t0) * 1000 parts.append(ch.text) elif ch.kind == ChunkKind.REASONING: if res.first_reasoning_ms is None: res.first_reasoning_ms = (time.perf_counter() - t0) * 1000 reason.append(ch.text) elif ch.kind == ChunkKind.TOOL_CALL: acc[ch.tool_index] = {"id": ch.tool_id, "name": ch.tool_name, "arguments": ch.tool_arguments} elif ch.kind == ChunkKind.DONE: res.finish_reason = ch.finish_reason if ch.usage: res.usage = ch.usage res.content = "".join(parts) res.reasoning = "".join(reason) res.tool_calls = [ToolCall(**v) for _, v in sorted(acc.items())] res.total_ms = (time.perf_counter() - t0) * 1000 return res async def health(self) -> tuple[bool, str]: """启动自检：确认模型名有效。方案 §10 要求记录实际版本。""" try: session = await self._ensure_session() url = self.cfg.base_url.rstrip("/") + "/models" async with session.get(url, headers=self._headers(), timeout=aiohttp.ClientTimeout(total=10)) as r: if r.status != 200: return False, f"HTTP {r.status}" data = await r.json() ids = [m.get("id") for m in data.get("data", [])] if self.cfg.model in ids: return True, f"{self.cfg.model} 可用（清单 {len(ids)} 个）" return False, f"模型 {self.cfg.model!r} 不在清单 {ids} 中" except Exception as e: # noqa: BLE001 — 自检不该因异常类型漏判 return False, f"{type(e).__name__}: {e}" class LLMError(RuntimeError): pass __all__ = ["LLMProvider", "LLMError", "Chunk", "ChunkKind", "ToolCall", "LLMResult"]