docling-serve 0.3.0__py3-none-any.whl

docling_serve/__main__.py

@@ -0,0 +1,302 @@
+import importlib
+import logging
+import platform
+import sys
+import warnings
+from pathlib import Path
+from typing import Annotated, Any, Optional, Union
+
+import typer
+import uvicorn
+from rich.console import Console
+
+from docling_serve.settings import docling_serve_settings, uvicorn_settings
+
+warnings.filterwarnings(action="ignore", category=UserWarning, module="pydantic|torch")
+warnings.filterwarnings(action="ignore", category=FutureWarning, module="easyocr")
+
+
+err_console = Console(stderr=True)
+console = Console()
+
+app = typer.Typer(
+    no_args_is_help=True,
+    rich_markup_mode="rich",
+)
+
+logger = logging.getLogger(__name__)
+
+
+def version_callback(value: bool) -> None:
+    if value:
+        docling_serve_version = importlib.metadata.version("docling_serve")
+        docling_version = importlib.metadata.version("docling")
+        docling_core_version = importlib.metadata.version("docling-core")
+        docling_ibm_models_version = importlib.metadata.version("docling-ibm-models")
+        docling_parse_version = importlib.metadata.version("docling-parse")
+        platform_str = platform.platform()
+        py_impl_version = sys.implementation.cache_tag
+        py_lang_version = platform.python_version()
+        console.print(f"Docling Serve version: {docling_serve_version}")
+        console.print(f"Docling version: {docling_version}")
+        console.print(f"Docling Core version: {docling_core_version}")
+        console.print(f"Docling IBM Models version: {docling_ibm_models_version}")
+        console.print(f"Docling Parse version: {docling_parse_version}")
+        console.print(f"Python: {py_impl_version} ({py_lang_version})")
+        console.print(f"Platform: {platform_str}")
+        raise typer.Exit()
+
+
+@app.callback()
+def callback(
+    version: Annotated[
+        Union[bool, None],
+        typer.Option(
+            "--version", help="Show the version and exit.", callback=version_callback
+        ),
+    ] = None,
+    verbose: Annotated[
+        int,
+        typer.Option(
+            "--verbose",
+            "-v",
+            count=True,
+            help="Set the verbosity level. -v for info logging, -vv for debug logging.",
+        ),
+    ] = 0,
+) -> None:
+    if verbose == 0:
+        logging.basicConfig(level=logging.WARNING)
+    elif verbose == 1:
+        logging.basicConfig(level=logging.INFO)
+    elif verbose == 2:
+        logging.basicConfig(level=logging.DEBUG)
+
+
+def _run(
+    *,
+    command: str,
+) -> None:
+    server_type = "development" if command == "dev" else "production"
+
+    console.print(f"Starting {server_type} server 🚀")
+
+    url = f"http://{uvicorn_settings.host}:{uvicorn_settings.port}"
+    url_docs = f"{url}/docs"
+    url_ui = f"{url}/ui"
+
+    console.print("")
+    console.print(f"Server started at [link={url}]{url}[/]")
+    console.print(f"Documentation at [link={url_docs}]{url_docs}[/]")
+    if docling_serve_settings.enable_ui:
+        console.print(f"UI at [link={url_ui}]{url_ui}[/]")
+
+    if command == "dev":
+        console.print("")
+        console.print(
+            "Running in development mode, for production use: "
+            "[bold]docling-serve run[/]",
+        )
+
+    console.print("")
+    console.print("Logs:")
+
+    uvicorn.run(
+        app="docling_serve.app:create_app",
+        factory=True,
+        host=uvicorn_settings.host,
+        port=uvicorn_settings.port,
+        reload=uvicorn_settings.reload,
+        workers=uvicorn_settings.workers,
+        root_path=uvicorn_settings.root_path,
+        proxy_headers=uvicorn_settings.proxy_headers,
+    )
+
+
+@app.command()
+def dev(
+    *,
+    # uvicorn options
+    host: Annotated[
+        str,
+        typer.Option(
+            help=(
+                "The host to serve on. For local development in localhost "
+                "use [blue]127.0.0.1[/blue]. To enable public access, "
+                "e.g. in a container, use all the IP addresses "
+                "available with [blue]0.0.0.0[/blue]."
+            )
+        ),
+    ] = "127.0.0.1",
+    port: Annotated[
+        int,
+        typer.Option(help="The port to serve on."),
+    ] = uvicorn_settings.port,
+    reload: Annotated[
+        bool,
+        typer.Option(
+            help=(
+                "Enable auto-reload of the server when (code) files change. "
+                "This is [bold]resource intensive[/bold], "
+                "use it only during development."
+            )
+        ),
+    ] = True,
+    root_path: Annotated[
+        str,
+        typer.Option(
+            help=(
+                "The root path is used to tell your app that it is being served "
+                "to the outside world with some [bold]path prefix[/bold] "
+                "set up in some termination proxy or similar."
+            )
+        ),
+    ] = uvicorn_settings.root_path,
+    proxy_headers: Annotated[
+        bool,
+        typer.Option(
+            help=(
+                "Enable/Disable X-Forwarded-Proto, X-Forwarded-For, "
+                "X-Forwarded-Port to populate remote address info."
+            )
+        ),
+    ] = uvicorn_settings.proxy_headers,
+    # docling options
+    artifacts_path: Annotated[
+        Optional[Path],
+        typer.Option(
+            help=(
+                "If set to a valid directory, "
+                "the model weights will be loaded from this path."
+            )
+        ),
+    ] = docling_serve_settings.artifacts_path,
+    enable_ui: Annotated[bool, typer.Option(help="Enable the development UI.")] = True,
+) -> Any:
+    """
+    Run a [bold]Docling Serve[/bold] app in [yellow]development[/yellow] mode. 🧪
+
+    This is equivalent to [bold]docling-serve run[/bold] but with [bold]reload[/bold]
+    enabled and listening on the [blue]127.0.0.1[/blue] address.
+
+    Options can be set also with the corresponding ENV variable, with the exception
+    of --enable-ui, --host and --reload.
+    """
+
+    uvicorn_settings.host = host
+    uvicorn_settings.port = port
+    uvicorn_settings.reload = reload
+    uvicorn_settings.root_path = root_path
+    uvicorn_settings.proxy_headers = proxy_headers
+
+    docling_serve_settings.artifacts_path = artifacts_path
+    docling_serve_settings.enable_ui = enable_ui
+
+    _run(
+        command="dev",
+    )
+
+
+@app.command()
+def run(
+    *,
+    host: Annotated[
+        str,
+        typer.Option(
+            help=(
+                "The host to serve on. For local development in localhost "
+                "use [blue]127.0.0.1[/blue]. To enable public access, "
+                "e.g. in a container, use all the IP addresses "
+                "available with [blue]0.0.0.0[/blue]."
+            )
+        ),
+    ] = uvicorn_settings.host,
+    port: Annotated[
+        int,
+        typer.Option(help="The port to serve on."),
+    ] = uvicorn_settings.port,
+    reload: Annotated[
+        bool,
+        typer.Option(
+            help=(
+                "Enable auto-reload of the server when (code) files change. "
+                "This is [bold]resource intensive[/bold], "
+                "use it only during development."
+            )
+        ),
+    ] = uvicorn_settings.reload,
+    workers: Annotated[
+        Union[int, None],
+        typer.Option(
+            help=(
+                "Use multiple worker processes. "
+                "Mutually exclusive with the --reload flag."
+            )
+        ),
+    ] = uvicorn_settings.workers,
+    root_path: Annotated[
+        str,
+        typer.Option(
+            help=(
+                "The root path is used to tell your app that it is being served "
+                "to the outside world with some [bold]path prefix[/bold] "
+                "set up in some termination proxy or similar."
+            )
+        ),
+    ] = uvicorn_settings.root_path,
+    proxy_headers: Annotated[
+        bool,
+        typer.Option(
+            help=(
+                "Enable/Disable X-Forwarded-Proto, X-Forwarded-For, "
+                "X-Forwarded-Port to populate remote address info."
+            )
+        ),
+    ] = uvicorn_settings.proxy_headers,
+    # docling options
+    artifacts_path: Annotated[
+        Optional[Path],
+        typer.Option(
+            help=(
+                "If set to a valid directory, "
+                "the model weights will be loaded from this path."
+            )
+        ),
+    ] = docling_serve_settings.artifacts_path,
+    enable_ui: Annotated[
+        bool, typer.Option(help="Enable the development UI.")
+    ] = docling_serve_settings.enable_ui,
+) -> Any:
+    """
+    Run a [bold]Docling Serve[/bold] app in [green]production[/green] mode. 🚀
+
+    This is equivalent to [bold]docling-serve dev[/bold] but with [bold]reload[/bold]
+    disabled and listening on the [blue]0.0.0.0[/blue] address.
+
+    Options can be set also with the corresponding ENV variable, e.g. UVICORN_PORT
+    or DOCLING_SERVE_ENABLE_UI.
+    """
+
+    uvicorn_settings.host = host
+    uvicorn_settings.port = port
+    uvicorn_settings.reload = reload
+    uvicorn_settings.workers = workers
+    uvicorn_settings.root_path = root_path
+    uvicorn_settings.proxy_headers = proxy_headers
+
+    docling_serve_settings.artifacts_path = artifacts_path
+    docling_serve_settings.enable_ui = enable_ui
+
+    _run(
+        command="run",
+    )
+
+
+def main() -> None:
+    app()
+
+
+# Launch the CLI when calling python -m docling_serve
+if __name__ == "__main__":
+
+    main()

