app.py

import os
import logging
import tempfile
import traceback
from typing import Optional, List, Union
from contextlib import asynccontextmanager

from fastapi import FastAPI, UploadFile, File, Form, HTTPException, status
from fastapi.responses import JSONResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
from faster_whisper import WhisperModel

# --- Configuration & Logging ---
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger("ASR_API")

# Model Settings
# Default to 'base' for best balance of speed and accuracy on CPU.
# 'tiny' is faster but less accurate. 'small' is slower but more accurate.
MODEL_SIZE = os.getenv("WHISPER_MODEL_SIZE", "base")
# int8 is recommended for CPU to improve speed
COMPUTE_TYPE = os.getenv("WHISPER_COMPUTE_TYPE", "int8") 
DEVICE = "cpu" # Enforcing CPU usage as requested

# Global model instance
model_instance = None

# --- Lifespan Management ---
@asynccontextmanager
async def lifespan(app: FastAPI):
    # Startup: Load the model into memory
    global model_instance
    logger.info(f"Loading Whisper Model (Size: {MODEL_SIZE}, Device: {DEVICE}, Compute: {COMPUTE_TYPE})...")
    try:
        # faster-whisper handles model downloading automatically if not present
        model_instance = WhisperModel(
            MODEL_SIZE, 
            device=DEVICE, 
            compute_type=COMPUTE_TYPE,
            download_root=".cache/models" # Cache in space directory
        )
        logger.info("Model loaded successfully.")
    except Exception as e:
        logger.error(f"Failed to load model: {e}")
        raise RuntimeError("Failed to initialize ASR model")
    
    yield
    
    # Shutdown: Cleanup (if necessary)
    logger.info("Shutting down application...")
    del model_instance

app = FastAPI(
    title="Enterprise ASR API",
    description="High-performance Speech-to-Text API using Faster Whisper on CPU",
    version="1.0.0",
    lifespan=lifespan
)

# Enable CORS for all origins
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # Allow all origins
    allow_credentials=True,
    allow_methods=["*"],  # Allow all methods (GET, POST, etc.)
    allow_headers=["*"],  # Allow all headers
)

# --- Pydantic Schemas ---
class TranscriptionSegment(BaseModel):
    id: int
    seek: int
    start: float
    end: float
    text: str
    tokens: List[int]
    temperature: float
    avg_logprob: float
    compression_ratio: float
    no_speech_prob: float

    class Config:
        # Ensure fields like tokens are serialized correctly
        json_schema_extra = {
            "example": {
                "id": 0,
                "seek": 0,
                "start": 0.0,
                "end": 2.0,
                "text": "Hello world",
                "tokens": [50364, 1917, 1002, 50664],
                "temperature": 0.0,
                "avg_logprob": -0.5,
                "compression_ratio": 1.0,
                "no_speech_prob": 0.1
            }
        }

class TranscriptionResponse(BaseModel):
    text: str = Field(..., description="The full transcribed text")
    language: str = Field(..., description="Detected language code")
    segments: List[TranscriptionSegment] = Field(default_factory=list, description="Detailed segments of transcription")
    duration: float = Field(..., description="Duration of audio in seconds")

# --- Helper Functions ---
def clean_up_temp_file(filepath: str):
    """Safely remove a temporary file."""
    try:
        if os.path.exists(filepath):
            os.remove(filepath)
            logger.debug(f"Deleted temporary file: {filepath}")
    except Exception as e:
        logger.warning(f"Failed to delete temporary file {filepath}: {e}")

# --- Endpoints ---

@app.get("/health", tags=["System"])
async def health_check():
    """Check if the API is running and the model is loaded."""
    if model_instance is None:
        raise HTTPException(status_code=503, detail="Model not loaded yet")
    return {"status": "healthy", "model": MODEL_SIZE, "device": DEVICE}

