"""vault.py —— 单一主库：上传 → 解析 → 提取 → 嵌入 → 加密落盘 → 索引。 明文纪律： - add() 原文只在函数内流转，函数末尾丢弃引用，绝不落盘明文 - 磁盘上只有 .enc 密文（Fernet） - 检索 meta 存 index.json（name/category/tags/description，不含正文） 后端形态： - 本地：Vault（本机磁盘，vault/ 密文 + index/ 检索索引） - 远程：RemoteVault（操作转 HTTP 指向远程 vstore 根，接口语义与本地一致） """ from __future__ import annotations import hashlib import json import logging import os import uuid from pathlib import Path from typing import Any, Callable import httpx from .crypto import CryptoError, dump_encrypted, load_encrypted, secure_delete from .embed import embed_batch from .parsing import ParsingError, extract, outline_to_embedding_chunks from .retrieval import RetrievalIndex, RetrievalError from .chunks import ChunkIndex, build_chunks from .settings import ( EMBED_BODY_CAP, EMBED_DIM, MAX_OUTLINE_CHUNKS, get_paths, ) logger = logging.getLogger(__name__) # 嵌入调用函数（测试可注入假 embed） _EMBED_FN: Callable[..., list[list[float]]] = embed_batch def _chunked_embed(texts: list[str]) -> list[list[float]]: """chunk 批量嵌入（经模块属性动态查找——测试 monkeypatch vstore.embed 的 embed_batch_chunked 后此处即时生效，不受静态 import 快照影响）。""" from . import embed as _embed_mod return _embed_mod.embed_batch_chunked(texts) class VaultError(Exception): """主库操作失败（限额/解析/加密/检索）。""" def set_embed_fn(fn: Callable[..., list[list[float]]]) -> None: """注入嵌入函数（测试用）。""" global _EMBED_FN _EMBED_FN = fn def get_embed_fn() -> Callable[..., list[list[float]]]: return _EMBED_FN def _now_iso() -> str: from datetime import datetime, timezone return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") def _serialize_doc(extracted: dict[str, Any], meta: dict[str, Any]) -> bytes: """序列化落盘结构（JSON，明文只在内存）。""" payload = { "meta": meta, "extracted": { "title": extracted.get("title") or "", "outline": extracted.get("outline") or [], "paragraphs": extracted.get("paragraphs") or [], "tables": extracted.get("tables") or [], "style_hints": extracted.get("style_hints") or {}, }, } return json.dumps(payload, ensure_ascii=False).encode("utf-8") def _deserialize_doc(raw: bytes) -> dict[str, Any]: try: payload = json.loads(raw.decode("utf-8")) if not isinstance(payload, dict) or "extracted" not in payload: raise VaultError("密文内容结构异常") return payload except (ValueError, UnicodeDecodeError) as exc: raise VaultError(f"密文内容解析失败：{exc}") from exc def _meta_to_keywords(meta: dict[str, Any]) -> list[str]: """meta → 关键词列表（检索二档用）：name/description/tags 切词。""" words: list[str] = [] for key in ("name", "title", "description", "tags"): value = str(meta.get(key) or "") if value: words.extend([w for w in value.replace(",", " ").replace("，", " ").split() if len(w) >= 2]) return words def _make_doc_json(path: str, title: str, outline: list[dict[str, Any]], paragraphs: list[str], tables: list[list[list[str]]], style_hints: dict[str, Any]) -> dict[str, Any]: """构造 /documents/{id} 返回体（明文全文，仅生成时取）。""" return { "title": title, "outline": outline, "paragraphs": paragraphs, "tables": tables, "style_hints": style_hints, "enc_path": path, } class Vault: """本地主库。add：限额检查 → 解析 → 嵌入 → 加密落盘 → 索引。""" def __init__(self, vault_path: str | Path, index_path: str | Path) -> None: self.vault_dir = Path(vault_path) self.index = RetrievalIndex(index_path) # 私有资料 chunk 级索引（FR-PRIV-002）：与 doc 级检索索引并存； # 路径来自 settings（VSTORE_CHUNK_PATH，默认 {root}/chunks），测试可用 set_paths 覆盖 chunk_path = get_paths().chunk_path self.chunks = ChunkIndex(chunk_path) self.chunks.load() # ---- 增删改查 ---- def _add_file( self, doc_id: str, file_bytes: bytes, filename: str, name: str, category: str, tags: str, description: str, as_template: bool, ) -> dict[str, Any]: """入库一条（覆盖语义已由调用方保证：同内容 doc 已被移除）。 与旧 add() 唯一差异是 doc_id 由调用方传入；函数内失败照常抛 VaultError， 由调用方决定失败回滚语义。 """ try: extracted = extract(file_bytes, filename) except ParsingError as exc: raise VaultError(str(exc)) from exc paragraphs = extracted.get("paragraphs") or [] outline = extracted.get("outline") or [] doc_meta = { "id": doc_id, "name": name, "category": category, "tags": tags, "description": description, "as_template": bool(as_template), "size": len(file_bytes), "content_sha": hashlib.sha256(bytes(file_bytes)).hexdigest(), "created_at": _now_iso(), "embed_chunks": 0, } # 嵌入块：模版用标题+大纲；普通资料用标题+描述+正文前 20 段 chunks = [extracted.get("title") or name, description] if as_template: chunks = chunks[:1] + outline_to_embedding_chunks(outline)[:MAX_OUTLINE_CHUNKS] else: body = "\n".join(paragraphs[:20]) if body: chunks.append(body[:EMBED_BODY_CAP]) chunks = [c for c in chunks if c and c.strip()] vectors = _EMBED_FN(chunks) enc_path = self.vault_dir / f"{doc_id}.enc" try: dump_encrypted(enc_path, _serialize_doc(extracted, doc_meta)) except (CryptoError, OSError) as exc: raise VaultError(f"密文落盘失败：{exc}") from exc try: doc_meta["enc_path"] = str(enc_path) doc_meta["embed_chunks"] = len(chunks) # keywords 只用用户 meta，不含正文派生的 title（明文纪律） doc_meta["keywords"] = _meta_to_keywords(doc_meta) self.index.add_item(doc_meta, vectors[0]) except (RetrievalError, ValueError, IndexError) as exc: # 索引失败：清理密文，不让库处于半一致状态 secure_delete(enc_path) raise VaultError(f"索引写入失败：{exc}") from exc # ---- 私有资料 chunk 级索引（FR-PRIV-002）：doc 级成功后追加；任何失败只告警不阻断 ---- chunk_count = 0 chunk_ok = True chunk_error = "" try: para_list = [str(p or "") for p in extracted.get("paragraphs") or []] table_list = extracted.get("tables") or [] para_chars = sum(len(p) for p in para_list) table_cells = sum(len(row) for table in table_list for row in table) if para_chars or table_cells: built = build_chunks( extracted.get("title") or name, extracted.get("outline") or [], para_list, table_list, ) else: built = [] if built: texts = [c.text for c in built] vecs = _chunked_embed(texts) items = [] for c, vec in zip(built, vecs): items.append( { "seq": c.seq, "text": c.text, "vector": list(vec), "doc_name": name, "doc_category": category, "doc_tags": tags, "as_template": bool(as_template), "created_at": doc_meta["created_at"], } ) self.chunks.upsert_document(doc_id, items) chunk_count = len(items) except Exception as exc: # chunk 任何失败只告警，doc 级照常返回 chunk_ok = False chunk_error = str(exc) or exc.__class__.__name__ logger.warning("vault: chunk 索引写入失败（不阻断 doc 级，%s）：%s", doc_id, exc) embed_chunks = doc_meta["embed_chunks"] embed_sha = hashlib.sha256( json.dumps(vectors[0], ensure_ascii=False).encode("utf-8") ).hexdigest() # 明文纪律：原文字节与解析结果在本函数末尾丢弃引用 del file_bytes, extracted, paragraphs, outline, chunks, vectors return { "document_id": doc_id, "name": name, "category": category, "tags": tags, "description": description, "as_template": bool(as_template), "enc_path": str(enc_path), "content_sha": doc_meta["content_sha"], "embed_sha": embed_sha, "embed_chunks": embed_chunks, # FR-PRIV-002：chunk 级索引结果（失败不阻断 doc 级） "chunk_count": chunk_count, "chunk_ok": chunk_ok, "chunk_error": chunk_error, } def _find_content_duplicate(self, content_sha: str) -> dict[str, Any] | None: """库中同内容哈希的 doc 行（无则 None）。 内容哈希存于 index.json 每行的 content_sha 字段（#52b 新增）；历史行 无该字段（字段缺失/不一致时绝不判重）。直接读 index.json 而非内存 self.index._rows——覆盖语义要求与磁盘一致（同内容密文先删后建）。 """ try: data = json.loads(self.index.json_path.read_text(encoding="utf-8")) items = data.get("items", []) if isinstance(data, dict) else [] except (ValueError, OSError): return None for item in items: if isinstance(item, dict) and str(item.get("content_sha") or "") == content_sha: return item return None def add( self, file_bytes: bytes, filename: str, meta: dict[str, Any], as_template: bool = False, ) -> dict[str, Any]: """入库一条（内容 SHA256 去重，缺陷 #52b）。原文只在函数内流转，末尾丢弃引用。 去重语义（验收 §6 第 2 条 + 回归 G3）： - 内容哈希 = SHA256(原文件字节)。同内容重复上传 → 移除旧 doc 的密文/ 检索行/chunk 行 → 落一条新 doc；该内容在库中始终仅 1 个 doc_id， chunk 总数不翻倍。 - 旧行清理先于新 doc 落盘；新 doc 建库失败时旧 doc 已删（覆盖语义， 失败由调用方重传即可，绝无同内容双份）。 - 新 doc 的 meta 以本次上传为准（name/category/tags/description 覆盖， created_at 刷新）；as_template 以本次为准。 """ doc_id = str(uuid.uuid4()) name = str(meta.get("name") or filename or "未命名") category = str(meta.get("category") or "") tags = str(meta.get("tags") or "") description = str(meta.get("description") or "") if not file_bytes: raise VaultError("上传文件为空") # ---- 内容去重（#52b）：同内容哈希 → 先清旧 doc 再落新 doc ---- content_sha = hashlib.sha256(bytes(file_bytes)).hexdigest() dup = self._find_content_duplicate(content_sha) if dup: old_doc_id = str(dup.get("id") or "") if old_doc_id and old_doc_id != doc_id: try: self.index.remove_item(old_doc_id) except Exception as exc: logger.warning("vault: 去重移除旧检索行失败（%s）：%s", old_doc_id, exc) old_enc = self.vault_dir / f"{old_doc_id}.enc" if old_enc.exists(): try: secure_delete(old_enc) except OSError as exc: logger.warning("vault: 去重移除旧密文失败（%s）：%s", old_enc, exc) try: self.chunks.remove_document(old_doc_id) except Exception as exc: logger.warning("vault: 去重移除旧 chunk 行失败（%s）：%s", old_doc_id, exc) return self._add_file( doc_id, file_bytes, filename, name, category, tags, description, as_template ) def get_decrypted(self, doc_id: str) -> dict[str, Any] | None: """取解密后的文档结构（meta + extracted）。不存在返回 None。""" enc_path = self.vault_dir / f"{doc_id}.enc" if not enc_path.exists(): return None try: return _deserialize_doc(load_encrypted(enc_path)) except (CryptoError, VaultError, OSError) as exc: logger.warning("vault: 解密失败 %s：%s", doc_id, exc) return None def get_document_json(self, doc_id: str) -> dict[str, Any] | None: """GET /documents/{id} 响应体（明文全文 + outline + meta）。""" doc = self.get_decrypted(doc_id) if doc is None: return None extracted = doc.get("extracted") or {} return _make_doc_json( path=f"vault/{doc_id}.enc", title=extracted.get("title") or "", outline=extracted.get("outline") or [], paragraphs=extracted.get("paragraphs") or [], tables=extracted.get("tables") or [], style_hints=extracted.get("style_hints") or {}, ) def get_style(self, doc_id: str) -> dict[str, Any] | None: """docx 样式提取（/documents/{id}/style）。""" doc = self.get_decrypted(doc_id) if doc is None: return None return (doc.get("extracted") or {}).get("style_hints") or {} def delete(self, doc_id: str) -> bool: """删除：索引行 + 密文 secure_delete + chunk 行联动（FR-PRIV-006）。""" removed = self.index.remove_item(doc_id) enc_path = self.vault_dir / f"{doc_id}.enc" if enc_path.exists(): secure_delete(enc_path) # chunk 索引联动（失败只告警，不改变 doc 级结果） try: self.chunks.remove_document(doc_id) except Exception as exc: logger.warning("vault: chunk 行删除失败（%s）：%s", doc_id, exc) return removed def list_rows(self) -> list[dict[str, Any]]: """主库行列表（无向量无正文）。""" return self.index.rows() def update_metadata(self, doc_id: str, meta: dict[str, Any]) -> bool: """更新 meta。""" return self.index.update_meta(doc_id, meta) # ---- 检索 ---- def search( self, query: str, top_k: int = 5, category: str | None = None, as_template: bool | None = None, ) -> dict[str, Any]: """三档降级检索（向量 → meta 关键词 → 最新）。query 为空走 latest。""" keywords = _meta_to_keywords({"description": query, "name": query}) if query else None vector = _EMBED_FN([query or ""])[0] if query else [0.0] * EMBED_DIM return self.index.search( vector, top_k=top_k, category=category, as_template=as_template, keywords=keywords ) def rebuild_index(self) -> dict[str, Any]: """整库重建：全部密文重新解密 → 嵌入 → 原子切换。 明文只在内存流转；任何一条失败（密钥不符/嵌入失败）整库拒绝切换。 """ items: list[tuple[dict[str, Any], list[float]]] = [] enc_files = sorted(self.vault_dir.glob("*.enc")) for enc in enc_files: doc_id = enc.stem try: doc = _deserialize_doc(load_encrypted(enc)) except (CryptoError, VaultError, OSError) as exc: return {"status": "error", "message": f"密文 {doc_id} 解密失败：{exc}"} meta = doc.get("meta") or {} extracted = doc.get("extracted") or {} title = extracted.get("title") or str(meta.get("name") or doc_id) chunks = [title, str(meta.get("description") or "")] if meta.get("as_template"): chunks = chunks[:1] + outline_to_embedding_chunks( extracted.get("outline") or [] )[:MAX_OUTLINE_CHUNKS] else: body = "\n".join((extracted.get("paragraphs") or [])[:20]) if body: chunks.append(body[:EMBED_BODY_CAP]) chunks = [c for c in chunks if c and c.strip()] vectors = _EMBED_FN(chunks) if not vectors: return {"status": "error", "message": f"文档 {doc_id} 嵌入失败"} row = dict(meta) # keywords 只用 meta，不含正文派生的 title（明文纪律） row["keywords"] = _meta_to_keywords(meta) # content_sha（#52b）：历史密文 meta 若缺失则以密文内容补算（与入库时一致）， # 保证重建后去重依据不丢 if not row.get("content_sha"): try: row["content_sha"] = hashlib.sha256(load_encrypted(enc)).hexdigest() except (CryptoError, VaultError, OSError): row["content_sha"] = "" items.append((row, vectors[0])) try: count = self.index.rebuild(items) except (RetrievalError, ValueError) as exc: return {"status": "error", "message": f"索引重建失败：{exc}"} return {"status": "ok", "count": count} def rebuild_chunks(self) -> dict[str, Any]: """整库重建 chunk 索引（FR-PRIV-007）：遍历密文解密 → build_chunks → 批量嵌入 → rebuild_all 覆盖三文件。返回 {docs, chunks, duration_ms}。 """ from time import perf_counter started = perf_counter() docs: dict[str, list[dict[str, Any]]] = {} enc_files = sorted(self.vault_dir.glob("*.enc")) for enc in enc_files: doc_id = enc.stem try: doc = _deserialize_doc(load_encrypted(enc)) except (CryptoError, VaultError, OSError) as exc: logger.warning("vault: 重建 chunk 解密失败（跳过，%s）：%s", doc_id, exc) continue meta = doc.get("meta") or {} extracted = doc.get("extracted") or {} name = str(meta.get("name") or doc_id) built = build_chunks( extracted.get("title") or name, extracted.get("outline") or [], [str(p or "") for p in extracted.get("paragraphs") or []], extracted.get("tables") or [], ) if not built: continue vecs = _chunked_embed([c.text for c in built]) docs[doc_id] = [ { "seq": c.seq, "text": c.text, "vector": list(vec), "doc_name": name, "doc_category": str(meta.get("category") or ""), "doc_tags": str(meta.get("tags") or ""), "as_template": bool(meta.get("as_template")), "created_at": str(meta.get("created_at") or ""), } for c, vec in zip(built, vecs) ] count = self.chunks.rebuild_all(docs) return { "docs": len(docs), "chunks": count, "duration_ms": int((perf_counter() - started) * 1000), } class RemoteVault: """远程主库：所有操作转 HTTP 指向远程 vstore 根（接口语义与本地一致）。 vault/index 操作统一打到 index_url 的 /api/vstore/*（远程端自身 再按它的 vault/index 配置落盘）。同步 httpx（远程形态 V1 低频使用）。 """ def __init__(self, vault_url: str, index_url: str) -> None: self.vault_url = vault_url.rstrip("/") self.index_url = index_url.rstrip("/") def _request(self, method: str, path: str, **kwargs: Any) -> httpx.Response: resp = httpx.request(method, f"{self.index_url}{path}", timeout=60.0, **kwargs) resp.raise_for_status() return resp def add(self, file_bytes: bytes, filename: str, meta: dict[str, Any], as_template: bool = False) -> dict[str, Any]: resp = self._request( "POST", "/api/vstore/ingest", data={k: v for k, v in meta.items() if k in ("name", "category", "description", "tags")}, files={"file": (filename, file_bytes)}, params={"as_template": bool(as_template)}, ) payload = resp.json() return payload.get("item") or {} def search(self, query: str, top_k: int = 5, category: str | None = None, as_template: bool | None = None) -> dict[str, Any]: params: dict[str, Any] = {"q": query, "top_k": top_k} if category: params["category"] = category if as_template is not None: params["as_template"] = as_template resp = self._request("GET", "/api/vstore/search", params=params) payload = resp.json() return {"mode": payload.get("mode", "remote"), "items": payload.get("items", [])} def get_document_json(self, doc_id: str) -> dict[str, Any] | None: try: resp = self._request("GET", f"/api/vstore/documents/{doc_id}") except httpx.HTTPStatusError as exc: if exc.response.status_code == 404: return None raise return resp.json().get("document") or None def get_style(self, doc_id: str) -> dict[str, Any] | None: try: resp = self._request("GET", f"/api/vstore/documents/{doc_id}/style") except httpx.HTTPStatusError as exc: if exc.response.status_code == 404: return None raise return resp.json().get("style_hints") or None def delete(self, doc_id: str) -> bool: resp = self._request("DELETE", f"/api/vstore/documents/{doc_id}") return bool(resp.json().get("removed")) def update_metadata(self, doc_id: str, meta: dict[str, Any]) -> bool: # 远程端无 PATCH 端点；V1 远程形态以 BFF 落库为准（meta 更新容错忽略） logger.warning("vstore: 远程形态 update_metadata 暂不支持（%s）", doc_id) return False def rebuild_index(self) -> dict[str, Any]: resp = self._request("POST", "/api/vstore/rebuild-index") payload = resp.json() return {"status": "ok" if payload.get("ok") else "error", "count": payload.get("count", 0)} def list_rows(self) -> list[dict[str, Any]]: resp = self._request("GET", "/api/vstore/search", params={"q": "", "top_k": 50}) return resp.json().get("items", []) # ---- 全局单例管理（api.py 启动钩子经 set_backend/init_vault 注入）---- _g_vault: Vault | RemoteVault | None = None _g_backend: Any = None def set_backend(backend: Any) -> None: """注入后端（api.py 启动钩子调用）。""" global _g_backend _g_backend = backend def init_vault() -> None: """按后端构造全局主库（本地/远程）。""" global _g_vault backend = _g_backend if backend is None: raise RuntimeError("vstore: 后端未初始化（先 set_backend + init_vault）") if getattr(backend, "kind", "") == "remote": _g_vault = RemoteVault(backend.vault_url, backend.index_url) else: vault = Vault(str(backend.vault_path), str(backend.index_path)) vault.index.load() _g_vault = vault def get_vault() -> Vault | RemoteVault: """全局主库单例。""" if _g_vault is None: init_vault() return _g_vault # type: ignore[return-value]