jerryandClaude Fable 5
5db3ead73e
- vstore 独立安装包(8094):解析/嵌入/加密 vault/faiss 检索,存储路径支持本地 file 与 http(s):// URL,启动握手校验,集群化预留(embed 独立模块) - BFF docgen 路由:上传入库/generate/编辑/导出/下载/主库管理,意图护栏(与前端同词典),vstore 失败兜底,quota 507 与 TTL 惰性清理 - 前端:自然触发意图路由、拖拽上传 chips、能力徽标、右侧文档 TAB(WYSIWYG 编辑/prompt 修改/PDF+DOCX 导出下载)、右栏拖宽+手柄仿左栏 - 批量导入工具 docgen_import.py(并发 8、导入即加密+删源明文) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
267 lines
10 KiB
Python
267 lines
10 KiB
Python
"""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)
|