mirror of
https://github.com/safishamsi/graphify.git
synced 2026-08-27 08:46:43 +00:00
fix(llm): slice oversized PDFs through the converter (#2906)
An oversized PDF was sliced as raw container bytes, so slice offsets indexed binary data instead of the extracted text the prompt carries. Introduce unit_source_text so both the oversized-file expansion and per-slice reads index the same extracted text, and make PDFs splittable through the converter.
This commit is contained in:
committed by
safishamsi
parent
b6b2fb17c8
commit
4faa86c8d5
+72
-6
@@ -40,6 +40,57 @@ _SPLITTABLE_TEXT_SUFFIXES = frozenset({
|
||||
".qmd", ".skill", ".html", ".yaml", ".yml",
|
||||
})
|
||||
|
||||
# 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.
|
||||
@@ -72,8 +123,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:
|
||||
@@ -131,7 +190,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
|
||||
@@ -146,9 +208,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:
|
||||
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user