LLM/docling-serve
1
0
Fork 0
mirror of https://github.com/docling-project/docling-serve.git synced 2025-11-29 00:23:36 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: add support for file upload and return as file in async endpoints (#152)

Browse Source 
Signed-off-by: Michele Dolfi <dol@zurich.ibm.com>
This commit is contained in:
Michele Dolfi
2025-04-28 11:18:19 +02:00
committed by GitHub
parent 829effec1a
commit c65f3c654c
21 changed files with 516 additions and 151 deletions
Download Patch File Download Diff File Expand all files Collapse all files

198
docling_serve/app.py
View File

@@ -1,11 +1,11 @@
import asyncio
import importlib.metadata
import logging
import tempfile
import shutil
import time
from contextlib import asynccontextmanager
from io import BytesIO
from pathlib import Path
from typing import Annotated, Any, Optional, Union
from typing import Annotated
from fastapi import (
BackgroundTasks,
@@ -35,6 +35,7 @@ from docling_serve.datamodel.callback import (
from docling_serve.datamodel.convert import ConvertDocumentsOptions
from docling_serve.datamodel.requests import (
ConvertDocumentFileSourcesRequest,
ConvertDocumentHttpSourcesRequest,
ConvertDocumentsRequest,
)
from docling_serve.datamodel.responses import (
@@ -44,11 +45,7 @@ from docling_serve.datamodel.responses import (
TaskStatusResponse,
WebsocketMessage,
)
from docling_serve.docling_conversion import (
convert_documents,
get_converter,
get_pdf_pipeline_opts,
)
from docling_serve.datamodel.task import Task, TaskSource
from docling_serve.engines.async_orchestrator import (
BaseAsyncOrchestrator,
ProgressInvalid,
@@ -56,8 +53,8 @@ from docling_serve.engines.async_orchestrator import (
from docling_serve.engines.async_orchestrator_factory import get_async_orchestrator
from docling_serve.engines.base_orchestrator import TaskNotFoundError
from docling_serve.helper_functions import FormDepends
from docling_serve.response_preparation import process_results
from docling_serve.settings import docling_serve_settings
from docling_serve.storage import get_scratch
# Set up custom logging as we'll be intermixes with FastAPI/Uvicorn's logging
@@ -95,11 +92,11 @@ _log = logging.getLogger(__name__)
# Context manager to initialize and clean up the lifespan of the FastAPI app
@asynccontextmanager
async def lifespan(app: FastAPI):
# Converter with default options
pdf_format_option = get_pdf_pipeline_opts(ConvertDocumentsOptions())
get_converter(pdf_format_option)
orchestrator = get_async_orchestrator()
scratch_dir = get_scratch()
# Warm up processing cache
await orchestrator.warm_up_caches()
# Start the background queue processor
queue_task = asyncio.create_task(orchestrator.process_queue())
@@ -113,6 +110,10 @@ async def lifespan(app: FastAPI):
except asyncio.CancelledError:
_log.info("Queue processor cancelled.")
# Remove scratch directory in case it was a tempfile
if docling_serve_settings.scratch_path is not None:
shutil.rmtree(scratch_dir, ignore_errors=True)
##################################
# App creation and configuration #
@@ -162,7 +163,8 @@ def create_app(): # noqa: C901
from docling_serve.gradio_ui import ui as gradio_ui
tmp_output_dir = Path(tempfile.mkdtemp())
tmp_output_dir = get_scratch() / "gradio"
tmp_output_dir.mkdir(exist_ok=True, parents=True)
gradio_ui.gradio_output_dir = tmp_output_dir
app = gr.mount_gradio_app(
app,
@@ -210,6 +212,56 @@ def create_app(): # noqa: C901
redoc_js_url="/static/redoc.standalone.js",
)
########################
# Async / Sync helpers #
########################
async def _enque_source(
orchestrator: BaseAsyncOrchestrator, conversion_request: ConvertDocumentsRequest
) -> Task:
sources: list[TaskSource] = []
if isinstance(conversion_request, ConvertDocumentFileSourcesRequest):
sources.extend(conversion_request.file_sources)
if isinstance(conversion_request, ConvertDocumentHttpSourcesRequest):
sources.extend(conversion_request.http_sources)
task = await orchestrator.enqueue(
sources=sources, options=conversion_request.options
)
return task
async def _enque_file(
orchestrator: BaseAsyncOrchestrator,
files: list[UploadFile],
options: ConvertDocumentsOptions,
) -> Task:
_log.info(f"Received {len(files)} files for processing.")
# Load the uploaded files to Docling DocumentStream
file_sources: list[TaskSource] = []
for i, file in enumerate(files):
buf = BytesIO(file.file.read())
suffix = "" if len(file_sources) == 1 else f"_{i}"
name = file.filename if file.filename else f"file{suffix}.pdf"
file_sources.append(DocumentStream(name=name, stream=buf))
task = await orchestrator.enqueue(sources=file_sources, options=options)
return task
async def _wait_task_complete(
orchestrator: BaseAsyncOrchestrator, task_id: str
) -> bool:
MAX_WAIT = 120
start_time = time.monotonic()
while True:
task = await orchestrator.task_status(task_id=task_id)
if task.is_completed():
return True
await asyncio.sleep(5)
elapsed_time = time.monotonic() - start_time
if elapsed_time > MAX_WAIT:
return False
#############################
# API Endpoints definitions #
#############################
@@ -243,33 +295,33 @@ def create_app(): # noqa: C901
}
},
)
def process_url(
background_tasks: BackgroundTasks, conversion_request: ConvertDocumentsRequest
async def process_url(
background_tasks: BackgroundTasks,
orchestrator: Annotated[BaseAsyncOrchestrator, Depends(get_async_orchestrator)],
conversion_request: ConvertDocumentsRequest,
):
sources: list[Union[str, DocumentStream]] = []
headers: Optional[dict[str, Any]] = None
if isinstance(conversion_request, ConvertDocumentFileSourcesRequest):
for file_source in conversion_request.file_sources:
sources.append(file_source.to_document_stream())
else:
for http_source in conversion_request.http_sources:
sources.append(http_source.url)
if headers is None and http_source.headers:
headers = http_source.headers
# Note: results are only an iterator->lazy evaluation
results = convert_documents(
sources=sources, options=conversion_request.options, headers=headers
task = await _enque_source(
orchestrator=orchestrator, conversion_request=conversion_request
)
success = await _wait_task_complete(
orchestrator=orchestrator, task_id=task.task_id
)
# The real processing will happen here
response = process_results(
background_tasks=background_tasks,
conversion_options=conversion_request.options,
conv_results=results,
)
if not success:
# TODO: abort task!
return HTTPException(
status_code=504, detail="Conversion is taking too long."
)
return response
result = await orchestrator.task_result(
task_id=task.task_id, background_tasks=background_tasks
)
if result is None:
raise HTTPException(
status_code=404,
detail="Task result not found. Please wait for a completion status.",
)
return result
# Convert a document from file(s)
@app.post(
@@ -283,29 +335,34 @@ def create_app(): # noqa: C901
)
async def process_file(
background_tasks: BackgroundTasks,
orchestrator: Annotated[BaseAsyncOrchestrator, Depends(get_async_orchestrator)],
files: list[UploadFile],
options: Annotated[
ConvertDocumentsOptions, FormDepends(ConvertDocumentsOptions)
],
):
_log.info(f"Received {len(files)} files for processing.")
# Load the uploaded files to Docling DocumentStream
file_sources = []
for file in files:
buf = BytesIO(file.file.read())
name = file.filename if file.filename else "file.pdf"
file_sources.append(DocumentStream(name=name, stream=buf))
results = convert_documents(sources=file_sources, options=options)
response = process_results(
background_tasks=background_tasks,
conversion_options=options,
conv_results=results,
task = await _enque_file(
orchestrator=orchestrator, files=files, options=options
)
success = await _wait_task_complete(
orchestrator=orchestrator, task_id=task.task_id
)
return response
if not success:
# TODO: abort task!
return HTTPException(
status_code=504, detail="Conversion is taking too long."
)
result = await orchestrator.task_result(
task_id=task.task_id, background_tasks=background_tasks
)
if result is None:
raise HTTPException(
status_code=404,
detail="Task result not found. Please wait for a completion status.",
)
return result
# Convert a document from URL(s) using the async api
@app.post(
@@ -316,7 +373,35 @@ def create_app(): # noqa: C901
orchestrator: Annotated[BaseAsyncOrchestrator, Depends(get_async_orchestrator)],
conversion_request: ConvertDocumentsRequest,
):
task = await orchestrator.enqueue(request=conversion_request)
task = await _enque_source(
orchestrator=orchestrator, conversion_request=conversion_request
)
task_queue_position = await orchestrator.get_queue_position(
task_id=task.task_id
)
return TaskStatusResponse(
task_id=task.task_id,
task_status=task.task_status,
task_position=task_queue_position,
task_meta=task.processing_meta,
)
# Convert a document from file(s) using the async api
@app.post(
"/v1alpha/convert/file/async",
response_model=TaskStatusResponse,
)
async def process_file_async(
orchestrator: Annotated[BaseAsyncOrchestrator, Depends(get_async_orchestrator)],
background_tasks: BackgroundTasks,
files: list[UploadFile],
options: Annotated[
ConvertDocumentsOptions, FormDepends(ConvertDocumentsOptions)
],
):
task = await _enque_file(
orchestrator=orchestrator, files=files, options=options
)
task_queue_position = await orchestrator.get_queue_position(
task_id=task.task_id
)
@@ -426,9 +511,12 @@ def create_app(): # noqa: C901
)
async def task_result(
orchestrator: Annotated[BaseAsyncOrchestrator, Depends(get_async_orchestrator)],
background_tasks: BackgroundTasks,
task_id: str,
):
result = await orchestrator.task_result(task_id=task_id)
result = await orchestrator.task_result(
task_id=task_id, background_tasks=background_tasks
)
if result is None:
raise HTTPException(
status_code=404,

4
docling_serve/datamodel/requests.py
View File

@@ -2,7 +2,7 @@ import base64
from io import BytesIO
from typing import Annotated, Any, Union
from pydantic import BaseModel, Field
from pydantic import AnyHttpUrl, BaseModel, Field
from docling.datamodel.base_models import DocumentStream
@@ -15,7 +15,7 @@ class DocumentsConvertBase(BaseModel):
class HttpSource(BaseModel):
url: Annotated[
str,
AnyHttpUrl,
Field(
description="HTTP url to process",
examples=["https://arxiv.org/pdf/2206.01062"],

21
docling_serve/datamodel/task.py
View File

@@ -1,18 +1,29 @@
from typing import Optional
from pathlib import Path
from typing import Optional, Union
from pydantic import BaseModel
from fastapi.responses import FileResponse
from pydantic import BaseModel, ConfigDict
from docling.datamodel.base_models import DocumentStream
from docling_serve.datamodel.convert import ConvertDocumentsOptions
from docling_serve.datamodel.engines import TaskStatus
from docling_serve.datamodel.requests import ConvertDocumentsRequest
from docling_serve.datamodel.requests import FileSource, HttpSource
from docling_serve.datamodel.responses import ConvertDocumentResponse
from docling_serve.datamodel.task_meta import TaskProcessingMeta
TaskSource = Union[HttpSource, FileSource, DocumentStream]
class Task(BaseModel):
model_config = ConfigDict(arbitrary_types_allowed=True)
task_id: str
task_status: TaskStatus = TaskStatus.PENDING
request: Optional[ConvertDocumentsRequest]
result: Optional[ConvertDocumentResponse] = None
sources: list[TaskSource] = []
options: Optional[ConvertDocumentsOptions]
result: Optional[Union[ConvertDocumentResponse, FileResponse]] = None
scratch_dir: Optional[Path] = None
processing_meta: Optional[TaskProcessingMeta] = None
def is_completed(self) -> bool:

19
docling_serve/engines/async_kfp/orchestrator.py
View File

@@ -14,10 +14,11 @@ from docling_serve.datamodel.callback import (
ProgressSetNumDocs,
ProgressUpdateProcessed,
)
from docling_serve.datamodel.convert import ConvertDocumentsOptions
from docling_serve.datamodel.engines import TaskStatus
from docling_serve.datamodel.kfp import CallbackSpec
from docling_serve.datamodel.requests import ConvertDocumentsRequest
from docling_serve.datamodel.task import Task
from docling_serve.datamodel.requests import HttpSource
from docling_serve.datamodel.task import Task, TaskSource
from docling_serve.datamodel.task_meta import TaskProcessingMeta
from docling_serve.engines.async_kfp.kfp_pipeline import process
from docling_serve.engines.async_orchestrator import (
@@ -71,7 +72,9 @@ class AsyncKfpOrchestrator(BaseAsyncOrchestrator):
# verify_ssl=False,
)
async def enqueue(self, request: ConvertDocumentsRequest) -> Task:
async def enqueue(
self, sources: list[TaskSource], options: ConvertDocumentsOptions
) -> Task:
callbacks = []
if docling_serve_settings.eng_kfp_self_callback_endpoint is not None:
headers = {}
@@ -92,6 +95,8 @@ class AsyncKfpOrchestrator(BaseAsyncOrchestrator):
)
CallbacksType = TypeAdapter(list[CallbackSpec])
SourcesListType = TypeAdapter(list[HttpSource])
http_sources = [s for s in sources if isinstance(s, HttpSource)]
# hack: since the current kfp backend is not resolving the job_id placeholder,
# we set the run_name and pass it as argument to the job itself.
run_name = f"docling-job-{uuid.uuid4()}"
@@ -99,7 +104,8 @@ class AsyncKfpOrchestrator(BaseAsyncOrchestrator):
process,
arguments={
"batch_size": 10,
"request": request.model_dump(mode="json"),
"sources": SourcesListType.dump_python(http_sources, mode="json"),
"options": options.model_dump(mode="json"),
"callbacks": CallbacksType.dump_python(callbacks, mode="json"),
"run_name": run_name,
},
@@ -107,7 +113,7 @@ class AsyncKfpOrchestrator(BaseAsyncOrchestrator):
)
task_id = kfp_run.run_id
task = Task(task_id=task_id, request=request)
task = Task(task_id=task_id, sources=sources, options=options)
await self.init_task_tracking(task)
return task
@@ -185,6 +191,9 @@ class AsyncKfpOrchestrator(BaseAsyncOrchestrator):
async def process_queue(self):
return
async def warm_up_caches(self):
return
async def _get_run_id(self, run_name: str) -> str:
res = self._client.list_runs(
filter=json.dumps(

16
docling_serve/engines/async_local/orchestrator.py
View File

@@ -3,8 +3,9 @@ import logging
import uuid
from typing import Optional
from docling_serve.datamodel.requests import ConvertDocumentsRequest
from docling_serve.datamodel.task import Task
from docling_serve.datamodel.convert import ConvertDocumentsOptions
from docling_serve.datamodel.task import Task, TaskSource
from docling_serve.docling_conversion import get_converter, get_pdf_pipeline_opts
from docling_serve.engines.async_local.worker import AsyncLocalWorker
from docling_serve.engines.async_orchestrator import BaseAsyncOrchestrator
from docling_serve.settings import docling_serve_settings
@@ -18,9 +19,11 @@ class AsyncLocalOrchestrator(BaseAsyncOrchestrator):
self.task_queue = asyncio.Queue()
self.queue_list: list[str] = []
async def enqueue(self, request: ConvertDocumentsRequest) -> Task:
async def enqueue(
self, sources: list[TaskSource], options: ConvertDocumentsOptions
) -> Task:
task_id = str(uuid.uuid4())
task = Task(task_id=task_id, request=request)
task = Task(task_id=task_id, sources=sources, options=options)
await self.init_task_tracking(task)
self.queue_list.append(task_id)
@@ -47,3 +50,8 @@ class AsyncLocalOrchestrator(BaseAsyncOrchestrator):
# Wait for all workers to complete (they won't, as they run indefinitely)
await asyncio.gather(*workers)
_log.debug("All workers completed.")
async def warm_up_caches(self):
# Converter with default options
pdf_format_option = get_pdf_pipeline_opts(ConvertDocumentsOptions())
get_converter(pdf_format_option)

60
docling_serve/engines/async_local/worker.py
View File

@@ -1,17 +1,18 @@
import asyncio
import logging
import shutil
import time
from typing import TYPE_CHECKING, Any, Optional, Union
from fastapi import BackgroundTasks
from fastapi.responses import FileResponse
from docling.datamodel.base_models import DocumentStream
from docling_serve.datamodel.engines import TaskStatus
from docling_serve.datamodel.requests import ConvertDocumentFileSourcesRequest
from docling_serve.datamodel.responses import ConvertDocumentResponse
from docling_serve.datamodel.requests import FileSource, HttpSource
from docling_serve.docling_conversion import convert_documents
from docling_serve.response_preparation import process_results
from docling_serve.storage import get_scratch
if TYPE_CHECKING:
from docling_serve.engines.async_local.orchestrator import AsyncLocalOrchestrator
@@ -44,59 +45,66 @@ class AsyncLocalWorker:
# Notify clients about queue updates
await self.orchestrator.notify_queue_positions()
# Get the current event loop
asyncio.get_event_loop()
# Define a callback function to send progress updates to the client.
# TODO: send partial updates, e.g. when a document in the batch is done
def run_conversion():
sources: list[Union[str, DocumentStream]] = []
convert_sources: list[Union[str, DocumentStream]] = []
headers: Optional[dict[str, Any]] = None
if isinstance(task.request, ConvertDocumentFileSourcesRequest):
for file_source in task.request.file_sources:
sources.append(file_source.to_document_stream())
else:
for http_source in task.request.http_sources:
sources.append(http_source.url)
if headers is None and http_source.headers:
headers = http_source.headers
for source in task.sources:
if isinstance(source, DocumentStream):
convert_sources.append(source)
elif isinstance(source, FileSource):
convert_sources.append(source.to_document_stream())
elif isinstance(source, HttpSource):
convert_sources.append(str(source.url))
if headers is None and source.headers:
headers = source.headers
# Note: results are only an iterator->lazy evaluation
results = convert_documents(
sources=sources,
options=task.request.options,
sources=convert_sources,
options=task.options,
headers=headers,
)
# The real processing will happen here
work_dir = get_scratch() / task_id
response = process_results(
background_tasks=BackgroundTasks(),
conversion_options=task.request.options,
conversion_options=task.options,
conv_results=results,
work_dir=work_dir,
)
if work_dir.exists():
task.scratch_dir = work_dir
if not isinstance(response, FileResponse):
_log.warning(
f"Task {task_id=} produced content in {work_dir=} but the response is not a file."
)
shutil.rmtree(work_dir, ignore_errors=True)
return response
# Run the prediction in a thread to avoid blocking the event loop.
start_time = time.monotonic()
# Run the prediction in a thread to avoid blocking the event loop.
# Get the current event loop
# loop = asyncio.get_event_loop()
# future = asyncio.run_coroutine_threadsafe(
# run_conversion(),
# loop=loop
# )
# response = future.result()
# Run in a thread
response = await asyncio.to_thread(
run_conversion,
)
processing_time = time.monotonic() - start_time
if not isinstance(response, ConvertDocumentResponse):
_log.error(
f"Worker {self.worker_id} got un-processable "
"result for {task_id}: {type(response)}"
)
task.result = response
task.request = None
task.sources = []
task.options = None
task.task_status = TaskStatus.SUCCESS
_log.info(

17
docling_serve/engines/async_orchestrator.py
View File

@@ -1,8 +1,13 @@
from fastapi import WebSocket
import shutil
from typing import Union
from fastapi import BackgroundTasks, WebSocket
from fastapi.responses import FileResponse
from docling_serve.datamodel.callback import ProgressCallbackRequest
from docling_serve.datamodel.engines import TaskStatus
from docling_serve.datamodel.responses import (
ConvertDocumentResponse,
MessageKind,
TaskStatusResponse,
WebsocketMessage,
@@ -13,6 +18,7 @@ from docling_serve.engines.base_orchestrator import (
OrchestratorError,
TaskNotFoundError,
)
from docling_serve.settings import docling_serve_settings
class ProgressInvalid(OrchestratorError):
@@ -37,8 +43,15 @@ class BaseAsyncOrchestrator(BaseOrchestrator):
async def task_status(self, task_id: str, wait: float = 0.0) -> Task:
return await self.get_raw_task(task_id=task_id)
async def task_result(self, task_id: str):
async def task_result(
self, task_id: str, background_tasks: BackgroundTasks
) -> Union[ConvertDocumentResponse, FileResponse, None]:
task = await self.get_raw_task(task_id=task_id)
if task.is_completed() and task.scratch_dir is not None:
if docling_serve_settings.single_use_results:
background_tasks.add_task(
shutil.rmtree, task.scratch_dir, ignore_errors=True
)
return task.result
async def notify_task_subscribers(self, task_id: str):

22
docling_serve/engines/base_orchestrator.py
View File

@@ -1,8 +1,12 @@
from abc import ABC, abstractmethod
from typing import Optional
from typing import Optional, Union
from docling_serve.datamodel.requests import ConvertDocumentsRequest
from docling_serve.datamodel.task import Task
from fastapi import BackgroundTasks
from fastapi.responses import FileResponse
from docling_serve.datamodel.convert import ConvertDocumentsOptions
from docling_serve.datamodel.responses import ConvertDocumentResponse
from docling_serve.datamodel.task import Task, TaskSource
class OrchestratorError(Exception):
@@ -15,7 +19,9 @@ class TaskNotFoundError(OrchestratorError):
class BaseOrchestrator(ABC):
@abstractmethod
async def enqueue(self, request: ConvertDocumentsRequest) -> Task:
async def enqueue(
self, sources: list[TaskSource], options: ConvertDocumentsOptions
) -> Task:
pass
@abstractmethod
@@ -31,9 +37,15 @@ class BaseOrchestrator(ABC):
pass
@abstractmethod
async def task_result(self, task_id: str):
async def task_result(
self, task_id: str, background_tasks: BackgroundTasks
) -> Union[ConvertDocumentResponse, FileResponse, None]:
pass
@abstractmethod
async def process_queue(self):
pass
@abstractmethod
async def warm_up_caches(self):
pass

58
docling_serve/gradio_ui.py
View File

@@ -6,6 +6,7 @@ import ssl
import tempfile
import time
from pathlib import Path
from typing import Optional
import certifi
import gradio as gr
@@ -203,12 +204,16 @@ def clear_file_input():
return None
def auto_set_return_as_file(url_input, file_input, image_export_mode):
def auto_set_return_as_file(
url_input_value: str,
file_input_value: Optional[list[str]],
image_export_mode_value: str,
):
# If more than one input source is provided, return as file
if (
(len(url_input.split(",")) > 1)
or (file_input and len(file_input) > 1)
or (image_export_mode == "referenced")
(len(url_input_value.split(",")) > 1)
or (file_input_value and len(file_input_value) > 1)
or (image_export_mode_value == "referenced")
):
return True
else:
@@ -344,7 +349,7 @@ def file_to_base64(file):
def process_file(
file,
files,
to_formats,
image_export_mode,
pipeline,
@@ -361,10 +366,12 @@ def process_file(
do_picture_classification,
do_picture_description,
):
if not file or file == "":
if not files or len(files) == 0:
logger.error("No files provided.")
raise gr.Error("No files provided.", print_exception=False)
files_data = [{"base64_string": file_to_base64(file), "filename": file.name}]
files_data = [
{"base64_string": file_to_base64(file), "filename": file.name} for file in files
]
parameters = {
"file_sources": files_data,
@@ -552,7 +559,7 @@ with gr.Blocks(
".png",
".gif",
],
file_count="single",
file_count="multiple",
scale=4,
)
with gr.Column(scale=1):
@@ -625,9 +632,7 @@ with gr.Blocks(
)
with gr.Column(scale=1):
abort_on_error = gr.Checkbox(label="Abort on Error", value=False)
return_as_file = gr.Checkbox(
label="Return as File", visible=False, value=False
) # Disable until async handle output as file
return_as_file = gr.Checkbox(label="Return as File", value=False)
with gr.Row():
with gr.Column():
do_code_enrichment = gr.Checkbox(
@@ -677,23 +682,22 @@ with gr.Blocks(
# UI Actions #
##############
# Disable until async handle output as file
# Handle Return as File
# url_input.change(
# auto_set_return_as_file,
# inputs=[url_input, file_input, image_export_mode],
# outputs=[return_as_file],
# )
# file_input.change(
# auto_set_return_as_file,
# inputs=[url_input, file_input, image_export_mode],
# outputs=[return_as_file],
# )
# image_export_mode.change(
# auto_set_return_as_file,
# inputs=[url_input, file_input, image_export_mode],
# outputs=[return_as_file],
# )
url_input.change(
auto_set_return_as_file,
inputs=[url_input, file_input, image_export_mode],
outputs=[return_as_file],
)
file_input.change(
auto_set_return_as_file,
inputs=[url_input, file_input, image_export_mode],
outputs=[return_as_file],
)
image_export_mode.change(
auto_set_return_as_file,
inputs=[url_input, file_input, image_export_mode],
outputs=[return_as_file],
)
# URL processing
url_process_btn.click(

9
docling_serve/response_preparation.py
View File

@@ -1,13 +1,12 @@
import logging
import os
import shutil
import tempfile
import time
from collections.abc import Iterable
from pathlib import Path
from typing import Union
from fastapi import BackgroundTasks, HTTPException
from fastapi import HTTPException
from fastapi.responses import FileResponse
from docling.datamodel.base_models import OutputFormat
@@ -124,9 +123,9 @@ def _export_documents_as_files(
def process_results(
background_tasks: BackgroundTasks,
conversion_options: ConvertDocumentsOptions,
conv_results: Iterable[ConversionResult],
work_dir: Path,
) -> Union[ConvertDocumentResponse, FileResponse]:
# Let's start by processing the documents
try:
@@ -183,7 +182,6 @@ def process_results(
# Multiple documents were processed, or we are forced returning as a file
else:
# Temporary directory to store the outputs
work_dir = Path(tempfile.mkdtemp(prefix="docling_"))
output_dir = work_dir / "output"
output_dir.mkdir(parents=True, exist_ok=True)
@@ -203,7 +201,6 @@ def process_results(
)
files = os.listdir(output_dir)
if len(files) == 0:
raise HTTPException(status_code=500, detail="No documents were exported.")
@@ -216,7 +213,7 @@ def process_results(
# Other cleanups after the response is sent
# Output directory
background_tasks.add_task(shutil.rmtree, work_dir, ignore_errors=True)
# background_tasks.add_task(shutil.rmtree, work_dir, ignore_errors=True)
response = FileResponse(
file_path, filename=file_path.name, media_type="application/zip"

2
docling_serve/settings.py
View File

@@ -38,6 +38,8 @@ class DoclingServeSettings(BaseSettings):
api_host: str = "localhost"
artifacts_path: Optional[Path] = None
static_path: Optional[Path] = None
scratch_path: Optional[Path] = None
single_use_results: bool = True
options_cache_size: int = 2
enable_remote_services: bool = False
allow_external_plugins: bool = False

16
docling_serve/storage.py Normal file
View File

@@ -0,0 +1,16 @@
import tempfile
from functools import lru_cache
from pathlib import Path
from docling_serve.settings import docling_serve_settings
@lru_cache
def get_scratch() -> Path:
scratch_dir = (
docling_serve_settings.scratch_path
if docling_serve_settings.scratch_path is not None
else Path(tempfile.mkdtemp(prefix="docling_"))
)
scratch_dir.mkdir(exist_ok=True, parents=True)
return scratch_dir