scribe/scraibe/mcp_server.py

"""
MCP-style HTTP server for ScrAIbe.

- Exposes an OpenAPI-compliant endpoint for external LLMs to:
  - Submit audio (as base64, URL, or internal path) via JSON.
  - 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 logging
from typing import Optional

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

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 submit audio and receive transcript JSON."
    ),
    openapi_tags=[
        {"name": "transcription", "description": "Transcription endpoints"}
    ],
)

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


class TranscribeRequest(BaseModel):
    """
    Input for transcription.

    Exactly one of audio_base64, audio_url, or audio_path must be provided.
    """
    audio_base64: Optional[str] = Field(
        None,
        description="Base64-encoded audio file content."
    )
    audio_url: Optional[str] = Field(
        None,
        description="Public URL to the audio file."
    )
    audio_path: Optional[str] = Field(
        None,
        description="Internal file path on the server (for internal use only)."
    )
    language: Optional[str] = Field(
        None,
        description="Optional language hint (e.g., 'english', 'german')."
    )
    num_speakers: Optional[int] = Field(
        None,
        description="Optional number of speakers for diarization."
    )


class TranscribeResponse(BaseModel):
    job_id: str
    status: str
    message: str


class JobStatusResponse(BaseModel):
    job_id: str
    status: str
    message: str


class TranscriptJSONResponse(BaseModel):
    job_id: str
    transcript: str
    segments: list


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:
        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:
        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:
        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 exactly one of: audio_base64, audio_url, or audio_path",
    )


@app.get("/health", tags=["transcription"])
async def health():
    return {"status": "ok"}


@app.post(
    "/transcribe",
    tags=["transcription"],
    operation_id="transcribe",
    response_model=TranscribeResponse,
)
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 TranscribeResponse(
                    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 TranscribeResponse(
                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 TranscribeResponse(
        job_id=job_id,
        status="processing",
        message=_mcp_jobs[job_id]["message"],
    )


@app.get(
    "/transcribe/{job_id}/status",
    tags=["transcription"],
    operation_id="get_status",
    response_model=JobStatusResponse,
)
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 JobStatusResponse(
        job_id=job_id,
        status=job["status"],
        message=job.get("message", ""),
    )


@app.get(
    "/transcribe/{job_id}/json",
    tags=["transcription"],
    operation_id="get_json",
    response_model=TranscriptJSONResponse,
)
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 TranscriptJSONResponse(
        job_id=job_id,
        transcript=transcript_text,
        segments=segments,
    )