diff --git a/vstore/tests/test_vault_dedup.py b/vstore/tests/test_vault_dedup.py new file mode 100644 index 0000000..8526516 --- /dev/null +++ b/vstore/tests/test_vault_dedup.py @@ -0,0 +1,123 @@ +"""test_vault_dedup.py —— 内容 SHA256 去重（缺陷 #52b，验收 §6 第 2 条 + 回归 G3）。 + +运行：cd vstore && /home/xzq/5090/teamWorkspace/zh/plantai-rag-new-venv/bin/python -m pytest tests/test_vault_dedup.py -v +（或 vstore/pyproject.toml testpaths 下全量 pytest；venv = plantai-rag-new-venv） + +G3 断言： +- 同内容 add 两次 → 库中该内容仅 1 个 doc_id、chunk 总数不翻倍； +- 不同内容 add → 2 个 doc_id 互不覆盖； +- 同名不同内容 → 按内容哈希不误判（保留 2 个 doc_id）。 +""" +from __future__ import annotations + +import hashlib + +import pytest + +from vstore.vault import Vault + + +@pytest.fixture +def vault(tmp_path, embed_fn, monkeypatch): + from vstore.settings import VstorePaths, set_paths + + set_paths(VstorePaths(root=str(tmp_path), vault_path="", index_path="", chunk_path="")) + # 桩掉 chunk 批量嵌入（vault._chunked_embed 动态查找），不触网真 Ollama + dim = 1024 + monkeypatch.setattr( + "vstore.embed.embed_batch_chunked", + lambda texts, *a, **k: [[1.0 / dim] * dim for _ in texts], + ) + return Vault(tmp_path / "vault", tmp_path / "index") + + +def _doc_ids_for_content(vault, content: bytes) -> list[str]: + """库中 content_sha == sha256(content) 的 doc 行 id 列表（磁盘事实）。""" + sha = hashlib.sha256(content).hexdigest() + import json + + rows = [] + try: + data = json.loads(vault.index.json_path.read_text(encoding="utf-8")) + items = data.get("items", []) if isinstance(data, dict) else [] + except (ValueError, OSError): + items = [] + for item in items: + if isinstance(item, dict) and str(item.get("content_sha") or "") == sha: + rows.append(item) + return [r.get("id") for r in rows] + + +def _make_pdf(text: str) -> bytes: + """真实最小 PDF 字节（pymupdf 可解析、含目标文本）。""" + import pymupdf + + doc = pymupdf.open() + page = doc.new_page() + page.insert_text((72, 72), text) + return doc.tobytes() + + +def test_same_content_added_twice_keeps_one_doc_id(vault): + """#52b：同内容 add 两次 → 1 个 doc_id、chunk 数不翻倍、旧密文被清。""" + content = _make_pdf("同一份论文内容。") + first = vault.add(content, "paper.pdf", {"name": "论文", "category": "资料"}) + assert first["content_sha"] == hashlib.sha256(content).hexdigest() + chunks_after_first = vault.chunks.count() + assert chunks_after_first >= 1 + + second = vault.add(content, "paper.pdf", {"name": "论文", "category": "资料"}) + + ids = _doc_ids_for_content(vault, content) + assert ids == [second["document_id"]], f"库中该内容 doc_id 应为仅新 doc：{ids}" + assert second["document_id"] != first["document_id"] + # 旧密文被清理（覆盖语义不留两份） + assert not (vault.vault_dir / f"{first['document_id']}.enc").exists() + assert (vault.vault_dir / f"{second['document_id']}.enc").exists() + # chunk 总数不翻倍（旧 doc_id 的 chunk 已清理） + assert vault.chunks.count() == chunks_after_first + # chunk 行全部属于新 doc_id（旧 doc_id 不再出现在检索结果） + assert {r["doc_id"] for r in vault.chunks.rows()} == {second["document_id"]} + + +def test_duplicate_reuses_latest_meta(vault): + """#52b：覆盖语义——重复上传 meta 以本次为准（created_at 刷新、name 覆盖）。""" + content = "内容甲。".encode("utf-8") + first = vault.add(content, "a.txt", {"name": "旧名"}) + second = vault.add(content, "a.txt", {"name": "新名"}) + rows = vault.list_rows() + assert len(rows) == 1 + assert rows[0]["name"] == "新名" + assert rows[0]["id"] == second["document_id"] + assert first["document_id"] not in [r["id"] for r in rows] + + +def test_different_content_keeps_two_doc_ids(vault): + """G3-1：不同内容上传 → 2 个 doc_id 互不覆盖。""" + a = vault.add("内容甲。".encode("utf-8"), "a.txt", {"name": "甲"}) + b = vault.add("内容乙。".encode("utf-8"), "b.txt", {"name": "乙"}) + assert a["document_id"] != b["document_id"] + assert len(vault.list_rows()) == 2 + assert set(_doc_ids_for_content(vault, "内容甲。".encode("utf-8"))) == {a["document_id"]} + assert set(_doc_ids_for_content(vault, "内容乙。".encode("utf-8"))) == {b["document_id"]} + + +def test_same_name_different_content_not_deduped(vault): + """G3-2：同名不同内容 → 按 SHA256 不误判（保留 2 个 doc_id）。""" + a = vault.add(_make_pdf("内容甲。"), "same.pdf", {"name": "同名"}) + b = vault.add(_make_pdf("内容乙。"), "same.pdf", {"name": "同名"}) + assert a["document_id"] != b["document_id"] + assert len(vault.list_rows()) == 2 + assert vault.chunks.count() >= 2 + + +def test_duplicate_dedup_persists_across_reload(vault, tmp_path, monkeypatch): + """#52b：去重依据（content_sha）落盘持久——重启后（新实例 load）仍不判重。""" + content = "持久化内容。".encode("utf-8") + vault.add(content, "a.txt", {"name": "甲"}) + vault2 = Vault(tmp_path / "vault", tmp_path / "index") + vault2.index.load() + vault2.chunks.load() + item = vault2.add(content, "a.txt", {"name": "甲"}) + assert _doc_ids_for_content(vault2, content) == [item["document_id"]] + assert vault2.chunks.count() >= 1 diff --git a/vstore/vault.py b/vstore/vault.py index 2a60a42..d9d32a7 100644 --- a/vstore/vault.py +++ b/vstore/vault.py @@ -130,22 +130,22 @@ class Vault: # ---- 增删改查 ---- - def add( + def _add_file( self, + doc_id: str, file_bytes: bytes, filename: str, - meta: dict[str, Any], - as_template: bool = False, + name: str, + category: str, + tags: str, + description: str, + as_template: bool, ) -> 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 "") + """入库一条（覆盖语义已由调用方保证：同内容 doc 已被移除）。 - if not file_bytes: - raise VaultError("上传文件为空") + 与旧 add() 唯一差异是 doc_id 由调用方传入；函数内失败照常抛 VaultError， + 由调用方决定失败回滚语义。 + """ try: extracted = extract(file_bytes, filename) except ParsingError as exc: @@ -162,6 +162,7 @@ class Vault: "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, } @@ -252,6 +253,7 @@ class Vault: "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 级） @@ -260,6 +262,75 @@ class Vault: "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" @@ -362,6 +433,13 @@ class Vault: 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)