import json from pathlib import Path from ingestion.chunk import ( CHUNK_KIND_BLOCK_DESCRIPTOR, 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 ( PART_PROSE, Heading, Monograph, SectionPart, SectionSpan, TableBlock, ) from ingestion.tables import SHAPE_FORMULA_2D, SHAPE_MULTI_HEADER, SHAPE_SIMPLE def _section(key, display, text, page=202): return SectionSpan( key=key, display_name=display, 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 []), ) def _monograph(sections, tables=()): return Monograph( drug_id="ampicilin_va_sulbactam", drug_name="AMPICILIN VÀ SULBACTAM", source_page_range=[200, 203], sections={s.key: s for s in sections}, atc_codes=["J01CR01"], tables=list(tables), ) def _block(block_id="p202_t0", shape=SHAPE_SIMPLE, section_key="lieu_luong_va_cach_dung"): return TableBlock( table_id=block_id, shape=shape, physical_page=202, bbox=[299.0, 189.6, 552.4, 300.5], section_key=section_key, text="Độ thanh thải creatinin Nửa đời Liều 1,5 - 3,0 g", quarantined=True, ) def test_a_section_whose_table_was_lifted_says_so(): """The defect this exists to prevent is silent, not visible. Without the reference, this chunk is grammatical, complete-looking prose with the renal-dosing table absent and nothing marking the absence. """ section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "Liều thường dùng cho người lớn là 1,5 - 3 g mỗi 6 giờ.") monograph = _monograph([section], [_block()]) chunks = chunk_monograph(monograph) prose = [c for c in chunks if c.chunk_kind == CHUNK_KIND_PROSE] assert len(prose) == 1 assert prose[0].has_quarantined_content is True assert [a.block_id for a in prose[0].attachments] == ["p202_t0"] assert prose[0].attachments[0].physical_page == 202 assert prose[0].attachments[0].bbox 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, 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(): section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "Prose.") block = _block() monograph = _monograph([section], [block]) descriptors = [c for c in chunk_monograph(monograph, {"p202_t0": []}) if c.chunk_kind == CHUNK_KIND_BLOCK_DESCRIPTOR] assert "1,5 - 3,0 g" not in descriptors[0].text def test_a_header_row_carrying_a_number_is_refused(): """AMIODARON, physical page 183 — a real case, caught by a gate. pdfplumber reported the first row as "Thời gian liệu pháp tĩnh mạch Liều 720 mg/ngày (0,5 mg/phút)", i.e. a dose inside what it called a header, from an extraction never verified by eye. Measured: 42 of 124 simple-table headers (34%) contain a digit. """ assert _is_label_row(["Các Statin", "Khởi đầu", "Liều duy trì"]) is True assert _is_label_row(["Liều 720 mg/ngày (0,5 mg/phút)"]) is False assert _is_label_row(["x" * 45]) is False assert _is_label_row([]) is False section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "Prose.") monograph = _monograph([section], [_block()]) chunks = chunk_monograph( monograph, {"p202_t0": ["Liều 720 mg/ngày (0,5 mg/phút)"]}) descriptor = next(c for c in chunks if c.chunk_kind == CHUNK_KIND_BLOCK_DESCRIPTOR) assert "720" not in descriptor.text assert descriptor.attachments[0].header_row == [] 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": ["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 == [] 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(): section = _section("than_trong", "Thận trọng", "Prose.") block = _block(block_id="p1042_f0", shape=SHAPE_FORMULA_2D, section_key="than_trong") monograph = _monograph([section], [block]) descriptor = next(c for c in chunk_monograph(monograph) if c.chunk_kind == CHUNK_KIND_BLOCK_DESCRIPTOR) assert "công thức" in descriptor.text assert "bảng" not in descriptor.text def test_attachments_do_not_change_the_prose_text(): """The condition under which this feature was accepted at all.""" section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "Liều thường dùng cho người lớn là 1,5 - 3 g mỗi 6 giờ.") with_block = chunk_section(_monograph([section], [_block()]), section, [_block()]) without = chunk_section(_monograph([section]), section) prose_with = [c for c in with_block if c.chunk_kind == CHUNK_KIND_PROSE] assert [c.text for c in prose_with] == [c.text for c in without] assert [c.chunk_id for c in prose_with] == [c.chunk_id for c in without] def test_a_section_with_no_text_but_a_block_still_yields_the_descriptor(): section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "") chunks = chunk_monograph(_monograph([section], [_block()])) assert [c.chunk_kind for c in chunks] == [CHUNK_KIND_BLOCK_DESCRIPTOR] def test_written_chunks_declare_their_schema_version(tmp_path: Path): section = _section("chi_dinh", "Chỉ định", "Nhiễm khuẩn.") chunks = chunk_monograph(_monograph([section])) out = tmp_path / "chunks.jsonl" assert write_chunks_jsonl(chunks, out) == 1 record = json.loads(out.read_text(encoding="utf-8").splitlines()[0]) assert record["schema_version"] == SCHEMA_VERSION assert record["chunk_kind"] == CHUNK_KIND_PROSE assert record["has_quarantined_content"] is False def test_describe_block_names_the_page_even_with_no_header(): section = _section("lieu_luong_va_cach_dung", "Liều lượng và cách dùng", "Prose.") monograph = _monograph([section], [_block()]) 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, 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)