"""retrieval.py —— faiss 向量检索 + index.json 单一事实源。 - IndexFlatIP（内积，配合 L2 归一化向量即余弦相似度） - index.json：检索元数据单一事实源（行号 == json 下标）；faiss 内存索引 可随时由 index.json 重建 - search 三档降级：向量 → meta 关键词 → 最新 N 条 - 增量写锁（threading.Lock）；rebuild 临时索引 + os.replace 原子切换 """ from __future__ import annotations import json import logging import os import threading from pathlib import Path from typing import Any import faiss import numpy as np from .settings import ( DEFAULT_TOP_K, EMBED_DIM, FALLBACK_LATEST_N, MAX_TOP_K, ) logger = logging.getLogger(__name__) INDEX_FILE = "index.json" INDEX_TMP_SUFFIX = ".rebuild.tmp" class RetrievalError(Exception): """检索/索引操作失败。""" def _now_iso() -> str: from datetime import datetime, timezone return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") def _normalize(vec: list[float]) -> np.ndarray: """L2 归一化（IndexFlatIP 下等价余弦）。零向量原样返回。""" arr = np.asarray(vec, dtype=np.float32) norm = float(np.linalg.norm(arr)) if norm > 0: arr = arr / norm return arr class RetrievalIndex: """检索索引：index.json + faiss 内存索引（同一目录，单一事实源）。""" def __init__(self, index_path: str | Path, dim: int = EMBED_DIM) -> None: self.dir = Path(index_path) self.dim = dim self._lock = threading.Lock() self._faiss: faiss.IndexFlatIP = faiss.IndexFlatIP(dim) self._rows: list[dict[str, Any]] = [] # ---- 加载 / 落盘 ---- @property def json_path(self) -> Path: return self.dir / INDEX_FILE def load(self) -> None: """从 index.json 加载行并重建 faiss。文件不存在视为空库。""" with self._lock: self._rows = [] self._faiss = faiss.IndexFlatIP(self.dim) path = self.json_path if not path.exists(): return try: data = json.loads(path.read_text(encoding="utf-8")) items = data.get("items", []) if isinstance(data, dict) else [] except (ValueError, OSError) as exc: logger.warning("retrieval: index.json 读取失败，按空库处理：%s", exc) return for row_id, item in enumerate(items): if not isinstance(item, dict): continue row = dict(item) row["row_id"] = row_id self._rows.append(row) vec = row.get("vector") if isinstance(vec, list) and vec: self._faiss.add(_normalize(vec).reshape(1, self.dim)) def _dump_locked(self, path: Path) -> None: """写 json（需持有锁）。""" tmp = path.with_name(path.name + INDEX_TMP_SUFFIX) payload = {"version": 1, "items": [{k: v for k, v in r.items() if k != "row_id"} for r in self._rows]} tmp.write_text(json.dumps(payload, ensure_ascii=False), encoding="utf-8") os.replace(tmp, path) def save(self) -> None: """落盘（原子）。""" with self._lock: self.dir.mkdir(parents=True, exist_ok=True) self._dump_locked(self.json_path) # ---- 增删 ---- def add_item(self, item: dict[str, Any], vector: list[float]) -> int: """追加一行（json 下标 == 行号）。返回 row_id。""" with self._lock: self.dir.mkdir(parents=True, exist_ok=True) row_id = len(self._rows) row = dict(item) row["row_id"] = row_id row["vector"] = list(vector) self._rows.append(row) self._faiss.add(_normalize(vector).reshape(1, self.dim)) self._dump_locked(self.json_path) return row_id def remove_item(self, document_id: str) -> bool: """按 document_id 删除（faiss 不支持逐行删，整库重建——库规模小可接受）。""" with self._lock: before = len(self._rows) self._rows = [r for r in self._rows if str(r.get("id")) != document_id] self._faiss = faiss.IndexFlatIP(self.dim) for row in self._rows: vec = row.get("vector") if isinstance(vec, list) and vec: self._faiss.add(_normalize(vec).reshape(1, self.dim)) if self.dir.exists(): self._dump_locked(self.json_path) return len(self._rows) < before def update_meta(self, document_id: str, meta: dict[str, Any]) -> bool: """更新行的 meta 字段（向量不变）。""" with self._lock: for row in self._rows: if str(row.get("id")) == document_id: for key, value in meta.items(): if key != "vector": row[key] = value if self.dir.exists(): self._dump_locked(self.json_path) return True return False # ---- 查询 ---- def rows(self) -> list[dict[str, Any]]: """行快照（不含 vector，供列表接口）。""" with self._lock: return [{k: v for k, v in r.items() if k != "vector"} for r in self._rows] def count(self) -> int: with self._lock: return len(self._rows) def _match_meta(self, row: dict[str, Any], keywords: list[str]) -> bool: """meta 关键词匹配：name/category/tags/description 任一命中任一关键词。""" if not keywords: return True haystack = " ".join( str(row.get(k) or "") for k in ("name", "category", "tags", "description") ).lower() return any(kw.lower() in haystack for kw in keywords) def search( self, query_vector: list[float], top_k: int = DEFAULT_TOP_K, category: str | None = None, as_template: bool | None = None, keywords: list[str] | None = None, ) -> dict[str, Any]: """三档降级检索：向量 → meta 关键词 → 最新 N 条。 返回 {"mode": "vector"|"keyword"|"latest"|"empty", "items": [...]}。 """ top_k = max(1, min(int(top_k), MAX_TOP_K)) with self._lock: rows = self._rows def _filtered(candidates: list[dict[str, Any]]) -> list[dict[str, Any]]: out = [] for r in candidates: if category and str(r.get("category", "")) != category: continue if as_template is not None and bool(r.get("as_template")) != bool(as_template): continue out.append(r) return out # 一档：向量 if rows: try: q = _normalize(query_vector).reshape(1, self.dim) dist, idx = self._faiss.search(q, min(top_k, len(rows))) vector_hits = [] for score, row_idx in zip(dist[0].tolist(), idx[0].tolist()): if 0 <= row_idx < len(rows): vector_hits.append((float(score), rows[row_idx])) filtered = _filtered([r for _, r in vector_hits]) # 分数阈值：零向量查询（或过滤后无相关命中）时 # 最高分仍 ≈ 0，视为无命中 → 降级下一档 best_score = vector_hits[0][0] if vector_hits else 0.0 if filtered and best_score > 0.01: return { "mode": "vector", "items": [ {k: v for k, v in r.items() if k != "vector"} for r in filtered[:top_k] ], } except Exception as exc: logger.warning("retrieval: 向量检索失败，降级关键词：%s", exc) # 二档：meta 关键词 if keywords: kw_hits = [r for r in rows if self._match_meta(r, keywords)] filtered = _filtered(kw_hits) if filtered: return { "mode": "keyword", "items": [ {k: v for k, v in r.items() if k != "vector"} for r in filtered[:top_k] ], } # 三档：最新 N 条（created_at 倒序） latest = sorted( rows, key=lambda r: str(r.get("created_at") or ""), reverse=True ) filtered = _filtered(latest) if filtered: return { "mode": "latest", "items": [ {k: v for k, v in r.items() if k != "vector"} for r in filtered[: min(top_k, FALLBACK_LATEST_N)] ], } return {"mode": "empty", "items": []} def rebuild(self, items: list[tuple[dict[str, Any], list[float]]]) -> int: """整库重建：items = [(meta_row, vector), ...]。 临时文件写好后 os.replace 原子切换到 index.json；faiss 同步重建。 返回行数。 """ with self._lock: self.dir.mkdir(parents=True, exist_ok=True) new_rows: list[dict[str, Any]] = [] new_faiss = faiss.IndexFlatIP(self.dim) for row_id, (item, vector) in enumerate(items): row = dict(item) row["row_id"] = row_id row["vector"] = list(vector) new_rows.append(row) new_faiss.add(_normalize(vector).reshape(1, self.dim)) # 原子切换：写临时文件 → os.replace self._rows = new_rows self._faiss = new_faiss self._dump_locked(self.json_path) return len(new_rows)