CRM_Montzar/apps/stt/main.py

"""Speech-to-text sidecar — local Faster-Whisper (Pipecat stack), no GapGPT."""

from __future__ import annotations

import io
import logging
import os
import shutil
import subprocess
import tempfile
import threading
import wave
from contextlib import asynccontextmanager
from pathlib import Path

import numpy as np
from dotenv import load_dotenv
from fastapi import FastAPI, File, Header, HTTPException, UploadFile
from fastapi.responses import JSONResponse
from faster_whisper import WhisperModel

load_dotenv(Path(__file__).resolve().parent / ".env")

logger = logging.getLogger("mento-stt")
logging.basicConfig(level=logging.INFO)

MAX_BYTES = 8 * 1024 * 1024
WAV_MAGIC = (b"RIFF", b"WAVE")

_model: WhisperModel | None = None
_model_lock = threading.Lock()
_model_error: str | None = None


def _env(name: str, default: str = "") -> str:
    return (os.getenv(name) or default).strip()


def _require_token(header_token: str | None) -> None:
    expected = _env("STT_INTERNAL_TOKEN")
    if not expected:
        return
    if (header_token or "").strip() != expected:
        raise HTTPException(status_code=401, detail="توکن سرویس تبدیل صوت نامعتبر است")


def _model_name() -> str:
    """Map env / legacy cloud names to Faster-Whisper local sizes."""
    raw = (_env("STT_MODEL", "small") or "small").lower()
    legacy = {
        "whisper-1": "small",
        "whisper": "small",
        "gpt-4o-transcribe": "small",
        "gpt-4o-mini-transcribe": "small",
        # tiny/base are too weak for Persian dictation
        "tiny": "small",
        "tiny.en": "small",
        "base": "small",
        "base.en": "small",
    }
    if raw in legacy:
        return legacy[raw]
    allowed = {
        "small",
        "small.en",
        "medium",
        "medium.en",
        "large",
        "large-v1",
        "large-v2",
        "large-v3",
        "distil-large-v2",
        "distil-medium.en",
        "distil-large-v3",
        "large-v3-turbo",
    }
    return raw if raw in allowed else "small"


def _get_model() -> WhisperModel:
    """Load once, thread-safe. Only cache after success."""
    global _model, _model_error
    if _model is not None:
        return _model

    with _model_lock:
        if _model is not None:
            return _model
        if _model_error:
            raise HTTPException(
                status_code=503,
                detail=f"مدل Whisper آماده نیست: {_model_error}",
            )

        name = _model_name()
        device = _env("STT_DEVICE", "cpu") or "cpu"
        compute = _env("STT_COMPUTE_TYPE", "int8") or "int8"
        logger.info("Loading Faster-Whisper model=%s device=%s compute=%s", name, device, compute)
        try:
            loaded = WhisperModel(name, device=device, compute_type=compute)
        except Exception as exc:
            _model_error = str(exc)
            logger.exception("Whisper model load failed")
            raise HTTPException(
                status_code=503,
                detail=(
                    "بارگذاری مدل Whisper محلی ناموفق بود. "
                    "اگر اولین بار است، دانلود طول می‌کشد؛ یا STT_MODEL=tiny/base بگذارید. "
                    f"جزئیات: {exc}"
                ),
            ) from exc

        _model = loaded
        _model_error = None
        logger.info("Whisper model ready: %s", name)
        return _model


def _is_wav(data: bytes) -> bool:
    return len(data) >= 12 and data[:4] == WAV_MAGIC[0] and data[8:12] == WAV_MAGIC[1]


def _to_wav(data: bytes, filename: str) -> bytes:
    if _is_wav(data):
        return data

    ffmpeg = shutil.which("ffmpeg")
    if not ffmpeg:
        raise HTTPException(
            status_code=400,
            detail="فرمت صوت پشتیبانی نشد. فایل WAV ارسال کنید یا ffmpeg را نصب کنید.",
        )

    suffix = Path(filename or "audio.webm").suffix or ".webm"
    with tempfile.TemporaryDirectory() as tmp:
        src = Path(tmp) / f"in{suffix}"
        dst = Path(tmp) / "out.wav"
        src.write_bytes(data)
        result = subprocess.run(
            [
                ffmpeg,
                "-y",
                "-i",
                str(src),
                "-ac",
                "1",
                "-ar",
                "16000",
                "-f",
                "wav",
                str(dst),
            ],
            capture_output=True,
            check=False,
        )
        if result.returncode != 0 or not dst.exists():
            raise HTTPException(status_code=400, detail="تبدیل فایل صوتی ناموفق بود")
        return dst.read_bytes()


def _wav_to_float32(wav_bytes: bytes) -> tuple[np.ndarray, int]:
    with wave.open(io.BytesIO(wav_bytes), "rb") as wf:
        channels = wf.getnchannels()
        sample_width = wf.getsampwidth()
        rate = wf.getframerate()
        frames = wf.readframes(wf.getnframes())

    if sample_width == 2:
        audio = np.frombuffer(frames, dtype=np.int16).astype(np.float32) / 32768.0
    elif sample_width == 4:
        audio = np.frombuffer(frames, dtype=np.int32).astype(np.float32) / 2147483648.0
    else:
        raise HTTPException(status_code=400, detail="فرمت نمونهٔ صوتی پشتیبانی نمی‌شود")

    if channels > 1:
        audio = audio.reshape(-1, channels).mean(axis=1)

    return audio, rate


