fix(livekit-agents): propagate prewarm through the STT fallback and stream adapters

livekit-agents/livekit/agents/stt/fallback_adapter.py

@@ -280,6 +280,15 @@ def stream(
     ) -> RecognizeStream:
         return FallbackRecognizeStream(stt=self, language=language, conn_options=conn_options)
 
+    def prewarm(self) -> None:
+        """Pre-warm the primary STT.
+
+        Only the first instance is prewarmed; the remaining instances are not expected to
+        serve traffic unless the primary fails.
+        """
+        if self._stt_instances:
+            self._stt_instances[0].prewarm()
+
     async def aclose(self) -> None:
         for stt_status in self._status:
             if stt_status.recovering_recognize_task is not None:

livekit-agents/livekit/agents/stt/stream_adapter.py

@@ -78,6 +78,9 @@ def stream(
             conn_options=conn_options,
         )
 
+    def prewarm(self) -> None:
+        self._stt.prewarm()
+
     def _on_metrics_collected(self, *args: Any, **kwargs: Any) -> None:
         self.emit("metrics_collected", *args, **kwargs)
 

tests/test_stt_prewarm.py

@@ -0,0 +1,70 @@
+from __future__ import annotations
+
+import pytest
+
+from livekit.agents.stt import FallbackAdapter, StreamAdapter, STTCapabilities
+
+from .fake_stt import FakeSTT
+from .fake_vad import FakeVAD
+
+pytestmark = [pytest.mark.unit]
+
+
+class PrewarmableSTT(FakeSTT):
+    """FakeSTT that records prewarm calls, optionally without streaming support."""
+
+    def __init__(self, *, streaming: bool = True) -> None:
+        super().__init__()
+        self._capabilities = STTCapabilities(streaming=streaming, interim_results=False)
+        self.prewarm_count = 0
+
+    def prewarm(self) -> None:
+        self.prewarm_count += 1
+
+
+async def test_fallback_adapter_prewarms_primary_stt() -> None:
+    primary = PrewarmableSTT()
+    fallback = PrewarmableSTT()
+
+    fallback_adapter = FallbackAdapter([primary, fallback])
+    try:
+        fallback_adapter.prewarm()
+
+        assert primary.prewarm_count == 1, "expected the primary STT to be prewarmed"
+        assert fallback.prewarm_count == 0, (
+            "expected only the primary STT to be prewarmed, the fallbacks should stay cold"
+        )
+    finally:
+        await fallback_adapter.aclose()
+
+
+async def test_stream_adapter_forwards_prewarm() -> None:
+    wrapped = PrewarmableSTT(streaming=False)
+
+    stream_adapter = StreamAdapter(stt=wrapped, vad=FakeVAD())
+    try:
+        stream_adapter.prewarm()
+
+        assert wrapped.prewarm_count == 1, (
+            "expected StreamAdapter to forward prewarm to the wrapped STT"
+        )
+    finally:
+        await stream_adapter.aclose()
+
+
+async def test_prewarm_reaches_non_streaming_primary() -> None:
+    # a non-streaming STT is transparently wrapped in a StreamAdapter by the fallback
+    # adapter, so prewarm has to survive both hops to reach the provider
+    primary = PrewarmableSTT(streaming=False)
+    fallback = PrewarmableSTT()
+
+    fallback_adapter = FallbackAdapter([primary, fallback], vad=FakeVAD())
+    try:
+        fallback_adapter.prewarm()
+
+        assert primary.prewarm_count == 1, (
+            "expected prewarm to reach the primary STT through the wrapping StreamAdapter"
+        )
+        assert fallback.prewarm_count == 0
+    finally:
+        await fallback_adapter.aclose()