fix(livekit-agents): forward prewarm to the primary LLM in the fallback adapter

livekit-agents/livekit/agents/llm/fallback_adapter.py

@@ -108,6 +108,19 @@ def chat(
             extra_kwargs=extra_kwargs,
         )
 
+    def prewarm(self, *, loop: asyncio.AbstractEventLoop | None = None) -> None:
+        """Pre-warm the primary LLM.
+
+        Only the first instance is prewarmed; the remaining instances are not expected to
+        serve traffic unless the primary fails.
+
+        Args:
+            loop: Event loop to schedule the prewarm request on. Defaults to the
+                running event loop.
+        """
+        if self._llm_instances:
+            self._llm_instances[0].prewarm(loop=loop)
+
     async def aclose(self) -> None:
         for llm_instance in self._llm_instances:
             llm_instance.off("metrics_collected", self._on_metrics_collected)

tests/test_llm_fallback.py

@@ -0,0 +1,69 @@
+from __future__ import annotations
+
+import asyncio
+
+import pytest
+
+from livekit.agents.llm import FallbackAdapter
+
+from .fake_llm import FakeLLM
+
+pytestmark = [pytest.mark.unit]
+
+
+class PrewarmableLLM(FakeLLM):
+    """FakeLLM that opts into prewarming by overriding ``_prewarm_impl``."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.prewarmed = asyncio.Event()
+
+    async def _prewarm_impl(self) -> None:
+        self.prewarmed.set()
+
+
+class RecordingLLM(FakeLLM):
+    """FakeLLM that records the event loop it was asked to prewarm on."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.prewarm_loop: asyncio.AbstractEventLoop | None = None
+
+    def prewarm(self, *, loop: asyncio.AbstractEventLoop | None = None) -> None:
+        self.prewarm_loop = loop
+
+
+async def test_prewarm_forwarded_to_primary_llm() -> None:
+    primary = PrewarmableLLM()
+    fallback = PrewarmableLLM()
+
+    fallback_adapter = FallbackAdapter([primary, fallback])
+    try:
+        fallback_adapter.prewarm()
+
+        await asyncio.wait_for(primary.prewarmed.wait(), timeout=5)
+        assert not fallback.prewarmed.is_set(), (
+            "expected only the primary LLM to be prewarmed, the fallbacks should stay cold"
+        )
+    finally:
+        await fallback_adapter.aclose()
+        await primary.aclose()
+        await fallback.aclose()
+
+
+async def test_prewarm_forwards_event_loop() -> None:
+    primary = RecordingLLM()
+
+    fallback_adapter = FallbackAdapter([primary])
+    # a loop distinct from the running one, so the assertion fails if `loop` is dropped
+    # and the wrapped LLM falls back to the running loop
+    supplied_loop = asyncio.new_event_loop()
+    try:
+        fallback_adapter.prewarm(loop=supplied_loop)
+
+        assert primary.prewarm_loop is supplied_loop, (
+            "expected the provided event loop to be forwarded to the primary LLM"
+        )
+    finally:
+        supplied_loop.close()
+        await fallback_adapter.aclose()