STT-Assistant/tests/test_transcriber.py

"""Tests for the Whisper segment confidence gate."""

import numpy as np
import pytest

from stt_daemon import transcriber as transcriber_mod
from stt_daemon.transcriber import Transcriber, accept_segment


class _Seg:
    def __init__(self, text="hello", no_speech_prob=0.1, avg_logprob=-0.3):
        self.text = text
        self.no_speech_prob = no_speech_prob
        self.avg_logprob = avg_logprob


def test_accepts_confident_speech():
    assert accept_segment(_Seg(), 0.6, -1.0)


def test_rejects_high_no_speech_prob():
    assert not accept_segment(_Seg(no_speech_prob=0.9), 0.6, -1.0)


def test_rejects_low_avg_logprob():
    assert not accept_segment(_Seg(avg_logprob=-1.5), 0.6, -1.0)


def test_threshold_boundary_is_accepted():
    assert accept_segment(_Seg(no_speech_prob=0.6, avg_logprob=-1.0), 0.6, -1.0)


def test_missing_attributes_default_to_accept():
    class Bare:
        text = "x"

    assert accept_segment(Bare(), 0.6, -1.0)


class _Info:
    def __init__(self, language="en", language_probability=1.0):
        self.language = language
        self.language_probability = language_probability


class _FakeWhisperModel:
    """Records constructor/transcribe kwargs without loading a real model."""

    last = None

    def __init__(self, model, device, compute_type, cpu_threads=0):
        self.init_kwargs = {
            "model": model,
            "device": device,
            "compute_type": compute_type,
            "cpu_threads": cpu_threads,
        }
        self.transcribe_kwargs = None
        self.transcribe_calls = []
        self.segments_to_return = []
        self.detected_language = "en"
        self.detected_language_probability = 1.0
        _FakeWhisperModel.last = self

    def transcribe(self, audio, **kwargs):
        self.transcribe_kwargs = kwargs
        self.transcribe_calls.append(kwargs)
        return self.segments_to_return, _Info(
            self.detected_language, self.detected_language_probability
        )


@pytest.fixture
def fake_model(monkeypatch):
    monkeypatch.setattr(transcriber_mod, "WhisperModel", _FakeWhisperModel)
    return _FakeWhisperModel


def _audio():
    return np.zeros(16000, dtype=np.float32)


def test_cpu_threads_forwarded_to_model(fake_model):
    t = Transcriber(
        "small", device="cpu", compute_type="int8", language=None, cpu_threads=16
    )
    assert t.cpu_threads == 16
    assert fake_model.last.init_kwargs["cpu_threads"] == 16


def test_cpu_threads_default_is_auto(fake_model):
    Transcriber("small", device="cpu", compute_type="int8", language=None)
    assert fake_model.last.init_kwargs["cpu_threads"] == 0


def test_transcribe_disables_timestamps(fake_model):
    t = Transcriber("small", device="cpu", compute_type="int8", language=None)
    t.transcribe(_audio())
    kwargs = fake_model.last.transcribe_kwargs
    assert kwargs["without_timestamps"] is True
    assert kwargs["condition_on_previous_text"] is False


def test_auto_language_stays_none_for_bilingual(fake_model):
    """language=auto must reach the engine as None so es/en both keep working."""
    t = Transcriber("small", device="cpu", compute_type="int8", language=None)
    t.transcribe(_audio())
    assert fake_model.last.transcribe_kwargs["language"] is None


def test_explicit_language_is_preserved(fake_model):
    t = Transcriber("small", device="cpu", compute_type="int8", language="es")
    t.transcribe(_audio())
    assert fake_model.last.transcribe_kwargs["language"] == "es"


def test_sticky_pins_language_after_first_detection(fake_model):
    """With sticky strategy, language is detected once and pinned."""
    t = Transcriber(
        "small",
        device="cpu",
        compute_type="int8",
        language=None,
        language_strategy="sticky",
    )
    fake_model.last.detected_language = "es"
    fake_model.last.detected_language_probability = 0.95

    # First call: auto-detect (language=None sent)
    t.transcribe(_audio())
    assert fake_model.last.transcribe_calls[0]["language"] is None
    assert t._pinned_language == "es"
    assert t._pinned_language_probability == 0.95

    # Second call: uses pinned language (language="es" sent)
    t.transcribe(_audio())
    assert fake_model.last.transcribe_calls[1]["language"] == "es"


