qa_flow.py

"""Server-side Q&A flow helpers for the Gradio Ask interaction."""
from __future__ import annotations

import logging
import os
import threading
import time
import uuid
from pathlib import Path
from typing import Callable, Iterable, Protocol

logger = logging.getLogger(__name__)
DEFAULT_SAMPLE_RATE = 24000

_RECENT_AUDIO_LOCK = threading.Lock()
_RECENT_AUDIO_SUBMISSIONS: dict[str, float] = {}
_DUPLICATE_WINDOW_SEC = 30.0
_GENERATED_QA_MAX_FILES = int(os.environ.get("MOMSVOICE_QA_AUDIO_MAX_FILES", "30"))


class AudioWriter(Protocol):
    def __call__(self, path: Path, waveform, sample_rate: int) -> None:
        ...


def _normalize_audio_path(audio_path: str | Path | None) -> str:
    if not audio_path:
        return ""
    try:
        path = Path(audio_path).resolve()
        stat = path.stat()
        return f"{path}:{stat.st_size}:{stat.st_mtime_ns}"
    except OSError:
        return str(audio_path)


def claim_audio_submission(audio_path: str | Path | None) -> bool:
    """Return False when the same recorded audio was already claimed recently."""
    normalized = _normalize_audio_path(audio_path)
    if not normalized:
        return True

    now = time.monotonic()
    cutoff = now - _DUPLICATE_WINDOW_SEC
    with _RECENT_AUDIO_LOCK:
        stale = [
            key for key, seen_at in _RECENT_AUDIO_SUBMISSIONS.items()
            if seen_at < cutoff
        ]
        for key in stale:
            _RECENT_AUDIO_SUBMISSIONS.pop(key, None)

        if normalized in _RECENT_AUDIO_SUBMISSIONS:
            return False

        _RECENT_AUDIO_SUBMISSIONS[normalized] = now
        return True


def release_audio_submission(audio_path: str | Path | None) -> None:
    """Allow a recording to be submitted again, used after failed generation."""
    normalized = _normalize_audio_path(audio_path)
    if not normalized:
        return
    with _RECENT_AUDIO_LOCK:
        _RECENT_AUDIO_SUBMISSIONS.pop(normalized, None)


def _default_audio_writer(path: Path, waveform, sample_rate: int) -> None:
    import soundfile as sf
    import numpy as np

    wav = np.asarray(waveform, dtype=np.float32)
    # Clip to [-1, 1] before writing as PCM_16 for universal compatibility
    wav = np.clip(wav, -1.0, 1.0)
    sf.write(str(path), wav, sample_rate, subtype='PCM_16')


def _prune_generated_answers(output_dir: Path) -> None:
    try:
        files = sorted(
            output_dir.glob("qa_answer_*.wav"),
            key=lambda path: path.stat().st_mtime,
            reverse=True,
        )
    except OSError:
        return

    for stale in files[_GENERATED_QA_MAX_FILES:]:
        try:
            stale.unlink()
        except OSError:
            continue


def build_qa_response(
    *,
    question_text: str | None,
    question_audio_path: str | Path | None,
    paragraphs: Iterable[str] | None,
    output_dir: Path,
    answer_fn: Callable[..., tuple[str, object | None, int]],
    audio_writer: AudioWriter = _default_audio_writer,
    max_new_tokens: int = 512,
) -> dict[str, object]:
    """Run Q&A and return display-ready data without importing Gradio."""
    q_txt = (question_text or "").strip()
    audio_path = str(question_audio_path) if question_audio_path else None
    has_audio = bool(audio_path)

    if not q_txt and not has_audio:
        return {
            "ok": False,
            "answer_text": "Please type a question or record one with the microphone.",
            "display_question": "",
            "audio_path": None,
            "error": None,
        }

    story_context = "\n\n".join(paragraphs or [])
    error = None
    try:
        answer_text, waveform, sr = answer_fn(
            question_audio_path=audio_path if has_audio else None,
            question_text=q_txt if q_txt and not has_audio else None,
            story_context=story_context,
            max_new_tokens=max_new_tokens,
        )
        if not answer_text:
            answer_text = "Hmm, I'm not sure about that! Let's keep listening to find out."
    except Exception as exc:
        logger.exception("LFM Q&A failed: %s", exc)
        error = type(exc).__name__
        answer_text = (
            "Oops, I couldn't think of an answer right now. Let's keep reading! "
            f"({type(exc).__name__})"
        )
        waveform = None
        sr = DEFAULT_SAMPLE_RATE

    answer_audio_path = None
    if waveform is not None:
        output_dir.mkdir(parents=True, exist_ok=True)
        answer_audio_path = output_dir / f"qa_answer_{uuid.uuid4().hex[:8]}.wav"
        try:
            audio_writer(answer_audio_path, waveform, sr)
            _prune_generated_answers(output_dir)
        except Exception as exc:
            logger.exception("Failed to write answer audio: %s", exc)
            error = error or type(exc).__name__
            answer_audio_path = None

    return {
        "ok": True,
        "answer_text": answer_text,
        "display_question": q_txt or "(audio question)",
        "audio_path": answer_audio_path,
        "error": error,
    }