def _resample_audio(audio: np.ndarray, src_rate: int, dst_rate: int = 16000) -> np.ndarray:
    if src_rate == dst_rate or audio.size == 0:
        return audio.astype(np.float32, copy=False)
    duration = audio.shape[0] / float(src_rate)
    new_len = max(1, int(round(duration * dst_rate)))
    x_old = np.linspace(0.0, 1.0, num=audio.shape[0], endpoint=False)
    x_new = np.linspace(0.0, 1.0, num=new_len, endpoint=False)
    return np.interp(x_new, x_old, audio).astype(np.float32)


def _ensure_16k(wav_bytes: bytes) -> np.ndarray:
    audio, rate = _wav_to_float32(wav_bytes)
    if rate == 16000:
        return audio

    # Prefer ffmpeg when available; otherwise pure-numpy resample (browser often sends 44.1/48 kHz)
    ffmpeg = shutil.which("ffmpeg")
    if ffmpeg:
        with tempfile.TemporaryDirectory() as tmp:
            src = Path(tmp) / "in.wav"
            dst = Path(tmp) / "out.wav"
            src.write_bytes(wav_bytes)
            result = subprocess.run(
                [ffmpeg, "-y", "-i", str(src), "-ac", "1", "-ar", "16000", "-f", "wav", str(dst)],
                capture_output=True,
                check=False,
            )
            if result.returncode == 0 and dst.exists():
                audio, _ = _wav_to_float32(dst.read_bytes())
                return audio

    return _resample_audio(audio, rate, 16000)


def _normalize(audio: np.ndarray) -> np.ndarray:
    if audio.size == 0:
        return audio
    peak = float(np.max(np.abs(audio)))
    if peak < 1e-4:
        return audio
    # Bring quiet mic levels up without clipping
    gain = min(8.0, 0.9 / peak)
    if gain <= 1.05:
        return audio
    return np.clip(audio * gain, -1.0, 1.0).astype(np.float32)


def _transcribe_local(wav_bytes: bytes) -> str:
    model = _get_model()
    audio = _normalize(_ensure_16k(wav_bytes))
    initial_prompt = _env(
        "STT_PROMPT",
        "این یک یادداشت فارسی است. لطفاً این مورد را بررسی کنید.",
    ) or None

    segments, _info = model.transcribe(
        audio,
        language="fa",
        task="transcribe",
        initial_prompt=initial_prompt,
        beam_size=5,
        best_of=5,
        temperature=0.0,
        condition_on_previous_text=False,
        vad_filter=False,
        no_speech_threshold=0.45,
    )
    parts = [seg.text.strip() for seg in segments if getattr(seg, "text", None)]
    return " ".join(p for p in parts if p).strip()


def _friendly_stt_error(exc: BaseException) -> tuple[int, str]:
    text = str(exc)
    lowered = text.lower()
    if "timeout" in lowered or "timed out" in lowered:
        return 504, "زمان تبدیل صوت به پایان رسید. دوباره با صدای کوتاه‌تر تلاش کنید."
    if "memory" in lowered or "oom" in lowered:
        return 503, "حافظهٔ سرور برای مدل Whisper کافی نیست. STT_MODEL=tiny یا base بگذارید."
    if len(text) > 220:
        return 502, "تبدیل صوت به متن محلی ناموفق بود. لاگ mento-stt را بررسی کنید."
    return 502, f"تبدیل صوت به متن ناموفق بود: {text}"


@asynccontextmanager
async def lifespan(_app: FastAPI):
    # Warm model in background thread so first user request is faster
    def warm() -> None:
        try:
            _get_model()
        except Exception:
            logger.exception("Startup model warm failed (will retry on first request)")

    threading.Thread(target=warm, name="stt-warm", daemon=True).start()
    yield


app = FastAPI(title="Mento STT", docs_url=None, redoc_url=None, lifespan=lifespan)


@app.get("/health")
async def health() -> dict[str, object]:
    return {
        "ok": True,
        "engine": "faster-whisper",
        "service": "local",
        "mode": "local",
        "model": _model_name(),
        "device": _env("STT_DEVICE", "cpu") or "cpu",
        "ready": _model is not None,
        "cloud": False,
        "gapgpt": False,
        "lastError": _model_error,
    }


@app.post("/transcribe")
async def transcribe(
    audio: UploadFile = File(...),
    x_stt_token: str | None = Header(default=None, alias="X-STT-Token"),
) -> JSONResponse:
    _require_token(x_stt_token)

    data = await audio.read()
    if not data:
        raise HTTPException(status_code=400, detail="فایل صوتی خالی است")
    if len(data) > MAX_BYTES:
        raise HTTPException(status_code=400, detail="حجم فایل صوتی بیش از حد مجاز است")

    wav = _to_wav(data, audio.filename or "audio.wav")

    try:
        text = _transcribe_local(wav)
    except HTTPException:
        raise
    except Exception as exc:
        status, detail = _friendly_stt_error(exc)
        raise HTTPException(status_code=status, detail=detail) from exc

    return JSONResponse({"text": text})