def test_sticky_reports_stored_detection_probability_when_pinned(fake_model):
    """When pinned, report the stored detection probability, not Whisper's 1.0."""
    t = Transcriber(
        "small",
        device="cpu",
        compute_type="int8",
        language=None,
        language_strategy="sticky",
    )
    fake_model.last.detected_language = "es"
    fake_model.last.detected_language_probability = 0.92

    # First call: auto-detect and pin
    result1 = t.transcribe(_audio())
    assert result1.language_probability == 0.92

    # Second call: pinned; Whisper returns 1.0 (explicit language), but we report stored
    fake_model.last.detected_language = "es"
    fake_model.last.detected_language_probability = 1.0
    fake_model.last.segments_to_return = [_Seg()]  # Provide a segment so pin isn't cleared
    result2 = t.transcribe(_audio())
    assert result2.language_probability == 0.92  # Stored, not 1.0


def test_sticky_clears_pin_when_best_logprob_too_low(fake_model):
    """Sticky clears the pin when the best segment's avg_logprob is below threshold."""
    t = Transcriber(
        "small",
        device="cpu",
        compute_type="int8",
        language=None,
        language_strategy="sticky",
        sticky_redetect_logprob=-0.8,
        avg_logprob_min=-1.0,
    )
    fake_model.last.detected_language = "en"
    fake_model.last.detected_language_probability = 0.98

    # Pin the language on first call with good segment
    fake_model.last.segments_to_return = [_Seg(avg_logprob=-0.5)]
    t.transcribe(_audio())
    assert t._pinned_language == "en"

    # Second call: return a segment with avg_logprob below threshold (-0.85 < -0.8)
    # but still accepted (above avg_logprob_min of -1.0)
    seg_low = _Seg(avg_logprob=-0.85)
    fake_model.last.segments_to_return = [seg_low]

    t.transcribe(_audio())
    # Pin should be cleared because max logprob is below sticky_redetect_logprob
    assert t._pinned_language is None


def test_sticky_clears_pin_when_no_segments_accepted(fake_model):
    """Sticky clears the pin when no segments pass the confidence gate."""
    t = Transcriber(
        "small",
        device="cpu",
        compute_type="int8",
        language=None,
        language_strategy="sticky",
    )
    fake_model.last.detected_language = "en"
    fake_model.last.detected_language_probability = 0.99

    # Pin the language
    t.transcribe(_audio())
    assert t._pinned_language == "en"

    # Second call: return a segment that fails the gate (high no_speech_prob)
    seg_bad = _Seg(no_speech_prob=0.9)
    fake_model.last.segments_to_return = [seg_bad]

    t.transcribe(_audio())
    # Pin should be cleared because no segments were accepted
    assert t._pinned_language is None


def test_static_strategy_never_pins(fake_model):
    """Static (default) strategy never pins; always sends language=None for auto-detect."""
    t = Transcriber(
        "small",
        device="cpu",
        compute_type="int8",
        language=None,
        language_strategy="static",
    )
    fake_model.last.detected_language = "es"

    # First call
    t.transcribe(_audio())
    assert fake_model.last.transcribe_calls[0]["language"] is None
    assert t._pinned_language is None

    # Second call
    t.transcribe(_audio())
    assert fake_model.last.transcribe_calls[1]["language"] is None
    assert t._pinned_language is None


def test_explicit_language_with_sticky_ignores_sticky(fake_model):
    """When language is explicit, sticky strategy has no effect; no pinning occurs."""
    t = Transcriber(
        "small",
        device="cpu",
        compute_type="int8",
        language="es",
        language_strategy="sticky",
    )
    fake_model.last.detected_language = "es"

    # First call: sends language="es" (explicit)
    t.transcribe(_audio())
    assert fake_model.last.transcribe_calls[0]["language"] == "es"
    assert t._pinned_language is None  # No pinning for explicit language

    # Second call: still sends language="es"
    t.transcribe(_audio())
    assert fake_model.last.transcribe_calls[1]["language"] == "es"
    assert t._pinned_language is None


def test_sticky_tolerates_missing_avg_logprob_attribute(fake_model):
    """Sticky pin clearing uses getattr(..., 0.0) to tolerate missing avg_logprob."""
    t = Transcriber(
        "small",
        device="cpu",
        compute_type="int8",
        language=None,
        language_strategy="sticky",
        sticky_redetect_logprob=-0.8,
    )
    fake_model.last.detected_language = "en"
    fake_model.last.detected_language_probability = 0.99

    # Pin the language with a normal segment
    fake_model.last.segments_to_return = [_Seg(avg_logprob=-0.5)]
    t.transcribe(_audio())
    assert t._pinned_language == "en"

    # Second call: return a segment with no avg_logprob attribute (only text/no_speech_prob)
    class SegNoLogprob:
        text = "hello"
        no_speech_prob = 0.1
    fake_model.last.segments_to_return = [SegNoLogprob()]

    # Should not raise; defaulted 0.0 is above threshold -0.8, so pin is RETAINED
    t.transcribe(_audio())
    assert t._pinned_language == "en"  # Pin retained because 0.0 >= -0.8