kplam
51dc8dc4f6
- FastAPI + Qdrant + Redis + Ollama 技术栈 - L1→L2→L3→chunk 四层分层检索（dense + sparse RRF 融合） - 文档三级总结与 2.5 级回退 - query 解析路由与分类 - /admin 管理页面
274 lines
10 KiB
Python
274 lines
10 KiB
Python
"""Ingester 全链路单元测试（Summarizer/Classifier/Embedding/Qdrant 均为假实现，不真实联网）"""
|
||
|
||
from typing import Any
|
||
|
||
import pytest
|
||
|
||
from app.core.chunker import Chunker
|
||
from app.core.ingestion import Ingester, IngestionError
|
||
from app.core.sparse import SparseEncoder
|
||
from app.models.document import DocumentInput, DocumentSummary, SummaryLevel
|
||
from app.models.knowledge import CategoryResult
|
||
from app.services.qdrant import COLLECTION_L2, COLLECTION_L3
|
||
|
||
|
||
class FakeSummarizer:
|
||
"""返回固定总结结果的假 Summarizer"""
|
||
|
||
def __init__(self, summary: DocumentSummary) -> None:
|
||
self.summary = summary
|
||
self.calls: list[tuple[str, str]] = []
|
||
|
||
async def summarize(self, text: str, *, title: str = "") -> DocumentSummary:
|
||
self.calls.append((text, title))
|
||
return self.summary
|
||
|
||
|
||
class FakeClassifier:
|
||
"""返回固定分类结果的假 Classifier"""
|
||
|
||
def __init__(self, result: CategoryResult) -> None:
|
||
self.result = result
|
||
self.calls: list[tuple[str, str]] = []
|
||
|
||
async def classify(self, l1_summary: str, title: str = "") -> CategoryResult:
|
||
self.calls.append((l1_summary, title))
|
||
return self.result
|
||
|
||
|
||
class FakeEmbedding:
|
||
"""按输入数量返回伪向量的假 EmbeddingService，记录每次调用的文本"""
|
||
|
||
def __init__(self) -> None:
|
||
self.calls: list[list[str]] = []
|
||
|
||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||
self.calls.append(list(texts))
|
||
return [[float(i), 1.0] for i in range(len(texts))]
|
||
|
||
|
||
class FakeQdrant:
|
||
"""内存版 QdrantService，记录各层 upsert 调用；可配置在某一层抛错"""
|
||
|
||
def __init__(self, fail_on: str = "") -> None:
|
||
self.fail_on = fail_on
|
||
self.l1_calls: list[dict[str, Any]] = []
|
||
self.nodes_calls: list[tuple[str, list[dict[str, Any]]]] = []
|
||
self.chunks_calls: list[list[dict[str, Any]]] = []
|
||
|
||
async def upsert_l1(
|
||
self,
|
||
doc_id: str,
|
||
title: str,
|
||
summary: str,
|
||
category: str,
|
||
tags: list[str],
|
||
dense_vector: list[float],
|
||
sparse_vector: Any = None,
|
||
) -> None:
|
||
if self.fail_on == "l1":
|
||
raise RuntimeError("qdrant down")
|
||
self.l1_calls.append(
|
||
{
|
||
"doc_id": doc_id,
|
||
"title": title,
|
||
"summary": summary,
|
||
"category": category,
|
||
"tags": tags,
|
||
"dense_vector": dense_vector,
|
||
"sparse_vector": sparse_vector,
|
||
}
|
||
)
|
||
|
||
async def upsert_nodes(self, collection: str, nodes: list[dict[str, Any]]) -> None:
|
||
if self.fail_on == collection:
|
||
raise RuntimeError("qdrant down")
|
||
self.nodes_calls.append((collection, nodes))
|
||
|
||
async def upsert_chunks(self, chunks: list[dict[str, Any]]) -> None:
|
||
if self.fail_on == "chunks":
|
||
raise RuntimeError("qdrant down")
|
||
self.chunks_calls.append(chunks)
|
||
|
||
|
||
def _make_ingester(
|
||
summary: DocumentSummary,
|
||
category: CategoryResult,
|
||
qdrant: FakeQdrant,
|
||
embedding: FakeEmbedding | None = None,
|
||
) -> Ingester:
|
||
return Ingester(
|
||
summarizer=FakeSummarizer(summary), # type: ignore[arg-type]
|
||
classifier=FakeClassifier(category), # type: ignore[arg-type]
|
||
chunker=Chunker(),
|
||
embedding=embedding or FakeEmbedding(), # type: ignore[arg-type]
|
||
sparse=SparseEncoder(),
|
||
qdrant=qdrant, # type: ignore[arg-type]
|
||
)
|
||
|
||
|
||
def _structured_doc() -> DocumentInput:
|
||
"""带标题结构的长文档（切出多个 chunk，L2 走标题树）"""
|
||
paragraph = "这是章节正文内容，包含足够多的信息量，用于测试切分与向量化流程。" * 10
|
||
text = f"# 安装指南\n{paragraph}\n\n## 环境准备\n{paragraph}\n\n## 安装步骤\n{paragraph}"
|
||
return DocumentInput(text=text, title="安装文档")
|
||
|
||
|
||
def _structured_summary() -> DocumentSummary:
|
||
return DocumentSummary(
|
||
l1_summary="本文介绍软件的安装流程。",
|
||
l2_outline="- 安装指南\n - 环境准备\n - 安装步骤",
|
||
l3_content_outline=(
|
||
"## 安装指南\n整体安装流程说明。\n## 环境准备\n准备依赖环境。\n## 安装步骤\n执行安装命令。"
|
||
),
|
||
level=SummaryLevel.L3,
|
||
)
|
||
|
||
|
||
def _category() -> CategoryResult:
|
||
return CategoryResult(main_category="技术文档", tags=["安装", "运维"], confidence=0.9)
|
||
|
||
|
||
class TestStructuredDocument:
|
||
"""结构化长文档：四层集合 upsert 均被调用，结果字段透传"""
|
||
|
||
async def test_full_pipeline(self):
|
||
doc = _structured_doc()
|
||
summary = _structured_summary()
|
||
qdrant = FakeQdrant()
|
||
embedding = FakeEmbedding()
|
||
ingester = _make_ingester(summary, _category(), qdrant, embedding)
|
||
|
||
result = await ingester.ingest(doc)
|
||
|
||
# 结果字段透传
|
||
assert result.document_id
|
||
assert result.summary is summary
|
||
assert result.category == "技术文档"
|
||
assert result.tags == ["安装", "运维"]
|
||
assert result.category_confidence == 0.9
|
||
|
||
# chunk 数与 Chunker 直出一致
|
||
expected_chunks = Chunker().chunk(doc.text, "expected")
|
||
assert result.chunks_count == len(expected_chunks) > 1
|
||
|
||
# 批量 embedding 恰好一次：L1 + L2 节点 + L3 节点 + chunks
|
||
assert len(embedding.calls) == 1
|
||
l2_calls = [nodes for c, nodes in qdrant.nodes_calls if c == COLLECTION_L2]
|
||
l3_calls = [nodes for c, nodes in qdrant.nodes_calls if c == COLLECTION_L3]
|
||
assert len(embedding.calls[0]) == 1 + len(l2_calls[0]) + len(l3_calls[0]) + result.chunks_count
|
||
|
||
# L1：category/tags 透传，sparse 已启用
|
||
assert len(qdrant.l1_calls) == 1
|
||
l1 = qdrant.l1_calls[0]
|
||
assert l1["doc_id"] == result.document_id
|
||
assert l1["title"] == "安装文档"
|
||
assert l1["summary"] == summary.l1_summary
|
||
assert l1["category"] == "技术文档"
|
||
assert l1["tags"] == ["安装", "运维"]
|
||
assert l1["sparse_vector"] is not None
|
||
|
||
# L2：每个标题一个节点，text 与 section_path 均为祖先标题链
|
||
assert len(l2_calls) == 1
|
||
l2_nodes = l2_calls[0]
|
||
assert len(l2_nodes) == 3
|
||
assert [n["section_path"] for n in l2_nodes] == [
|
||
"安装指南",
|
||
"安装指南 / 环境准备",
|
||
"安装指南 / 安装步骤",
|
||
]
|
||
assert all(n["text"] == n["section_path"] for n in l2_nodes)
|
||
assert all(n["category"] == "技术文档" and n["tags"] == ["安装", "运维"] for n in l2_nodes)
|
||
|
||
# L3：按 "## " 分块，section_path 精确匹配到标题链
|
||
assert len(l3_calls) == 1
|
||
l3_nodes = l3_calls[0]
|
||
assert len(l3_nodes) == 3
|
||
assert [n["section_path"] for n in l3_nodes] == [
|
||
"安装指南",
|
||
"安装指南 / 环境准备",
|
||
"安装指南 / 安装步骤",
|
||
]
|
||
assert l3_nodes[0]["text"].startswith("## 安装指南")
|
||
|
||
# chunks：携带 doc_summary 与 sparse 向量
|
||
assert len(qdrant.chunks_calls) == 1
|
||
chunk_dicts = qdrant.chunks_calls[0]
|
||
assert len(chunk_dicts) == result.chunks_count
|
||
assert all(c["doc_summary"] == summary.l1_summary for c in chunk_dicts)
|
||
assert all(c["sparse_vector"] is not None for c in chunk_dicts)
|
||
assert all(c["category"] == "技术文档" for c in chunk_dicts)
|
||
assert [c["chunk_index"] for c in chunk_dicts] == list(range(result.chunks_count))
|
||
|
||
|
||
class TestFallbackDocuments:
|
||
"""2.5 级文档与无结构文档的 L2/L3 节点构建"""
|
||
|
||
async def test_l2_half_document_skips_l2_upsert(self):
|
||
"""2.5 级文档（l2_outline=None）→ 不写 L2 集合，L3 整块一个节点"""
|
||
summary = DocumentSummary(
|
||
l1_summary="一条简短通知。",
|
||
l2_outline=None,
|
||
l3_content_outline="要点一：明天放假。\n要点二：注意安全。",
|
||
level=SummaryLevel.L2_HALF,
|
||
)
|
||
doc = DocumentInput(text="简短通知正文，无标题结构。", title="通知")
|
||
qdrant = FakeQdrant()
|
||
ingester = _make_ingester(summary, _category(), qdrant)
|
||
|
||
result = await ingester.ingest(doc)
|
||
|
||
collections = [c for c, _ in qdrant.nodes_calls]
|
||
assert COLLECTION_L2 not in collections
|
||
# L3 内容大纲无 "## " → 整块一个节点，section_path 为空
|
||
l3_calls = [nodes for c, nodes in qdrant.nodes_calls if c == COLLECTION_L3]
|
||
assert len(l3_calls) == 1
|
||
assert len(l3_calls[0]) == 1
|
||
assert l3_calls[0][0]["section_path"] == ""
|
||
assert l3_calls[0][0]["text"] == summary.l3_content_outline
|
||
# 其余层级正常写入
|
||
assert len(qdrant.l1_calls) == 1
|
||
assert result.chunks_count == len(qdrant.chunks_calls[0])
|
||
|
||
async def test_l2_from_llm_outline_lines(self):
|
||
"""无标题结构的 L3 级文档 → L2 节点来自 LLM 大纲行，section_path 为空"""
|
||
summary = DocumentSummary(
|
||
l1_summary="本文介绍两个主题。",
|
||
l2_outline="1. 主题一\n2. 主题二",
|
||
l3_content_outline="详细摘要内容，无分块标题。",
|
||
level=SummaryLevel.L3,
|
||
)
|
||
text = "这是一段没有标题结构的正文内容，" * 40
|
||
doc = DocumentInput(text=text, title="")
|
||
qdrant = FakeQdrant()
|
||
ingester = _make_ingester(summary, _category(), qdrant)
|
||
|
||
await ingester.ingest(doc)
|
||
|
||
l2_calls = [nodes for c, nodes in qdrant.nodes_calls if c == COLLECTION_L2]
|
||
assert len(l2_calls) == 1
|
||
l2_nodes = l2_calls[0]
|
||
assert [n["text"] for n in l2_nodes] == ["1. 主题一", "2. 主题二"]
|
||
assert all(n["section_path"] == "" for n in l2_nodes)
|
||
|
||
|
||
class TestQdrantFailure:
|
||
"""Qdrant 写入失败：抛 IngestionError，stage=qdrant，总结不丢可重试"""
|
||
|
||
async def test_chunks_upsert_failure(self):
|
||
doc = _structured_doc()
|
||
summary = _structured_summary()
|
||
qdrant = FakeQdrant(fail_on="chunks")
|
||
ingester = _make_ingester(summary, _category(), qdrant)
|
||
|
||
with pytest.raises(IngestionError) as exc_info:
|
||
await ingester.ingest(doc)
|
||
|
||
err = exc_info.value
|
||
assert err.stage == "qdrant"
|
||
assert err.summary is summary
|
||
# L1/L2/L3 已写入，失败发生在 chunks 层
|
||
assert len(qdrant.l1_calls) == 1
|
||
assert {c for c, _ in qdrant.nodes_calls} == {COLLECTION_L2, COLLECTION_L3}
|
||
assert qdrant.chunks_calls == []
|