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"]