jerryandClaude Fable 5
4749f77544
第一批：上传去限制+进度展示、模板HTML预览+按需导出、检索TAB实时更新
第二批：edit护栏移除、导出按钮纯化+file_name落库、编辑失焦还原、
二次预览白底根治（防竞态+渲染兜底）
模板体系：7套公文模板(templates/)+templating渲染器、websearch联网
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
414 lines
16 KiB
Python
414 lines
16 KiB
Python
"""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]
|