ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Studio: extract text from PDF web results (#7154)

Browse source
This commit is contained in:
oobabooga 2026-07-16 19:48:12 -03:00 • committed by GitHub
commit c4e6dd4f6c
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 324 additions and 36 deletions
Download patch file Download diff file Expand all files Collapse all files

117
studio/backend/core/inference/tools.py
View file

@ -3600,14 +3600,18 @@ _MAX_PAGE_CHARS = 16000 # cap fetched page text (after HTML-to-MD conversion)
# Raw download cap > _MAX_PAGE_CHARS since SSR pages embed large <head> sections
# stripped during conversion; 512 KB still reaches article content.
_MAX_FETCH_BYTES = 512 * 1024
# PDF cross-reference data lives at EOF, so extraction needs the whole body.
_MAX_PDF_FETCH_BYTES = 10 * 1024 * 1024
_MAX_WEB_PDF_PAGES = 50
# Control/undecodable chars, excluding text whitespace and ESC (for ANSI logs).
# Binary when they exceed 12.5%, after allowing 16 minor encoding glitches.
_BINARY_CHAR_RE = re.compile("[\\x00-\\x08\\x0b\\x0c\\x0e-\\x1a\\x1c-\\x1f\\x7f-\\x9f\\ufffd]")
_MIN_BINARY_CHARS = 16
_BINARY_CHAR_DIVISOR = 8
# Common binary signatures that can otherwise look text-heavy when mislabeled.
_PDF_MAGIC = b"%PDF-"
_BINARY_MAGIC = (
b"%PDF-", # PDF
_PDF_MAGIC,
b"PK\x03\x04", # zip / docx / xlsx / pptx / epub / jar
b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1", # OLE / legacy Office
b"\x89PNG\r\n\x1a\n", # PNG
@ -3641,14 +3645,22 @@ def _looks_binary(text: str) -> bool:
)
def _has_binary_magic(data: bytes) -> bool:
"""Whether a common binary signature follows optional BOM or whitespace."""
def _magic_head(data: bytes) -> bytes:
head = data[:1024].lstrip()
for bom, _codec in _UNICODE_BOM_CODECS:
if head.startswith(bom):
head = head.removeprefix(bom).lstrip()
break
return head.startswith(_BINARY_MAGIC)
return head
def _has_pdf_magic(data: bytes) -> bool:
return _magic_head(data).startswith(_PDF_MAGIC)
def _has_binary_magic(data: bytes) -> bool:
"""Whether a common binary signature follows optional BOM or whitespace."""
return _magic_head(data).startswith(_BINARY_MAGIC)
def _has_single_byte_text_evidence(data: bytes) -> bool:
@ -3659,6 +3671,45 @@ def _has_single_byte_text_evidence(data: bytes) -> bool:
return ascii_text_bytes / len(data) >= _MIN_SINGLE_BYTE_ASCII_RATIO
def _extract_pdf_text(data: bytes) -> str:
"""Extract page-delimited text with the same parser used by RAG ingestion."""
from ..rag.parsers import parse_pdf_bytes
pages, total_pages = parse_pdf_bytes(data, max_pages = _MAX_WEB_PDF_PAGES)
page_limit_reached = total_pages > _MAX_WEB_PDF_PAGES
parts: list[str] = []
length = 0
text_limited = False
for page in pages:
page_text = page.text.strip()
if not page_text:
continue
section = f"## Page {page.page_number}\n\n{page_text}"
piece = ("\n\n" if parts else "") + section
remaining = _MAX_PAGE_CHARS - length
if len(piece) > remaining:
parts.append(piece[:remaining])
text_limited = True
break
parts.append(piece)
length += len(piece)
text = "".join(parts).rstrip()
if not text:
if page_limit_reached:
return f"(PDF contains no extractable text in the first {_MAX_WEB_PDF_PAGES} pages)"
return ""
limits = []
if text_limited:
limits.append(f"text limited to {_MAX_PAGE_CHARS:,} characters")
if page_limit_reached:
limits.append(f"page processing capped at {_MAX_WEB_PDF_PAGES} pages")
if limits:
marker = f"\n\n... (PDF extraction {'; '.join(limits)})"
text = text[: _MAX_PAGE_CHARS - len(marker)].rstrip() + marker
return text
_USER_AGENTS = (
"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
"Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/131.0.0.0 Safari/537.36",
@ -4054,29 +4105,71 @@ def _fetch_url_raw(
return reason2, "", ""
current_host = rp.hostname
continue
# get_content_type() defaults to "text/plain" when the header is
# absent (RFC 2045); report "" instead so callers can tell a missing
# header apart from a server that really declared text/plain.
if resp.headers.get("Content-Type") is None:
content_type = ""
else:
content_type = (resp.headers.get_content_type() or "").lower()
# Success: read the capped body enforcing the budget between chunks
# (see _read_capped_body), so a slow-drip server can't stretch a
# single resp.read past the deadline.
declared_pdf = content_type == "application/pdf"
read_limit = _MAX_PDF_FETCH_BYTES + 1 if declared_pdf else max_bytes
body_error, raw_bytes = _read_capped_body(
resp,
max_bytes,
read_limit,
timeout,
deadline,
cancel_event,
)
if body_error is not None:
return body_error, "", ""
# A missing or wrong PDF MIME type is common: once the initial text-sized
# read identifies PDF magic, finish the bounded download to reach the EOF xref.
if not declared_pdf and len(raw_bytes) == max_bytes and _has_pdf_magic(raw_bytes):
tail_error, tail = _read_capped_body(
resp,
_MAX_PDF_FETCH_BYTES - max_bytes + 1,
timeout,
deadline,
cancel_event,
)
if tail_error is not None:
return tail_error, "", ""
raw_bytes += tail
break
else:
return "Failed to fetch URL: too many redirects.", "", ""
# get_content_type() defaults to "text/plain" when the header is
# absent (RFC 2045); report "" instead so callers can tell a missing
# header apart from a server that really declared text/plain.
if resp.headers.get("Content-Type") is None:
content_type = ""
else:
content_type = (resp.headers.get_content_type() or "").lower()
is_pdf = declared_pdf or _has_pdf_magic(raw_bytes)
if is_pdf:
if len(raw_bytes) > _MAX_PDF_FETCH_BYTES:
return (
"(PDF content exceeds the download limit; not readable as text)",
"",
content_type,
)
budget_error = _fetch_budget_exceeded(deadline, cancel_event)
if budget_error is not None:
return budget_error, "", content_type
try:
pdf_text = _extract_pdf_text(raw_bytes)
except Exception as exc:
logger.debug("web PDF text extraction failed (%s)", type(exc).__name__)
return "(PDF content could not be read as text)", "", content_type
budget_error = _fetch_budget_exceeded(deadline, cancel_event)
if budget_error is not None:
return budget_error, "", content_type
if not pdf_text:
pdf_text = "(PDF contains no extractable text)"
# Report the true type even for a mislabeled body so the caller's "html"
# check routes the extracted text to the plain-text path, not html_to_markdown.
return None, pdf_text, "application/pdf"
# Reject known-binary MIME types before decoding. Binary is returned as the
# error string so the caller surfaces the placeholder, not replacement chars.

55
studio/backend/core/rag/parsers.py
View file

@ -103,7 +103,7 @@ def _markdown_incomplete(markdown: str, plain: str) -> bool:
return markdown_letters < _PDF_INCOMPLETE_RATIO * plain_letters
def _pdf_markdown(doc) -> list[str] | None:
def _pdf_markdown(doc, pages: range | None = None) -> list[str] | None:
"""Per-page layout-aware Markdown (tables, headings, lists) via pymupdf4llm; index
i maps to page i+1. Returns None when the lib is missing, extraction fails, or the
page count does not line up, so the caller falls back to plain PyMuPDF text."""
@ -112,28 +112,44 @@ def _pdf_markdown(doc) -> list[str] | None:
except Exception:
return None
try:
chunks = pymupdf4llm.to_markdown(
doc,
page_chunks = True,
show_progress = False,
)
kwargs = {"page_chunks": True, "show_progress": False}
if pages is not None:
kwargs["pages"] = list(pages)
chunks = pymupdf4llm.to_markdown(doc, **kwargs)
except Exception: # noqa: BLE001 - never let Markdown extraction break ingestion
logger.warning("pymupdf4llm extraction failed; using plain text", exc_info = True)
return None
if not isinstance(chunks, list) or len(chunks) != doc.page_count:
expected_pages = doc.page_count if pages is None else len(pages)
if not isinstance(chunks, list) or len(chunks) != expected_pages:
return None
return [str(c.get("text") or "") for c in chunks]
def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
def _pdf(
source: str | bytes,
want_images: bool,
max_pages: int | None = None,
) -> tuple[list[Page], list[ParsedImage], int]:
import fitz # PyMuPDF
pages: list[Page] = []
images: list[ParsedImage] = []
doc = fitz.open(path)
doc = (
fitz.open(stream = source, filetype = "pdf") if isinstance(source, bytes) else fitz.open(source)
)
try:
md = _pdf_markdown(doc) if config.PDF_MARKDOWN else None
for i, page in enumerate(doc):
if doc.needs_pass:
raise ValueError("encrypted PDF requires a password")
total_pages = doc.page_count
page_numbers = range(total_pages if max_pages is None else min(total_pages, max_pages))
if not config.PDF_MARKDOWN:
md = None
elif max_pages is None:
md = _pdf_markdown(doc)
else:
md = _pdf_markdown(doc, page_numbers)
for i, page_number in enumerate(page_numbers):
page = doc[page_number]
plain = page.get_text("text") or ""
candidate = md[i] if md else ""
# Prefer layout-aware Markdown (keeps tables/headings legible for retrieval),
@ -147,7 +163,7 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
text = candidate
else:
text = plain
pages.append(_page(text, i + 1))
pages.append(_page(text, page_number + 1))
if want_images:
for img in page.get_images(full = True):
xref = img[0]
@ -161,13 +177,22 @@ def _pdf(path: str, want_images: bool) -> tuple[list[Page], list[ParsedImage]]:
images.append(
ParsedImage(
image_bytes = image_bytes,
page_number = i + 1,
page_number = page_number + 1,
xref = xref,
)
)
finally:
doc.close()
return pages, images
return pages, images, total_pages
def parse_pdf_bytes(data: bytes, *, max_pages: int | None = None) -> tuple[list[Page], int]:
"""Extract PDF pages from an in-memory download using the ingestion parser.
Returns the (capped) pages plus the document's full page count, so a caller
that set ``max_pages`` can tell a fully-read short PDF from a truncated one."""
pages, _images, total_pages = _pdf(data, want_images = False, max_pages = max_pages)
return pages, total_pages
def _merge_rects(boxes: list) -> list:
@ -416,7 +441,7 @@ def parse(path: str, *, want_images: bool = False):
ext = os.path.splitext(path)[1].lower()
if ext == ".pdf":
pages, images = _pdf(path, want_images)
pages, images, _total = _pdf(path, want_images)
return (pages, images) if want_images else pages
if ext == ".docx":

50
studio/backend/tests/test_rag_parsing.py
View file

@ -54,6 +54,56 @@ def test_pdf_markdown_off_uses_plain_text(tmp_path, monkeypatch):
assert "#" not in text and "|" not in text # plain text path emits no Markdown markup
def test_pdf_bytes_use_same_extraction_path(tmp_path, monkeypatch):
from core.rag import config, parsers
monkeypatch.setattr(config, "PDF_MARKDOWN", False)
pdf = tmp_path / "table.pdf"
_table_pdf(pdf)
from_file = parsers.parse(str(pdf))
from_bytes, total_pages = parsers.parse_pdf_bytes(pdf.read_bytes())
assert [page.text for page in from_bytes] == [page.text for page in from_file]
assert total_pages == len(from_file)
def test_pdf_bytes_limit_pages_before_extraction(monkeypatch):
import pymupdf
from core.rag import config, parsers
monkeypatch.setattr(config, "PDF_MARKDOWN", False)
doc = pymupdf.open()
for marker in ("page one", "page two", "page three"):
page = doc.new_page()
page.insert_text((40, 40), marker)
data = doc.tobytes()
doc.close()
pages, total_pages = parsers.parse_pdf_bytes(data, max_pages = 2)
assert len(pages) == 2
assert "page two" in pages[-1].text
assert total_pages == 3 # full count, not the 2 extracted
def test_pdf_markdown_receives_page_limit(monkeypatch):
from core.rag import parsers
captured = {}
class _FakePymupdf4llm:
@staticmethod
def to_markdown(doc, **kwargs):
captured.update(kwargs)
return [{"text": "page"} for _ in kwargs["pages"]]
class _Doc:
page_count = 100
monkeypatch.setitem(__import__("sys").modules, "pymupdf4llm", _FakePymupdf4llm)
assert parsers._pdf_markdown(_Doc(), range(2)) == ["page", "page"]
assert captured == {"page_chunks": True, "show_progress": False, "pages": [0, 1]}
def test_pdf_markdown_passes_only_supported_legacy_kwargs(monkeypatch):
# The pinned PyMuPDF4LLM legacy path ignores unknown kwargs; do not pass the
# newer layout-only OCR knobs or Markdown extraction silently loses policy control.

138
studio/backend/tests/test_web_fetch_binary_guard.py
View file

@ -59,6 +59,18 @@ def _fetch_with(monkeypatch, body: bytes, content_type: str | None) -> str:
return tools._fetch_page_text("https://example.com/thing", timeout = 5)
def _pdf_bytes(*page_texts: str) -> bytes:
pymupdf = pytest.importorskip("pymupdf")
doc = pymupdf.open()
for text in page_texts:
page = doc.new_page()
if text:
page.insert_textbox(pymupdf.Rect(40, 40, 550, 750), text, fontsize = 11)
data = doc.tobytes()
doc.close()
return data
@pytest.mark.parametrize(
"content_type,expected",
[
@ -90,10 +102,119 @@ def test_is_text_candidate_content_type(content_type, expected):
assert tools._is_text_candidate_content_type(content_type) is expected
def test_pdf_rejected_by_content_type(monkeypatch):
out = _fetch_with(monkeypatch, b"%PDF-1.7\n\xff\xd8\xff\x00\x89PNG" * 200, "application/pdf")
assert "<EFBFBD>" not in out
assert "non-text content" in out and "application/pdf" in out
@pytest.mark.parametrize(
"content_type",
["application/pdf", "application/octet-stream", "text/html", "text/plain", None],
)
def test_pdf_text_extracted(monkeypatch, content_type):
out = _fetch_with(
monkeypatch,
_pdf_bytes("First page marker", "Second page marker"),
content_type,
)
assert "## Page 1\n\nFirst page marker" in out
assert "## Page 2" in out and "Second page marker" in out
assert "binary content" not in out and "non-text content" not in out
@pytest.mark.parametrize("content_type", ["application/pdf", "text/plain"])
def test_malformed_pdf_returns_safe_placeholder(monkeypatch, content_type):
out = _fetch_with(monkeypatch, b"%PDF-1.7\nnot a complete PDF", content_type)
assert out == "(PDF content could not be read as text)"
def test_pdf_without_text_layer_reported(monkeypatch):
out = _fetch_with(monkeypatch, _pdf_bytes(""), "application/pdf")
assert out == "(PDF contains no extractable text)"
def test_encrypted_pdf_returns_safe_placeholder(monkeypatch):
pymupdf = pytest.importorskip("pymupdf")
doc = pymupdf.open()
doc.new_page().insert_text((40, 40), "private text")
data = doc.tobytes(
encryption = pymupdf.PDF_ENCRYPT_AES_256,
owner_pw = "owner",
user_pw = "secret",
)
doc.close()
out = _fetch_with(monkeypatch, data, "application/pdf")
assert out == "(PDF content could not be read as text)"
def test_pdf_download_limit_enforced(monkeypatch):
monkeypatch.setattr(tools, "_MAX_PDF_FETCH_BYTES", 256)
out = _fetch_with(monkeypatch, _pdf_bytes("Readable but oversized"), "application/pdf")
assert out == "(PDF content exceeds the download limit; not readable as text)"
def test_mislabeled_pdf_is_read_past_text_download_cap(monkeypatch):
body = _pdf_bytes("Cross-reference data was fetched")
monkeypatch.setattr(tools, "_MAX_FETCH_BYTES", 128)
monkeypatch.setattr(tools, "_MAX_PDF_FETCH_BYTES", len(body) + 100)
out = _fetch_with(monkeypatch, body, "text/plain")
assert "Cross-reference data was fetched" in out
def test_pdf_extraction_caps_pages_and_intermediate_text(monkeypatch):
from core.rag.parsers import Page
seen = {}
def fake_parse(data, *, max_pages = None):
seen["max_pages"] = max_pages
pages = [Page(text = "x" * 1000, page_number = i, char_count = 1000) for i in range(1, 51)]
return pages, 60 # document actually has more pages than the cap
monkeypatch.setattr("core.rag.parsers.parse_pdf_bytes", fake_parse)
text = tools._extract_pdf_text(b"unused")
assert seen["max_pages"] == tools._MAX_WEB_PDF_PAGES
assert len(text) <= tools._MAX_PAGE_CHARS
assert "text limited to 16,000 characters" in text
assert "page processing capped at 50 pages" in text
def test_pdf_exactly_at_page_cap_not_marked_capped(monkeypatch):
from core.rag.parsers import Page
# Exactly _MAX_WEB_PDF_PAGES pages are fully read, so no "capped" marker.
monkeypatch.setattr(
"core.rag.parsers.parse_pdf_bytes",
lambda data, *, max_pages = None: (
[Page(text = "short", page_number = i, char_count = 5) for i in range(1, 51)],
50,
),
)
text = tools._extract_pdf_text(b"unused")
assert "page processing capped" not in text
assert "## Page 50\n\nshort" in text
def test_pdf_page_cap_does_not_claim_later_pages_are_textless(monkeypatch):
from core.rag.parsers import Page
monkeypatch.setattr(
"core.rag.parsers.parse_pdf_bytes",
lambda data, *, max_pages = None: (
[Page(text = "", page_number = i, char_count = 0) for i in range(1, 51)],
60,
),
)
assert tools._extract_pdf_text(b"unused") == (
"(PDF contains no extractable text in the first 50 pages)"
)
def test_pdf_result_discarded_after_fetch_deadline(monkeypatch):
clock = {"time": 1000.0}
monkeypatch.setattr(tools.time, "monotonic", lambda: clock["time"])
def slow_extract(data):
clock["time"] += 10.0
return "late PDF text"
monkeypatch.setattr(tools, "_extract_pdf_text", slow_extract)
out = _fetch_with(monkeypatch, _pdf_bytes("Readable text"), "application/pdf")
assert out == "Failed to fetch URL: timed out."
def test_text_octet_stream_kept_after_sniffing(monkeypatch):
@ -154,7 +275,6 @@ def test_valid_utf8_binary_caught_by_control_chars(monkeypatch):
@pytest.mark.parametrize(
"magic",
[
b"%PDF-",
b"PK\x03\x04",
b"\xd0\xcf\x11\xe0\xa1\xb1\x1a\xe1",
b"\x1f\x8b",
@ -180,8 +300,8 @@ def test_text_labeled_binary_caught_by_magic(monkeypatch, magic):
b"\t\xef\xbb\xbf ",
],
)
def test_pdf_magic_after_harmless_prefix(monkeypatch, prefix):
body = prefix + b"%PDF-1.7\n" + b"1 0 obj<</Type/Catalog>>endobj\n" * 100
def test_binary_magic_after_harmless_prefix(monkeypatch, prefix):
body = prefix + b"\x1f\x8b" + b" printable text-heavy body" * 100
out = _fetch_with(monkeypatch, body, "text/plain")
assert "binary content" in out
@ -243,10 +363,10 @@ def test_html_page_unaffected(monkeypatch):
def test_content_type_sanitized_in_message(monkeypatch):
# Do not echo obs-folded header content into the model response.
out = _fetch_with(monkeypatch, b"\x00\x01\x02" * 500, "application/pdf\r\n data: injected")
out = _fetch_with(monkeypatch, b"PK\x03\x04" * 500, "application/zip\r\n data: injected")
assert "\n" not in out and "\r" not in out
assert "injected" not in out
assert "application/pdf" in out
assert "application/zip" in out
@pytest.mark.parametrize(