Wire the guarded conversational RAG answer layer end-to-end

This commit is contained in:
2026-08-05 14:33:13 +07:00
parent 834d9e51b0
commit ef08b4929e
127 changed files with 37921 additions and 169 deletions
@@ -0,0 +1,17 @@
{
"source_id": "dtqgvn_2_2018",
"title": "Dược thư Quốc gia Việt Nam - lần xuất bản thứ hai",
"edition": 2,
"publication_year": 2018,
"pdf_path": "data/raw/duoc-thu-quoc-gia-viet-nam-2018.pdf",
"sha256": "2aa81c846a5e760f82658c46816ab63174204a7d95b3e2755b6565288e53d0d1",
"superseded_by": {
"edition": 3,
"publication_year": 2022,
"decision": "3445/QĐ-BYT",
"decision_date": "2022-12-23"
},
"production_use_rights_documented": false,
"clinical_production_eligible": false,
"note": "Suitable for parser development and historical comparison only; not the sole source for clinical production."
}
File diff suppressed because it is too large Load Diff
+4 -1
View File
@@ -1,4 +1,5 @@
from .chunker import chunk_all, chunk_monograph, chunk_section, estimate_tokens
from .chunker import chunk_all, chunk_monograph, chunk_section
from .tokens import count_tokens, estimate_tokens, tokenizer_available
from .io import read_monographs_jsonl, write_chunks_jsonl
from .models import (
CHUNK_KIND_BLOCK_DESCRIPTOR,
@@ -17,7 +18,9 @@ __all__ = [
"chunk_all",
"chunk_monograph",
"chunk_section",
"count_tokens",
"estimate_tokens",
"tokenizer_available",
"read_monographs_jsonl",
"write_chunks_jsonl",
]
+388 -60
View File
@@ -7,10 +7,11 @@ sub-chunked, with a sentence-boundary-aware sliding window.
from __future__ import annotations
import re
from typing import Dict, Iterable, Iterator, List, Sequence
from dataclasses import dataclass
from typing import Dict, Iterable, Iterator, List, Mapping, Optional, Sequence
from ..segment.models import Monograph, SectionSpan, TableBlock
from ..tables.classify import SHAPE_FORMULA_2D, SHAPE_SIMPLE
from ..tables.classify import SHAPE_FORMULA_2D
from .models import (
CHUNK_KIND_BLOCK_DESCRIPTOR,
CHUNK_KIND_PROSE,
@@ -18,16 +19,12 @@ from .models import (
ChunkAttachment,
)
from .sentences import split_sentences
from .tokens import TokenCounter, count_tokens
CEILING_TOKENS = 800
TARGET_TOKENS = 650
OVERLAP_TOKENS = 65
# Physical -> printed page. Empirically constant across every tested
# milestone page (extract/page_map.py, ADR 0003); the descriptor quotes the
# printed number because that is what a reader holding the book looks for.
PRINTED_PAGE_OFFSET = 1
KIND_TABLE = "table"
KIND_FORMULA = "formula"
@@ -39,6 +36,15 @@ KIND_FORMULA = "formula"
# eye. A label carrying no digit cannot be mistaken for a dose; a long cell is
# content rather than a label.
_DIGIT = re.compile(r"\d")
_DOSE_VALUE = re.compile(
r"\d+(?:[.,]\d+)?\s*(?:mg|g|ml|microgam|mcg|µg|%|đơn vị|iu)\b",
re.IGNORECASE,
)
_SUBGROUP_LABEL = re.compile(
r"\b(?:trẻ|người lớn|người cao tuổi|sơ sinh|thiếu tháng|bệnh nhân|"
r"suy thận|suy gan|clcr|cân nặng)\b",
re.IGNORECASE,
)
HEADER_CELL_MAX_CHARS = 40
@@ -52,61 +58,293 @@ def _is_label_row(cells: Sequence[str]) -> bool:
)
def estimate_tokens(text: str) -> int:
"""ADR 0004's chars/4 estimate — an estimate, not a tokenizer count."""
return len(text) // 4
def _split_trailing_label(atom: str) -> List[str]:
"""Separate a label suffix from clinical text that precedes it.
The sentence splitter deliberately ends atoms at ``:``. In list-like dose
prose that can yield ``"7,5 mg ... .\nBước 5:"`` as one atom. Treating the
whole atom as the new label makes the dose at its beginning lose ``Bước 4``
when copied across a seam. Split only at an explicit newline, sentence, or
semicolon boundary and preserve every character exactly.
"""
stripped = atom.rstrip()
if not stripped.endswith(":"):
return [atom]
body = stripped[:-1]
cuts = []
newline = body.rfind("\n")
if newline >= 0:
cuts.append(newline + 1)
for marker in (". ", "; "):
position = body.rfind(marker)
if position >= 0:
cuts.append(position + len(marker))
if not cuts:
return [atom]
cut = max(cuts)
if not atom[cut:].strip():
return [atom]
return [atom[:cut], atom[cut:]]
def _pack(sentences: List[str]) -> List[List[str]]:
def _label_has_embedded_content(atom: str) -> bool:
"""Whether a colon-ending atom contains content before its final label."""
prefix = atom.rstrip()[:-1]
return (
"\n" in prefix
or ". " in prefix
or "; " in prefix
or bool(_DOSE_VALUE.search(prefix))
)
def _is_subgroup_label(atom: str) -> bool:
"""Population/organ-function labels that can sit under a route/indication."""
return bool(_SUBGROUP_LABEL.search(atom))
def _atoms(text: str, measure: TokenCounter) -> List[str]:
"""Smallest units the packer may not split.
Normally a sentence. A drug-interaction list, though, is one "sentence"
hundreds of names long: VORICONAZOL's `tương tác thuốc` produced two parts
of 981 and 888 tokens even after sentence packing. Left that size the
embedding truncates them, and a truncated interaction list reads as "this
drug is not listed" — a false negative in exactly the direction that
matters. Such a run is comma-separated by construction, so a comma is a
lossless place to break it.
"""
atoms: List[str] = []
sentences = [
atom
for sentence in split_sentences(text)
for atom in _split_trailing_label(sentence)
]
for sentence in sentences:
# Split against the packing target, not the ceiling: an atom sized
# right up to the ceiling leaves no room for the overlap prepended to
# it, which is how a 710-token atom became a 981-token chunk.
if measure(sentence) <= TARGET_TOKENS or "," not in sentence:
atoms.append(sentence)
continue
piece = ""
for fragment in sentence.split(","):
candidate = f"{piece},{fragment}" if piece else fragment
if piece and measure(candidate) > TARGET_TOKENS:
atoms.append(piece + ",")
piece = fragment
else:
piece = candidate
if piece:
atoms.append(piece)
return atoms
@dataclass(frozen=True)
class _PackedPart:
atoms: List[str]
text: str
source_text: str
context_labels: List[str]
def _pack_parts(sentences: List[str], measure: TokenCounter) -> List[_PackedPart]:
"""Greedily pack sentences up to TARGET_TOKENS, overlapping by OVERLAP_TOKENS.
A single sentence longer than the target becomes its own part rather than
being cut mid-sentence — the caller flags it instead of splitting it.
"""
parts: List[List[str]] = []
current: List[str] = []
IndexedAtom = tuple[int, str]
PackedAtom = tuple[int, str, bool] # index, text, retrieval-only context
parts: List[List[PackedAtom]] = []
current: List[PackedAtom] = []
current_tokens = 0
for sentence in sentences:
tokens = estimate_tokens(sentence)
def is_label(atom: str) -> bool:
return atom.rstrip().endswith(":")
def contains(items: Sequence[PackedAtom], atom: IndexedAtom) -> bool:
return any((index, text) == atom for index, text, _ in items)
# Record the active label for every atom before packing. Looking only inside
# the current window is insufficient: a population may span several parts,
# so its label can have fallen outside both the 650-token buffer and the
# 65-token overlap by the time another dose reaches a seam.
contexts: List[IndexedAtom | None] = []
scope_contexts: List[IndexedAtom | None] = []
active_label: IndexedAtom | None = None
scope_label: IndexedAtom | None = None
for index, atom in enumerate(sentences):
# Context at the *start* of an atom comes from the preceding label. An
# atom may contain a dose and only end with the next label (Bisoprolol:
# "7,5 mg ... Bước 5:"); assigning that atom to its own final label
# makes the dose at its beginning context-free.
contexts.append(active_label)
next_scope = scope_label
if is_label(atom):
if _is_subgroup_label(atom):
if active_label is not None and not _is_subgroup_label(active_label[1]):
next_scope = active_label
else:
next_scope = None
# A pure subgroup label needs its parent scope when it itself lands at
# the start of a continuation, not only when the following dose arrives.
scope_contexts.append(next_scope)
if is_label(atom):
scope_label = next_scope
active_label = (index, atom)
def context_chain(index: int) -> List[IndexedAtom]:
"""Labels needed to make atom ``index`` independently interpretable.
Some malformed sentence atoms contain a dose and only end with the next
label. In that case the atom itself is context for what follows, but it
still needs its own preceding label. Walk only until a plain label;
this keeps the chain clinically complete without dragging the entire
section into every overlap.
"""
chain: List[IndexedAtom] = []
def add(atom: IndexedAtom | None) -> None:
if atom is not None and atom not in chain:
chain.append(atom)
add(scope_contexts[index])
target = sentences[index]
# A plain label introduces new sibling context, so it needs only its
# retained parent scope. A compound colon-ending atom still contains
# clinical material before that new label and therefore needs the
# preceding active label as well.
if is_label(target) and not _label_has_embedded_content(target):
return chain
nested: List[IndexedAtom] = []
context = contexts[index]
while context is not None and context not in nested:
nested.append(context)
add(scope_contexts[context[0]])
if not _label_has_embedded_content(context[1]):
break
context = contexts[context[0]]
for atom in reversed(nested):
add(atom)
return chain
for index, sentence in enumerate(sentences):
tokens = measure(sentence)
if current and current_tokens + tokens > TARGET_TOKENS:
parts.append(current)
overlap: List[str] = []
acc = 0
for prev in reversed(current):
overlap.insert(0, prev)
acc += estimate_tokens(prev)
if acc >= OVERLAP_TOKENS:
break
current = list(overlap)
current_tokens = sum(estimate_tokens(s) for s in current)
current.append(sentence)
# Never end a part on a label. `split_sentences` treats ':' as a
# boundary, so "Người lớn: 500 mg mỗi 8 giờ." splits after the
# colon; flushing there leaves a chunk ending "Người lớn:" with
# the dose in the next one. Measured before this rule: 38 chunks,
# AMOXICILIN's ending on a Lyme-disease indication followed by a
# bare "Người lớn:". Outlier item 17 counted population markers on
# 1,121 of ~1,400 monograph pages, so this is the common shape,
# and a dose separated from the population it applies to is a
# patient-safety defect rather than a cosmetic one.
carried: List[PackedAtom] = []
while current and is_label(current[-1][1]):
carried.insert(0, current.pop())
if not current:
# A label is safety context, not a useful standalone retrieval
# unit. Keep it with the following atom even when that makes a
# synthetic pathological atom oversized; the caller will flag
# the oversize instead of publishing an unqualified dose.
current = carried
current_tokens = sum(measure(atom) for _, atom, _ in current)
else:
parts.append(current)
overlap: List[PackedAtom] = []
acc = 0
for prev_index, prev, is_context in reversed(current):
# Stop *before* exceeding the budget. A governing label
# longer than the budget is retained by itself because
# clinical context wins over the overlap target.
size = measure(prev)
if acc + size > OVERLAP_TOKENS:
break
overlap.insert(0, (prev_index, prev, is_context))
acc += size
# If dose/detail atoms are copied, their own governing label
# must precede them even when a newer trailing label is being
# carried for the incoming sentence.
if overlap:
prefix = [
(context_index, context_text, True)
for context_index, context_text in context_chain(overlap[0][0])
if not contains(overlap, (context_index, context_text))
]
overlap = prefix + overlap
current = overlap + carried
current_tokens = sum(measure(atom) for _, atom, _ in current)
# The incoming atom may belong to a label that was flushed several
# chunks ago. Repeat that label before the new material. Do not add
# it twice when the incoming atom is itself the label or it was
# already carried/overlapped.
incoming_scope = scope_contexts[index]
for incoming_context in context_chain(index):
if not contains(current, incoming_context):
packed_context = (*incoming_context, True)
if incoming_context == incoming_scope:
current.insert(0, packed_context)
else:
current.append(packed_context)
current_tokens += measure(incoming_context[1])
current.append((index, sentence, False))
current_tokens += tokens
if current:
parts.append(current)
return parts
return [
_PackedPart(
atoms=[atom for _, atom, _ in part],
text="".join(atom for _, atom, _ in part).strip(),
source_text="".join(
atom for _, atom, is_context in part if not is_context
).strip(),
context_labels=[
atom.strip() for _, atom, is_context in part if is_context
],
)
for part in parts
]
def _pack(sentences: List[str], measure: TokenCounter) -> List[List[str]]:
"""Compatibility surface used by focused packer tests."""
return [part.atoms for part in _pack_parts(sentences, measure)]
def _block_kind(block: TableBlock) -> str:
return KIND_FORMULA if block.shape == SHAPE_FORMULA_2D else KIND_TABLE
def _attachment(block: TableBlock, header_row: List[str]) -> ChunkAttachment:
def _attachment(
block: TableBlock,
_header_row: List[str],
printed_page: int | None = None,
) -> ChunkAttachment:
return ChunkAttachment(
block_id=block.table_id,
kind=_block_kind(block),
shape=block.shape,
physical_page=block.physical_page,
bbox=list(block.bbox),
printed_page=printed_page,
quarantined=block.quarantined,
# Only a simple table's first row can be a row of plain labels, and
# only when it actually reads like one. A multi-level or merged header
# is the shape whose extraction is least trustworthy, so it
# contributes nothing rather than something wrong.
header_row=(list(header_row)
if block.shape == SHAPE_SIMPLE and _is_label_row(header_row)
else []),
# Header extraction is not human-verified and continuation tables can
# begin with a body row. Two ARSENIC TRIOXYD ADR rows were previously
# shipped as "Cột:" metadata. Keep every cell value out of both the
# retrieval text and vector payload until a reviewed logical-table
# artifact can prove which row is a header.
header_row=[],
)
@@ -119,7 +357,8 @@ def _blocks_by_section(monograph: Monograph) -> Dict[str, List[TableBlock]]:
def describe_block(monograph: Monograph, section: SectionSpan,
attachment: ChunkAttachment) -> str:
attachment: ChunkAttachment,
printed_page: int | None = None) -> str:
"""Retrieval text for a block, built only from metadata.
No cell value ever appears here. A header row is a row of labels;
@@ -127,28 +366,91 @@ def describe_block(monograph: Monograph, section: SectionSpan,
linearising a body row does.
"""
noun = "công thức" if attachment.kind == KIND_FORMULA else "bảng"
printed = attachment.physical_page + PRINTED_PAGE_OFFSET
page_label = (f"trang {printed_page}" if printed_page is not None
else "chưa xác định trang in")
text = (f"{monograph.drug_name}{section.display_name}{noun}, "
f"trang {printed}.")
if attachment.header_row:
columns = " | ".join(c.replace("\n", " ").strip()
for c in attachment.header_row if c and c.strip())
if columns:
text += f" Cột: {columns}."
f"{page_label}.")
text += (" Nội dung chỉ tra cứu được trên ảnh trang gốc, "
"không trích dẫn được dưới dạng văn bản.")
return text
def _supporting_pages(section: SectionSpan, source_text: str) -> List[int]:
"""Physical pages whose prose parts intersect one contiguous chunk span."""
section_text = section.text.strip()
if not source_text or section_text.count(source_text) != 1:
raise ValueError(
f"cannot map {section.key!r} chunk source text uniquely to its section"
)
chunk_start = section_text.index(source_text)
chunk_end = chunk_start + len(source_text)
pages: List[int] = []
cursor = 0
for part in section.parts:
if part.kind != "prose" or part.quarantined or not part.text:
continue
part_start = section_text.find(part.text, cursor)
if part_start < 0:
raise ValueError(
f"cannot map {section.key!r} part on physical page "
f"{part.physical_page} back to section text"
)
part_end = part_start + len(part.text)
cursor = part_end
if part_start < chunk_end and part_end > chunk_start:
pages.append(part.physical_page)
if not pages:
raise ValueError(f"no page provenance supports section {section.key!r} chunk")
return sorted(set(pages))
def _page_ranges(
physical_pages: Sequence[int],
printed_page_map: Mapping[int, Optional[int]] | None,
label: str,
) -> tuple[List[int], List[int]]:
physical_range = [min(physical_pages), max(physical_pages)]
if printed_page_map is None:
return physical_range, []
printed_pages = [printed_page_map.get(page) for page in physical_pages]
if any(page is None for page in printed_pages):
missing = [physical for physical, printed in zip(physical_pages, printed_pages, strict=True)
if printed is None]
raise ValueError(
f"cannot cite {label}: printed folio missing for physical pages {missing}"
)
verified = [page for page in printed_pages if page is not None]
return physical_range, [min(verified), max(verified)]
def chunk_section(monograph: Monograph, section: SectionSpan,
blocks: Sequence[TableBlock] = (),
header_rows: Dict[str, List[str]] | None = None) -> List[Chunk]:
header_rows: Dict[str, List[str]] | None = None,
printed_page_map: Mapping[int, Optional[int]] | None = None,
measure: TokenCounter = count_tokens) -> List[Chunk]:
header_rows = header_rows or {}
attachments = [_attachment(b, header_rows.get(b.table_id, [])) for b in blocks]
attachments = []
for block in blocks:
_, attachment_printed_range = _page_ranges(
[block.physical_page], printed_page_map,
f"{monograph.drug_id}/{block.table_id}",
)
attachments.append(_attachment(
block,
header_rows.get(block.table_id, []),
attachment_printed_range[0] if attachment_printed_range else None,
))
quarantined = any(a.quarantined for a in attachments)
def build(body: str, part_index: int, part_count: int) -> Chunk:
tokens = estimate_tokens(body)
def build(body: str, source_body: str, context_labels: List[str],
part_index: int, part_count: int) -> Chunk:
tokens = measure(body)
source_page_range, printed_page_range = _page_ranges(
_supporting_pages(section, source_body), printed_page_map,
f"{monograph.drug_id}/{section.key}/{part_index}",
)
return Chunk(
chunk_id=f"{monograph.drug_id}__{section.key}__{part_index}",
drug_id=monograph.drug_id,
@@ -156,8 +458,11 @@ def chunk_section(monograph: Monograph, section: SectionSpan,
section_key=section.key,
section_display_name=section.display_name,
text=body,
source_text=source_body,
context_labels=list(context_labels),
heading_physical_page=section.heading.physical_page,
source_page_range=list(monograph.source_page_range),
source_page_range=source_page_range,
printed_page_range=list(printed_page_range),
atc_codes=list(monograph.atc_codes),
part_index=part_index,
part_count=part_count,
@@ -171,16 +476,26 @@ def chunk_section(monograph: Monograph, section: SectionSpan,
text = section.text.strip()
prose: List[Chunk] = []
if text:
if estimate_tokens(text) <= CEILING_TOKENS:
prose = [build(text, 0, 1)]
if measure(text) <= CEILING_TOKENS:
prose = [build(text, text, [], 0, 1)]
else:
parts = _pack(split_sentences(text))
bodies = [b for b in ("".join(p).strip() for p in parts) if b]
prose = [build(b, i, len(bodies)) for i, b in enumerate(bodies)]
parts = [part for part in _pack_parts(_atoms(text, measure), measure)
if part.text]
prose = [
build(part.text, part.source_text, part.context_labels,
index, len(parts))
for index, part in enumerate(parts)
]
descriptors = []
for attachment in attachments:
body = describe_block(monograph, section, attachment)
source_page_range, attachment_printed_range = _page_ranges(
[attachment.physical_page], printed_page_map,
f"{monograph.drug_id}/{attachment.block_id}",
)
printed_page = (attachment_printed_range[0]
if attachment_printed_range else None)
body = describe_block(monograph, section, attachment, printed_page)
descriptors.append(Chunk(
chunk_id=f"{monograph.drug_id}__{section.key}__block__{attachment.block_id}",
drug_id=monograph.drug_id,
@@ -188,10 +503,12 @@ def chunk_section(monograph: Monograph, section: SectionSpan,
section_key=section.key,
section_display_name=section.display_name,
text=body,
source_text=body,
heading_physical_page=section.heading.physical_page,
source_page_range=list(monograph.source_page_range),
source_page_range=source_page_range,
printed_page_range=attachment_printed_range,
atc_codes=list(monograph.atc_codes),
est_tokens=estimate_tokens(body),
est_tokens=measure(body),
chunk_kind=CHUNK_KIND_BLOCK_DESCRIPTOR,
attachments=[attachment],
has_quarantined_content=attachment.quarantined,
@@ -200,16 +517,27 @@ def chunk_section(monograph: Monograph, section: SectionSpan,
def chunk_monograph(monograph: Monograph,
header_rows: Dict[str, List[str]] | None = None) -> List[Chunk]:
header_rows: Dict[str, List[str]] | None = None,
printed_page_map: Mapping[int, Optional[int]] | None = None,
measure: TokenCounter = count_tokens) -> List[Chunk]:
grouped = _blocks_by_section(monograph)
chunks: List[Chunk] = []
for section in monograph.sections.values():
chunks.extend(chunk_section(monograph, section,
grouped.get(section.key, ()), header_rows))
grouped.get(section.key, ()), header_rows,
printed_page_map, measure))
return chunks
def chunk_all(monographs: Iterable[Monograph],
header_rows: Dict[str, List[str]] | None = None) -> Iterator[Chunk]:
header_rows: Dict[str, List[str]] | None = None,
*,
printed_page_map: Mapping[int, Optional[int]] | None,
measure: TokenCounter = count_tokens) -> Iterator[Chunk]:
if printed_page_map is None:
raise ValueError(
"chunk_all requires a verified printed_page_map; refusing to emit "
"an embedding corpus without printed-page provenance"
)
for monograph in monographs:
yield from chunk_monograph(monograph, header_rows)
yield from chunk_monograph(monograph, header_rows, printed_page_map, measure)
+24 -1
View File
@@ -6,7 +6,13 @@ from dataclasses import asdict
from pathlib import Path
from typing import Iterable, Iterator
from ..segment.models import Heading, Monograph, SectionSpan, TableBlock
from ..segment.models import (
Heading,
Monograph,
SectionPart,
SectionSpan,
TableBlock,
)
from .models import SCHEMA_VERSION, Chunk
@@ -31,6 +37,23 @@ def read_monographs_jsonl(path: Path) -> Iterator[Monograph]:
section_key=h.get("section_key"),
),
text=s["text"],
# Carried, not dropped: CLAUDE.md's provenance rule is
# explicit that a stage boundary must not shed fields.
parts=[
SectionPart(
kind=p["kind"],
text=p["text"],
physical_page=p["physical_page"],
bbox=p["bbox"],
source_span_ids=p.get("source_span_ids", []),
table_id=p.get("table_id"),
table_part_id=p.get("table_part_id"),
continuation_group=p.get("continuation_group"),
shape=p.get("shape"),
quarantined=p.get("quarantined", False),
)
for p in s.get("parts", [])
],
)
yield Monograph(
drug_id=raw["drug_id"],
+9 -1
View File
@@ -18,7 +18,7 @@ from __future__ import annotations
from dataclasses import dataclass, field
from typing import List
SCHEMA_VERSION = 2
SCHEMA_VERSION = 4
CHUNK_KIND_PROSE = "prose"
CHUNK_KIND_BLOCK_DESCRIPTOR = "block_descriptor"
@@ -37,6 +37,8 @@ class ChunkAttachment:
shape: str
physical_page: int
bbox: List[float]
printed_page: int | None = None
source_crop: str | None = None
quarantined: bool = True
# First row of a `simple_table`, used to make the block findable. Comes
# from pdfplumber and has NOT been verified by eye — the 180 real tables'
@@ -54,6 +56,12 @@ class Chunk:
text: str
heading_physical_page: int
source_page_range: List[int]
# Exact contiguous source material. `text` may prepend retrieval-only
# context labels at a seam; provenance/reassembly must never mistake those
# repetitions for a literal source span.
source_text: str = ""
context_labels: List[str] = field(default_factory=list)
printed_page_range: List[int] = field(default_factory=list)
atc_codes: List[str] = field(default_factory=list)
part_index: int = 0
part_count: int = 1
+65
View File
@@ -0,0 +1,65 @@
"""Token counting for the chunk stage.
ADR 0004 sized chunks with `len(text) // 4`, describing it honestly as an
estimate. Measured against the real tokenizer on this corpus, that estimate is
wrong by about a factor of two for Vietnamese: real/estimate is **1.95 at the
median, 2.50 at p95, 6.0 at worst**, because accented Vietnamese characters
cost multiple byte-pair tokens each where English prose costs about four
characters per token.
The consequence was not academic. Under the estimate the pipeline reported
**0 chunks over the 800-token ceiling**; counted properly, **1,884 of 12,838
(14.7%)** were over it, the largest at 1,645 tokens — twice the ceiling. A
reassuring number that was simply false.
The counter is injectable so the chunking logic stays testable without the
tokenizer installed, and so a different embedding model's tokenizer can be
substituted without touching chunk shaping.
"""
from __future__ import annotations
from typing import Callable
# What OpenAI's text-embedding-3-* models use. Named here rather than inside
# the function so swapping models is a one-line, visible change.
ENCODING_NAME = "cl100k_base"
# Only for the no-tokenizer fallback. Derived from the measurement above
# (median 1.95 real tokens per chars/4 unit), i.e. ~2 characters per token —
# still an estimate, but one that errs on the side of smaller chunks instead
# of larger ones.
FALLBACK_CHARS_PER_TOKEN = 2
TokenCounter = Callable[[str], int]
_encoder = None
def _load_encoder():
global _encoder
if _encoder is None:
import tiktoken
_encoder = tiktoken.get_encoding(ENCODING_NAME)
return _encoder
def count_tokens(text: str) -> int:
"""Real token count, falling back to a conservative estimate."""
try:
return len(_load_encoder().encode(text))
except Exception:
return estimate_tokens(text)
def estimate_tokens(text: str) -> int:
"""Character-ratio fallback. An estimate — never report it as a count."""
return len(text) // FALLBACK_CHARS_PER_TOKEN
def tokenizer_available() -> bool:
try:
_load_encoder()
return True
except Exception:
return False
+21 -4
View File
@@ -16,6 +16,7 @@ from pathlib import Path
import fitz
from .extract import (
build_page_map,
extract_spans,
load_transcribed_runs,
merge_outlined_runs,
@@ -26,6 +27,7 @@ from .extract import (
)
from .chunk import (
CHUNK_KIND_PROSE,
tokenizer_available,
chunk_all,
read_monographs_jsonl,
write_chunks_jsonl,
@@ -151,9 +153,12 @@ def _cmd_validate(args: argparse.Namespace) -> int:
return 1
doc = fitz.open(pdf_path)
spans = list(extract_spans(doc))
# Same stream and same regions as `run`. Measuring recall against a
# pipeline that is not the one producing the output is how `coverage`
# ended up describing a different build earlier today.
spans = _extracted_and_repaired_spans(doc)
try:
monographs = list(assemble(spans))
monographs = list(assemble(spans, table_index=_region_index(args.tables)))
except DuplicateDrugIdError as e:
print(f"error: {e}", file=sys.stderr)
return 1
@@ -332,8 +337,17 @@ def _cmd_chunk(args: argparse.Namespace) -> int:
header_rows = {r.table_id: r.first_row
for r in read_regions_json(regions_path)}
pdf_path = Path(args.pdf)
if not pdf_path.exists():
print(f"error: no source PDF at {pdf_path}", file=sys.stderr)
return 1
with fitz.open(pdf_path) as document:
printed_page_map = build_page_map(document)
monographs = list(read_monographs_jsonl(monographs_path))
chunks = list(chunk_all(monographs, header_rows))
chunks = list(chunk_all(
monographs, header_rows, printed_page_map=printed_page_map,
))
kinds = Counter(c.chunk_kind for c in chunks)
with_attachments = sum(1 for c in chunks
@@ -347,7 +361,7 @@ def _cmd_chunk(args: argparse.Namespace) -> int:
print(f" {kind:<22}{n:>7}")
print(f"prose chunks carrying a lifted block: {with_attachments}")
print(f"oversized (over the {800}-token ceiling): {oversized}")
print(f"estimated tokens (chars/4, an estimate): {tokens:,}")
print(f"tokens ({'cl100k_base' if tokenizer_available() else 'ESTIMATED — tokenizer missing'}): {tokens:,}")
out_path = Path(args.out)
written = write_chunks_jsonl(chunks, out_path)
@@ -444,6 +458,7 @@ def build_parser() -> argparse.ArgumentParser:
p_validate = sub.add_parser("validate", help="Whole-book recall/precision vs. back-of-book index")
p_validate.add_argument("--pdf", required=True)
p_validate.add_argument("--tables", default="data/processed/table_regions.json")
p_validate.set_defaults(func=_cmd_validate)
p_tables = sub.add_parser(
@@ -491,6 +506,8 @@ def build_parser() -> argparse.ArgumentParser:
"chunk", help="Build retrieval chunks (ADR 0004/0005/0006)")
p_chunk.add_argument("--monographs", default="data/processed/monographs.jsonl")
p_chunk.add_argument("--tables", default="data/processed/table_regions.json")
p_chunk.add_argument(
"--pdf", default="data/raw/duoc-thu-quoc-gia-viet-nam-2018.pdf")
p_chunk.add_argument("--out", default="data/processed/chunks.jsonl")
p_chunk.set_defaults(func=_cmd_chunk)
+54
View File
@@ -0,0 +1,54 @@
"""Embedding stage: text in, vectors out, with the model behind an interface.
Import from here rather than from a provider module. Nothing outside this
package should name boto3, `sentence_transformers`, or a model id — that is
what lets `load/` and the retrieval side stay testable without a live account,
and what makes swapping the benchmark winner a one-line change.
"""
from .bedrock_runtime import DEFAULT_REGION, BedrockInvoker, Boto3BedrockInvoker
from .cache import (
CacheStats,
CachingEmbeddingProvider,
EmbeddingCache,
cache_key,
)
from .ports import (
INPUT_DOCUMENT,
INPUT_KINDS,
INPUT_QUERY,
EmbeddingBatch,
EmbeddingProvider,
EmbeddingVector,
text_digest,
)
from .registry import (
BGE_M3,
CLOUD_PROVIDERS,
COHERE_V4,
TITAN_V2,
build_provider,
provider_names,
)
__all__ = [
"BGE_M3",
"CLOUD_PROVIDERS",
"COHERE_V4",
"DEFAULT_REGION",
"INPUT_DOCUMENT",
"INPUT_KINDS",
"INPUT_QUERY",
"TITAN_V2",
"BedrockInvoker",
"Boto3BedrockInvoker",
"CacheStats",
"CachingEmbeddingProvider",
"EmbeddingBatch",
"EmbeddingCache",
"EmbeddingProvider",
"EmbeddingVector",
"build_provider",
"cache_key",
"provider_names",
"text_digest",
]
+152
View File
@@ -0,0 +1,152 @@
"""Cohere Embed v4 (`cohere.embed-v4:0`).
Request/response shape taken from the AWS Bedrock user guide page "Cohere
Embed v4" (read 2026-08-03):
request {"input_type": "search_document|search_query|classification|
clustering",
"texts": [str], # max 96 per call
"embedding_types": ["float"|"int8"|"uint8"|"binary"|"ubinary"],
"output_dimension": 256|512|1024|1536,
"truncate": "NONE|LEFT|RIGHT"}
The response has two documented shapes and this adapter accepts both. Asking
for one or more `embedding_types` returns
`{"response_type": "embeddings_by_type", "embeddings": {"float": [[...]]}}`;
omitting the field returns
`{"response_type": "embeddings_floats", "embeddings": [[...]]}`.
Three defaults are chosen rather than inherited:
- `output_dimension` is set explicitly. The documented default is 1536, and a
collection built at one width cannot absorb vectors of another.
- `input_type` is derived from the caller's input kind. This is the model
whose asymmetry the `ports` contract exists for: corpus records go in as
`search_document`, queries as `search_query`.
- `truncate` is `NONE`, which makes an over-length input an error instead of a
silently shortened one. A dosing section that lost its tail and embedded
anyway is exactly the failure this project's rules are written against.
"""
from __future__ import annotations
from typing import Any, List, Sequence
from .bedrock_runtime import BedrockInvoker
from .ports import (
INPUT_DOCUMENT,
INPUT_QUERY,
EmbeddingProvider,
EmbeddingVector,
text_digest,
)
MODEL_ID = "cohere.embed-v4:0"
PROVIDER_NAME = "cohere-v4"
SUPPORTED_DIMENSIONS = (256, 512, 1024, 1536)
# The documented per-request ceiling for `texts`.
MAX_TEXTS_PER_REQUEST = 96
COHERE_INPUT_TYPES = {
INPUT_DOCUMENT: "search_document",
INPUT_QUERY: "search_query",
}
# Cohere's docs do not state whether float vectors are unit-length, so this
# stays unset rather than being asserted either way.
NORMALIZED_UNKNOWN = None
class CohereEmbedV4(EmbeddingProvider):
def __init__(
self,
invoker: BedrockInvoker,
dimensions: int = 1024,
truncate: str = "NONE",
batch_size: int = MAX_TEXTS_PER_REQUEST,
):
if dimensions not in SUPPORTED_DIMENSIONS:
raise ValueError(
f"{MODEL_ID} supports {SUPPORTED_DIMENSIONS}, got {dimensions}"
)
if not 1 <= batch_size <= MAX_TEXTS_PER_REQUEST:
raise ValueError(
f"batch_size must be 1..{MAX_TEXTS_PER_REQUEST}, got {batch_size}"
)
self._invoker = invoker
self._dimensions = dimensions
self._truncate = truncate
self._batch_size = batch_size
@property
def name(self) -> str:
return PROVIDER_NAME
@property
def model_id(self) -> str:
return MODEL_ID
@property
def dimensions(self) -> int:
return self._dimensions
@property
def max_batch_size(self) -> int:
return self._batch_size
def _embed_batch(
self, texts: Sequence[str], input_kind: str
) -> List[EmbeddingVector]:
body = self._invoker.invoke_json(
MODEL_ID,
{
"texts": list(texts),
"input_type": COHERE_INPUT_TYPES[input_kind],
"embedding_types": ["float"],
"output_dimension": self._dimensions,
"truncate": self._truncate,
},
# The AWS code example for this model sends `*/*`.
accept="*/*",
)
rows = _float_rows(body)
if len(rows) != len(texts):
raise ValueError(
f"{MODEL_ID} returned {len(rows)} vectors for {len(texts)} texts"
)
vectors: List[EmbeddingVector] = []
for text, values in zip(texts, rows, strict=True):
self._check_dimensions(values)
vectors.append(
EmbeddingVector(
values=list(values),
text_sha256=text_digest(text),
provider=PROVIDER_NAME,
model_id=MODEL_ID,
dimensions=self._dimensions,
input_kind=input_kind,
normalized=NORMALIZED_UNKNOWN,
)
)
return vectors
def _float_rows(body: Any) -> List[List[float]]:
"""Pull the float vectors out of either documented response shape."""
embeddings = body.get("embeddings")
if embeddings is None:
raise ValueError(
f"{MODEL_ID} response has no 'embeddings' field; "
f"keys were {sorted(body)}"
)
if isinstance(embeddings, dict):
rows = embeddings.get("float")
if rows is None:
raise ValueError(
f"{MODEL_ID} returned no float embeddings; "
f"types present: {sorted(embeddings)}"
)
return rows
return embeddings
@@ -0,0 +1,75 @@
"""The only module in this package that knows boto3 exists.
Keeping the SDK behind `BedrockInvoker` is what makes the two Bedrock
adapters testable with no AWS account, no credentials and no spend: a test
passes a stub that returns a canned response body, and the adapter's request
shaping and response parsing — the parts that can actually be wrong — are
exercised in full.
boto3 is imported inside the method rather than at module scope so that
`ingestion.embed` imports cleanly in an environment that never talks to AWS
(the local bge-m3 control, or the mocked tests).
"""
from __future__ import annotations
import json
from typing import Any, Mapping, Optional, Protocol
DEFAULT_REGION = "us-east-1"
BEDROCK_RUNTIME_SERVICE = "bedrock-runtime"
class BedrockInvoker(Protocol):
def invoke_json(
self,
model_id: str,
payload: Mapping[str, Any],
accept: str = "application/json",
) -> dict:
"""POST `payload` as JSON to a Bedrock model, return the parsed body."""
class Boto3BedrockInvoker:
"""`InvokeModel` over boto3, with the JSON/stream plumbing hidden."""
def __init__(self, region: str = DEFAULT_REGION, client: Optional[Any] = None):
self._region = region
self._client = client
@property
def region(self) -> str:
return self._region
def _runtime(self) -> Any:
if self._client is None:
import boto3
from botocore.config import Config
# Without an explicit read timeout a single stalled response hangs
# the whole corpus run: observed 2026-08-04, one request held an
# open socket for over five minutes while boto3's default waited.
self._client = boto3.client(
BEDROCK_RUNTIME_SERVICE,
region_name=self._region,
config=Config(
connect_timeout=10,
read_timeout=60,
retries={"max_attempts": 3, "mode": "standard"},
),
)
return self._client
def invoke_json(
self,
model_id: str,
payload: Mapping[str, Any],
accept: str = "application/json",
) -> dict:
response = self._runtime().invoke_model(
modelId=model_id,
body=json.dumps(payload),
accept=accept,
contentType="application/json",
)
return json.loads(response["body"].read())
@@ -0,0 +1,96 @@
"""Amazon Titan Text Embeddings V2 (`amazon.titan-embed-text-v2:0`).
Request/response shape taken from the AWS Bedrock user guide page
"Amazon Titan Embeddings G1 - Text", V2 tabs (read 2026-08-03):
request {"inputText": str, "dimensions": int, "normalize": bool,
"embeddingTypes": list}
response {"embedding": [float], "inputTextTokenCount": int,
"embeddingsByType": {...}}
`embeddingTypes` is left unset so the response keeps the plain float
`embedding` field — the documentation notes that field disappears when
`embeddingTypes` contains only `binary`.
Titan draws no distinction between a corpus record and a query, so both input
kinds produce a byte-identical request. The kind is still recorded on each
vector, because "this model ignores it" is a fact worth being able to read
back off the data rather than infer.
"""
from __future__ import annotations
from typing import List, Sequence
from .bedrock_runtime import BedrockInvoker
from .ports import EmbeddingProvider, EmbeddingVector, text_digest
MODEL_ID = "amazon.titan-embed-text-v2:0"
PROVIDER_NAME = "titan-v2"
# Per the V2 request documentation.
SUPPORTED_DIMENSIONS = (256, 512, 1024)
class TitanTextEmbeddingsV2(EmbeddingProvider):
def __init__(
self,
invoker: BedrockInvoker,
dimensions: int = 1024,
normalize: bool = True,
):
if dimensions not in SUPPORTED_DIMENSIONS:
raise ValueError(
f"{MODEL_ID} supports {SUPPORTED_DIMENSIONS}, got {dimensions}"
)
self._invoker = invoker
self._dimensions = dimensions
self._normalize = normalize
@property
def name(self) -> str:
return PROVIDER_NAME
@property
def model_id(self) -> str:
return MODEL_ID
@property
def dimensions(self) -> int:
return self._dimensions
@property
def max_batch_size(self) -> int:
# InvokeModel takes a single `inputText`; there is no text array.
return 1
def _embed_batch(
self, texts: Sequence[str], input_kind: str
) -> List[EmbeddingVector]:
text = texts[0]
body = self._invoker.invoke_json(
MODEL_ID,
{
"inputText": text,
"dimensions": self._dimensions,
"normalize": self._normalize,
},
)
values = body.get("embedding")
if values is None:
raise ValueError(
f"{MODEL_ID} response has no 'embedding' field; "
f"keys were {sorted(body)}"
)
self._check_dimensions(values)
return [
EmbeddingVector(
values=list(values),
text_sha256=text_digest(text),
provider=PROVIDER_NAME,
model_id=MODEL_ID,
dimensions=self._dimensions,
input_kind=input_kind,
normalized=self._normalize,
input_token_count=body.get("inputTextTokenCount"),
)
]
+247
View File
@@ -0,0 +1,247 @@
"""A disk cache so a corpus is never paid for twice.
**Why the key is content-addressed.** A vector is a pure function of three
things: the model, the input kind, and the exact bytes embedded. Nothing else
about the record changes the answer. `docs/v1-delivery-plan.md` §4.A proposed
keying on `chunk_id` + sha256; measured against the real corpus that would
charge twice for identical text — `chunks.jsonl` holds 15,066 records but only
14,869 distinct texts, so 197 records (1.31%) are repeats of a text already
embedded. The key here is `(model_id, input_kind, text_sha256)`, which collapses
those and, more importantly, cannot silently serve a stale vector after a chunk's
text is edited: an edit changes the digest, so it is a miss.
Traceability is not lost by dropping `chunk_id` from the key. Every cached
record carries the same sha256 rule (`ports.text_digest`) that produced it, so a
vector is matched back to its chunk by re-digesting that chunk's text. Pairing
vectors to chunk records is `load/`'s job, not the cache's.
**Why the index holds offsets, not vectors.** 15,066 vectors of 1,024 floats do
not belong in memory all at once; a list of that many Python floats is tens of
kilobytes each. Startup scans the file once to map key to byte offset, and a
`get` seeks and parses exactly one line.
The file is append-only. A key already present is never rewritten, so the file
is a log that can be inspected, truncated, or resumed after an interrupted run
without a repair step.
"""
from __future__ import annotations
import json
import os
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Optional, Sequence, Tuple
from .ports import (
INPUT_KINDS,
EmbeddingBatch,
EmbeddingProvider,
EmbeddingVector,
text_digest,
)
CacheKey = Tuple[str, str, str]
_KEY_FIELDS = ("model_id", "input_kind", "text_sha256")
def cache_key(model_id: str, input_kind: str, text: str) -> CacheKey:
return (model_id, input_kind, text_digest(text))
@dataclass(frozen=True)
class CacheStats:
hits: int = 0
misses: int = 0
@property
def lookups(self) -> int:
return self.hits + self.misses
@property
def hit_rate(self) -> float:
return self.hits / self.lookups if self.lookups else 0.0
class EmbeddingCache:
"""Append-only JSONL of vectors, indexed by byte offset."""
def __init__(self, path: os.PathLike | str) -> None:
self._path = Path(path)
self._offsets: Dict[CacheKey, int] = {}
self._hits = 0
self._misses = 0
if self._path.exists():
self._build_index()
@property
def path(self) -> Path:
return self._path
@property
def stats(self) -> CacheStats:
return CacheStats(hits=self._hits, misses=self._misses)
def __len__(self) -> int:
return len(self._offsets)
def __contains__(self, key: CacheKey) -> bool:
return key in self._offsets
def _build_index(self) -> None:
with self._path.open("rb") as handle:
offset = 0
for raw in handle:
line = raw.decode("utf-8").strip()
if line:
record = json.loads(line)
self._offsets[self._key_of(record)] = offset
offset += len(raw)
@staticmethod
def _key_of(record: dict) -> CacheKey:
missing = [f for f in _KEY_FIELDS if not record.get(f)]
if missing:
raise ValueError(f"cache record is missing key fields: {missing}")
return (
record["model_id"],
record["input_kind"],
record["text_sha256"],
)
def get(self, key: CacheKey) -> Optional[EmbeddingVector]:
offset = self._offsets.get(key)
if offset is None:
self._misses += 1
return None
with self._path.open("rb") as handle:
handle.seek(offset)
record = json.loads(handle.readline().decode("utf-8"))
self._hits += 1
return _vector_from_record(record)
def put(self, vector: EmbeddingVector) -> bool:
"""Append a vector. Returns False if the key was already stored."""
key = (vector.model_id, vector.input_kind, vector.text_sha256)
if key in self._offsets:
return False
self._path.parent.mkdir(parents=True, exist_ok=True)
line = json.dumps(_record_from_vector(vector), ensure_ascii=False) + "\n"
encoded = line.encode("utf-8")
with self._path.open("ab") as handle:
offset = handle.tell()
handle.write(encoded)
self._offsets[key] = offset
return True
def put_many(self, vectors: Iterable[EmbeddingVector]) -> int:
return sum(1 for vector in vectors if self.put(vector))
def _record_from_vector(vector: EmbeddingVector) -> dict:
return {
"model_id": vector.model_id,
"input_kind": vector.input_kind,
"text_sha256": vector.text_sha256,
"provider": vector.provider,
"dimensions": vector.dimensions,
"normalized": vector.normalized,
"input_token_count": vector.input_token_count,
"values": vector.values,
}
def _vector_from_record(record: dict) -> EmbeddingVector:
values: List[float] = record["values"]
declared = record["dimensions"]
if len(values) != declared:
raise ValueError(
f"cached vector for {record['text_sha256'][:12]} has {len(values)} "
f"values but declares {declared} dimensions"
)
return EmbeddingVector(
values=values,
text_sha256=record["text_sha256"],
provider=record["provider"],
model_id=record["model_id"],
dimensions=declared,
input_kind=record["input_kind"],
normalized=record.get("normalized"),
input_token_count=record.get("input_token_count"),
)
class CachingEmbeddingProvider(EmbeddingProvider):
"""Wraps a provider so only uncached texts reach it.
A decorator rather than a change to each adapter: the three existing
providers stay unaware that a cache exists, and a fourth needs no cache code
to benefit. `request_count` counts requests the *inner* provider actually
made, which is what makes "a second run costs nothing" a checkable claim
rather than an assertion — a fully cached run reports zero.
"""
def __init__(self, inner: EmbeddingProvider, cache: EmbeddingCache) -> None:
self._inner = inner
self._cache = cache
@property
def name(self) -> str:
return f"cached:{self._inner.name}"
@property
def model_id(self) -> str:
return self._inner.model_id
@property
def dimensions(self) -> int:
return self._inner.dimensions
@property
def max_batch_size(self) -> int:
return self._inner.max_batch_size
@property
def cache(self) -> EmbeddingCache:
return self._cache
def _embed_batch(
self, texts: Sequence[str], input_kind: str
) -> List[EmbeddingVector]:
return list(self._inner.embed(list(texts), input_kind).vectors)
def embed(self, texts: Sequence[str], input_kind: str) -> EmbeddingBatch:
if input_kind not in INPUT_KINDS:
raise ValueError(
f"input_kind must be one of {INPUT_KINDS}, got {input_kind!r}"
)
if any(not t.strip() for t in texts):
raise ValueError("refusing to embed an empty or whitespace-only text")
resolved: List[Optional[EmbeddingVector]] = []
pending: Dict[str, List[int]] = {}
for position, text in enumerate(texts):
hit = self._cache.get(cache_key(self.model_id, input_kind, text))
resolved.append(hit)
if hit is None:
pending.setdefault(text, []).append(position)
requests = 0
latency_ms = 0.0
if pending:
wanted = list(pending)
batch = self._inner.embed(wanted, input_kind)
requests = batch.request_count
latency_ms = batch.latency_ms
for text, vector in zip(wanted, batch.vectors, strict=True):
self._cache.put(vector)
for position in pending[text]:
resolved[position] = vector
if any(vector is None for vector in resolved):
raise ValueError("cache resolution left a text without a vector")
return EmbeddingBatch(
vectors=[vector for vector in resolved if vector is not None],
request_count=requests,
latency_ms=latency_ms,
)
+103
View File
@@ -0,0 +1,103 @@
"""BAAI/bge-m3 running locally — the zero-API-cost control in the benchmark.
Its job is to answer "how much is the paid model actually buying us on
Vietnamese medical prose?". Without a free baseline in the same harness, a
cloud model's recall number has nothing to be better *than*.
Two properties are taken from the published model card and have **not** been
verified on this machine (no local run has happened yet — see the coordination
handoff): the dense vector is 1024-dimensional, and bge-m3 needs no
instruction prefix on either the corpus or the query side, unlike the earlier
English bge models. Both are asserted at runtime rather than trusted: the
dimension is checked on every vector by `EmbeddingProvider._check_dimensions`,
so a wrong assumption fails on the first call instead of producing a
quietly unusable collection.
`sentence-transformers` is imported lazily and the encoder is injectable, so
this module costs nothing to import and can be tested without the ~2 GB of
model weights.
"""
from __future__ import annotations
from typing import Callable, List, Optional, Sequence
from .ports import EmbeddingProvider, EmbeddingVector, text_digest
MODEL_ID = "BAAI/bge-m3"
PROVIDER_NAME = "bge-m3"
DENSE_DIMENSIONS = 1024
Encoder = Callable[[Sequence[str]], Sequence[Sequence[float]]]
class BgeM3Local(EmbeddingProvider):
def __init__(
self,
encoder: Optional[Encoder] = None,
batch_size: int = 16,
device: Optional[str] = None,
):
if batch_size < 1:
raise ValueError(f"batch_size must be >= 1, got {batch_size}")
self._encoder = encoder
self._batch_size = batch_size
self._device = device
# Only the encoder built below is known to normalize. An injected one
# is somebody else's function, so its output is recorded as unknown.
self._normalized: Optional[bool] = None if encoder is not None else True
@property
def name(self) -> str:
return PROVIDER_NAME
@property
def model_id(self) -> str:
return MODEL_ID
@property
def dimensions(self) -> int:
return DENSE_DIMENSIONS
@property
def max_batch_size(self) -> int:
return self._batch_size
def _load_encoder(self) -> Encoder:
if self._encoder is None:
from sentence_transformers import SentenceTransformer
model = SentenceTransformer(MODEL_ID, device=self._device)
def encode(texts: Sequence[str]) -> Sequence[Sequence[float]]:
return model.encode(
list(texts), normalize_embeddings=True
).tolist()
self._encoder = encode
return self._encoder
def _embed_batch(
self, texts: Sequence[str], input_kind: str
) -> List[EmbeddingVector]:
rows = self._load_encoder()(texts)
if len(rows) != len(texts):
raise ValueError(
f"{MODEL_ID} returned {len(rows)} vectors for {len(texts)} texts"
)
vectors: List[EmbeddingVector] = []
for text, values in zip(texts, rows, strict=True):
self._check_dimensions(values)
vectors.append(
EmbeddingVector(
values=list(values),
text_sha256=text_digest(text),
provider=PROVIDER_NAME,
model_id=MODEL_ID,
dimensions=DENSE_DIMENSIONS,
input_kind=input_kind,
normalized=self._normalized,
)
)
return vectors
+144
View File
@@ -0,0 +1,144 @@
"""The embedding boundary: what a provider must do, and what it must record.
Two things drive this design.
**Asymmetric models make the input kind part of the contract.** Cohere embeds
a corpus record and a search query into deliberately different subspaces —
the same string sent as `search_document` and as `search_query` does not come
back as the same vector. Getting that backwards raises no error; recall just
quietly drops. So `input_kind` is a required argument of `embed()`, not an
optional keyword a caller can forget, and the value used is recorded on every
vector so a mismatch is detectable after the fact.
**Provenance applies to vectors too.** CLAUDE.md requires an extracted unit to
stay traceable to its source; a vector is no different. `model_id`,
`dimensions`, `input_kind` and the sha256 of the exact text embedded are what
let a collection be checked for the one mistake that is invisible from the
outside — vectors from two different models mixed into one Qdrant collection,
where every query still returns *something*.
`normalized` is deliberately three-valued. Titan is asked to normalize and
says so; a local encoder is told to; Cohere's Bedrock documentation does not
state whether its float vectors are unit-length, so the field stays `None`
rather than guessing. An unmeasured claim does not get written down as a fact.
"""
from __future__ import annotations
import hashlib
import time
from abc import ABC, abstractmethod
from dataclasses import dataclass, field
from typing import List, Optional, Sequence
INPUT_DOCUMENT = "document"
INPUT_QUERY = "query"
INPUT_KINDS = (INPUT_DOCUMENT, INPUT_QUERY)
def text_digest(text: str) -> str:
"""sha256 of the exact string sent to the provider.
Shared by every adapter so a cached vector can be matched to its text by
the same rule that produced it.
"""
return hashlib.sha256(text.encode("utf-8")).hexdigest()
@dataclass(frozen=True)
class EmbeddingVector:
values: List[float]
text_sha256: str
provider: str
model_id: str
dimensions: int
input_kind: str
# None means the provider does not document it — not "no".
normalized: Optional[bool] = None
# Only some providers report it (Titan does, Cohere's documented text
# response does not).
input_token_count: Optional[int] = None
@dataclass(frozen=True)
class EmbeddingBatch:
"""Vectors plus what a benchmark needs to compare providers fairly."""
vectors: List[EmbeddingVector] = field(default_factory=list)
request_count: int = 0
latency_ms: float = 0.0
class EmbeddingProvider(ABC):
"""One embedding model, reachable without the caller knowing its SDK.
Subclasses implement `_embed_batch` for a single request; `embed` owns
input validation, splitting into provider-sized requests, and timing, so
that logic exists once rather than per adapter.
"""
@property
@abstractmethod
def name(self) -> str:
"""Short registry key, e.g. `titan-v2`."""
@property
@abstractmethod
def model_id(self) -> str:
"""Provider-side identifier, e.g. `amazon.titan-embed-text-v2:0`."""
@property
@abstractmethod
def dimensions(self) -> int:
"""Vector length this instance is configured to produce."""
@property
def max_batch_size(self) -> int:
"""Texts accepted per request. Default is the safest possible value."""
return 1
@abstractmethod
def _embed_batch(
self, texts: Sequence[str], input_kind: str
) -> List[EmbeddingVector]:
"""Embed at most `max_batch_size` texts in one provider request."""
def embed(self, texts: Sequence[str], input_kind: str) -> EmbeddingBatch:
if input_kind not in INPUT_KINDS:
raise ValueError(
f"input_kind must be one of {INPUT_KINDS}, got {input_kind!r}"
)
if any(not t.strip() for t in texts):
raise ValueError("refusing to embed an empty or whitespace-only text")
vectors: List[EmbeddingVector] = []
requests = 0
started = time.perf_counter()
for start in range(0, len(texts), self.max_batch_size):
window = texts[start : start + self.max_batch_size]
vectors.extend(self._embed_batch(window, input_kind))
requests += 1
elapsed_ms = (time.perf_counter() - started) * 1000.0
if len(vectors) != len(texts):
raise ValueError(
f"{self.name} returned {len(vectors)} vectors for "
f"{len(texts)} texts"
)
return EmbeddingBatch(
vectors=vectors, request_count=requests, latency_ms=elapsed_ms
)
def embed_documents(self, texts: Sequence[str]) -> EmbeddingBatch:
return self.embed(texts, INPUT_DOCUMENT)
def embed_queries(self, texts: Sequence[str]) -> EmbeddingBatch:
return self.embed(texts, INPUT_QUERY)
def _check_dimensions(self, values: Sequence[float]) -> None:
"""A wrong-length vector is a corpus-wide defect; fail on the first."""
if len(values) != self.dimensions:
raise ValueError(
f"{self.model_id} returned {len(values)} dimensions, "
f"expected {self.dimensions}"
)
+71
View File
@@ -0,0 +1,71 @@
"""One live call to one provider — the smallest thing that proves it works.
Deliberately not wired into `ingestion.cli`: that module is being edited for
the parser/chunking work, and this task has no business touching it. Run as
python -m ingestion.embed.probe --provider titan-v2
Each run embeds a **single short string** and makes exactly one request, so it
answers "are the credentials, the model access and my request shape all
right?" without approaching the cost or the risk of a corpus run. It reports
the request key shape, the returned dimension, the measured L2 norm and the
latency — the L2 norm because whether a provider returns unit-length vectors
is a property worth measuring rather than reading off a documentation page.
"""
from __future__ import annotations
import argparse
import math
import sys
from typing import Sequence
from .ports import INPUT_DOCUMENT, INPUT_KINDS, EmbeddingVector
from .registry import DEFAULT_REGION, build_provider, provider_names
# Vietnamese, diacritics, and a dose string: the shape of the real corpus, not
# "hello world". If an encoding path is broken this is what shows it.
DEFAULT_TEXT = "Liều dùng: uống 500 mg mỗi 8 giờ, không quá 4 g mỗi ngày."
def _l2_norm(values: Sequence[float]) -> float:
return math.sqrt(sum(v * v for v in values))
def _report(vector: EmbeddingVector, latency_ms: float, requests: int) -> None:
print(f"provider : {vector.provider}")
print(f"model_id : {vector.model_id}")
print(f"input_kind : {vector.input_kind}")
print(f"requests : {requests}")
print(f"dimensions : {len(vector.values)} (expected {vector.dimensions})")
print(f"measured L2 norm: {_l2_norm(vector.values):.6f}")
print(f"normalized (per provider docs): {vector.normalized}")
print(f"input_token_count: {vector.input_token_count}")
print(f"latency_ms : {latency_ms:.1f}")
print(f"first 5 values : {[round(v, 6) for v in vector.values[:5]]}")
def main(argv=None) -> int:
parser = argparse.ArgumentParser(
prog="python -m ingestion.embed.probe",
description="Make ONE live embedding call and report what came back.",
)
parser.add_argument("--provider", required=True, choices=provider_names())
parser.add_argument("--region", default=DEFAULT_REGION)
parser.add_argument("--dimensions", type=int, default=None)
parser.add_argument("--text", default=DEFAULT_TEXT)
parser.add_argument("--input-kind", default=INPUT_DOCUMENT, choices=INPUT_KINDS)
args = parser.parse_args(argv)
if hasattr(sys.stdout, "reconfigure"):
sys.stdout.reconfigure(encoding="utf-8")
provider = build_provider(
args.provider, region=args.region, dimensions=args.dimensions
)
batch = provider.embed([args.text], args.input_kind)
_report(batch.vectors[0], batch.latency_ms, batch.request_count)
return 0
if __name__ == "__main__":
raise SystemExit(main())
+78
View File
@@ -0,0 +1,78 @@
"""Name -> provider, so nothing downstream has to import an SDK to pick one.
Adding a fourth model to the benchmark means adding one entry to `_BUILDERS`,
not editing a caller — the open/closed rule CLAUDE.md applies to the section
taxonomy, applied here for the same reason.
"""
from __future__ import annotations
from typing import Callable, Dict, Optional
from . import bedrock_cohere, bedrock_titan, local_bge_m3
from .bedrock_runtime import DEFAULT_REGION, BedrockInvoker, Boto3BedrockInvoker
from .ports import EmbeddingProvider
TITAN_V2 = bedrock_titan.PROVIDER_NAME
COHERE_V4 = bedrock_cohere.PROVIDER_NAME
BGE_M3 = local_bge_m3.PROVIDER_NAME
CLOUD_PROVIDERS = (TITAN_V2, COHERE_V4)
def _build_titan(
invoker: Optional[BedrockInvoker], region: str, dimensions: Optional[int]
) -> EmbeddingProvider:
return bedrock_titan.TitanTextEmbeddingsV2(
invoker or Boto3BedrockInvoker(region=region),
dimensions=dimensions or 1024,
)
def _build_cohere(
invoker: Optional[BedrockInvoker], region: str, dimensions: Optional[int]
) -> EmbeddingProvider:
return bedrock_cohere.CohereEmbedV4(
invoker or Boto3BedrockInvoker(region=region),
dimensions=dimensions or 1024,
)
def _build_bge_m3(
_invoker: Optional[BedrockInvoker], _region: str, dimensions: Optional[int]
) -> EmbeddingProvider:
# Runs on this machine: there is no invoker and no region to honour.
if dimensions not in (None, local_bge_m3.DENSE_DIMENSIONS):
raise ValueError(
f"{BGE_M3} produces {local_bge_m3.DENSE_DIMENSIONS} dimensions; "
f"{dimensions} was requested"
)
return local_bge_m3.BgeM3Local()
Builder = Callable[[Optional[BedrockInvoker], str, Optional[int]], EmbeddingProvider]
_BUILDERS: Dict[str, Builder] = {
TITAN_V2: _build_titan,
COHERE_V4: _build_cohere,
BGE_M3: _build_bge_m3,
}
def provider_names() -> tuple:
return tuple(_BUILDERS)
def build_provider(
name: str,
*,
invoker: Optional[BedrockInvoker] = None,
region: str = DEFAULT_REGION,
dimensions: Optional[int] = None,
) -> EmbeddingProvider:
try:
builder = _BUILDERS[name]
except KeyError:
raise ValueError(
f"unknown embedding provider {name!r}; known: {provider_names()}"
) from None
return builder(invoker, region, dimensions)
+1
View File
@@ -0,0 +1 @@
"""Verified drug-entity artifact builders."""
+156
View File
@@ -0,0 +1,156 @@
from __future__ import annotations
import argparse
import json
import re
import unicodedata
from difflib import SequenceMatcher
from pathlib import Path
import fitz
from ingestion.validation.back_index import parse_back_index_see_aliases
WORD_RE = re.compile(r"\w+", re.UNICODE)
PAREN_RE = re.compile(r"\(([^()]*)\)")
def normalize_name(text: str) -> str:
decomposed = unicodedata.normalize("NFKD", text.casefold()).replace("đ", "d")
plain = "".join(char for char in decomposed if not unicodedata.combining(char))
return " ".join(WORD_RE.findall(plain))
def _canonical_aliases(drug_name: str) -> set[str]:
aliases = {drug_name.strip()}
without_parentheses = PAREN_RE.sub("", drug_name).strip()
if without_parentheses:
aliases.add(without_parentheses)
aliases.update(
value.strip() for value in PAREN_RE.findall(drug_name) if value.strip()
)
return aliases
def _trade_aliases(monograph: dict) -> set[str]:
section = monograph.get("sections", {}).get("ten_thuong_mai")
if not section:
return set()
text = section.get("text", "")
return {
value.strip().strip(".")
for value in re.split(r"[,;\n]", text)
if value.strip().strip(".")
}
def _read_monographs(path: Path) -> list[dict]:
with path.open(encoding="utf-8") as handle:
return [json.loads(line) for line in handle if line.strip()]
def build_entities(monographs_path: Path, pdf_path: Path) -> dict:
monographs = _read_monographs(monographs_path)
entities: dict[str, dict] = {}
lookup: list[tuple[str, str, tuple[int, int]]] = []
for monograph in monographs:
drug_id = monograph["drug_id"]
aliases = _canonical_aliases(monograph["drug_name"])
aliases.update(_trade_aliases(monograph))
aliases.add(drug_id.replace("_", " "))
entities[drug_id] = {
"drug_id": drug_id,
"canonical_name": monograph["drug_name"],
"aliases": aliases,
"atc_codes": sorted(set(monograph.get("atc_codes", []))),
"source_page_range": monograph.get("source_page_range"),
}
page_range = tuple(monograph["source_page_range"])
lookup.extend(
(normalize_name(alias), drug_id, page_range) for alias in _canonical_aliases(
monograph["drug_name"],
) if normalize_name(alias)
)
unresolved = []
ambiguous = []
with fitz.open(pdf_path) as doc:
index_aliases = parse_back_index_see_aliases(doc)
for relation in index_aliases:
target = normalize_name(relation.target)
candidates = {
drug_id for alias, drug_id, page_range in lookup
if page_range[0] <= relation.printed_page - 1 <= page_range[1]
and (
target == alias
or target.startswith(f"{alias} ")
or target.endswith(f" {alias}")
or f" {alias} " in target
)
}
if not candidates:
fuzzy = sorted(
(
SequenceMatcher(None, target, alias).ratio(),
drug_id,
)
for alias, drug_id, page_range in lookup
if page_range[0] <= relation.printed_page - 1 <= page_range[1]
)
if fuzzy and fuzzy[-1][0] >= 0.9:
runner_up = fuzzy[-2][0] if len(fuzzy) > 1 else 0.0
if fuzzy[-1][0] - runner_up >= 0.05:
candidates = {fuzzy[-1][1]}
if len(candidates) == 1:
entities[next(iter(candidates))]["aliases"].add(relation.alias)
elif candidates:
ambiguous.append(relation.alias)
else:
unresolved.append(relation.alias)
output_entities = []
for entity in entities.values():
output_entities.append({
**entity,
"aliases": sorted(entity["aliases"], key=lambda value: value.casefold()),
})
return {
"schema_version": 1,
"entities": sorted(output_entities, key=lambda item: item["drug_id"]),
"stats": {
"entity_count": len(output_entities),
"back_index_see_relations": len(index_aliases),
"back_index_aliases_mapped": len(index_aliases) - len(unresolved) - len(ambiguous),
"back_index_aliases_unresolved": len(unresolved),
"back_index_aliases_ambiguous": len(ambiguous),
"trade_name_sections": sum(
"ten_thuong_mai" in row.get("sections", {}) for row in monographs
),
"total_aliases": sum(len(row["aliases"]) for row in output_entities),
},
"unresolved_back_index_aliases": sorted(unresolved),
"ambiguous_back_index_aliases": sorted(ambiguous),
}
def write_entities(payload: dict, path: Path) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(
json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
encoding="utf-8",
)
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--monographs", type=Path, required=True)
parser.add_argument("--pdf", type=Path, required=True)
parser.add_argument("--output", type=Path, required=True)
args = parser.parse_args()
payload = build_entities(args.monographs, args.pdf)
write_entities(payload, args.output)
print(json.dumps(payload["stats"], ensure_ascii=False))
if __name__ == "__main__":
main()
+11 -1
View File
@@ -26,6 +26,12 @@ from ..tables.models import TableRegion
# One line of body type on this book measures ~10.5pt; a fraction spans the
# numerator line, the bar and the denominator line.
FORMULA_BAND_HEIGHT_PT = 13.0
# ADENOSIN p147 is printed as numerator and denominator lines with no rule.
# Its synthetic anchor sits on the wrapped numerator baseline, so symmetric
# growth covers only the numerator. The denominator ends ~29pt below the
# anchor; the next "Ví dụ:" line has its centre ~32.3pt below. 31pt captures
# the denominator while excluding that prose and the following table.
BARLESS_FORMULA_BOTTOM_PT = 31.0
# Wide on purpose. The bar is often narrower than the numerator above it, and
# a numerator span can carry leading spaces that push its box's centre well to
@@ -51,12 +57,16 @@ def load_formula_regions(path: Path | None = None) -> List[TableRegion]:
regions = []
for index, entry in enumerate(payload["regions"]):
x0, y0, x1, y1 = entry["bar_bbox"]
bottom_margin = (
BARLESS_FORMULA_BOTTOM_PT
if entry.get("source_prints_no_bar") else FORMULA_BAND_HEIGHT_PT
)
regions.append(
TableRegion(
table_id=f"p{entry['physical_page']}_f{index}",
physical_page=entry["physical_page"],
bbox=(x0 - FORMULA_SIDE_MARGIN_PT, y0 - FORMULA_BAND_HEIGHT_PT,
x1 + FORMULA_SIDE_MARGIN_PT, y1 + FORMULA_BAND_HEIGHT_PT),
x1 + FORMULA_SIDE_MARGIN_PT, y1 + bottom_margin),
n_rows=2,
n_cols=1,
shape=SHAPE_FORMULA_2D,
+4
View File
@@ -32,6 +32,10 @@ class Span:
def bold(self) -> bool:
return "Bold" in self.font
@property
def italic(self) -> bool:
return "Italic" in self.font or "Oblique" in self.font
@property
def span_id(self) -> str:
"""Stable identifier for one source span.
+68
View File
@@ -0,0 +1,68 @@
"""Load stage: chunks plus vectors into a vector store, idempotently.
Import from here rather than from a module. Nothing outside this package should
name `qdrant_client` — that dependency is confined to `qdrant_repo`, so the
whole stage runs and is tested with no server, exactly as `embed/` confines
boto3 to `bedrock_runtime`.
The two rules this package exists to enforce are worth naming at the front
door. A point id is derived from `chunk_id`, so loading twice converges instead
of duplicating. And a collection records the corpus digest and model it was
built from, so a second corpus or a second model cannot be mixed into it — a
state that produces no error at query time and would otherwise be invisible.
"""
from .corpus import corpus_sha256, count_chunks, iter_chunk_records
from .in_memory import InMemoryVectorStore
from .manifest import (
MANIFEST_SUFFIX,
CorpusMismatch,
assert_compatible,
manifest_collection,
read_manifest,
write_manifest,
)
from .models import (
COSINE,
INDEXED_PAYLOAD_FIELDS,
REQUIRED_CHUNK_FIELDS,
CollectionSpec,
CorpusManifest,
VectorPoint,
build_point,
point_id_for,
validate_chunk_record,
)
from .ports import VectorStore
from .upsert import (
DEFAULT_BATCH_SIZE,
ChunkLoader,
LoadReport,
PointCountMismatch,
)
__all__ = [
"COSINE",
"DEFAULT_BATCH_SIZE",
"INDEXED_PAYLOAD_FIELDS",
"MANIFEST_SUFFIX",
"REQUIRED_CHUNK_FIELDS",
"ChunkLoader",
"CollectionSpec",
"CorpusManifest",
"CorpusMismatch",
"InMemoryVectorStore",
"LoadReport",
"PointCountMismatch",
"VectorPoint",
"VectorStore",
"assert_compatible",
"build_point",
"corpus_sha256",
"count_chunks",
"iter_chunk_records",
"manifest_collection",
"point_id_for",
"read_manifest",
"validate_chunk_record",
"write_manifest",
]
+48
View File
@@ -0,0 +1,48 @@
"""Reading the chunk artifact, and computing the identity it is loaded under.
The digest is taken over content rather than over parsed records, so it still
changes on a field reordering that leaves every record semantically identical —
a normalised, semantic digest would let a regenerated corpus pass the A6 gate
while its point payloads no longer match what `chunk/` produces.
**Line endings are the one thing normalised, and only because not doing it was
a bug.** A raw-byte digest makes the same JSONL hash differently after a Windows
checkout with CRLF than after a Linux one with LF. The A6 gate would then refuse
a load in CI against a corpus that is byte-for-byte the same data, which is a
false rejection with a confusing message — the failure mode of a safety gate
that cries wolf is that someone turns it off. Each line is digested with its
terminator normalised to `\\n`, which costs none of the strictness that matters.
"""
from __future__ import annotations
import hashlib
import json
from pathlib import Path
from typing import Any, Dict, Iterator
def corpus_sha256(path: Path | str) -> str:
"""Content digest, identical across CRLF and LF checkouts of the same data."""
digest = hashlib.sha256()
with Path(path).open("rb") as handle:
for raw in handle:
digest.update(raw.rstrip(b"\r\n"))
digest.update(b"\n")
return digest.hexdigest()
def iter_chunk_records(path: Path | str) -> Iterator[Dict[str, Any]]:
with Path(path).open("r", encoding="utf-8") as handle:
for line_number, line in enumerate(handle, start=1):
stripped = line.strip()
if not stripped:
continue
try:
yield json.loads(stripped)
except json.JSONDecodeError as exc:
raise ValueError(
f"{Path(path).name} line {line_number} is not valid JSON: {exc}"
) from exc
def count_chunks(path: Path | str) -> int:
return sum(1 for _ in iter_chunk_records(path))
+93
View File
@@ -0,0 +1,93 @@
"""A `VectorStore` that keeps everything in a dict.
Not only a test double. It is the reference implementation of the port: the
idempotency rule (same point id overwrites, never appends) and the
unknown-collection failures are stated here once, so a test that passes against
this store is testing the contract rather than an accident of Qdrant's
behaviour. `apps/ai-service/rag` already keeps an in-memory store for the same
reason.
"""
from __future__ import annotations
from typing import Any, Dict, List, Mapping, Optional, Sequence
from .models import CollectionSpec, VectorPoint
class InMemoryVectorStore:
def __init__(self) -> None:
self._collections: Dict[str, CollectionSpec] = {}
self._points: Dict[str, Dict[str, VectorPoint]] = {}
self._indexes: Dict[str, List[tuple]] = {}
def collection_exists(self, name: str) -> bool:
return name in self._collections
def create_collection(self, spec: CollectionSpec) -> None:
if spec.name in self._collections:
raise ValueError(f"collection {spec.name!r} already exists")
self._collections[spec.name] = spec
self._points[spec.name] = {}
self._indexes[spec.name] = []
def delete_collection(self, name: str) -> None:
self._collections.pop(name, None)
self._points.pop(name, None)
self._indexes.pop(name, None)
def create_payload_index(
self, name: str, field_name: str, field_schema: str
) -> None:
self._require(name)
self._indexes[name].append((field_name, field_schema))
def upsert(self, name: str, points: Sequence[VectorPoint]) -> int:
self._require(name)
spec = self._collections[name]
for point in points:
if len(point.vector) != spec.vector_size:
raise ValueError(
f"point {point.id} has {len(point.vector)} dimensions, "
f"collection {name!r} expects {spec.vector_size}"
)
self._points[name][point.id] = point
return len(points)
def count(self, name: str) -> int:
self._require(name)
return len(self._points[name])
def retrieve(self, name: str, point_id: str) -> Optional[VectorPoint]:
self._require(name)
return self._points[name].get(point_id)
def find_by_payload(
self, name: str, equals: Mapping[str, Any]
) -> List[VectorPoint]:
self._require(name)
if not equals:
raise ValueError("find_by_payload needs at least one condition")
return [
point
for point in self._points[name].values()
if all(_matches(point.payload.get(k), v) for k, v in equals.items())
]
def indexed_fields(self, name: str) -> List[tuple]:
self._require(name)
return list(self._indexes[name])
def spec(self, name: str) -> CollectionSpec:
self._require(name)
return self._collections[name]
def _require(self, name: str) -> None:
if name not in self._collections:
raise KeyError(f"collection {name!r} does not exist")
def _matches(stored: Any, wanted: Any) -> bool:
"""Qdrant's rule: a list field matches when any element equals the value."""
if isinstance(stored, list):
return wanted in stored
return stored == wanted
+86
View File
@@ -0,0 +1,86 @@
"""Binding a collection to the exact corpus and model it was built from (A6).
**Why a sidecar collection rather than a reserved point.** Qdrant has no
collection-level metadata field, so the manifest has to live in a point. Putting
that point inside the data collection would make `count()` one larger than the
chunk count — and `qdrant_point_count != chunk_count` is a v1 acceptance gate
(`docs/v1-delivery-plan.md` §6). A gate that needs an "except the manifest"
footnote is a gate that will eventually be read wrong. A `<name>__manifest`
collection keeps the data collection's count exactly equal to the number of
chunks, and keeps the manifest out of every search result by construction
rather than by remembering to filter it.
The vector on that point is a single zero. It is never searched; the point
exists only to carry a payload.
"""
from __future__ import annotations
from typing import Optional
from .models import CollectionSpec, CorpusManifest, VectorPoint
from .ports import VectorStore
MANIFEST_SUFFIX = "__manifest"
MANIFEST_POINT_ID = "00000000-0000-5000-8000-000000000001"
class CorpusMismatch(RuntimeError):
"""Raised instead of upserting a corpus into a collection built elsewhere."""
def manifest_collection(name: str) -> str:
return f"{name}{MANIFEST_SUFFIX}"
def write_manifest(store: VectorStore, name: str, manifest: CorpusManifest) -> None:
sidecar = manifest_collection(name)
if not store.collection_exists(sidecar):
store.create_collection(CollectionSpec(name=sidecar, vector_size=1))
store.upsert(
sidecar,
[
VectorPoint(
id=MANIFEST_POINT_ID,
vector=[0.0],
payload=manifest.to_payload(),
)
],
)
def read_manifest(store: VectorStore, name: str) -> Optional[CorpusManifest]:
sidecar = manifest_collection(name)
if not store.collection_exists(sidecar):
return None
point = store.retrieve(sidecar, MANIFEST_POINT_ID)
if point is None:
return None
return CorpusManifest.from_payload(point.payload)
def assert_compatible(
store: VectorStore, name: str, incoming: CorpusManifest
) -> Optional[CorpusManifest]:
"""Refuse the load unless the collection was built from the same corpus.
Returns the stored manifest, or None when the collection is new. An
existing data collection with no manifest is itself a refusal: it was
loaded by something that did not record what it loaded, so nothing can be
said about what is already in there.
"""
stored = read_manifest(store, name)
if stored is None:
if store.collection_exists(name) and store.count(name) > 0:
raise CorpusMismatch(
f"collection {name!r} already holds {store.count(name)} points but "
f"has no manifest; refusing to mix an unknown corpus with "
f"{incoming.corpus_sha256[:12]}"
)
return None
reasons = stored.conflicts_with(incoming)
if reasons:
raise CorpusMismatch(
f"refusing to load into {name!r}: " + "; ".join(reasons)
)
return stored
+243
View File
@@ -0,0 +1,243 @@
"""What goes into the vector store, and the rules a record must satisfy first.
Two decisions are worth stating because their alternatives look reasonable.
**Point ids are derived, never generated.** `uuid5` of `chunk_id` means the
same chunk always lands on the same point, so a re-run overwrites rather than
duplicates. A random id would make `cli load` non-idempotent, and the damage
would be invisible — the collection would simply hold two copies of a dose and
return whichever ranked higher.
**Payload carries the whole chunk record, not a chosen subset.** CLAUDE.md
requires provenance to survive a stage boundary; a whitelist here would silently
drop any field a later chunker adds, which is exactly the failure it warns
about. So the record passes through intact and only a required core is
*checked* — `population_tags` or `printed_page_range` will flow through the day
`chunk/` starts emitting them, with no edit to this module.
"""
from __future__ import annotations
import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, List, Mapping, Optional, Sequence
# Fixed namespace: changing it would re-id the entire corpus and orphan every
# point already loaded. It is a constant of the project, not a tunable.
POINT_NAMESPACE = uuid.UUID("6f0d6d1e-4c2a-5f6b-9a3d-2f8e1c7b4a90")
COSINE = "Cosine"
# Schema v4 separates retrieval-enriched `text` from contiguous `source_text`.
# Clinicians cite the printed folio, and
# `citation_uses_physical_page = 0` is a v1 acceptance gate, so a chunk without
# one cannot be cited honestly. Refusing it here is the difference between
# noticing before an embedding run and noticing after paying for one, when every
# answer abstains for missing provenance.
SUPPORTED_SCHEMA_VERSION = 4
# Absence means the loader was handed something other than the chunk artifact,
# and it should say so loudly rather than write a point that cannot be traced
# back to a page.
REQUIRED_CHUNK_FIELDS = (
"chunk_id",
"drug_id",
"drug_name",
"section_key",
"text",
"source_text",
"heading_physical_page",
"source_page_range",
"printed_page_range",
"chunk_kind",
)
# Both must be a real two-page span, not merely present. An empty list is the
# shape a defaulted field takes, and `[] in (None, "")` is False — which is how
# an earlier version of this check passed one.
PAGE_RANGE_FIELDS = ("source_page_range", "printed_page_range")
# Fields the retrieval side filters on. Qdrant needs an explicit index per
# field; without it a filtered query still works but scans.
INDEXED_PAYLOAD_FIELDS = (
("chunk_id", "keyword"),
("drug_id", "keyword"),
("section_key", "keyword"),
("atc_codes", "keyword"),
("chunk_kind", "keyword"),
("has_quarantined_content", "bool"),
)
def point_id_for(chunk_id: str) -> str:
if not chunk_id or not chunk_id.strip():
raise ValueError("chunk_id is required to derive a point id")
return str(uuid.uuid5(POINT_NAMESPACE, chunk_id))
@dataclass(frozen=True)
class VectorPoint:
id: str
vector: List[float]
payload: Dict[str, Any]
@dataclass(frozen=True)
class CollectionSpec:
name: str
vector_size: int
distance: str = COSINE
indexed_fields: Sequence[tuple] = INDEXED_PAYLOAD_FIELDS
def __post_init__(self) -> None:
if not self.name.strip():
raise ValueError("collection name is required")
if self.vector_size <= 0:
raise ValueError(f"vector_size must be positive, got {self.vector_size}")
@dataclass(frozen=True)
class CorpusManifest:
"""What a collection was built from — the thing A6 compares against.
Mixing two generations of the corpus, or two models' vectors, into one
collection produces no error at query time: every search still returns
something. This record is what makes that state detectable instead.
"""
corpus_sha256: str
chunk_count: int
model_id: str
dimensions: int
input_kind: str
provider: Optional[str] = None
distance: str = COSINE
extras: Dict[str, Any] = field(default_factory=dict)
def to_payload(self) -> Dict[str, Any]:
return {
"corpus_sha256": self.corpus_sha256,
"chunk_count": self.chunk_count,
"model_id": self.model_id,
"dimensions": self.dimensions,
"input_kind": self.input_kind,
"provider": self.provider,
"distance": self.distance,
**self.extras,
}
@classmethod
def from_payload(cls, payload: Mapping[str, Any]) -> "CorpusManifest":
known = {
"corpus_sha256",
"chunk_count",
"model_id",
"dimensions",
"input_kind",
"provider",
"distance",
}
return cls(
corpus_sha256=payload["corpus_sha256"],
chunk_count=payload["chunk_count"],
model_id=payload["model_id"],
dimensions=payload["dimensions"],
input_kind=payload["input_kind"],
provider=payload.get("provider"),
distance=payload.get("distance", COSINE),
extras={k: v for k, v in payload.items() if k not in known},
)
def conflicts_with(self, other: "CorpusManifest") -> List[str]:
"""Every reason these two must not share a collection."""
reasons = []
if self.corpus_sha256 != other.corpus_sha256:
reasons.append(
f"corpus sha256 {other.corpus_sha256[:12]} does not match the "
f"{self.corpus_sha256[:12]} this collection was built from"
)
if self.model_id != other.model_id:
reasons.append(
f"model {other.model_id} does not match the collection's "
f"{self.model_id}"
)
if self.dimensions != other.dimensions:
reasons.append(
f"{other.dimensions} dimensions do not match the collection's "
f"{self.dimensions}"
)
if self.input_kind != other.input_kind:
reasons.append(
f"input kind {other.input_kind} does not match the collection's "
f"{self.input_kind}"
)
return reasons
def _is_missing(value: Any) -> bool:
"""`0` and `False` are values; an empty string or empty list is not.
Physical page 0 and `has_quarantined_content=False` are both legitimate, so
a plain falsiness test would reject real records. Only `None` and empty
collections count as absent.
"""
if value is None:
return True
return isinstance(value, (str, list, tuple, dict, set)) and len(value) == 0
def validate_chunk_record(record: Mapping[str, Any]) -> None:
label = record.get("chunk_id", "<no chunk_id>")
missing = [f for f in REQUIRED_CHUNK_FIELDS if _is_missing(record.get(f))]
if missing:
raise ValueError(
f"chunk record {label!r} is missing "
f"required provenance fields: {missing}"
)
version = record.get("schema_version")
if type(version) is not int or version != SUPPORTED_SCHEMA_VERSION:
raise ValueError(
f"chunk record {label!r} declares schema_version {version!r}; "
f"this loader supports exactly v{SUPPORTED_SCHEMA_VERSION}; "
"unknown older or newer schemas are refused until compatibility "
"is implemented explicitly"
)
for field_name in PAGE_RANGE_FIELDS:
_validate_page_range(label, field_name, record.get(field_name))
def _validate_page_range(label: str, field_name: str, value: Any) -> None:
if not isinstance(value, (list, tuple)) or len(value) != 2:
raise ValueError(
f"chunk record {label!r} has {field_name}={value!r}; expected a "
f"[start, end] pair"
)
start, end = value
if type(start) is not int or type(end) is not int:
raise ValueError(
f"chunk record {label!r} has a non-integer page in "
f"{field_name}={value!r}"
)
if start > end:
raise ValueError(
f"chunk record {label!r} has {field_name}={value!r} running backwards"
)
minimum = 1 if field_name == "printed_page_range" else 0
if start < minimum:
raise ValueError(
f"chunk record {label!r} has {field_name}={value!r}; "
f"pages must start at {minimum} or later"
)
def build_point(
record: Mapping[str, Any], vector: Sequence[float]
) -> VectorPoint:
validate_chunk_record(record)
return VectorPoint(
id=point_id_for(record["chunk_id"]),
vector=list(vector),
payload=dict(record),
)
+52
View File
@@ -0,0 +1,52 @@
"""The vector-store boundary.
The loader, the corpus gate and every test in this package talk to this
Protocol and never to a database SDK. That is what lets the whole load stage be
exercised offline against `InMemoryVectorStore`, and what keeps `qdrant_client`
named in exactly one module — the same arrangement `embed/` uses to confine
boto3 to `bedrock_runtime`.
Deliberately primitive: collections, points, payload indexes. Anything with an
opinion — how a corpus is bound to a collection, how a chunk becomes a point —
is a layer above this, so swapping the store does not drag those rules with it.
"""
from __future__ import annotations
from typing import Any, Mapping, Optional, Protocol, Sequence
from .models import CollectionSpec, VectorPoint
class VectorStore(Protocol):
def collection_exists(self, name: str) -> bool: ...
def create_collection(self, spec: CollectionSpec) -> None: ...
def delete_collection(self, name: str) -> None: ...
def create_payload_index(
self, name: str, field_name: str, field_schema: str
) -> None: ...
def upsert(self, name: str, points: Sequence[VectorPoint]) -> int: ...
def count(self, name: str) -> int: ...
def retrieve(self, name: str, point_id: str) -> Optional[VectorPoint]: ...
def find_by_payload(
self, name: str, equals: Mapping[str, Any]
) -> Sequence[VectorPoint]:
"""Every point whose payload matches all of `equals`. No vector, no ranking.
This is mode A of `docs/v1-delivery-plan.md` §3, and it is a `scroll`
rather than a `search` on purpose. The rule the plan refuses to bend is
"return the whole section, not a top-k of fragments" — returning two of
five contraindications is more dangerous than returning none, because a
partial list reads as a complete one. A ranked search cannot express
that; an exhaustive filter can.
A list-valued field (`atc_codes`) matches when any element equals the
given value.
"""
...
+163
View File
@@ -0,0 +1,163 @@
"""The only module in this package that knows `qdrant_client` exists.
Same arrangement as `embed/bedrock_runtime.py` and for the same reason: the
loader's rules — corpus binding, derived point ids, batching, the point-count
gate — are exercised in full against `InMemoryVectorStore` with no server
running, and this file is the thin edge where those rules meet a real database.
The import sits inside `_client()` so `ingestion.load` imports cleanly on a
machine with no `qdrant-client` installed and no Qdrant reachable, which is the
state this repository is in by default.
What this adapter deliberately does not do is retry, shard, or tune. A load is
an offline batch run a human starts and watches; a failure should surface, not
be smoothed over into a partially loaded collection.
"""
from __future__ import annotations
from typing import Any, List, Mapping, Optional, Sequence
from .models import CollectionSpec, VectorPoint
DEFAULT_URL = "http://localhost:6333"
# Points fetched per scroll round-trip. Only affects round trips, never the
# result: `find_by_payload` pages until the collection says there is no more.
SCROLL_PAGE = 256
_DISTANCES = {"cosine": "Cosine", "dot": "Dot", "euclid": "Euclid"}
class QdrantVectorStore:
def __init__(
self,
url: str = DEFAULT_URL,
api_key: Optional[str] = None,
timeout: float = 60.0,
client: Optional[Any] = None,
) -> None:
self._url = url
self._api_key = api_key
self._timeout = timeout
self._client = client
@property
def url(self) -> str:
return self._url
def _client_or_connect(self) -> Any:
if self._client is None:
from qdrant_client import QdrantClient
self._client = QdrantClient(
url=self._url, api_key=self._api_key, timeout=self._timeout
)
return self._client
def collection_exists(self, name: str) -> bool:
client = self._client_or_connect()
existing = client.get_collections().collections
return any(collection.name == name for collection in existing)
def create_collection(self, spec: CollectionSpec) -> None:
from qdrant_client.models import Distance, VectorParams
distance = _DISTANCES.get(spec.distance.lower())
if distance is None:
raise ValueError(
f"unsupported distance {spec.distance!r}; "
f"expected one of {sorted(_DISTANCES)}"
)
self._client_or_connect().create_collection(
collection_name=spec.name,
vectors_config=VectorParams(
size=spec.vector_size, distance=Distance(distance)
),
)
def delete_collection(self, name: str) -> None:
self._client_or_connect().delete_collection(collection_name=name)
def create_payload_index(
self, name: str, field_name: str, field_schema: str
) -> None:
self._client_or_connect().create_payload_index(
collection_name=name,
field_name=field_name,
field_schema=field_schema,
)
def upsert(self, name: str, points: Sequence[VectorPoint]) -> int:
from qdrant_client.models import PointStruct
if not points:
return 0
self._client_or_connect().upsert(
collection_name=name,
points=[
PointStruct(id=point.id, vector=point.vector, payload=point.payload)
for point in points
],
wait=True,
)
return len(points)
def count(self, name: str) -> int:
return self._client_or_connect().count(
collection_name=name, exact=True
).count
def find_by_payload(
self, name: str, equals: Mapping[str, Any]
) -> List[VectorPoint]:
from qdrant_client.models import FieldCondition, Filter, MatchValue
if not equals:
raise ValueError("find_by_payload needs at least one condition")
scroll_filter = Filter(
must=[
FieldCondition(key=key, match=MatchValue(value=value))
for key, value in equals.items()
]
)
client = self._client_or_connect()
found: List[VectorPoint] = []
offset = None
while True:
page, offset = client.scroll(
collection_name=name,
scroll_filter=scroll_filter,
limit=SCROLL_PAGE,
offset=offset,
with_payload=True,
with_vectors=False,
)
found.extend(
VectorPoint(
id=str(record.id),
vector=[],
payload=dict(record.payload or {}),
)
for record in page
)
# Paging must not stop early. A section split into more parts than
# one page holds would come back truncated, and a truncated section
# is the one failure mode mode A exists to prevent.
if offset is None:
return found
def retrieve(self, name: str, point_id: str) -> Optional[VectorPoint]:
found = self._client_or_connect().retrieve(
collection_name=name,
ids=[point_id],
with_payload=True,
with_vectors=True,
)
if not found:
return None
record = found[0]
return VectorPoint(
id=str(record.id),
vector=list(record.vector or []),
payload=dict(record.payload or {}),
)
+124
View File
@@ -0,0 +1,124 @@
"""Corpus embed + load entry point.
`cli.py` belongs to the parser/chunking work, so this is the standalone entry
point promised to Codex in `coordination/CLAUDE_TASK_2026-08-04.md` §4.1:
python -m ingestion.load.run --provider cohere-v4 --collection duocthu_v1
Two phases, deliberately separable. Embedding is the only part that leaves the
machine, so it goes through the disk cache: an interrupted run resumes from
whatever it already paid for instead of re-embedding it. Loading then reads
that cache and never calls a provider at all.
Vectors are keyed by `(model_id, input_kind, sha256(text))`, so a re-run after
an unrelated chunk edit re-embeds only the texts that actually changed.
"""
from __future__ import annotations
import argparse
import time
from pathlib import Path
from typing import Iterator, List, Sequence, Tuple
from ..embed.cache import CachingEmbeddingProvider, EmbeddingCache
from ..embed.ports import INPUT_DOCUMENT
from ..embed.registry import DEFAULT_REGION, build_provider
from .corpus import corpus_sha256, count_chunks, iter_chunk_records
from .models import CollectionSpec, CorpusManifest
from .qdrant_repo import QdrantVectorStore
from .upsert import ChunkLoader
DEFAULT_SLICE = 960
DEFAULT_ATTEMPTS = 3
def _slices(items: Sequence, size: int) -> Iterator[Tuple[int, Sequence]]:
for start in range(0, len(items), size):
yield start, items[start : start + size]
def _embed_with_retry(provider, texts: Sequence[str], attempts: int) -> List:
last: Exception | None = None
for attempt in range(1, attempts + 1):
try:
return list(provider.embed(texts, INPUT_DOCUMENT).vectors)
except Exception as exc: # noqa: BLE001 — a flaky link is the norm here
last = exc
if attempt == attempts:
break
backoff = 2**attempt
print(f" attempt {attempt} failed ({exc}); retrying in {backoff}s")
time.sleep(backoff)
raise RuntimeError(f"embedding failed after {attempts} attempts") from last
def main(argv=None) -> int:
parser = argparse.ArgumentParser(prog="python -m ingestion.load.run")
parser.add_argument("--chunks", type=Path, default=Path("data/processed/chunks.jsonl"))
parser.add_argument("--cache", type=Path, default=Path("data/processed/embeddings"))
parser.add_argument("--provider", required=True)
parser.add_argument("--collection", required=True)
parser.add_argument("--region", default=DEFAULT_REGION)
parser.add_argument("--qdrant-url", default="http://localhost:6333")
parser.add_argument("--slice-size", type=int, default=DEFAULT_SLICE)
parser.add_argument("--attempts", type=int, default=DEFAULT_ATTEMPTS)
parser.add_argument("--embed-only", action="store_true")
args = parser.parse_args(argv)
records = list(iter_chunk_records(args.chunks))
sha = corpus_sha256(args.chunks)
print(f"corpus : {args.chunks}")
print(f"chunks : {len(records)} (count_chunks={count_chunks(args.chunks)})")
print(f"sha256 : {sha}")
inner = build_provider(args.provider, region=args.region)
cache_path = args.cache / f"{args.provider}.jsonl"
cache_path.parent.mkdir(parents=True, exist_ok=True)
provider = CachingEmbeddingProvider(inner, EmbeddingCache(cache_path))
print(f"provider : {provider.model_id} ({provider.dimensions}d)")
print(f"cache : {cache_path}")
texts = [str(record["text"]) for record in records]
vectors: List = []
started = time.time()
for start, slice_texts in _slices(texts, args.slice_size):
vectors.extend(_embed_with_retry(provider, slice_texts, args.attempts))
done = len(vectors)
rate = done / max(time.time() - started, 1e-9)
remaining = (len(texts) - done) / rate if rate else 0
print(
f" embedded {done}/{len(texts)} "
f"({rate:.1f}/s, ~{remaining/60:.1f} min left)",
flush=True,
)
stats = provider.cache.stats
print(f"cache : {stats.hits} hits, {stats.misses} misses")
if args.embed_only:
print("embed-only: stopping before the vector store")
return 0
manifest = CorpusManifest(
corpus_sha256=sha,
chunk_count=len(records),
model_id=provider.model_id,
dimensions=provider.dimensions,
input_kind=INPUT_DOCUMENT,
provider=args.provider,
)
spec = CollectionSpec(name=args.collection, vector_size=provider.dimensions)
store = QdrantVectorStore(url=args.qdrant_url)
loader = ChunkLoader(store, spec, manifest)
report = loader.load(zip(records, (v.values for v in vectors)))
print(
f"loaded : {report.points_upserted} points in {report.batches} batches; "
f"collection holds {report.collection_count}; "
f"count gate {'PASS' if report.count_matches else 'FAIL'}"
)
return 0 if report.count_matches else 1
if __name__ == "__main__":
raise SystemExit(main())
+129
View File
@@ -0,0 +1,129 @@
"""Loading chunks into a collection, idempotently and only into the right one.
The loader does three things in a fixed order, and the order is the point.
It checks the corpus binding *before* creating or writing anything (A6), so a
refused load leaves the store untouched rather than half-overwritten. It derives
every point id from `chunk_id` (A5), so the same corpus loaded twice converges
instead of doubling. And it reports the collection's point count against the
number of chunks it was given, which is the v1 gate
`qdrant_point_count != chunk_count`.
Vectors are validated against the collection's declared size before the first
upsert. A wrong-sized vector is a whole-run defect, not a bad record, and
failing on the first one is cheaper than discovering it after 15,066 upserts.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Iterable, List, Mapping, Sequence, Tuple
from .manifest import assert_compatible, write_manifest
from .models import CollectionSpec, CorpusManifest, VectorPoint, build_point
from .ports import VectorStore
ChunkVectorPair = Tuple[Mapping[str, Any], Sequence[float]]
DEFAULT_BATCH_SIZE = 256
class PointCountMismatch(RuntimeError):
"""The collection does not hold exactly one point per chunk."""
@dataclass(frozen=True)
class LoadReport:
collection: str
collection_created: bool
points_upserted: int
batches: int
collection_count: int
corpus_sha256: str
@property
def count_matches(self) -> bool:
return self.collection_count == self.points_upserted
class ChunkLoader:
def __init__(
self,
store: VectorStore,
spec: CollectionSpec,
manifest: CorpusManifest,
batch_size: int = DEFAULT_BATCH_SIZE,
) -> None:
if batch_size <= 0:
raise ValueError(f"batch_size must be positive, got {batch_size}")
if manifest.dimensions != spec.vector_size:
raise ValueError(
f"manifest declares {manifest.dimensions} dimensions but the "
f"collection spec declares {spec.vector_size}"
)
self._store = store
self._spec = spec
self._manifest = manifest
self._batch_size = batch_size
def prepare(self) -> bool:
"""Gate on the corpus binding, then ensure the collection exists.
Returns True when the collection was created by this call.
"""
assert_compatible(self._store, self._spec.name, self._manifest)
created = False
if not self._store.collection_exists(self._spec.name):
self._store.create_collection(self._spec)
for field_name, field_schema in self._spec.indexed_fields:
self._store.create_payload_index(
self._spec.name, field_name, field_schema
)
created = True
write_manifest(self._store, self._spec.name, self._manifest)
return created
def load(self, pairs: Iterable[ChunkVectorPair]) -> LoadReport:
created = self.prepare()
upserted = 0
batches = 0
buffer: List[VectorPoint] = []
for record, vector in pairs:
buffer.append(self._point(record, vector))
if len(buffer) >= self._batch_size:
upserted += self._store.upsert(self._spec.name, buffer)
batches += 1
buffer = []
if buffer:
upserted += self._store.upsert(self._spec.name, buffer)
batches += 1
return LoadReport(
collection=self._spec.name,
collection_created=created,
points_upserted=upserted,
batches=batches,
collection_count=self._store.count(self._spec.name),
corpus_sha256=self._manifest.corpus_sha256,
)
def assert_point_count(self, expected_chunks: int) -> int:
"""v1 gate: exactly one point per chunk, no more and no fewer."""
actual = self._store.count(self._spec.name)
if actual != expected_chunks:
raise PointCountMismatch(
f"collection {self._spec.name!r} holds {actual} points but the "
f"corpus has {expected_chunks} chunks"
)
return actual
def _point(
self, record: Mapping[str, Any], vector: Sequence[float]
) -> VectorPoint:
if len(vector) != self._spec.vector_size:
raise ValueError(
f"chunk {record.get('chunk_id', '<no chunk_id>')!r} has a "
f"{len(vector)}-dimension vector, collection "
f"{self._spec.name!r} expects {self._spec.vector_size}"
)
return build_point(record, vector)
+187 -18
View File
@@ -30,7 +30,7 @@ from typing import Iterator, List, Optional, Union
from ..extract.models import Span
from ..extract.page_map import HEADER_BAND_Y
from ..normalize import join_spans, substitute_pua
from ..tables.classify import QUARANTINE_SHAPES
from ..tables.classify import QUARANTINE_SHAPES, SHAPE_FORMULA_2D
from .atc import extract_atc_codes
from .detector import in_monograph_range, is_monograph_title_candidate
from .merge import merge_multiline_headings, merge_same_line_bold_fragments
@@ -51,6 +51,9 @@ from .vocab import (
)
_QUALIFIER_RE = re.compile(r"^\(.+\)$")
_DOSING_TABLE_CAPTION_RE = re.compile(
r"(?:^|\b)Bảng\s+\d+\s*[.:]\s*Điều chỉnh liều\b", re.IGNORECASE,
)
def _is_page_boilerplate(span: Span) -> bool:
@@ -94,6 +97,31 @@ def _slugify(text: str) -> str:
return re.sub(r"[^a-z0-9]+", "_", ascii_text.lower()).strip("_")
def _starts_its_visual_line(span: Span, previous: Span | None) -> bool:
"""True when nothing else was printed to the left of this span on its line.
A real section heading opens a line. Confirmed content loss when this was
not checked: CISPLATIN (physical page 402) prints `Suy thận: Chống chỉ
định.` inside its dosing section, and the second half is a section name.
Matched as a heading, `Chống chỉ định.` vanished from the dosing text and
the section ended on a bare `Suy thận:` — a renal-impairment
contraindication silently dropped. ISOPRENALIN had the same shape.
"To the left" is decided on `x0`, not on `previous.x1 <= span.x0`: NEVIRAPIN
(physical page 1045) prints `Xem thêm mục ` at x1=104.89 immediately before
an italic `Liều lượng và cách dùng` at x0=104.88 — a 0.01pt overlap from the
trailing space's advance width, enough to make an end-before-start test call
a mid-line cross-reference a heading. Comparing the left edges cannot be
defeated by glyph-advance rounding and still reads False for the synthetic
fixtures, which place every span at identical coordinates.
"""
if previous is None:
return True
same_line = (previous.physical_page, previous.block, previous.line) == (
span.physical_page, span.block, span.line)
return not (same_line and previous.x0 < span.x0)
def _is_body_line_that_reads_like_a_label(span: Span, items: List) -> bool:
"""A plain line that repeats a section name, sitting under a heading.
@@ -114,7 +142,48 @@ def _is_body_line_that_reads_like_a_label(span: Span, items: List) -> bool:
return bool(items) and isinstance(items[-1], _SectionEvent)
def _classify(spans: List[Span]) -> List[Union[Span, _SectionEvent, _TextEvent]]:
def _is_mid_line_label(span: Span, previous: Span | None) -> bool:
"""A non-bold section name printed part-way along a line is body text."""
return not span.bold and not _starts_its_visual_line(span, previous)
def _continues_previous_visual_line(span: Span, previous: Span | None) -> bool:
"""A plain label-shaped span can be the wrapped tail of body prose.
Confirmed across the corpus: sentences such as ``không phải là`` /
``chống chỉ định.`` remain in the same PDF block on adjacent visual
lines. Exact vocabulary matching used to consume the second line as a
heading. A real plain heading observed in this book follows completed
prose; a non-bold adjacent continuation after an unterminated line does
not. Styling and PDF block geometry make this deliberately narrower than
a text-only language heuristic.
"""
if span.bold or previous is None:
return False
same_block = (previous.physical_page, previous.block) == (
span.physical_page, span.block)
adjacent_line = span.line == previous.line + 1
previous_text = previous.text.rstrip()
terminal = previous_text.endswith((".", "!", "?", ":", ";"))
return same_block and adjacent_line and bool(previous_text) and not terminal
def _is_italic_cross_reference(span: Span) -> bool:
"""The book italicises references to other sections; headings are never italic.
Measured over the whole monograph range (physical 99-1496): 11,916 spans
carrying a section name are bold, 29 are plain, and **7 are italic — none
of them a heading**. Two of the seven wrap onto a line of their own, where
the "opens its line" test cannot help: CALCI LACTAT (p296) breaks `xem thêm
về nhu cầu hàng ngày trong mục ` / `Dược lý và cơ chế tác dụng)`, and
physical page 432 breaks ordinary prose — `hoặc khi không được ` /
`chỉ định.` — across a line, so `chỉ định.` alone opened a spurious
"Chỉ định" section.
"""
return span.italic
def _classify(spans: List[Span], table_index=None) -> List[Union[Span, _SectionEvent, _TextEvent]]:
"""Pass 1: tag each span. Title candidates are left as raw Span objects
(pass 2 groups + merges them); everything else becomes a typed event.
@@ -127,29 +196,110 @@ def _classify(spans: List[Span]) -> List[Union[Span, _SectionEvent, _TextEvent]]
same lesson as "don't gate on font size" (ADR 0003 item 10) applied to
boldness instead.
"""
# PyMuPDF's block order is not guaranteed to keep every cell of a table
# together. A visually later section heading can therefore occur between
# cells of one physical region in the extracted stream (confirmed on
# CAPECITABIN pp. 308-309 and IMATINIB p. 795). Gather each region first,
# then emit it atomically at its first occurrence. Besides preserving the
# section active at the top of the table, this guarantees one run/block per
# physical region instead of duplicate fragments with conflicting owners.
region_spans: dict[int, List[Span]] = {}
region_for_span: dict[int, object] = {}
spans_by_page: dict[int, List[Span]] = {}
for span in spans:
spans_by_page.setdefault(span.physical_page, []).append(span)
region = _region_for(table_index, span)
if region is None:
continue
region_for_span[id(span)] = region
region_spans.setdefault(id(region), []).append(span)
# A table can continue after the monograph's ordinary final headings.
# CAPECITABIN p. 309 does exactly that: two dose-adjustment tables follow
# the trade-name line, with their own explicit captions but without a
# repeated "Liều lượng và cách dùng" heading. Use only the narrow,
# unambiguous caption signal and its geometry; a generic keyword rule
# would be unsafe in clinical prose.
caption_for_region: dict[int, Span] = {}
region_for_caption: dict[int, object] = {}
for page_regions in (table_index or {}).values():
for region in page_regions:
candidates = [
span for span in spans_by_page.get(region.physical_page, ())
if 0 <= region.bbox[1] - span.y1 <= 80
and _DOSING_TABLE_CAPTION_RE.search(span.text.strip())
]
if candidates:
caption = max(candidates, key=lambda candidate: candidate.y1)
caption_for_region[id(region)] = caption
region_for_caption[id(caption)] = region
items: List[Union[Span, _SectionEvent, _TextEvent]] = []
emitted_regions: set[int] = set()
emitted_captions: set[int] = set()
previous: Span | None = None
for span in spans:
if not span.text.strip():
continue
if _is_page_boilerplate(span):
previous = span
continue
caption_region = region_for_caption.get(id(span))
if caption_region is not None:
caption_key = id(span)
if caption_key not in emitted_captions:
dose_section = match_section("Liều lượng và cách dùng")
assert dose_section is not None
items.append(_SectionEvent(dose_section, span, inline_value=span.text))
emitted_captions.add(caption_key)
previous = span
continue
# Out-of-scope index text must never create section events, and cells
# inside a known table are content even when a cell says "Chỉ định".
# Classification previously happened without either context, letting
# the WARFARIN/IOBITRIDOL table header change the owning section.
region = region_for_span.get(id(span))
if region is not None:
region_key = id(region)
if region_key not in emitted_regions:
caption = caption_for_region.get(region_key)
if caption is not None and id(caption) not in emitted_captions:
dose_section = match_section("Liều lượng và cách dùng")
assert dose_section is not None
items.append(_SectionEvent(dose_section, caption,
inline_value=caption.text))
emitted_captions.add(id(caption))
items.extend(_TextEvent(cell) for cell in region_spans[region_key]
if cell.text.strip() and not _is_page_boilerplate(cell))
emitted_regions.add(region_key)
previous = span
continue
if not in_monograph_range(span):
items.append(_TextEvent(span))
previous = span
continue
if is_monograph_title_candidate(span):
items.append(span)
previous = span
continue
section_def = match_section(span.text)
if section_def is not None and not _is_body_line_that_reads_like_a_label(
span, items
):
items.append(_SectionEvent(section_def, span))
continue
if section_def is not None:
items.append(_TextEvent(span))
is_body = (_is_body_line_that_reads_like_a_label(span, items)
or _is_mid_line_label(span, previous)
or _continues_previous_visual_line(span, previous)
or _is_italic_cross_reference(span))
items.append(_TextEvent(span) if is_body
else _SectionEvent(section_def, span))
previous = span
continue
inline = match_section_with_inline_value(span.text)
if inline is not None:
if (inline is not None and not _is_mid_line_label(span, previous)
and not _continues_previous_visual_line(span, previous)
and not _is_italic_cross_reference(span)):
items.append(_SectionEvent(inline[0], span, inline_value=inline[1]))
else:
items.append(_TextEvent(span))
previous = span
return items
@@ -228,6 +378,17 @@ def _region_for(table_index, span: Span):
if not table_index:
return None
for region in table_index.get(span.physical_page, ()):
if region.shape == SHAPE_FORMULA_2D and span.column in ("left", "right"):
# Formula bands are intentionally widened enough to reach past
# the gutter (some numerator spans have misleading leading-space
# boxes), so geometry alone can swallow prose from the opposite
# column. Keep the wide band but require its source column to
# agree with the span's extracted block column. This preserves
# AMPICILIN's gutter-adjacent "Cl" while excluding NETILMICIN's
# right-column cross-reference from a left-column formula.
region_column = "left" if (region.bbox[0] + region.bbox[2]) / 2 < 303.5 else "right"
if span.column != region_column:
continue
if region.contains(span.x0, span.y0, span.x1, span.y1):
return region
return None
@@ -259,7 +420,7 @@ def assemble(spans: List[Span], table_index=None, ledger: Optional[list] = None)
"""
raw_chars = sum(len(s.text) for s in spans)
spans = merge_same_line_bold_fragments(spans)
events = _filter_false_positive_titles(_coalesce_titles(_classify(spans)))
events = _filter_false_positive_titles(_coalesce_titles(_classify(spans, table_index)))
# Span-level coverage ledger. Character counts alone cannot balance here
# (normalization joins, substitutes and drops characters), so every span
@@ -291,7 +452,7 @@ def assemble(spans: List[Span], table_index=None, ledger: Optional[list] = None)
current: Optional[Monograph] = None
current_section_key: Optional[str] = None
runs: List[tuple] = [] # ordered [(region_or_None, [spans])]
inline_prefix: str = ""
inline_prefix: Optional[tuple[str, Span]] = None
awaiting_qualifier = False
def append_span(span: Span, region):
@@ -338,10 +499,12 @@ def assemble(spans: List[Span], table_index=None, ledger: Optional[list] = None)
if region is not None else False,
))
if inline_prefix:
inline_text, inline_span = inline_prefix
head = SectionPart(
kind=PART_PROSE, text=inline_prefix,
physical_page=parts[0].physical_page if parts else 0,
bbox=parts[0].bbox if parts else [0.0, 0.0, 0.0, 0.0],
kind=PART_PROSE, text=inline_text,
physical_page=inline_span.physical_page,
bbox=[inline_span.x0, inline_span.y0, inline_span.x1, inline_span.y1],
source_span_ids=[inline_span.span_id],
)
parts.insert(0, head)
return parts
@@ -387,7 +550,7 @@ def assemble(spans: List[Span], table_index=None, ledger: Optional[list] = None)
source_span_ids=list(part.source_span_ids),
))
runs = []
inline_prefix = ""
inline_prefix = None
def finalize(monograph: Monograph) -> Monograph:
# Duplicate check happens here, not at title-detection time: the
@@ -401,8 +564,14 @@ def assemble(spans: List[Span], table_index=None, ledger: Optional[list] = None)
f"line (outlier item 18) before assuming this is a real collision"
)
seen_ids.add(monograph.drug_id)
if "ma_atc" in monograph.sections:
result = extract_atc_codes(monograph.sections["ma_atc"].text)
atc_source = monograph.sections.get("ma_atc")
# One real class monograph combines the two labels as "Tên chung
# quốc tế và mã ATC". It remains the required title anchor, and
# its ATC codes are still extracted rather than silently discarded.
if atc_source is None:
atc_source = monograph.sections.get("ten_chung_quoc_te")
if atc_source is not None:
result = extract_atc_codes(atc_source.text)
monograph.atc_codes = result.codes
monograph.atc_stated_absent = result.stated_absent
return monograph
@@ -440,7 +609,7 @@ def assemble(spans: List[Span], table_index=None, ledger: Optional[list] = None)
text="",
)
if event.inline_value:
inline_prefix = event.inline_value
inline_prefix = (event.inline_value, event.span)
awaiting_qualifier = False
mark(event.span, SPAN_STATE_HEADING)
current.source_page_range[1] = max(current.source_page_range[1], event.span.physical_page)
+9
View File
@@ -30,6 +30,13 @@ from .vocab import is_part_divider, match_section
MONOGRAPH_PRINTED_PAGE_START = 99
MONOGRAPH_PRINTED_PAGE_END = 1496
# Printed folios are metadata inferred from page headers and can be wrong in
# the back index (confirmed: physical page 1655 was inferred as printed page
# 1496). The physical bounds are source-document invariants and prevent an
# index entry that looks like a section label from extending ZOLPIDEM by 161
# pages. Keep both checks: either signal alone has known failure modes.
MONOGRAPH_PHYSICAL_PAGE_START = 99
MONOGRAPH_PHYSICAL_PAGE_END = 1496
_MIN_TITLE_LEN = 3
_MAX_TITLE_LEN = 60
_MAX_LOWERCASE_RATIO = 0.10 # HMG-CoA: 1/27 = 3.7% (real title) vs "Mã ATC:": 1/5 = 20% (real
@@ -49,6 +56,8 @@ def _is_mostly_upper(text: str) -> bool:
def in_monograph_range(span: Span) -> bool:
return (
MONOGRAPH_PHYSICAL_PAGE_START <= span.physical_page <= MONOGRAPH_PHYSICAL_PAGE_END
and
span.printed_page is not None
and MONOGRAPH_PRINTED_PAGE_START <= span.printed_page <= MONOGRAPH_PRINTED_PAGE_END
)
+5 -3
View File
@@ -37,11 +37,12 @@ class SectionDef:
SECTION_DEFS = [
SectionDef("ten_chung_quoc_te", "Tên chung quốc tế", ("Ten chung quốc tế",)),
SectionDef("ten_chung_quoc_te", "Tên chung quốc tế",
("Ten chung quốc tế", "Tên chung quốc tế và mã ATC")),
SectionDef("ma_atc", "Mã ATC", ("Mã ACT",)),
SectionDef("loai_thuoc", "Loại thuốc", ("Loại thuôc", "Lọai thuốc", "Phân loại thuốc")),
SectionDef("dang_thuoc_va_ham_luong", "Dạng thuốc và hàm lượng",
("Dạng dùng và hàm lượng",)),
("Dạng dùng và hàm lượng", "Dạng bào chế và hàm lượng")),
SectionDef("duoc_ly_va_co_che_tac_dung", "Dược lý và cơ chế tác dụng",
("Dược lí và cơ chế tác dụng", "Dược lý học và cơ chế tác dụng")),
SectionDef("chi_dinh", "Chỉ định"),
@@ -56,7 +57,8 @@ SECTION_DEFS = [
"Hướng dẫn cách xử trí các ADR")),
SectionDef("lieu_luong_va_cach_dung", "Liều lượng và cách dùng",
("Liều lượng cách dùng", "Liều lượng, cách dùng",
"Liều dùng và cách dùng", "Liều lượng và cách sử dụng")),
"Liều dùng và cách dùng", "Liều lượng và cách sử dụng",
"Liều lượng và cách dùng giải độc tố uốn ván hấp phụ đơn giá")),
SectionDef("tuong_tac_thuoc", "Tương tác thuốc"),
SectionDef("do_on_dinh_va_bao_quan", "Độ ổn định và bảo quản"),
SectionDef("tuong_ky", "Tương kỵ"),
@@ -8,6 +8,7 @@ from .readiness import (
read_chunks,
read_monographs,
)
from .clinical_readiness import evaluate_clinical
from .residual_ink import (
ANTIALIAS_SPECK,
FRACTION_BAR_CANDIDATE,
@@ -35,6 +36,7 @@ __all__ = [
"read_chunks",
"corpus_size",
"read_monographs",
"evaluate_clinical",
"PageContext",
"ResidualRegion",
"classify",
+95 -8
View File
@@ -23,7 +23,6 @@ from __future__ import annotations
import re
from dataclasses import dataclass
from typing import List
import fitz
@@ -32,6 +31,26 @@ BACK_INDEX_START_PHYSICAL = 1530 # printed 1531 — first page of real entries
_ENTRY_RE = re.compile(r"^(.+?),\s*(\d+)\s*$")
_CROSS_REF_MARKER = " - "
# Page furniture inside the index: the running header, the index's own title,
# and the single-letter section dividers. These used to be dropped implicitly
# by not matching the entry pattern — once wrapped lines are rejoined they
# would instead be glued onto the entry below them, so they now have to be
# named.
_FURNITURE_RE = re.compile(r"^(DTQGVN\s*\d*|Mục lục tra cứu|[A-ZĐÀ-Ỹ])$")
# A bare number is ambiguous: the page's own folio, or the tail of an entry
# whose page number wrapped onto the next line (`Bromhexine hydrochloride -
# Bromhexin hydroclorid,` / `269`). Discarding it unconditionally cost real
# entries — the unterminated fragment then swallowed the *following* entry, so
# `Bromocriptin, 270` was consumed into a cross-reference and vanished from the
# ground truth. It is furniture only when no fragment is waiting for it.
_BARE_NUMBER_RE = re.compile(r"^\d{1,4}$")
# An index entry never wraps into a paragraph; the longest genuine one measured
# in this book is well under this. A buffer that grows past it means the join
# has lost the thread, so it is dropped rather than emitted as a bogus name.
_MAX_JOINED_ENTRY_CHARS = 200
@dataclass(frozen=True)
class GroundTruthEntry:
@@ -39,14 +58,82 @@ class GroundTruthEntry:
printed_page: int
def parse_back_index(doc: fitz.Document, start_physical_page: int = BACK_INDEX_START_PHYSICAL) -> List[GroundTruthEntry]:
entries: List[GroundTruthEntry] = []
@dataclass(frozen=True)
class BackIndexAlias:
alias: str
target: str
printed_page: int
physical_page: int
def _index_page_lines(doc: fitz.Document, start_physical_page: int):
"""Index lines page by page, furniture removed. An entry never wraps across
a page, so the caller can drop a pending fragment at each page break."""
for pno in range(start_physical_page, doc.page_count):
for line in doc[pno].get_text().split("\n"):
line = line.strip()
if not line or _CROSS_REF_MARKER in line:
lines = [line.strip() for line in doc[pno].get_text().split("\n")]
yield [line for line in lines if line and not _FURNITURE_RE.match(line)]
def parse_back_index(doc: fitz.Document, start_physical_page: int = BACK_INDEX_START_PHYSICAL) -> list[GroundTruthEntry]:
"""Entries in reading order, cross-references excluded.
Long entries wrap across two printed lines, and each fragment was read as
an entry of its own before this was handled: `Acinet 10 - xem Atorvastatin
- Các chất ức chế HMG - ` / `CoA reductase, 285` produced a phantom
ground-truth drug called "- CoA reductase". Measured on the whole index,
that shape accounted for 50 of the 76 entries `cli validate` could not
match — noise that inflates the denominator and hides real misses inside
it. A fragment is a line that does not yet end in ", <page>", so lines are
accumulated until they do.
"""
entries: list[GroundTruthEntry] = []
for page_lines in _index_page_lines(doc, start_physical_page):
buffer = ""
for line in page_lines:
if not buffer and _BARE_NUMBER_RE.match(line):
continue
match = _ENTRY_RE.match(line)
buffer = f"{buffer} {line}".strip() if buffer else line
match = _ENTRY_RE.match(buffer)
if match:
entries.append(GroundTruthEntry(name=match.group(1).strip(), printed_page=int(match.group(2))))
if _CROSS_REF_MARKER not in buffer:
entries.append(GroundTruthEntry(
name=match.group(1).strip(), printed_page=int(match.group(2))))
buffer = ""
elif len(buffer) > _MAX_JOINED_ENTRY_CHARS:
buffer = ""
return entries
def parse_back_index_see_aliases(
doc: fitz.Document,
start_physical_page: int = BACK_INDEX_START_PHYSICAL,
) -> list[BackIndexAlias]:
"""Extract the book's explicit ``X - xem Y`` alias relations.
This is intentionally separate from :func:`parse_back_index`, whose
validation semantics must remain unchanged.
"""
aliases: list[BackIndexAlias] = []
marker = re.compile(r"\s+-\s+xem\s+", re.IGNORECASE)
for physical_page, page_lines in enumerate(
_index_page_lines(doc, start_physical_page), start_physical_page,
):
buffer = ""
for line in page_lines:
if not buffer and _BARE_NUMBER_RE.match(line):
continue
buffer = f"{buffer} {line}".strip() if buffer else line
match = _ENTRY_RE.match(buffer)
if match:
parts = marker.split(match.group(1), maxsplit=1)
if len(parts) == 2 and all(part.strip() for part in parts):
aliases.append(BackIndexAlias(
alias=parts[0].strip(),
target=parts[1].strip(),
printed_page=int(match.group(2)),
physical_page=physical_page,
))
buffer = ""
elif len(buffer) > _MAX_JOINED_ENTRY_CHARS:
buffer = ""
return aliases
@@ -0,0 +1,106 @@
"""Executable gates for a clinical-production release.
These gates deliberately sit above parser/chunk readiness. Passing extraction
tests cannot establish that the source is current, licensed, clinically
reviewed, or safe to operate in a care setting.
"""
from __future__ import annotations
import hashlib
import json
from pathlib import Path
from typing import List
from .readiness import Gate
CURRENT_FORMULARY_EDITION = 3
REQUIRED_APPROVAL_ROLES = {
"physician",
"clinical_pharmacist",
"clinical_safety_owner",
"regulatory_owner",
}
REQUIRED_RELEASE_ARTIFACTS = {
"logical_tables": "data/processed/logical_tables.jsonl",
"clinical_eval": "data/clinical/clinical_eval_report.json",
"risk_management": "data/clinical/risk_management.json",
"security_privacy": "data/clinical/security_privacy_review.json",
"operations": "data/clinical/operations_readiness.json",
}
def read_json(path: Path) -> dict:
return json.loads(path.read_text(encoding="utf-8"))
def file_sha256(path: Path) -> str:
digest = hashlib.sha256()
with path.open("rb") as handle:
for chunk in iter(lambda: handle.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def evaluate_clinical(
project_root: Path,
source_manifest: dict,
approval_manifest: dict | None = None,
) -> List[Gate]:
pdf_path = project_root / source_manifest.get("pdf_path", "")
expected_sha = source_manifest.get("sha256", "")
actual_sha = file_sha256(pdf_path) if pdf_path.is_file() else ""
approvals = approval_manifest or {}
approved_roles = {
item.get("role") for item in approvals.get("approvals", [])
if item.get("approved") and item.get("reviewer") and item.get("date")
}
artifacts = approvals.get("artifacts", {})
gates = [
Gate(
"current_national_formulary_edition",
int(source_manifest.get("edition") == CURRENT_FORMULARY_EDITION),
target=1,
detail=(
f"found edition {source_manifest.get('edition')}; "
f"required edition {CURRENT_FORMULARY_EDITION}"
),
),
Gate(
"source_pdf_sha256_matches_manifest",
int(bool(expected_sha) and actual_sha == expected_sha),
target=1,
detail="source PDF missing or hash mismatch" if actual_sha != expected_sha else "",
),
Gate(
"production_use_rights_documented",
int(bool(source_manifest.get("production_use_rights_documented"))),
target=1,
),
Gate(
"source_marked_clinical_production_eligible",
int(bool(source_manifest.get("clinical_production_eligible"))),
target=1,
),
]
for name, default_path in REQUIRED_RELEASE_ARTIFACTS.items():
configured = artifacts.get(name, default_path)
path = project_root / configured
gates.append(Gate(f"artifact_{name}", int(path.is_file()), target=1,
detail=str(configured)))
missing_roles = sorted(REQUIRED_APPROVAL_ROLES - approved_roles)
gates.append(Gate(
"required_clinical_release_approvals",
len(approved_roles & REQUIRED_APPROVAL_ROLES),
target=len(REQUIRED_APPROVAL_ROLES),
detail="missing: " + ", ".join(missing_roles) if missing_roles else "",
))
gates.append(Gate(
"intended_use_and_regulatory_classification_approved",
int(bool(approvals.get("intended_use_approved")) and
bool(approvals.get("regulatory_classification_approved"))),
target=1,
))
return gates
+161 -1
View File
@@ -17,6 +17,13 @@ from dataclasses import dataclass
from pathlib import Path
from typing import Dict, Iterable, List, Sequence
import re as _re
_WHITESPACE = _re.compile(r"\s+")
# Head of a source line, enough to identify it without demanding that an
# intentionally comma-split list match in full.
SENTENCE_PROBE_CHARS = 60
PUA_RANGE = (0xE000, 0xF8FF)
REPLACEMENT_CHAR = ""
@@ -66,11 +73,12 @@ def _count_pua(text: str) -> int:
def evaluate(monographs: Sequence[dict],
transcribed_runs: Sequence[dict] = ()) -> List[Gate]:
"""Compute every readiness gate over the whole corpus."""
pua = replacement = empty = no_provenance = 0
pua = replacement = empty = no_provenance = part_no_provenance = 0
corruptions: Dict[str, int] = {c: 0 for c in KNOWN_CORRUPTIONS}
formula_leaks: Dict[str, int] = {f: 0 for f in FORMULA_FRAGMENTS}
unflagged_blocks = 0
ids: Dict[str, int] = {}
table_ids: Dict[str, int] = {}
no_page_range = 0
corpus = []
@@ -85,6 +93,9 @@ def evaluate(monographs: Sequence[dict],
empty += 1
if not section.get("parts"):
no_provenance += 1
for part in section.get("parts") or []:
if not part.get("source_span_ids"):
part_no_provenance += 1
pua += _count_pua(text)
replacement += text.count(REPLACEMENT_CHAR)
for phrase in KNOWN_CORRUPTIONS:
@@ -92,6 +103,9 @@ def evaluate(monographs: Sequence[dict],
for phrase in FORMULA_FRAGMENTS:
formula_leaks[phrase] += text.count(phrase)
for block in monograph.get("tables") or []:
table_id = block.get("table_id")
if table_id:
table_ids[table_id] = table_ids.get(table_id, 0) + 1
if not block.get("quarantined"):
unflagged_blocks += 1
@@ -113,7 +127,9 @@ def evaluate(monographs: Sequence[dict],
Gate("replacement_char_ufffd", replacement),
Gate("empty_section", empty),
Gate("section_without_provenance", no_provenance),
Gate("part_without_source_span_ids", part_no_provenance),
Gate("unflagged_quarantine_block", unflagged_blocks),
Gate("duplicate_table_id", sum(1 for n in table_ids.values() if n > 1)),
Gate("duplicate_drug_id", sum(1 for n in ids.values() if n > 1)),
Gate("monograph_without_page_range", no_page_range),
]
@@ -147,7 +163,10 @@ def evaluate_chunks(monographs: Sequence[dict],
blocks_by_section: Dict[tuple, list] = {}
block_ids: Dict[str, str] = {}
block_texts: Dict[str, str] = {}
sections_by_key: Dict[tuple, dict] = {}
for monograph in monographs:
for section_key, section in (monograph.get("sections") or {}).items():
sections_by_key[(monograph["drug_id"], section_key)] = section
for block in monograph.get("tables") or []:
key = (monograph["drug_id"], block.get("section_key"))
blocks_by_section.setdefault(key, []).append(block)
@@ -159,15 +178,76 @@ def evaluate_chunks(monographs: Sequence[dict],
unknown_id = missing_provenance = leaked = 0
descriptors = 0
descriptor_without_attachment = 0
attachment_header_row_present = 0
descriptor_with_unverified_columns = 0
unsupported_schema = 0
prose_without_source_text = 0
source_text_not_unique = 0
physical_range_not_exact = 0
descriptor_range_not_attachment_page = 0
attachment_without_printed_page = 0
context_label_missing_from_text = 0
invalid_printed_page_range = 0
for chunk in chunks:
if chunk.get("schema_version") != 4:
unsupported_schema += 1
printed_range = chunk.get("printed_page_range")
if (
not isinstance(printed_range, list)
or len(printed_range) != 2
or not all(isinstance(page, int) for page in printed_range)
or printed_range[0] > printed_range[1]
):
invalid_printed_page_range += 1
attachments = chunk.get("attachments") or []
if chunk.get("chunk_kind") == "block_descriptor":
descriptors += 1
if not attachments:
descriptor_without_attachment += 1
if "Cột:" in (chunk.get("text") or ""):
descriptor_with_unverified_columns += 1
if attachments:
page = attachments[0].get("physical_page")
if chunk.get("source_page_range") != [page, page]:
descriptor_range_not_attachment_page += 1
else:
source_body = chunk.get("source_text") or ""
if not source_body:
prose_without_source_text += 1
section = sections_by_key.get((chunk["drug_id"], chunk["section_key"]))
section_text = (section or {}).get("text", "").strip()
if not source_body or section_text.count(source_body) != 1:
source_text_not_unique += 1
else:
chunk_start = section_text.index(source_body)
chunk_end = chunk_start + len(source_body)
cursor = 0
pages = []
for part in (section or {}).get("parts") or []:
if (part.get("kind") != "prose" or part.get("quarantined")
or not part.get("text")):
continue
part_start = section_text.find(part["text"], cursor)
if part_start < 0:
pages = []
break
part_end = part_start + len(part["text"])
cursor = part_end
if part_start < chunk_end and part_end > chunk_start:
pages.append(part["physical_page"])
expected = [min(pages), max(pages)] if pages else None
if chunk.get("source_page_range") != expected:
physical_range_not_exact += 1
body = chunk.get("text") or ""
if any(label not in body for label in chunk.get("context_labels") or []):
context_label_missing_from_text += 1
key = (chunk["drug_id"], chunk["section_key"])
for attachment in attachments:
if attachment.get("header_row"):
attachment_header_row_present += 1
if attachment.get("printed_page") is None:
attachment_without_printed_page += 1
referenced.setdefault(key, set()).add(attachment["block_id"])
if block_ids.get(attachment["block_id"]) != chunk["drug_id"]:
unknown_id += 1
@@ -180,6 +260,9 @@ def evaluate_chunks(monographs: Sequence[dict],
if len(probe) > 20 and probe in body:
leaked += 1
over_ceiling = [c for c in chunks if c.get("oversized")]
uncovered = _sections_not_covered(monographs, chunks)
unreferenced = 0
for key, blocks in blocks_by_section.items():
seen = referenced.get(key, set())
@@ -187,10 +270,25 @@ def evaluate_chunks(monographs: Sequence[dict],
total_blocks = sum(len(v) for v in blocks_by_section.values())
return [
Gate("chunk_over_token_ceiling", len(over_ceiling),
detail="; ".join(c["chunk_id"] for c in over_ceiling[:3])),
Gate("chunk_without_printed_page_range", invalid_printed_page_range),
Gate("chunk_schema_version_not_supported", unsupported_schema),
Gate("prose_without_source_text", prose_without_source_text),
Gate("chunk_source_text_not_unique", source_text_not_unique),
Gate("chunk_physical_range_not_exact", physical_range_not_exact),
Gate("descriptor_range_not_attachment_page",
descriptor_range_not_attachment_page),
Gate("attachment_without_printed_page", attachment_without_printed_page),
Gate("context_label_missing_from_text", context_label_missing_from_text),
Gate("section_not_reassemblable_from_chunks", len(uncovered),
detail="; ".join(uncovered[:3])),
Gate("section_block_without_chunk_reference", unreferenced),
Gate("attachment_block_id_unknown", unknown_id),
Gate("attachment_without_page_or_bbox", missing_provenance),
Gate("block_text_leaked_into_chunk_text", leaked),
Gate("attachment_header_row_present", attachment_header_row_present),
Gate("descriptor_with_unverified_columns", descriptor_with_unverified_columns),
Gate("descriptor_chunk_without_attachment", descriptor_without_attachment),
Gate("descriptor_count_vs_block_count", descriptors, target=total_blocks,
detail=f"{descriptors} descriptors for {total_blocks} blocks"),
@@ -200,3 +298,65 @@ def evaluate_chunks(monographs: Sequence[dict],
def read_chunks(path: Path) -> List[dict]:
with path.open(encoding="utf-8") as handle:
return [json.loads(line) for line in handle if line.strip()]
def _reassemble(parts: Sequence[str]) -> str:
"""Glue chunk parts back together, removing the deliberate overlap."""
if not parts:
return ""
text = parts[0]
for part in parts[1:]:
overlap = 0
for size in range(min(len(text), len(part)), 0, -1):
if text.endswith(part[:size]):
overlap = size
break
text += part[overlap:]
return text
def _sections_not_covered(monographs: Sequence[dict],
chunks: Sequence[dict]) -> List[str]:
"""Sections that cannot be rebuilt exactly from their own chunks.
Stronger than asking whether each line still appears somewhere: it proves
the chunks are a faithful partition of the section, so nothing was dropped
*and* nothing was reordered or duplicated beyond the intended overlap.
Whitespace is removed from both sides rather than normalised, because
each split seam legitimately loses one space: a part is built with
`"".join(...).strip()`, and the space that sat between two sentences falls
on the boundary. Measured on ABACAVIR: exactly two single spaces across a
4,232-character section, nothing else. The gate exists to prove no
character of *content* is lost, reordered, or duplicated beyond the
intended overlap; it is not a formatting check.
Two weaker versions were tried first and both were instrument bugs, not
data bugs. Joining chunk texts with a newline meant a paragraph split
across parts could never match, reporting 734 sections missing when the
first one it named was present. Probing a 60-character head then failed on
NAPROXEN alone, because that probe straddled an overlap seam where the
repeated text legitimately appears twice.
"""
by_section: Dict[tuple, List[dict]] = {}
for chunk in chunks:
if chunk.get("chunk_kind") == "block_descriptor":
continue
key = (chunk["drug_id"], chunk["section_key"])
by_section.setdefault(key, []).append(chunk)
broken = []
for monograph in monographs:
for key, section in (monograph.get("sections") or {}).items():
source = _WHITESPACE.sub("", section.get("text") or "")
if not source:
continue
parts = sorted(by_section.get((monograph["drug_id"], key), []),
key=lambda c: c.get("part_index", 0))
rebuilt = _WHITESPACE.sub(
"", _reassemble([
c.get("source_text") or c.get("text") or "" for c in parts
]))
if rebuilt != source:
broken.append(f"{monograph['drug_id']}/{key}")
return broken
+13 -1
View File
@@ -3,10 +3,22 @@ name = "ingestion"
version = "0.0.0"
description = "Offline batch pipeline: PDF -> monographs -> chunks -> embeddings -> Qdrant"
requires-python = ">=3.11"
dependencies = ["pymupdf>=1.24", "numpy>=1.26", "scipy>=1.11"]
dependencies = ["pymupdf>=1.24", "numpy>=1.26", "scipy>=1.11", "tiktoken>=0.7"]
[project.optional-dependencies]
dev = ["pytest>=7.4"]
# Only the live probe and a real embedding run need these. The adapters and
# their tests import neither, so the default install stays offline-capable.
bedrock = ["boto3>=1.34"]
local-embed = ["sentence-transformers>=3.0"]
# `load/` talks to a VectorStore port; only `load.qdrant_repo` imports this,
# lazily, so the load stage is tested in full with no server running.
qdrant = ["qdrant-client>=1.7"]
[tool.pytest.ini_options]
markers = [
"integration: needs a live service (a local Qdrant); skips when absent",
]
[build-system]
requires = ["setuptools>=68"]
+267 -8
View File
@@ -6,11 +6,19 @@ from ingestion.chunk import (
CHUNK_KIND_PROSE,
SCHEMA_VERSION,
chunk_monograph,
chunk_all,
chunk_section,
write_chunks_jsonl,
)
from ingestion.chunk.chunker import _is_label_row, describe_block
from ingestion.segment.models import Heading, Monograph, SectionSpan, TableBlock
from ingestion.segment.models import (
PART_PROSE,
Heading,
Monograph,
SectionPart,
SectionSpan,
TableBlock,
)
from ingestion.tables import SHAPE_FORMULA_2D, SHAPE_MULTI_HEADER, SHAPE_SIMPLE
@@ -20,6 +28,15 @@ def _section(key, display, text, page=202):
heading=Heading(text=display, physical_page=page, y0=100.0,
is_monograph_title=False, section_key=key),
text=text,
parts=([
SectionPart(
kind=PART_PROSE,
text=text,
physical_page=page,
bbox=[50.0, 120.0, 550.0, 700.0],
source_span_ids=[f"p{page}_s0"],
)
] if text else []),
)
@@ -65,13 +82,17 @@ def test_a_section_whose_table_was_lifted_says_so():
def test_a_lifted_block_gets_its_own_retrievable_descriptor():
section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "Prose.")
monograph = _monograph([section], [_block()])
descriptors = [c for c in chunk_monograph(monograph)
descriptors = [c for c in chunk_monograph(
monograph, printed_page_map={200: 201, 202: 203, 203: 204})
if c.chunk_kind == CHUNK_KIND_BLOCK_DESCRIPTOR]
assert len(descriptors) == 1
assert "AMPICILIN VÀ SULBACTAM" in descriptors[0].text
assert "Liều lượng và cách dùng" in descriptors[0].text
# printed page, which is what a reader holding the book looks for
assert "trang 203" in descriptors[0].text
assert descriptors[0].source_page_range == [202, 202]
assert descriptors[0].printed_page_range == [203, 203]
assert descriptors[0].attachments[0].printed_page == 203
def test_no_cell_value_ever_reaches_the_descriptor_text():
@@ -106,15 +127,17 @@ def test_a_header_row_carrying_a_number_is_refused():
assert descriptor.attachments[0].header_row == []
def test_only_a_simple_table_contributes_a_header():
def test_unverified_header_rows_are_embargoed_for_every_table_shape():
section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "Prose.")
header = {"p202_t0": ["Nhóm", "Liều"]}
for shape, expected in ((SHAPE_SIMPLE, ["Nhóm", "Liều"]),
(SHAPE_MULTI_HEADER, [])):
header = {"p202_t0": ["Ngoại tâm thu thất", "Thường gặp", "Không rõ tần suất"]}
for shape in (SHAPE_SIMPLE, SHAPE_MULTI_HEADER):
monograph = _monograph([section], [_block(shape=shape)])
descriptor = next(c for c in chunk_monograph(monograph, header)
if c.chunk_kind == CHUNK_KIND_BLOCK_DESCRIPTOR)
assert descriptor.attachments[0].header_row == expected
assert descriptor.attachments[0].header_row == []
assert "Cột:" not in descriptor.text
assert "Ngoại tâm thu thất" not in descriptor.text
assert "Không rõ tần suất" not in descriptor.text
def test_a_formula_block_is_described_as_a_formula():
@@ -163,6 +186,242 @@ def test_describe_block_names_the_page_even_with_no_header():
chunks = chunk_monograph(monograph)
attachment = next(c for c in chunks
if c.chunk_kind == CHUNK_KIND_BLOCK_DESCRIPTOR).attachments[0]
text = describe_block(monograph, section, attachment)
text = describe_block(monograph, section, attachment, printed_page=203)
assert "trang 203" in text
assert "không trích dẫn được dưới dạng văn bản" in text
def test_chunk_carries_only_verified_printed_page_range():
section = _section("chi_dinh", "Chỉ định", "Nhiễm khuẩn.")
monograph = _monograph([section])
chunk = chunk_monograph(
monograph,
printed_page_map={202: 203},
)[0]
assert chunk.source_page_range == [202, 202]
assert chunk.printed_page_range == [203, 203]
def test_chunk_refuses_an_unmapped_printed_folio():
section = _section("chi_dinh", "Chỉ định", "Nhiễm khuẩn.")
monograph = _monograph([section])
try:
chunk_monograph(monograph, printed_page_map={202: None})
except ValueError as exc:
assert "printed folio missing" in str(exc)
else:
raise AssertionError("missing printed folio must fail closed")
def test_a_chunk_spanning_two_section_parts_cites_only_those_pages():
text = "Nội dung trang một.\nNội dung trang hai."
section = SectionSpan(
key="chi_dinh",
display_name="Chỉ định",
heading=Heading(
text="Chỉ định", physical_page=201, y0=100.0,
is_monograph_title=False, section_key="chi_dinh",
),
text=text,
parts=[
SectionPart(PART_PROSE, "Nội dung trang một.", 201,
[50.0, 100.0, 550.0, 200.0], ["p201_s0"]),
SectionPart(PART_PROSE, "Nội dung trang hai.", 202,
[50.0, 100.0, 550.0, 200.0], ["p202_s0"]),
],
)
chunk = chunk_monograph(
_monograph([section]), printed_page_map={201: 202, 202: 203}
)[0]
assert chunk.source_page_range == [201, 202]
assert chunk.printed_page_range == [202, 203]
def test_whole_corpus_chunking_refuses_to_run_without_a_printed_page_map():
section = _section("chi_dinh", "Chỉ định", "Nhiễm khuẩn.")
try:
list(chunk_all([_monograph([section])], printed_page_map=None))
except ValueError as exc:
assert "requires a verified printed_page_map" in str(exc)
else:
raise AssertionError("whole-corpus chunking must fail closed without folios")
def test_the_char_ratio_estimate_is_never_used_as_a_token_count():
"""ADR 0004 sized chunks with len(text)//4 and reported 0 over the ceiling.
Counted with the real tokenizer, 1,884 of 12,838 chunks (14.7%) were over
it, the largest at 1,645 tokens — twice the ceiling. Vietnamese diacritics
cost multiple byte-pair tokens each; measured ratio real/estimate is 1.95
at the median and 6.0 at worst.
"""
from ingestion.chunk.tokens import count_tokens, estimate_tokens
vietnamese = "Liều thường dùng cho người lớn là 1,5 - 3 g mỗi 6 giờ."
assert count_tokens(vietnamese) > len(vietnamese) // 4
# the fallback errs small, so it can never certify an oversized chunk as safe
assert estimate_tokens(vietnamese) > len(vietnamese) // 4
def test_a_long_comma_list_is_split_at_commas_not_left_oversized():
"""VORICONAZOL's interaction list is one 'sentence' hundreds of names long.
Truncated by an embedding model it reads as "this drug is not listed" — a
false negative in the direction that matters. A comma is a lossless break.
"""
from ingestion.chunk.chunker import CEILING_TOKENS, _atoms
drugs = ", ".join(f"thuốc {n}" for n in range(400))
atoms = _atoms(drugs + ".", lambda t: len(t) // 2)
assert len(atoms) > 1
assert all(len(a) // 2 <= CEILING_TOKENS for a in atoms)
assert "".join(atoms).replace(",", "") == (drugs + ".").replace(",", "")
def test_the_overlap_never_exceeds_its_budget():
"""A 251-token atom produced a 273-token overlap against a 65-token
setting, because the loop added whole atoms until the total passed it.
That was most of how a 981-token chunk came about."""
from ingestion.chunk.chunker import OVERLAP_TOKENS, _pack
measure = lambda t: len(t) # noqa: E731 - one-line stub for the test
atoms = ["a" * 300, "b" * 300, "c" * 300]
parts = _pack(atoms, measure)
assert len(parts) > 1
for part in parts[1:]:
carried = part[:-1]
assert sum(measure(a) for a in carried) <= OVERLAP_TOKENS
def test_a_continuation_repeats_the_label_governing_its_dose():
"""A budget-only overlap used to strand population labels.
The two 30-token dose atoms fit the 65-token overlap, while the preceding
label did not. The continuation was therefore independently retrievable
as a bare dose even though its source context was population-specific.
"""
from ingestion.chunk.chunker import _pack
measure = len
label = "Trẻ đẻ thiếu tháng:"
atoms = [
"p" * 570,
label,
"Uống liều 2 mg/kg q12h. " + "a" * 5,
"Nếu không uống được: " + "b" * 8,
"Theo dõi đáp ứng và điều chỉnh liều. " + "c" * 70,
]
parts = _pack(atoms, measure)
assert len(parts) == 2
assert label in parts[1]
assert parts[1].index(label) < next(
index for index, atom in enumerate(parts[1]) if "liều" in atom
)
def test_a_single_long_label_is_never_emitted_without_its_dose():
"""When the current buffer held only one long label, the old loop could
not carry it and emitted a label-only retrievable chunk."""
from ingestion.chunk.chunker import _pack
label = "Trẻ sơ sinh có tình trạng lâm sàng cần hiệu chỉnh đặc biệt " * 2 + ":"
dose = "Dùng liều khởi đầu " + "x" * 640
parts = _pack([label, dose], len)
assert all(part != [label] for part in parts)
assert any(label in part and dose in part for part in parts)
def test_a_new_trailing_label_does_not_orphan_the_previous_population_dose():
"""Real shape: a neonatal dose is followed by ``Suy thận:`` at the seam.
Carrying only the new trailing label is insufficient: any dose atoms copied
into the overlap must retain the older population label that governs them.
"""
from ingestion.chunk.chunker import _pack
population = "Trẻ đẻ thiếu tháng và trẻ sơ sinh dưới 8 ngày tuổi:"
renal = "Suy thận:"
first_dose = "100 mg/kg/ngày, chia hai lần. "
dose_limit = "Liều tối đa 10 mg/kg/ngày. "
atoms = [
"p" * 520,
population,
first_dose,
dose_limit,
renal,
"Điều chỉnh theo độ thanh thải creatinin. " + "x" * 80,
]
parts = _pack(atoms, len)
assert len(parts) == 2
assert parts[1][0] == population
assert renal in parts[1]
copied_doses = [atom for atom in parts[1] if atom in (first_dose, dose_limit)]
if copied_doses:
assert parts[1].index(population) < min(parts[1].index(atom) for atom in copied_doses)
def test_an_atom_ending_in_the_next_label_keeps_the_previous_dose_context():
"""Bisoprolol has atoms shaped ``dose for step 4 ... Step 5:``.
Ending in a colon does not make the dose at the beginning of that same atom
belong to the new label.
"""
from ingestion.chunk.chunker import _pack, _split_trailing_label
previous = "Bước 4:"
compound = "7,5 mg/lần/ngày trong 4 tuần; chuyển bước 5.\nBước 5:"
next_dose = "10 mg/lần/ngày để duy trì. "
split_compound = _split_trailing_label(compound)
assert "".join(split_compound) == compound
assert split_compound == [
"7,5 mg/lần/ngày trong 4 tuần; chuyển bước 5.\n",
"Bước 5:",
]
atoms = [
"p" * 540,
previous,
*split_compound,
next_dose,
"Theo dõi dung nạp và điều chỉnh. " + "x" * 80,
]
parts = _pack(atoms, len)
assert len(parts) == 2
dose_atom = split_compound[0]
if dose_atom in parts[1]:
assert previous in parts[1]
assert parts[1].index(previous) < parts[1].index(dose_atom)
assert split_compound[1] in parts[1]
assert parts[1].index(split_compound[1]) < parts[1].index(next_dose)
def test_a_population_continuation_retains_its_parent_route():
"""PARACETAMOL: age-band labels are children of ``Đường trực tràng:``."""
from ingestion.chunk.chunker import _pack
route = "Đường trực tràng:"
population = "Trẻ em 1 - 3 tháng tuổi:"
dose = "30 mg/kg một liều duy nhất. "
next_population = "Trẻ em 3 tháng - 6 tuổi:"
atoms = [
"p" * 540,
route,
population,
dose,
next_population,
"30 - 40 mg/kg một liều duy nhất. " + "x" * 80,
]
parts = _pack(atoms, len)
assert len(parts) == 2
assert route in parts[1]
assert next_population in parts[1]
assert parts[1].index(route) < parts[1].index(next_population)
+233
View File
@@ -0,0 +1,233 @@
"""The embedding cache, exercised with a counting stub and no network.
The claim these tests exist to make checkable is narrow and financial: running
the corpus a second time must cost nothing. That is asserted by counting calls
the *inner* provider received, not by trusting a hit counter.
The other half is the inverse — the cases where a hit would be wrong. Serving a
vector after its text was edited, across two models, or across Cohere's
document/query subspaces would each be silent: no error, just worse recall or a
corpus of mixed vectors. There is a test per direction.
"""
import json
import pytest
from ingestion.embed import INPUT_DOCUMENT, EmbeddingVector
from ingestion.embed.cache import (
CachingEmbeddingProvider,
EmbeddingCache,
cache_key,
)
from ingestion.embed.ports import EmbeddingProvider
DIMENSIONS = 8
class CountingProvider(EmbeddingProvider):
"""Deterministic vectors, and a record of every text it was asked for."""
def __init__(self, model_id="stub-model-v1", dimensions=DIMENSIONS, batch=96):
self._model_id = model_id
self._dimensions = dimensions
self._batch = batch
self.embedded_texts = []
self.batch_calls = 0
@property
def name(self):
return "stub"
@property
def model_id(self):
return self._model_id
@property
def dimensions(self):
return self._dimensions
@property
def max_batch_size(self):
return self._batch
def _embed_batch(self, texts, input_kind):
self.batch_calls += 1
self.embedded_texts.extend(texts)
return [self._vector(text, input_kind) for text in texts]
def _vector(self, text, input_kind):
from ingestion.embed import text_digest
seed = len(text) + (0 if input_kind == INPUT_DOCUMENT else 1000)
return EmbeddingVector(
values=[float(seed + i) for i in range(self._dimensions)],
text_sha256=text_digest(text),
provider="stub",
model_id=self._model_id,
dimensions=self._dimensions,
input_kind=input_kind,
normalized=True,
input_token_count=len(text.split()),
)
@pytest.fixture()
def cache_path(tmp_path):
return tmp_path / "embeddings.jsonl"
def test_second_run_over_the_same_texts_makes_zero_provider_requests(cache_path):
texts = ["paracetamol", "chống chỉ định", "liều dùng cho trẻ em"]
inner = CountingProvider()
first = CachingEmbeddingProvider(inner, EmbeddingCache(cache_path))
cold = first.embed_documents(texts)
assert cold.request_count == 1
assert inner.embedded_texts == texts
reopened = EmbeddingCache(cache_path)
warm = CachingEmbeddingProvider(inner, reopened).embed_documents(texts)
assert warm.request_count == 0
assert inner.embedded_texts == texts, "no text reached the provider twice"
assert reopened.stats.hits == len(texts)
assert reopened.stats.misses == 0
assert reopened.stats.hit_rate == 1.0
assert [v.values for v in warm.vectors] == [v.values for v in cold.vectors]
def test_a_repeated_text_in_one_call_is_embedded_once(cache_path):
inner = CountingProvider()
provider = CachingEmbeddingProvider(inner, EmbeddingCache(cache_path))
batch = provider.embed_documents(["Abacavir.", "Abacavir.", "Abacavir."])
assert inner.embedded_texts == ["Abacavir."]
assert len(batch.vectors) == 3
assert batch.vectors[0].values == batch.vectors[2].values
def test_editing_the_text_is_a_miss_not_a_stale_hit(cache_path):
inner = CountingProvider()
cache = EmbeddingCache(cache_path)
CachingEmbeddingProvider(inner, cache).embed_documents(["liều 500 mg"])
CachingEmbeddingProvider(inner, cache).embed_documents(["liều 250 mg"])
assert inner.embedded_texts == ["liều 500 mg", "liều 250 mg"]
def test_a_second_model_never_reuses_the_first_models_vectors(cache_path):
cache = EmbeddingCache(cache_path)
titan = CountingProvider(model_id="amazon.titan-embed-text-v2:0")
cohere = CountingProvider(model_id="cohere.embed-v4:0")
CachingEmbeddingProvider(titan, cache).embed_documents(["metformin"])
CachingEmbeddingProvider(cohere, cache).embed_documents(["metformin"])
assert titan.embedded_texts == ["metformin"]
assert cohere.embedded_texts == ["metformin"]
assert len(cache) == 2
def test_query_and_document_kinds_are_cached_separately(cache_path):
inner = CountingProvider()
cache = EmbeddingCache(cache_path)
provider = CachingEmbeddingProvider(inner, cache)
as_document = provider.embed_documents(["aspirin"])
as_query = provider.embed_queries(["aspirin"])
assert inner.embedded_texts == ["aspirin", "aspirin"]
assert as_document.vectors[0].values != as_query.vectors[0].values
assert len(cache) == 2
def test_index_and_values_survive_reopening_the_file(cache_path):
inner = CountingProvider()
original = CachingEmbeddingProvider(
inner, EmbeddingCache(cache_path)
).embed_documents(["ACETAZOLAMID", "ADENOSIN"])
reopened = EmbeddingCache(cache_path)
assert len(reopened) == 2
restored = reopened.get(cache_key(inner.model_id, INPUT_DOCUMENT, "ADENOSIN"))
assert restored is not None
assert restored.values == original.vectors[1].values
assert restored.input_kind == INPUT_DOCUMENT
assert restored.normalized is True
assert restored.input_token_count == 1
def test_putting_a_key_twice_does_not_append_a_second_record(cache_path):
cache = EmbeddingCache(cache_path)
inner = CountingProvider()
vector = inner.embed_documents(["digoxin"]).vectors[0]
assert cache.put(vector) is True
assert cache.put(vector) is False
lines = cache_path.read_text(encoding="utf-8").strip().splitlines()
assert len(lines) == 1
assert len(cache) == 1
def test_a_cached_record_whose_length_contradicts_its_dimensions_is_rejected(
cache_path,
):
record = {
"model_id": "stub-model-v1",
"input_kind": INPUT_DOCUMENT,
"text_sha256": "a" * 64,
"provider": "stub",
"dimensions": 1024,
"normalized": True,
"input_token_count": 3,
"values": [0.1, 0.2],
}
cache_path.write_text(json.dumps(record) + "\n", encoding="utf-8")
cache = EmbeddingCache(cache_path)
with pytest.raises(ValueError, match="declares 1024 dimensions"):
cache.get(("stub-model-v1", INPUT_DOCUMENT, "a" * 64))
def test_a_record_missing_key_fields_is_rejected_at_index_time(cache_path):
cache_path.write_text(
json.dumps({"model_id": "stub", "values": []}) + "\n", encoding="utf-8"
)
with pytest.raises(ValueError, match="missing key fields"):
EmbeddingCache(cache_path)
def test_the_cache_wrapper_still_refuses_empty_text_and_bad_input_kind(cache_path):
provider = CachingEmbeddingProvider(CountingProvider(), EmbeddingCache(cache_path))
with pytest.raises(ValueError, match="empty or whitespace-only"):
provider.embed_documents(["paracetamol", " "])
with pytest.raises(ValueError, match="input_kind must be one of"):
provider.embed(["paracetamol"], "search_document")
def test_a_missing_cache_file_starts_empty_and_is_created_on_first_put(cache_path):
cache = EmbeddingCache(cache_path)
assert len(cache) == 0
assert not cache_path.exists()
CachingEmbeddingProvider(CountingProvider(), cache).embed_documents(["insulin"])
assert cache_path.exists()
assert len(cache) == 1
def test_wrapper_reports_the_inner_models_identity_not_its_own(cache_path):
inner = CountingProvider(model_id="cohere.embed-v4:0", dimensions=DIMENSIONS)
provider = CachingEmbeddingProvider(inner, EmbeddingCache(cache_path))
assert provider.model_id == "cohere.embed-v4:0"
assert provider.dimensions == DIMENSIONS
assert provider.max_batch_size == inner.max_batch_size
assert provider.name == "cached:stub"
+324
View File
@@ -0,0 +1,324 @@
"""Provider adapters, exercised with no AWS account and no network.
Every Bedrock call goes through a recording stub, so what is under test is the
part that can actually be wrong offline: the request body we send, and our
reading of the response bodies AWS documents. The one thing these tests cannot
establish is whether AWS accepts that body — that needs the live probe, and
the coordination handoff says so explicitly.
"""
import io
import json
import math
import pytest
from ingestion.embed import (
BGE_M3,
COHERE_V4,
TITAN_V2,
INPUT_DOCUMENT,
INPUT_QUERY,
Boto3BedrockInvoker,
EmbeddingVector,
build_provider,
provider_names,
text_digest,
)
from ingestion.embed import probe
from ingestion.embed.bedrock_cohere import CohereEmbedV4
from ingestion.embed.bedrock_titan import TitanTextEmbeddingsV2
from ingestion.embed.local_bge_m3 import BgeM3Local
class RecordingInvoker:
"""Stands in for Bedrock; remembers every request it was handed."""
def __init__(self, responses):
self._responses = list(responses)
self.calls = []
def invoke_json(self, model_id, payload, accept="application/json"):
self.calls.append(
{"model_id": model_id, "payload": payload, "accept": accept}
)
return self._responses.pop(0)
def _titan_response(dimensions=1024, token_count=12):
return {
"embedding": [0.01] * dimensions,
"inputTextTokenCount": token_count,
"embeddingsByType": {"float": [0.01] * dimensions},
}
def _cohere_by_type_response(rows, dimensions=1024):
return {
"id": "stub-id",
"response_type": "embeddings_by_type",
"embeddings": {"float": [[0.02] * dimensions for _ in range(rows)]},
"texts": ["stub"] * rows,
}
def _cohere_floats_response(rows, dimensions=1024):
return {
"id": "stub-id",
"response_type": "embeddings_floats",
"embeddings": [[0.02] * dimensions for _ in range(rows)],
}
def test_titan_request_body_matches_the_documented_v2_shape():
invoker = RecordingInvoker([_titan_response()])
provider = TitanTextEmbeddingsV2(invoker, dimensions=1024, normalize=True)
provider.embed_documents(["paracetamol"])
payload = invoker.calls[0]["payload"]
assert invoker.calls[0]["model_id"] == "amazon.titan-embed-text-v2:0"
assert payload == {
"inputText": "paracetamol",
"dimensions": 1024,
"normalize": True,
}
def test_titan_records_provenance_and_reported_token_count():
invoker = RecordingInvoker([_titan_response(token_count=7)])
provider = TitanTextEmbeddingsV2(invoker)
vector = provider.embed_documents(["paracetamol"]).vectors[0]
assert vector.model_id == "amazon.titan-embed-text-v2:0"
assert vector.provider == TITAN_V2
assert vector.dimensions == 1024
assert vector.input_kind == INPUT_DOCUMENT
assert vector.normalized is True
assert vector.input_token_count == 7
assert vector.text_sha256 == text_digest("paracetamol")
def test_titan_sends_one_request_per_text():
invoker = RecordingInvoker([_titan_response(), _titan_response()])
provider = TitanTextEmbeddingsV2(invoker)
batch = provider.embed_documents(["a", "b"])
assert batch.request_count == 2
assert len(batch.vectors) == 2
def test_titan_rejects_a_dimension_the_model_does_not_offer():
with pytest.raises(ValueError, match="supports"):
TitanTextEmbeddingsV2(RecordingInvoker([]), dimensions=768)
def test_cohere_uses_search_document_for_corpus_and_search_query_for_queries():
invoker = RecordingInvoker(
[_cohere_by_type_response(1), _cohere_by_type_response(1)]
)
provider = CohereEmbedV4(invoker)
provider.embed_documents(["metformin"])
provider.embed_queries(["liều metformin"])
assert invoker.calls[0]["payload"]["input_type"] == "search_document"
assert invoker.calls[1]["payload"]["input_type"] == "search_query"
def test_cohere_request_body_pins_dimension_float_type_and_no_truncation():
invoker = RecordingInvoker([_cohere_by_type_response(2)])
provider = CohereEmbedV4(invoker, dimensions=1024)
provider.embed_documents(["a", "b"])
payload = invoker.calls[0]["payload"]
assert invoker.calls[0]["model_id"] == "cohere.embed-v4:0"
assert payload["texts"] == ["a", "b"]
assert payload["embedding_types"] == ["float"]
# Left unset the model would return 1536, which no 1024-wide collection
# can accept.
assert payload["output_dimension"] == 1024
# An over-length input must fail, not arrive silently shortened.
assert payload["truncate"] == "NONE"
assert invoker.calls[0]["accept"] == "*/*"
def test_cohere_reads_the_embeddings_by_type_response():
invoker = RecordingInvoker([_cohere_by_type_response(2)])
provider = CohereEmbedV4(invoker)
batch = provider.embed_documents(["a", "b"])
assert len(batch.vectors) == 2
assert all(len(v.values) == 1024 for v in batch.vectors)
assert batch.request_count == 1
def test_cohere_also_reads_the_plain_embeddings_floats_response():
invoker = RecordingInvoker([_cohere_floats_response(2)])
provider = CohereEmbedV4(invoker)
batch = provider.embed_documents(["a", "b"])
assert len(batch.vectors) == 2
assert all(len(v.values) == 1024 for v in batch.vectors)
def test_cohere_leaves_normalization_unknown_because_the_docs_do_not_say():
invoker = RecordingInvoker([_cohere_by_type_response(1)])
vector = CohereEmbedV4(invoker).embed_documents(["a"]).vectors[0]
assert vector.normalized is None
def test_cohere_splits_at_the_documented_96_text_ceiling():
invoker = RecordingInvoker(
[_cohere_by_type_response(96), _cohere_by_type_response(4)]
)
provider = CohereEmbedV4(invoker)
batch = provider.embed_documents([f"t{i}" for i in range(100)])
assert batch.request_count == 2
assert len(invoker.calls[0]["payload"]["texts"]) == 96
assert len(invoker.calls[1]["payload"]["texts"]) == 4
assert len(batch.vectors) == 100
def test_cohere_rejects_a_batch_size_above_the_documented_ceiling():
with pytest.raises(ValueError, match="batch_size"):
CohereEmbedV4(RecordingInvoker([]), batch_size=97)
def test_a_wrong_width_vector_fails_instead_of_entering_the_corpus():
invoker = RecordingInvoker([_titan_response(dimensions=512)])
provider = TitanTextEmbeddingsV2(invoker, dimensions=1024)
with pytest.raises(ValueError, match="512 dimensions"):
provider.embed_documents(["a"])
def test_a_response_missing_its_vectors_fails_loudly():
invoker = RecordingInvoker([{"id": "stub", "response_type": "x"}])
with pytest.raises(ValueError, match="no 'embeddings' field"):
CohereEmbedV4(invoker).embed_documents(["a"])
def test_a_count_mismatch_between_texts_and_vectors_fails():
invoker = RecordingInvoker([_cohere_by_type_response(1)])
with pytest.raises(ValueError, match="1 vectors for 2 texts"):
CohereEmbedV4(invoker).embed_documents(["a", "b"])
def test_an_unknown_input_kind_is_refused_before_any_request_is_made():
invoker = RecordingInvoker([])
with pytest.raises(ValueError, match="input_kind"):
CohereEmbedV4(invoker).embed(["a"], "search_document")
assert invoker.calls == []
def test_empty_text_is_refused_before_any_request_is_made():
invoker = RecordingInvoker([])
with pytest.raises(ValueError, match="empty"):
CohereEmbedV4(invoker).embed_documents(["a", " "])
assert invoker.calls == []
def test_bge_m3_runs_through_an_injected_encoder_with_no_weights_loaded():
seen = []
def encoder(texts):
seen.append(list(texts))
unit = 1.0 / math.sqrt(1024)
return [[unit] * 1024 for _ in texts]
provider = BgeM3Local(encoder=encoder, batch_size=2)
batch = provider.embed_queries(["a", "b", "c"])
assert seen == [["a", "b"], ["c"]]
assert batch.request_count == 2
assert len(batch.vectors) == 3
assert batch.vectors[0].input_kind == INPUT_QUERY
assert batch.vectors[0].model_id == "BAAI/bge-m3"
# Injected encoder: we did not set normalize_embeddings, so we do not claim it.
assert batch.vectors[0].normalized is None
def test_registry_builds_every_provider_without_touching_an_sdk():
assert set(provider_names()) == {TITAN_V2, COHERE_V4, BGE_M3}
titan = build_provider(TITAN_V2, invoker=RecordingInvoker([]))
cohere = build_provider(COHERE_V4, invoker=RecordingInvoker([]))
local = build_provider(BGE_M3)
assert (titan.dimensions, cohere.dimensions, local.dimensions) == (
1024,
1024,
1024,
)
assert titan.max_batch_size == 1
assert cohere.max_batch_size == 96
def test_registry_rejects_an_unknown_provider_name():
with pytest.raises(ValueError, match="unknown embedding provider"):
build_provider("text-embedding-3-small")
class FakeBotoClient:
"""The shape boto3's bedrock-runtime client returns: a streaming body."""
def __init__(self, response_body):
self._response_body = response_body
self.kwargs = None
def invoke_model(self, **kwargs):
self.kwargs = kwargs
return {"body": io.BytesIO(json.dumps(self._response_body).encode())}
def test_boto3_invoker_serialises_the_request_and_reads_the_streamed_body():
client = FakeBotoClient({"embedding": [0.5]})
invoker = Boto3BedrockInvoker(region="us-east-1", client=client)
body = invoker.invoke_json("some.model", {"inputText": "à"}, accept="*/*")
assert body == {"embedding": [0.5]}
assert client.kwargs["modelId"] == "some.model"
assert client.kwargs["contentType"] == "application/json"
assert client.kwargs["accept"] == "*/*"
# Vietnamese must survive the round trip as characters, not \\u escapes
# the model would then embed literally.
assert json.loads(client.kwargs["body"]) == {"inputText": "à"}
def test_probe_measures_the_l2_norm_rather_than_trusting_the_docs():
unit = 1.0 / math.sqrt(4)
assert probe._l2_norm([unit] * 4) == pytest.approx(1.0)
assert probe._l2_norm([3.0, 4.0]) == pytest.approx(5.0)
def test_probe_reports_a_vector_without_raising(capsys):
vector = EmbeddingVector(
values=[0.5, 0.5, 0.5, 0.5],
text_sha256=text_digest("x"),
provider=TITAN_V2,
model_id="amazon.titan-embed-text-v2:0",
dimensions=4,
input_kind=INPUT_DOCUMENT,
normalized=True,
input_token_count=3,
)
probe._report(vector, latency_ms=12.5, requests=1)
out = capsys.readouterr().out
assert "amazon.titan-embed-text-v2:0" in out
assert "measured L2 norm: 1.000000" in out
+40
View File
@@ -0,0 +1,40 @@
from pathlib import Path
import pytest
from ingestion.entities.catalog import build_entities
DATA = Path(__file__).parents[1] / "data"
PDF = DATA / "raw/duoc-thu-quoc-gia-viet-nam-2018.pdf"
MONOGRAPHS = DATA / "processed/monographs.jsonl"
needs_source = pytest.mark.skipif(
not PDF.exists() or not MONOGRAPHS.exists(),
reason="source corpus artifacts are not available",
)
@needs_source
def test_verified_entity_catalog_maps_every_explicit_see_alias():
payload = build_entities(MONOGRAPHS, PDF)
stats = payload["stats"]
assert stats["entity_count"] == 684
assert stats["back_index_see_relations"] == 344
assert stats["back_index_aliases_mapped"] == 344
assert stats["back_index_aliases_unresolved"] == 0
assert stats["back_index_aliases_ambiguous"] == 0
assert stats["trade_name_sections"] == 492
@needs_source
def test_common_parenthesized_names_are_emitted_as_aliases():
payload = build_entities(MONOGRAPHS, PDF)
entities = {item["drug_id"]: item for item in payload["entities"]}
aliases = {
drug_id: {alias.casefold() for alias in entity["aliases"]}
for drug_id, entity in entities.items()
}
assert "paracetamol" in aliases["paracetamol_acetaminophen"]
assert "acetaminophen" in aliases["paracetamol_acetaminophen"]
assert "aspirin" in aliases["acid_acetylsalicylic_aspirin"]
assert "oresol" in aliases["thuoc_uong_bu_nuoc_va_ien_giai"]
+11
View File
@@ -2,6 +2,7 @@ import json
from pathlib import Path
from ingestion.extract.formulas import (
BARLESS_FORMULA_BOTTOM_PT,
FORMULA_BAND_HEIGHT_PT,
FORMULA_SIDE_MARGIN_PT,
load_formula_regions,
@@ -50,6 +51,16 @@ def test_the_barless_adenosin_formula_is_recorded_as_a_recall_limit():
assert "UNMEASURED" in payload["recall_limit"]
def test_barless_adenosin_band_reaches_its_printed_denominator():
payload = json.loads(VERIFIED.read_text(encoding="utf-8"))
entry = next(r for r in payload["regions"] if r.get("source_prints_no_bar"))
region = next(r for r in load_formula_regions() if r.physical_page == 147)
assert region.bbox[3] == entry["bar_bbox"][3] + BARLESS_FORMULA_BOTTOM_PT
# "Ví dụ:" begins immediately afterwards with its span centre at ~704.3;
# the formula band must stop before that prose and the following table.
assert region.bbox[3] < 704
def test_outlined_text_transcriptions_cover_every_detected_run():
payload = json.loads(TRANSCRIPTIONS.read_text(encoding="utf-8"))
runs = payload["runs"]
+589
View File
@@ -0,0 +1,589 @@
"""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"
@@ -0,0 +1,377 @@
"""`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")
)
+119
View File
@@ -1,3 +1,5 @@
from dataclasses import replace
import pytest
from ingestion.extract.models import Span
@@ -50,6 +52,55 @@ def test_non_bold_combined_heading_value_span_confirmed_real_amitriptylin_case()
m = list(assemble(spans))[0]
assert m.sections["ma_atc"].text == "N06AA09."
assert m.atc_codes == ["N06AA09"]
part = m.sections["ma_atc"].parts[0]
assert part.physical_page == 184
assert part.bbox != [0.0, 0.0, 0.0, 0.0]
assert part.source_span_ids == [spans[3].span_id]
def test_combined_international_name_and_atc_heading_is_a_title_anchor():
# Confirmed real GnRH class-monograph variant on physical page 1371.
spans = [
_span("THUỐC TƯƠNG TỰ HORMON GIẢI PHÓNG", 1371, 664.0),
_span("GONADOTROPIN", 1371, 676.0),
_span("Tên chung quốc tế và mã ATC", 1371, 690.0),
_span("Gonadorelin: H01CA01; Triptorelin: L02AE04.", 1371, 702.0, bold=False),
_span("Chỉ định", 1371, 714.0),
_span("Kích thích phóng noãn.", 1371, 726.0, bold=False),
]
m = list(assemble(spans))[0]
assert m.drug_name == "THUỐC TƯƠNG TỰ HORMON GIẢI PHÓNG GONADOTROPIN"
assert m.atc_codes == ["H01CA01", "L02AE04"]
def test_plain_wrapped_section_label_is_body_not_a_heading():
# NADROPARIN CALCI p1016: "không phải là" / "chống chỉ định."
# are adjacent lines of one sentence in the same PDF block.
spans = [
_span("NADROPARIN CALCI", 1016, 60.0),
_span("Tên chung quốc tế", 1016, 80.0),
_span("Nadroparin calcium.", 1016, 92.0, bold=False),
_span("Thời kỳ cho con bú", 1016, 110.0),
]
lead = replace(_span("Việc dùng thuốc không phải là", 1016, 122.0, bold=False),
block=4, line=7)
tail = replace(_span("chống chỉ định.", 1016, 134.0, bold=False),
block=4, line=8)
m = list(assemble(spans + [lead, tail]))[0]
assert m.sections["thoi_ky_cho_con_bu"].text.endswith("chống chỉ định.")
assert "chong_chi_dinh" not in m.sections
def test_plain_heading_after_completed_prose_still_opens_section():
spans = [
_span("TESTDRUG", 300, 60.0),
_span("Tên chung quốc tế", 300, 80.0),
replace(_span("Testdrug.", 300, 92.0, bold=False), block=2, line=0),
replace(_span("Chỉ định", 300, 104.0, bold=False), block=2, line=1),
replace(_span("Điều trị thử nghiệm.", 300, 116.0, bold=False), block=2, line=2),
]
m = list(assemble(spans))[0]
assert m.sections["chi_dinh"].text == "Điều trị thử nghiệm."
def test_atc_stated_absent_propagates():
@@ -330,3 +381,71 @@ def test_a_bold_label_line_still_opens_its_section():
]
monograph = list(assemble(spans))[0]
assert monograph.sections["chong_chi_dinh"].text == "Suy tủy nặng."
def test_a_section_name_printed_mid_line_is_body_not_a_heading():
"""CISPLATIN, physical page 402 — confirmed content loss.
The book prints "Suy thận: Chống chỉ định." inside the dosing section. The
second half is itself a section name, so it was matched as a heading: the
renal-impairment contraindication vanished from the dosing text and the
section ended on a bare "Suy thận:". ISOPRENALIN had the same shape. A
real heading opens its line; this one does not.
"""
spans = [
_span("CISPLATIN", 401, 60.0),
_span("Tên chung quốc tế", 401, 80.0),
_span("Cisplatinum.", 401, 92.0, bold=False),
_span("Liều lượng và cách dùng", 401, 110.0),
_span("Truyền tĩnh mạch mỗi 3 tuần.", 401, 122.0, bold=False),
]
label = _span("Suy thận: ", 401, 140.0, bold=False)
label = replace(label, block=4, line=0, x0=35.0, x1=70.0)
trailing = _span("Chống chỉ định.", 401, 140.0, bold=False)
trailing = replace(trailing, block=4, line=0, x0=70.0, x1=140.0)
monograph = list(assemble(spans + [label, trailing]))[0]
dosing = monograph.sections["lieu_luong_va_cach_dung"].text
assert "Suy thận: Chống chỉ định." in dosing
assert "chong_chi_dinh" not in monograph.sections
def test_italic_cross_reference_overlapping_its_neighbour_by_a_hairline_is_body():
"""NEVIRAPIN, physical page 1045 — confirmed misassignment, whole-corpus.
The book prints `Xem thêm mục ` (x1=104.89) immediately before an italic
`Liều lượng và cách dùng` (x0=104.88): the trailing space's advance width
makes the neighbour end 0.01pt *after* the cross-reference starts. An
end-before-start test therefore read a mid-line cross-reference as a
heading. Same shape, same cause, in CALCI LACTAT (p296, `xem thêm mục
Tương tác thuốc`, 0.02pt) and CEFAZOLIN (p344, `ghi ở mục: Dạng thuốc và
hàm lượng.`), where 4,533 characters of adult dosing were filed under
dosage forms.
"""
spans = [
_span("NEVIRAPIN", 1044, 60.0),
_span("Tên chung quốc tế", 1044, 80.0),
_span("Nevirapine.", 1044, 92.0, bold=False),
_span("Hướng dẫn cách xử trí ADR", 1044, 110.0),
_span("Điều trị các phản ứng bất lợi theo triệu chứng.", 1044, 122.0, bold=False),
]
lead = replace(_span("Xem thêm mục ", 1044, 140.0, bold=False),
block=4, line=0, x0=43.94, x1=104.89)
reference = replace(_span("Liều lượng và cách dùng", 1044, 140.0, bold=False),
block=4, line=0, x0=104.88, x1=199.57)
monograph = list(assemble(spans + [lead, reference]))[0]
assert "Xem thêm mục Liều lượng và cách dùng" in monograph.sections["huong_dan_xu_tri_adr"].text
assert "lieu_luong_va_cach_dung" not in monograph.sections
def test_a_section_name_opening_its_own_line_is_still_a_heading():
spans = [
_span("CISPLATIN", 401, 60.0),
_span("Tên chung quốc tế", 401, 80.0),
_span("Cisplatinum.", 401, 92.0, bold=False),
_span("Chống chỉ định", 401, 110.0),
_span("Suy tủy nặng.", 401, 122.0, bold=False),
]
monograph = list(assemble(spans))[0]
assert monograph.sections["chong_chi_dinh"].text == "Suy tủy nặng."
+14 -1
View File
@@ -1,5 +1,9 @@
from ingestion.extract.models import Span
from ingestion.segment.detector import detect_monograph_titles, detect_section_headings
from ingestion.segment.detector import (
detect_monograph_titles,
detect_section_headings,
in_monograph_range,
)
def _span(text, physical_page, printed_page, y0=100.0, bold=True, size=10.0):
@@ -86,3 +90,12 @@ def test_unknown_bold_text_not_matched_as_section():
def test_section_heading_outside_monograph_range_excluded():
spans = [_span("Chỉ định", 5, 6, bold=True, size=9.5)]
assert list(detect_section_headings(spans)) == []
def test_back_index_cannot_reenter_range_via_bad_inferred_printed_page():
# Confirmed real failure: physical page 1655 of the back index was mapped
# to printed page 1496, making its "Tương tác thuốc" entry extend
# ZOLPIDEM's source range from page 1494 through page 1655.
index_span = _span("Tương tác thuốc", 1655, 1496)
assert in_monograph_range(index_span) is False
assert list(detect_section_headings([index_span])) == []
+77 -8
View File
@@ -1,6 +1,12 @@
from ingestion.extract.models import Span
from ingestion.segment import assemble
from ingestion.tables import SHAPE_GRID_2D, SHAPE_SIMPLE, TableRegion, index_by_page
from ingestion.tables import (
SHAPE_FORMULA_2D,
SHAPE_GRID_2D,
SHAPE_SIMPLE,
TableRegion,
index_by_page,
)
def _span(text, page, y0, *, bold=False, x0=50.0, block=0, line=0, column="left"):
@@ -54,6 +60,40 @@ def test_table_spans_are_lifted_out_of_section_prose():
assert block.quarantined is True
def test_section_named_table_cell_does_not_change_owning_section():
# Confirmed in WARFARIN p1485 and IOBITRIDOL p826: a table column named
# "Chỉ định" belongs to the dosing table; it is not a document
# section heading and must not move the block into chi_dinh.
spans = [
_span("WARFARIN", 1485, 60.0, bold=True),
_span("Tên chung quốc tế", 1485, 80.0, bold=True),
_span("Warfarinum.", 1485, 92.0),
_span("Liều lượng và cách dùng", 1485, 200.0, bold=True),
_span("Chỉ định", 1485, 400.0, bold=True, block=5),
_span("INR 2,0 - 3,0", 1485, 412.0, block=5),
]
region = TableRegion("p1485_t0", 1485, (40.0, 380.0, 400.0, 460.0), 2, 2, SHAPE_SIMPLE)
m = list(assemble(spans, table_index=index_by_page([region])))[0]
assert "chi_dinh" not in m.sections
assert len(m.tables) == 1
assert m.tables[0].section_key == "lieu_luong_va_cach_dung"
def test_wide_formula_band_does_not_swallow_the_opposite_column():
spans = _monograph_spans([
_span("Công thức:", 109, 360.0),
_span("Cl", 109, 400.0, x0=280.0, column="left", block=5),
_span("Xem thêm Liều lượng và cách dùng", 109, 400.0,
x0=310.0, column="right", block=6),
])
# Deliberately extends across the gutter, as verified formula bands do.
region = TableRegion("p109_f0", 109, (40.0, 380.0, 390.0, 430.0),
2, 1, SHAPE_FORMULA_2D)
m = list(assemble(spans, table_index=index_by_page([region])))[0]
assert m.tables[0].text == "Cl"
assert "Xem thêm Liều lượng và cách dùng" in m.sections["dang_thuoc_va_ham_luong"].text
def test_without_a_region_map_behaviour_is_unchanged():
spans = _monograph_spans([
_span("Thuốc dùng đường uống.", 109, 220.0),
@@ -85,24 +125,53 @@ def test_non_table_regions_are_never_lifted():
def test_table_block_ids_stay_unique_when_a_section_resumes():
# a region flushed twice (section closes, then resumes) must not emit two
# blocks with the same table_id — provenance ids have to be unique
# Confirmed on CAPECITABIN pp. 308-309 and IMATINIB p. 795: PDF block
# order can place a visually later heading between cells from one physical
# table. The complete region must stay atomic and owned by the section
# active where the table first appears.
spans = [
_span("CEFAMANDOL", 339, 60.0, bold=True),
_span("Tên chung quốc tế", 339, 80.0, bold=True),
_span("Cefamandolum.", 339, 92.0),
_span("Liều lượng và cách dùng", 339, 200.0, bold=True),
_span("80 - 50", 339, 400.0, block=5),
_span("Liều lượng và cách dùng", 339, 500.0, bold=True),
# Visually below the table, but emitted before its final cell by the
# PDF's internal block order.
_span("Tương tác thuốc", 339, 640.0, bold=True),
_span("< 25 - 10", 339, 600.0, block=9),
_span("Không phối hợp với thuốc X.", 339, 660.0, block=10),
]
region = TableRegion("p339_t0", 339, (40.0, 380.0, 400.0, 620.0), 5, 2, SHAPE_SIMPLE)
m = list(assemble(spans, table_index=index_by_page([region])))[0]
# table_id is deterministic per REGION, so two parts of one table share
# it on purpose; table_part_id is the unique key, derived from the first
# source span rather than a counter (a counter would renumber whenever
# anything upstream shifted, hiding rather than identifying a duplicate)
assert len(m.tables) == 1
assert m.tables[0].section_key == "lieu_luong_va_cach_dung"
assert "80 - 50" in m.tables[0].text
assert "< 25 - 10" in m.tables[0].text
assert "Không phối hợp với thuốc X." in m.sections["tuong_tac_thuoc"].text
assert len({t.table_part_id for t in m.tables}) == len(m.tables)
assert {t.continuation_group for t in m.tables} == {"p339_t0"}
assert all(t.table_part_id.startswith("p339_t0@") for t in m.tables)
assert all(t.quarantined for t in m.tables)
def test_explicit_dose_adjustment_caption_reassigns_late_appendix_table():
# CAPECITABIN p. 309: the PDF puts dose-adjustment tables after the trade
# names and does not repeat the ordinary dosage section heading. Internal
# block order can even emit a cell before the visually preceding caption.
spans = [
_span("CAPECITABIN", 309, 40.0, bold=True),
_span("Tên chung quốc tế", 309, 50.0, bold=True),
_span("Capecitabinum.", 309, 60.0),
_span("Tên thương mại", 309, 70.0, bold=True),
_span("Xeloda.", 309, 80.0),
_span("Mức độ theo NCIC", 309, 120.0, block=5),
_span("Bảng 3. Điều chỉnh liều do độc tính.", 309, 100.0),
_span("Ngừng thuốc cho đến khi về mức 0.", 309, 140.0, block=5),
]
region = TableRegion("p309_t0", 309, (40.0, 115.0, 400.0, 180.0),
3, 4, SHAPE_SIMPLE)
m = list(assemble(spans, table_index=index_by_page([region])))[0]
assert len(m.tables) == 1
assert m.tables[0].section_key == "lieu_luong_va_cach_dung"
dosage = m.sections["lieu_luong_va_cach_dung"].text
assert dosage.count("Bảng 3. Điều chỉnh liều do độc tính.") == 1
+4
View File
@@ -53,6 +53,10 @@ def test_real_spelling_variants_found_in_the_book_all_match():
"Hướng dẫn cách sử trí ADR": "huong_dan_xu_tri_adr",
"Quá liều và xử lý": "qua_lieu_va_xu_tri",
"Lọai thuốc": "loai_thuoc",
"Tên chung quốc tế và mã ATC": "ten_chung_quoc_te",
"Dạng bào chế và hàm lượng": "dang_thuoc_va_ham_luong",
"Liều lượng và cách dùng giải độc tố uốn ván hấp phụ đơn giá":
"lieu_luong_va_cach_dung",
}
for text, expected_key in cases.items():
matched = match_section(text)
@@ -0,0 +1,72 @@
from ingestion.validation.readiness import evaluate, evaluate_chunks
def _count(monographs, gate_name):
return next(g.count for g in evaluate(monographs) if g.name == gate_name)
def test_readiness_rejects_a_part_without_source_span_provenance():
corpus = [{
"drug_id": "testdrug",
"source_page_range": [100, 100],
"sections": {
"ma_atc": {
"text": "N00AA00",
"parts": [{"kind": "prose", "text": "N00AA00", "source_span_ids": []}],
},
},
"tables": [],
}]
assert _count(corpus, "section_without_provenance") == 0
assert _count(corpus, "part_without_source_span_ids") == 1
def test_readiness_accepts_part_level_source_span_provenance():
corpus = [{
"drug_id": "testdrug",
"source_page_range": [100, 100],
"sections": {
"ma_atc": {
"text": "N00AA00",
"parts": [{
"kind": "prose", "text": "N00AA00",
"source_span_ids": ["p100_b1_l0_s0"],
}],
},
},
"tables": [],
}]
assert _count(corpus, "part_without_source_span_ids") == 0
def test_readiness_rejects_duplicate_physical_region_ids():
corpus = [{
"drug_id": "testdrug",
"source_page_range": [100, 100],
"sections": {},
"tables": [
{"table_id": "p100_t0", "quarantined": True},
{"table_id": "p100_t0", "quarantined": True},
],
}]
assert _count(corpus, "duplicate_table_id") == 1
def test_chunk_readiness_requires_a_verified_printed_page_range():
chunk = {
"chunk_id": "drug__dose__0",
"drug_id": "drug",
"section_key": "dose",
"chunk_kind": "prose",
"text": "Dose.",
"est_tokens": 2,
"attachments": [],
}
gates = evaluate_chunks([], [chunk])
missing = next(g for g in gates if g.name == "chunk_without_printed_page_range")
assert missing.count == 1
chunk["printed_page_range"] = [101, 102]
gates = evaluate_chunks([], [chunk])
present = next(g for g in gates if g.name == "chunk_without_printed_page_range")
assert present.count == 0