271 lines
10 KiB
Python
271 lines
10 KiB
Python
"""专家模型池：统一 Expert 接口，支持三种后端。
|
||
|
||
- MockExpert ：确定性模板输出（零依赖，离线可跑，便于测试与演示）
|
||
- HFExpert ：HuggingFace transformers 真实小模型（可选，需 ML 依赖）
|
||
- APIExpert ：OpenAI 兼容 API（可选，需 API Key，如 DeepSeek）
|
||
|
||
成本估计：cost_est 按参数量粗估（美元/百万 token 的近似比例）。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import re
|
||
from typing import Dict, List, Optional
|
||
|
||
from .models import ExpertResponse
|
||
|
||
# 按模型规模粗估的相对成本（$ / 1M output tokens，近似）
|
||
MODEL_COST_EST = {
|
||
"mock": 0.0,
|
||
"0.5b": 0.02,
|
||
"1b": 0.05,
|
||
"1.7b": 0.08,
|
||
"3b": 0.12,
|
||
"4b": 0.15,
|
||
"7b": 0.25,
|
||
"70b": 2.50,
|
||
"api": 1.00,
|
||
}
|
||
|
||
|
||
def _cost_for(model_name: str, default: str = "1b") -> float:
|
||
mn = model_name.lower()
|
||
for key in ("0.5b", "1.7b", "3b", "4b", "7b", "70b"):
|
||
if key in mn:
|
||
return MODEL_COST_EST[key]
|
||
if "api" in mn or mn in ("deepseek-chat", "gpt-4o-mini", "claude"):
|
||
return MODEL_COST_EST["api"]
|
||
return MODEL_COST_EST.get(default, 0.1)
|
||
|
||
|
||
def extract_content_terms(query: str) -> List[str]:
|
||
"""抽取查询中的"内容词"（中文词/英文单词），用于 Judge 覆盖度与 Mock 回显。"""
|
||
q = query.lower()
|
||
terms: List[str] = []
|
||
# 英文单词（>=2 字符）
|
||
for w in re.findall(r"[a-z][a-z0-9_]{1,}", q):
|
||
if w not in _STOPWORDS_EN and w not in terms:
|
||
terms.append(w)
|
||
# 中文：按 2-4 字窗口切分，保留含中文字符的片段
|
||
cn = re.findall(r"[\u4e00-\u9fff]{2,8}", q)
|
||
for c in cn:
|
||
terms.append(c)
|
||
return terms
|
||
|
||
|
||
_STOPWORDS_EN = {
|
||
"the", "a", "an", "is", "are", "to", "of", "in", "on", "for", "with", "and",
|
||
"or", "do", "does", "can", "could", "would", "should", "please", "me", "my",
|
||
"this", "that", "it", "be", "was", "were", "have", "has", "had", "will",
|
||
"not", "no", "yes", "i", "you", "he", "she", "we", "they",
|
||
}
|
||
|
||
|
||
class Expert:
|
||
name: str = "expert"
|
||
|
||
async def generate(self, query: str, difficulty: str) -> ExpertResponse:
|
||
raise NotImplementedError
|
||
|
||
|
||
class MockExpert(Expert):
|
||
"""确定性模板专家：零依赖，离线可跑。
|
||
|
||
输出会回显查询中的内容词以提高 Judge 覆盖度，并带领域结构，
|
||
使端到端管线（分类 -> 专家 -> Judge -> 缓存）可被稳定测试与演示。
|
||
"""
|
||
|
||
def __init__(self, name: str, domain: str, model: str = "mock"):
|
||
self.name = name
|
||
self.domain = domain
|
||
self.model = model
|
||
|
||
async def generate(self, query: str, difficulty: str) -> ExpertResponse:
|
||
await asyncio.sleep(0.001) # 模拟极短推理延迟
|
||
terms = extract_content_terms(query)
|
||
body = self._template(query, terms, difficulty)
|
||
# 预估 token 数：中文约 1.5 字符/token，英文约 4 字符/token
|
||
tokens = max(8, int(len(body) / 2.2))
|
||
return ExpertResponse(
|
||
text=body,
|
||
model_used=self.model,
|
||
latency_ms=1.0,
|
||
tokens=tokens,
|
||
cost_est=_cost_for(self.model) * tokens / 1_000_000,
|
||
)
|
||
|
||
def _template(self, query: str, terms: List[str], difficulty: str) -> str:
|
||
kw = "、".join(terms[:6]) if terms else "该主题"
|
||
if self.domain == "code":
|
||
return (
|
||
f"（mock 代码专家）针对「{query}」的实现思路如下：\n\n"
|
||
f"```python\n"
|
||
f"def solve() -> None:\n"
|
||
f" # 关键点：{kw}\n"
|
||
f" # 1. 明确输入输出约束\n"
|
||
f" # 2. 选择合适数据结构\n"
|
||
f" # 3. 处理边界条件（空输入、极端值）\n"
|
||
f" # 4. 补充单元测试\n"
|
||
f" pass\n"
|
||
f"```\n\n"
|
||
f"复杂度：平均 O(n)。请按上述步骤补充具体实现。"
|
||
)
|
||
if self.domain == "math":
|
||
return (
|
||
f"（mock 数学专家）求解「{query}」的步骤：\n\n"
|
||
f"1. 明确已知条件与目标：{kw}\n"
|
||
f"2. 选择合适的方法（代数变形 / 积分 / 归纳等）\n"
|
||
f"3. 逐步推导并验证中间结果\n"
|
||
f"4. 检查边界与特殊情况\n\n"
|
||
f"结论：在标准假设下，结果可化简为闭合形式。完整推导见正式解答。"
|
||
)
|
||
if self.domain == "legal":
|
||
return (
|
||
f"（mock 法律专家）关于「{query}」的初步法律分析：\n\n"
|
||
f"相关要点：{kw}\n"
|
||
f"1. 适用法规：请以现行有效法条为准（建议核对最新修订版）\n"
|
||
f"2. 合同/合规风险点识别\n"
|
||
f"3. 责任划分与救济途径\n\n"
|
||
f"⚠️ 提示：以上为一般性分析，不构成正式法律意见，个案请咨询执业律师。"
|
||
)
|
||
if self.domain == "medical":
|
||
return (
|
||
f"（mock 医学专家）关于「{query}」的科普性说明：\n\n"
|
||
f"相关关键词：{kw}\n"
|
||
f"1. 常见表现与可能原因\n"
|
||
f"2. 一般处理建议与注意事项\n"
|
||
f"3. 何时需要就医（警示信号）\n\n"
|
||
f"⚠️ 提示：内容仅供健康科普，不能替代医生诊断；如有不适请及时就医。"
|
||
)
|
||
return (
|
||
f"（mock 通用专家）关于「{query}」的回答：\n\n"
|
||
f"核心要点：{kw}\n"
|
||
f"1. 背景与定义\n"
|
||
f"2. 主要分类/维度\n"
|
||
f"3. 实际应用与注意事项\n\n"
|
||
f"如需更深入的分析，可以补充更多上下文。"
|
||
)
|
||
|
||
|
||
class HFExpert(Expert):
|
||
"""可选：HuggingFace 真实小模型（需 requirements-ml.txt）。"""
|
||
|
||
def __init__(self, name: str, domain: str, model: str):
|
||
self.name = name
|
||
self.domain = domain
|
||
self.model = model
|
||
self._loaded = False
|
||
self._model = None
|
||
self._tokenizer = None
|
||
|
||
def _ensure_loaded(self):
|
||
if self._loaded:
|
||
return
|
||
try:
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
except ImportError as e:
|
||
raise RuntimeError("HFExpert 需要安装 ML 依赖：pip install -r requirements-ml.txt") from e
|
||
self._tokenizer = AutoTokenizer.from_pretrained(self.model)
|
||
self._model = AutoModelForCausalLM.from_pretrained(
|
||
self.model, device_map="auto", torch_dtype="auto"
|
||
)
|
||
self._loaded = True
|
||
|
||
async def generate(self, query: str, difficulty: str) -> ExpertResponse:
|
||
self._ensure_loaded()
|
||
return await asyncio.to_thread(self._generate_sync, query)
|
||
|
||
def _generate_sync(self, query: str) -> ExpertResponse:
|
||
messages = [{"role": "user", "content": query}]
|
||
text = self._tokenizer.apply_chat_template(messages, tokenize=False)
|
||
inputs = self._tokenizer(text, return_tensors="pt").to(self._model.device)
|
||
outputs = self._model.generate(
|
||
**inputs,
|
||
max_new_tokens=512,
|
||
temperature=0.2,
|
||
do_sample=True,
|
||
)
|
||
body = self._tokenizer.decode(outputs[0][inputs["input_ids"].shape[1]:], skip_special_tokens=True)
|
||
return ExpertResponse(
|
||
text=body,
|
||
model_used=self.model,
|
||
latency_ms=0.0,
|
||
tokens=512,
|
||
cost_est=_cost_for(self.model) * 512 / 1_000_000,
|
||
)
|
||
|
||
|
||
class APIExpert(Expert):
|
||
"""可选：OpenAI 兼容 Chat Completions（DeepSeek / OpenAI / 本地 vLLM）。"""
|
||
|
||
def __init__(self, name: str, domain: str, model: str, base_url: str, api_key: str):
|
||
self.name = name
|
||
self.domain = domain
|
||
self.model = model
|
||
self.base_url = base_url.rstrip("/")
|
||
self.api_key = api_key
|
||
self._client = None
|
||
|
||
def _get_client(self):
|
||
if self._client is None:
|
||
import httpx
|
||
self._client = httpx.AsyncClient(timeout=60.0)
|
||
return self._client
|
||
|
||
async def generate(self, query: str, difficulty: str) -> ExpertResponse:
|
||
client = self._get_client()
|
||
resp = await client.post(
|
||
f"{self.base_url}/chat/completions",
|
||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||
json={
|
||
"model": self.model,
|
||
"messages": [{"role": "user", "content": query}],
|
||
"temperature": 0.2,
|
||
"max_tokens": 1024,
|
||
},
|
||
)
|
||
resp.raise_for_status()
|
||
data = resp.json()
|
||
body = data["choices"][0]["message"]["content"]
|
||
usage = data.get("usage", {})
|
||
tokens = usage.get("completion_tokens", int(len(body) / 2.2))
|
||
return ExpertResponse(
|
||
text=body,
|
||
model_used=self.model,
|
||
latency_ms=0.0,
|
||
tokens=tokens,
|
||
cost_est=_cost_for(self.model) * tokens / 1_000_000,
|
||
)
|
||
|
||
|
||
def build_expert(domain: str, cfg: Dict) -> Expert:
|
||
"""根据配置构建领域专家。cfg 为 experts.<domain> 段配置。"""
|
||
etype = cfg.get("type", "mock")
|
||
model = cfg.get("model", "mock")
|
||
name = f"expert-{domain}"
|
||
if etype == "mock":
|
||
return MockExpert(name, domain, model)
|
||
if etype == "hf":
|
||
return HFExpert(name, domain, model)
|
||
if etype == "api":
|
||
base_url = cfg.get("base_url", "https://api.deepseek.com/v1")
|
||
api_key = cfg.get("api_key") or _env(cfg.get("api_key_env", ""))
|
||
if not api_key:
|
||
raise RuntimeError(f"APIExpert({domain}) 缺少 API Key（env: {cfg.get('api_key_env')}）")
|
||
return APIExpert(name, domain, model, base_url, api_key)
|
||
raise ValueError(f"未知专家后端类型: {etype}（支持 mock | hf | api）")
|
||
|
||
|
||
def _env(name: str) -> Optional[str]:
|
||
import os
|
||
return os.environ.get(name) if name else None
|
||
|
||
|
||
def build_expert_pool(experts_cfg: Dict[str, Dict], domains: List[str]) -> Dict[str, Expert]:
|
||
"""构建完整专家池。"""
|
||
pool: Dict[str, Expert] = {}
|
||
for domain in domains:
|
||
cfg = experts_cfg.get(domain, {"type": "mock", "model": "mock"})
|
||
pool[domain] = build_expert(domain, cfg)
|
||
return pool
|
||
|