From 2571fe1b11a553274ec37fa2b0b13d789583acd2 Mon Sep 17 00:00:00 2001 From: jannefleischer Date: Thu, 9 Jul 2026 16:13:08 +0200 Subject: [PATCH] feat: rework DoclingLoader as async submit/poll/retrieve with timeout MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- backend/open_webui/config.py | 5 + backend/open_webui/retrieval/loaders/main.py | 158 +++++++++++++++---- backend/open_webui/retrieval/utils.py | 1 + backend/open_webui/routers/retrieval.py | 9 ++ 4 files changed, 142 insertions(+), 31 deletions(-) diff --git a/backend/open_webui/config.py b/backend/open_webui/config.py index f89738b0db..b36c2093df 100644 --- a/backend/open_webui/config.py +++ b/backend/open_webui/config.py @@ -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, diff --git a/backend/open_webui/retrieval/loaders/main.py b/backend/open_webui/retrieval/loaders/main.py index c5eeb32e91..d61d7a974f 100644 --- a/backend/open_webui/retrieval/loaders/main.py +++ b/backend/open_webui/retrieval/loaders/main.py @@ -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 '' - 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", "") + 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' diff --git a/backend/open_webui/retrieval/utils.py b/backend/open_webui/retrieval/utils.py index 2788d4dd0c..f6c47f1d36 100644 --- a/backend/open_webui/retrieval/utils.py +++ b/backend/open_webui/retrieval/utils.py @@ -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', diff --git a/backend/open_webui/routers/retrieval.py b/backend/open_webui/routers/retrieval.py index 3f62af4af4..ee6afb72fb 100644 --- a/backend/open_webui/routers/retrieval.py +++ b/backend/open_webui/routers/retrieval.py @@ -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,