789 lines
35 KiB
Python
789 lines
35 KiB
Python
|
"""RAG 向量索引（collection=知识库），显式选择 Chroma 或 Qdrant。"""
|
||
|
|
import hashlib
|
||
|
|
import logging
|
||
|
|
import re
|
||
|
|
import shutil
|
||
|
|
import sqlite3
|
||
|
|
import json
|
||
|
|
import math
|
||
|
|
import time
|
||
|
|
import uuid
|
||
|
|
import threading
|
||
|
|
from ..config import CHROMA_DIR, DATA_DIR
|
||
|
|
from . import backend
|
||
|
|
|
||
|
|
logger = logging.getLogger("yawen.rag.store")
|
||
|
|
|
||
|
|
|
||
|
|
class VectorIndexDisabled(RuntimeError):
|
||
|
|
"""Raised when an operation requires a vector backend but it is disabled."""
|
||
|
|
|
||
|
|
|
||
|
|
_client = None
|
||
|
|
_client_init_lock = threading.Lock()
|
||
|
|
FALLBACK_DIR = DATA_DIR / "vector_fallback"
|
||
|
|
FALLBACK_DIR.mkdir(parents=True, exist_ok=True)
|
||
|
|
_count_cache: dict[int, tuple[str, dict[str, int]]] = {}
|
||
|
|
_search_metric_local = threading.local()
|
||
|
|
COUNT_CACHE_DIR = DATA_DIR / "vector_counts"
|
||
|
|
COUNT_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||
|
|
|
||
|
|
def _count_cache_path(kb_id: int):
|
||
|
|
return COUNT_CACHE_DIR / f"kb_{int(kb_id)}.json"
|
||
|
|
|
||
|
|
def kb_revision(kb_id: int) -> str:
|
||
|
|
"""Cross-process version of one library, namespaced to this data directory."""
|
||
|
|
marker = COUNT_CACHE_DIR / f"kb_{int(kb_id)}.revision"
|
||
|
|
parts = [backend.namespace(), str(CHROMA_DIR.resolve()), str(FALLBACK_DIR.resolve()), str(marker.resolve())]
|
||
|
|
try:
|
||
|
|
parts.append(marker.read_text(encoding="ascii"))
|
||
|
|
except FileNotFoundError:
|
||
|
|
parts.append("initial")
|
||
|
|
if not backend.is_qdrant():
|
||
|
|
fallback = _fallback_path(kb_id)
|
||
|
|
try:
|
||
|
|
stat = fallback.stat()
|
||
|
|
parts.append(f"{stat.st_ino}:{stat.st_size}:{stat.st_mtime_ns}:{stat.st_ctime_ns}")
|
||
|
|
except FileNotFoundError:
|
||
|
|
pass
|
||
|
|
return hashlib.sha256("\0".join(parts).encode()).hexdigest()
|
||
|
|
|
||
|
|
|
||
|
|
def _bump_revision(kb_id: int) -> None:
|
||
|
|
COUNT_CACHE_DIR.mkdir(parents=True, exist_ok=True)
|
||
|
|
marker = COUNT_CACHE_DIR / f"kb_{int(kb_id)}.revision"
|
||
|
|
temporary = marker.with_suffix(f".{uuid.uuid4().hex}.tmp")
|
||
|
|
try:
|
||
|
|
temporary.write_text(uuid.uuid4().hex, encoding="ascii")
|
||
|
|
temporary.replace(marker)
|
||
|
|
finally:
|
||
|
|
temporary.unlink(missing_ok=True)
|
||
|
|
|
||
|
|
|
||
|
|
def _invalidate_count_cache(kb_id: int) -> None:
|
||
|
|
_bump_revision(kb_id)
|
||
|
|
_count_cache.pop(int(kb_id), None)
|
||
|
|
try:
|
||
|
|
_count_cache_path(kb_id).unlink(missing_ok=True)
|
||
|
|
except OSError:
|
||
|
|
pass
|
||
|
|
|
||
|
|
|
||
|
|
def _fallback_path(kb_id: int):
|
||
|
|
return FALLBACK_DIR / f"kb_{int(kb_id)}.json"
|
||
|
|
|
||
|
|
|
||
|
|
def _read_fallback(kb_id: int) -> dict | None:
|
||
|
|
if backend.is_qdrant():
|
||
|
|
return None
|
||
|
|
path = _fallback_path(kb_id)
|
||
|
|
if not path.exists():
|
||
|
|
return None
|
||
|
|
try:
|
||
|
|
payload = json.loads(path.read_text(encoding="utf-8"))
|
||
|
|
if isinstance(payload, dict) and isinstance(payload.get("items"), list):
|
||
|
|
return payload
|
||
|
|
except (OSError, ValueError, TypeError):
|
||
|
|
return None
|
||
|
|
return None
|
||
|
|
|
||
|
|
|
||
|
|
def _write_fallback(kb_id: int, items: list[dict]) -> None:
|
||
|
|
path = _fallback_path(kb_id)
|
||
|
|
temp = path.with_suffix(".tmp")
|
||
|
|
payload = {"version": 1, "kb_id": int(kb_id), "items": items}
|
||
|
|
temp.write_text(json.dumps(payload, ensure_ascii=False, separators=(",", ":")), encoding="utf-8")
|
||
|
|
temp.replace(path)
|
||
|
|
_bump_revision(kb_id)
|
||
|
|
|
||
|
|
|
||
|
|
class _StreamingJSON:
|
||
|
|
"""只保留当前字段的 JSON 流；向量字段逐字符跳过，不构造 embedding 数组。"""
|
||
|
|
def __init__(self, handle):
|
||
|
|
self.handle, self.buffer, self.pos, self.eof = handle, "", 0, False
|
||
|
|
|
||
|
|
def peek(self):
|
||
|
|
if self.pos >= len(self.buffer) and not self.eof:
|
||
|
|
self.buffer, self.pos = self.handle.read(65536), 0
|
||
|
|
self.eof = not self.buffer
|
||
|
|
return self.buffer[self.pos:self.pos + 1]
|
||
|
|
|
||
|
|
def take(self):
|
||
|
|
value = self.peek()
|
||
|
|
self.pos += bool(value)
|
||
|
|
return value
|
||
|
|
|
||
|
|
def whitespace(self):
|
||
|
|
while self.peek() and self.peek().isspace():
|
||
|
|
self.take()
|
||
|
|
|
||
|
|
def expect(self, value):
|
||
|
|
self.whitespace()
|
||
|
|
if self.take() != value:
|
||
|
|
raise ValueError("索引 JSON 结构损坏")
|
||
|
|
|
||
|
|
def value(self, retain=True):
|
||
|
|
self.whitespace()
|
||
|
|
first = self.peek()
|
||
|
|
if not first:
|
||
|
|
raise ValueError("索引 JSON 意外结束")
|
||
|
|
output, depth, quoted, escaped = [], 0, False, False
|
||
|
|
while self.peek():
|
||
|
|
char = self.peek()
|
||
|
|
if not quoted and depth == 0 and output and (char in ",]}" or char.isspace()):
|
||
|
|
break
|
||
|
|
# 不保留的值也需要一个已开始标记，避免把后续逗号吞掉。
|
||
|
|
if not quoted and depth == 0 and not retain and output and (char in ",]}" or char.isspace()):
|
||
|
|
break
|
||
|
|
char = self.take()
|
||
|
|
if retain:
|
||
|
|
output.append(char)
|
||
|
|
elif not output:
|
||
|
|
output.append(True)
|
||
|
|
if quoted:
|
||
|
|
if escaped:
|
||
|
|
escaped = False
|
||
|
|
elif char == "\\":
|
||
|
|
escaped = True
|
||
|
|
elif char == '"':
|
||
|
|
quoted = False
|
||
|
|
elif char == '"':
|
||
|
|
quoted = True
|
||
|
|
elif char in "[{":
|
||
|
|
depth += 1
|
||
|
|
elif char in "]}":
|
||
|
|
depth -= 1
|
||
|
|
if depth == 0 and not quoted and first in '[{"':
|
||
|
|
break
|
||
|
|
if quoted or depth:
|
||
|
|
raise ValueError("索引 JSON 字段不完整")
|
||
|
|
return json.loads("".join(output)) if retain else None
|
||
|
|
|
||
|
|
|
||
|
|
def _stream_fallback(kb_id: int):
|
||
|
|
path = _fallback_path(kb_id)
|
||
|
|
before = path.stat()
|
||
|
|
with path.open(encoding="utf-8") as handle:
|
||
|
|
stream, found = _StreamingJSON(handle), False
|
||
|
|
stream.expect("{")
|
||
|
|
while True:
|
||
|
|
stream.whitespace()
|
||
|
|
if stream.peek() == "}":
|
||
|
|
stream.take()
|
||
|
|
break
|
||
|
|
key = stream.value()
|
||
|
|
stream.expect(":")
|
||
|
|
if key == "items":
|
||
|
|
if found:
|
||
|
|
raise ValueError("索引包含重复 items 字段")
|
||
|
|
found = True
|
||
|
|
stream.expect("[")
|
||
|
|
while True:
|
||
|
|
stream.whitespace()
|
||
|
|
if stream.peek() == "]":
|
||
|
|
stream.take()
|
||
|
|
break
|
||
|
|
stream.expect("{")
|
||
|
|
item = {}
|
||
|
|
while True:
|
||
|
|
stream.whitespace()
|
||
|
|
if stream.peek() == "}":
|
||
|
|
stream.take()
|
||
|
|
break
|
||
|
|
field = stream.value()
|
||
|
|
stream.expect(":")
|
||
|
|
if field in {"id", "text", "meta"}:
|
||
|
|
item[field] = stream.value()
|
||
|
|
else:
|
||
|
|
stream.value(retain=False)
|
||
|
|
stream.whitespace()
|
||
|
|
if stream.peek() == "}":
|
||
|
|
continue
|
||
|
|
stream.expect(",")
|
||
|
|
yield item
|
||
|
|
stream.whitespace()
|
||
|
|
if stream.peek() == "]":
|
||
|
|
continue
|
||
|
|
stream.expect(",")
|
||
|
|
elif key == "kb_id":
|
||
|
|
if stream.value() != kb_id:
|
||
|
|
raise ValueError("索引知识库编号不匹配")
|
||
|
|
else:
|
||
|
|
stream.value(retain=False)
|
||
|
|
stream.whitespace()
|
||
|
|
if stream.peek() == "}":
|
||
|
|
continue
|
||
|
|
stream.expect(",")
|
||
|
|
stream.whitespace()
|
||
|
|
if stream.peek() or not found:
|
||
|
|
raise ValueError("索引 JSON 缺少正文或包含多余数据")
|
||
|
|
after = path.stat()
|
||
|
|
if (before.st_ino, before.st_size, before.st_mtime_ns) != (after.st_ino, after.st_size, after.st_mtime_ns):
|
||
|
|
raise ValueError("读取过程中知识库索引发生变化，请在索引完成后重新评估")
|
||
|
|
|
||
|
|
|
||
|
|
def iter_chunks(kb_id: int, page_size: int = 256, file_names: set[str] | None = None):
|
||
|
|
"""顺序遍历完整索引正文；不做相似度筛选，不加载向量，不改变消费记录。"""
|
||
|
|
if backend.is_disabled():
|
||
|
|
return
|
||
|
|
if isinstance(kb_id, bool) or not isinstance(kb_id, int) or kb_id <= 0:
|
||
|
|
raise ValueError("知识库编号无效")
|
||
|
|
if not isinstance(page_size, int) or not 1 <= page_size <= 4000:
|
||
|
|
raise ValueError("索引分页大小无效")
|
||
|
|
if file_names is not None and not file_names:
|
||
|
|
return
|
||
|
|
if backend.is_qdrant():
|
||
|
|
revision = kb_revision(kb_id)
|
||
|
|
collection = get_client().get_collection(collection_name(kb_id))
|
||
|
|
total = collection.count()
|
||
|
|
seen = set()
|
||
|
|
groups = [None] if file_names is None else [sorted(file_names)[i:i + 128] for i in range(0, len(file_names), 128)]
|
||
|
|
for names in groups:
|
||
|
|
where = None if names is None else {"file": {"$in": names}}
|
||
|
|
for page in collection.iter_records(where=where, page_size=page_size):
|
||
|
|
for index, identifier in enumerate(page['ids']):
|
||
|
|
meta = page['metadatas'][index]
|
||
|
|
if identifier in seen:
|
||
|
|
raise RuntimeError("知识索引分页返回重复片段，请重新评估")
|
||
|
|
seen.add(identifier)
|
||
|
|
if meta.get('kb_id') not in (None, kb_id):
|
||
|
|
raise RuntimeError("知识索引返回了其他知识库的片段")
|
||
|
|
if file_names is not None and meta.get('file') not in file_names:
|
||
|
|
raise RuntimeError("知识索引返回了分类范围外的文件")
|
||
|
|
yield {"id": identifier, "text": page['documents'][index], "meta": meta}
|
||
|
|
if collection.count() != total or kb_revision(kb_id) != revision or (file_names is None and len(seen) != total):
|
||
|
|
raise RuntimeError("读取过程中知识索引数量变化，请重新评估")
|
||
|
|
return
|
||
|
|
if _fallback_path(kb_id).exists():
|
||
|
|
try:
|
||
|
|
for item in _stream_fallback(kb_id):
|
||
|
|
if file_names is None or (item.get('meta') or {}).get('file') in file_names:
|
||
|
|
yield item
|
||
|
|
except (OSError, ValueError, TypeError) as exc:
|
||
|
|
raise RuntimeError(f"知识库 {kb_id} 本地索引读取失败：{exc}") from exc
|
||
|
|
return
|
||
|
|
try:
|
||
|
|
collection = get_client().get_collection(collection_name(kb_id))
|
||
|
|
total, offset = collection.count(), 0
|
||
|
|
if file_names is not None:
|
||
|
|
# Filter in Chroma before paging; only selected-category documents are read.
|
||
|
|
names = sorted(file_names)
|
||
|
|
for start in range(0, len(names), 128):
|
||
|
|
offset = 0
|
||
|
|
while True:
|
||
|
|
page = collection.get(include=["documents", "metadatas"], limit=page_size, offset=offset,
|
||
|
|
where={"file": {"$in": names[start:start + 128]}})
|
||
|
|
ids = _as_list(page.get("ids"))
|
||
|
|
documents, metadata = _as_list(page.get("documents")), _as_list(page.get("metadatas"))
|
||
|
|
if len(documents) != len(ids) or len(metadata) != len(ids):
|
||
|
|
raise ValueError("分类索引正文或来源缺失")
|
||
|
|
for index, identifier in enumerate(ids):
|
||
|
|
meta = metadata[index] or {}
|
||
|
|
if meta.get('file') not in file_names:
|
||
|
|
raise ValueError("分类索引返回了范围外的文件")
|
||
|
|
yield {"id": identifier, "text": documents[index], "meta": meta}
|
||
|
|
offset += len(ids)
|
||
|
|
if len(ids) < page_size:
|
||
|
|
break
|
||
|
|
if collection.count() != total:
|
||
|
|
raise ValueError("读取过程中索引数量变化，请重新评估")
|
||
|
|
return
|
||
|
|
while offset < total:
|
||
|
|
page = collection.get(include=["documents", "metadatas"], limit=page_size, offset=offset)
|
||
|
|
ids = _as_list(page.get("ids"))
|
||
|
|
documents, metadata = _as_list(page.get("documents")), _as_list(page.get("metadatas"))
|
||
|
|
if not ids or len(documents) != len(ids) or len(metadata) != len(ids):
|
||
|
|
raise ValueError("索引分页正文或来源缺失")
|
||
|
|
for index, identifier in enumerate(ids):
|
||
|
|
yield {"id": identifier, "text": documents[index], "meta": metadata[index] or {}}
|
||
|
|
offset += len(ids)
|
||
|
|
if offset != total or collection.count() != total:
|
||
|
|
raise ValueError("读取过程中索引数量变化，请重新评估")
|
||
|
|
except Exception as exc:
|
||
|
|
raise RuntimeError(f"知识库 {kb_id} 索引读取失败：{exc}") from exc
|
||
|
|
|
||
|
|
def _as_list(value, default: list | None = None) -> list:
|
||
|
|
"""把 Chroma/NumPy 返回的序列安全转换为普通列表，避免 array 的布尔判断异常。"""
|
||
|
|
if value is None:
|
||
|
|
return list(default or [])
|
||
|
|
if isinstance(value, list):
|
||
|
|
return value
|
||
|
|
try:
|
||
|
|
return list(value)
|
||
|
|
except TypeError:
|
||
|
|
return list(default or [])
|
||
|
|
|
||
|
|
def get_client():
|
||
|
|
if backend.is_disabled():
|
||
|
|
configured = backend.configured_backend()
|
||
|
|
if configured == "chroma":
|
||
|
|
raise VectorIndexDisabled("本地向量索引未启用：未安装 chromadb，请使用轻量模式或安装可选 RAG 依赖")
|
||
|
|
raise VectorIndexDisabled("本地向量索引已禁用：设置 RAG_VECTOR_BACKEND=chroma/qdrant 后才可访问向量库")
|
||
|
|
if backend.is_qdrant():
|
||
|
|
from .qdrant_store import get_client as remote_client
|
||
|
|
return remote_client()
|
||
|
|
global _client
|
||
|
|
if _client is None:
|
||
|
|
with _client_init_lock:
|
||
|
|
if _client is None:
|
||
|
|
try:
|
||
|
|
import chromadb
|
||
|
|
except ImportError as exc:
|
||
|
|
raise VectorIndexDisabled("缺少 chromadb，本地向量索引不可用") from exc
|
||
|
|
_client = chromadb.PersistentClient(path=str(CHROMA_DIR))
|
||
|
|
return _client
|
||
|
|
|
||
|
|
def collection_name(kb_id: int) -> str:
|
||
|
|
return f"kb_{kb_id}"
|
||
|
|
|
||
|
|
def ensure_collection(kb_id: int):
|
||
|
|
started = time.perf_counter()
|
||
|
|
collection = get_client().get_or_create_collection(collection_name(kb_id))
|
||
|
|
metric = getattr(_search_metric_local, "metric", None)
|
||
|
|
if metric is not None:
|
||
|
|
metric["collection_get_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||
|
|
return collection
|
||
|
|
|
||
|
|
|
||
|
|
def last_search_metric() -> dict:
|
||
|
|
"""Return timings for the most recent vector search on this thread.
|
||
|
|
|
||
|
|
Search requests fan out across threads, so a thread-local envelope avoids
|
||
|
|
cross-library races while keeping the public ``search`` signature intact.
|
||
|
|
"""
|
||
|
|
return dict(getattr(_search_metric_local, "metric", {}) or {})
|
||
|
|
|
||
|
|
def _chunk_id(file_name: str, item: dict, occurrence: int) -> str:
|
||
|
|
"""用来源和正文生成稳定 ID，重建整个库也不会改变同一片段的消费记录。"""
|
||
|
|
meta = item.get("meta", {})
|
||
|
|
heading = str(meta.get("heading") or "")
|
||
|
|
text = " ".join(str(item.get("text") or "").split())
|
||
|
|
raw = f"{file_name}\0{heading}\0{text}\0{occurrence}".encode("utf-8")
|
||
|
|
return "chunk-" + hashlib.sha256(raw).hexdigest()
|
||
|
|
|
||
|
|
def _file_rows(collection, file_name: str) -> list[dict]:
|
||
|
|
if hasattr(collection, "iter_records"):
|
||
|
|
total = collection.count(where={"file": file_name})
|
||
|
|
rows = []
|
||
|
|
for data in collection.iter_records(where={"file": file_name}, include=["documents", "metadatas", "embeddings"]):
|
||
|
|
rows.extend({"id": identifier, "text": data["documents"][i], "embedding": data["embeddings"][i], "meta": data["metadatas"][i]}
|
||
|
|
for i, identifier in enumerate(data["ids"]))
|
||
|
|
if len(rows) != total or collection.count(where={"file": file_name}) != total:
|
||
|
|
raise RuntimeError("文件索引快照不完整，已取消修改")
|
||
|
|
return rows
|
||
|
|
rows, offset = [], 0
|
||
|
|
while True:
|
||
|
|
data = collection.get(where={"file": file_name}, include=["documents", "embeddings", "metadatas"],
|
||
|
|
limit=4000, offset=offset)
|
||
|
|
ids = _as_list(data.get("ids"))
|
||
|
|
texts, vectors, metas = (_as_list(data.get(key)) for key in ("documents", "embeddings", "metadatas"))
|
||
|
|
if any(len(values) != len(ids) for values in (texts, vectors, metas)):
|
||
|
|
raise RuntimeError("文件索引快照不完整，已取消修改")
|
||
|
|
rows.extend({"id": ident, "text": texts[i], "embedding": _as_list(vectors[i]), "meta": metas[i] or {}}
|
||
|
|
for i, ident in enumerate(ids))
|
||
|
|
offset += len(ids)
|
||
|
|
if len(ids) < 4000:
|
||
|
|
return rows
|
||
|
|
|
||
|
|
|
||
|
|
def snapshot_file(kb_id: int, file_name: str) -> dict:
|
||
|
|
"""Snapshot the file in each existing index; never seed a partial fallback."""
|
||
|
|
if backend.is_disabled():
|
||
|
|
return {"kb_id": kb_id, "file_name": file_name, "fallback_exists": False,
|
||
|
|
"backend_namespace": backend.namespace(), "fallback_items": [], "chroma_items": []}
|
||
|
|
exists = not backend.is_qdrant() and _fallback_path(kb_id).exists()
|
||
|
|
fallback = _read_fallback(kb_id) if exists else None
|
||
|
|
if exists and fallback is None:
|
||
|
|
raise RuntimeError("本地向量索引损坏，已取消修改")
|
||
|
|
return {"kb_id": kb_id, "file_name": file_name, "fallback_exists": exists,
|
||
|
|
"backend_namespace": backend.namespace(),
|
||
|
|
"fallback_items": [row for row in (fallback or {}).get("items", [])
|
||
|
|
if (row.get("meta") or {}).get("file") == file_name],
|
||
|
|
"chroma_items": _file_rows(ensure_collection(kb_id), file_name)}
|
||
|
|
|
||
|
|
|
||
|
|
def _insert_rows(collection, rows: list[dict]):
|
||
|
|
for start in range(0, len(rows), 4000):
|
||
|
|
batch = rows[start:start + 4000]
|
||
|
|
collection.add(ids=[row["id"] for row in batch], documents=[row["text"] for row in batch],
|
||
|
|
embeddings=[row["embedding"] for row in batch], metadatas=[row["meta"] for row in batch])
|
||
|
|
|
||
|
|
|
||
|
|
def restore_file(kb_id: int, file_name: str, snapshot: dict) -> None:
|
||
|
|
if snapshot.get("kb_id") != kb_id or snapshot.get("file_name") != file_name:
|
||
|
|
raise ValueError("文件索引快照归属不匹配")
|
||
|
|
if snapshot.get("backend_namespace", backend.namespace()) != backend.namespace():
|
||
|
|
raise ValueError("文件索引快照所属后端已改变，已取消恢复")
|
||
|
|
collection = ensure_collection(kb_id)
|
||
|
|
try:
|
||
|
|
collection.delete(where={"file": file_name})
|
||
|
|
_insert_rows(collection, snapshot["chroma_items"])
|
||
|
|
if snapshot["fallback_exists"]:
|
||
|
|
fallback = _read_fallback(kb_id)
|
||
|
|
if fallback is None:
|
||
|
|
raise RuntimeError("回滚时本地索引不可读取")
|
||
|
|
others = [row for row in fallback["items"] if (row.get("meta") or {}).get("file") != file_name]
|
||
|
|
_write_fallback(kb_id, others + snapshot["fallback_items"])
|
||
|
|
finally:
|
||
|
|
_invalidate_count_cache(kb_id)
|
||
|
|
|
||
|
|
|
||
|
|
def _prepared_rows(kb_id, chunks, embeddings, file_name, category_path):
|
||
|
|
if len(chunks) != len(embeddings):
|
||
|
|
raise ValueError("分段数量与向量数量不一致，已取消写入")
|
||
|
|
rows, occurrences = [], {}
|
||
|
|
for i, item in enumerate(chunks):
|
||
|
|
if not isinstance(item.get("text"), str) or not item["text"].strip():
|
||
|
|
raise ValueError("索引片段正文为空，已取消写入")
|
||
|
|
vector = _as_list(embeddings[i])
|
||
|
|
if not vector or any(not math.isfinite(float(x)) for x in vector):
|
||
|
|
raise ValueError("索引向量无效，已取消写入")
|
||
|
|
if backend.is_qdrant() and len(vector) != 512:
|
||
|
|
raise ValueError("Qdrant 索引向量必须是 512 维，已取消写入")
|
||
|
|
fingerprint = f"{item.get('meta', {}).get('heading', '')}\0{' '.join(item['text'].split())}"
|
||
|
|
occurrences[fingerprint] = occurrences.get(fingerprint, 0) + 1
|
||
|
|
ident = _chunk_id(file_name, item, occurrences[fingerprint])
|
||
|
|
meta = dict(item.get("meta") or {})
|
||
|
|
meta.setdefault("seq", i)
|
||
|
|
from .embed import EMBEDDING_POLICY
|
||
|
|
meta.update(kb_id=kb_id, file=file_name, chunk_hash=ident.removeprefix("chunk-"), embedding_policy=EMBEDDING_POLICY)
|
||
|
|
if category_path:
|
||
|
|
meta["category_path"] = category_path
|
||
|
|
rows.append({"id": ident, "text": item["text"], "embedding": vector, "meta": meta})
|
||
|
|
return rows
|
||
|
|
|
||
|
|
|
||
|
|
def replace_file_chunks(kb_id: int, chunks: list[dict], embeddings: list[list[float]],
|
||
|
|
file_name: str, category_path: str = "") -> None:
|
||
|
|
if backend.is_disabled():
|
||
|
|
raise VectorIndexDisabled("本地向量索引已禁用，不能写入向量片段")
|
||
|
|
"""Prepare before mutation; failed replacements restore that file in both stores.
|
||
|
|
|
||
|
|
Chroma-only libraries stay Chroma-only. A fallback, when already present,
|
||
|
|
continues to contain every other file rather than shadowing the whole library
|
||
|
|
with only this upload. Callers serialize writers with the per-KB import lock.
|
||
|
|
"""
|
||
|
|
rows = _prepared_rows(kb_id, chunks, embeddings, file_name, category_path)
|
||
|
|
previous = snapshot_file(kb_id, file_name)
|
||
|
|
collection = ensure_collection(kb_id)
|
||
|
|
_invalidate_count_cache(kb_id)
|
||
|
|
try:
|
||
|
|
collection.delete(where={"file": file_name})
|
||
|
|
_insert_rows(collection, rows)
|
||
|
|
if previous["fallback_exists"]:
|
||
|
|
fallback = _read_fallback(kb_id)
|
||
|
|
if fallback is None:
|
||
|
|
raise RuntimeError("本地向量索引无法读取")
|
||
|
|
others = [row for row in fallback["items"] if (row.get("meta") or {}).get("file") != file_name]
|
||
|
|
_write_fallback(kb_id, others + rows)
|
||
|
|
_invalidate_count_cache(kb_id)
|
||
|
|
except Exception as exc:
|
||
|
|
try:
|
||
|
|
restore_file(kb_id, file_name, previous)
|
||
|
|
except Exception as recovery:
|
||
|
|
raise RuntimeError(f"文件索引更新失败，恢复也失败，请保留源文件并检查索引：{recovery}") from exc
|
||
|
|
raise RuntimeError(f"文件索引更新失败，已恢复原索引：{exc}") from exc
|
||
|
|
|
||
|
|
|
||
|
|
def add_chunks(kb_id: int, chunks: list[dict], embeddings: list[list[float]],
|
||
|
|
file_name: str, category_path: str = ""):
|
||
|
|
return replace_file_chunks(kb_id, chunks, embeddings, file_name, category_path)
|
||
|
|
|
||
|
|
|
||
|
|
def delete_file(kb_id: int, file_name: str):
|
||
|
|
if backend.is_disabled():
|
||
|
|
_invalidate_count_cache(kb_id)
|
||
|
|
return
|
||
|
|
return replace_file_chunks(kb_id, [], [], file_name)
|
||
|
|
|
||
|
|
def count_by_file(kb_id: int) -> dict:
|
||
|
|
"""按当前后端及版本缓存文件片段数；Qdrant 故障不读取旧 Chroma 缓存。"""
|
||
|
|
if backend.is_disabled():
|
||
|
|
return {}
|
||
|
|
remote = backend.is_qdrant()
|
||
|
|
revision = kb_revision(kb_id)
|
||
|
|
cache_path = _count_cache_path(kb_id)
|
||
|
|
try:
|
||
|
|
# An exact remote count also verifies service availability before cached
|
||
|
|
# statistics are returned. The endpoint is cheap and does not load text.
|
||
|
|
coll = ensure_collection(kb_id) if remote else None
|
||
|
|
total = coll.count() if remote else None
|
||
|
|
cached = _count_cache.get(int(kb_id))
|
||
|
|
if cached and cached[0] == revision and (total is None or sum(cached[1].values()) == total):
|
||
|
|
return dict(cached[1])
|
||
|
|
try:
|
||
|
|
payload = json.loads(cache_path.read_text(encoding="utf-8"))
|
||
|
|
counts = payload.get("counts") if isinstance(payload, dict) else None
|
||
|
|
if payload.get("version") == 2 and payload.get("kb_id") == int(kb_id) and payload.get("revision") == revision and isinstance(counts, dict):
|
||
|
|
if any(not isinstance(k, str) or isinstance(v, bool) or not isinstance(v, int) or v < 0 for k, v in counts.items()):
|
||
|
|
raise ValueError("知识索引统计缓存无效")
|
||
|
|
if total is None or sum(counts.values()) == total:
|
||
|
|
_count_cache[int(kb_id)] = (revision, counts)
|
||
|
|
return dict(counts)
|
||
|
|
except (OSError, ValueError, TypeError, AttributeError):
|
||
|
|
pass
|
||
|
|
out: dict[str, int] = {}
|
||
|
|
if not remote and _fallback_path(kb_id).exists():
|
||
|
|
records = (item.get("meta") or {} for item in _stream_fallback(kb_id))
|
||
|
|
for meta in records:
|
||
|
|
name = meta.get("file") or "未知"
|
||
|
|
out[name] = out.get(name, 0) + 1
|
||
|
|
else:
|
||
|
|
coll = coll if coll is not None else ensure_collection(kb_id)
|
||
|
|
if coll is None:
|
||
|
|
return {}
|
||
|
|
if hasattr(coll, "iter_records"):
|
||
|
|
pages = coll.iter_records(include=["metadatas"])
|
||
|
|
for data in pages:
|
||
|
|
for meta in data["metadatas"]:
|
||
|
|
name = meta.get("file") or "未知"
|
||
|
|
out[name] = out.get(name, 0) + 1
|
||
|
|
else:
|
||
|
|
offset, page = 0, 5000
|
||
|
|
while True:
|
||
|
|
got = coll.get(include=["metadatas"], limit=page, offset=offset)
|
||
|
|
metas = got.get("metadatas") or []
|
||
|
|
for meta in metas:
|
||
|
|
name = (meta or {}).get("file") or "未知"
|
||
|
|
out[name] = out.get(name, 0) + 1
|
||
|
|
offset += len(metas)
|
||
|
|
if len(metas) < page:
|
||
|
|
break
|
||
|
|
if kb_revision(kb_id) != revision or (remote and (sum(out.values()) != total or coll.count() != total)):
|
||
|
|
raise RuntimeError("统计过程中知识索引发生变化，请重新加载")
|
||
|
|
_count_cache[int(kb_id)] = (revision, out)
|
||
|
|
temporary = cache_path.with_suffix(f".{uuid.uuid4().hex}.tmp")
|
||
|
|
try:
|
||
|
|
cache_path.parent.mkdir(parents=True, exist_ok=True)
|
||
|
|
temporary.write_text(json.dumps({"version": 2, "kb_id": int(kb_id), "revision": revision, "counts": out},
|
||
|
|
ensure_ascii=False, separators=(",", ":")), encoding="utf-8")
|
||
|
|
temporary.replace(cache_path)
|
||
|
|
except OSError:
|
||
|
|
pass
|
||
|
|
finally:
|
||
|
|
temporary.unlink(missing_ok=True)
|
||
|
|
return dict(out)
|
||
|
|
except Exception:
|
||
|
|
if remote:
|
||
|
|
raise
|
||
|
|
return {}
|
||
|
|
|
||
|
|
|
||
|
|
def count(kb_id: int) -> int:
|
||
|
|
if backend.is_disabled():
|
||
|
|
return 0
|
||
|
|
fallback = _read_fallback(kb_id)
|
||
|
|
if fallback is not None:
|
||
|
|
return len(fallback.get("items", []))
|
||
|
|
try:
|
||
|
|
return ensure_collection(kb_id).count()
|
||
|
|
except Exception:
|
||
|
|
if backend.is_qdrant():
|
||
|
|
raise
|
||
|
|
return 0
|
||
|
|
|
||
|
|
def search(kb_id: int, query_embedding: list[float], k: int = 6, file_names: set[str] | None = None) -> list[dict]:
|
||
|
|
if backend.is_disabled():
|
||
|
|
return []
|
||
|
|
metric = {"kb_id": int(kb_id), "metadata_filter_ms": 0.0,
|
||
|
|
"vector_ms": 0.0, "collection_get_ms": 0.0}
|
||
|
|
_search_metric_local.metric = metric
|
||
|
|
started = time.perf_counter()
|
||
|
|
if file_names is not None and not file_names:
|
||
|
|
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||
|
|
return []
|
||
|
|
fallback = _read_fallback(kb_id)
|
||
|
|
if fallback is not None:
|
||
|
|
filter_started = time.perf_counter()
|
||
|
|
if file_names is not None:
|
||
|
|
filtered = [item for item in fallback.get("items", [])
|
||
|
|
if (item.get("meta") or {}).get("file") in file_names]
|
||
|
|
else:
|
||
|
|
filtered = list(fallback.get("items", []))
|
||
|
|
metric["metadata_filter_ms"] = round((time.perf_counter() - filter_started) * 1000, 2)
|
||
|
|
vector_started = time.perf_counter()
|
||
|
|
def cosine_distance(vector: list[float]) -> float:
|
||
|
|
dot = sum(float(a) * float(b) for a, b in zip(query_embedding, vector))
|
||
|
|
left = math.sqrt(sum(float(a) * float(a) for a in query_embedding)) or 1.0
|
||
|
|
right = math.sqrt(sum(float(b) * float(b) for b in vector)) or 1.0
|
||
|
|
return 1.0 - dot / (left * right)
|
||
|
|
ranked = sorted(filtered, key=lambda item: cosine_distance(item.get("embedding", [])))[:k]
|
||
|
|
metric["vector_ms"] = round((time.perf_counter() - vector_started) * 1000, 2)
|
||
|
|
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||
|
|
return [{"id": item.get("id"), "text": item.get("text", ""),
|
||
|
|
"meta": item.get("meta", {}), "distance": cosine_distance(item.get("embedding", []))}
|
||
|
|
for item in ranked]
|
||
|
|
coll = ensure_collection(kb_id)
|
||
|
|
# Chroma and Qdrant both return an empty result for an empty collection.
|
||
|
|
# Do not call ``count()`` on the hot retrieval path: on an external disk
|
||
|
|
# it can scan collection metadata (over 1.7s cold for the 108k-chunk KB),
|
||
|
|
# while the bounded nearest-neighbor query itself is only a few ms.
|
||
|
|
# Keeping ``available`` unknown also removes a redundant remote count for
|
||
|
|
# Qdrant; the backend safely truncates to the number of stored points.
|
||
|
|
available = None
|
||
|
|
if file_names is not None:
|
||
|
|
candidates = []
|
||
|
|
names = sorted(file_names)
|
||
|
|
for start in range(0, len(names), 128):
|
||
|
|
vector_started = time.perf_counter()
|
||
|
|
result = coll.query(query_embeddings=[query_embedding], n_results=k,
|
||
|
|
where={"file": {"$in": names[start:start + 128]}})
|
||
|
|
metric["vector_ms"] += (time.perf_counter() - vector_started) * 1000
|
||
|
|
filter_started = time.perf_counter()
|
||
|
|
for i, identifier in enumerate((result.get('ids') or [[]])[0]):
|
||
|
|
meta = result['metadatas'][0][i] or {}
|
||
|
|
if meta.get('file') in file_names:
|
||
|
|
candidates.append({'id': identifier, 'text': result['documents'][0][i], 'meta': meta,
|
||
|
|
'distance': result['distances'][0][i]})
|
||
|
|
metric["metadata_filter_ms"] += (time.perf_counter() - filter_started) * 1000
|
||
|
|
metric["metadata_filter_ms"] = round(metric["metadata_filter_ms"], 2)
|
||
|
|
metric["vector_ms"] = round(metric["vector_ms"], 2)
|
||
|
|
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||
|
|
return sorted(candidates, key=lambda item: item['distance'])[:k]
|
||
|
|
vector_started = time.perf_counter()
|
||
|
|
res = coll.query(query_embeddings=[query_embedding], n_results=k)
|
||
|
|
metric["vector_ms"] = round((time.perf_counter() - vector_started) * 1000, 2)
|
||
|
|
out = []
|
||
|
|
if res and res.get("ids") and res["ids"][0]:
|
||
|
|
for i, doc_id in enumerate(res["ids"][0]):
|
||
|
|
out.append({
|
||
|
|
"id": doc_id,
|
||
|
|
"text": res["documents"][0][i],
|
||
|
|
"meta": res["metadatas"][0][i] if res.get("metadatas") else {},
|
||
|
|
"distance": res["distances"][0][i] if res.get("distances") else None,
|
||
|
|
})
|
||
|
|
metric["total_ms"] = round((time.perf_counter() - started) * 1000, 2)
|
||
|
|
return out
|
||
|
|
|
||
|
|
def _vector_segment_dirs(kb_id: int) -> list:
|
||
|
|
"""在删除 collection 前记录其 HNSW 段目录；只返回当前 collection 的 UUID 目录。"""
|
||
|
|
db_path = CHROMA_DIR / "chroma.sqlite3"
|
||
|
|
if not db_path.exists():
|
||
|
|
return []
|
||
|
|
with sqlite3.connect(str(db_path)) as conn:
|
||
|
|
rows = conn.execute(
|
||
|
|
"""SELECT s.id FROM collections c JOIN segments s ON s.collection = c.id
|
||
|
|
WHERE c.name = ? AND s.scope = 'VECTOR'""", (collection_name(kb_id),)
|
||
|
|
).fetchall()
|
||
|
|
root = CHROMA_DIR.resolve()
|
||
|
|
dirs = []
|
||
|
|
for (segment_id,) in rows:
|
||
|
|
if not re.fullmatch(r"[0-9a-f-]{36}", str(segment_id)):
|
||
|
|
continue
|
||
|
|
candidate = (root / str(segment_id)).resolve()
|
||
|
|
if candidate.parent == root:
|
||
|
|
dirs.append(candidate)
|
||
|
|
return dirs
|
||
|
|
|
||
|
|
|
||
|
|
def reset_kb(kb_id: int):
|
||
|
|
if backend.is_disabled():
|
||
|
|
_fallback_path(kb_id).unlink(missing_ok=True)
|
||
|
|
_invalidate_count_cache(kb_id)
|
||
|
|
return
|
||
|
|
if backend.is_qdrant():
|
||
|
|
_invalidate_count_cache(kb_id)
|
||
|
|
get_client().delete_collection(collection_name(kb_id))
|
||
|
|
_invalidate_count_cache(kb_id)
|
||
|
|
return
|
||
|
|
fallback_path = _fallback_path(kb_id)
|
||
|
|
if fallback_path.exists():
|
||
|
|
fallback_path.unlink()
|
||
|
|
_invalidate_count_cache(kb_id)
|
||
|
|
segment_dirs = _vector_segment_dirs(kb_id)
|
||
|
|
try:
|
||
|
|
get_client().delete_collection(collection_name(kb_id))
|
||
|
|
except Exception:
|
||
|
|
# collection 不存在时无需清理；其余失败交给调用方，避免误删仍被引用的段目录。
|
||
|
|
if segment_dirs:
|
||
|
|
raise
|
||
|
|
return
|
||
|
|
_invalidate_count_cache(kb_id)
|
||
|
|
for directory in segment_dirs:
|
||
|
|
if directory.exists():
|
||
|
|
shutil.rmtree(directory)
|
||
|
|
|
||
|
|
|
||
|
|
def snapshot_kb(kb_id: int) -> dict:
|
||
|
|
"""重建前导出当前 collection，供分块/嵌入异常时原样恢复。"""
|
||
|
|
if backend.is_disabled():
|
||
|
|
return {"ids": [], "documents": [], "embeddings": [], "metadatas": []}
|
||
|
|
fallback = _read_fallback(kb_id)
|
||
|
|
if fallback is not None:
|
||
|
|
return {"ids": [item.get("id") for item in fallback["items"]],
|
||
|
|
"documents": [item.get("text", "") for item in fallback["items"]],
|
||
|
|
"embeddings": [item.get("embedding", []) for item in fallback["items"]],
|
||
|
|
"metadatas": [item.get("meta", {}) for item in fallback["items"]]}
|
||
|
|
coll = ensure_collection(kb_id)
|
||
|
|
if coll.count() <= 0:
|
||
|
|
return {"ids": [], "documents": [], "embeddings": [], "metadatas": []}
|
||
|
|
if hasattr(coll, "iter_records"):
|
||
|
|
revision, total = kb_revision(kb_id), coll.count()
|
||
|
|
result = {"ids": [], "documents": [], "embeddings": [], "metadatas": []}
|
||
|
|
for page in coll.iter_records(include=["documents", "metadatas", "embeddings"]):
|
||
|
|
for key in result:
|
||
|
|
result[key].extend(page[key])
|
||
|
|
if len(result["ids"]) != total or coll.count() != total or kb_revision(kb_id) != revision:
|
||
|
|
raise RuntimeError("知识库索引快照不完整，已取消重建")
|
||
|
|
return result
|
||
|
|
# 大库分页导出（单次全量 get 会触发 sqlite 变量上限）
|
||
|
|
ids: list = []
|
||
|
|
documents: list = []
|
||
|
|
embeddings: list = []
|
||
|
|
metadatas: list = []
|
||
|
|
page = 4000
|
||
|
|
offset = 0
|
||
|
|
while True:
|
||
|
|
data = coll.get(include=["documents", "embeddings", "metadatas"], limit=page, offset=offset)
|
||
|
|
batch_ids = _as_list(data.get("ids"))
|
||
|
|
if not batch_ids:
|
||
|
|
break
|
||
|
|
ids.extend(batch_ids)
|
||
|
|
documents.extend(_as_list(data.get("documents")))
|
||
|
|
embeddings.extend(_as_list(data.get("embeddings")))
|
||
|
|
metadatas.extend(_as_list(data.get("metadatas")))
|
||
|
|
offset += len(batch_ids)
|
||
|
|
if len(batch_ids) < page:
|
||
|
|
break
|
||
|
|
return {
|
||
|
|
"ids": ids,
|
||
|
|
"documents": documents,
|
||
|
|
"embeddings": embeddings,
|
||
|
|
"metadatas": metadatas,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def restore_kb(kb_id: int, snapshot: dict) -> None:
|
||
|
|
"""仅用于重建失败回滚；恢复原 chunk id 和元数据，消费记录仍可对应。"""
|
||
|
|
ids = _as_list(snapshot.get("ids"))
|
||
|
|
if backend.is_disabled():
|
||
|
|
if ids:
|
||
|
|
raise VectorIndexDisabled("本地向量索引已禁用，无法恢复向量快照")
|
||
|
|
return
|
||
|
|
documents = _as_list(snapshot.get("documents"))
|
||
|
|
embeddings = _as_list(snapshot.get("embeddings"))
|
||
|
|
metadatas = _as_list(snapshot.get("metadatas"))
|
||
|
|
if any(len(values) != len(ids) for values in (documents, embeddings, metadatas)) or len(set(ids)) != len(ids):
|
||
|
|
raise ValueError("知识索引快照不完整，已取消恢复")
|
||
|
|
remote = backend.is_qdrant()
|
||
|
|
for identifier, text, vector, meta in zip(ids, documents, embeddings, metadatas):
|
||
|
|
if not isinstance(identifier, str) or not identifier or not isinstance(text, str) or not text.strip() or not isinstance(meta, dict):
|
||
|
|
raise ValueError("知识索引快照正文或来源无效，已取消恢复")
|
||
|
|
if not len(vector) or (remote and len(vector) != 512) or any(not math.isfinite(float(v)) for v in vector):
|
||
|
|
raise ValueError("知识索引快照向量无效，已取消恢复")
|
||
|
|
reset_kb(kb_id)
|
||
|
|
if not ids:
|
||
|
|
return
|
||
|
|
items = [{"id": ids[index], "text": documents[index],
|
||
|
|
"embedding": embeddings[index], "meta": metadatas[index]}
|
||
|
|
for index in range(len(ids))]
|
||
|
|
if not remote:
|
||
|
|
_write_fallback(kb_id, items)
|
||
|
|
ensure_collection(kb_id).add(
|
||
|
|
ids=ids,
|
||
|
|
documents=documents,
|
||
|
|
embeddings=embeddings,
|
||
|
|
metadatas=metadatas,
|
||
|
|
)
|
||
|
|
_invalidate_count_cache(kb_id)
|