graphify/tests/test_pdf_token_estimate.py

"""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