kplam
f92eff6f65
1. add `/api/v1/documents/{doc_id}/file` endpoint for downloading original document files
2. add file field to document detail API response based on metadata
3. add comprehensive tests for file download API and metadata integration
4. update API documentation in CLAUDE.md
386 lines
13 KiB
Python
386 lines
13 KiB
Python
"""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"
|
||
|
||
|
||
# ---------- L1 metadata 存储与读取 ----------
|
||
|
||
|
||
async def test_upsert_l1_metadata_roundtrip(service: QdrantService) -> None:
|
||
"""upsert_l1 写入 metadata 后，get_l1_metadata 返回写入的 dict"""
|
||
meta = {
|
||
"raw_file_path": "/data/uploads/report.pdf",
|
||
"original_filename": "年报.pdf",
|
||
"original_size_bytes": "2048",
|
||
}
|
||
await service.upsert_l1(
|
||
doc_id="doc-meta",
|
||
title="带元数据文档",
|
||
summary="总结",
|
||
category="tech",
|
||
tags=["x"],
|
||
dense_vector=_dense(0.9),
|
||
metadata=meta,
|
||
)
|
||
assert await service.get_l1_metadata("doc-meta") == meta
|
||
|
||
|
||
async def test_upsert_l1_without_metadata_stores_empty_dict(service: QdrantService) -> None:
|
||
"""upsert_l1 不传 metadata（旧签名）→ payload 存空 dict，get_l1_metadata 返回 {}"""
|
||
await service.upsert_l1(
|
||
doc_id="doc-no-meta",
|
||
title="无元数据文档",
|
||
summary="总结",
|
||
category="tech",
|
||
tags=["x"],
|
||
dense_vector=_dense(0.9),
|
||
)
|
||
assert await service.get_l1_metadata("doc-no-meta") == {}
|
||
|
||
|
||
async def test_get_l1_metadata_doc_not_found(service: QdrantService) -> None:
|
||
"""文档不存在 → get_l1_metadata 返回 None"""
|
||
assert await service.get_l1_metadata("doc-missing") is None
|
||
|
||
|
||
async def test_get_l1_metadata_legacy_payload_without_field(service: QdrantService) -> None:
|
||
"""存量旧文档：payload 完全没有 metadata 字段（绕过 upsert_l1 直接写点）→ 返回 None 不报错"""
|
||
import uuid
|
||
|
||
from qdrant_client import models
|
||
|
||
point = models.PointStruct(
|
||
id=str(uuid.uuid5(uuid.NAMESPACE_URL, "legacy-doc:l1")),
|
||
vector={"dense": _dense(0.4)},
|
||
payload={"doc_id": "legacy-doc", "title": "旧文档", "category": "tech", "tags": [], "text": "旧总结"},
|
||
)
|
||
await service.client.upsert(collection_name=COLLECTION_L1, points=[point])
|
||
assert await service.get_l1_metadata("legacy-doc") is None
|
||
|
||
|
||
# ---------- get_doc_detail file 字段 ----------
|
||
|
||
|
||
async def test_get_doc_detail_file_field_with_raw_path(service: QdrantService) -> None:
|
||
"""L1 metadata 含 raw_file_path → file 字段含 filename/size_bytes/url"""
|
||
await service.upsert_l1(
|
||
doc_id="doc-file",
|
||
title="文件文档",
|
||
summary="总结",
|
||
category="tech",
|
||
tags=["x"],
|
||
dense_vector=_dense(0.9),
|
||
metadata={
|
||
"raw_file_path": "/data/uploads/report.pdf",
|
||
"original_filename": "年报.pdf",
|
||
"original_size_bytes": "4096",
|
||
},
|
||
)
|
||
detail = await service.get_doc_detail("doc-file")
|
||
assert detail is not None
|
||
file_info = detail["file"]
|
||
assert file_info == {
|
||
"filename": "年报.pdf",
|
||
"size_bytes": 4096,
|
||
"url": "/api/v1/documents/doc-file/file",
|
||
}
|
||
|
||
|
||
async def test_get_doc_detail_file_field_filename_fallback_to_path_name(service: QdrantService) -> None:
|
||
"""metadata 缺 original_filename 时，filename 回退为 raw_file_path 的 Path.name"""
|
||
await service.upsert_l1(
|
||
doc_id="doc-file-fb",
|
||
title="文件文档",
|
||
summary="总结",
|
||
category="tech",
|
||
tags=["x"],
|
||
dense_vector=_dense(0.9),
|
||
metadata={"raw_file_path": "/data/uploads/notes.md", "original_size_bytes": "100"},
|
||
)
|
||
detail = await service.get_doc_detail("doc-file-fb")
|
||
assert detail is not None
|
||
assert detail["file"]["filename"] == "notes.md"
|
||
assert detail["file"]["size_bytes"] == 100
|
||
|
||
|
||
async def test_get_doc_detail_file_none_when_no_metadata(service: QdrantService) -> None:
|
||
"""无 metadata（文本入库，存空 dict）→ file=None"""
|
||
await service.upsert_l1(
|
||
doc_id="doc-text",
|
||
title="纯文本文档",
|
||
summary="总结",
|
||
category="tech",
|
||
tags=["x"],
|
||
dense_vector=_dense(0.9),
|
||
)
|
||
detail = await service.get_doc_detail("doc-text")
|
||
assert detail is not None
|
||
assert detail["file"] is None
|
||
|
||
|
||
async def test_get_doc_detail_file_none_when_no_raw_path(service: QdrantService) -> None:
|
||
"""metadata 无 raw_file_path → file=None"""
|
||
await service.upsert_l1(
|
||
doc_id="doc-meta-no-path",
|
||
title="文档",
|
||
summary="总结",
|
||
category="tech",
|
||
tags=["x"],
|
||
dense_vector=_dense(0.9),
|
||
metadata={"original_filename": "x.pdf", "original_size_bytes": "10"},
|
||
)
|
||
detail = await service.get_doc_detail("doc-meta-no-path")
|
||
assert detail is not None
|
||
assert detail["file"] is None
|
||
|
||
|
||
async def test_get_doc_detail_file_none_for_legacy_doc(service: QdrantService) -> None:
|
||
"""存量旧文档（payload 无 metadata 字段，绕过 upsert_l1 直接写点）→ file=None 不报错"""
|
||
import uuid
|
||
|
||
from qdrant_client import models
|
||
|
||
point = models.PointStruct(
|
||
id=str(uuid.uuid5(uuid.NAMESPACE_URL, "legacy-detail:l1")),
|
||
vector={"dense": _dense(0.3)},
|
||
payload={"doc_id": "legacy-detail", "title": "旧文档", "category": "tech", "tags": [], "text": "旧总结"},
|
||
)
|
||
await service.client.upsert(collection_name=COLLECTION_L1, points=[point])
|
||
detail = await service.get_doc_detail("legacy-detail")
|
||
assert detail is not None
|
||
assert detail["file"] is None
|
||
assert detail["l1"]["doc_id"] == "legacy-detail"
|