"""test_chunks.py —— 私有资料 chunk 级索引单元测试（T1 骨架，FR-PRIV-002）。 对应用例：TC-PRIV-L1-01 ~ TC-PRIV-L1-17、TC-PRIV-L1-07b（docs/team/TEST_CASES_PRIVATE.md）。 TL 裁决已并入：L1-05 表格原位聚合、L1-07b outline 不参与切块。 纪律： - 只依赖 vstore 已声明契约（docs/team/brief_priv_a1.md P1 签名）；A1 交付后 本文件应零修改直接可跑（若契约有出入按任务书/本用例表为准要求 A1 对齐）。 - 复用 vstore/tests/conftest.py 的假 embed 注入模式；embed_batch_chunked 的 测试注入钩子按「仿 set_embed_fn 模式」的约定名 `_CHUNKED_REQUEST_FN` 使用—— 若 A1 实际钩子名不同，仅需改本文件头部 _REQUEST_HOOK 一行（并在报告注明偏离）。 - 本文件不执行真实 ollama、不触网、不写产品代码。 运行：cd vstore && ../..venv/bin/pytest tests/test_chunks.py -v （venv = /home/xzq/5090/teamWorkspace/zh/plantai-rag-new-venv/bin/python） """ from __future__ import annotations import numpy as np import pytest from vstore.chunks import ( CHUNK_MAX_PARAS, CHUNK_MAX_PER_DOC, CHUNK_MIN_CHARS, CHUNK_MIN_SCORE, CHUNK_OVERLAP_PARAS, CHUNK_TARGET_CHARS, ChunkIndex, build_chunks, ) from vstore.settings import EMBED_DIM # 假嵌入注入钩子（契约约定名；A1 若改名仅改此常量并在测试报告注明偏离） _REQUEST_HOOK = "_CHUNKED_REQUEST_FN" DIM = EMBED_DIM # --------------------------------------------------------------------------- # 假嵌入工具（chunk 检索可控分数） # --------------------------------------------------------------------------- def fake_vec(seed: int) -> list[float]: """确定性单位向量：seed → 基向量方向。""" vec = np.zeros(DIM, dtype=np.float32) vec[seed % DIM] = 1.0 return vec.tolist() def fake_embed_controlled(texts: list[str]) -> list[list[float]]: """文本 → 向量：『种子SEED』前缀取 SEED 作方向；其余全零向量。""" out: list[list[float]] = [] for text in texts: seed = 0 if isinstance(text, str) and text.startswith("种子"): try: seed = int(text[2:].split()[0]) except ValueError: seed = 0 out.append(fake_vec(seed) if seed else [0.0] * DIM) return out @pytest.fixture def embed_fn(monkeypatch): """把可控假嵌入注入 vault 模块（与 conftest.embed_fn 同模式）。""" from vstore import vault as vault_mod monkeypatch.setattr(vault_mod, "_EMBED_FN", fake_embed_controlled) return fake_embed_controlled # --------------------------------------------------------------------------- # TC-PRIV-L1-01 ~ L1-07b：build_chunks 切块参数边界（含 TL 裁决版 L1-05/L1-07b） # --------------------------------------------------------------------------- class TestBuildChunksParams: def test_eight_short_paras_make_one_chunk(self): """TC-PRIV-L1-01：8 段短段落 → 恰好 1 块，首行 # 标题，seq==0。""" paras = ["段落内容。"] * 8 out = build_chunks("报告", [], paras, []) assert len(out) == 1 assert out[0].seq == 0 assert out[0].text.startswith("# 报告") for p in paras: assert p in out[0].text def test_char_target_600_fills_chunk(self): """TC-PRIV-L1-02：总长 ≥600 字先封块（不要求满 8 段）。""" paras = ["字" * 90] * 7 # 630 字 out = build_chunks("T", [], paras, []) assert len(out) >= 1 # 首块内容长度 ≥ 目标字符数（不含标题前缀） assert len(out[0].text.replace("# T", "", 1).strip()) >= CHUNK_TARGET_CHARS def test_overlap_one_para(self): """TC-PRIV-L1-03：相邻块重叠 1 段（前块末段 = 后块首段）。""" # 每段 140 字 ×12：首块 8 段封块（≥600 字），尾块 4 段（≥ MIN_CHARS 不并块）， # 2 块产出与 MAX_PARAS/字符目标语义均无歧义 paras = [f"段落{i}的内容。" * 20 for i in range(12)] out = build_chunks("T", [], paras, []) assert len(out) >= 2 # 重叠判定：块 B 的正文首段内容出现在块 A 的正文中（前块最后 1 段） b_lines = out[1].text.strip().splitlines() b_first = b_lines[1] if out[1].text.startswith("#") else b_lines[0] assert b_first in out[0].text # 参数常量自检（契约值） assert CHUNK_OVERLAP_PARAS == 1 def test_max_200_chunks(self): """TC-PRIV-L1-04：超长输入截断为 200 块，seq 0..199，不抛异常。""" paras = ["超长段落内容" * 100] * 300 out = build_chunks("T", [], paras, []) assert len(out) == CHUNK_MAX_PER_DOC assert out[0].seq == 0 assert out[-1].seq == CHUNK_MAX_PER_DOC - 1 def test_table_matrix_inserted_in_place(self): """TC-PRIV-L1-05（TL 裁决版）：表格矩阵化行段插入段落流原位参与聚合。""" pre = ["前置段落甲内容。", "前置段落乙内容。"] tables = [[["a", "b"], ["1", "2"]]] post = ["后续段落丙内容。", "后续段落丁内容。"] out = build_chunks("T", [], pre + post, tables) assert len(out) >= 1 joined = "\n".join(c.text for c in out) assert "|a|b|" in joined assert "|1|2|" in joined # 原位聚合：表格行段位于前置段落之后、后续段落之前 assert joined.find("前置段落甲") < joined.find("|a|b|") < joined.find("后续段落丙") # 段落与表格均完整参与（不丢内容） for p in pre + post: assert p in joined def test_first_chunk_title_prefix(self): """TC-PRIV-L1-06：首块以 `# {title}` 开头。""" out = build_chunks("小麦试验", [], ["正文段落。"], []) assert out[0].text.startswith("# 小麦试验") def test_empty_input_returns_no_chunks(self): """TC-PRIV-L1-07：全空输入 → []，不抛异常（chunk 失败不阻断 doc 级前提）。""" out = build_chunks("", [], [], []) assert out == [] def test_outline_does_not_affect_chunks(self): """TC-PRIV-L1-07b（TL 裁决版）：outline 不参与切块，不影响 chunk 输出。""" paras = ["段落一内容。", "段落二内容。", "段落三内容。"] outline_a = [{"level": 1, "text": "第一章 总则"}, {"level": 2, "text": "1.1 目的"}] out_a = build_chunks("T", outline_a, paras, []) out_b = build_chunks("T", [], paras, []) # 两版输出逐块一致（outline 不参与切块） assert [c.text for c in out_a] == [c.text for c in out_b] # chunk 正文中不含 outline 文本 joined = "\n".join(c.text for c in out_a) assert "第一章 总则" not in joined assert "1.1 目的" not in joined def test_param_constants_contract(self): """参数常量契约自检（A1 任务书取值）。""" assert CHUNK_TARGET_CHARS == 600 assert CHUNK_MIN_CHARS == 120 assert CHUNK_MAX_PARAS == 8 assert CHUNK_OVERLAP_PARAS == 1 assert CHUNK_MAX_PER_DOC == 200 assert CHUNK_MIN_SCORE == 0.30 # --------------------------------------------------------------------------- # TC-PRIV-L1-08 ~ L1-13：ChunkIndex 增删改查与原子写 # --------------------------------------------------------------------------- def make_chunk_rows(prefix: str, n: int, name: str = "资料") -> list[dict]: return [{"text": f"种子{prefix} 第{i}块正文", "doc_name": name} for i in range(n)] class TestChunkIndex: def test_upsert_and_rows(self, tmp_path): """TC-PRIV-L1-08：upsert 后 rows()/count()，chunk_id 为 {doc_id}:{seq}。""" idx = ChunkIndex(tmp_path) idx.upsert_document("A", make_chunk_rows("1", 3)) assert idx.count() == 3 rows = idx.rows() assert {r["chunk_id"] for r in rows} == {"A:0", "A:1", "A:2"} for r in rows: assert r["doc_id"] == "A" assert "seq" in r and "snippet" in r and "doc_name" in r def test_upsert_replace_semantics(self, tmp_path): """TC-PRIV-L1-09：同 doc 二次 upsert 为替换语义（非追加）。""" idx = ChunkIndex(tmp_path) idx.upsert_document("A", make_chunk_rows("1", 3)) idx.upsert_document("A", make_chunk_rows("2", 2)) assert idx.count() == 2 rows = idx.rows() assert {r["chunk_id"] for r in rows} == {"A:0", "A:1"} # 旧行内容彻底消失（替换语义：snippet 全部来自新 upsert 批次） assert all((r.get("snippet") or "").startswith("种子2") for r in rows) assert not any((r.get("snippet") or "").startswith("种子1") for r in rows) def test_remove_document(self, tmp_path): """TC-PRIV-L1-10：remove 只删目标 doc；不存在返回 False。""" idx = ChunkIndex(tmp_path) idx.upsert_document("A", make_chunk_rows("1", 2)) idx.upsert_document("B", make_chunk_rows("2", 3)) assert idx.remove_document("A") is True assert idx.count() == 3 assert all(r["doc_id"] == "B" for r in idx.rows()) assert idx.remove_document("A") is False def test_atomic_write_three_files(self, tmp_path): """TC-PRIV-L1-11：三文件落盘；重载一致；无 .tmp 残留。""" idx = ChunkIndex(tmp_path) idx.upsert_document("A", make_chunk_rows("1", 3)) files = sorted(p.name for p in tmp_path.iterdir() if p.is_file()) assert "index.json" in files assert "texts.jsonl" in files assert "index.faiss" in files # 原子写：无 tmp 残留 assert not [p for p in tmp_path.iterdir() if p.name.endswith(".tmp")] # 重载一致性 idx2 = ChunkIndex(tmp_path) idx2.load() assert idx2.count() == 3 assert {r["chunk_id"] for r in idx2.rows()} == {"A:0", "A:1", "A:2"} def test_load_missing_dir_as_empty(self, tmp_path): """load 对缺失目录按空库降级（可服务，不抛异常）。""" idx = ChunkIndex(tmp_path / "not-exist") idx.load() assert idx.count() == 0 def test_search_threshold_topk_and_filters(self, tmp_path): """TC-PRIV-L1-12：阈值两侧/top_k/doc_ids/exclude_template 组合过滤。""" idx = ChunkIndex(tmp_path) idx.upsert_document("A", make_chunk_rows("1", 4)) # 分数 0.90（种子1 高位） idx.upsert_document("B", make_chunk_rows("1", 2)) # 同种子，不同 doc idx.upsert_document("T", [{"text": "种子1 模板块", "doc_name": "模版", "as_template": True}]) q = fake_vec(1) # 阈值两侧：0.29 全丢；0.30 命中（分数 0.9 的行） assert idx.search(q, top_k=20, min_score=0.29)["mode"] == "vector" assert idx.search(q, top_k=20, min_score=0.29)["items"] # 高阈值（>0.9 不存在）→ empty assert idx.search(q, top_k=20, min_score=0.91)["mode"] == "empty" # top_k clamp：返回条数 ≤ top_k r = idx.search(q, top_k=2, min_score=0.30) assert len(r["items"]) <= 2 # doc_ids 白名单 r = idx.search(q, top_k=20, min_score=0.30, doc_ids=["A"]) assert all(i["doc_id"] == "A" for i in r["items"]) # exclude_template r = idx.search(q, top_k=20, min_score=0.30, exclude_template=True) assert all(not i.get("as_template") for i in r["items"]) def test_search_below_min_score_empty_no_fallback(self, tmp_path): """TC-PRIV-L1-13：低于阈值 → empty 且无降级档。""" idx = ChunkIndex(tmp_path) idx.upsert_document("A", make_chunk_rows("1", 3)) r = idx.search(fake_vec(99), top_k=5, min_score=CHUNK_MIN_SCORE) assert r["mode"] == "empty" assert r["items"] == [] def test_snippet_truncated(self, tmp_path): """snippet 截断契约（200 字符，假设 4）。""" idx = ChunkIndex(tmp_path) idx.upsert_document("A", [{"text": "字" * 500, "doc_name": "n"}]) rows = idx.rows() assert len(rows) == 1 assert len(rows[0]["snippet"]) <= 200 def test_rebuild_all(self, tmp_path): """rebuild_all：覆盖写入并返回行数（TC-PRIV-L2-07 的单元前置）。""" idx = ChunkIndex(tmp_path) idx.upsert_document("A", make_chunk_rows("1", 2)) n = idx.rebuild_all({"A": make_chunk_rows("2", 3)}) assert n == 3 idx2 = ChunkIndex(tmp_path) idx2.load() assert idx2.count() == 3 assert all("种子2" in (r.get("snippet") or "") for r in idx2.rows()) # --------------------------------------------------------------------------- # TC-PRIV-L1-14 ~ L1-17：embed_batch_chunked（input 数组批量/并发/重试/降级） # --------------------------------------------------------------------------- class TestEmbedBatchChunked: def _install_hook(self, monkeypatch, fn): """注入批量请求钩子；返回调用记录器（batches 记录每批 payload）。""" calls = {"batches": [], "peak": 0, "active": 0} def hook(payload): calls["batches"].append(payload) calls["active"] += 1 calls["peak"] = max(calls["peak"], calls["active"]) try: return fn(payload) finally: calls["active"] -= 1 # embed_batch_chunked 定义于 vstore.embed；钩子注入到 embed 模块 import vstore.embed as embed_mod monkeypatch.setattr(embed_mod, _REQUEST_HOOK, hook) return calls def test_batch16_array_input_and_order(self, monkeypatch): """TC-PRIV-L1-14：batch_size=16 分批；payload input 为数组；顺序保持。""" import vstore.embed as embed_mod def resp(payload): assert isinstance(payload["input"], list), "payload.input 必须是数组（批量契约）" return [fake_vec(1) for _ in payload["input"]] calls = self._install_hook(monkeypatch, resp) texts = [f"文本{i}" for i in range(40)] out = embed_mod.embed_batch_chunked(texts, batch_size=16, workers=4) assert len(out) == 40 assert all(len(v) == DIM for v in out) # 批次数 == ceil(40/16) == 3；每批条数 ≤ batch_size assert len(calls["batches"]) == 3 assert all(len(b["input"]) <= 16 for b in calls["batches"]) # 顺序保持：返回向量与请求文本一一对应（首条非零） assert out[0] == fake_vec(1) def test_workers_concurrency_capped(self, monkeypatch): """TC-PRIV-L1-15：并发峰值 ≤ workers=4。""" import time import vstore.embed as embed_mod def slow(payload): time.sleep(0.05) return [fake_vec(2) for _ in payload["input"]] calls = self._install_hook(monkeypatch, slow) out = embed_mod.embed_batch_chunked([f"文本{i}" for i in range(33)], batch_size=16, workers=4) assert len(out) == 33 assert calls["peak"] <= 4 def test_batch_retry_once(self, monkeypatch): """TC-PRIV-L1-16：单批失败重试一次成功。""" import vstore.embed as embed_mod state = {"fails": {1: 1}} # 第 2 批首调失败一次 def flaky(payload): text = payload["input"][0] batch_no = int(text[2:]) // 16 if state["fails"].get(batch_no, 0) > 0: state["fails"][batch_no] -= 1 raise RuntimeError("模拟批量失败") return [fake_vec(3) for _ in payload["input"]] calls = self._install_hook(monkeypatch, flaky) out = embed_mod.embed_batch_chunked([f"文本{i}" for i in range(40)], batch_size=16, workers=4) assert len(out) == 40 # 重试后全部成功：无零向量 assert all(v != [0.0] * DIM for v in out) def test_batch_failure_degrades_to_zero(self, monkeypatch): """TC-PRIV-L1-17：重试仍失败 → 该批降级零向量，函数不抛。""" import vstore.embed as embed_mod def dead_batch(payload): text = payload["input"][0] if int(text[2:]) < 16: raise RuntimeError("永远失败的第一批") return [fake_vec(4) for _ in payload["input"]] self._install_hook(monkeypatch, dead_batch) out = embed_mod.embed_batch_chunked([f"文本{i}" for i in range(32)], batch_size=16, workers=4) assert len(out) == 32 assert out[:16] == [[0.0] * DIM] * 16 # 失败批零向量 assert all(v != [0.0] * DIM for v in out[16:]) def test_legacy_embed_batch_untouched(self, monkeypatch): """回归护栏：既有 embed_batch 仍存在且签名不变（FR-PRIV-008）。""" import vstore.embed as embed_mod assert callable(embed_mod.embed_batch) assert callable(embed_mod.EmbedError) assert callable(embed_mod.embed_batch_chunked) # --------------------------------------------------------------------------- # vault 钩子（TC-PRIV-L1-18 ~ L1-21）：与 A1 交付对齐——Vault.__init__ 已加载 # ChunkIndex（get_paths().chunk_path）；chunk 嵌入经 vault._chunked_embed → # vstore.embed.embed_batch_chunked（模块属性动态查找）。 # 纪律：夹具必须 set_paths 指向 tmp_path（隔离 chunk 路径）并桩掉批量嵌入， # 否则测试会触网真 Ollama。 # --------------------------------------------------------------------------- FAKE_CHUNK_VEC = [1.0 / DIM] * DIM # 归一化后单位向量（非零即可召回） def _install_vault_paths(monkeypatch, tmp_path) -> None: """隔离 vstore 路径（vault/index/chunks 全部落到 tmp_path，防 env 泄漏）。""" from vstore.settings import VstorePaths, set_paths set_paths(VstorePaths(root=str(tmp_path), vault_path="", index_path="", chunk_path="")) class TestVaultChunkHooks: @pytest.fixture def vault(self, tmp_path, embed_fn, monkeypatch): from vstore.vault import Vault _install_vault_paths(monkeypatch, tmp_path) # 桩掉 chunk 批量嵌入（走 vault._chunked_embed 的动态查找），不触网真 Ollama monkeypatch.setattr( "vstore.embed.embed_batch_chunked", lambda texts, *a, **k: [FAKE_CHUNK_VEC for _ in texts], ) v = Vault(tmp_path / "vault", tmp_path / "index") return v def test_add_returns_chunk_fields(self, vault): """TC-PRIV-L1-18：add 返回体含 chunk_count/chunk_ok/chunk_error，且成功。""" item = vault.add("正文段落。".encode("utf-8"), "a.txt", {"name": "资料"}) assert "chunk_count" in item assert "chunk_ok" in item assert "chunk_error" in item # A1 契约：成功路径 chunk_ok==True、chunk_error==""、chunk_count==ChunkIndex 行数 assert item["chunk_ok"] is True assert item["chunk_error"] == "" assert item["chunk_count"] >= 1 assert vault.chunks.count() == item["chunk_count"] def test_add_chunk_failure_does_not_block(self, vault, monkeypatch): """TC-PRIV-L1-19：chunk 失败不阻断 doc 级；chunk_ok False + 错误说明。""" monkeypatch.setattr( "vstore.embed.embed_batch_chunked", lambda *a, **k: (_ for _ in ()).throw(RuntimeError("chunk 嵌入失败")), ) item = vault.add("正文段落。".encode("utf-8"), "a.txt", {"name": "资料"}) assert item["document_id"] assert item["chunk_ok"] is False assert item["chunk_error"] def test_delete_removes_chunks(self, vault): """TC-PRIV-L1-20：delete 联动删除 chunk 行。""" item = vault.add("正文段落。".encode("utf-8"), "a.txt", {"name": "资料"}) doc_id = item["document_id"] assert vault.chunks.count() >= 1 assert vault.delete(doc_id) is True # chunk 行联动删除（FR-PRIV-006）：该 doc 的 chunk 全部消失 assert vault.chunks.count() == 0 def test_rebuild_chunks(self, vault): """TC-PRIV-L1-21：rebuild_chunks 返回 {docs, chunks, duration_ms}。""" vault.add("正文段落一。".encode("utf-8"), "a.txt", {"name": "资料甲"}) vault.add("正文段落二。".encode("utf-8"), "b.txt", {"name": "资料乙"}) result = vault.rebuild_chunks() assert "docs" in result and result["docs"] == 2 assert "chunks" in result and result["chunks"] > 0 assert "duration_ms" in result and isinstance(result["duration_ms"], (int, float))