app/main.py

"""
nlp-asr-kit Python API
FastAPI backend: Whisper ASR (faster-whisper) + Qwen2.5-0.5B-Instruct (llama-cpp GGUF)

Simple interface: send audio + a prompt, get back a single text output.

Run: uvicorn app.main:app --host 0.0.0.0 --port 8000
"""

import tempfile
import os
from contextlib import asynccontextmanager
from fastapi import FastAPI, UploadFile, File, Form, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
import uvicorn

from app.pipeline import NLPASRPipeline, DEFAULT_PROMPT, ASR_MODEL, LLM_MODEL_NAME

# ── Global pipeline instance ──────────────────────────────────────────────────
pipeline: NLPASRPipeline = None

@asynccontextmanager
async def lifespan(app: FastAPI):
    global pipeline
    print("Loading models... (this may take a minute on first run)")
    pipeline = NLPASRPipeline()
    pipeline.load_models()
    print("✓ Models ready!")
    yield
    print("Shutting down...")

app = FastAPI(
    title="nlp-asr-kit API",
    description=(
        "Local ASR (Whisper small) + LLM (Qwen2.5-0.5B-Instruct Q4 GGUF) pipeline. "
        "Send audio plus a prompt, get back a single text output. "
        "Use the prompt to ask for a summary, translation, action items, "
        "sentiment, or any custom text task."
    ),
    version="5.0.0",
    lifespan=lifespan
)

# ── CORS ──────────────────────────────────────────────────────────────────────
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# ── Routes ────────────────────────────────────────────────────────────────────

@app.get("/health")
def health():
    return {
        "status": "ok",
        "models_loaded": pipeline is not None and pipeline.ready,
        "asr_model": f"whisper-{ASR_MODEL} (faster-whisper, int8)",
        "llm_model": LLM_MODEL_NAME,
        "default_prompt": DEFAULT_PROMPT,
    }

@app.post("/api/process")
async def process(
    audio: UploadFile = File(...),
    prompt: str = Form(default=""),
    max_tokens: int = Form(default=384),
):
    """
    Send audio + a prompt, get back a single text output from the LLM.

    - audio: audio file (webm, wav, mp3, m4a, ogg)
    - prompt: instruction telling the LLM what to do with the transcript
              (e.g. "Summarize this in 2 sentences.", "Translate to French.",
              "Extract action items as a numbered list.")
              If empty, defaults to cleaning up + summarizing the transcript.
    - max_tokens: max tokens the LLM should generate (default 384)
    """
    if not pipeline or not pipeline.ready:
        raise HTTPException(status_code=503, detail="Models not loaded yet")

    audio_bytes = await audio.read()
    suffix = os.path.splitext(audio.filename or "audio.webm")[1] or ".webm"

    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(audio_bytes)
        tmp_path = tmp.name

    try:
        output = pipeline.process(tmp_path, prompt=prompt, max_tokens=max_tokens)
        return JSONResponse({"output": output})
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        os.unlink(tmp_path)

@app.post("/api/transcribe")
async def transcribe(audio: UploadFile = File(...)):
    """Transcribe audio only — no LLM step. Returns the raw transcript."""
    if not pipeline or not pipeline.ready:
        raise HTTPException(status_code=503, detail="Models not loaded yet")

    audio_bytes = await audio.read()
    suffix = os.path.splitext(audio.filename or "audio.webm")[1] or ".webm"

    with tempfile.NamedTemporaryFile(suffix=suffix, delete=False) as tmp:
        tmp.write(audio_bytes)
        tmp_path = tmp.name

    try:
        transcript = pipeline.transcribe_only(tmp_path)
        return JSONResponse({"transcript": transcript})
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))
    finally:
        os.unlink(tmp_path)

if __name__ == "__main__":
    uvicorn.run("app.main:app", host="0.0.0.0", port=8000, reload=False)