speaker_sample.py

"""
Speaker sample verification and storage utilities.

Provides functions for:
- Verifying and transcribing speech samples
- Downloading samples from GCS
- Deleting samples from GCS
"""

from typing import Any, Dict, List, Optional, Tuple, cast

from config.stt_provider_policy import normalized_stt_language
from utils.executors import sync_executor, run_blocking
from utils.other.storage import delete_speech_profile_blob, download_speech_profile_bytes
from utils.stt.pre_recorded import prerecorded_from_bytes as deepgram_prerecorded_from_bytes
from utils.text_utils import compute_text_containment

MIN_WORDS = 5
MIN_CONTAINMENT = 0.9
MIN_DOMINANT_SPEAKER_RATIO = 0.7

# Chinese, Japanese and Thai don't put spaces between words, so `.split()` treats a whole
# segment-granular entry (Parakeet, Modulate) as a single "word" no matter how long it is.
# Counting characters instead keeps the same gate meaningful for these languages. 12 chars is
# a rough parity with MIN_WORDS=5 English words' worth of speech content.
NON_SPACE_DELIMITED_LANGUAGES = frozenset({'zh', 'ja', 'th'})
MIN_CJK_CHARS = 12


async def verify_and_transcribe_sample(
    audio_bytes: bytes,
    sample_rate: int,
    expected_text: Optional[str] = None,
    language: Optional[str] = None,
) -> Tuple[Optional[str], bool, str]:
    """
    Transcribe audio and verify quality using PR #4291 rules.

    Checks:
    1. Transcription has at least MIN_WORDS words (MIN_CJK_CHARS characters for
       non-space-delimited languages, see NON_SPACE_DELIMITED_LANGUAGES)
    2. Dominant speaker accounts for >= MIN_DOMINANT_SPEAKER_RATIO of words (via diarization)
    3. Transcribed text has >= MIN_CONTAINMENT containment in expected text (if provided)

    Args:
        audio_bytes: WAV format audio bytes
        sample_rate: Audio sample rate in Hz
        expected_text: Expected text from the segment for comparison (optional)
        language: Language of the sample, so it's transcribed in the right language instead of
            silently defaulting to English (optional)

    Returns:
        (transcript, is_valid, reason): Tuple of (str or None, bool, str)
        - reason "transcription_failed" indicates transient API error (sample should be kept)
        - other reasons indicate quality issues (sample may be dropped)
    """
    try:
        raw_words = await run_blocking(
            sync_executor,
            cast(Any, deepgram_prerecorded_from_bytes),
            audio_bytes,
            sample_rate,
            True,
            language=language,
        )
    except RuntimeError as e:
        # Transient transcription failure - distinguish from quality issues
        return None, False, f"transcription_failed: {e}"

    # deepgram_prerecorded_from_bytes returns List[dict] or (when return_language=True) Tuple[List[dict], str].
    # return_language defaults to False, so the runtime value is always the list; narrow for the type system.
    if isinstance(raw_words, tuple):
        raw_words = cast(Any, raw_words[0])
    words: List[Dict[str, Any]] = cast(List[Dict[str, Any]], raw_words)

    # Count words in the text, not entries in the list. A transcriber is free to return either
    # granularity: Deepgram emits one entry per word, so the two are the same number there, but
    # parakeet emits one entry per SEGMENT with the whole utterance inside. Measured on a 24.9s
    # sample: parakeet returns segments=1 and no word field, so the entry count read 1 and every
    # sample was rejected as `insufficient_words` no matter what it contained. The same miscount
    # made the multi-speaker guard inert rather than strict — one entry means ratio 1.0, so a
    # sample carrying two voices passed. Both are unchanged for a word-granular provider.
    #
    # For CJK/Thai, `.split()` undercounts even on a word-granular provider, since those
    # languages don't use whitespace between words at all — count characters instead.
    count_chars = normalized_stt_language(language) in NON_SPACE_DELIMITED_LANGUAGES

    def _words_in(entry: Dict[str, Any]) -> int:
        text = entry.get('text') or ''
        return len(text.strip()) if count_chars else len(text.split())

    total_words = sum(_words_in(word) for word in words)
    min_required = MIN_CJK_CHARS if count_chars else MIN_WORDS

    if total_words < min_required:
        return None, False, f"insufficient_words: {total_words}/{min_required}"

    speaker_counts: Dict[str, int] = {}
    for word in words:
        speaker = word.get('speaker', 'SPEAKER_00')
        speaker_counts[speaker] = speaker_counts.get(speaker, 0) + _words_in(word)

    dominant_count = max(speaker_counts.values()) if speaker_counts else 0
    dominant_ratio = dominant_count / total_words if total_words > 0 else 0

    if dominant_ratio < MIN_DOMINANT_SPEAKER_RATIO:
        return None, False, f"multi_speaker: ratio={dominant_ratio:.2f}"

    transcript = ' '.join(w.get('text', '') for w in words)

    if expected_text:
        containment = compute_text_containment(transcript, expected_text)
        if containment < MIN_CONTAINMENT:
            return transcript, False, f"text_mismatch: containment={containment:.2f}"

    return transcript, True, "ok"


def download_sample_audio(sample_path: str) -> bytes:
    """
    Download speech sample audio from GCS.

    Args:
        sample_path: GCS path to the sample (e.g., '{uid}/people_profiles/{person_id}/{filename}.wav')

    Returns:
        Audio bytes (WAV format)

    Raises:
        NotFound: If the sample doesn't exist
    """
    return download_speech_profile_bytes(sample_path)


def delete_sample_from_storage(sample_path: str) -> bool:
    """
    Delete speech sample from GCS.

    Args:
        sample_path: GCS path to the sample

    Returns:
        True if deleted, False if not found
    """
    return delete_speech_profile_blob(sample_path)