melbinjp/DocQA

1"""Selecting pages to look at, and folding what was read back in.2 3The rules that matter here are about not losing or corrupting the document:4a transcription is what a model believes a page says, and it must never quietly5replace what the PDF actually says. The exception is a page with no text, where6there is nothing to protect and without the transcription the page is gone.7"""8import fitz9import pytest10 11from utils.vision import (12    LOW_TEXT_THRESHOLD,13    MAX_VISION_PAGES,14    is_transcript,15    merge,16    pages_for_vision,17    wrap_transcript,18)19 20 21def make_pdf(page_texts):22    """`insert_text` draws one unwrapped line that runs off the page edge, so23    only the part inside the page is extractable. A text box wraps, which is24    what a page of real prose looks like."""25    doc = fitz.open()26    for text in page_texts:27        page = doc.new_page()28        if text:29            page.insert_textbox(fitz.Rect(50, 50, 545, 790), text, fontsize=11)30    data = doc.tobytes()31    doc.close()32    return data33 34 35def test_the_fixture_really_makes_a_dense_page():36    """Guard the fixture itself: an earlier version silently produced 9737    characters and made the dense-page test meaningless."""38    raw = make_pdf(["word " * 200])39    doc = fitz.open(stream=raw, filetype="pdf")40    assert len(doc[0].get_text().strip()) > LOW_TEXT_THRESHOLD41    doc.close()42 43 44def test_a_page_with_no_text_is_selected():45    raw = make_pdf(["", "x" * 400])46    chosen = {n: reason for n, _, reason in pages_for_vision(raw)}47    assert chosen.get(1) == "no-text"48 49 50def test_a_dense_page_with_nothing_on_it_is_left_alone():51    raw = make_pdf(["word " * 200])52    assert pages_for_vision(raw) == []53 54 55def test_pages_with_no_text_are_prioritised_over_the_rest():56    raw = make_pdf([""] * 3 + ["word " * 200] * 3)57    chosen = pages_for_vision(raw, max_pages=2)58    assert [reason for _, _, reason in chosen] == ["no-text", "no-text"]59 60 61def test_selection_is_capped():62    raw = make_pdf([""] * (MAX_VISION_PAGES + 6))63    assert len(pages_for_vision(raw)) == MAX_VISION_PAGES64 65 66def test_rendered_pages_come_back_as_images_in_page_order():67    raw = make_pdf(["", "", ""])68    chosen = pages_for_vision(raw)69    assert [n for n, _, _ in chosen] == [1, 2, 3]70    for _, image, _ in chosen:71        assert image[:2] == b"\xff\xd8", "not a JPEG"72        # Keeps a page small enough that several in one ingest is not73        # itself the failure.74        assert len(image) < 2_000_000, f"{len(image)} bytes is too big to send"75 76 77def test_bytes_that_are_not_a_pdf_select_nothing_rather_than_raising():78    assert pages_for_vision(b"not a pdf at all") == []79    assert pages_for_vision(b"") == []80 81 82def test_a_transcript_is_marked_as_one():83    assert is_transcript(wrap_transcript("|a|b|"))84    assert not is_transcript("ordinary extracted text")85    assert wrap_transcript("   ") == ""86 87 88def test_merging_keeps_the_documents_own_text():89    pages = [(1, "The real extracted text.")]90    out = dict(merge(pages, {1: "|col|col|"}))91    assert "The real extracted text." in out[1], "extraction was overwritten"92    assert "|col|col|" in out[1]93 94 95def test_a_page_that_had_no_text_is_created_from_the_transcript():96    """The scanned case. Without this the page does not exist at all."""97    out = dict(merge([], {3: "Transcribed from the image."}))98    assert out[3].endswith("Transcribed from the image.")99    assert is_transcript(out[3])100 101 102def test_an_empty_transcript_changes_nothing():103    pages = [(1, "Original.")]104    assert merge(pages, {1: "   "}) == [(1, "Original.")]105    assert merge(pages, {}) == [(1, "Original.")]106 107 108def test_pages_come_back_in_order_after_merging():109    out = merge([(5, "five"), (1, "one")], {3: "three"})110    assert [n for n, _ in out] == [1, 3, 5]111 112 113def test_the_low_text_threshold_is_small_enough_not_to_catch_real_pages():114    """A page of genuine prose must not be mistaken for a scan."""115    assert LOW_TEXT_THRESHOLD < 400116 117 118@pytest.mark.asyncio119async def test_no_targets_returns_an_empty_result_and_no_errors():120    from utils.vision import transcribe_pages121    assert await transcribe_pages(None, "m", []) == ({}, [])122 123 124@pytest.mark.asyncio125async def test_a_failing_page_reports_why_instead_of_vanishing():126    """The first version swallowed these, every page failed, and the only127    symptom was a document that would not ingest for no stated reason."""128    from utils.vision import transcribe_pages129 130    class Boom:131        class aio:132            class models:133                @staticmethod134                async def generate_content(**kwargs):135                    raise RuntimeError("model refused the image")136 137    out, errors = await transcribe_pages(Boom(), "m", [(1, b"jpegbytes", "no-text")])138    assert out == {}139    assert errors and "model refused the image" in errors[0]140 141 142@pytest.mark.asyncio143async def test_the_reader_walks_the_model_ladder_on_failure():144    """A quota error on the first model is a queue, not a dead end.145 146    Measured 2026-09-01: every page came back "429 RESOURCE_EXHAUSTED" from the147    primary model while ordinary questions were still being answered fine on the148    same key, because the answering path falls back and the reader did not.149    """150    from utils.vision import transcribe_pages151 152    tried = []153 154    class Ladder:155        class aio:156            class models:157                @staticmethod158                async def generate_content(model=None, **kwargs):159                    tried.append(model)160                    if model == "first":161                        raise RuntimeError("429 RESOURCE_EXHAUSTED")162                    class R:163                        text = "transcribed by the second model"164                    return R()165 166    out, errors = await transcribe_pages(167        Ladder(), ["first", "second"], [(1, b"img", "no-text")], timeout=5168    )169    assert tried == ["first", "second"]170    assert out == {1: "transcribed by the second model"}171    assert errors == []172 173 174@pytest.mark.asyncio175async def test_a_single_model_name_still_works():176    from utils.vision import transcribe_pages177 178    class Ok:179        class aio:180            class models:181                @staticmethod182                async def generate_content(**kwargs):183                    class R:184                        text = "read"185                    return R()186 187    out, _ = await transcribe_pages(Ok(), "only-model", [(2, b"img", "no-text")])188    assert out == {2: "read"}189 190 191def test_by_default_only_scanned_pages_are_sent():192    """Tables and figures answered 16 of 16 from text, so spending a request on193    them buys nothing on the evidence available."""194    raw = make_pdf(["", "word " * 200])195    reasons = [r for _, _, r in pages_for_vision(raw)]196    assert reasons == ["no-text"]197 198 199def test_every_chunk_of_a_transcript_is_marked_not_just_the_first():200    """Marking the page and then splitting it marks one chunk in four.201 202    Measured on the live Space: a scanned page became four chunks and only the203    first carried "[Read from the page image]", so three of them were204    indistinguishable from the document's own words. The guarantee is that a205    reader can always tell, which means the marker belongs on the chunk.206    """207    from utils.splitter import split_pages208 209    long_transcript = ("The transcribed sentence repeats. " * 200)210    chunks = split_pages([(1, long_transcript)])211    assert len(chunks) > 1, "this fixture needs to produce several chunks"212 213    marked = [wrap_transcript(c["text"]) for c in chunks]214    assert all(is_transcript(m) for m in marked)215    assert all(m.count("[Read from the page image]") == 1 for m in marked)216 217 218def test_marking_a_chunk_does_not_change_what_it_says():219    body = "Row: base 6 512 2048 8 64 64 0.1 0.1 100K 4.92 25.8 65"220    assert body in wrap_transcript(body)221 222 223# --- a document is a scan, or it is not; bare pages inside a text document224#     are figures, and figures already answer from their captions ---225 226def test_a_text_document_with_a_few_bare_pages_is_not_a_scan():227    """The regression that made a 75-page paper stop ingesting at all.228 229    GPT-3 has figure pages with almost no text, so a per-page rule sent three of230    them to be read. Three pages against a three-model ladder with a long231    timeout on each rung turned a thirty second ingest into a 504.232    """233    raw = make_pdf(["word " * 200] * 8 + [""] * 2)234    assert pages_for_vision(raw) == []235 236 237def test_a_document_that_really_is_scanned_still_qualifies():238    raw = make_pdf([""] * 6 + ["word " * 200] * 2)239    assert len(pages_for_vision(raw)) == 6240 241 242def test_the_all_scope_ignores_the_document_level_gate():243    from utils.vision import SCOPE_ALL244    raw = make_pdf(["word " * 200] * 8 + [""] * 2)245    assert len(pages_for_vision(raw, scope=SCOPE_ALL)) == 2246 247 248@pytest.mark.asyncio249async def test_the_reader_gives_up_at_its_budget_rather_than_hanging():250    """Ingest happens inside one HTTP request. A request that never returns is251    worse than a document that is merely harder to search."""252    import asyncio253    from utils.vision import transcribe_pages254 255    class Slow:256        class aio:257            class models:258                @staticmethod259                async def generate_content(**kwargs):260                    await asyncio.sleep(30)261 262    out, errors = await transcribe_pages(263        Slow(), "m", [(1, b"img", "no-text")], timeout=20, budget=0.4264    )265    assert out == {}266    assert errors and "budget" in errors[0]267