211 lines
8.8 KiB
Python
211 lines
8.8 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""推理后端实现 —— 本地 70B 模型接入与推理封装（issue #44）。
|
||
|
||
在 `gateway.InferenceBackend` 抽象之上交付**真实可用的本地后端**：
|
||
- OpenAI 兼容接口（vLLM / TGI 等本地推理服务，`/v1/chat/completions`），
|
||
仅用标准库 urllib，无第三方依赖；
|
||
- 参数化：endpoint / model / timeout / max_tokens / temperature / context 引用注入；
|
||
- **数据不出厂**（PRD 5.4）：敏感/核心内容走本地后端，云端仅接收脱敏内容；
|
||
- 未配置 endpoint 时进入 dry-run 占位模式（保持与旧 LocalBackend 一致的
|
||
可测试行为，供端到端演示与联调）。
|
||
|
||
业务代码只依赖 `gateway.InferenceBackend.generate(prompt, context)`，
|
||
切换后端 = 换实现（见 `LLMGateway(local=...)`）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import os
|
||
import time
|
||
import urllib.request
|
||
from typing import Callable, Optional, Sequence
|
||
|
||
from .gateway import InferenceBackend
|
||
|
||
|
||
class Local70BBackend(InferenceBackend):
|
||
"""本地 70B 推理后端（OpenAI 兼容 vLLM/TGI，参数化）。"""
|
||
|
||
name = "local-70b"
|
||
|
||
def __init__(
|
||
self,
|
||
endpoint: str = "",
|
||
model: str = "iaop-local-70b",
|
||
timeout_seconds: float = 60.0,
|
||
max_tokens: int = 1024,
|
||
temperature: float = 0.1,
|
||
echo_context: bool = True,
|
||
) -> None:
|
||
self.endpoint = (endpoint or "").rstrip("/")
|
||
self.model = model
|
||
self.timeout = float(timeout_seconds)
|
||
self.max_tokens = int(max_tokens)
|
||
self.temperature = float(temperature)
|
||
self.echo_context = echo_context
|
||
|
||
# ------------------------------------------------------------------
|
||
def generate(self, prompt: str, context: Sequence[str]) -> str:
|
||
"""根据 prompt 与 RAG 上下文生成回答。
|
||
|
||
- 未配置 endpoint：dry-run 占位（回显 prompt 前 40 字符 + 来源引用）；
|
||
- 已配置：调用本地 OpenAI 兼容服务（/v1/chat/completions）。
|
||
"""
|
||
if not self.endpoint:
|
||
return self._dry_run(prompt, context)
|
||
|
||
payload = {
|
||
"model": self.model,
|
||
"messages": [
|
||
{"role": "system", "content": self._system_prompt(context)},
|
||
{"role": "user", "content": prompt},
|
||
],
|
||
"max_tokens": self.max_tokens,
|
||
"temperature": self.temperature,
|
||
}
|
||
body = self._post_json("/v1/chat/completions", payload)
|
||
try:
|
||
return body["choices"][0]["message"]["content"]
|
||
except (KeyError, IndexError, TypeError):
|
||
raise RuntimeError(
|
||
f"本地推理服务响应格式异常: {str(body)[:200]}")
|
||
|
||
# ------------------------------------------------------------------
|
||
def _system_prompt(self, context: Sequence[str]) -> str:
|
||
"""把 RAG 引用注入 system 提示（引用溯源，PRD 5.4）。"""
|
||
refs = "\n".join(f"- {c}" for c in (context or []))
|
||
base = "你是工业 AI 助手。回答须基于给定资料并标注来源。"
|
||
return f"{base}\n参考资料：\n{refs}" if refs else base
|
||
|
||
def _dry_run(self, prompt: str, context: Sequence[str]) -> str:
|
||
head = f"[本地70B占位] {prompt[:40]}"
|
||
if self.echo_context:
|
||
for i, src in enumerate(context[:3], 1):
|
||
head += f"\n[来源: {src}]"
|
||
return head
|
||
|
||
def _post_json(self, path: str, payload: dict) -> dict:
|
||
"""向后端推理服务发起 JSON POST（标准库 urllib）。"""
|
||
url = self.endpoint + path
|
||
data = json.dumps(payload).encode("utf-8")
|
||
req = urllib.request.Request(
|
||
url, data=data,
|
||
headers={"Content-Type": "application/json"})
|
||
with urllib.request.urlopen(req, timeout=self.timeout) as resp:
|
||
raw = resp.read().decode("utf-8")
|
||
return json.loads(raw) if raw else {}
|
||
|
||
def health(self) -> dict:
|
||
"""后端健康信息（本地推理服务可探测 /health）。"""
|
||
base = {
|
||
"backend": self.name, "model": self.model,
|
||
"endpoint": self.endpoint or "(dry-run)",
|
||
}
|
||
if not self.endpoint:
|
||
base["status"] = "dry-run"
|
||
return base
|
||
try:
|
||
started = time.monotonic()
|
||
with urllib.request.urlopen(
|
||
self.endpoint + "/health", timeout=self.timeout) as resp:
|
||
base["status"] = "ok" if resp.status == 200 else f"http-{resp.status}"
|
||
base["latency_ms"] = round((time.monotonic() - started) * 1000, 2)
|
||
except Exception as exc: # noqa: BLE001 - 健康探测失败仅记录
|
||
base["status"] = f"error: {exc}"
|
||
return base
|
||
|
||
|
||
class CloudApiBackend(InferenceBackend):
|
||
"""云端 API 推理后端（Qwen / DeepSeek 等 OpenAI 兼容）—— issue #45。
|
||
|
||
**安全网关约束（PRD 5.4）**：
|
||
- 仅接收 **DLP 放行**的脱敏/通用内容（上游 `LLMGateway` 主编排出站检查 +
|
||
cloud 分支输出 DLP 复查）；
|
||
- API Key 从**环境变量**读取（`api_key_env`），不硬编码、不落日志；
|
||
- 可选 `safety_checker` 出站复查钩子（fail-closed：复查拒绝 → 拦截占位，
|
||
不调用上游）。
|
||
"""
|
||
|
||
name = "cloud-api"
|
||
|
||
def __init__(
|
||
self,
|
||
endpoint: str = "",
|
||
api_key_env: str = "",
|
||
model: str = "deepseek-chat",
|
||
timeout_seconds: float = 60.0,
|
||
max_tokens: int = 1024,
|
||
temperature: float = 0.1,
|
||
safety_checker: Optional[Callable[[str], bool]] = None,
|
||
) -> None:
|
||
self.endpoint = (endpoint or "").rstrip("/")
|
||
self.api_key_env = api_key_env
|
||
self.model = model
|
||
self.timeout = float(timeout_seconds)
|
||
self.max_tokens = int(max_tokens)
|
||
self.temperature = float(temperature)
|
||
# 出站安全复查：返回 False 即拦截（fail-closed）
|
||
self.safety_checker = safety_checker
|
||
self._api_key = os.environ.get(api_key_env, "") if api_key_env else ""
|
||
|
||
# ------------------------------------------------------------------
|
||
def generate(self, prompt: str, context: Sequence[str]) -> str:
|
||
"""生成回答。安全网关：safety_checker 拒绝 → 拦截占位，不调用上游。"""
|
||
if self.safety_checker is not None and not self.safety_checker(prompt):
|
||
return "[云端安全网关拦截] 出站复查未通过，已拦截（数据不出厂）。"
|
||
|
||
if not self.endpoint:
|
||
return self._dry_run(prompt, context)
|
||
|
||
payload = {
|
||
"model": self.model,
|
||
"messages": [
|
||
{"role": "system", "content": self._system_prompt(context)},
|
||
{"role": "user", "content": prompt},
|
||
],
|
||
"max_tokens": self.max_tokens,
|
||
"temperature": self.temperature,
|
||
}
|
||
body = self._post_json("/v1/chat/completions", payload)
|
||
try:
|
||
return body["choices"][0]["message"]["content"]
|
||
except (KeyError, IndexError, TypeError):
|
||
raise RuntimeError(
|
||
f"云端 API 响应格式异常: {str(body)[:200]}")
|
||
|
||
# ------------------------------------------------------------------
|
||
def _system_prompt(self, context: Sequence[str]) -> str:
|
||
refs = "\n".join(f"- {c}" for c in (context or []))
|
||
base = "你是工业 AI 助手。回答须基于给定资料并标注来源。"
|
||
return f"{base}\n参考资料：\n{refs}" if refs else base
|
||
|
||
def _dry_run(self, prompt: str, context: Sequence[str]) -> str:
|
||
head = f"[云端API占位] {prompt[:40]}"
|
||
for i, src in enumerate(context[:3], 1):
|
||
head += f"\n[来源: {src}]"
|
||
if self.safety_checker is not None:
|
||
head += "\n[安全网关: 已复查放行]"
|
||
return head
|
||
|
||
def _post_json(self, path: str, payload: dict) -> dict:
|
||
"""向后端推理服务发起 JSON POST（Bearer 认证，Key 来自环境变量）。"""
|
||
url = self.endpoint + path
|
||
data = json.dumps(payload).encode("utf-8")
|
||
headers = {"Content-Type": "application/json"}
|
||
if self._api_key:
|
||
headers["Authorization"] = f"Bearer {self._api_key}"
|
||
req = urllib.request.Request(url, data=data, headers=headers)
|
||
with urllib.request.urlopen(req, timeout=self.timeout) as resp:
|
||
raw = resp.read().decode("utf-8")
|
||
return json.loads(raw) if raw else {}
|
||
|
||
def health(self) -> dict:
|
||
"""后端健康信息（含安全网关状态，不含密钥）。"""
|
||
return {
|
||
"backend": self.name, "model": self.model,
|
||
"endpoint": self.endpoint or "(dry-run)",
|
||
"api_key_configured": bool(self._api_key),
|
||
"safety_checker": self.safety_checker is not None,
|
||
"status": "dry-run" if not self.endpoint else "configured",
|
||
}
|