"""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)