root
0cd51ddd4c
Adds optional save_to_s3=True to convert_document. After conversion, the text output is uploaded to chat-artifacts S3 and s3_url is returned alongside the result. Also fixes boto3 1.35+ checksum regression (request_checksum_calculation=when_required) and adds boto3 dep. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
235 lines
8.5 KiB
Python
235 lines
8.5 KiB
Python
"""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")
|