135 lines
5.5 KiB
Python
135 lines
5.5 KiB
Python
|
"""统一 LLM 客户端：本地 Ollama / OpenAI 兼容 API（api_key）双模式。
|
|||
|
|
|
|||
|
|
设计：两个模式的底层都是 OpenAI 兼容 /v1 端点，切换只改配置（LLM_MODE），
|
|||
|
|
业务代码一律调用 `chat()` / `embed()`，不感知底层。失败自动指数退避重试，
|
|||
|
|
主模型失败切换备用模型（fallback_chat_model），再失败抛 LLMFailError（§1001），
|
|||
|
|
调用方兜底话术见 FALLBACK_REPLY（降级链与原需求文档一致）。
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import asyncio
|
|||
|
|
import logging
|
|||
|
|
|
|||
|
|
import httpx
|
|||
|
|
|
|||
|
|
from config.settings import LLMCfg, settings
|
|||
|
|
from utils.exceptions import LLMFailError
|
|||
|
|
|
|||
|
|
logger = logging.getLogger("tool.llm")
|
|||
|
|
|
|||
|
|
# 兜底话术：主模型 + 备用模型均失败时返回
|
|||
|
|
FALLBACK_REPLY = "抱歉，服务暂时不可用，请稍后再试，或拨打客服热线 400-XXX-XXXX。"
|
|||
|
|
|
|||
|
|
|
|||
|
|
class LLMClient:
|
|||
|
|
def __init__(self, cfg: LLMCfg | None = None):
|
|||
|
|
self.cfg = cfg or settings.llm
|
|||
|
|
|
|||
|
|
# ---- 模式解析：端点与模型由 LLM_MODE 决定 ---------------------------
|
|||
|
|
@property
|
|||
|
|
def is_ollama(self) -> bool:
|
|||
|
|
return self.cfg.mode.lower() == "ollama"
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def base_url(self) -> str:
|
|||
|
|
return (self.cfg.ollama_base + "/v1") if self.is_ollama else self.cfg.api_base
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def chat_model(self) -> str:
|
|||
|
|
return self.cfg.ollama_chat_model if self.is_ollama else self.cfg.api_chat_model
|
|||
|
|
|
|||
|
|
@property
|
|||
|
|
def embed_model(self) -> str:
|
|||
|
|
return self.cfg.ollama_embed_model if self.is_ollama else self.cfg.api_embed_model
|
|||
|
|
|
|||
|
|
def _headers(self) -> dict:
|
|||
|
|
h = {"Content-Type": "application/json"}
|
|||
|
|
if not self.is_ollama:
|
|||
|
|
h["Authorization"] = f"Bearer {self.cfg.api_key}"
|
|||
|
|
return h
|
|||
|
|
|
|||
|
|
def describe(self) -> str:
|
|||
|
|
"""健康检查/日志用：当前模式 + 端点 + 模型。"""
|
|||
|
|
return f"{self.cfg.mode} [{self.base_url}] chat={self.chat_model} embed={self.embed_model}"
|
|||
|
|
|
|||
|
|
# ---- 对话 -----------------------------------------------------------
|
|||
|
|
async def chat(
|
|||
|
|
self,
|
|||
|
|
messages: list[dict],
|
|||
|
|
*,
|
|||
|
|
temperature: float | None = None,
|
|||
|
|
max_tokens: int | None = None,
|
|||
|
|
model: str | None = None,
|
|||
|
|
) -> str:
|
|||
|
|
"""chat/completions，返回 assistant 文本。
|
|||
|
|
|
|||
|
|
模型降级链：model > chat_model > fallback_chat_model；每模型内部退避重试。
|
|||
|
|
"""
|
|||
|
|
models = [model or self.chat_model]
|
|||
|
|
if self.cfg.fallback_chat_model and self.cfg.fallback_chat_model not in models:
|
|||
|
|
models.append(self.cfg.fallback_chat_model)
|
|||
|
|
|
|||
|
|
last_err: Exception | None = None
|
|||
|
|
for m in models:
|
|||
|
|
try:
|
|||
|
|
return await self._chat_once(m, messages, temperature, max_tokens)
|
|||
|
|
except Exception as e: # noqa: BLE001 重试/切换模型，异常向上收敛
|
|||
|
|
logger.warning("chat model=%r failed: %s: %s", m, type(e).__name__, e)
|
|||
|
|
last_err = e
|
|||
|
|
raise LLMFailError() from last_err
|
|||
|
|
|
|||
|
|
async def _chat_once(
|
|||
|
|
self,
|
|||
|
|
model: str,
|
|||
|
|
messages: list[dict],
|
|||
|
|
temperature: float | None,
|
|||
|
|
max_tokens: int | None,
|
|||
|
|
) -> str:
|
|||
|
|
payload = {
|
|||
|
|
"model": model,
|
|||
|
|
"messages": messages,
|
|||
|
|
"temperature": self.cfg.temperature if temperature is None else temperature,
|
|||
|
|
"max_tokens": self.cfg.max_tokens if max_tokens is None else max_tokens,
|
|||
|
|
"stream": False,
|
|||
|
|
}
|
|||
|
|
url = f"{self.base_url}/chat/completions"
|
|||
|
|
# 指数退避：retry_backoff_sec → 翻倍 → …，最多 max_retries 次
|
|||
|
|
for attempt in range(self.cfg.max_retries):
|
|||
|
|
try:
|
|||
|
|
async with httpx.AsyncClient(timeout=self.cfg.timeout) as client:
|
|||
|
|
r = await client.post(url, headers=self._headers(), json=payload)
|
|||
|
|
r.raise_for_status()
|
|||
|
|
return r.json()["choices"][0]["message"]["content"]
|
|||
|
|
except Exception:
|
|||
|
|
if attempt == self.cfg.max_retries - 1:
|
|||
|
|
raise
|
|||
|
|
await asyncio.sleep(self.cfg.retry_backoff_sec * (2**attempt))
|
|||
|
|
raise LLMFailError() # 理论不可达，防御
|
|||
|
|
|
|||
|
|
# ---- Embedding（RAG 入口的统一向量化） -------------------------------
|
|||
|
|
async def embed(self, texts: list[str]) -> list[list[float]]:
|
|||
|
|
"""批量向量化，返回与 input 一一对应的向量列表（维度由模型决定）。"""
|
|||
|
|
url = f"{self.base_url}/embeddings"
|
|||
|
|
payload = {"model": self.embed_model, "input": texts}
|
|||
|
|
async with httpx.AsyncClient(timeout=self.cfg.timeout) as client:
|
|||
|
|
r = await client.post(url, headers=self._headers(), json=payload)
|
|||
|
|
r.raise_for_status()
|
|||
|
|
data = r.json()
|
|||
|
|
return [item["embedding"] for item in data["data"]]
|
|||
|
|
|
|||
|
|
async def embed_one(self, text: str) -> list[float]:
|
|||
|
|
return (await self.embed([text]))[0]
|
|||
|
|
|
|||
|
|
# ---- 就绪探测（启动/健康检查可选接入） --------------------------------
|
|||
|
|
async def check_health(self) -> None:
|
|||
|
|
"""验证当前模式端点可达。ollama 走 /api/tags，API 走 /models（需鉴权，401 也算可达）。"""
|
|||
|
|
base = self.cfg.ollama_base if self.is_ollama else self.cfg.api_base
|
|||
|
|
path = "/api/tags" if self.is_ollama else "/models"
|
|||
|
|
async with httpx.AsyncClient(timeout=min(self.cfg.timeout, 10)) as client:
|
|||
|
|
r = await client.get(f"{base}{path}", headers=self._headers())
|
|||
|
|
if r.status_code >= 500:
|
|||
|
|
r.raise_for_status()
|
|||
|
|
|
|||
|
|
|
|||
|
|
# 全局单例：Agent 统一引入
|
|||
|
|
llm = LLMClient()
|