"""语义向量通道(可选增强):本地 BGE-M3 权重 或 联网 Embedding API 双模式。 - provider=local:sentence-transformers 加载本地权重(如 models/bge-m3); - provider=api:OpenAI 兼容 /embeddings 接口(默认阿里 DashScope 兼容模式, 模型 text-embedding-v4,填 API Key 即可,无需下载权重); - 任一模式就绪即启用语义召回;均不可用时回退词表归一化,行为与无此模块一致; - 依赖与权重均为可选安装,不进入 requirements("零重依赖、离线优先"约束)。 """ from __future__ import annotations import importlib.util import logging from pathlib import Path import httpx logger = logging.getLogger(__name__) # BGE-M3 目录的就绪判定:配置文件 + 任一权重文件 + 分词器 _REQUIRED_CONFIG = "config.json" _WEIGHT_FILES = ("model.safetensors", "pytorch_model.bin") _TOKENIZER_FILES = ("tokenizer.json", "vocab.txt") DEFAULT_API_BASE = "https://dashscope.aliyuncs.com/compatible-mode/v1" DEFAULT_API_MODEL = "text-embedding-v4" _model = None # 懒加载的本地 SentenceTransformer 实例 _model_path = "" _skill_emb = None # 词表技能向量缓存 _skill_names: list[str] = [] _emb_cache_key = "" # 缓存对应的通道标识(provider+路径/模型) def deps_installed() -> bool: """sentence-transformers 是否已安装。""" return importlib.util.find_spec("sentence_transformers") is not None def resolve_weights_path(path: str | Path) -> Path: """权重路径解析:相对路径相对工程根目录解析,绝对路径原样返回(兼容既有配置)。""" from app.config import ROOT p = Path(str(path or "")) return p if p.is_absolute() else ROOT / p def weights_present(path: str | Path) -> bool: """权重目录是否包含 BGE-M3 所需文件。""" p = Path(str(path or "")) if not p.is_dir() or not (p / _REQUIRED_CONFIG).exists(): return False if not any((p / f).exists() for f in _WEIGHT_FILES): return False return any((p / f).exists() for f in _TOKENIZER_FILES) def status(cfg: dict) -> dict: """语义通道状态:供 /api/health 与 /api/config 展示。""" provider = str(cfg.get("embed_provider") or "local").strip() or "local" path = str(cfg.get("embed_model_path") or "").strip() api_key = str(cfg.get("embed_api_key") or "").strip() present = weights_present(resolve_weights_path(path)) if path else False deps = deps_installed() local_ready = bool(path and present and deps) api_ready = bool(api_key) ready = api_ready if provider == "api" else local_ready return { "provider": provider, "configured": bool(path) if provider == "local" else api_ready, "path": path, "weights_present": present, "deps_installed": deps, "local_ready": local_ready, "api_key_masked": ("••••" + api_key[-4:]) if api_key else "", "api_base_url": str(cfg.get("embed_api_base_url") or DEFAULT_API_BASE), "api_model": str(cfg.get("embed_api_model") or DEFAULT_API_MODEL), "ready": bool(ready), } # ---------------------------------------------------------------- 本地模式 def _load_local(path: str): """懒加载本地模型;不可用时返回 None。""" global _model, _model_path if _model is not None and _model_path == path: return _model try: from sentence_transformers import SentenceTransformer # 可选依赖 _model = SentenceTransformer(path) _model_path = path logger.info("语义向量模型已加载(本地): %s", path) except Exception as exc: # noqa: BLE001 - 模型加载失败必须可回退 logger.warning("本地语义模型加载失败,回退词表归一化: %s", exc) _model = None return None return _model def _embed_local(texts: list[str], path: str): model = _load_local(path) if model is None: return None return model.encode(texts, normalize_embeddings=True) # ---------------------------------------------------------------- 联网模式 def _embed_api(texts: list[str], cfg: dict): """调用 OpenAI 兼容 embeddings 接口(DashScope 等),返回向量列表。""" st = status(cfg) try: resp = httpx.post( f"{st['api_base_url'].rstrip('/')}/embeddings", headers={"Authorization": f"Bearer {cfg.get('embed_api_key')}"}, json={"model": st["api_model"], "input": texts}, timeout=30, ) resp.raise_for_status() data = resp.json()["data"] data.sort(key=lambda d: d["index"]) import math vecs = [] for d in data: v = d["embedding"] norm = math.sqrt(sum(x * x for x in v)) or 1.0 vecs.append([x / norm for x in v]) return vecs except Exception as exc: # noqa: BLE001 - 网络/鉴权失败必须可回退 logger.warning("Embedding API 调用失败,回退词表归一化: %s", exc) return None # ---------------------------------------------------------------- 统一入口 def _embed(texts: list[str], cfg: dict): """按 provider 取向量(已归一化);失败返回 None。""" provider = str(cfg.get("embed_provider") or "local").strip() or "local" if provider == "api": return _embed_api(texts, cfg) return _embed_local(texts, str(resolve_weights_path( str(cfg.get("embed_model_path") or "").strip()))) def most_similar(name: str, cfg: dict, lexicon: dict, threshold: float = 0.62) -> tuple[str, float] | None: """返回词表中与 name 语义最近的技能 (标准名, 相似度);未过阈值或不可用返回 None。""" global _skill_emb, _skill_names, _emb_cache_key st = status(cfg) if not st["ready"]: return None cache_key = f"{st['provider']}|{st['path'] or st['api_model']}" if _skill_emb is None or cache_key != _emb_cache_key: _skill_names = [s["name"] for s in lexicon["skills"]] _skill_emb = _embed(_skill_names, cfg) if _skill_emb is None: return None _emb_cache_key = cache_key vecs = _embed([name], cfg) if vecs is None: return None vec = vecs[0] scores = [sum(a * b for a, b in zip(row, vec)) for row in _skill_emb] best_i = max(range(len(scores)), key=lambda i: scores[i]) if scores[best_i] >= threshold: return _skill_names[best_i], round(scores[best_i], 4) return None