fix(plugins/sarvam): thread language_probability into SpeechData.confidence (#5830)

Author: hashirventhodi

livekit-plugins/livekit-plugins-sarvam/livekit/plugins/sarvam/stt.py

@@ -22,6 +22,7 @@
 import asyncio
 import enum
 import json
+import logging
 import os
 import platform
 import weakref
@@ -343,6 +344,32 @@ def _get_urls_for_model(model: str) -> tuple[str, str]:
     return SARVAM_STT_BASE_URL, SARVAM_STT_STREAMING_URL
 
 
+def _extract_confidence(
+    payload: dict,
+    instance_logger: logging.Logger,
+) -> float:
+    """Read Sarvam's ``language_probability`` from a response payload.
+
+    Returns the value as a float when present and numeric. Falls back to
+    ``1.0`` when the field is absent, ``None``, or has an unexpected type
+    (defensive — the field is documented for the REST endpoint but not
+    explicitly for streaming, so contract drift is logged for visibility).
+    """
+    value = payload.get("language_probability")
+    # bool is a subclass of int — exclude explicitly so that an accidental
+    # JSON `false` doesn't silently become ``confidence=0.0``. Same pattern
+    # as livekit-plugins-slng/.../stt.py.
+    if isinstance(value, (int, float)) and not isinstance(value, bool):
+        return float(value)
+    if value is not None:
+        instance_logger.debug(
+            "Unexpected language_probability type: %s (value=%r); falling back to confidence=1.0",
+            type(value).__name__,
+            value,
+        )
+    return 1.0
+
+
 def _calculate_audio_duration(
     buffer: AudioBuffer,
 ) -> float:  # TODO: Copied from livekit/agents/utils/audio.py, check if it can be reused
@@ -698,7 +725,7 @@ async def _recognize_impl(
                         text=transcript_text,
                         start_time=start_time,
                         end_time=end_time,
-                        confidence=1.0,  # Sarvam doesn't provide confidence score in this response
+                        confidence=_extract_confidence(response_json, self._logger),
                     )
                 ]
 
@@ -1477,6 +1504,7 @@ async def _handle_transcript_data(self, data: dict) -> None:
                 text=transcript_text,
                 start_time=transcript_data.get("speech_start", 0.0),
                 end_time=transcript_data.get("speech_end", 0.0),
+                confidence=_extract_confidence(transcript_data, self._logger),
             )
 
             # Create final transcript event with request_id

livekit-plugins/livekit-plugins-sarvam/tests/test_language_probability.py

