398 lines
16 KiB
Python
398 lines
16 KiB
Python
"""LLM 客户端抽象层。
|
||
|
||
目标：为直播间弹幕回复提供统一的对话接口，屏蔽底层各家协议差异。
|
||
|
||
结构：
|
||
- BaseLLMProvider 抽象基类，定义 async chat(messages) -> str
|
||
- OpenAICompatProvider OpenAI /chat/completions 兼容协议
|
||
（DeepSeek / 通义千问DashScope兼容模式 / Kimi / 硅基流动 /
|
||
Ollama / vLLM / LM Studio 等几乎所有主流服务）
|
||
- OllamaNativeProvider Ollama 原生 /api/chat 协议
|
||
- AnthropicProvider Anthropic /v1/messages 协议（Claude 系列）
|
||
- create_provider() 按 config["llm_reply"]["provider"] 工厂构造
|
||
- LLMReplyEngine 弹幕回复业务引擎：过滤、限流、上下文、重试
|
||
|
||
扩展新协议只需：继承 BaseLLMProvider 实现 chat()，并在 PROVIDERS 注册。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import random
|
||
import re
|
||
import time
|
||
from abc import ABC, abstractmethod
|
||
from collections import deque
|
||
|
||
try:
|
||
import aiohttp
|
||
except ImportError: # pragma: no cover - aiohttp 为运行时必需依赖
|
||
aiohttp = None
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Provider 抽象与实现
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class LLMError(Exception):
|
||
"""LLM 调用失败（网络/协议/内容拒绝等）。"""
|
||
|
||
|
||
class BaseLLMProvider(ABC):
|
||
"""所有 LLM 协议适配器的公共基类。"""
|
||
|
||
name = "base"
|
||
|
||
def __init__(self, base_url: str, api_key: str = "", model: str = "",
|
||
timeout_sec: float = 20.0, logger: logging.Logger | None = None):
|
||
self.base_url = (base_url or "").rstrip("/")
|
||
self.api_key = api_key or ""
|
||
self.model = model or ""
|
||
self.timeout_sec = max(3.0, float(timeout_sec or 20.0))
|
||
self.logger = logger or logging.getLogger(__name__)
|
||
|
||
@abstractmethod
|
||
async def chat(self, messages: list[dict], *, temperature: float = 0.8,
|
||
max_tokens: int = 256) -> str:
|
||
"""发送一轮对话，返回助手回复文本；失败抛 LLMError。"""
|
||
|
||
async def _post_json(self, url: str, payload: dict,
|
||
headers: dict | None = None) -> dict:
|
||
if aiohttp is None:
|
||
raise LLMError("缺少 aiohttp 依赖")
|
||
timeout = aiohttp.ClientTimeout(total=self.timeout_sec)
|
||
try:
|
||
async with aiohttp.ClientSession(timeout=timeout) as session:
|
||
async with session.post(url, json=payload,
|
||
headers=headers or {}) as resp:
|
||
body = await resp.text()
|
||
if resp.status != 200:
|
||
raise LLMError(
|
||
f"HTTP {resp.status}: {body[:200]}")
|
||
import json as _json
|
||
return _json.loads(body)
|
||
except asyncio.TimeoutError as e:
|
||
raise LLMError(f"请求超时({self.timeout_sec}s)") from e
|
||
except LLMError:
|
||
raise
|
||
except Exception as e:
|
||
raise LLMError(f"请求失败: {e}") from e
|
||
|
||
|
||
class OpenAICompatProvider(BaseLLMProvider):
|
||
"""OpenAI /v1/chat/completions 兼容协议。
|
||
|
||
覆盖：OpenAI、DeepSeek、Kimi(Moonshot)、通义千问(DashScope兼容模式)、
|
||
智谱GLM、硅基流动、OpenRouter、Ollama/VoidAI/vLLM/LM Studio 本地端点等。
|
||
"""
|
||
|
||
name = "openai"
|
||
|
||
DEFAULT_PATH = "/chat/completions"
|
||
|
||
async def chat(self, messages: list[dict], *, temperature: float = 0.8,
|
||
max_tokens: int = 256) -> str:
|
||
url = self.base_url + self.DEFAULT_PATH
|
||
headers = {"Content-Type": "application/json"}
|
||
if self.api_key:
|
||
headers["Authorization"] = f"Bearer {self.api_key}"
|
||
payload = {
|
||
"model": self.model,
|
||
"messages": messages,
|
||
"temperature": max(0.0, min(2.0, temperature)),
|
||
"max_tokens": max(16, int(max_tokens)),
|
||
"stream": False,
|
||
}
|
||
data = await self._post_json(url, payload, headers)
|
||
try:
|
||
content = data["choices"][0]["message"]["content"]
|
||
except (KeyError, IndexError, TypeError) as e:
|
||
raise LLMError(f"响应格式异常: {str(data)[:200]}") from e
|
||
if not isinstance(content, str):
|
||
raise LLMError("响应 content 非文本")
|
||
return content.strip()
|
||
|
||
|
||
class OllamaNativeProvider(BaseLLMProvider):
|
||
"""Ollama 原生 /api/chat 协议（无需鉴权）。"""
|
||
|
||
name = "ollama"
|
||
|
||
async def chat(self, messages: list[dict], *, temperature: float = 0.8,
|
||
max_tokens: int = 256) -> str:
|
||
url = self.base_url + "/api/chat"
|
||
payload = {
|
||
"model": self.model,
|
||
"messages": messages,
|
||
"stream": False,
|
||
"options": {
|
||
"temperature": max(0.0, min(2.0, temperature)),
|
||
"num_predict": max(16, int(max_tokens)),
|
||
},
|
||
}
|
||
data = await self._post_json(url, payload)
|
||
content = data.get("message", {}).get("content", "")
|
||
if not isinstance(content, str):
|
||
raise LLMError("响应 message.content 非文本")
|
||
return content.strip()
|
||
|
||
|
||
class AnthropicProvider(BaseLLMProvider):
|
||
"""Anthropic /v1/messages 协议（Claude 系列）。"""
|
||
|
||
name = "anthropic"
|
||
|
||
ANTHROPIC_VERSION = "2023-06-01"
|
||
|
||
async def chat(self, messages: list[dict], *, temperature: float = 0.8,
|
||
max_tokens: int = 256) -> str:
|
||
url = self.base_url + "/v1/messages"
|
||
headers = {
|
||
"Content-Type": "application/json",
|
||
"x-api-key": self.api_key,
|
||
"anthropic-version": self.ANTHROPIC_VERSION,
|
||
}
|
||
# Anthropic 协议中 system 是顶层字段，不放在 messages 里。
|
||
system_text = "\n".join(
|
||
m["content"] for m in messages if m.get("role") == "system")
|
||
chat_msgs = [m for m in messages if m.get("role") != "system"]
|
||
payload = {
|
||
"model": self.model,
|
||
"system": system_text or None,
|
||
"messages": [
|
||
{"role": m["role"], "content": m["content"]}
|
||
for m in chat_msgs
|
||
],
|
||
"temperature": max(0.0, min(1.0, temperature)),
|
||
"max_tokens": max(16, int(max_tokens)),
|
||
}
|
||
data = await self._post_json(url, payload, headers)
|
||
try:
|
||
parts = data["content"]
|
||
text = "".join(p.get("text", "") for p in parts if isinstance(p, dict))
|
||
except (KeyError, TypeError) as e:
|
||
raise LLMError(f"响应格式异常: {str(data)[:200]}") from e
|
||
return text.strip()
|
||
|
||
|
||
# 协议注册表：新增协议在此注册即可。
|
||
PROVIDERS: dict[str, type[BaseLLMProvider]] = {
|
||
OpenAICompatProvider.name: OpenAICompatProvider,
|
||
OllamaNativeProvider.name: OllamaNativeProvider,
|
||
AnthropicProvider.name: AnthropicProvider,
|
||
}
|
||
|
||
|
||
def create_provider(cfg: dict, logger: logging.Logger | None = None) -> BaseLLMProvider:
|
||
"""根据 llm_reply 配置节构造 provider。"""
|
||
cfg = cfg or {}
|
||
name = str(cfg.get("provider", "openai") or "openai").strip().lower()
|
||
cls = PROVIDERS.get(name)
|
||
if cls is None:
|
||
raise LLMError(
|
||
f"未知 LLM 协议: {name}（可选: {', '.join(sorted(PROVIDERS))}）")
|
||
base_url = str(cfg.get("base_url", "") or "").strip()
|
||
if not base_url:
|
||
raise LLMError("未配置 base_url")
|
||
return cls(
|
||
base_url=base_url,
|
||
api_key=str(cfg.get("api_key", "") or ""),
|
||
model=str(cfg.get("model", "") or "").strip(),
|
||
timeout_sec=float(cfg.get("timeout_sec", 20) or 20),
|
||
logger=logger,
|
||
)
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# 弹幕回复业务引擎
|
||
# ---------------------------------------------------------------------------
|
||
|
||
class LLMReplyEngine:
|
||
"""把观众弹幕交给 LLM 生成回复。
|
||
|
||
职责边界：本类只负责"生成回复文本"，不负责发送——
|
||
由调用方（CommandHandler）决定走弹幕/TTS 通道。
|
||
内置：
|
||
- 开关与触发过滤（长度、忽略前缀、黑名单）
|
||
- 频率限制（全局每分钟上限 + 单人冷却）
|
||
- 短期上下文（最近 N 条观众发言与机器人回复）
|
||
- 失败重试（指数退避，最多 retry_attempts 次）
|
||
- 并发保护（同一时刻只允许一个在途请求，忙时直接跳过该条）
|
||
"""
|
||
|
||
def __init__(self, cfg: dict | None = None, logger: logging.Logger | None = None,
|
||
config_getter=None):
|
||
"""config_getter: 可选的可调用对象，返回最新 llm_reply 配置 dict，
|
||
用于跟随 Config 热更新（reload 后 data 节点可能是新对象）。"""
|
||
self.logger = logger or logging.getLogger(__name__)
|
||
self._cfg = cfg or {}
|
||
self._config_getter = config_getter
|
||
self._provider: BaseLLMProvider | None = None
|
||
self._provider_built_at = 0.0
|
||
self._context: deque[tuple[str, str]] = deque() # (speaker, text)
|
||
self._reply_times: deque[float] = deque() # 全局滑动窗口
|
||
self._user_last_reply: dict[int, float] = {}
|
||
self._inflight = asyncio.Lock()
|
||
# 配置指纹：热更新后自动重建 provider。
|
||
self._cfg_fingerprint = ""
|
||
|
||
# -- 配置 ---------------------------------------------------------------
|
||
|
||
def _config(self) -> dict:
|
||
if self._config_getter is not None:
|
||
try:
|
||
cfg = self._config_getter()
|
||
if cfg:
|
||
self._cfg = cfg
|
||
except Exception:
|
||
pass
|
||
return self._cfg or {}
|
||
|
||
def update_config(self, cfg: dict):
|
||
"""配置热更新入口。"""
|
||
self._cfg = cfg or {}
|
||
|
||
@property
|
||
def enabled(self) -> bool:
|
||
return bool(self._config().get("enabled", False))
|
||
|
||
def _provider_fingerprint(self) -> str:
|
||
c = self._config()
|
||
return "|".join(str(c.get(k, "")) for k in
|
||
("provider", "base_url", "api_key", "model", "timeout_sec"))
|
||
|
||
def _get_provider(self) -> BaseLLMProvider:
|
||
fp = self._provider_fingerprint()
|
||
if self._provider is None or fp != self._cfg_fingerprint:
|
||
self._provider = create_provider(self._config(), self.logger)
|
||
self._cfg_fingerprint = fp
|
||
self.logger.info(f"[LLM] 使用协议 {self._provider.name} "
|
||
f"模型 {self._provider.model or '(默认)'}")
|
||
return self._provider
|
||
|
||
# -- 过滤 ---------------------------------------------------------------
|
||
|
||
def should_reply(self, text: str, uid: int, uname: str) -> str | None:
|
||
"""返回需要回复的规范化文本；不需要回复返回 None。"""
|
||
cfg = self._config()
|
||
if not self.enabled:
|
||
return None
|
||
text = (text or "").strip()
|
||
raw = cfg.get("min_text_len", 2)
|
||
min_len = int(raw) if raw is not None else 2
|
||
if len(text) < min_len:
|
||
return None
|
||
for prefix in cfg.get("ignore_prefixes", []) or []:
|
||
prefix = str(prefix or "").strip()
|
||
if prefix and text.startswith(prefix):
|
||
return None
|
||
blacklist = {str(u) for u in (cfg.get("blacklist_uids", []) or [])}
|
||
if str(uid) in blacklist:
|
||
return None
|
||
# 单人冷却（0 表示禁用冷却）
|
||
raw = cfg.get("per_user_cooldown_sec", 30)
|
||
cooldown = float(raw) if raw is not None else 30.0
|
||
if cooldown > 0:
|
||
last = self._user_last_reply.get(int(uid), 0.0)
|
||
if time.time() - last < cooldown:
|
||
return None
|
||
# 全局限流（0 表示不限流）
|
||
raw = cfg.get("replies_per_minute", 6)
|
||
rpm = int(raw) if raw is not None else 6
|
||
if rpm > 0:
|
||
now = time.time()
|
||
while self._reply_times and now - self._reply_times[0] > 60:
|
||
self._reply_times.popleft()
|
||
if len(self._reply_times) >= rpm:
|
||
return None
|
||
return text
|
||
|
||
# -- 上下文 ---------------------------------------------------------------
|
||
|
||
def _build_messages(self, text: str, uname: str) -> list[dict]:
|
||
cfg = self._config()
|
||
system_prompt = str(cfg.get("system_prompt", "") or "").strip()
|
||
messages: list[dict] = []
|
||
emotion_instruction = (
|
||
"每次回复必须先输出一个情绪标签，只能从[开心]、[很开心]、[平静]、[生气]、[很生气]中选一个，"
|
||
"紧接着输出给观众的回复正文。标签只用于桌宠状态，不属于口播内容。"
|
||
)
|
||
messages.append({"role": "system", "content": "\n".join(
|
||
part for part in (system_prompt, emotion_instruction) if part
|
||
)})
|
||
for speaker, content in list(self._context):
|
||
role = "assistant" if speaker == "bot" else "user"
|
||
content = f"{speaker}: {content}" if role == "user" else content
|
||
messages.append({"role": role, "content": content})
|
||
messages.append({"role": "user", "content": f"{uname}: {text}"})
|
||
return messages
|
||
|
||
def _remember(self, uname: str, text: str, reply: str | None):
|
||
max_turns = int(self._config().get("context_turns", 6) or 6)
|
||
if max_turns <= 0:
|
||
return
|
||
self._context.append((uname, text))
|
||
if reply:
|
||
self._context.append(("bot", reply))
|
||
while len(self._context) > max_turns * 2:
|
||
self._context.popleft()
|
||
|
||
# -- 主入口 ---------------------------------------------------------------
|
||
|
||
async def generate_reply(self, text: str, uid: int, uname: str) -> str | None:
|
||
"""生成回复；被过滤/限流/失败返回 None。非阻塞——忙时立即放弃。"""
|
||
text = self.should_reply(text, uid, uname)
|
||
if text is None:
|
||
return None
|
||
if self._inflight.locked():
|
||
# 已有请求在途：丢弃本条，避免弹幕高峰打爆 LLM。
|
||
return None
|
||
async with self._inflight:
|
||
cfg = self._config()
|
||
provider = self._get_provider()
|
||
messages = self._build_messages(text, uname or "观众")
|
||
temperature = float(cfg.get("temperature", 0.8) or 0.8)
|
||
max_tokens = int(cfg.get("max_tokens", 128) or 128)
|
||
attempts = max(1, int(cfg.get("retry_attempts", 2) or 1))
|
||
last_err: Exception | None = None
|
||
for i in range(attempts):
|
||
try:
|
||
raw = await provider.chat(
|
||
messages, temperature=temperature, max_tokens=max_tokens)
|
||
reply = self._sanitize(raw)
|
||
if reply:
|
||
now = time.time()
|
||
self._reply_times.append(now)
|
||
self._user_last_reply[int(uid)] = now
|
||
self._remember(uname or "观众", text, reply)
|
||
return reply
|
||
return None # 模型返回空内容，不重试
|
||
except LLMError as e:
|
||
last_err = e
|
||
if i < attempts - 1:
|
||
await asyncio.sleep(0.8 * (2 ** i) + random.random() * 0.4)
|
||
self.logger.warning(f"[LLM] 回复失败({last_err})")
|
||
return None
|
||
|
||
def _sanitize(self, raw: str) -> str:
|
||
"""清洗模型输出：去空白、截断到 max_reply_chars、去 markdown 粗体。
|
||
|
||
开头的情绪标签（如 [开心]）豁免截断长度，保证不会被截坏；
|
||
标签本身由调用方（danmu_queue）在弹幕/语音输出前剥离。
|
||
"""
|
||
cfg = self._config()
|
||
text = str(raw or "").strip()
|
||
# 直播口播场景不需要 markdown 标记
|
||
for token in ("**", "__", "```"):
|
||
text = text.replace(token, "")
|
||
text = " ".join(text.split())
|
||
m = re.match(r"^(\[[^\[\]]{1,8}\])\s*", text)
|
||
tag = m.group(1) if m else ""
|
||
body = text[m.end():] if m else text
|
||
raw = cfg.get("max_reply_chars", 60)
|
||
max_chars = int(raw) if raw is not None else 60
|
||
if max_chars > 0 and len(tag) + len(body) > max_chars:
|
||
body = body[:max(1, max_chars - len(tag))].rstrip() + "…"
|
||
return tag + body
|