data/api.py

#!/usr/bin/env python3
"""
FastAPI backend for Audio Transcription Pipeline.
Serves on port 8081 with endpoints for transcription, ASCII viz, and health check.
"""

import json
import logging
import os
import shutil
import tempfile
import time
from typing import Optional

import librosa
import uvicorn
from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse, PlainTextResponse

logging.basicConfig(
    level=logging.INFO,
    format="%(levelname)s:%(name)s:%(message)s",
)
logger = logging.getLogger(__name__)

# Quiet down libraries
logging.getLogger("faster_whisper").setLevel(logging.WARNING)
logging.getLogger("transformers").setLevel(logging.WARNING)
logging.getLogger("librosa").setLevel(logging.WARNING)
logging.getLogger("pipeline").setLevel(logging.WARNING)
logging.getLogger("httpx").setLevel(logging.WARNING)

app = FastAPI(
    title="Audio Transcription Pipeline",
    description="Transcribe, classify, diarize, summarize, and visualize audio",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# Lazy-loaded pipeline
_pipeline = None


def _get_pipeline():
    global _pipeline
    if _pipeline is None:
        logger.info("Initializing AudioPipeline...")
        from pipeline.orchestrator import AudioPipeline

        _pipeline = AudioPipeline(
            enable_summarizer=True,
            enable_ascii=False,  # ASCII on-demand
        )
        logger.info("AudioPipeline initialized.")
    return _pipeline


def _get_ascii_viz():
    """Get AsciiSpectrogram instance (no caching - lightweight)."""
    from pipeline.ascii_spectrogram import AsciiSpectrogram
    return AsciiSpectrogram()


def _clean_for_json(obj):
    """Recursively clean for JSON serialization."""
    if isinstance(obj, dict):
        return {k: _clean_for_json(v) for k, v in obj.items() if not k.startswith("_")}
    elif isinstance(obj, list):
        return [_clean_for_json(item) for item in obj]
    elif isinstance(obj, float):
        if obj != obj:  # NaN
            return None
        return obj
    return obj


@app.get("/health")
async def health_check():
    """Health check endpoint."""
    return {
        "status": "ok",
        "service": "audio-transcription-pipeline",
        "version": "1.0.0",
    }


@app.post("/transcribe")
async def transcribe(
    file: UploadFile = File(...),
    language: Optional[str] = Form(None),
    vad_filter: bool = Form(True),
    include_ascii: bool = Form(False),
    include_summary: bool = Form(True),
):
    """Transcribe an uploaded audio file through the full pipeline."""
    start_time = time.time()

    # Save uploaded file to temp
    suffix = os.path.splitext(file.filename or "audio.wav")[1] or ".wav"
    with tempfile.NamedTemporaryFile(delete=False, suffix=suffix) as tmp:
        shutil.copyfileobj(file.file, tmp)
        tmp_path = tmp.name

    try:
        pipe = _get_pipeline()
        result = pipe.process_file(
            tmp_path,
            language=language,
            vad_filter=vad_filter,
        )

        # Add ASCII if requested
        if include_ascii:
            try:
                audio, sr = librosa.load(tmp_path, sr=16000, mono=True)
                viz = _get_ascii_viz()
                frames = []
                for frame_text, timestamp in viz.generate_frames(audio, sr, fps=10):
                    frames.append({"timestamp": round(timestamp, 2), "frame": frame_text})
                result["ascii_frames"] = frames
            except Exception as e:
                logger.warning(f"ASCII viz failed: {e}")
                result["ascii_frames"] = []

        # Add processing time
        elapsed = time.time() - start_time
        result.setdefault("metadata", {})["api_processing_time_seconds"] = round(elapsed, 2)

        clean = _clean_for_json(result)
        return JSONResponse(content=clean)

    except Exception as e:
        logger.exception(f"Transcription failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))

    finally:
        # Clean up temp file
        try:
            os.unlink(tmp_path)
        except OSError:
            pass


@app.get("/ascii-viz")
async def ascii_viz(
    audio_path: str = Form(...),
    columns: int = 80,
    rows: int = 20,
    fps: int = 10,
    mode: str = "spectrogram",
):
    """Generate ASCII spectrogram frames from an audio file path."""
    if not os.path.exists(audio_path):
        raise HTTPException(status_code=404, detail=f"Audio file not found: {audio_path}")

    try:
        audio, sr = librosa.load(audio_path, sr=16000, mono=True)
        viz = _get_ascii_viz()
        frames = []
        for frame_text, timestamp in viz.generate_frames(audio, sr, fps=fps, mode=mode):
            frames.append({"timestamp": round(timestamp, 2), "frame": frame_text})
        return {"frames": frames, "count": len(frames), "duration": round(len(audio) / sr, 2)}
    except Exception as e:
        logger.exception(f"ASCII viz failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


@app.get("/ascii-viz-file")
async def ascii_viz_file(
    audio_path: str,
    columns: int = 80,
    rows: int = 20,
    fps: int = 10,
    mode: str = "spectrogram",
):
    """Generate ASCII spectrogram and return as plain text."""
    if not os.path.exists(audio_path):
        raise HTTPException(status_code=404, detail=f"Audio file not found: {audio_path}")

    try:
        audio, sr = librosa.load(audio_path, sr=16000, mono=True)
        viz = _get_ascii_viz()
        lines = []
        for frame_text, timestamp in viz.generate_frames(audio, sr, fps=fps, mode=mode):
            lines.append(f"--- Frame @ {timestamp:.1f}s ---")
            lines.append(frame_text)
        return PlainTextResponse("\n".join(lines))
    except Exception as e:
        logger.exception(f"ASCII viz failed: {e}")
        raise HTTPException(status_code=500, detail=str(e))


if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=8081, log_level="info")