diff --git a/graphify/llm.py b/graphify/llm.py index a05c9d4632..a96d164084 100644 --- a/graphify/llm.py +++ b/graphify/llm.py @@ -1884,6 +1884,33 @@ def extract_files_direct( return result +# Estimating a PDF means extracting its text, and packing asks for the same +# file repeatedly while it decides where a chunk ends. Memoise on +# (path, size, mtime) so a corpus of papers is parsed once per run rather than +# once per packing probe, and so a file rewritten mid-run is not served a stale +# estimate. Bounded because a huge corpus should not pin every paper's text in +# memory; the entries are cheap (an int) but the dict should not grow forever. +_PDF_ESTIMATE_CACHE: "dict[tuple, str]" = {} +_PDF_ESTIMATE_CACHE_MAX = 512 + + +def _pdf_text_for_estimate(path: Path) -> str: + """Extracted text of a PDF, memoised for the packing pass.""" + try: + st = path.stat() + key = (str(path), st.st_size, st.st_mtime_ns) + except OSError: + return "" + hit = _PDF_ESTIMATE_CACHE.get(key) + if hit is not None: + return hit + text = _file_to_text(path) + if len(_PDF_ESTIMATE_CACHE) >= _PDF_ESTIMATE_CACHE_MAX: + _PDF_ESTIMATE_CACHE.clear() + _PDF_ESTIMATE_CACHE[key] = text + return text + + def _estimate_file_tokens(unit: "Path | FileSlice") -> int: """Estimate the prompt-token cost of a file or slice under `_read_files` rules. @@ -1908,18 +1935,36 @@ def _estimate_file_tokens(unit: "Path | FileSlice") -> int: # fixed token cost, so estimate by image count rather than (binary) byte size. if _is_vision_image(path): return _IMAGE_TOKEN_ESTIMATE - if _TOKENIZER is None: + + # A PDF's bytes are not what the prompt carries. `_read_files` sends it + # through `_file_to_text` -> `extract_pdf_text`, so estimating from the file + # instead measures a compressed binary: every real PDF Flate-compresses its + # text streams, so the estimate came out several times too SMALL and packing + # overfilled the chunk. On a 400-line fixture the same document estimated at + # 1,334 tokens uncompressed-vs-4,598 actual, and 1,334 vs 4,599 once + # FlateDecode was applied — a 3.45x undercount, which is what a real PDF + # looks like. The chunk then blows the context window and falls into + # adaptive bisection, paying for the same content several times (#2903). + if path.suffix.lower() == ".pdf": + try: + content = _pdf_text_for_estimate(path)[:_FILE_CHAR_CAP] + except Exception: + return 0 + elif _TOKENIZER is None: try: size = path.stat().st_size except OSError: return 0 chars = min(size, _FILE_CHAR_CAP) + _PER_FILE_OVERHEAD_CHARS return chars // _CHARS_PER_TOKEN + else: + try: + content = path.read_text(encoding="utf-8", errors="replace")[:_FILE_CHAR_CAP] + except OSError: + return 0 - try: - content = path.read_text(encoding="utf-8", errors="replace")[:_FILE_CHAR_CAP] - except OSError: - return 0 + if _TOKENIZER is None: + return (len(content) + _PER_FILE_OVERHEAD_CHARS) // _CHARS_PER_TOKEN return len(_TOKENIZER.encode(content, disallowed_special=())) + (_PER_FILE_OVERHEAD_CHARS // _CHARS_PER_TOKEN) diff --git a/tests/test_pdf_token_estimate.py b/tests/test_pdf_token_estimate.py new file mode 100644 index 0000000000..0ad3abca8b --- /dev/null +++ b/tests/test_pdf_token_estimate.py @@ -0,0 +1,168 @@ +"""A PDF's token estimate must measure the text the prompt actually carries. + +`_read_files` sends a PDF through `_file_to_text` -> `extract_pdf_text`, but +`_estimate_file_tokens` read the file with `read_text` and tokenised *that*. A +PDF's bytes are not its text: every real PDF Flate-compresses its content +streams, so the estimate measured a compressed binary and came out several times +too small. + +`_pack_chunks_by_tokens` sizes chunks from that estimate, so an undercount +overfills the chunk, the request blows the model's context window, and +`_extract_with_adaptive_retry` bisects — paying for the same content two, four +or eight times before it fits (#2903). + +Measured on one 400-line fixture, same text, stored both ways: + + uncompressed est=5559 actual=4598 actual/est=0.83x + FlateDecode (real PDFs) est=1334 actual=4599 actual/est=3.45x +""" +import zlib +from pathlib import Path + +import pytest + +from graphify.llm import ( + _FILE_CHAR_CAP, + _estimate_file_tokens, + _file_to_text, + _get_tokenizer, + _read_files, +) + +LINES = [f"Section {i}: the parser calls the tokenizer and emits a node." + for i in range(400)] + + +def _make_pdf(path: Path, lines, *, compress: bool): + """A minimal one-page PDF with a real text layer, stored raw or FlateDecode.""" + ops = "BT /F1 10 Tf 20 780 Td 12 TL\n" + "".join(f"({ln}) Tj T*\n" for ln in lines) + "ET" + raw = ops.encode("latin-1", "replace") + stream = zlib.compress(raw) if compress else raw + filt = b" /Filter /FlateDecode" if compress else b"" + objs = [ + b"<< /Type /Catalog /Pages 2 0 R >>", + b"<< /Type /Pages /Kids [3 0 R] /Count 1 >>", + b"<< /Type /Page /Parent 2 0 R /MediaBox [0 0 612 792] " + b"/Resources << /Font << /F1 5 0 R >> >> /Contents 4 0 R >>", + b"<< /Length " + str(len(stream)).encode() + filt + b" >>\nstream\n" + stream + b"\nendstream", + b"<< /Type /Font /Subtype /Type1 /BaseFont /Helvetica >>", + ] + out = bytearray(b"%PDF-1.4\n") + offsets = [] + for i, body in enumerate(objs, 1): + offsets.append(len(out)) + out += f"{i} 0 obj\n".encode() + body + b"\nendobj\n" + xref = len(out) + out += f"xref\n0 {len(objs) + 1}\n0000000000 65535 f \n".encode() + for off in offsets: + out += f"{off:010d} 00000 n \n".encode() + out += (f"trailer\n<< /Size {len(objs) + 1} /Root 1 0 R >>\n" + f"startxref\n{xref}\n%%EOF\n").encode() + path.write_bytes(bytes(out)) + + +def _actual_tokens(prompt: str) -> int: + tok = _get_tokenizer() + return len(tok.encode(prompt, disallowed_special=())) if tok else len(prompt) // 4 + + +@pytest.fixture +def pdfs(tmp_path): + raw = tmp_path / "raw.pdf" + flate = tmp_path / "flate.pdf" + _make_pdf(raw, LINES, compress=False) + _make_pdf(flate, LINES, compress=True) + if not _file_to_text(flate).strip(): + pytest.skip("pypdf not available or cannot read the fixture") + return tmp_path, raw, flate + + +def test_the_two_encodings_hold_the_same_text(pdfs): + """Precondition: only the storage differs, so any estimate gap is the bug.""" + _, raw, flate = pdfs + assert _file_to_text(raw) == _file_to_text(flate) + assert flate.stat().st_size < raw.stat().st_size / 4, "fixture is not really compressed" + + +@pytest.mark.parametrize("which", ["raw", "flate"]) +def test_estimate_tracks_the_prompt_it_will_build(pdfs, which): + root, raw, flate = pdfs + p = raw if which == "raw" else flate + est = _estimate_file_tokens(p) + actual = _actual_tokens(_read_files([p], root)) + assert est > 0 + ratio = actual / est + assert 0.7 <= ratio <= 1.4, ( + f"{which}: estimate {est} vs actual {actual} ({ratio:.2f}x) — packing " + f"will size chunks from a number that does not describe the prompt" + ) + + +def test_compression_does_not_change_the_estimate(pdfs): + """The tell: identical text stored two ways must estimate the same. It used + to differ by 4x purely because one was compressed.""" + _, raw, flate = pdfs + e_raw, e_flate = _estimate_file_tokens(raw), _estimate_file_tokens(flate) + assert abs(e_raw - e_flate) <= max(8, e_raw * 0.02), ( + f"raw={e_raw} flate={e_flate} — the estimate is reading the container, " + f"not the text" + ) + + +def test_estimate_is_not_derived_from_file_size(pdfs): + """A compressed PDF is a fraction of its text's size; the estimate must not + follow the byte count.""" + _, _, flate = pdfs + assert _estimate_file_tokens(flate) > flate.stat().st_size / 4 + + +def test_repeated_estimates_agree(pdfs): + """Packing probes the same file repeatedly; the memo must not drift.""" + _, _, flate = pdfs + assert len({_estimate_file_tokens(flate) for _ in range(4)}) == 1 + + +def test_a_rewritten_pdf_is_re_estimated(tmp_path): + """The memo keys on size and mtime, so editing a paper mid-run must not + serve the previous estimate. + + Both fixtures stay well under `_FILE_CHAR_CAP` on purpose: past the cap + `_read_files` truncates, so a bigger file legitimately estimates the same + and the test would prove nothing about the memo. + """ + p = tmp_path / "growing.pdf" + _make_pdf(p, LINES[:80], compress=True) + if not _file_to_text(p).strip(): + pytest.skip("pypdf not available") + before = _estimate_file_tokens(p) + _make_pdf(p, LINES[:160], compress=True) + after = _estimate_file_tokens(p) + assert after > before * 1.5, f"stale estimate served: {before} -> {after}" + + +def test_an_unreadable_pdf_estimates_zero_instead_of_raising(tmp_path): + """Packing must not blow up on a corrupt file.""" + bad = tmp_path / "corrupt.pdf" + bad.write_bytes(b"%PDF-1.4\nnot really a pdf\n") + assert _estimate_file_tokens(bad) >= 0 + + +def test_non_pdf_estimates_are_unchanged(tmp_path): + """The fix is scoped to PDFs; a text file still measures its own bytes.""" + f = tmp_path / "notes.md" + f.write_text("# Heading\n\n" + ("word " * 2000), encoding="utf-8") + est = _estimate_file_tokens(f) + actual = _actual_tokens(_read_files([f], tmp_path)) + assert 0.7 <= actual / est <= 1.4 + + +def test_the_cap_still_applies(tmp_path): + """A very long PDF is truncated at _FILE_CHAR_CAP by _read_files, so the + estimate must not exceed what that cap allows.""" + p = tmp_path / "huge.pdf" + _make_pdf(p, LINES * 12, compress=True) + if not _file_to_text(p).strip(): + pytest.skip("pypdf not available") + tok = _get_tokenizer() + ceiling = (_FILE_CHAR_CAP if tok is None else _FILE_CHAR_CAP) + 1000 + assert _estimate_file_tokens(p) <= ceiling