fix(llm): slice oversized PDFs through the converter (#2906)

graphify/file_slice.py

@@ -28,6 +28,57 @@
 # docs (.pdf) — those are never sliced.
 _SPLITTABLE_TEXT_SUFFIXES = frozenset({".md", ".mdx", ".markdown", ".txt", ".rst"})
 
+# Document types whose BYTES are not what the model is shown. `llm._file_to_text`
+# routes these through a converter, so a character range has to be taken over the
+# converted text, never over the file. Kept separate from the set above because
+# they are splittable for a different reason and via a different reader.
+_CONVERTED_TEXT_SUFFIXES = frozenset({".pdf"})
+
+
+def _pdf_text(path: Path) -> str:
+    """Extracted text of a PDF — the same string `llm._file_to_text` builds.
+
+    Imported lazily from ``detect`` so this module keeps no import-time
+    dependency on the extraction stack (``llm`` imports *this* module, so the
+    reverse direction would be a cycle).
+    """
+    from graphify.detect import extract_pdf_text
+    return extract_pdf_text(path)
+
+
+# Slicing a PDF means extracting its text, and the slicing pass asks for the same
+# file several times: once to measure it, then once per slice as the prompt is
+# built. Memoised on (path, size, mtime_ns) so a corpus of papers is parsed once
+# rather than once per slice, and so a file rewritten mid-run is re-read instead
+# of served a stale body. Bounded: the entries are whole documents, and a large
+# corpus should not pin all of them in memory.
+_CONVERTED_TEXT_CACHE: "dict[tuple, str]" = {}
+_CONVERTED_TEXT_CACHE_MAX = 64
+
+
+def unit_source_text(path: Path) -> str:
+    """The text a unit contributes to the prompt, whatever its container.
+
+    Plain-text files are read directly; converted types (PDF) go through their
+    converter. Both `expand_oversized_files` and `read_slice_text` use this, so
+    the offsets a slice carries always index the same string the model sees.
+    """
+    if path.suffix.lower() not in _CONVERTED_TEXT_SUFFIXES:
+        return path.read_text(encoding="utf-8", errors="replace")
+    try:
+        st = path.stat()
+        key = (str(path), st.st_size, st.st_mtime_ns)
+    except OSError:
+        return _pdf_text(path)
+    hit = _CONVERTED_TEXT_CACHE.get(key)
+    if hit is not None:
+        return hit
+    text = _pdf_text(path)
+    if len(_CONVERTED_TEXT_CACHE) >= _CONVERTED_TEXT_CACHE_MAX:
+        _CONVERTED_TEXT_CACHE.clear()
+    _CONVERTED_TEXT_CACHE[key] = text
+    return text
+
 # Boundary preferences, strongest first. A Markdown heading (``\n#``) keeps a
 # section with its title; a blank line keeps a paragraph intact; a bare newline
 # avoids cutting mid-line. If none is found in the window we hard-cut.
@@ -60,8 +111,16 @@ def unit_path(unit: "Path | FileSlice") -> Path:
 
 
 def is_splittable_text(path: Path) -> bool:
-    """True for plain-text document types that may be sliced."""
-    return path.suffix.lower() in _SPLITTABLE_TEXT_SUFFIXES
+    """True for document types that may be sliced.
+
+    Covers plain text read straight off disk and converted types (PDF) whose
+    text is produced by a converter. Both are sliceable because
+    :func:`unit_source_text` gives the slicing pass the same string the prompt
+    will carry; what disqualifies a type is having no text at all (an image) or
+    text the reader cannot address by character offset.
+    """
+    suffix = path.suffix.lower()
+    return suffix in _SPLITTABLE_TEXT_SUFFIXES or suffix in _CONVERTED_TEXT_SUFFIXES
 
 
 def _best_cut(text: str, start: int, end: int) -> int:
@@ -119,7 +178,10 @@ def expand_oversized_files(
             out.append(f)
             continue
         try:
-            text = f.read_text(encoding="utf-8", errors="replace")
+            # The CONVERTED text for a PDF, so the boundaries below index the
+            # same string read_slice_text will later slice and the prompt will
+            # carry — not the container's bytes (#2906).
+            text = unit_source_text(f)
         except OSError:
             out.append(f)
             continue
@@ -134,9 +196,13 @@ def expand_oversized_files(
 
 
 def read_slice_text(fs: FileSlice) -> str:
-    """Read just this slice's characters from its parent file."""
-    text = fs.path.read_text(encoding="utf-8", errors="replace")
-    return text[fs.start:fs.end]
+    """Read just this slice's characters from its parent file.
+
+    Goes through :func:`unit_source_text`, so a PDF slice indexes the extracted
+    text rather than the container's bytes — the offsets `expand_oversized_files`
+    computed and the string the prompt carries are then the same string (#2906).
+    """
+    return unit_source_text(fs.path)[fs.start:fs.end]
 
 
 def bisect_slice(fs: FileSlice) -> tuple[FileSlice, FileSlice] | None:

tests/test_pdf_slicing.py

@@ -0,0 +1,185 @@
+"""A PDF over the character cap must be sliced, not truncated.
+
+`_read_files` caps every unit at `_FILE_CHAR_CAP` (20,000 characters).
+`expand_oversized_files` slices oversized documents so the whole file still
+reaches the model — but it read candidates with `path.read_text()` and
+`read_slice_text` sliced the same way, while `llm._file_to_text` routes a PDF
+through `extract_pdf_text`. Slicing a PDF would therefore have indexed the
+container's bytes rather than its text, so PDFs were excluded from slicing
+altogether and simply lost everything past 20,000 characters.
+
+Papers are the longest documents anyone points graphify at, and a compressed
+PDF gives no hint of its text length: the fixture here is 3,094 bytes on disk
+and 55,690 characters of text (#2906).
+
+`unit_source_text` is the fix: one reader that returns the string the prompt
+will carry, used by both the boundary pass and the slice reader, so the offsets
+a `FileSlice` holds always index the same string.
+"""
+import zlib
+from pathlib import Path
+
+import pytest
+
+from graphify.file_slice import (
+    FileSlice,
+    expand_oversized_files,
+    is_splittable_text,
+    read_slice_text,
+)
+from graphify.llm import _FILE_CHAR_CAP, _file_to_text, _read_files
+
+try:  # the reader this fix introduces
+    from graphify.file_slice import unit_source_text
+except ImportError:  # pre-fix tree — describe the expected text the same way
+    unit_source_text = _file_to_text  # type: ignore[assignment]
+
+LINES = [f"Section {i}: the parser calls the tokenizer and emits a node."
+         for i in range(900)]
+
+
+def _make_pdf(path: Path, lines, *, compress: bool = True):
+    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))
+
+
+@pytest.fixture
+def big_pdf(tmp_path):
+    p = tmp_path / "paper.pdf"
+    _make_pdf(p, LINES)
+    if not _file_to_text(p).strip():
+        pytest.skip("pypdf not available or cannot read the fixture")
+    return tmp_path, p
+
+
+def test_the_fixture_is_the_shape_the_bug_needs(big_pdf):
+    """Precondition: small on disk, large in text — which is why file size never
+    revealed the problem."""
+    _, p = big_pdf
+    text = unit_source_text(p)
+    assert len(text) > 2 * _FILE_CHAR_CAP
+    assert p.stat().st_size < len(text) / 4
+
+
+def test_a_pdf_is_splittable(big_pdf):
+    _, p = big_pdf
+    assert is_splittable_text(p)
+
+
+def test_an_oversized_pdf_is_sliced(big_pdf):
+    _, p = big_pdf
+    units = expand_oversized_files([p], _FILE_CHAR_CAP)
+    assert len(units) > 1
+    assert all(isinstance(u, FileSlice) for u in units)
+
+
+def test_the_slices_tile_the_extracted_text_exactly(big_pdf):
+    """Gap-free and non-overlapping: rejoining reproduces the document."""
+    _, p = big_pdf
+    units = expand_oversized_files([p], _FILE_CHAR_CAP)
+    assert "".join(read_slice_text(u) for u in units) == unit_source_text(p)
+
+
+def test_no_slice_exceeds_the_cap(big_pdf):
+    _, p = big_pdf
+    for u in expand_oversized_files([p], _FILE_CHAR_CAP):
+        assert len(read_slice_text(u)) <= _FILE_CHAR_CAP
+
+
+def test_a_slice_indexes_text_not_bytes(big_pdf):
+    """The bug in one assertion: slicing used to address the container. A slice
+    must not contain PDF structure."""
+    _, p = big_pdf
+    first = expand_oversized_files([p], _FILE_CHAR_CAP)[0]
+    body = read_slice_text(first)
+    assert "%PDF" not in body
+    assert "endstream" not in body
+    assert "Section 0" in body
+
+
+def test_the_tail_reaches_the_prompt(big_pdf):
+    """The symptom a user would notice: content past 20k could never become a
+    node because the model never saw it."""
+    root, p = big_pdf
+    text = unit_source_text(p)
+    tail = text[-80:].strip()[:40]
+    units = expand_oversized_files([p], _FILE_CHAR_CAP)
+    joined = "".join(_read_files([u], root) for u in units)
+    assert tail and tail in joined
+
+
+# ---------------------------------------------------------------------------
+# What must NOT change
+# ---------------------------------------------------------------------------
+
+def test_a_small_pdf_still_passes_through_whole(tmp_path):
+    p = tmp_path / "note.pdf"
+    _make_pdf(p, LINES[:20])
+    if not _file_to_text(p).strip():
+        pytest.skip("pypdf not available")
+    assert expand_oversized_files([p], _FILE_CHAR_CAP) == [p]
+
+
+def test_plain_text_slicing_is_unchanged(tmp_path):
+    f = tmp_path / "doc.md"
+    body = "# H\n\n" + ("word " * 12000)
+    f.write_text(body, encoding="utf-8")
+    units = expand_oversized_files([f], _FILE_CHAR_CAP)
+    assert len(units) > 1
+    assert "".join(read_slice_text(u) for u in units) == body
+
+
+def test_images_and_code_are_still_not_sliced(tmp_path):
+    img = tmp_path / "a.png"
+    img.write_bytes(b"\x89PNG\r\n\x1a\n" + b"x" * 40_000)
+    code = tmp_path / "m.py"
+    code.write_text("def f():\n    pass\n" * 4000, encoding="utf-8")
+    assert not is_splittable_text(img)
+    assert not is_splittable_text(code)
+    assert expand_oversized_files([img, code], _FILE_CHAR_CAP) == [img, code]
+
+
+def test_a_corrupt_pdf_does_not_break_the_pass(tmp_path):
+    bad = tmp_path / "corrupt.pdf"
+    bad.write_bytes(b"%PDF-1.4\nnot really a pdf\n")
+    assert expand_oversized_files([bad], _FILE_CHAR_CAP) == [bad]
+
+
+def test_a_rewritten_pdf_is_re_read(tmp_path):
+    """The reader memoises on (path, size, mtime), so a paper replaced mid-run
+    must not be sliced against the previous text."""
+    p = tmp_path / "growing.pdf"
+    _make_pdf(p, LINES[:60])
+    if not _file_to_text(p).strip():
+        pytest.skip("pypdf not available")
+    before = len(unit_source_text(p))
+    _make_pdf(p, LINES[:600])
+    after = len(unit_source_text(p))
+    assert after > before * 5, f"stale text served: {before} -> {after}"
+
+
+def test_repeated_reads_agree(big_pdf):
+    _, p = big_pdf
+    assert len({unit_source_text(p) for _ in range(3)}) == 1