scriptoriumadmin/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui.git synced 2026-10-25 05:54:29 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 3

feat: rework DoclingLoader as async submit/poll/retrieve with timeout

Browse source 
The existing docling extraction engine called Docling's synchronous
/v1/convert/file endpoint and blocked until conversion finished, with
no timeout and no visibility into queue position while waiting.

Rework DoclingLoader to use the async Docling API instead:
- _submit_file(): POST /v1/convert/file/async, returns a task_id
- _poll_task_until_done(): long-polls /v1/status/poll/{task_id} until
  success/failure, respecting an optional overall timeout
- _retrieve_result(): GET /v1/result/{task_id} once the task succeeds
- load_from_task_id(): resume polling+retrieval for an
  already-submitted task_id, e.g. after a server restart while
  docling-serve is still processing it
- optional status_callback(dict) hook, invoked with task_id and queue
  position as they become available, so a caller can persist/display
  progress (inert unless a caller supplies one)

New DOCLING_SERVE_TIMEOUT setting (env var + rag config API) controls
the overall wait; unset means wait indefinitely, matching the previous
behavior.

Deliberately excludes: the DoclingLoaderJson structured-JSON extraction
engine and its chunk-mode config, and the unrelated RAG pipeline
improvements (citation extraction, JSON context format, etc.) — those
remain on a separate branch to be split out on their own.
This commit is contained in:
jannefleischer 2026-07-09 16:13:08 +02:00
parent 10724d057a
commit 2571fe1b11
4 changed files with 142 additions and 31 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

5
backend/open_webui/config.py
View file