docling_serve/app.py

@@ -0,0 +1,230 @@
+import logging
+import tempfile
+from contextlib import asynccontextmanager
+from io import BytesIO
+from pathlib import Path
+from typing import Annotated, Any, Dict, List, Optional, Union
+
+from docling.datamodel.base_models import DocumentStream, InputFormat
+from docling.document_converter import DocumentConverter
+from fastapi import BackgroundTasks, FastAPI, UploadFile
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import RedirectResponse
+from pydantic import BaseModel
+
+from docling_serve.docling_conversion import (
+    ConvertDocumentFileSourcesRequest,
+    ConvertDocumentsOptions,
+    ConvertDocumentsRequest,
+    convert_documents,
+    converters,
+    get_pdf_pipeline_opts,
+)
+from docling_serve.helper_functions import FormDepends
+from docling_serve.response_preparation import ConvertDocumentResponse, process_results
+from docling_serve.settings import docling_serve_settings
+
+
+# Set up custom logging as we'll be intermixes with FastAPI/Uvicorn's logging
+class ColoredLogFormatter(logging.Formatter):
+    COLOR_CODES = {
+        logging.DEBUG: "\033[94m",  # Blue
+        logging.INFO: "\033[92m",  # Green
+        logging.WARNING: "\033[93m",  # Yellow
+        logging.ERROR: "\033[91m",  # Red
+        logging.CRITICAL: "\033[95m",  # Magenta
+    }
+    RESET_CODE = "\033[0m"
+
+    def format(self, record):
+        color = self.COLOR_CODES.get(record.levelno, "")
+        record.levelname = f"{color}{record.levelname}{self.RESET_CODE}"
+        return super().format(record)
+
+
+logging.basicConfig(
+    level=logging.INFO,  # Set the logging level
+    format="%(levelname)s:\t%(asctime)s - %(name)s - %(message)s",
+    datefmt="%H:%M:%S",
+)
+
+# Override the formatter with the custom ColoredLogFormatter
+root_logger = logging.getLogger()  # Get the root logger
+for handler in root_logger.handlers:  # Iterate through existing handlers
+    if handler.formatter:
+        handler.setFormatter(ColoredLogFormatter(handler.formatter._fmt))
+
+_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, options_hash = get_pdf_pipeline_opts(ConvertDocumentsOptions())
+    converters[options_hash] = DocumentConverter(
+        format_options={
+            InputFormat.PDF: pdf_format_option,
+            InputFormat.IMAGE: pdf_format_option,
+        }
+    )
+
+    converters[options_hash].initialize_pipeline(InputFormat.PDF)
+
+    yield
+
+    converters.clear()
+    # if WITH_UI:
+    #     gradio_ui.close()
+
+
+##################################
+# App creation and configuration #
+##################################
+
+
+def create_app():
+    app = FastAPI(
+        title="Docling Serve",
+        lifespan=lifespan,
+    )
+
+    origins = ["*"]
+    methods = ["*"]
+    headers = ["*"]
+
+    app.add_middleware(
+        CORSMiddleware,
+        allow_origins=origins,
+        allow_credentials=True,
+        allow_methods=methods,
+        allow_headers=headers,
+    )
+
+    # Mount the Gradio app
+    if docling_serve_settings.enable_ui:
+
+        try:
+            import gradio as gr
+
+            from docling_serve.gradio_ui import ui as gradio_ui
+
+            tmp_output_dir = Path(tempfile.mkdtemp())
+            gradio_ui.gradio_output_dir = tmp_output_dir
+            app = gr.mount_gradio_app(
+                app,
+                gradio_ui,
+                path="/ui",
+                allowed_paths=["./logo.png", tmp_output_dir],
+                root_path="/ui",
+            )
+        except ImportError:
+            _log.warning(
+                "Docling Serve enable_ui is activated, but gradio is not installed. "
+                "Install it with `pip install docling-serve[ui]` "
+                "or `pip install gradio`"
+            )
+
+    #############################
+    # API Endpoints definitions #
+    #############################
+
+    # Favicon
+    @app.get("/favicon.ico", include_in_schema=False)
+    async def favicon():
+        response = RedirectResponse(
+            url="https://ds4sd.github.io/docling/assets/logo.png"
+        )
+        return response
+
+    # Status
+    class HealthCheckResponse(BaseModel):
+        status: str = "ok"
+
+    @app.get("/health")
+    def health() -> HealthCheckResponse:
+        return HealthCheckResponse()
+
+    # API readiness compatibility for OpenShift AI Workbench
+    @app.get("/api", include_in_schema=False)
+    def api_check() -> HealthCheckResponse:
+        return HealthCheckResponse()
+
+    # Convert a document from URL(s)
+    @app.post(
+        "/v1alpha/convert/source",
+        response_model=ConvertDocumentResponse,
+        responses={
+            200: {
+                "content": {"application/zip": {}},
+                # "description": "Return the JSON item or an image.",
+            }
+        },
+    )
+    def process_url(
+        background_tasks: BackgroundTasks, 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
+        )
+
+        # The real processing will happen here
+        response = process_results(
+            background_tasks=background_tasks,
+            conversion_options=conversion_request.options,
+            conv_results=results,
+        )
+
+        return response
+
+    # Convert a document from file(s)
+    @app.post(
+        "/v1alpha/convert/file",
+        response_model=ConvertDocumentResponse,
+        responses={
+            200: {
+                "content": {"application/zip": {}},
+            }
+        },
+    )
+    async def process_file(
+        background_tasks: BackgroundTasks,
+        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,
+        )
+
+        return response
+
+    return app