scribe/scraibe/mcp_server.py

"""
MCP-style HTTP server for ScrAIbe.

- Exposes an OpenAPI-compliant endpoint for external LLMs to:
  - Upload audio (as base64 or via URL) in a JSON body.
  - Receive transcript JSON (no summary).
- WebUI remains always enabled; this is additive.

Configuration (env):
- MCP_SERVER_ENABLED: "true"/"false" (default: false)
- MCP_SERVER_HOST: bind address (default: 0.0.0.0)
- MCP_SERVER_PORT: port (default: 8000)
- MCP_USE_CELERY: "true"/"false" (default: true)
  - If true, uses Celery tasks; if false, runs synchronously.
"""

import os
import time
import uuid
import base64
import tempfile
import logging
from typing import Optional

import httpx
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel

from .autotranscript import Scraibe

logger = logging.getLogger("scraibe.mcp_server")

app = FastAPI(
    title="ScrAIbe MCP Transcription API",
    version="0.1.0",
    description=(
        "MCP-style HTTP API for ScrAIbe. "
        "Allows external LLMs to upload audio and receive transcript JSON."
    ),
)

# In-memory job store for MCP (simple; can be replaced with Redis later)
_mcp_jobs: dict = {}


class TranscribeRequest(BaseModel):
    audio_base64: Optional[str] = None
    audio_url: Optional[str] = None
    audio_path: Optional[str] = None
    language: Optional[str] = None
    num_speakers: Optional[int] = None


def _job_id() -> str:
    return str(uuid.uuid4())


def _save_audio_from_request(req: TranscribeRequest) -> str:
    """
    Save audio to a temporary file from base64, URL, or path.
    Returns the local file path.
    """
    upload_dir = os.getenv("SCRAIBE_UPLOAD_DIR", "/tmp/scraibe_uploads")
    os.makedirs(upload_dir, exist_ok=True)

    if req.audio_base64:
        # base64 upload
        try:
            data = base64.b64decode(req.audio_base64)
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Invalid base64 audio: {e}")

        ts = time.strftime("%Y%m%d%H%M%S")
        tmp_name = f"mcp_upload_{ts}_{uuid.uuid4().hex[:8]}.wav"
        file_path = os.path.join(upload_dir, tmp_name)
        with open(file_path, "wb") as f:
            f.write(data)
        return file_path

    if req.audio_url:
        # download from URL
        try:
            with httpx.stream("GET", req.audio_url, timeout=60) as resp:
                if resp.status_code != 200:
                    raise HTTPException(
                        status_code=400,
                        detail=f"Failed to download audio from URL: {resp.status_code}",
                    )
                ts = time.strftime("%Y%m%d%H%M%S")
                tmp_name = f"mcp_url_{ts}_{uuid.uuid4().hex[:8]}.wav"
                file_path = os.path.join(upload_dir, tmp_name)
                with open(file_path, "wb") as f:
                    for chunk in resp.iter_bytes():
                        f.write(chunk)
            return file_path
        except HTTPException:
            raise
        except Exception as e:
            raise HTTPException(status_code=400, detail=f"Error downloading audio: {e}")

    if req.audio_path:
        # use provided path (for internal use)
        path = req.audio_path
        if not os.path.isfile(path):
            raise HTTPException(status_code=400, detail="audio_path does not exist")
        return path

    raise HTTPException(
        status_code=400,
        detail="Provide one of: audio_base64, audio_url, or audio_path",
    )


@app.get("/health")
async def health():
    return {"status": "ok"}


@app.post("/transcribe")
async def transcribe(req: TranscribeRequest):
    """
    Submit an audio file for transcription.

    Input (JSON body):
      - audio_base64: base64-encoded audio file
      - audio_url: URL to audio file
      - audio_path: local file path (for internal use)
      - language: (optional)
      - num_speakers: (optional)

    Returns:
      {
        "job_id": "<id>",
        "status": "queued" | "processing",
        "message": "..."
      }

    Use GET /transcribe/{job_id}/status and /json to retrieve results.
    """
    use_celery = os.getenv("MCP_USE_CELERY", "true").strip().lower() in ("true", "1", "yes")

    # Save audio to a temporary file
    try:
        file_path = _save_audio_from_request(req)
    except HTTPException:
        raise
    except Exception as e:
        logger.error("Error saving MCP upload: %s", e)
        raise HTTPException(status_code=500, detail=f"Error saving file: {e}")

    job_id = _job_id()

    if use_celery:
        try:
            from .tasks import process_mcp_transcribe_task
        except ImportError:
            # Fallback: run synchronously
            use_celery = False

        if use_celery:
            try:
                process_mcp_transcribe_task.delay(
                    audio_path=file_path,
                    job_id=job_id,
                    language=req.language or None,
                    num_speakers=int(req.num_speakers) if req.num_speakers else None,
                )
            except Exception as e:
                logger.error("Error enqueuing MCP job: %s", e)
                _mcp_jobs[job_id] = {
                    "status": "error",
                    "message": f"Error enqueuing job: {e}",
                    "file_path": file_path,
                }
                return {
                    "job_id": job_id,
                    "status": "error",
                    "message": _mcp_jobs[job_id]["message"],
                }

            _mcp_jobs[job_id] = {
                "status": "queued",
                "message": "Job queued for processing.",
                "file_path": file_path,
            }
            return {
                "job_id": job_id,
                "status": "queued",
                "message": _mcp_jobs[job_id]["message"],
            }

    # Synchronous path
    _mcp_jobs[job_id] = {
        "status": "processing",
        "message": "Transcription started (synchronous).",
        "file_path": file_path,
    }

    def _run_sync():
        try:
            scraibe = Scraibe(verbose=False)
            result = scraibe.transcribe(
                audio_file=file_path,
                language=req.language or None,
                num_speakers=int(req.num_speakers) if req.num_speakers else None,
                verbose=False,
                for_export=True,
            )
            transcript_text = result.get("transcript", "")
            segments = result.get("segments", [])
            _mcp_jobs[job_id]["status"] = "completed"
            _mcp_jobs[job_id]["transcript"] = transcript_text
            _mcp_jobs[job_id]["segments"] = segments
            _mcp_jobs[job_id]["message"] = "Transcription completed."
        except Exception as e:
            logger.error("MCP sync transcription error: %s", e)
            _mcp_jobs[job_id]["status"] = "error"
            _mcp_jobs[job_id]["message"] = f"Transcription error: {e}"

    import threading
    t = threading.Thread(target=_run_sync, daemon=True)
    t.start()

    return {
        "job_id": job_id,
        "status": "processing",
        "message": _mcp_jobs[job_id]["message"],
    }


@app.get("/transcribe/{job_id}/status")
async def get_status(job_id: str):
    job = _mcp_jobs.get(job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")
    return {
        "job_id": job_id,
        "status": job["status"],
        "message": job.get("message", ""),
    }


@app.get("/transcribe/{job_id}/json")
async def get_json(job_id: str):
    job = _mcp_jobs.get(job_id)
    if not job:
        raise HTTPException(status_code=404, detail="Job not found")

    if job["status"] != "completed":
        raise HTTPException(
            status_code=400,
            detail=f"Job not completed. Current status: {job['status']}",
        )

    transcript_text = job.get("transcript", "")
    segments = job.get("segments", [])

    return {
        "job_id": job_id,
        "transcript": transcript_text,
        "segments": segments,
    }