diff --git a/requirements.txt b/requirements.txt index c0db0b6..98185a9 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ mcp[cli]>=1.6.0 httpx>=0.27 +boto3>=1.34 diff --git a/server.py b/server.py index 8bbced6..d1d5576 100644 --- a/server.py +++ b/server.py @@ -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)