"""QdrantService 测试 使用 AsyncQdrantClient(location=":memory:") 本地模式，无需 Docker。 """ import pytest from qdrant_client import AsyncQdrantClient from app.config import settings from app.services.qdrant import ( ALL_COLLECTIONS, COLLECTION_CHUNKS, COLLECTION_L1, COLLECTION_L2, COLLECTION_L3, QdrantService, ) DIM = settings.embedding_dimension def _dense(seed: float) -> list[float]: """构造确定性 dense 向量：前 4 维取特征值，便于区分不同文档""" vec = [0.0] * DIM vec[0] = seed vec[1] = 1.0 - seed vec[2] = seed * 0.5 vec[3] = 0.1 return vec @pytest.fixture async def service() -> QdrantService: svc = QdrantService(client=AsyncQdrantClient(location=":memory:")) await svc.ensure_collections() return svc async def test_ensure_collections_idempotent(service: QdrantService) -> None: """重复调用 ensure_collections 不报错，且 4 个集合均存在""" await service.ensure_collections() # fixture 中已调一次，这里第二次 collections = await service.client.get_collections() names = {c.name for c in collections.collections} assert set(ALL_COLLECTIONS) <= names # L1 与 chunks 应配置 sparse 命名向量 for name in (COLLECTION_L1, COLLECTION_CHUNKS): info = await service.client.get_collection(name) assert info.config.params.sparse_vectors is not None assert "sparse" in info.config.params.sparse_vectors async def test_upsert_l1_and_filter_by_doc_id(service: QdrantService) -> None: await service.upsert_l1( doc_id="doc-1", title="文档一", summary="这是文档一的总结", category="tech", tags=["ai", "rag"], dense_vector=_dense(0.9), sparse_vector=([1, 2, 3], [0.5, 0.3, 0.2]), ) # 重复写入（幂等覆盖）不报错 await service.upsert_l1( doc_id="doc-1", title="文档一", summary="这是文档一的总结（更新）", category="tech", tags=["ai", "rag"], dense_vector=_dense(0.9), ) results = await service.search_dense( COLLECTION_L1, _dense(0.9), limit=5, query_filter=QdrantService.build_filter(doc_ids=["doc-1"]), ) assert len(results) == 1 assert results[0].payload is not None assert results[0].payload["doc_id"] == "doc-1" assert results[0].payload["text"] == "这是文档一的总结（更新）" async def test_upsert_nodes(service: QdrantService) -> None: nodes = [ { "doc_id": "doc-2", "section_path": "1", "text": "第一章大纲", "category": "tech", "tags": ["db"], "dense_vector": _dense(0.8), }, { "doc_id": "doc-2", "section_path": "2", "text": "第二章大纲", "category": "tech", "tags": ["db"], "dense_vector": _dense(0.7), }, ] await service.upsert_nodes(COLLECTION_L2, nodes) await service.upsert_nodes(COLLECTION_L3, nodes) for collection in (COLLECTION_L2, COLLECTION_L3): results = await service.search_dense( collection, _dense(0.8), limit=10, query_filter=QdrantService.build_filter(doc_ids=["doc-2"]), ) assert len(results) == 2 paths = {r.payload["section_path"] for r in results if r.payload} assert paths == {"1", "2"} # 非法集合应抛 ValueError with pytest.raises(ValueError): await service.upsert_nodes(COLLECTION_L1, nodes) async def test_upsert_chunks(service: QdrantService) -> None: chunks = [ { "doc_id": "doc-3", "chunk_index": 0, "text": "第一段原文", "section_path": "1", "title": "文档三", "category": "finance", "tags": ["stock"], "dense_vector": _dense(0.6), "sparse_vector": ([10, 20], [1.0, 0.8]), }, { "doc_id": "doc-3", "chunk_index": 1, "text": "第二段原文", "section_path": "2", "title": "文档三", "category": "finance", "tags": ["stock"], "dense_vector": _dense(0.4), }, ] await service.upsert_chunks(chunks) results = await service.search_dense( COLLECTION_CHUNKS, _dense(0.6), limit=10, query_filter=QdrantService.build_filter(doc_ids=["doc-3"]), ) assert len(results) == 2 indices = {r.payload["chunk_index"] for r in results if r.payload} assert indices == {0, 1} def test_build_filter_empty() -> None: assert QdrantService.build_filter() is None assert QdrantService.build_filter(categories=[], doc_ids=[], section_paths=[]) is None async def test_build_filter_categories(service: QdrantService) -> None: """categories 过滤：主类命中与标签命中的文档都能召回，无关类目被排除""" await service.upsert_l1("doc-cat", "主类命中", "总结", category="tech", tags=["x"], dense_vector=_dense(0.9)) await service.upsert_l1( "doc-tag", "标签命中", "总结", category="life", tags=["tech", "y"], dense_vector=_dense(0.8) ) await service.upsert_l1("doc-none", "无关文档", "总结", category="finance", tags=["z"], dense_vector=_dense(0.7)) query_filter = QdrantService.build_filter(categories=["tech"]) assert query_filter is not None results = await service.search_dense(COLLECTION_L1, _dense(0.9), limit=10, query_filter=query_filter) doc_ids = {r.payload["doc_id"] for r in results if r.payload} assert doc_ids == {"doc-cat", "doc-tag"} async def test_build_filter_doc_ids(service: QdrantService) -> None: await service.upsert_l1("doc-a", "A", "总结A", category="t", tags=[], dense_vector=_dense(0.9)) await service.upsert_l1("doc-b", "B", "总结B", category="t", tags=[], dense_vector=_dense(0.8)) results = await service.search_dense( COLLECTION_L1, _dense(0.9), limit=10, query_filter=QdrantService.build_filter(doc_ids=["doc-b"]), ) assert len(results) == 1 assert results[0].payload is not None assert results[0].payload["doc_id"] == "doc-b" async def test_search_hybrid_rrf(service: QdrantService) -> None: """hybrid 查询：dense + sparse 两路 prefetch 走服务端 RRF 融合，应正常返回结果""" await service.upsert_l1( doc_id="doc-h1", title="混合一", summary="混合检索文档一", category="tech", tags=["ai"], dense_vector=_dense(0.9), sparse_vector=([100, 200], [1.0, 0.5]), ) await service.upsert_l1( doc_id="doc-h2", title="混合二", summary="混合检索文档二", category="tech", tags=["ai"], dense_vector=_dense(0.3), sparse_vector=([100, 300], [0.9, 0.7]), ) results = await service.search_hybrid( COLLECTION_L1, dense_vector=_dense(0.9), sparse=([100, 200], [1.0, 0.5]), limit=5, ) assert len(results) >= 1 doc_ids = {r.payload["doc_id"] for r in results if r.payload} assert "doc-h1" in doc_ids # hybrid 带过滤也应正常工作 filtered = await service.search_hybrid( COLLECTION_L1, dense_vector=_dense(0.9), sparse=([100], [1.0]), limit=5, query_filter=QdrantService.build_filter(doc_ids=["doc-h2"]), ) assert len(filtered) == 1 assert filtered[0].payload is not None assert filtered[0].payload["doc_id"] == "doc-h2"