"""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 .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 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) # ---- 增删改查 ---- def add( self, file_bytes: bytes, filename: str, meta: dict[str, Any], as_template: bool = False, ) -> dict[str, Any]: """入库一条。原文只在函数内流转，末尾丢弃引用。""" 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("上传文件为空") 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), "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 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), "embed_sha": embed_sha, "embed_chunks": embed_chunks, } 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。""" removed = self.index.remove_item(doc_id) enc_path = self.vault_dir / f"{doc_id}.enc" if enc_path.exists(): secure_delete(enc_path) 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) 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} 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]