"""SAIA Docling MCP server — document → Markdown/HTML/JSON conversion. Wraps GWDG SAIA's Docling endpoint (POST /v1/documents/convert). Async queue: large PDFs (OCR + image extraction) can take minutes and would blow past LLM client tool-call timeouts. `convert_document` submits the job and waits up to CONVERT_DEADLINE seconds inline; if it's still running it returns a job_id to poll with `get_job_status`. The conversion keeps running in the background regardless (asyncio.shield), so the result is never lost. Defaults chosen for LLM/RAG ingestion: response_type = markdown (best for embedding / chat context) extract_tables_as_images = False (tables stay structured, not pictures) image_resolution_scale = 2 (images kept, reasonable size) """ import asyncio import base64 import os import time import uuid from datetime import datetime, timezone from urllib.parse import urlparse import boto3 import httpx from botocore.config import Config from mcp.server.fastmcp import FastMCP API_KEY = os.environ.get("SAIA_API_KEY", "") API_ENDPOINT = os.environ.get( "DOCLING_ENDPOINT", "https://chat-ai.academiccloud.de/v1/documents/convert" ) CONVERT_DEADLINE = int(os.environ.get("CONVERT_DEADLINE", 60)) _VALID_TYPES = ("markdown", "html", "json", "tokens") _S3_ENDPOINT = os.environ.get("CHAT_ARTIFACTS_ENDPOINT", "") _S3_ACCESS = os.environ.get("CHAT_ARTIFACTS_ACCESS_KEY", "") _S3_SECRET = os.environ.get("CHAT_ARTIFACTS_SECRET_KEY", "") _S3_BUCKET = os.environ.get("CHAT_ARTIFACTS_BUCKET", "chat-artifacts") _S3_PUBLIC = "https://chat-artifacts.s3.nuclide.systems" _S3_ENABLED = bool(_S3_ENDPOINT and _S3_ACCESS and _S3_SECRET) _EXT = {"markdown": "md", "html": "html", "json": "json", "tokens": "txt"} _CT = { "markdown": "text/markdown; charset=utf-8", "html": "text/html; charset=utf-8", "json": "application/json", "tokens": "text/plain; charset=utf-8", } mcp = FastMCP("saia-docling", host="0.0.0.0", port=8000, stateless_http=True) # job_id -> {status, result, error, submitted_at, filename} _jobs: dict[str, dict] = {} _tasks: set[asyncio.Task] = set() def _s3_upload(text: str, job_id: str, response_type: str) -> str | None: if not _S3_ENABLED: return None ext = _EXT.get(response_type, "txt") key = datetime.now(timezone.utc).strftime("%Y-%m-%d/") + f"docling-{job_id}.{ext}" try: boto3.client( "s3", endpoint_url=_S3_ENDPOINT, aws_access_key_id=_S3_ACCESS, aws_secret_access_key=_S3_SECRET, region_name="garage", config=Config( signature_version="s3v4", request_checksum_calculation="when_required", response_checksum_validation="when_required", ), ).put_object( Bucket=_S3_BUCKET, Key=key, Body=text.encode(), ContentType=_CT.get(response_type, "text/plain"), ) return f"{_S3_PUBLIC}/{key}" except Exception as exc: print(f"S3 upload failed for {job_id}: {exc}", flush=True) return None async def _do_convert(filename: str, content: bytes, params: dict) -> dict: headers = {"accept": "application/json"} if API_KEY: headers["Authorization"] = f"Bearer {API_KEY}" async with httpx.AsyncClient(timeout=httpx.Timeout(600, connect=10)) as c: r = await c.post( API_ENDPOINT, params=params, headers=headers, files={"document": (filename, content)}, ) if r.status_code >= 400: raise RuntimeError(f"Docling API {r.status_code}: {r.text[:400]}") return r.json() async def _run_job(job_id, filename, content, params, response_type, save_to_s3=True): try: data = await _do_convert(filename, content, params) text = data.get(response_type, data.get("markdown", "")) _jobs[job_id].update( status="done", result={ "response_type": data.get("response_type", response_type.upper()), "filename": data.get("filename", filename), response_type: text, "images": data.get("images", []), }, ) if save_to_s3 and text: s3_url = await asyncio.get_event_loop().run_in_executor( None, _s3_upload, text, job_id, response_type ) _jobs[job_id]["s3_url"] = s3_url except Exception as e: # noqa: BLE001 _jobs[job_id].update(status="error", error=str(e)[:500]) def _result(job_id: str) -> dict: j = _jobs[job_id] if j["status"] == "done": out = {"job_id": job_id, "status": "done", **j["result"]} if j.get("s3_url"): out["s3_url"] = j["s3_url"] return out if j["status"] == "error": return {"job_id": job_id, "status": "error", "error": j["error"]} return { "job_id": job_id, "status": "running", "elapsed_s": int(time.time() - j["submitted_at"]), } @mcp.tool() async def convert_document( source: str, filename: str = "document", response_type: str = "markdown", extract_tables_as_images: bool = False, image_resolution_scale: int = 2, save_to_s3: bool = True, ) -> dict: """Convert a document (PDF, DOCX, PPTX, images, …) to Markdown/HTML/JSON. Submits the job and waits inline up to ~60s. If conversion takes longer (big/scanned PDFs), returns a job_id — call `get_job_status(job_id)` to retrieve the result when ready (it keeps processing in the background). Args: source: An http(s) URL to the document, OR its base64-encoded bytes. filename: Name (with extension) for the upload; aids format detection. response_type: "markdown" (default), "html", "json", or "tokens". extract_tables_as_images: False (default) keeps tables as structured text; True returns them as images. image_resolution_scale: Image scaling factor 1–4 (default 2). save_to_s3: Upload the converted text to chat-artifacts S3 and include an s3_url in the response (default True). Set False to skip upload. Returns: dict with status "done" (incl. the converted text under its format key, "images", and optionally "s3_url"), or status "running" with a job_id. """ if response_type not in _VALID_TYPES: raise ValueError(f"response_type must be one of {_VALID_TYPES}") if image_resolution_scale not in (1, 2, 3, 4): raise ValueError("image_resolution_scale must be 1, 2, 3, or 4") if source.startswith(("http://", "https://")): async with httpx.AsyncClient(timeout=60, follow_redirects=True) as c: resp = await c.get(source) resp.raise_for_status() content = resp.content if filename == "document": filename = urlparse(source).path.split("/")[-1] or "document" else: try: content = base64.b64decode(source, validate=True) except Exception as e: raise ValueError( "source is neither an http(s) URL nor valid base64" ) from e params = { "response_type": response_type, "extract_tables_as_images": str(extract_tables_as_images).lower(), "image_resolution_scale": str(image_resolution_scale), } job_id = f"doc-{uuid.uuid4().hex[:8]}" _jobs[job_id] = { "status": "running", "result": None, "error": None, "s3_url": None, "submitted_at": time.time(), "filename": filename, } if len(_jobs) > 200: for k in sorted(_jobs, key=lambda k: _jobs[k]["submitted_at"])[:50]: _jobs.pop(k, None) task = asyncio.create_task( _run_job(job_id, filename, content, params, response_type, save_to_s3) ) _tasks.add(task) task.add_done_callback(_tasks.discard) try: await asyncio.wait_for(asyncio.shield(task), timeout=CONVERT_DEADLINE) except asyncio.TimeoutError: return { "job_id": job_id, "status": "running", "note": ( f"'{filename}' still converting after {CONVERT_DEADLINE}s. " f"Call get_job_status('{job_id}') to retrieve it when ready." ), } return _result(job_id) @mcp.tool() async def get_job_status(job_id: str) -> dict: """Retrieve a previously submitted conversion job by its job_id.""" if job_id not in _jobs: return {"error": f"unknown job_id '{job_id}'"} return _result(job_id) if __name__ == "__main__": mcp.run(transport="streamable-http")