STT-Assistant/tests/test_pywhispercpp_transcriber.py

"""Tests for the pywhispercpp transcriber backend."""

import json
import logging
import math

import numpy as np
import pytest

from stt_daemon.pywhispercpp_transcriber import PyWhisperCppTranscriber, accept_segment


class _Seg:
    def __init__(self, text="hello", probability=0.9):
        self.text = text
        self.probability = probability


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

    last = None

    def __init__(self, model, **kwargs):
        self.init_kwargs = {"model": model, **kwargs}
        self.transcribe_calls = []
        self.auto_detect_calls = 0
        self.segments_to_return = []
        self.detected_language = "en"
        self.detected_language_probability: float | np.float32 = 1.0
        _FakeModel.last = self

    def auto_detect_language(self, audio):
        self.auto_detect_calls += 1
        return (self.detected_language, self.detected_language_probability), {}

    def transcribe(self, audio, **kwargs):
        self.transcribe_calls.append(kwargs)
        return self.segments_to_return


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


def _transcriber(**kwargs):
    kwargs.setdefault("model", "small")
    kwargs.setdefault("language", None)
    kwargs.setdefault("model_class", _FakeModel)
    return PyWhisperCppTranscriber(**kwargs)


# --- accept_segment gate ---------------------------------------------------


def test_accepts_probability_above_floor():
    assert accept_segment(_Seg(probability=0.9), -1.0)


def test_rejects_probability_below_floor():
    assert not accept_segment(_Seg(probability=math.exp(-1.0) - 0.01), -1.0)


def test_probability_at_floor_boundary_is_accepted():
    assert accept_segment(_Seg(probability=math.exp(-1.0)), -1.0)


def test_nan_probability_is_accepted():
    assert accept_segment(_Seg(probability=math.nan), -1.0)


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

    assert accept_segment(Bare(), -1.0)


# --- constructor wiring ----------------------------------------------------


def test_cpu_threads_forwarded_as_n_threads():
    _transcriber(cpu_threads=16)
    assert _FakeModel.last.init_kwargs["n_threads"] == 16


def test_cpu_threads_default_is_auto():
    _transcriber()
    # Omitted so pywhispercpp picks its own thread count (it rejects None).
    assert "n_threads" not in _FakeModel.last.init_kwargs


def test_language_omitted_from_model_constructor():
    """pywhispercpp rejects 'auto' ("unknown language"); language is set per call."""
    _transcriber()
    assert "language" not in _FakeModel.last.init_kwargs


def test_beam_size_five_enables_beam_search():
    _transcriber(beam_size=5)
    kwargs = _FakeModel.last.init_kwargs
    assert kwargs["params_sampling_strategy"] == 1
    assert kwargs["beam_search"] == {"beam_size": 5, "patience": -1.0}


def test_beam_size_one_uses_greedy():
    _transcriber(beam_size=1)
    kwargs = _FakeModel.last.init_kwargs
    assert kwargs["params_sampling_strategy"] == 0
    # Omitted so pywhispercpp keeps its default (it rejects an empty dict).
    assert "beam_search" not in kwargs


def test_vad_filter_warns_and_is_ignored(caplog):
    with caplog.at_level(logging.WARNING):
        _transcriber(vad_filter=True)
    assert "vad_filter" in caplog.text
    assert _FakeModel.last.init_kwargs["model"] == "small"


def test_close_is_noop():
    assert _transcriber().close() is None


# --- transcribe: language handling -----------------------------------------


def test_forced_language_skips_detection():
    t = _transcriber(language="es")
    result = t.transcribe(_audio())
    assert _FakeModel.last.auto_detect_calls == 0
    assert _FakeModel.last.transcribe_calls[0]["language"] == "es"
    assert result.language == "es"
    assert result.language_probability == 1.0


def test_auto_detect_propagates_language_and_probability():
    t = _transcriber(language=None)
    _FakeModel.last.detected_language = "es"
    _FakeModel.last.detected_language_probability = 0.93

    result = t.transcribe(_audio())

    assert _FakeModel.last.auto_detect_calls == 1
    assert _FakeModel.last.transcribe_calls[0]["language"] == "es"
    assert result.language == "es"
    assert result.language_probability == 0.93


def test_auto_detect_probability_is_builtin_float():
    """np.float32 from pywhispercpp is not JSON-serializable by the event bus."""
    t = _transcriber(language=None)
    _FakeModel.last.detected_language = "es"
    _FakeModel.last.detected_language_probability = np.float32(0.93)

    result = t.transcribe(_audio())

    assert type(result.language) is str
    assert type(result.language_probability) is float
    json.dumps({"lang": result.language, "lang_prob": result.language_probability})


def test_forced_language_probability_is_builtin_float():
    result = _transcriber(language="es").transcribe(_audio())

    assert type(result.language_probability) is float
    json.dumps({"lang_prob": result.language_probability})


def test_explicit_language_with_sticky_ignores_sticky():
    t = _transcriber(language="es", language_strategy="sticky")
    t.transcribe(_audio())
    t.transcribe(_audio())
    assert _FakeModel.last.auto_detect_calls == 0
    assert t._pinned_language is None


# --- transcribe: segment gate and text -------------------------------------


