fkrebs/mcp-docling
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: S3 upload on conversion complete (save_to_s3 param)

Browse Source 
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>
This commit is contained in:
root
2026-05-21 09:29:09 +02:00
parent cf451ad784
commit 0cd51ddd4c
2 changed files with 67 additions and 7 deletions
Download Patch File Download Diff File Expand all files Collapse all files
server.py
+66 -7
View File
@@ -18,9 +18,12 @@ 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", "")
@@ -30,6 +33,21 @@ API_ENDPOINT = os.environ.get(
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}
@@ -37,6 +55,35 @@ _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:
@@ -53,18 +100,24 @@ async def _do_convert(filename: str, content: bytes, params: dict) -> dict:
return r.json()
async def _run_job(job_id, filename, content, params, response_type):
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: data.get(response_type, data.get("markdown", "")),
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])
@@ -72,7 +125,10 @@ async def _run_job(job_id, filename, content, params, response_type):
def _result(job_id: str) -> dict:
j = _jobs[job_id]
if j["status"] == "done":
return {"job_id": job_id, "status": "done", **j["result"]}
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 {
@@ -89,6 +145,7 @@ async def convert_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.
@@ -103,10 +160,12 @@ async def convert_document(
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
plus "images"), or status "running" with a job_id to poll.
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}")
@@ -136,7 +195,7 @@ async def convert_document(
job_id = f"doc-{uuid.uuid4().hex[:8]}"
_jobs[job_id] = {
"status": "running", "result": None, "error": None,
"status": "running", "result": None, "error": None, "s3_url": None,
"submitted_at": time.time(), "filename": filename,
}
if len(_jobs) > 200:
@@ -144,7 +203,7 @@ async def convert_document(
_jobs.pop(k, None)
task = asyncio.create_task(
_run_job(job_id, filename, content, params, response_type)
_run_job(job_id, filename, content, params, response_type, save_to_s3)
)
_tasks.add(task)
task.add_done_callback(_tasks.discard)