From c4e6dd4f6c60c95ac6d2951693a2ffc40abcde6a Mon Sep 17 00:00:00 2001 From: oobabooga Date: Thu, 16 Jul 2026 19:48:12 -0300 Subject: [PATCH] Studio: extract text from PDF web results (#7154) --- studio/backend/core/inference/tools.py | 117 +++++++++++++-- studio/backend/core/rag/parsers.py | 55 +++++-- studio/backend/tests/test_rag_parsing.py | 50 +++++++ .../tests/test_web_fetch_binary_guard.py | 138 ++++++++++++++++-- 4 files changed, 324 insertions(+), 36 deletions(-) diff --git a/studio/backend/core/inference/tools.py b/studio/backend/core/inference/tools.py index dd268a6bb7..5fd57e1b2c 100644 --- a/studio/backend/core/inference/tools.py +++ b/studio/backend/core/inference/tools.py @@ -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 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. diff --git a/studio/backend/core/rag/parsers.py b/studio/backend/core/rag/parsers.py index 9afddf1d9e..0b42906b85 100644 --- a/studio/backend/core/rag/parsers.py +++ b/studio/backend/core/rag/parsers.py @@ -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": diff --git a/studio/backend/tests/test_rag_parsing.py b/studio/backend/tests/test_rag_parsing.py index 14ab0efe2e..3e259f6bd0 100644 --- a/studio/backend/tests/test_rag_parsing.py +++ b/studio/backend/tests/test_rag_parsing.py @@ -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. diff --git a/studio/backend/tests/test_web_fetch_binary_guard.py b/studio/backend/tests/test_web_fetch_binary_guard.py index 3041ed5c34..10db953913 100644 --- a/studio/backend/tests/test_web_fetch_binary_guard.py +++ b/studio/backend/tests/test_web_fetch_binary_guard.py @@ -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 "�" 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<>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(