@ -953,6 +953,10 @@ ENABLE_RAG_HYBRID_SEARCH_ENRICHED_TEXTS = (
os.getenv('ENABLE_RAG_HYBRID_SEARCH_ENRICHED_TEXTS', 'False').lower() == 'true'
)
DOCLING_SERVE_TIMEOUT = (
int(os.environ.get("DOCLING_SERVE_TIMEOUT")) if os.environ.get("DOCLING_SERVE_TIMEOUT") else None
)
RAG_FULL_CONTEXT = os.getenv('RAG_FULL_CONTEXT', 'False').lower() == 'true'
RAG_FILE_MAX_COUNT = int(os.getenv('RAG_FILE_MAX_COUNT')) if os.getenv('RAG_FILE_MAX_COUNT') else None
@ -2807,6 +2811,7 @@ DEFAULT_CONFIG = {
'rag.docling_server_url': DOCLING_SERVER_URL,
'rag.docling_api_key': DOCLING_API_KEY,
'rag.docling_params': DOCLING_PARAMS,
'rag.docling_serve_timeout': DOCLING_SERVE_TIMEOUT,
'rag.document_intelligence_endpoint': DOCUMENT_INTELLIGENCE_ENDPOINT,
'rag.document_intelligence_key': DOCUMENT_INTELLIGENCE_KEY,
'rag.document_intelligence_model': DOCUMENT_INTELLIGENCE_MODEL,

158
backend/open_webui/retrieval/loaders/main.py
View file

@ -2,6 +2,7 @@ import asyncio
import json
import logging
import sys
import time
import ftfy
import requests
@ -179,23 +180,27 @@ class TikaLoader:
class DoclingLoader:
def __init__(self, url, api_key=None, file_path=None, mime_type=None, params=None):
def __init__(self, url, api_key=None, file_path=None, mime_type=None, params=None, timeout=None, status_callback=None):
self.url = url.rstrip('/')
self.api_key = api_key
self.file_path = file_path
self.mime_type = mime_type
self.params = params or {}
self.timeout = timeout # total seconds to wait; None = infinite
self.status_callback = status_callback # optional callable(dict) to persist queue state
def load(self) -> list[Document]:
def _build_headers(self) -> dict:
headers = {}
if self.api_key:
headers["X-Api-Key"] = f"{self.api_key}"
return headers
def _submit_file(self) -> tuple[str, int | None]:
headers = self._build_headers()
page_break_marker = '\f'
with open(self.file_path, 'rb') as f:
headers = {}
if self.api_key:
headers['X-Api-Key'] = f'{self.api_key}'
with open(self.file_path, "rb") as f:
r = requests.post(
f'{self.url}/v1/convert/file',
f"{self.url}/v1/convert/file/async",
files={
'files': (
self.file_path,
@ -210,28 +215,11 @@ class DoclingLoader:
},
headers=headers,
verify=AIOHTTP_CLIENT_SESSION_SSL,
timeout=30,
)
if r.ok:
result = r.json()
document_data = result.get('document', {})
md_content = document_data.get('md_content', '')
text = md_content or '<No text content found>'
metadata = {'Content-Type': self.mime_type} if self.mime_type else {}
if page_break_marker in md_content:
documents = [
Document(page_content=page.strip(), metadata={**metadata, 'page': page_idx})
for page_idx, page in enumerate(md_content.split(page_break_marker))
if page.strip()
]
if documents:
log.debug('Docling extracted text: %s', text)
return documents
log.debug('Docling extracted text: %s', text)
return [Document(page_content=text, metadata=metadata)]
else:
error_msg = f'Error calling Docling API: {r.reason}'
if not r.ok:
error_msg = f"Error calling Docling API: {r.reason}"
if r.text:
try:
error_data = r.json()
@ -241,6 +229,105 @@ class DoclingLoader:
error_msg += f' - {r.text}'
raise Exception(f'Error calling Docling: {error_msg}')
submit_data = r.json()
task_id = submit_data.get("task_id")
task_position = submit_data.get("task_position")
if not task_id:
raise Exception("Docling async submit did not return a task_id")
log.info(
"Docling task submitted: %s, queue position: %s",
task_id,
task_position,
)
if self.status_callback:
self.status_callback({"task_id": task_id, "task_position": task_position})
return task_id, task_position
def _poll_task_until_done(self, task_id: str) -> dict:
headers = self._build_headers()
deadline = time.monotonic() + self.timeout if self.timeout is not None else None
poll_wait = 30 # long-poll window per request (seconds)
while True:
if deadline is not None:
remaining = deadline - time.monotonic()
if remaining <= 0:
raise Exception(
f"Docling conversion timed out after {self.timeout}s (task_id={task_id})"
)
poll_wait = min(poll_wait, int(remaining) + 1)
poll_start = time.monotonic()
try:
status_r = requests.get(
f"{self.url}/v1/status/poll/{task_id}",
params={"wait": poll_wait},
headers=headers,
timeout=poll_wait + 10,
)
except requests.Timeout:
log.warning("Docling status poll timed out for task %s, retrying", task_id)
elapsed = time.monotonic() - poll_start
if elapsed < poll_wait:
time.sleep(poll_wait - elapsed)
continue
if not status_r.ok:
raise Exception(f"Error polling Docling task status: {status_r.reason}")
status_data = status_r.json()
task_status = status_data.get("task_status", "")
log.debug(
"Docling task %s: status=%s, queue_position=%s",
task_id,
task_status,
status_data.get("task_position"),
)
if self.status_callback and status_data.get("task_position") is not None:
self.status_callback({"task_position": status_data["task_position"]})
if task_status == "success":
return status_data
elif task_status == "failure":
error_msg = status_data.get("error_message") or "Unknown error"
raise Exception(f"Docling conversion failed: {error_msg}")
# else "pending" or "started" – keep polling
elapsed = time.monotonic() - poll_start
if elapsed < poll_wait:
time.sleep(poll_wait - elapsed)
def _retrieve_result(self, task_id: str) -> dict:
headers = self._build_headers()
result_r = requests.get(f"{self.url}/v1/result/{task_id}", headers=headers, timeout=30)
if not result_r.ok:
raise Exception(f"Error retrieving Docling result: {result_r.reason}")
return result_r.json()
def format_result(self, result_json: dict) -> list[Document]:
document_data = result_json.get("document", {})
text = document_data.get("md_content", "<No text content found>")
metadata = {"Content-Type": self.mime_type} if self.mime_type else {}
log.debug("Docling extracted text: %s", text)
return [Document(page_content=text, metadata=metadata)]
def load(self) -> list[Document]:
task_id, _ = self._submit_file()
self._poll_task_until_done(task_id)
result_json = self._retrieve_result(task_id)
return self.format_result(result_json)
def load_from_task_id(self, task_id: str) -> list[Document]:
"""Resume processing from an already-submitted docling task_id.
Used on server restart when docling-serve is still running the task:
skips re-submission and goes straight to polling → retrieve → format.
"""
self._poll_task_until_done(task_id)
result_json = self._retrieve_result(task_id)
return self.format_result(result_json)
class Loader:
def __init__(self, engine: str = '', **kwargs):
@ -494,12 +581,21 @@ class Loader:
log.error('Invalid DOCLING_PARAMS format, expected JSON object')
params = {}
docling_timeout = self.kwargs.get("DOCLING_SERVE_TIMEOUT")
if docling_timeout is not None:
try:
docling_timeout = int(docling_timeout)
except (ValueError, TypeError):
docling_timeout = None
loader = DoclingLoader(
url=self.kwargs.get('DOCLING_SERVER_URL'),
api_key=self.kwargs.get('DOCLING_API_KEY', None),
url=self.kwargs.get("DOCLING_SERVER_URL"),
api_key=self.kwargs.get("DOCLING_API_KEY", None),
file_path=file_path,
mime_type=file_content_type,
params=params,
timeout=docling_timeout,
status_callback=self.kwargs.get("DOCLING_STATUS_CALLBACK"),
)
elif (
self.engine == 'document_intelligence'

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

@ -90,6 +90,7 @@ LOADER_CONFIG_KEYS = {
'DOCLING_SERVER_URL': 'rag.docling_server_url',
'DOCLING_API_KEY': 'rag.docling_api_key',
'DOCLING_PARAMS': 'rag.docling_params',
'DOCLING_SERVE_TIMEOUT': 'rag.docling_serve_timeout',
'PDF_EXTRACT_IMAGES': 'rag.pdf_extract_images',
'PDF_LOADER_MODE': 'rag.pdf_loader_mode',
'DOCUMENT_INTELLIGENCE_ENDPOINT': 'rag.document_intelligence_endpoint',

9
backend/open_webui/routers/retrieval.py
View file

@ -278,6 +278,7 @@ RETRIEVAL_CONFIG_KEYS = {
'DOCLING_API_KEY': 'rag.docling_api_key',
'DOCLING_PARAMS': 'rag.docling_params',
'DOCLING_SERVER_URL': 'rag.docling_server_url',
'DOCLING_SERVE_TIMEOUT': 'rag.docling_serve_timeout',
'DOCUMENT_INTELLIGENCE_ENDPOINT': 'rag.document_intelligence_endpoint',
'DOCUMENT_INTELLIGENCE_KEY': 'rag.document_intelligence_key',
'DOCUMENT_INTELLIGENCE_MODEL': 'rag.document_intelligence_model',
@ -646,6 +647,7 @@ async def get_rag_config(request: Request, user=Depends(get_admin_user)):
'DOCLING_SERVER_URL': config.DOCLING_SERVER_URL,
'DOCLING_API_KEY': config.DOCLING_API_KEY,
'DOCLING_PARAMS': config.DOCLING_PARAMS,
'DOCLING_SERVE_TIMEOUT': config.DOCLING_SERVE_TIMEOUT,
'DOCUMENT_INTELLIGENCE_ENDPOINT': config.DOCUMENT_INTELLIGENCE_ENDPOINT,
'DOCUMENT_INTELLIGENCE_KEY': config.DOCUMENT_INTELLIGENCE_KEY,
'DOCUMENT_INTELLIGENCE_MODEL': config.DOCUMENT_INTELLIGENCE_MODEL,
@ -878,6 +880,7 @@ class ConfigForm(BaseModel):
DOCLING_SERVER_URL: str | None = None
DOCLING_API_KEY: str | None = None
DOCLING_PARAMS: dict | None = None
DOCLING_SERVE_TIMEOUT: int | None = None
DOCUMENT_INTELLIGENCE_ENDPOINT: str | None = None
DOCUMENT_INTELLIGENCE_KEY: str | None = None
DOCUMENT_INTELLIGENCE_MODEL: str | None = None
@ -1053,6 +1056,11 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend
form_data.DOCLING_API_KEY if form_data.DOCLING_API_KEY is not None else config.DOCLING_API_KEY
)
config.DOCLING_PARAMS = form_data.DOCLING_PARAMS if form_data.DOCLING_PARAMS is not None else config.DOCLING_PARAMS
config.DOCLING_SERVE_TIMEOUT = (
form_data.DOCLING_SERVE_TIMEOUT
if form_data.DOCLING_SERVE_TIMEOUT is not None
else config.DOCLING_SERVE_TIMEOUT
)
config.DOCUMENT_INTELLIGENCE_ENDPOINT = (
form_data.DOCUMENT_INTELLIGENCE_ENDPOINT
if form_data.DOCUMENT_INTELLIGENCE_ENDPOINT is not None
@ -1345,6 +1353,7 @@ async def update_rag_config(request: Request, form_data: ConfigForm, user=Depend
'DOCLING_SERVER_URL': config.DOCLING_SERVER_URL,
'DOCLING_API_KEY': config.DOCLING_API_KEY,
'DOCLING_PARAMS': config.DOCLING_PARAMS,
'DOCLING_SERVE_TIMEOUT': config.DOCLING_SERVE_TIMEOUT,
'DOCUMENT_INTELLIGENCE_ENDPOINT': config.DOCUMENT_INTELLIGENCE_ENDPOINT,
'DOCUMENT_INTELLIGENCE_KEY': config.DOCUMENT_INTELLIGENCE_KEY,
'DOCUMENT_INTELLIGENCE_MODEL': config.DOCUMENT_INTELLIGENCE_MODEL,