297 lines
12 KiB
Python
297 lines
12 KiB
Python
"""统一 LLM 客户端：本地 Ollama / OpenAI 兼容 API（api_key）双后端，按优先级降级。
|
||
|
||
设计：两个后端的底层都是 OpenAI 兼容 /v1 端点，业务代码一律调用 `chat()` / `embed()`，
|
||
不感知底层。后端链由 .env 决定（LLM_MODE）：
|
||
- auto（默认）：本地 Ollama 参数齐全则优先本地，LLM_API_KEY 齐全则作为其后的降级后端；
|
||
- ollama / api：强制只用该后端。
|
||
chat 降级链：后端内 主模型 → fallback_chat_model（每模型指数退避重试）→ 下一后端 → 全部失败抛
|
||
LLMFailError（§1001），调用方兜底话术见 FALLBACK_REPLY。
|
||
embed 只走首选后端、不跨后端降级：不同模型的向量空间不兼容，混入 Milvus 会污染检索。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
from dataclasses import dataclass
|
||
|
||
import httpx
|
||
|
||
from config.settings import LLMCfg, settings
|
||
from utils.exceptions import LLMFailError
|
||
|
||
logger = logging.getLogger("tool.llm")
|
||
|
||
# 兜底话术：所有后端 + 备用模型均失败时返回
|
||
FALLBACK_REPLY = "抱歉，服务暂时不可用，请稍后再试，或拨打客服热线 400-XXX-XXXX。"
|
||
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class Backend:
|
||
"""一个 OpenAI 兼容端点及其模型；name 为 ollama / api。"""
|
||
|
||
name: str
|
||
base_url: str # 以 /v1 结尾
|
||
chat_model: str
|
||
embed_model: str
|
||
api_key: str = ""
|
||
api_trust_env: bool = True
|
||
|
||
@property
|
||
def is_ollama(self) -> bool:
|
||
return self.name == "ollama"
|
||
|
||
@property
|
||
def headers(self) -> dict:
|
||
h = {"Content-Type": "application/json"}
|
||
if self.api_key:
|
||
h["Authorization"] = f"Bearer {self.api_key}"
|
||
return h
|
||
|
||
def client(self, timeout: float) -> httpx.AsyncClient:
|
||
"""ollama 走本机端点，必须绕过系统代理：httpx 会读取 Windows 系统代理，
|
||
但不识别其 ProxyOverride 白名单，localhost 请求会被转发到代理并返回 502。
|
||
API 后端保留 trust_env，远端接口可能依赖代理。"""
|
||
|
||
return httpx.AsyncClient(timeout=timeout, trust_env=self.api_trust_env if not self.is_ollama else False)
|
||
|
||
|
||
def resolve_backends(cfg: LLMCfg) -> list[Backend]:
|
||
"""按 LLM_MODE 与 .env 完整度解析后端链，首个为首选后端。"""
|
||
ollama = (
|
||
Backend(
|
||
name="ollama",
|
||
base_url=cfg.ollama_base.rstrip("/") + "/v1",
|
||
chat_model=cfg.ollama_chat_model,
|
||
embed_model=cfg.ollama_embed_model,
|
||
api_trust_env=False,
|
||
)
|
||
if cfg.ollama_base and cfg.ollama_chat_model
|
||
else None
|
||
)
|
||
api = (
|
||
Backend(
|
||
name="api",
|
||
base_url=cfg.api_base.rstrip("/"),
|
||
chat_model=cfg.api_chat_model,
|
||
embed_model=cfg.api_embed_model,
|
||
api_key=cfg.api_key,
|
||
api_trust_env=cfg.api_trust_env,
|
||
)
|
||
if cfg.api_key and cfg.api_base and cfg.api_chat_model
|
||
else None
|
||
)
|
||
mode = cfg.mode.lower()
|
||
if mode == "ollama":
|
||
chain = [ollama]
|
||
elif mode == "api":
|
||
chain = [api]
|
||
elif mode == "auto":
|
||
chain = [ollama, api] # 本地优先，API 兜底
|
||
else:
|
||
raise ValueError(f"LLM_MODE 无效: {cfg.mode!r}，应为 auto / ollama / api")
|
||
chain = [b for b in chain if b is not None]
|
||
if not chain:
|
||
raise ValueError(
|
||
f"LLM_MODE={cfg.mode} 下没有可用后端：本地需 LLM_OLLAMA_BASE + LLM_OLLAMA_CHAT_MODEL，"
|
||
"API 需 LLM_API_KEY + LLM_API_BASE + LLM_API_CHAT_MODEL"
|
||
)
|
||
return chain
|
||
|
||
|
||
class LLMClient:
|
||
def __init__(self, cfg: LLMCfg | None = None):
|
||
self.cfg = cfg or settings.llm
|
||
self.backends = resolve_backends(self.cfg)
|
||
|
||
# ---- 首选后端（健康检查/日志/embed 使用） -----------------------------
|
||
@property
|
||
def primary(self) -> Backend:
|
||
return self.backends[0]
|
||
|
||
@property
|
||
def is_ollama(self) -> bool:
|
||
return self.primary.is_ollama
|
||
|
||
@property
|
||
def base_url(self) -> str:
|
||
return self.primary.base_url
|
||
|
||
@property
|
||
def chat_model(self) -> str:
|
||
return self.primary.chat_model
|
||
|
||
@property
|
||
def embed_model(self) -> str:
|
||
return self.primary.embed_model
|
||
|
||
def describe(self) -> str:
|
||
"""健康检查/日志用：后端链（首选在前）+ 各自模型。"""
|
||
chain = " -> ".join(
|
||
f"{b.name}[{b.base_url}] chat={b.chat_model} embed={b.embed_model or '-'}"
|
||
for b in self.backends
|
||
)
|
||
return f"mode={self.cfg.mode} {chain}"
|
||
|
||
# ---- 对话 -----------------------------------------------------------
|
||
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] → 后续后端 [chat_model >
|
||
fallback_chat_model]。显式 model 只作用于首选后端——模型名与后端绑定，跨后端无意义。
|
||
"""
|
||
last_err: Exception | None = None
|
||
for index, backend in enumerate(self.backends):
|
||
models = [(model or backend.chat_model) if index == 0 else backend.chat_model]
|
||
if self.cfg.fallback_chat_model and self.cfg.fallback_chat_model not in models:
|
||
models.append(self.cfg.fallback_chat_model)
|
||
for m in models:
|
||
try:
|
||
return await self._chat_once(backend, m, messages, temperature, max_tokens)
|
||
except Exception as e: # noqa: BLE001 重试/切换模型或后端，异常向上收敛
|
||
logger.warning(
|
||
"chat backend=%s model=%r failed: %s: %s",
|
||
backend.name, m, type(e).__name__, e,
|
||
)
|
||
last_err = e
|
||
if index < len(self.backends) - 1:
|
||
logger.warning("backend=%s exhausted, falling back to %s",
|
||
backend.name, self.backends[index + 1].name)
|
||
raise LLMFailError() from last_err
|
||
|
||
async def _chat_once(
|
||
self,
|
||
backend: Backend,
|
||
model: str,
|
||
messages: list[dict],
|
||
temperature: float | None,
|
||
max_tokens: int | None,
|
||
) -> str:
|
||
request_max_tokens = self.cfg.max_tokens if max_tokens is None else max_tokens
|
||
url = f"{backend.base_url}/chat/completions"
|
||
# 指数退避：retry_backoff_sec → 翻倍 → …，最多 max_retries 次
|
||
for attempt in range(self.cfg.max_retries):
|
||
payload = {
|
||
"model": model,
|
||
"messages": messages,
|
||
"temperature": self.cfg.temperature if temperature is None else temperature,
|
||
"max_tokens": request_max_tokens,
|
||
"stream": False,
|
||
}
|
||
try:
|
||
async with backend.client(self.cfg.timeout) as client:
|
||
r = await client.post(url, headers=backend.headers, json=payload)
|
||
r.raise_for_status()
|
||
choice = r.json()["choices"][0]
|
||
message = choice["message"]
|
||
content = message.get("content")
|
||
if isinstance(content, str) and content.strip():
|
||
return content
|
||
|
||
# DeepSeek reasoning models may consume the entire small token budget
|
||
# in reasoning_content, leaving content empty and finish_reason=length.
|
||
if choice.get("finish_reason") == "length" and attempt < self.cfg.max_retries - 1:
|
||
request_max_tokens = max(request_max_tokens * 4, 256)
|
||
logger.warning(
|
||
"chat backend=%s model=%r returned empty content after token limit; "
|
||
"retrying with max_tokens=%d",
|
||
backend.name,
|
||
model,
|
||
request_max_tokens,
|
||
)
|
||
continue
|
||
raise ValueError("LLM 返回空内容")
|
||
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 一一对应的向量列表。
|
||
|
||
只走首选后端（见模块说明）。按 OpenAI 兼容协议传 `dimensions`（来自
|
||
LLM_EMBED_DIMENSIONS），MRL 模型（qwen3-embedding 等）会截断并重新归一化到目标维度。
|
||
"""
|
||
backend = self.primary
|
||
embed_model = (
|
||
getattr(self.cfg, "api_embed_model", "")
|
||
if backend.name == "api"
|
||
else backend.embed_model
|
||
) or backend.embed_model
|
||
if not embed_model:
|
||
raise LLMFailError(f"backend={backend.name} 未配置 embed 模型")
|
||
is_dashscope_multimodal = "vision" in embed_model.lower() or "multimodal" in embed_model.lower()
|
||
if is_dashscope_multimodal:
|
||
base_url = backend.base_url.rstrip("/")
|
||
marker = "/compatible-mode/v1"
|
||
if base_url.lower().endswith(marker):
|
||
base_url = base_url[: -len(marker)]
|
||
url = f"{base_url}/api/v1/services/embeddings/multimodal-embedding/multimodal-embedding"
|
||
payload = {
|
||
"model": embed_model,
|
||
"input": {"contents": [{"text": value} for value in texts]},
|
||
}
|
||
else:
|
||
url = f"{backend.base_url}/embeddings"
|
||
payload = {
|
||
"model": embed_model,
|
||
"input": texts,
|
||
"dimensions": self.cfg.embed_dimensions,
|
||
}
|
||
response = None
|
||
max_retries = max(1, int(getattr(self.cfg, "max_retries", 1)))
|
||
retry_backoff_sec = float(getattr(self.cfg, "retry_backoff_sec", 0))
|
||
async with backend.client(self.cfg.timeout) as client:
|
||
for attempt in range(max_retries):
|
||
try:
|
||
response = await client.post(
|
||
url,
|
||
headers=backend.headers,
|
||
json=payload,
|
||
)
|
||
response.raise_for_status()
|
||
break
|
||
except Exception:
|
||
if attempt == max_retries - 1:
|
||
raise
|
||
await asyncio.sleep(retry_backoff_sec * (2**attempt))
|
||
if response is None: # pragma: no cover - defensive guard
|
||
raise LLMFailError("Embedding 请求未返回响应")
|
||
data = response.json()
|
||
if is_dashscope_multimodal:
|
||
embeddings = data["output"]["embeddings"]
|
||
return [
|
||
item["embedding"]
|
||
for item in sorted(embeddings, key=lambda item: item["index"])
|
||
]
|
||
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 也算可达）。"""
|
||
backend = self.primary
|
||
if backend.is_ollama:
|
||
url = backend.base_url.removesuffix("/v1") + "/api/tags"
|
||
else:
|
||
url = f"{backend.base_url}/models"
|
||
async with backend.client(min(self.cfg.timeout, 10)) as client:
|
||
r = await client.get(url, headers=backend.headers)
|
||
if r.status_code >= 500:
|
||
r.raise_for_status()
|
||
|
||
|
||
# 全局单例：Agent 统一引入
|
||
llm = LLMClient()
|