@@ -0,0 +1,144 @@
+"""Tests for SpeechData.confidence threading from Sarvam's language_probability.
+
+Verifies that both the REST and WS paths thread ``language_probability`` from
+Sarvam's response into ``SpeechData.confidence`` (instead of the previous
+hardcoded ``1.0``), with a defensive fallback when the field is absent or has
+an unexpected type.
+"""
+
+from __future__ import annotations
+
+from typing import Any
+from unittest.mock import MagicMock
+
+import pytest
+
+from livekit.agents import stt
+from livekit.plugins.sarvam.stt import SpeechStream
+
+# ---------------------------------------------------------------------------
+# Helpers — build a minimal STT instance + fake the channel/logger/state that
+# `_handle_transcript_data` touches. We bypass __init__ so the test doesn't
+# need an API key, an HTTP session, or a real WebSocket.
+# ---------------------------------------------------------------------------
+
+
+def _make_stream_under_test() -> tuple[SpeechStream, list[Any]]:
+    """Construct a minimal SpeechStream and collect its emitted events.
+
+    Returns ``(stream_instance, captured_events)`` where each event sent via
+    ``send_nowait`` is appended to ``captured_events``. We bypass ``__init__``
+    so the test doesn't need an API key, an HTTP session, or a real WebSocket.
+    """
+    instance = SpeechStream.__new__(SpeechStream)
+    captured: list[Any] = []
+    event_ch = MagicMock()
+    event_ch.send_nowait = captured.append
+    instance._event_ch = event_ch  # type: ignore[attr-defined]
+    instance._logger = MagicMock()  # type: ignore[attr-defined]
+    instance._build_log_context = lambda: {}  # type: ignore[attr-defined]
+    instance._server_request_id = None  # type: ignore[attr-defined]
+    instance._opts = MagicMock(language="en-IN")  # type: ignore[attr-defined]
+    return instance, captured
+
+
+def _ws_message(**transcript_overrides: Any) -> dict:
+    """Build the outer WS message dict expected by ``_handle_transcript_data``.
+
+    Default shape mirrors a Saaras v3 streaming final-transcript chunk.
+    """
+    transcript_data: dict[str, Any] = {
+        "transcript": "नमस्ते",
+        "language_code": "hi-IN",
+        "speech_start": 0.0,
+        "speech_end": 1.2,
+        "metrics": {"audio_duration": 1.2},
+        "request_id": "req-test",
+    }
+    transcript_data.update(transcript_overrides)
+    return {"type": "data", "data": transcript_data}
+
+
+def _final_event(captured: list[Any]) -> stt.SpeechEvent:
+    finals = [ev for ev in captured if ev.type == stt.SpeechEventType.FINAL_TRANSCRIPT]
+    assert finals, "no FINAL_TRANSCRIPT event emitted"
+    return finals[0]
+
+
+# ---------------------------------------------------------------------------
+# WS path — happy cases
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize(
+    "language_probability, expected_confidence",
+    [
+        (0.87, 0.87),
+        (1.0, 1.0),
+        (0.0, 0.0),
+        (0.5, 0.5),
+        (0.123, 0.123),
+    ],
+)
+async def test_ws_threads_language_probability_into_confidence(
+    language_probability: float, expected_confidence: float
+) -> None:
+    """WS path must thread Sarvam's language_probability into SpeechData.confidence."""
+    instance, captured = _make_stream_under_test()
+    await instance._handle_transcript_data(_ws_message(language_probability=language_probability))
+    final = _final_event(captured)
+    assert final.alternatives[0].confidence == pytest.approx(expected_confidence)
+
+
+# ---------------------------------------------------------------------------
+# WS path — defensive fallback cases (absent / null / wrong type)
+# ---------------------------------------------------------------------------
+
+
+async def test_ws_missing_language_probability_falls_back_to_1_0() -> None:
+    """When the field is absent, confidence falls back to 1.0 (no crash)."""
+    instance, captured = _make_stream_under_test()
+    # _ws_message() helper omits language_probability by default
+    await instance._handle_transcript_data(_ws_message())
+    final = _final_event(captured)
+    assert final.alternatives[0].confidence == 1.0
+
+
+async def test_ws_null_language_probability_falls_back_to_1_0() -> None:
+    """Explicit null also falls back to 1.0."""
+    instance, captured = _make_stream_under_test()
+    await instance._handle_transcript_data(_ws_message(language_probability=None))
+    final = _final_event(captured)
+    assert final.alternatives[0].confidence == 1.0
+
+
+@pytest.mark.parametrize("bad_value", ["0.95", [], {}, object(), True, False])
+async def test_ws_unexpected_type_falls_back_to_1_0(bad_value: Any) -> None:
+    """String / list / dict / object / bool → confidence falls back to 1.0 with a debug log.
+
+    bool is included because Python's ``bool`` is a subclass of ``int``;
+    without an explicit guard a JSON ``false`` from Sarvam would silently
+    become ``confidence=0.0`` and wrongly flag a valid transcript as low
+    confidence. Same defensive pattern as ``livekit-plugins-slng``.
+    """
+    instance, captured = _make_stream_under_test()
+    await instance._handle_transcript_data(_ws_message(language_probability=bad_value))
+    final = _final_event(captured)
+    assert final.alternatives[0].confidence == 1.0
+    # The defensive branch logs a debug warning so contract drift is visible.
+    assert instance._logger.debug.called  # type: ignore[attr-defined]
+
+
+# ---------------------------------------------------------------------------
+# WS path — out-of-range values pass through verbatim (clamping is not this
+# layer's job; downstream consumers can clamp if they need to).
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("value", [-0.5, 1.5, 2.0])
+async def test_ws_out_of_range_values_passed_through(value: float) -> None:
+    """Out-of-[0,1] values are passed through verbatim (no clamping)."""
+    instance, captured = _make_stream_under_test()
+    await instance._handle_transcript_data(_ws_message(language_probability=value))
+    final = _final_event(captured)
+    assert final.alternatives[0].confidence == pytest.approx(value)