Files
2026-09-07 22:41:36 +08:00

460 lines
20 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""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))