kplam
51dc8dc4f6
- FastAPI + Qdrant + Redis + Ollama 技术栈 - L1→L2→L3→chunk 四层分层检索（dense + sparse RRF 融合） - 文档三级总结与 2.5 级回退 - query 解析路由与分类 - /admin 管理页面
329 lines
14 KiB
Python
329 lines
14 KiB
Python
"""端到端集成测试
|
||
|
||
真实 Ingester + Retriever + 内存 Qdrant（location=":memory:"）串联分层 RAG 全链路，
|
||
仅替换两个外部边界：
|
||
- FakeOllama：按 prompt 内容返回确定性的 L1/L2/L3 总结、分类 JSON、query 解析 JSON
|
||
- DeterministicEmbedding：稳定哈希生成固定维度向量（无语义，仅保证流程可跑）
|
||
|
||
本机无需 Docker / Ollama / Redis，CI 可直接运行。
|
||
"""
|
||
|
||
import json
|
||
import random
|
||
import uuid
|
||
from dataclasses import dataclass
|
||
from hashlib import sha256
|
||
from typing import Any
|
||
|
||
import pytest
|
||
from fastapi.testclient import TestClient
|
||
from qdrant_client import AsyncQdrantClient
|
||
|
||
from app.api.v1 import document as document_module
|
||
from app.api.v1 import search as search_module
|
||
from app.config import Settings, settings
|
||
from app.core.chunker import Chunker
|
||
from app.core.classifier import Classifier
|
||
from app.core.ingest_tasks import IngestTaskManager, IngestTaskStatus
|
||
from app.core.ingestion import Ingester
|
||
from app.core.query_parser import QueryParser
|
||
from app.core.retriever import Retriever
|
||
from app.core.sparse import SparseEncoder
|
||
from app.core.summarizer import Summarizer
|
||
from app.main import app
|
||
from app.models.document import DocumentInput
|
||
from app.models.knowledge import load_taxonomy
|
||
from app.models.search import SearchRequest
|
||
from app.services.qdrant import (
|
||
ALL_COLLECTIONS,
|
||
COLLECTION_CHUNKS,
|
||
COLLECTION_L1,
|
||
COLLECTION_L2,
|
||
COLLECTION_L3,
|
||
QdrantService,
|
||
)
|
||
|
||
# FakeOllama 的确定性输出
|
||
FAKE_L1_SUMMARY = "本文介绍安装指南、环境准备与安装步骤，是一篇集成测试文档。"
|
||
FAKE_L3_OUTLINE = (
|
||
"## 安装指南\n安装指南的整体流程说明。\n"
|
||
"## 环境准备\n环境准备的依赖与注意事项。\n"
|
||
"## 安装步骤\n安装步骤的命令与验证方法。"
|
||
)
|
||
FAKE_L2_HALF = "要点一：短文档的核心通知内容。\n要点二：需要关注的事项。"
|
||
FAKE_CATEGORY = "技术文档"
|
||
FAKE_TAGS = ["安装", "运维"]
|
||
|
||
|
||
class FakeOllama:
|
||
"""按 prompt 内容返回确定性结果的假 OllamaClient
|
||
|
||
覆盖 Summarizer / Classifier / QueryParser 三类调用方；
|
||
分类与 query 解析的置信度可配置，用于构造路由兜底场景。
|
||
"""
|
||
|
||
def __init__(self, classify_confidence: float = 0.9, query_confidence: float = 0.9) -> None:
|
||
self.classify_confidence = classify_confidence
|
||
self.query_confidence = query_confidence
|
||
|
||
async def generate(self, prompt: str, json_mode: bool = False) -> str:
|
||
if "你是知识库分类助手" in prompt:
|
||
return json.dumps(
|
||
{"main_category": FAKE_CATEGORY, "tags": FAKE_TAGS, "confidence": self.classify_confidence},
|
||
ensure_ascii=False,
|
||
)
|
||
if "你是搜索查询分析助手" in prompt:
|
||
return json.dumps(
|
||
{
|
||
"categories": [{"name": FAKE_CATEGORY, "confidence": self.query_confidence}],
|
||
"rewrite": "安装指南 环境准备 安装步骤",
|
||
"keywords": ["安装", "环境准备"],
|
||
},
|
||
ensure_ascii=False,
|
||
)
|
||
if "请用一句话对以下文档内容进行高度概括" in prompt:
|
||
return FAKE_L1_SUMMARY
|
||
if "请提取以下文档的主要章节结构" in prompt:
|
||
return "1. 主题一\n2. 主题二"
|
||
if "请对以下文档的每个章节/主题进行详细的内容摘要" in prompt:
|
||
return FAKE_L3_OUTLINE
|
||
if "请对以下文档内容进行详细摘要" in prompt:
|
||
return FAKE_L2_HALF
|
||
raise AssertionError(f"FakeOllama 收到未识别的 prompt: {prompt[:100]}")
|
||
|
||
|
||
class DeterministicEmbedding:
|
||
"""稳定哈希伪向量：同文本恒同向量，维度等于 settings.embedding_dimension"""
|
||
|
||
async def embed(self, texts: list[str]) -> list[list[float]]:
|
||
vectors: list[list[float]] = []
|
||
for text in texts:
|
||
seed = int.from_bytes(sha256(text.encode("utf-8")).digest()[:8], "big")
|
||
rng = random.Random(seed)
|
||
vectors.append([rng.random() for _ in range(settings.embedding_dimension)])
|
||
return vectors
|
||
|
||
|
||
class FakeCache:
|
||
"""无缓存行为的假 RedisCache（get 恒未命中，set 恒成功）"""
|
||
|
||
async def get_json(self, key: str) -> dict[str, Any] | None:
|
||
return None
|
||
|
||
async def set_json(self, key: str, value: dict[str, Any], ttl: int | None = None) -> bool:
|
||
return True
|
||
|
||
|
||
@dataclass
|
||
class _Env:
|
||
"""一套共享内存 Qdrant 的真实 Ingester + Retriever 环境"""
|
||
|
||
qdrant: QdrantService
|
||
ollama: FakeOllama
|
||
ingester: Ingester
|
||
retriever: Retriever
|
||
|
||
|
||
async def _make_env(classify_confidence: float = 0.9, query_confidence: float = 0.9) -> _Env:
|
||
"""构建集成环境：真实组件 + 内存 Qdrant + FakeOllama + 确定性向量"""
|
||
qdrant = QdrantService(client=AsyncQdrantClient(location=":memory:"))
|
||
await qdrant.ensure_collections()
|
||
ollama = FakeOllama(classify_confidence=classify_confidence, query_confidence=query_confidence)
|
||
taxonomy = load_taxonomy()
|
||
embedding = DeterministicEmbedding()
|
||
ingester = Ingester(
|
||
summarizer=Summarizer(ollama=ollama), # type: ignore[arg-type]
|
||
classifier=Classifier(ollama=ollama, taxonomy=taxonomy), # type: ignore[arg-type]
|
||
chunker=Chunker(),
|
||
embedding=embedding,
|
||
sparse=SparseEncoder(),
|
||
qdrant=qdrant,
|
||
)
|
||
retriever = Retriever(
|
||
qdrant=qdrant,
|
||
query_parser=QueryParser(ollama=ollama, taxonomy=taxonomy, cache=FakeCache()), # type: ignore[arg-type]
|
||
embedding=embedding,
|
||
sparse_encoder=SparseEncoder(),
|
||
)
|
||
return _Env(qdrant=qdrant, ollama=ollama, ingester=ingester, retriever=retriever)
|
||
|
||
|
||
def _structured_doc() -> DocumentInput:
|
||
"""带 Markdown 标题结构的中文长文档（>500 字符，切出 3 个 section chunk）"""
|
||
paragraph = "这是章节正文内容，包含足够多的信息量，用于测试切分与向量化流程。" * 20
|
||
text = f"# 安装指南\n{paragraph}\n\n## 环境准备\n{paragraph}\n\n## 安装步骤\n{paragraph}"
|
||
return DocumentInput(text=text, title="安装文档")
|
||
|
||
|
||
def _short_doc() -> DocumentInput:
|
||
"""短文档（< summary_min_text_length），触发 2.5 级回退"""
|
||
return DocumentInput(text="# 维护通知\n明天凌晨系统维护，请提前保存工作。", title="维护通知")
|
||
|
||
|
||
async def _counts(qdrant: QdrantService) -> dict[str, int]:
|
||
"""四个集合的点数"""
|
||
return {name: (await qdrant.client.count(collection_name=name)).count for name in ALL_COLLECTIONS}
|
||
|
||
|
||
class TestIngestIntegration:
|
||
"""入库链路：四层集合点数与 chunk payload 完整性"""
|
||
|
||
async def test_ingest_structured_document(self):
|
||
"""结构化长文档 → L1=1、L2=标题数、L3/chunks ≥1，chunk payload 字段齐全"""
|
||
env = await _make_env()
|
||
doc = _structured_doc()
|
||
|
||
result = await env.ingester.ingest(doc)
|
||
|
||
# 结果透传
|
||
assert result.document_id
|
||
assert result.category == FAKE_CATEGORY
|
||
assert result.tags == FAKE_TAGS
|
||
assert result.category_confidence == 0.9
|
||
assert result.summary.l1_summary == FAKE_L1_SUMMARY
|
||
assert result.chunks_count >= 1
|
||
|
||
counts = await _counts(env.qdrant)
|
||
assert counts[COLLECTION_L1] == 1
|
||
assert counts[COLLECTION_L2] == 3 # 标题数：安装指南/环境准备/安装步骤
|
||
assert counts[COLLECTION_L3] >= 1
|
||
assert counts[COLLECTION_CHUNKS] == result.chunks_count >= 1
|
||
|
||
# chunk payload：doc_summary/category/tags/section_path 齐全且与分类结果一致
|
||
points, _ = await env.qdrant.client.scroll(
|
||
collection_name=COLLECTION_CHUNKS, limit=100, with_payload=True
|
||
)
|
||
assert len(points) == result.chunks_count
|
||
expected_paths = {"安装指南", "安装指南 / 环境准备", "安装指南 / 安装步骤"}
|
||
for point in points:
|
||
payload = point.payload or {}
|
||
assert payload["doc_id"] == result.document_id
|
||
assert payload["doc_summary"] == FAKE_L1_SUMMARY
|
||
assert payload["category"] == FAKE_CATEGORY
|
||
assert payload["tags"] == FAKE_TAGS
|
||
assert payload["section_path"] in expected_paths
|
||
assert payload["text"] in doc.text
|
||
|
||
async def test_ingest_short_document_l2_half(self):
|
||
"""短文档 → 2.5 级：doc_l2 无该 doc 节点，其余层正常写入"""
|
||
env = await _make_env()
|
||
|
||
result = await env.ingester.ingest(_short_doc())
|
||
|
||
assert result.summary.l2_outline is None
|
||
counts = await _counts(env.qdrant)
|
||
assert counts[COLLECTION_L1] == 1
|
||
assert counts[COLLECTION_L2] == 0 # 2.5 级文档无 L2 大纲节点
|
||
assert counts[COLLECTION_L3] >= 1
|
||
assert counts[COLLECTION_CHUNKS] >= 1
|
||
|
||
async def test_reingest_idempotent(self, monkeypatch: pytest.MonkeyPatch):
|
||
"""同 doc_id 重复入库：幂等覆盖不报错，各层点数不翻倍"""
|
||
env = await _make_env()
|
||
doc = _structured_doc()
|
||
# 固定 doc_id，两次入库写入同一组确定性 point id（uuid5 覆盖）
|
||
fixed_uuid = uuid.UUID("12345678-1234-5678-1234-567812345678")
|
||
monkeypatch.setattr("app.core.ingestion.uuid.uuid4", lambda: fixed_uuid)
|
||
|
||
first = await env.ingester.ingest(doc)
|
||
counts_after_first = await _counts(env.qdrant)
|
||
second = await env.ingester.ingest(doc)
|
||
counts_after_second = await _counts(env.qdrant)
|
||
|
||
assert first.document_id == second.document_id == fixed_uuid.hex
|
||
assert counts_after_second == counts_after_first
|
||
assert counts_after_second[COLLECTION_L1] == 1
|
||
assert counts_after_second[COLLECTION_L2] == 3
|
||
|
||
|
||
class TestSearchIntegration:
|
||
"""检索链路：正常路由 / 路由兜底 / 空库"""
|
||
|
||
async def test_search_normal_query(self):
|
||
"""正常 query：hits 非空，text 为原文片段，doc_summary 非空，routed_categories 来自分类"""
|
||
env = await _make_env()
|
||
doc = _structured_doc()
|
||
await env.ingester.ingest(doc)
|
||
|
||
resp = await env.retriever.search(SearchRequest(query="安装步骤有哪些注意事项？"))
|
||
|
||
assert resp.fallback is False
|
||
assert resp.routed_categories == [FAKE_CATEGORY]
|
||
assert resp.hits
|
||
hit = resp.hits[0]
|
||
assert hit.text in doc.text
|
||
assert hit.doc_summary == FAKE_L1_SUMMARY
|
||
assert hit.score > 0
|
||
|
||
async def test_search_route_fallback(self):
|
||
"""query 解析低置信 → 路由兜底：fallback=True 且仍有 hits"""
|
||
env = await _make_env(query_confidence=0.3)
|
||
doc = _structured_doc()
|
||
await env.ingester.ingest(doc)
|
||
|
||
resp = await env.retriever.search(SearchRequest(query="随便问点什么"))
|
||
|
||
assert resp.fallback is True
|
||
assert resp.routed_categories == []
|
||
assert resp.hits
|
||
assert resp.hits[0].text in doc.text
|
||
|
||
async def test_search_empty_db(self):
|
||
"""空库 search：hits 为空、不报错、fallback=True"""
|
||
env = await _make_env()
|
||
|
||
resp = await env.retriever.search(SearchRequest(query="空库查询"))
|
||
|
||
assert resp.hits == []
|
||
assert resp.fallback is True
|
||
|
||
|
||
class TestApiIntegration:
|
||
"""API 层冒烟：documents → search → knowledge/categories，统一响应格式 code=0"""
|
||
|
||
async def test_api_smoke(self, monkeypatch: pytest.MonkeyPatch):
|
||
env = await _make_env()
|
||
doc = _structured_doc()
|
||
# 入库走异步任务：测试自建纯内存任务管理器（内存 Qdrant + FakeOllama 的 Ingester）
|
||
manager = IngestTaskManager(env.ingester, None, Settings())
|
||
|
||
# lifespan 的 ensure_collections 替换为空操作（真实 Qdrant 不可达时也能启动）
|
||
async def _noop_ensure_collections(self: QdrantService) -> None:
|
||
return None
|
||
|
||
monkeypatch.setattr(QdrantService, "ensure_collections", _noop_ensure_collections)
|
||
# 模块级单例替换为集成环境实例
|
||
monkeypatch.setattr(document_module, "_task_manager", manager)
|
||
monkeypatch.setattr(search_module, "_retriever", env.retriever)
|
||
monkeypatch.setattr(search_module, "get_cache", lambda: FakeCache())
|
||
|
||
with TestClient(app) as client:
|
||
# 入库：202 拿 task_id，等待任务终态后断言入库结果
|
||
resp_doc = client.post("/api/v1/documents", json={"text": doc.text, "title": doc.title})
|
||
assert resp_doc.status_code == 202
|
||
body_doc = resp_doc.json()
|
||
assert body_doc["code"] == 0
|
||
assert body_doc["data"]["status"] == "pending"
|
||
task_id = body_doc["data"]["task_id"]
|
||
assert task_id
|
||
|
||
final = await manager.wait_done(task_id)
|
||
assert final["status"] == IngestTaskStatus.DONE
|
||
result = final["result"]
|
||
assert result["document_id"]
|
||
assert result["category"] == FAKE_CATEGORY
|
||
assert result["chunks_count"] >= 1
|
||
|
||
# 检索
|
||
resp_search = client.post("/api/v1/search", json={"query": "安装步骤有哪些注意事项？"})
|
||
body_search = resp_search.json()
|
||
assert body_search["code"] == 0
|
||
assert body_search["data"]["hits"]
|
||
assert body_search["data"]["routed_categories"] == [FAKE_CATEGORY]
|
||
|
||
# 知识分类
|
||
resp_categories = client.get("/api/v1/knowledge/categories")
|
||
body_categories = resp_categories.json()
|
||
assert body_categories["code"] == 0
|
||
assert body_categories["data"]["count"] > 0
|