feat(gradium): optional VAD-based start of speech in STT

livekit-plugins/livekit-plugins-gradium/livekit/plugins/gradium/stt.py

@@ -65,6 +65,9 @@ class STTOptions:
     # When set, we flush the stt state on the first time the VAD triggers
     # in order to recover the text currently being processed as soon as possible.
     vad_flush: bool = True
+    # When set, a VAD inactivity probability below this threshold signals the start of
+    # speech before any text is transcribed, which allows a faster barge-in.
+    vad_start_threshold: float | None = None
 
 
 class STT(stt.STT):
@@ -81,6 +84,7 @@ def __init__(
         vad_threshold: float = 0.9,
         vad_bucket: int | None = 2,
         vad_flush: bool = True,
+        vad_start_threshold: float | None = None,
         temperature: float | None = None,
         language: str = "en",
     ):
@@ -122,6 +126,7 @@ def __init__(
             vad_threshold=vad_threshold,
             vad_bucket=vad_bucket,
             vad_flush=vad_flush,
+            vad_start_threshold=vad_start_threshold,
             temperature=temperature,
             language=LanguageCode(language),
         )
@@ -399,9 +404,23 @@ async def recv_task(ws: aiohttp.ClientWebSocketResponse) -> None:
 
                     elif type_ == "step":
                         if not speaking:
+                            start_threshold = self._opts.vad_start_threshold
+                            vad_bucket = self._opts.vad_bucket
+                            # steps sent before any audio was processed carry a warm-up
+                            # estimate that reads as activity, ignore them
+                            if (
+                                vad_bucket is not None
+                                and start_threshold is not None
+                                and data.get("total_duration_s", 0.0) > 0
+                                and data["vad"][vad_bucket]["inactivity_prob"] < start_threshold
+                            ):
+                                speaking = True
+                                self._event_ch.send_nowait(
+                                    stt.SpeechEvent(type=stt.SpeechEventType.START_OF_SPEECH)
+                                )
                             self._mark_provider_progress()
                             continue