@app.post(
    "/v1/transcribe", 
    response_model=TranscriptionResponse, 
    tags=["ASR Operations"],
    summary="Transcribe or Translate Audio"
)
async def transcribe_audio(
    file: UploadFile = File(..., description="Audio file (mp3, wav, m4a, etc.)"),
    task: str = Form("transcribe", description="Task: 'transcribe' or 'translate'"),
    language: Optional[str] = Form(None, description="Language code (e.g., 'en', 'es'). Auto-detect if None."),
    initial_prompt: Optional[str] = Form(None, description="Optional initial text to guide the model."),
    vad_filter: bool = Form(True, description="Enable Voice Activity Detection (VAD) to silence gaps."),
    word_timestamps: bool = Form(False, description="Extract word-level timestamps."),
    temperature: float = Form(0.0, description="Sampling temperature (0.0 for greedy).")
):
    """
    Endpoint to process audio files for transcription or translation.
    
    - **file**: The audio file to process.
    - **task**: Use 'transcribe' for same-language text, 'translate' for English.
    - **language**: Specifying a language improves speed and accuracy.
    - **vad_filter**: Strongly recommended for long audio to cut out silence.
    """
    
    if model_instance is None:
        raise HTTPException(status_code=503, detail="Model is initializing or failed to load.")

    # 1. Validate input file extension (basic check)
    filename = file.filename or "audio.blob"
    logger.info(f"Received request: File='{filename}', Task='{task}', Lang='{language}'")

    # Create a temporary file to store the upload
    temp_filepath = None
    try:
        # We write the uploaded file to a temp disk because faster-whisper expects a file path
        # and handles loading optimized chunks from disk better than loading full byte-strings in memory
        with tempfile.NamedTemporaryFile(delete=False, suffix=os.path.splitext(filename)[1]) as tmp:
            content = await file.read()
            tmp.write(content)
            temp_filepath = tmp.name

        # 2. Perform Inference
        # We run the blocking CPU task in a thread pool to avoid blocking the async event loop
        segments_generator, info = await run_inference(
            temp_filepath, 
            task, 
            language, 
            initial_prompt, 
            vad_filter, 
            word_timestamps, 
            temperature
        )

        # 3. Format Results
        full_text = ""
        segment_objects = []
        
        for segment in segments_generator:
            full_text += segment.text
            segment_objects.append(TranscriptionSegment(
                id=segment.id,
                seek=segment.seek,
                start=segment.start,
                end=segment.end,
                text=segment.text.strip(),
                tokens=segment.tokens,
                temperature=segment.temperature,
                avg_logprob=segment.avg_logprob,
                compression_ratio=segment.compression_ratio,
                no_speech_prob=segment.no_speech_prob
            ))

        return TranscriptionResponse(
            text=full_text.strip(),
            language=info.language,
            duration=info.duration,
            segments=segment_objects
        )

    except HTTPException:
        raise
    except Exception as e:
        logger.error(f"Error during processing: {str(e)}")
        logger.error(traceback.format_exc())
        raise HTTPException(
            status_code=500, 
            detail=f"Internal Server Error during processing: {str(e)}"
        )
    finally:
        # 4. Cleanup
        if temp_filepath:
            clean_up_temp_file(temp_filepath)

async def run_inference(
    filepath: str, 
    task: str, 
    language: Optional[str], 
    initial_prompt: Optional[str], 
    vad_filter: bool, 
    word_timestamps: bool, 
    temperature: float
):
    """
    Wrapper to run the synchronous faster-whisper model in a thread pool executor.
    This prevents the API from freezing while the CPU crunches numbers.
    """
    def _compute():
        return model_instance.transcribe(
            audio=filepath,
            task=task,
            language=language,
            initial_prompt=initial_prompt,
            vad_filter=vad_filter,
            word_timestamps=word_timestamps,
            temperature=temperature
        )
    
    # Run in a separate thread
    from fastapi.concurrency import run_in_threadpool
    return await run_in_threadpool(_compute)