jerryandClaude Fable 5
5db3ead73e
- vstore 独立安装包(8094):解析/嵌入/加密 vault/faiss 检索,存储路径支持本地 file 与 http(s):// URL,启动握手校验,集群化预留(embed 独立模块) - BFF docgen 路由:上传入库/generate/编辑/导出/下载/主库管理,意图护栏(与前端同词典),vstore 失败兜底,quota 507 与 TTL 惰性清理 - 前端:自然触发意图路由、拖拽上传 chips、能力徽标、右侧文档 TAB(WYSIWYG 编辑/prompt 修改/PDF+DOCX 导出下载)、右栏拖宽+手柄仿左栏 - 批量导入工具 docgen_import.py(并发 8、导入即加密+删源明文) Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
120 lines
4.1 KiB
Python
120 lines
4.1 KiB
Python
"""test_retrieval.py —— faiss 索引 + index.json 单一事实源 + 三档降级。"""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
|
|
import numpy as np
|
|
|
|
from vstore.retrieval import INDEX_FILE, RetrievalIndex
|
|
|
|
|
|
def _vec(idx: int, dim: int = 1024) -> list[float]:
|
|
v = np.zeros(dim, dtype=np.float32)
|
|
v[idx % dim] = 1.0
|
|
return v.tolist()
|
|
|
|
|
|
def _add(ri: RetrievalIndex, doc_id: str, name: str, idx: int, **meta):
|
|
row = {
|
|
"id": doc_id, "name": name, "category": meta.get("category", ""),
|
|
"tags": meta.get("tags", ""), "description": meta.get("description", ""),
|
|
"as_template": meta.get("as_template", False),
|
|
"created_at": meta.get("created_at", "2026-09-01T00:00:00Z"),
|
|
}
|
|
return ri.add_item(row, _vec(idx))
|
|
|
|
|
|
def test_add_and_json_row_order(tmp_path):
|
|
ri = RetrievalIndex(tmp_path)
|
|
_add(ri, "a", "甲", 0)
|
|
_add(ri, "b", "乙", 1)
|
|
data = json.loads((tmp_path / INDEX_FILE).read_text(encoding="utf-8"))
|
|
assert [it["id"] for it in data["items"]] == ["a", "b"] # 行号 == 下标
|
|
assert ri.count() == 2
|
|
|
|
|
|
def test_load_rebuilds_faiss(tmp_path):
|
|
ri = RetrievalIndex(tmp_path)
|
|
_add(ri, "a", "甲", 0)
|
|
ri2 = RetrievalIndex(tmp_path)
|
|
ri2.load()
|
|
result = ri2.search(_vec(0), top_k=1)
|
|
assert result["mode"] == "vector"
|
|
assert result["items"][0]["id"] == "a"
|
|
|
|
|
|
def test_search_vector_mode(tmp_path):
|
|
ri = RetrievalIndex(tmp_path)
|
|
_add(ri, "a", "小麦", 0)
|
|
_add(ri, "b", "水稻", 1)
|
|
result = ri.search(_vec(1), top_k=1)
|
|
assert result["mode"] == "vector"
|
|
assert result["items"][0]["id"] == "b"
|
|
|
|
|
|
def test_search_filters_category_and_template(tmp_path):
|
|
ri = RetrievalIndex(tmp_path)
|
|
_add(ri, "a", "模板甲", 0, category="公文", as_template=True)
|
|
_add(ri, "b", "资料乙", 1, category="资料", as_template=False)
|
|
result = ri.search(_vec(0), top_k=5, category="公文")
|
|
assert all(i["category"] == "公文" for i in result["items"])
|
|
result2 = ri.search(_vec(0), top_k=5, as_template=True)
|
|
assert all(i["as_template"] for i in result2["items"])
|
|
|
|
|
|
def test_search_keyword_fallback(tmp_path):
|
|
"""向量无命中时（零向量+过滤）→ meta 关键词档。"""
|
|
ri = RetrievalIndex(tmp_path)
|
|
_add(ri, "a", "小麦育种", 0, description="关于小麦的内容")
|
|
_add(ri, "b", "水稻栽培", 1, description="关于水稻的内容")
|
|
result = ri.search([0.0] * 1024, top_k=2, keywords=["小麦"])
|
|
assert result["mode"] == "keyword"
|
|
assert result["items"][0]["id"] == "a"
|
|
|
|
|
|
def test_search_latest_fallback(tmp_path):
|
|
ri = RetrievalIndex(tmp_path)
|
|
_add(ri, "a", "甲", 0, created_at="2026-09-01T00:00:00Z")
|
|
_add(ri, "b", "乙", 1, created_at="2026-09-02T00:00:00Z")
|
|
result = ri.search([0.0] * 1024, top_k=2) # 无关键词
|
|
assert result["mode"] == "latest"
|
|
assert result["items"][0]["id"] == "b" # 最新在前
|
|
|
|
|
|
def test_search_empty(tmp_path):
|
|
ri = RetrievalIndex(tmp_path)
|
|
result = ri.search(_vec(0))
|
|
assert result == {"mode": "empty", "items": []}
|
|
|
|
|
|
def test_remove_item(tmp_path):
|
|
ri = RetrievalIndex(tmp_path)
|
|
_add(ri, "a", "甲", 0)
|
|
_add(ri, "b", "乙", 1)
|
|
assert ri.remove_item("a") is True
|
|
assert ri.count() == 1
|
|
assert ri.remove_item("nope") is False
|
|
|
|
|
|
def test_update_meta(tmp_path):
|
|
ri = RetrievalIndex(tmp_path)
|
|
_add(ri, "a", "甲", 0)
|
|
assert ri.update_meta("a", {"name": "甲改", "category": "新类"}) is True
|
|
rows = ri.rows()
|
|
assert rows[0]["name"] == "甲改"
|
|
assert rows[0]["category"] == "新类"
|
|
assert ri.update_meta("nope", {"name": "x"}) is False
|
|
|
|
|
|
def test_rebuild_atomic_switch(tmp_path):
|
|
ri = RetrievalIndex(tmp_path)
|
|
_add(ri, "old", "旧", 0)
|
|
items = [({"id": "n1", "name": "新一", "created_at": "2026-09-03T00:00:00Z"}, _vec(2)),
|
|
({"id": "n2", "name": "新二", "created_at": "2026-09-03T00:00:00Z"}, _vec(3))]
|
|
count = ri.rebuild(items)
|
|
assert count == 2
|
|
data = json.loads((tmp_path / INDEX_FILE).read_text(encoding="utf-8"))
|
|
assert [it["id"] for it in data["items"]] == ["n1", "n2"]
|
|
assert not (tmp_path / (INDEX_FILE + ".rebuild.tmp")).exists() # 临时文件已切换
|
|
assert ri.count() == 2
|