590 lines
19 KiB
Python
590 lines
19 KiB
Python
"""The load stage, exercised against the in-memory store with no server.
|
|
|
|
Two classes of failure are silent in a vector database and are what most of
|
|
these tests aim at. Loading the same corpus twice can leave two copies of a
|
|
dose, and every query still succeeds — so idempotency is asserted by point
|
|
count, not by inspecting the upsert calls. Mixing two corpus generations or two
|
|
models into one collection also raises nothing at query time; every search
|
|
returns *something*, just from the wrong material. That is what the manifest
|
|
gate exists to make loud, and there is a test per way it can be violated.
|
|
|
|
The last test runs over the real `chunks.jsonl` when it is present, because a
|
|
provenance rule that only holds for hand-written records is not evidence.
|
|
"""
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
from ingestion.load import (
|
|
ChunkLoader,
|
|
CollectionSpec,
|
|
CorpusManifest,
|
|
CorpusMismatch,
|
|
InMemoryVectorStore,
|
|
PointCountMismatch,
|
|
build_point,
|
|
corpus_sha256,
|
|
count_chunks,
|
|
iter_chunk_records,
|
|
manifest_collection,
|
|
point_id_for,
|
|
read_manifest,
|
|
validate_chunk_record,
|
|
)
|
|
|
|
DIMENSIONS = 4
|
|
COLLECTION = "duoc_thu_chunks"
|
|
CORPUS_SHA = "a" * 64
|
|
OTHER_SHA = "b" * 64
|
|
MODEL = "amazon.titan-embed-text-v2:0"
|
|
|
|
REAL_CHUNKS = (
|
|
Path(__file__).resolve().parents[1] / "data" / "processed" / "chunks.jsonl"
|
|
)
|
|
|
|
|
|
def chunk_record(chunk_id="abacavir__lieu_luong__0", **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 vector(seed=0.1):
|
|
return [seed] * DIMENSIONS
|
|
|
|
|
|
def manifest(**overrides):
|
|
values = {
|
|
"corpus_sha256": CORPUS_SHA,
|
|
"chunk_count": 3,
|
|
"model_id": MODEL,
|
|
"dimensions": DIMENSIONS,
|
|
"input_kind": "document",
|
|
"provider": "titan-v2",
|
|
}
|
|
values.update(overrides)
|
|
return CorpusManifest(**values)
|
|
|
|
|
|
def spec(**overrides):
|
|
values = {"name": COLLECTION, "vector_size": DIMENSIONS}
|
|
values.update(overrides)
|
|
return CollectionSpec(**values)
|
|
|
|
|
|
def loader(store, **overrides):
|
|
return ChunkLoader(
|
|
store,
|
|
overrides.pop("spec", spec()),
|
|
overrides.pop("manifest", manifest()),
|
|
**overrides,
|
|
)
|
|
|
|
|
|
def pairs(count=3):
|
|
return [
|
|
(chunk_record(chunk_id=f"drug__section__{i}"), vector(0.1 * (i + 1)))
|
|
for i in range(count)
|
|
]
|
|
|
|
|
|
# --- A5: derived ids and idempotency -------------------------------------
|
|
|
|
|
|
def test_point_id_is_derived_from_chunk_id_and_is_stable():
|
|
first = point_id_for("abacavir__lieu_luong__0")
|
|
second = point_id_for("abacavir__lieu_luong__0")
|
|
|
|
assert first == second
|
|
assert first != point_id_for("abacavir__lieu_luong__1")
|
|
|
|
|
|
def test_point_id_refuses_an_empty_chunk_id():
|
|
with pytest.raises(ValueError, match="chunk_id is required"):
|
|
point_id_for(" ")
|
|
|
|
|
|
def test_loading_the_same_corpus_twice_leaves_the_point_count_unchanged():
|
|
store = InMemoryVectorStore()
|
|
data = pairs(3)
|
|
|
|
first = loader(store).load(data)
|
|
second = loader(store).load(data)
|
|
|
|
assert first.collection_created is True
|
|
assert second.collection_created is False
|
|
assert first.collection_count == 3
|
|
assert second.collection_count == 3, "a re-run duplicated points"
|
|
assert second.points_upserted == 3
|
|
|
|
|
|
def test_a_reloaded_chunk_overwrites_its_own_point_rather_than_adding_one():
|
|
store = InMemoryVectorStore()
|
|
record = chunk_record()
|
|
loader(store).load([(record, vector(0.1))])
|
|
|
|
edited = chunk_record(text="Người lớn: 600 mg, một lần mỗi ngày.")
|
|
loader(store).load([(edited, vector(0.9))])
|
|
|
|
assert store.count(COLLECTION) == 1
|
|
stored = store.retrieve(COLLECTION, point_id_for(record["chunk_id"]))
|
|
assert stored.payload["text"] == "Người lớn: 600 mg, một lần mỗi ngày."
|
|
assert stored.vector == vector(0.9)
|
|
|
|
|
|
def test_records_are_upserted_in_batches_of_the_configured_size():
|
|
store = InMemoryVectorStore()
|
|
|
|
report = loader(store, batch_size=2).load(pairs(5))
|
|
|
|
assert report.batches == 3
|
|
assert report.points_upserted == 5
|
|
assert report.collection_count == 5
|
|
|
|
|
|
def test_batch_size_must_be_positive():
|
|
with pytest.raises(ValueError, match="batch_size must be positive"):
|
|
ChunkLoader(InMemoryVectorStore(), spec(), manifest(), batch_size=0)
|
|
|
|
|
|
# --- A4: collection shape and payload ------------------------------------
|
|
|
|
|
|
def test_collection_is_created_with_the_declared_size_and_payload_indexes():
|
|
store = InMemoryVectorStore()
|
|
|
|
loader(store).load(pairs(1))
|
|
|
|
assert store.spec(COLLECTION).vector_size == DIMENSIONS
|
|
assert store.spec(COLLECTION).distance == "Cosine"
|
|
indexed = dict(store.indexed_fields(COLLECTION))
|
|
assert indexed["drug_id"] == "keyword"
|
|
assert indexed["section_key"] == "keyword"
|
|
assert indexed["atc_codes"] == "keyword"
|
|
assert indexed["chunk_kind"] == "keyword"
|
|
assert indexed["has_quarantined_content"] == "bool"
|
|
|
|
|
|
def test_payload_carries_every_provenance_field_of_the_chunk_record():
|
|
store = InMemoryVectorStore()
|
|
record = chunk_record()
|
|
|
|
loader(store).load([(record, vector())])
|
|
|
|
payload = store.retrieve(COLLECTION, point_id_for(record["chunk_id"])).payload
|
|
assert payload == record
|
|
|
|
|
|
def test_a_field_added_by_a_future_chunker_flows_through_untouched():
|
|
store = InMemoryVectorStore()
|
|
record = chunk_record(
|
|
population_tags=["Người lớn", "Suy thận"], printed_page_range=[101, 103]
|
|
)
|
|
|
|
loader(store).load([(record, vector())])
|
|
|
|
payload = store.retrieve(COLLECTION, point_id_for(record["chunk_id"])).payload
|
|
assert payload["population_tags"] == ["Người lớn", "Suy thận"]
|
|
assert payload["printed_page_range"] == [101, 103]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"field",
|
|
[
|
|
"chunk_id",
|
|
"drug_id",
|
|
"section_key",
|
|
"source_page_range",
|
|
"printed_page_range",
|
|
"text",
|
|
],
|
|
)
|
|
def test_a_chunk_missing_a_required_provenance_field_is_refused(field):
|
|
with pytest.raises(ValueError, match="missing required provenance fields"):
|
|
validate_chunk_record(chunk_record(**{field: None}))
|
|
|
|
|
|
# --- failing closed on incomplete provenance ------------------------------
|
|
#
|
|
# Every case below passed an earlier version of this validator. The cost of
|
|
# that is not an exception at load time — it is paying for an embedding run and
|
|
# then discovering every answer abstains because the chunks cannot be cited.
|
|
|
|
|
|
@pytest.mark.parametrize("field", ["source_page_range", "printed_page_range"])
|
|
def test_an_empty_page_range_is_missing_not_present(field):
|
|
"""`[] in (None, "")` is False, which is exactly how this slipped through."""
|
|
with pytest.raises(ValueError, match="missing required provenance fields"):
|
|
validate_chunk_record(chunk_record(**{field: []}))
|
|
|
|
|
|
def test_an_unknown_old_or_future_schema_is_refused_fail_closed():
|
|
for version in (3, 5):
|
|
with pytest.raises(ValueError, match="supports exactly v4"):
|
|
validate_chunk_record(chunk_record(schema_version=version))
|
|
|
|
|
|
def test_a_chunk_with_no_schema_version_at_all_is_refused():
|
|
record = chunk_record()
|
|
del record["schema_version"]
|
|
with pytest.raises(ValueError, match="declares schema_version None"):
|
|
validate_chunk_record(record)
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"value", [[101], [101, 102, 103], "101-103", 101, {"start": 101}]
|
|
)
|
|
def test_a_page_range_that_is_not_a_pair_is_refused(value):
|
|
with pytest.raises(ValueError, match=r"expected a \[start, end\] pair"):
|
|
validate_chunk_record(chunk_record(printed_page_range=value))
|
|
|
|
|
|
def test_a_page_range_running_backwards_is_refused():
|
|
with pytest.raises(ValueError, match="running backwards"):
|
|
validate_chunk_record(chunk_record(printed_page_range=[103, 101]))
|
|
|
|
|
|
def test_a_non_integer_page_is_refused():
|
|
with pytest.raises(ValueError, match="non-integer page"):
|
|
validate_chunk_record(chunk_record(printed_page_range=[101.5, 103]))
|
|
|
|
|
|
def test_boolean_pages_are_not_accepted_as_python_integers():
|
|
with pytest.raises(ValueError, match="non-integer page"):
|
|
validate_chunk_record(chunk_record(printed_page_range=[False, True]))
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("field", "value"),
|
|
[("source_page_range", [-1, 0]), ("printed_page_range", [0, 1])],
|
|
)
|
|
def test_page_ranges_reject_impossible_lower_bounds(field, value):
|
|
with pytest.raises(ValueError, match="pages must start"):
|
|
validate_chunk_record(chunk_record(**{field: value}))
|
|
|
|
|
|
def test_page_zero_and_false_are_values_not_absences():
|
|
"""Physical pages are 0-indexed; a falsiness test would reject real records."""
|
|
validate_chunk_record(
|
|
chunk_record(
|
|
heading_physical_page=0,
|
|
source_page_range=[0, 0],
|
|
printed_page_range=[1, 1],
|
|
has_quarantined_content=False,
|
|
oversized=False,
|
|
part_index=0,
|
|
)
|
|
)
|
|
|
|
|
|
def test_a_wrong_sized_vector_is_refused_before_anything_is_upserted():
|
|
store = InMemoryVectorStore()
|
|
|
|
with pytest.raises(ValueError, match="expects 4"):
|
|
loader(store).load([(chunk_record(), [0.1, 0.2])])
|
|
|
|
assert store.count(COLLECTION) == 0
|
|
|
|
|
|
def test_manifest_dimensions_must_agree_with_the_collection_spec():
|
|
with pytest.raises(ValueError, match="manifest declares 8 dimensions"):
|
|
ChunkLoader(InMemoryVectorStore(), spec(), manifest(dimensions=8))
|
|
|
|
|
|
def test_collection_spec_refuses_a_nonpositive_vector_size():
|
|
with pytest.raises(ValueError, match="vector_size must be positive"):
|
|
CollectionSpec(name=COLLECTION, vector_size=0)
|
|
|
|
|
|
# --- A6: the corpus binding gate -----------------------------------------
|
|
|
|
|
|
def test_the_manifest_lives_beside_the_data_so_the_point_count_stays_exact():
|
|
store = InMemoryVectorStore()
|
|
|
|
loader(store).load(pairs(3))
|
|
|
|
assert store.count(COLLECTION) == 3, "the manifest must not inflate the count"
|
|
assert store.count(manifest_collection(COLLECTION)) == 1
|
|
stored = read_manifest(store, COLLECTION)
|
|
assert stored.corpus_sha256 == CORPUS_SHA
|
|
assert stored.model_id == MODEL
|
|
assert stored.provider == "titan-v2"
|
|
|
|
|
|
def test_a_second_corpus_generation_is_refused_and_nothing_is_written():
|
|
store = InMemoryVectorStore()
|
|
loader(store).load(pairs(3))
|
|
|
|
with pytest.raises(CorpusMismatch, match="does not match"):
|
|
loader(store, manifest=manifest(corpus_sha256=OTHER_SHA)).load(pairs(2))
|
|
|
|
assert store.count(COLLECTION) == 3
|
|
assert read_manifest(store, COLLECTION).corpus_sha256 == CORPUS_SHA
|
|
|
|
|
|
def test_a_second_model_is_refused_even_when_the_corpus_matches():
|
|
store = InMemoryVectorStore()
|
|
loader(store).load(pairs(1))
|
|
|
|
with pytest.raises(CorpusMismatch, match="cohere.embed-v4:0"):
|
|
loader(store, manifest=manifest(model_id="cohere.embed-v4:0")).load(pairs(1))
|
|
|
|
|
|
def test_a_query_subspace_vector_is_refused_for_a_document_collection():
|
|
store = InMemoryVectorStore()
|
|
loader(store).load(pairs(1))
|
|
|
|
with pytest.raises(CorpusMismatch, match="input kind query"):
|
|
loader(store, manifest=manifest(input_kind="query")).load(pairs(1))
|
|
|
|
|
|
def test_a_dimension_change_is_refused():
|
|
store = InMemoryVectorStore()
|
|
loader(store).load(pairs(1))
|
|
|
|
eight = CollectionSpec(name=COLLECTION, vector_size=8)
|
|
with pytest.raises(CorpusMismatch, match="8 dimensions"):
|
|
ChunkLoader(store, eight, manifest(dimensions=8)).load(
|
|
[(chunk_record(), [0.1] * 8)]
|
|
)
|
|
|
|
|
|
def test_an_existing_collection_with_no_manifest_is_refused():
|
|
store = InMemoryVectorStore()
|
|
store.create_collection(spec())
|
|
store.upsert(COLLECTION, [build_point(chunk_record(), vector())])
|
|
|
|
with pytest.raises(CorpusMismatch, match="has no manifest"):
|
|
loader(store).load(pairs(1))
|
|
|
|
|
|
def test_the_same_corpus_and_model_is_allowed_through():
|
|
store = InMemoryVectorStore()
|
|
loader(store).load(pairs(3))
|
|
|
|
report = loader(store).load(pairs(3))
|
|
|
|
assert report.collection_count == 3
|
|
assert report.count_matches is True
|
|
|
|
|
|
# --- the v1 point-count gate ---------------------------------------------
|
|
|
|
|
|
def test_assert_point_count_passes_when_every_chunk_has_exactly_one_point():
|
|
store = InMemoryVectorStore()
|
|
active = loader(store)
|
|
active.load(pairs(3))
|
|
|
|
assert active.assert_point_count(3) == 3
|
|
|
|
|
|
def test_assert_point_count_raises_when_the_collection_is_short():
|
|
store = InMemoryVectorStore()
|
|
active = loader(store)
|
|
active.load(pairs(3))
|
|
|
|
with pytest.raises(PointCountMismatch, match="holds 3 points but the corpus has 4"):
|
|
active.assert_point_count(4)
|
|
|
|
|
|
# --- mode A: exhaustive filter retrieval ---------------------------------
|
|
|
|
|
|
def section_pairs(drug_id, section_key, parts):
|
|
return [
|
|
(
|
|
chunk_record(
|
|
chunk_id=f"{drug_id}__{section_key}__{i}",
|
|
drug_id=drug_id,
|
|
section_key=section_key,
|
|
part_index=i,
|
|
part_count=parts,
|
|
),
|
|
vector(0.1),
|
|
)
|
|
for i in range(parts)
|
|
]
|
|
|
|
|
|
def test_a_filter_returns_every_part_of_a_section_not_a_top_k():
|
|
"""The rule mode A exists for: two of five contraindications is worse than none."""
|
|
store = InMemoryVectorStore()
|
|
loader(store).load(
|
|
section_pairs("metformin", "chong_chi_dinh", 5)
|
|
+ section_pairs("metformin", "lieu_luong_va_cach_dung", 3)
|
|
+ section_pairs("pantoprazol", "chong_chi_dinh", 2)
|
|
)
|
|
|
|
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"}
|
|
|
|
|
|
def test_a_filter_never_leaks_a_neighbouring_drugs_section():
|
|
store = InMemoryVectorStore()
|
|
loader(store).load(
|
|
section_pairs("pantoprazol", "chong_chi_dinh", 2)
|
|
+ section_pairs("omeprazol", "chong_chi_dinh", 2)
|
|
)
|
|
|
|
found = store.find_by_payload(
|
|
COLLECTION, {"drug_id": "pantoprazol", "section_key": "chong_chi_dinh"}
|
|
)
|
|
|
|
assert len(found) == 2
|
|
assert {p.payload["drug_id"] for p in found} == {"pantoprazol"}
|
|
|
|
|
|
def test_a_list_valued_field_matches_on_any_element():
|
|
store = InMemoryVectorStore()
|
|
record = chunk_record(atc_codes=["A10BA02", "A10BD20"])
|
|
loader(store).load([(record, vector())])
|
|
|
|
assert len(store.find_by_payload(COLLECTION, {"atc_codes": "A10BD20"})) == 1
|
|
assert len(store.find_by_payload(COLLECTION, {"atc_codes": "J05AF06"})) == 0
|
|
|
|
|
|
def test_a_filter_matching_nothing_returns_empty_rather_than_raising():
|
|
store = InMemoryVectorStore()
|
|
loader(store).load(pairs(2))
|
|
|
|
assert store.find_by_payload(COLLECTION, {"drug_id": "khong_ton_tai"}) == []
|
|
|
|
|
|
def test_a_filter_with_no_condition_is_refused():
|
|
store = InMemoryVectorStore()
|
|
loader(store).load(pairs(1))
|
|
|
|
with pytest.raises(ValueError, match="at least one condition"):
|
|
store.find_by_payload(COLLECTION, {})
|
|
|
|
|
|
def test_parts_reassemble_in_order_into_the_whole_section():
|
|
store = InMemoryVectorStore()
|
|
bodies = ["Phần một.", "Phần hai.", "Phần ba."]
|
|
records = [
|
|
(
|
|
chunk_record(
|
|
chunk_id=f"metformin__chong_chi_dinh__{i}",
|
|
drug_id="metformin",
|
|
section_key="chong_chi_dinh",
|
|
text=body,
|
|
part_index=i,
|
|
part_count=len(bodies),
|
|
),
|
|
vector(),
|
|
)
|
|
for i, body in enumerate(bodies)
|
|
]
|
|
loader(store).load(records)
|
|
|
|
found = store.find_by_payload(
|
|
COLLECTION, {"drug_id": "metformin", "section_key": "chong_chi_dinh"}
|
|
)
|
|
ordered = sorted(found, key=lambda p: p.payload["part_index"])
|
|
|
|
assert [p.payload["text"] for p in ordered] == bodies
|
|
assert {p.payload["part_count"] for p in found} == {3}
|
|
|
|
|
|
# --- corpus digest and reading -------------------------------------------
|
|
|
|
|
|
def test_corpus_sha256_changes_when_a_single_byte_changes(tmp_path):
|
|
path = tmp_path / "chunks.jsonl"
|
|
path.write_text(json.dumps(chunk_record()) + "\n", encoding="utf-8")
|
|
before = corpus_sha256(path)
|
|
|
|
path.write_text(
|
|
json.dumps(chunk_record(text="Người lớn: 301 mg.")) + "\n", encoding="utf-8"
|
|
)
|
|
|
|
assert corpus_sha256(path) != before
|
|
assert len(before) == 64
|
|
|
|
|
|
def test_the_same_data_hashes_the_same_under_crlf_and_lf(tmp_path):
|
|
"""A gate that cries wolf gets switched off.
|
|
|
|
A raw-byte digest made a Windows CRLF checkout and a Linux LF checkout of
|
|
identical data disagree, so A6 would refuse a CI load against the very
|
|
corpus it was built from.
|
|
"""
|
|
body = json.dumps(chunk_record()) + "\n" + json.dumps(chunk_record("b")) + "\n"
|
|
lf = tmp_path / "lf.jsonl"
|
|
crlf = tmp_path / "crlf.jsonl"
|
|
lf.write_bytes(body.encode("utf-8"))
|
|
crlf.write_bytes(body.replace("\n", "\r\n").encode("utf-8"))
|
|
|
|
assert corpus_sha256(lf) == corpus_sha256(crlf)
|
|
assert crlf.stat().st_size > lf.stat().st_size, "the files really do differ"
|
|
|
|
|
|
def test_blank_lines_are_skipped_and_bad_json_names_its_line(tmp_path):
|
|
path = tmp_path / "chunks.jsonl"
|
|
path.write_text(
|
|
json.dumps(chunk_record()) + "\n\n" + json.dumps(chunk_record("b")) + "\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
assert count_chunks(path) == 2
|
|
|
|
broken = tmp_path / "broken.jsonl"
|
|
broken.write_text(json.dumps(chunk_record()) + "\n{oops\n", encoding="utf-8")
|
|
with pytest.raises(ValueError, match="line 2 is not valid JSON"):
|
|
list(iter_chunk_records(broken))
|
|
|
|
|
|
# --- the real artifact ----------------------------------------------------
|
|
|
|
|
|
@pytest.mark.skipif(
|
|
not REAL_CHUNKS.exists(), reason="chunks.jsonl has not been generated"
|
|
)
|
|
def test_every_real_chunk_satisfies_the_loader_provenance_contract():
|
|
"""Whole-artifact scope: all records in `data/processed/chunks.jsonl`."""
|
|
seen_ids = set()
|
|
seen_points = set()
|
|
total = 0
|
|
for record in iter_chunk_records(REAL_CHUNKS):
|
|
validate_chunk_record(record)
|
|
point = point_id_for(record["chunk_id"])
|
|
assert point not in seen_points, f"point id collision on {record['chunk_id']}"
|
|
seen_points.add(point)
|
|
seen_ids.add(record["chunk_id"])
|
|
total += 1
|
|
|
|
assert total == len(seen_ids), "duplicate chunk_id in the artifact"
|
|
assert total == len(seen_points)
|
|
# A floor, not the exact count: the corpus is regenerated as `segment/`
|
|
# changes, but a truncated or half-written artifact must not pass as
|
|
# whole-artifact evidence. Measured 15,066 records on 2026-08-04.
|
|
assert total > 10_000, f"chunks.jsonl looks truncated: only {total} records"
|