scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-09-16 23:43:03 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

fix: extract text from PDF URLs in fetch_url tool

Browse source 
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 <noreply@anthropic.com>
This commit is contained in:
mia 2026-04-19 01:57:25 +05:30
parent a4d62253df
commit 9d3d80b9f8
3 changed files with 485 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

111
backend/open_webui/retrieval/utils.py
View file

@ -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

33
backend/open_webui/retrieval/web/utils.py
View file

@ -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'<html><body><pre>{escaped}</pre></body></html>'
except Exception as e:
log.warning(f'Failed to extract text from PDF at {url}: {e}.')
return f'<html><body><pre>[Error: Unable to extract text from PDF at {url}]</pre></body></html>'
return await response.text()
except aiohttp.ClientConnectionError as e:
if i == retries - 1:

343
backend/tests/retrieval/test_pdf_handling.py Normal file
View file

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