-                        if vad_bucket := self._opts.vad_bucket:
+                        if (vad_bucket := self._opts.vad_bucket) is not None:
                             positive_vad = (
                                 data["vad"][vad_bucket]["inactivity_prob"]
                                 > self._opts.vad_threshold

tests/test_plugin_gradium_stt.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import asyncio
 import json
 from unittest.mock import MagicMock, patch
 
@@ -225,6 +226,199 @@ def _fake_create_task(coro, *args, **kwargs):
     assert setup_msg["json_config"] == {"language": "pt", "temp": 0.3}
 
 
+class _FakeMessage:
+    def __init__(self, payload: dict) -> None:
+        import aiohttp
+
+        self.type = aiohttp.WSMsgType.TEXT
+        self.data = json.dumps(payload)
+
+
+class _FakeWebSocket:
+    def __init__(self) -> None:
+        self.incoming: asyncio.Queue[dict] = asyncio.Queue()
+        self.close_code = None
+
+    async def send_str(self, message: str) -> None:
+        pass
+
+    async def ping(self) -> None:
+        pass
+
+    async def receive(self) -> _FakeMessage:
+        return _FakeMessage(await self.incoming.get())
+
+    async def close(self) -> None:
+        pass
+
+
+class _FakeSession:
+    def __init__(self) -> None:
+        self.ws = _FakeWebSocket()
+
+    async def ws_connect(self, url, headers, **kwargs) -> _FakeWebSocket:
+        return self.ws
+
+
+def _step(inactivity_prob: float, total_duration_s: float = 1.0) -> dict:
+    return {
+        "type": "step",
+        "total_duration_s": total_duration_s,
+        "vad": [
+            {"horizon_s": horizon_s, "inactivity_prob": inactivity_prob}
+            for horizon_s in (0.5, 1.0, 2.0)
+        ],
+    }
+
+
+class _EventCollector:
+    """Collects speech events from a stream in the background."""
+
+    def __init__(self, stream) -> None:
+        self.events: list = []
+        self._task = asyncio.create_task(self._collect(stream))
+
+    async def _collect(self, stream) -> None:
+        from livekit.agents.stt import SpeechEventType
+
+        async for ev in stream:
+            if ev.type != SpeechEventType.RECOGNITION_USAGE:
+                self.events.append(ev)
+
+    async def wait_for(self, count: int) -> list:
+        for _ in range(500):
+            if len(self.events) >= count:
+                return self.events
+            await asyncio.sleep(0.01)
+        raise AssertionError(f"expected {count} events, got {self.events}")
+
+    async def drain(self, ws: _FakeWebSocket) -> None:
+        while not ws.incoming.empty():
+            await asyncio.sleep(0.01)
+        await asyncio.sleep(0.05)
+
+
+@pytest.mark.asyncio
+async def test_vad_start_of_speech_opens_a_turn_before_any_text():
+    """With vad_start_threshold, VAD activity starts speech, ignoring warm-up steps."""
+    from livekit.agents.stt import SpeechEventType
+    from livekit.plugins.gradium import STT
+
+    session = _FakeSession()
+    stt = STT(api_key="test-key", http_session=session, vad_start_threshold=0.4)
+    stream = stt.stream()
+    collector = _EventCollector(stream)
+    try:
+        # warm-up steps sent before any audio was processed read as activity
+        for prob in (0.1, 0.1):
+            session.ws.incoming.put_nowait(_step(prob, total_duration_s=0.0))
+        await collector.drain(session.ws)
+        assert collector.events == []
+
+        # speech right from the first processed audio opens the turn
+        session.ws.incoming.put_nowait(_step(0.2, total_duration_s=0.08))
+        events = await collector.wait_for(1)
+        assert [ev.type for ev in events] == [SpeechEventType.START_OF_SPEECH]
+
+        # text in the open turn does not start it again
+        session.ws.incoming.put_nowait({"type": "text", "text": "hello", "start_s": 0.5})
+        events = await collector.wait_for(2)
+        assert [ev.type for ev in events] == [
+            SpeechEventType.START_OF_SPEECH,
+            SpeechEventType.INTERIM_TRANSCRIPT,
+        ]
+    finally:
+        await stream.aclose()
+
+
+@pytest.mark.asyncio
+async def test_vad_start_of_speech_without_text_still_ends_the_turn():
+    """A VAD-opened turn without text ends on inactivity and the next one can open."""
+    from livekit.agents.stt import SpeechEventType
+    from livekit.plugins.gradium import STT
+
+    session = _FakeSession()
+    stt = STT(api_key="test-key", http_session=session, vad_start_threshold=0.4)
+    stream = stt.stream()
+    collector = _EventCollector(stream)
+    try:
+        session.ws.incoming.put_nowait(_step(0.2))
+        # the first positive step starts the countdown, delay_in_tokens (6) more end the turn
+        for _ in range(7):
+            session.ws.incoming.put_nowait(_step(0.95))
+        session.ws.incoming.put_nowait(_step(0.2))
+
+        events = await collector.wait_for(4)
+        assert [ev.type for ev in events] == [
+            SpeechEventType.START_OF_SPEECH,
+            SpeechEventType.FINAL_TRANSCRIPT,
+            SpeechEventType.END_OF_SPEECH,
+            SpeechEventType.START_OF_SPEECH,
+        ]
+        assert events[1].alternatives[0].text == ""
+    finally:
+        await stream.aclose()
+
+
+@pytest.mark.asyncio
+async def test_vad_start_of_speech_is_off_by_default():
+    """Without vad_start_threshold, speech only starts on the first transcribed text."""
+    from livekit.agents.stt import SpeechEventType
+    from livekit.plugins.gradium import STT
+
+    session = _FakeSession()
+    stt = STT(api_key="test-key", http_session=session)
+    stream = stt.stream()
+    collector = _EventCollector(stream)
+    try:
+        for prob in (0.9, 0.2, 0.2):
+            session.ws.incoming.put_nowait(_step(prob))
+        await collector.drain(session.ws)
+        assert collector.events == []
+
+        session.ws.incoming.put_nowait({"type": "text", "text": "hello", "start_s": 0.5})
+        events = await collector.wait_for(2)
+        assert [ev.type for ev in events] == [
+            SpeechEventType.START_OF_SPEECH,
+            SpeechEventType.INTERIM_TRANSCRIPT,
+        ]
+    finally:
+        await stream.aclose()
+
+
+@pytest.mark.asyncio
+async def test_vad_start_of_speech_with_vad_bucket_zero():
+    """vad_bucket=0 selects the first horizon for both the start and the end of a turn."""
+    from livekit.agents.stt import SpeechEventType
+    from livekit.plugins.gradium import STT
+
+    session = _FakeSession()
+    stt = STT(api_key="test-key", http_session=session, vad_bucket=0, vad_start_threshold=0.4)
+    stream = stt.stream()
+    collector = _EventCollector(stream)
+    try:
+
+        def bucket_zero_step(inactivity_prob: float) -> dict:
+            # the other buckets disagree, so only bucket 0 can drive the turn
+            step = _step(1.0 - inactivity_prob)
+            step["vad"][0]["inactivity_prob"] = inactivity_prob
+            return step
+
+        session.ws.incoming.put_nowait(bucket_zero_step(0.2))
+        # the first positive step starts the countdown, delay_in_tokens (6) more end the turn
+        for _ in range(7):
+            session.ws.incoming.put_nowait(bucket_zero_step(0.95))
+
+        events = await collector.wait_for(3)
+        assert [ev.type for ev in events] == [
+            SpeechEventType.START_OF_SPEECH,
+            SpeechEventType.FINAL_TRANSCRIPT,
+            SpeechEventType.END_OF_SPEECH,
+        ]
+    finally:
+        await stream.aclose()
+
+
 @pytest.mark.no_concurrent
 async def test_internal_reconnect_reanchors_audio_clock(monkeypatch):
     from livekit.agents import APIConnectOptions