Emil_Shanaty/open-webui-ai4me
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2026-04-21 15:47:32 +09:00
parent c7b6de6ca4
commit b9fc3f367a
2 changed files with 117 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files
backend/open_webui/retrieval/utils.py
+113 -4
View File
@@ -83,11 +83,120 @@ def get_loader(request, url: str):
)
def build_loader_from_config(request):
"""Build a Loader instance with the admin's configured extraction engine settings."""
from open_webui.retrieval.loaders.main import Loader
config = request.app.state.config
return Loader(
engine=config.CONTENT_EXTRACTION_ENGINE,
DATALAB_MARKER_API_KEY=config.DATALAB_MARKER_API_KEY,
DATALAB_MARKER_API_BASE_URL=config.DATALAB_MARKER_API_BASE_URL,
DATALAB_MARKER_ADDITIONAL_CONFIG=config.DATALAB_MARKER_ADDITIONAL_CONFIG,
DATALAB_MARKER_SKIP_CACHE=config.DATALAB_MARKER_SKIP_CACHE,
DATALAB_MARKER_FORCE_OCR=config.DATALAB_MARKER_FORCE_OCR,
DATALAB_MARKER_PAGINATE=config.DATALAB_MARKER_PAGINATE,
DATALAB_MARKER_STRIP_EXISTING_OCR=config.DATALAB_MARKER_STRIP_EXISTING_OCR,
DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION=config.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION,
DATALAB_MARKER_FORMAT_LINES=config.DATALAB_MARKER_FORMAT_LINES,
DATALAB_MARKER_USE_LLM=config.DATALAB_MARKER_USE_LLM,
DATALAB_MARKER_OUTPUT_FORMAT=config.DATALAB_MARKER_OUTPUT_FORMAT,
EXTERNAL_DOCUMENT_LOADER_URL=config.EXTERNAL_DOCUMENT_LOADER_URL,
EXTERNAL_DOCUMENT_LOADER_API_KEY=config.EXTERNAL_DOCUMENT_LOADER_API_KEY,
TIKA_SERVER_URL=config.TIKA_SERVER_URL,
DOCLING_SERVER_URL=config.DOCLING_SERVER_URL,
DOCLING_API_KEY=config.DOCLING_API_KEY,
DOCLING_PARAMS=config.DOCLING_PARAMS,
PDF_EXTRACT_IMAGES=config.PDF_EXTRACT_IMAGES,
PDF_LOADER_MODE=config.PDF_LOADER_MODE,
DOCUMENT_INTELLIGENCE_ENDPOINT=config.DOCUMENT_INTELLIGENCE_ENDPOINT,
DOCUMENT_INTELLIGENCE_KEY=config.DOCUMENT_INTELLIGENCE_KEY,
DOCUMENT_INTELLIGENCE_MODEL=config.DOCUMENT_INTELLIGENCE_MODEL,
MISTRAL_OCR_API_BASE_URL=config.MISTRAL_OCR_API_BASE_URL,
MISTRAL_OCR_API_KEY=config.MISTRAL_OCR_API_KEY,
MINERU_API_MODE=config.MINERU_API_MODE,
MINERU_API_URL=config.MINERU_API_URL,
MINERU_API_KEY=config.MINERU_API_KEY,
MINERU_API_TIMEOUT=config.MINERU_API_TIMEOUT,
MINERU_PARAMS=config.MINERU_PARAMS,
)
def _extract_text_from_binary_response(
request, response: requests.Response, url: str
) -> tuple[str, list]:
"""Download response body to a temp file and extract text using the Loader pipeline."""
import mimetypes
import tempfile
import urllib.parse
content_type = response.headers.get('Content-Type', '').split(';')[0].strip()
# Derive filename from URL path, falling back to Content-Disposition or mime guess
url_path = urllib.parse.urlparse(url).path
filename = os.path.basename(url_path) if url_path else ''
if not filename or '.' not in filename:
# Try Content-Disposition header
cd = response.headers.get('Content-Disposition', '')
if 'filename=' in cd:
filename = cd.split('filename=')[-1].strip('"\'')
if not filename or '.' not in filename:
ext = mimetypes.guess_extension(content_type) or ''
filename = f'download{ext}'
suffix = '.' + filename.split('.')[-1].lower() if '.' in filename else ''
with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
tmp.write(response.content)
tmp_path = tmp.name
try:
loader = build_loader_from_config(request)
docs = loader.load(filename, content_type, tmp_path)
for doc in docs:
doc.metadata['source'] = url
content = ' '.join([doc.page_content for doc in docs])
return content, docs
finally:
os.remove(tmp_path)
def _is_text_content_type(content_type: str) -> bool:
"""Return True if the content type should be handled by the web loader."""
ct = content_type.split(';')[0].strip().lower()
if ct.startswith('text/'):
return True
if any(t in ct for t in ['xml', 'json', 'javascript']):
return True
return not ct # empty / missing → assume HTML
def get_content_from_url(request, url: str) -> str:
loader = get_loader(request, url)
docs = loader.load()
content = ' '.join([doc.page_content for doc in docs])
return content, docs
# Streamed GET to check Content-Type without downloading the body.
try:
response = requests.get(url, stream=True, timeout=30)
response.raise_for_status()
content_type = response.headers.get('Content-Type', '')
except Exception:
content_type = ''
response = None
# Text / HTML / unknown — use the configured web loader
if response is None or _is_text_content_type(content_type):
if response is not None:
response.close()
loader = get_loader(request, url)
docs = loader.load()
content = ' '.join([doc.page_content for doc in docs])
return content, docs
# Binary content (PDF, DOCX, XLSX, PPTX, etc.) — download and extract
try:
return _extract_text_from_binary_response(request, response, url)
finally:
response.close()
CHUNK_HASH_KEY = '_chunk_hash'
backend/open_webui/routers/retrieval.py
+4 -34
View File
@@ -48,7 +48,7 @@ from open_webui.retrieval.vector.factory import VECTOR_DB_CLIENT
from open_webui.retrieval.vector.async_client import ASYNC_VECTOR_DB_CLIENT
# Document loaders
from open_webui.retrieval.loaders.main import Loader
from open_webui.retrieval.loaders.youtube import YoutubeLoader
# Web search engines
@@ -82,6 +82,7 @@ from open_webui.retrieval.web.yandex import search_yandex
from open_webui.retrieval.web.ydc import search_youcom
from open_webui.retrieval.utils import (
build_loader_from_config,
filter_accessible_collections,
get_content_from_url,
get_embedding_function,
@@ -1623,39 +1624,8 @@ async def process_file(
file_path = file.path
if file_path:
file_path = await asyncio.to_thread(Storage.get_file, file_path)
loader = Loader(
engine=request.app.state.config.CONTENT_EXTRACTION_ENGINE,
user=user,
DATALAB_MARKER_API_KEY=request.app.state.config.DATALAB_MARKER_API_KEY,
DATALAB_MARKER_API_BASE_URL=request.app.state.config.DATALAB_MARKER_API_BASE_URL,
DATALAB_MARKER_ADDITIONAL_CONFIG=request.app.state.config.DATALAB_MARKER_ADDITIONAL_CONFIG,
DATALAB_MARKER_SKIP_CACHE=request.app.state.config.DATALAB_MARKER_SKIP_CACHE,
DATALAB_MARKER_FORCE_OCR=request.app.state.config.DATALAB_MARKER_FORCE_OCR,
DATALAB_MARKER_PAGINATE=request.app.state.config.DATALAB_MARKER_PAGINATE,
DATALAB_MARKER_STRIP_EXISTING_OCR=request.app.state.config.DATALAB_MARKER_STRIP_EXISTING_OCR,
DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION=request.app.state.config.DATALAB_MARKER_DISABLE_IMAGE_EXTRACTION,
DATALAB_MARKER_FORMAT_LINES=request.app.state.config.DATALAB_MARKER_FORMAT_LINES,
DATALAB_MARKER_USE_LLM=request.app.state.config.DATALAB_MARKER_USE_LLM,
DATALAB_MARKER_OUTPUT_FORMAT=request.app.state.config.DATALAB_MARKER_OUTPUT_FORMAT,
EXTERNAL_DOCUMENT_LOADER_URL=request.app.state.config.EXTERNAL_DOCUMENT_LOADER_URL,
EXTERNAL_DOCUMENT_LOADER_API_KEY=request.app.state.config.EXTERNAL_DOCUMENT_LOADER_API_KEY,
TIKA_SERVER_URL=request.app.state.config.TIKA_SERVER_URL,
DOCLING_SERVER_URL=request.app.state.config.DOCLING_SERVER_URL,
DOCLING_API_KEY=request.app.state.config.DOCLING_API_KEY,
DOCLING_PARAMS=request.app.state.config.DOCLING_PARAMS,
PDF_EXTRACT_IMAGES=request.app.state.config.PDF_EXTRACT_IMAGES,
PDF_LOADER_MODE=request.app.state.config.PDF_LOADER_MODE,
DOCUMENT_INTELLIGENCE_ENDPOINT=request.app.state.config.DOCUMENT_INTELLIGENCE_ENDPOINT,
DOCUMENT_INTELLIGENCE_KEY=request.app.state.config.DOCUMENT_INTELLIGENCE_KEY,
DOCUMENT_INTELLIGENCE_MODEL=request.app.state.config.DOCUMENT_INTELLIGENCE_MODEL,
MISTRAL_OCR_API_BASE_URL=request.app.state.config.MISTRAL_OCR_API_BASE_URL,
MISTRAL_OCR_API_KEY=request.app.state.config.MISTRAL_OCR_API_KEY,
MINERU_API_MODE=request.app.state.config.MINERU_API_MODE,
MINERU_API_URL=request.app.state.config.MINERU_API_URL,
MINERU_API_KEY=request.app.state.config.MINERU_API_KEY,
MINERU_API_TIMEOUT=request.app.state.config.MINERU_API_TIMEOUT,
MINERU_PARAMS=request.app.state.config.MINERU_PARAMS,
)
loader = build_loader_from_config(request)
loader.user = user
docs = await loader.aload(file.filename, file.meta.get('content_type'), file_path)
docs = [