def test_no_speech_thold_forwards_no_speech_prob_max():
    t = _transcriber(no_speech_prob_max=0.75)
    t.transcribe(_audio())
    kwargs = _FakeModel.last.transcribe_calls[0]
    assert kwargs["no_speech_thold"] == 0.75
    assert kwargs["extract_probability"] is True


def test_probability_gate_drops_low_and_keeps_high():
    t = _transcriber(avg_logprob_min=-1.0)
    _FakeModel.last.segments_to_return = [
        _Seg(text="bad ", probability=math.exp(-1.0) - 0.05),
        _Seg(text="good", probability=0.9),
    ]
    assert t.transcribe(_audio()).text == "good"


def test_text_is_joined_and_stripped():
    t = _transcriber()
    _FakeModel.last.segments_to_return = [
        _Seg(text=" hola ", probability=0.9),
        _Seg(text="mundo ", probability=0.9),
    ]
    assert t.transcribe(_audio()).text == "hola mundo"


def test_text_joins_spaceless_segments_with_space():
    t = _transcriber()
    _FakeModel.last.segments_to_return = [
        _Seg(text="most", probability=0.9),
        _Seg(text="matter", probability=0.9),
    ]
    assert t.transcribe(_audio()).text == "most matter"


# --- sticky strategy --------------------------------------------------------


def test_sticky_pins_on_first_call_and_reuses_pin():
    t = _transcriber(language=None, language_strategy="sticky")
    _FakeModel.last.detected_language = "es"
    _FakeModel.last.detected_language_probability = 0.95
    _FakeModel.last.segments_to_return = [_Seg(probability=0.9)]

    first = t.transcribe(_audio())
    assert first.language_probability == 0.95
    assert t._pinned_language == "es"
    assert t._pinned_language_probability == 0.95

    second = t.transcribe(_audio())
    assert _FakeModel.last.auto_detect_calls == 1  # pin reused, no second detection
    assert _FakeModel.last.transcribe_calls[1]["language"] == "es"
    assert second.language == "es"
    assert second.language_probability == 0.95


def test_sticky_pinned_probability_is_builtin_float():
    """The pin must store a builtin float, not the np.float32 detection result."""
    t = _transcriber(language=None, language_strategy="sticky")
    _FakeModel.last.detected_language = "es"
    _FakeModel.last.detected_language_probability = np.float32(0.95)
    _FakeModel.last.segments_to_return = [_Seg(probability=0.9)]

    first = t.transcribe(_audio())
    second = t.transcribe(_audio())

    assert type(first.language_probability) is float
    assert type(t._pinned_language_probability) is float
    assert type(second.language_probability) is float
    json.dumps({"lang_prob": second.language_probability})


def test_sticky_clears_pin_on_low_probability():
    t = _transcriber(
        language=None,
        language_strategy="sticky",
        sticky_redetect_logprob=-0.8,
        avg_logprob_min=-1.0,
    )
    _FakeModel.last.detected_language = "en"
    _FakeModel.last.detected_language_probability = 0.98
    _FakeModel.last.segments_to_return = [_Seg(probability=0.9)]

    t.transcribe(_audio())
    assert t._pinned_language == "en"

    # Accepted (0.40 >= exp(-1.0)) but below exp(-0.8) -> pin cleared
    _FakeModel.last.segments_to_return = [_Seg(probability=0.40)]
    t.transcribe(_audio())
    assert t._pinned_language is None


def test_sticky_clears_pin_when_no_segments_accepted():
    t = _transcriber(language=None, language_strategy="sticky")
    _FakeModel.last.detected_language = "en"
    _FakeModel.last.detected_language_probability = 0.99
    _FakeModel.last.segments_to_return = [_Seg(probability=0.9)]

    t.transcribe(_audio())
    assert t._pinned_language == "en"

    _FakeModel.last.segments_to_return = [_Seg(probability=0.1)]
    t.transcribe(_audio())
    assert t._pinned_language is None


def test_sticky_retains_pin_when_probability_missing():
    """A NaN probability is accepted and does not trigger re-detection."""
    t = _transcriber(language=None, language_strategy="sticky")
    _FakeModel.last.detected_language = "en"
    _FakeModel.last.detected_language_probability = 0.99
    _FakeModel.last.segments_to_return = [_Seg(probability=0.9)]

    t.transcribe(_audio())
    assert t._pinned_language == "en"

    _FakeModel.last.segments_to_return = [_Seg(probability=math.nan)]
    t.transcribe(_audio())
    assert t._pinned_language == "en"


def test_static_strategy_never_pins():
    t = _transcriber(language=None, language_strategy="static")
    _FakeModel.last.detected_language = "es"

    t.transcribe(_audio())
    t.transcribe(_audio())

    assert _FakeModel.last.auto_detect_calls == 2
    assert t._pinned_language is None


@pytest.mark.parametrize("strategy", ["static", "sticky"])
def test_transcribe_receives_detected_language(strategy):
    t = _transcriber(language=None, language_strategy=strategy)
    _FakeModel.last.detected_language = "en"
    t.transcribe(_audio())
    assert _FakeModel.last.transcribe_calls[0]["language"] == "en"