378 lines
13 KiB
Python
378 lines
13 KiB
Python
"""`QdrantVectorStore` against a real Qdrant, skipped when none is running.
|
|
|
|
The rest of the load suite runs against `InMemoryVectorStore` and proves the
|
|
loader's rules. It cannot prove the adapter: whether Qdrant accepts a uuid5
|
|
string as a point id, whether `create_payload_index` takes a bare `"keyword"`,
|
|
whether an upsert of an existing id replaces rather than appends. Those are
|
|
claims about another system, and the same class of claim as the Bedrock request
|
|
bodies that are still documentation-derived and unproven — so they get a live
|
|
check, against a free local container rather than a paid API.
|
|
|
|
Start one with `docker compose -f infra/docker/docker-compose.yml up -d qdrant`.
|
|
Without it these tests skip; they never fail for being offline.
|
|
|
|
Every test works in its own collection named after the test and deletes it
|
|
afterwards, so a shared local Qdrant is not left holding fixtures.
|
|
"""
|
|
import os
|
|
import uuid
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from ingestion.load import (
|
|
ChunkLoader,
|
|
CollectionSpec,
|
|
CorpusManifest,
|
|
CorpusMismatch,
|
|
manifest_collection,
|
|
point_id_for,
|
|
read_manifest,
|
|
)
|
|
from ingestion.load.qdrant_repo import DEFAULT_URL, SCROLL_PAGE, QdrantVectorStore
|
|
|
|
DIMENSIONS = 4
|
|
QDRANT_URL = os.environ.get("QDRANT_URL", DEFAULT_URL)
|
|
REAL_CHUNKS = (
|
|
Path(__file__).resolve().parents[1] / "data" / "processed" / "chunks.jsonl"
|
|
)
|
|
|
|
pytestmark = pytest.mark.integration
|
|
|
|
|
|
def _server_is_up() -> bool:
|
|
try:
|
|
from qdrant_client import QdrantClient
|
|
except ImportError:
|
|
return False
|
|
try:
|
|
QdrantClient(url=QDRANT_URL, timeout=3.0).get_collections()
|
|
except Exception:
|
|
return False
|
|
return True
|
|
|
|
|
|
requires_qdrant = pytest.mark.skipif(
|
|
not _server_is_up(), reason=f"no Qdrant reachable at {QDRANT_URL}"
|
|
)
|
|
|
|
|
|
@pytest.fixture()
|
|
def store():
|
|
return QdrantVectorStore(url=QDRANT_URL, timeout=10.0)
|
|
|
|
|
|
@pytest.fixture()
|
|
def collection(store):
|
|
name = f"test_load_{uuid.uuid4().hex[:10]}"
|
|
yield name
|
|
for target in (manifest_collection(name), name):
|
|
if store.collection_exists(target):
|
|
store.delete_collection(target)
|
|
|
|
|
|
def chunk_record(chunk_id, **overrides):
|
|
record = {
|
|
"schema_version": 4,
|
|
"chunk_id": chunk_id,
|
|
"drug_id": "abacavir",
|
|
"drug_name": "ABACAVIR",
|
|
"section_key": "lieu_luong_va_cach_dung",
|
|
"section_display_name": "Liều lượng và cách dùng",
|
|
"text": "Người lớn: 300 mg, hai lần mỗi ngày.",
|
|
"source_text": "Người lớn: 300 mg, hai lần mỗi ngày.",
|
|
"heading_physical_page": 100,
|
|
"source_page_range": [100, 102],
|
|
"printed_page_range": [101, 103],
|
|
"atc_codes": ["J05AF06"],
|
|
"part_index": 0,
|
|
"part_count": 1,
|
|
"est_tokens": 14,
|
|
"oversized": False,
|
|
"chunk_kind": "prose",
|
|
"attachments": [],
|
|
"has_quarantined_content": False,
|
|
}
|
|
record.update(overrides)
|
|
return record
|
|
|
|
|
|
def manifest(**overrides):
|
|
values = {
|
|
"corpus_sha256": "a" * 64,
|
|
"chunk_count": 2,
|
|
"model_id": "amazon.titan-embed-text-v2:0",
|
|
"dimensions": DIMENSIONS,
|
|
"input_kind": "document",
|
|
"provider": "titan-v2",
|
|
}
|
|
values.update(overrides)
|
|
return CorpusManifest(**values)
|
|
|
|
|
|
def pairs(count):
|
|
return [
|
|
(chunk_record(f"abacavir__section__{i}"), [0.1 * (i + 1)] * DIMENSIONS)
|
|
for i in range(count)
|
|
]
|
|
|
|
|
|
@requires_qdrant
|
|
def test_a_real_load_creates_the_collection_indexes_and_points(store, collection):
|
|
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
|
|
|
report = ChunkLoader(store, spec, manifest()).load(pairs(2))
|
|
|
|
assert report.collection_created is True
|
|
assert report.collection_count == 2
|
|
assert store.collection_exists(collection)
|
|
assert store.count(collection) == 2
|
|
|
|
|
|
@requires_qdrant
|
|
def test_qdrant_accepts_the_derived_uuid5_point_id_and_returns_the_payload(
|
|
store, collection
|
|
):
|
|
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
|
record = chunk_record("abacavir__lieu_luong__0")
|
|
ChunkLoader(store, spec, manifest()).load([(record, [0.5] * DIMENSIONS)])
|
|
|
|
stored = store.retrieve(collection, point_id_for(record["chunk_id"]))
|
|
|
|
assert stored is not None
|
|
assert stored.id == point_id_for(record["chunk_id"])
|
|
assert stored.payload["chunk_id"] == record["chunk_id"]
|
|
assert stored.payload["source_page_range"] == [100, 102]
|
|
assert stored.payload["atc_codes"] == ["J05AF06"]
|
|
assert stored.payload["has_quarantined_content"] is False
|
|
assert len(stored.vector) == DIMENSIONS
|
|
|
|
|
|
@requires_qdrant
|
|
def test_loading_twice_against_a_real_server_does_not_duplicate(store, collection):
|
|
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
|
data = pairs(2)
|
|
|
|
ChunkLoader(store, spec, manifest()).load(data)
|
|
second = ChunkLoader(store, spec, manifest()).load(data)
|
|
|
|
assert second.collection_created is False
|
|
assert store.count(collection) == 2, "a re-run duplicated points in Qdrant"
|
|
|
|
|
|
@requires_qdrant
|
|
def test_the_manifest_sidecar_round_trips_and_leaves_the_count_exact(
|
|
store, collection
|
|
):
|
|
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
|
ChunkLoader(store, spec, manifest()).load(pairs(2))
|
|
|
|
stored = read_manifest(store, collection)
|
|
|
|
assert stored is not None
|
|
assert stored.corpus_sha256 == "a" * 64
|
|
assert stored.model_id == "amazon.titan-embed-text-v2:0"
|
|
assert stored.dimensions == DIMENSIONS
|
|
assert store.count(collection) == 2
|
|
assert store.count(manifest_collection(collection)) == 1
|
|
|
|
|
|
@requires_qdrant
|
|
def test_a_second_corpus_is_refused_against_a_real_collection(store, collection):
|
|
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
|
ChunkLoader(store, spec, manifest()).load(pairs(2))
|
|
|
|
with pytest.raises(CorpusMismatch, match="does not match"):
|
|
ChunkLoader(store, spec, manifest(corpus_sha256="b" * 64)).load(pairs(2))
|
|
|
|
assert store.count(collection) == 2
|
|
|
|
|
|
@requires_qdrant
|
|
def test_retrieve_returns_none_for_an_id_that_was_never_loaded(store, collection):
|
|
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
|
ChunkLoader(store, spec, manifest()).load(pairs(1))
|
|
|
|
assert store.retrieve(collection, point_id_for("never__loaded__0")) is None
|
|
|
|
|
|
@requires_qdrant
|
|
def test_bbox_floats_lose_precision_in_qdrant_but_nothing_else_does(store, collection):
|
|
"""Pins a measured round-trip loss so it cannot silently get worse.
|
|
|
|
Scrolling all 15,066 points of a full load on 2026-08-04 found 86 chunks
|
|
whose payload did not compare equal to its source record. Every one of the
|
|
96 differing leaf values was a float inside `attachments[].bbox`, the
|
|
largest delta was 5.684e-14, and **no** text, id, page number, page range,
|
|
token count or boolean differed at all. A PDF point is 1/72 inch, so that
|
|
delta cannot move a rendered crop; what would matter is the loss spreading
|
|
to another field, or growing. This test fails if either happens.
|
|
"""
|
|
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
|
# 17 significant digits: the real corpus carries these, and they are what
|
|
# does not survive a float64 -> JSON -> float64 round trip.
|
|
bbox = [44.45098876953125, 397.45245361328125, 278.09100341796875, 463.4044494628906]
|
|
record = chunk_record(
|
|
"cefazolin__lieu_luong_va_cach_dung__0",
|
|
has_quarantined_content=True,
|
|
attachments=[
|
|
{
|
|
"block_id": "p344_t2",
|
|
"kind": "table",
|
|
"shape": "simple_table",
|
|
"physical_page": 344,
|
|
"bbox": bbox,
|
|
"quarantined": True,
|
|
"header_row": ["Cỡ lọ", "Lượng\ndung môi"],
|
|
}
|
|
],
|
|
)
|
|
ChunkLoader(store, spec, manifest()).load([(record, [0.3] * DIMENSIONS)])
|
|
|
|
stored = store.retrieve(collection, point_id_for(record["chunk_id"])).payload
|
|
attachment = stored["attachments"][0]
|
|
|
|
for value, original in zip(attachment["bbox"], bbox, strict=True):
|
|
assert abs(value - original) < 1e-9, "bbox drift grew beyond rounding"
|
|
|
|
assert attachment["block_id"] == "p344_t2"
|
|
assert attachment["physical_page"] == 344
|
|
assert attachment["quarantined"] is True
|
|
assert attachment["header_row"] == ["Cỡ lọ", "Lượng\ndung môi"]
|
|
for field in (
|
|
"chunk_id",
|
|
"drug_id",
|
|
"drug_name",
|
|
"section_key",
|
|
"text",
|
|
"heading_physical_page",
|
|
"source_page_range",
|
|
"atc_codes",
|
|
"est_tokens",
|
|
"chunk_kind",
|
|
"has_quarantined_content",
|
|
):
|
|
assert stored[field] == record[field], f"{field} must round-trip exactly"
|
|
|
|
|
|
@requires_qdrant
|
|
def test_the_payload_index_actually_serves_a_filtered_query(store, collection):
|
|
"""Creating an index proves nothing; querying through it does.
|
|
|
|
Mode A of the delivery plan never ranks by vector — it filters on
|
|
`drug_id` + `section_key` and returns the whole section. Until this test
|
|
existed the loader had only established that `create_payload_index`
|
|
returned without error.
|
|
"""
|
|
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
|
data = (
|
|
_section("metformin", "chong_chi_dinh", 5)
|
|
+ _section("metformin", "lieu_luong_va_cach_dung", 3)
|
|
+ _section("pantoprazol", "chong_chi_dinh", 4)
|
|
)
|
|
ChunkLoader(store, spec, manifest()).load(data)
|
|
|
|
found = store.find_by_payload(
|
|
collection, {"drug_id": "metformin", "section_key": "chong_chi_dinh"}
|
|
)
|
|
|
|
assert len(found) == 5
|
|
assert sorted(p.payload["part_index"] for p in found) == [0, 1, 2, 3, 4]
|
|
assert {p.payload["drug_id"] for p in found} == {"metformin"}
|
|
|
|
|
|
@requires_qdrant
|
|
def test_a_section_longer_than_one_scroll_page_comes_back_whole(store, collection):
|
|
"""Paging must not truncate a section — that is the mode A failure mode.
|
|
|
|
Sized deliberately above `SCROLL_PAGE` (256) so a single-page implementation
|
|
fails here rather than in production on the one drug with a long section.
|
|
"""
|
|
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
|
parts = SCROLL_PAGE + 44
|
|
ChunkLoader(store, spec, manifest()).load(
|
|
_section("insulin", "lieu_luong_va_cach_dung", parts)
|
|
)
|
|
|
|
found = store.find_by_payload(
|
|
collection,
|
|
{"drug_id": "insulin", "section_key": "lieu_luong_va_cach_dung"},
|
|
)
|
|
|
|
assert len(found) == parts
|
|
assert sorted(p.payload["part_index"] for p in found) == list(range(parts))
|
|
|
|
|
|
@requires_qdrant
|
|
def test_a_list_valued_atc_field_matches_on_any_element_in_qdrant(store, collection):
|
|
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
|
record = chunk_record("metformin__lieu_luong__0", atc_codes=["A10BA02", "A10BD20"])
|
|
ChunkLoader(store, spec, manifest()).load([(record, [0.4] * DIMENSIONS)])
|
|
|
|
assert len(store.find_by_payload(collection, {"atc_codes": "A10BD20"})) == 1
|
|
assert len(store.find_by_payload(collection, {"atc_codes": "J05AF06"})) == 0
|
|
|
|
|
|
@requires_qdrant
|
|
@pytest.mark.skipif(
|
|
not REAL_CHUNKS.exists(), reason="chunks.jsonl has not been generated"
|
|
)
|
|
def test_a_real_multipart_section_round_trips_through_the_filter(store, collection):
|
|
"""Against the real artifact, not fixtures: every part, and only those."""
|
|
from ingestion.load import iter_chunk_records
|
|
|
|
wanted = None
|
|
records = []
|
|
for record in iter_chunk_records(REAL_CHUNKS):
|
|
if wanted is None and record["part_count"] >= 4:
|
|
wanted = (record["drug_id"], record["section_key"])
|
|
records.append(record)
|
|
assert wanted is not None, "no multi-part section in the artifact"
|
|
|
|
drug_id, section_key = wanted
|
|
expected = {
|
|
r["chunk_id"]
|
|
for r in records
|
|
if r["drug_id"] == drug_id and r["section_key"] == section_key
|
|
}
|
|
subset = [
|
|
r for r in records
|
|
if r["drug_id"] == drug_id or r["chunk_id"].startswith("abacavir")
|
|
]
|
|
|
|
spec = CollectionSpec(name=collection, vector_size=DIMENSIONS)
|
|
ChunkLoader(store, spec, manifest()).load(
|
|
(r, [0.2] * DIMENSIONS) for r in subset
|
|
)
|
|
|
|
found = store.find_by_payload(
|
|
collection, {"drug_id": drug_id, "section_key": section_key}
|
|
)
|
|
|
|
assert {p.payload["chunk_id"] for p in found} == expected
|
|
assert len(expected) >= 4
|
|
|
|
|
|
def _section(drug_id, section_key, parts):
|
|
return [
|
|
(
|
|
chunk_record(
|
|
f"{drug_id}__{section_key}__{i}",
|
|
drug_id=drug_id,
|
|
section_key=section_key,
|
|
part_index=i,
|
|
part_count=parts,
|
|
),
|
|
[0.1] * DIMENSIONS,
|
|
)
|
|
for i in range(parts)
|
|
]
|
|
|
|
|
|
@requires_qdrant
|
|
def test_an_unsupported_distance_is_rejected_before_the_server_is_called(store):
|
|
with pytest.raises(ValueError, match="unsupported distance"):
|
|
store.create_collection(
|
|
CollectionSpec(name="never_created", vector_size=4, distance="manhattan")
|
|
)
|