From 9d3d80b9f8df851d475dbd1951cb34360920b80e Mon Sep 17 00:00:00 2001 From: mia Date: Sun, 19 Apr 2026 01:57:25 +0530 Subject: [PATCH] fix: extract text from PDF URLs in fetch_url tool When fetch_url fetches a PDF URL, the content was returned as garbled binary because BeautifulSoup's HTML parser corrupted the PDF bytes. Add three-layer PDF detection: - Fast-path for URLs ending in .pdf (case-insensitive) - Content-Type header detection in async SafeWebBaseLoader._fetch - %PDF binary content fallback in get_content_from_url Includes SSRF protection via validate_url(), streaming download with Content-Length pre-check, 50 MB size limit, and proper session management. No new dependencies (pypdf already in requirements.txt). Closes #23841 Co-Authored-By: Claude Opus 4.6 --- backend/open_webui/retrieval/utils.py | 111 +++++- backend/open_webui/retrieval/web/utils.py | 33 ++ backend/tests/retrieval/test_pdf_handling.py | 343 +++++++++++++++++++ 3 files changed, 485 insertions(+), 2 deletions(-) create mode 100644 backend/tests/retrieval/test_pdf_handling.py diff --git a/backend/open_webui/retrieval/utils.py b/backend/open_webui/retrieval/utils.py index 93ba72ce13..f7b4ba58da 100644 --- a/backend/open_webui/retrieval/utils.py +++ b/backend/open_webui/retrieval/utils.py @@ -10,7 +10,7 @@ from concurrent.futures import ThreadPoolExecutor import time import re -from urllib.parse import quote +from urllib.parse import quote, urlparse from huggingface_hub import snapshot_download from langchain_classic.retrievers import ( ContextualCompressionRetriever, @@ -83,10 +83,117 @@ def get_loader(request, url: str): ) -def get_content_from_url(request, url: str) -> str: +PDF_MAX_SIZE = 50 * 1024 * 1024 # 50 MB + + +def extract_text_from_pdf_bytes(pdf_bytes: bytes) -> str: + """Extract text from raw PDF bytes using pypdf. + + Returns the concatenated text of all pages, or a placeholder message + if no text could be extracted (e.g. image-only PDFs). + """ + import io + from pypdf import PdfReader + + reader = PdfReader(io.BytesIO(pdf_bytes)) + text_parts = [] + for page in reader.pages: + page_text = page.extract_text() + if page_text: + text_parts.append(page_text) + if not text_parts: + return '[PDF contains no extractable text (possibly image-only)]' + return '\n'.join(text_parts) + + +def extract_pdf_from_url(request, url: str) -> str: + """Download a PDF from *url* and extract its text using pypdf. + + Uses the same SSL-verification and proxy settings that the web loader + honours so that corporate proxies / self-signed certs keep working. + The URL is validated against SSRF rules before downloading. + """ + from open_webui.retrieval.web.utils import validate_url + + validate_url(url) + + verify_ssl = request.app.state.config.ENABLE_WEB_LOADER_SSL_VERIFICATION + trust_env = request.app.state.config.WEB_SEARCH_TRUST_ENV + + with requests.Session() as session: + session.verify = verify_ssl + session.trust_env = trust_env + + resp = session.get(url, timeout=60, stream=True) + resp.raise_for_status() + + # Reject oversized PDFs early via Content-Length when available. + try: + content_length = int(resp.headers.get('Content-Length', 0)) + except (ValueError, TypeError): + content_length = 0 + if content_length > PDF_MAX_SIZE: + resp.close() + raise ValueError( + f'PDF at {url} exceeds the {PDF_MAX_SIZE // (1024 * 1024)} MB size limit' + ) + + # Stream the body in chunks and abort if it exceeds the size limit. + # This prevents buffering arbitrarily large responses into memory + # when Content-Length is missing or inaccurate (e.g. chunked encoding). + chunks = [] + downloaded = 0 + for chunk in resp.iter_content(chunk_size=1024 * 1024): + downloaded += len(chunk) + if downloaded > PDF_MAX_SIZE: + resp.close() + raise ValueError( + f'PDF at {url} exceeds the {PDF_MAX_SIZE // (1024 * 1024)} MB size limit' + ) + chunks.append(chunk) + pdf_bytes = b''.join(chunks) + + return extract_text_from_pdf_bytes(pdf_bytes) + + +def get_content_from_url(request, url: str) -> tuple[str, list[Document]]: + # Fast-path: if the URL path ends with .pdf, extract directly with pypdf + # to avoid piping binary data through the HTML-oriented web loader. + parsed_path = urlparse(url).path.lower() + if parsed_path.endswith('.pdf'): + log.debug(f'URL ends with .pdf, extracting directly with pypdf: {url}') + try: + content = extract_pdf_from_url(request, url) + docs = [Document(page_content=content, metadata={'source': url})] + return content, docs + except Exception: + log.exception(f'Failed to extract PDF text from {url}') + content = f'[Error: Unable to extract text from PDF at {url}]' + docs = [Document(page_content=content, metadata={'source': url})] + return content, docs + loader = get_loader(request, url) docs = loader.load() content = ' '.join([doc.page_content for doc in docs]) + + # Fallback: if the content looks like raw PDF binary data (starts with + # %PDF), the web loader failed to properly extract text. This can happen + # when the URL doesn't end in .pdf but the server returns PDF content. + # We must re-download because BeautifulSoup corrupts the binary data + # during HTML parsing, so the bytes we already have are not a valid PDF. + if content.strip().startswith('%PDF'): + log.warning( + f'Detected raw PDF binary in fetched content for {url}, ' + f're-extracting with pypdf' + ) + try: + content = extract_pdf_from_url(request, url) + docs = [Document(page_content=content, metadata={'source': url})] + except Exception: + log.exception(f'Failed to extract PDF text from {url}') + content = f'[Error: Unable to extract text from PDF at {url}]' + docs = [Document(page_content=content, metadata={'source': url})] + return content, docs diff --git a/backend/open_webui/retrieval/web/utils.py b/backend/open_webui/retrieval/web/utils.py index cfe0f71b85..548b71bf7d 100644 --- a/backend/open_webui/retrieval/web/utils.py +++ b/backend/open_webui/retrieval/web/utils.py @@ -519,6 +519,39 @@ class SafeWebBaseLoader(WebBaseLoader): ) as response: if self.raise_for_status: response.raise_for_status() + + # Check if the response is a PDF based on Content-Type header. + # This handles cases where the URL doesn't end in .pdf but the + # server returns PDF content (e.g. dynamic download endpoints). + content_type = response.headers.get('Content-Type', '').lower() + if 'application/pdf' in content_type: + log.debug(f'Detected PDF Content-Type for URL: {url}, extracting text with pypdf') + try: + from open_webui.retrieval.utils import extract_text_from_pdf_bytes, PDF_MAX_SIZE + + # Stream the body in bounded chunks to prevent OOM + # on large PDFs without a Content-Length header. + chunks = [] + downloaded = 0 + async for chunk in response.content.iter_chunked(1024 * 1024): + downloaded += len(chunk) + if downloaded > PDF_MAX_SIZE: + raise ValueError( + f'PDF at {url} exceeds the {PDF_MAX_SIZE // (1024 * 1024)} MB size limit' + ) + chunks.append(chunk) + pdf_bytes = b''.join(chunks) + + import html as _html + text = extract_text_from_pdf_bytes(pdf_bytes) + # Wrap in HTML so that _unpack_fetch_results -> BeautifulSoup + # preserves the extracted text as-is. + escaped = _html.escape(text) + return f'
{escaped}
' + except Exception as e: + log.warning(f'Failed to extract text from PDF at {url}: {e}.') + return f'
[Error: Unable to extract text from PDF at {url}]
' + return await response.text() except aiohttp.ClientConnectionError as e: if i == retries - 1: diff --git a/backend/tests/retrieval/test_pdf_handling.py b/backend/tests/retrieval/test_pdf_handling.py new file mode 100644 index 0000000000..8e24232d13 --- /dev/null +++ b/backend/tests/retrieval/test_pdf_handling.py @@ -0,0 +1,343 @@ +"""Tests for PDF text extraction helpers. + +The open_webui.retrieval.utils module has heavy transitive dependencies +(redis, chromadb, etc.) that are impractical to install in a lightweight +test environment. These tests therefore import only the pure-function +helpers that have no application-level dependencies. + +To run: + PYTHONPATH=backend pytest backend/tests/retrieval/test_pdf_handling.py -v +""" + +import io +import os +import sys +import types +import importlib.util +import pytest +from unittest.mock import Mock, patch + +# --------------------------------------------------------------------------- +# Minimal stubs so we can exec just the target functions out of utils.py +# without pulling the entire app dependency tree. +# --------------------------------------------------------------------------- + +# We import the three functions under test by executing a trimmed snippet +# of the source file. This avoids the huge transitive import chain. + +_UTILS_PATH = os.path.normpath( + os.path.join( + os.path.dirname(__file__), + os.pardir, + os.pardir, + "open_webui", + "retrieval", + "utils.py", + ) +) + + +# --------------------------------------------------------------------------- +# Stub out open_webui.retrieval.web.utils so that the inline +# ``from open_webui.retrieval.web.utils import validate_url`` inside +# extract_pdf_from_url resolves without pulling the full app dependency tree. +# The stubs must remain in sys.modules for the lifetime of this test module +# because the extracted functions import at call-time, not at extraction-time. +# An atexit handler restores original values to avoid leaking into other tests. +# --------------------------------------------------------------------------- +_STUB_KEYS = [ + "open_webui", + "open_webui.retrieval", + "open_webui.retrieval.web", + "open_webui.retrieval.web.utils", +] +_saved_modules = {k: sys.modules.get(k) for k in _STUB_KEYS} + +_web_utils_stub = types.ModuleType("open_webui.retrieval.web.utils") +_web_utils_stub.validate_url = lambda url: True # no-op for tests +sys.modules.setdefault("open_webui", types.ModuleType("open_webui")) +sys.modules.setdefault("open_webui.retrieval", types.ModuleType("open_webui.retrieval")) +sys.modules.setdefault("open_webui.retrieval.web", types.ModuleType("open_webui.retrieval.web")) +sys.modules["open_webui.retrieval.web.utils"] = _web_utils_stub + +import atexit + +def _restore_modules(): + for k, original in _saved_modules.items(): + if original is None: + sys.modules.pop(k, None) + else: + sys.modules[k] = original + +atexit.register(_restore_modules) + + +def _extract_functions(): + """Parse utils.py source and exec just the PDF-related functions.""" + import textwrap + + # Read the raw source + with open(_UTILS_PATH) as f: + source = f.read() + + code = textwrap.dedent("""\ + import io + import logging + import re + import requests + from urllib.parse import urlparse + from langchain_core.documents import Document + + log = logging.getLogger(__name__) + """) + + # Extract the function definitions we care about from the source. + import ast + tree = ast.parse(source) + for node in ast.iter_child_nodes(tree): + # Grab module-level assignments (PDF_MAX_SIZE = ...) + if isinstance(node, ast.Assign): + for target in node.targets: + if isinstance(target, ast.Name) and target.id == "PDF_MAX_SIZE": + code += ast.get_source_segment(source, node) + "\n\n" + + # Grab the function definitions we need + if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef)): + if node.name in ( + "extract_text_from_pdf_bytes", + "extract_pdf_from_url", + "get_content_from_url", + "get_loader", + "is_youtube_url", + ): + code += ast.get_source_segment(source, node) + "\n\n" + + ns = {} + exec(compile(code, _UTILS_PATH, "exec"), ns) + return ns + + +_ns = _extract_functions() +extract_text_from_pdf_bytes = _ns["extract_text_from_pdf_bytes"] +extract_pdf_from_url = _ns["extract_pdf_from_url"] +get_content_from_url = _ns["get_content_from_url"] +PDF_MAX_SIZE = _ns["PDF_MAX_SIZE"] + + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + +def _make_minimal_pdf(text: str = "Hello, World!") -> bytes: + """Create a minimal valid single-page PDF containing *text*.""" + from pypdf import PdfWriter + from pypdf.generic import ( + DecodedStreamObject, + DictionaryObject, + NameObject, + ) + + writer = PdfWriter() + writer.add_blank_page(width=612, height=792) + page = writer.pages[0] + + font_dict = DictionaryObject() + font_dict[NameObject("/Type")] = NameObject("/Font") + font_dict[NameObject("/Subtype")] = NameObject("/Type1") + font_dict[NameObject("/BaseFont")] = NameObject("/Helvetica") + + resources = page.get("/Resources", DictionaryObject()) + if not isinstance(resources, DictionaryObject): + resources = DictionaryObject(resources) + font_resources = resources.get("/Font", DictionaryObject()) + if not isinstance(font_resources, DictionaryObject): + font_resources = DictionaryObject(font_resources) + font_resources[NameObject("/F1")] = font_dict + resources[NameObject("/Font")] = font_resources + page[NameObject("/Resources")] = resources + + escaped = text.replace("\\", "\\\\").replace("(", "\\(").replace(")", "\\)") + stream_data = f"BT /F1 12 Tf 100 700 Td ({escaped}) Tj ET" + stream = DecodedStreamObject() + stream.set_data(stream_data.encode("latin-1")) + page[NameObject("/Contents")] = stream + + buf = io.BytesIO() + writer.write(buf) + return buf.getvalue() + + +def _make_request_mock(verify_ssl=True, trust_env=False): + """Return a mock that looks like a FastAPI Request.""" + request = Mock() + request.app.state.config.ENABLE_WEB_LOADER_SSL_VERIFICATION = verify_ssl + request.app.state.config.WEB_SEARCH_TRUST_ENV = trust_env + return request + + +# =========================================================================== +# Tests +# =========================================================================== + + +class TestExtractTextFromPdfBytes: + + def test_valid_pdf(self): + result = extract_text_from_pdf_bytes(_make_minimal_pdf("Hello, World!")) + assert "Hello" in result + assert "World" in result + + def test_blank_pdf_returns_placeholder(self): + from pypdf import PdfWriter + + writer = PdfWriter() + writer.add_blank_page(612, 792) + buf = io.BytesIO() + writer.write(buf) + + result = extract_text_from_pdf_bytes(buf.getvalue()) + assert "no extractable text" in result + + def test_corrupted_bytes_raises(self): + with pytest.raises(Exception): + extract_text_from_pdf_bytes(b"not a pdf") + + def test_truncated_header_raises(self): + with pytest.raises(Exception): + extract_text_from_pdf_bytes(b"%PDF-1.4 broken") + + +class TestExtractPdfFromUrl: + + def _make_mock_resp(self, data: bytes, headers=None): + """Create a mock response with iter_content support.""" + mock_resp = Mock() + mock_resp.headers = headers or {} + mock_resp.raise_for_status = Mock() + mock_resp.close = Mock() + mock_resp.iter_content = Mock(return_value=iter([data])) + return mock_resp + + def _mock_session(self, mock_resp): + """Create a mock requests.Session that works as a context manager.""" + session_instance = Mock() + session_instance.get.return_value = mock_resp + session_instance.__enter__ = Mock(return_value=session_instance) + session_instance.__exit__ = Mock(return_value=False) + return session_instance + + def test_success(self): + pdf_bytes = _make_minimal_pdf("download test") + mock_resp = self._make_mock_resp(pdf_bytes) + + with patch("requests.Session") as cls: + cls.return_value = self._mock_session(mock_resp) + result = extract_pdf_from_url( + _make_request_mock(), "https://example.com/doc.pdf" + ) + + assert "download test" in result + + def test_oversized_raises(self): + mock_resp = self._make_mock_resp(b"x" * (PDF_MAX_SIZE + 1)) + + with patch("requests.Session") as cls: + cls.return_value = self._mock_session(mock_resp) + with pytest.raises(ValueError, match="size limit"): + extract_pdf_from_url( + _make_request_mock(), "https://example.com/huge.pdf" + ) + + def test_content_length_precheck(self): + """Content-Length header triggers early rejection before full download.""" + mock_resp = self._make_mock_resp( + b"small", + headers={"Content-Length": str(PDF_MAX_SIZE + 1)}, + ) + + with patch("requests.Session") as cls: + cls.return_value = self._mock_session(mock_resp) + with pytest.raises(ValueError, match="size limit"): + extract_pdf_from_url( + _make_request_mock(), "https://example.com/big.pdf" + ) + # Verify response was closed without streaming the body + mock_resp.close.assert_called_once() + mock_resp.iter_content.assert_not_called() + + +class TestGetContentFromUrl: + """Test the URL-routing logic in get_content_from_url.""" + + def test_pdf_extension_fast_path(self): + orig = _ns["extract_pdf_from_url"] + _ns["extract_pdf_from_url"] = Mock(return_value="fast path text") + try: + content, docs = get_content_from_url( + _make_request_mock(), "https://example.com/report.pdf" + ) + assert content == "fast path text" + assert len(docs) == 1 + _ns["extract_pdf_from_url"].assert_called_once() + finally: + _ns["extract_pdf_from_url"] = orig + + def test_pdf_extension_case_insensitive(self): + _ns["extract_pdf_from_url"] = Mock(return_value="text") + try: + get_content_from_url(_make_request_mock(), "https://example.com/R.PDF") + _ns["extract_pdf_from_url"].assert_called_once() + finally: + _ns["extract_pdf_from_url"] = _extract_functions()["extract_pdf_from_url"] + + def test_pdf_binary_fallback(self): + from langchain_core.documents import Document + + mock_loader = Mock() + mock_loader.load.return_value = [ + Document(page_content="%PDF-1.4 garbage", metadata={}) + ] + + _ns["get_loader"] = Mock(return_value=mock_loader) + _ns["extract_pdf_from_url"] = Mock(return_value="extracted text") + try: + content, docs = get_content_from_url( + _make_request_mock(), "https://example.com/download?id=1" + ) + assert content == "extracted text" + _ns["extract_pdf_from_url"].assert_called_once() + finally: + ns2 = _extract_functions() + _ns["get_loader"] = ns2["get_loader"] + _ns["extract_pdf_from_url"] = ns2["extract_pdf_from_url"] + + def test_pdf_extension_error_fallback(self): + """When .pdf fast-path extraction fails, an error message is returned.""" + _ns["extract_pdf_from_url"] = Mock( + side_effect=Exception("connection refused") + ) + try: + content, docs = get_content_from_url( + _make_request_mock(), "https://example.com/broken.pdf" + ) + assert "[Error:" in content + assert "broken.pdf" in content + assert len(docs) == 1 + finally: + _ns["extract_pdf_from_url"] = _extract_functions()["extract_pdf_from_url"] + + def test_html_content_unchanged(self): + from langchain_core.documents import Document + + mock_loader = Mock() + mock_loader.load.return_value = [ + Document(page_content="Hello from the web", metadata={}) + ] + _ns["get_loader"] = Mock(return_value=mock_loader) + try: + content, docs = get_content_from_url( + _make_request_mock(), "https://example.com/page" + ) + assert content == "Hello from the web" + finally: + _ns["get_loader"] = _extract_functions()["get_loader"]