forked from admin/PlantAI-rag-new
adminandClaude Fable 5
60f7360cec
缺陷 #52b(验收 §6 第 2 条 + 回归 G3):同一 PDF 重复上传 4 次 → 4 个 doc_id × 200 chunk = 800 条,FAISS 检索 top_k 被同内容不同 doc_id 的块占满。 修复(vstore/vault.py): - add() 先算内容哈希 SHA256(原始文件字节),查 index.json 的 content_sha 字段判重;同内容 → 移除旧 doc 的检索行/密文/chunk 行,再落一条新 doc (新 doc_id,meta 以本次上传为准,created_at 刷新) - 去重后该内容在库中恒仅 1 个 doc_id、chunk 总数不翻倍;同名不同内容 按内容哈希不误判(G3-2);不同内容互不覆盖(G3-1) - content_sha 落 index.json(doc_meta 新字段);rebuild_index 对历史密文 缺 content_sha 时以密文内容补算(重建后去重依据不丢) - 与现有删除机制一致:index.remove_item 过滤重建 faiss、chunks.remove_document 过滤重建 chunk faiss(IndexFlatIP 不支持逐行删,现有机制即重建) 测试:vstore/tests/test_vault_dedup.py(5 用例):同内容两次 → 1 个 doc_id/ chunk 不翻倍/旧密文清理;meta 以本次为准;不同内容 → 2 个 doc_id; 同名不同内容不误判;重启(新实例 load)后仍去重。 Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
605 lines
24 KiB
Python
605 lines
24 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 .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]
|