AgentLauncher: track the active agents and stop those that are idle

Author: dangusev

agents-core/vision_agents/core/agents/__init__.py

@@ -7,9 +7,11 @@
 from .agents import Agent as Agent
 from .conversation import Conversation as Conversation
 from .agent_launcher import AgentLauncher as AgentLauncher
+from .agent_types import AgentOptions as AgentOptions
 
 __all__ = [
     "Agent",
     "Conversation",
     "AgentLauncher",
+    "AgentOptions",
 ]

agents-core/vision_agents/core/agents/agent_launcher.py

@@ -1,9 +1,10 @@
 import asyncio
 import logging
-from typing import TYPE_CHECKING, Awaitable, Callable
+import weakref
+from typing import TYPE_CHECKING, Awaitable, Callable, Optional
 
-from vision_agents.core.utils.utils import await_or_run
-from vision_agents.core.warmup import WarmupCache, Warmable
+from vision_agents.core.utils.utils import await_or_run, cancel_and_wait
+from vision_agents.core.warmup import Warmable, WarmupCache
 
 if TYPE_CHECKING:
     from .agents import Agent
@@ -23,33 +24,72 @@ def __init__(
         self,
         create_agent: Callable[..., "Agent" | Awaitable["Agent"]],
         join_call: Callable[..., None | Awaitable[None]] | None = None,
+        agent_idle_timeout: float = 10.0,
+        agent_idle_cleanup_interval: float = 5.0,
     ):
         """
         Initialize the agent launcher.
 
         Args:
             create_agent: A function that creates and returns an Agent instance
             join_call: Optional function that handles joining a call with the agent
+            agent_idle_timeout: Optional timeout in seconds for agent to stay alone on the call. Default - `30.0`.
+                `0` means idle agents won't leave the call until it's ended.
+
         """
         self.create_agent = create_agent
         self.join_call = join_call
         self._warmup_lock = asyncio.Lock()
         self._warmup_cache = WarmupCache()
 
+        if agent_idle_timeout < 0:
+            raise ValueError("agent_idle_timeout must be >= 0")
+        self._agent_idle_timeout = agent_idle_timeout
+
+        if agent_idle_cleanup_interval <= 0:
+            raise ValueError("agent_idle_cleanup_interval must be > 0")
+        self._agent_idle_cleanup_interval = agent_idle_cleanup_interval
+
+        self._active_agents: weakref.WeakSet[Agent] = weakref.WeakSet()
+
+        self._running = False
+        self._cleanup_task: Optional[asyncio.Task] = None
+        self._warmed_up: bool = False
+
+    async def start(self):
+        if self._running:
+            raise RuntimeError("AgentLauncher is already running")
+        logger.debug("Starting AgentLauncher")
+        self._running = True
+        await self.warmup()
+        self._cleanup_task = asyncio.create_task(self._cleanup_idle_agents())
+        logger.debug("AgentLauncher started")
+
+    async def stop(self):
+        logger.debug("Stopping AgentLauncher")
+        self._running = False
+        if self._cleanup_task:
+            await cancel_and_wait(self._cleanup_task)
+        logger.debug("AgentLauncher stopped")
+
     async def warmup(self) -> None:
         """
         Warm up all agent components.
 
         This method creates the agent and calls warmup() on LLM, TTS, STT,
         and turn detection components if they exist.
         """
+        if self._warmed_up or self._warmup_lock.locked():
+            return
+
         async with self._warmup_lock:
             logger.info("Creating agent...")
 
             # Create a dry-run Agent instance and warmup its components for the first time.
             agent: "Agent" = await await_or_run(self.create_agent)
             logger.info("Warming up agent components...")
             await self._warmup_agent(agent)
+            self._warmed_up = True
 
             logger.info("Agent warmup completed")
 
@@ -65,6 +105,7 @@ async def launch(self, **kwargs) -> "Agent":
         """
         agent: "Agent" = await await_or_run(self.create_agent, **kwargs)
         await self._warmup_agent(agent)
+        self._active_agents.add(agent)
         return agent
 
     async def _warmup_agent(self, agent: "Agent") -> None:
@@ -99,10 +140,46 @@ async def _warmup_agent(self, agent: "Agent") -> None:
             warmup_tasks.append(agent.turn_detection.warmup(self._warmup_cache))
 
         # Warmup processors
-        if agent.processors:
-            for processor in agent.processors:
-                if isinstance(processor, Warmable):
-                    warmup_tasks.append(processor.warmup(self._warmup_cache))
+        for processor in agent.processors:
+            if isinstance(processor, Warmable):
+                warmup_tasks.append(processor.warmup(self._warmup_cache))
 
         if warmup_tasks:
             await asyncio.gather(*warmup_tasks)
+
+    async def _cleanup_idle_agents(self) -> None:
+        if not self._agent_idle_timeout:
+            return
+
+        while self._running:
+            # Collect idle agents first to close them all at once
+            idle_agents = []
+            for agent in self._active_agents:
+                agent_idle_for = agent.idle_for()
+                if agent_idle_for >= self._agent_idle_timeout:
+                    logger.info(
+                        f'Agent with user_id "{agent.agent_user.id}" is idle for {round(agent_idle_for, 2)}s, '
+                        f"closing it after {self._agent_idle_timeout}s timeout"
+                    )
+                    idle_agents.append(agent)
+
+            if idle_agents:
+                coros = [asyncio.shield(a.close()) for a in idle_agents]
+                result = await asyncio.shield(
+                    asyncio.gather(*coros, return_exceptions=True)
+                )
+                for agent, r in zip(idle_agents, result):
+                    if isinstance(r, Exception):
+                        logger.error(
+                            f"Failed to close idle agent with user_id {agent.agent_user.id}",
+                            exc_info=r,
+                        )
+
+            await asyncio.sleep(self._agent_idle_cleanup_interval)
+
+    async def __aenter__(self):
+        await self.start()
+        return self
+
+    async def __aexit__(self, exc_type, exc_val, exc_tb):
+        await self.stop()

agents-core/vision_agents/core/agents/agents.py

@@ -690,6 +690,10 @@ def _end_tracing(self):
             otel_context.detach(self._context_token)
             self._context_token = None
 
+    @property
+    def closed(self) -> bool:
+        return self._closed
+
     async def close(self):
         """
         Clean up all connections and resources.

agents-core/vision_agents/core/cli/cli_runner.py

@@ -101,8 +101,8 @@ async def _run():
             logger.info("🚀 Launching agent...")
 
             try:
-                # Warmup the agent's dependencies
-                await launcher.warmup()
+                # Start the agent launcher.
+                await launcher.start()
 
                 # Create the agent
                 agent = await launcher.launch()
@@ -135,6 +135,8 @@ async def _run():
             except Exception as e:
                 logger.error(f"❌ Error running agent: {e}", exc_info=True)
                 raise
+            finally:
+                await launcher.stop()
 
         asyncio_logger_level = asyncio_logger.level
 

tests/test_agents/__init__.py

@@ -0,0 +1,163 @@
+import asyncio
+from typing import Any
+from unittest.mock import MagicMock, patch
+
+import pytest
+from vision_agents.core import Agent, AgentLauncher, User
+from vision_agents.core.events import EventManager
+from vision_agents.core.llm import LLM
+from vision_agents.core.llm.llm import LLMResponseEvent
+from vision_agents.core.tts import TTS
+from vision_agents.core.warmup import Warmable
+
+
+class DummyTTS(TTS):
+    async def stream_audio(self, *_, **__):
+        return b""
+
+    async def stop_audio(self) -> None: ...
+
+
+class DummyLLM(LLM, Warmable[bool]):
+    def __init__(self):
+        super(DummyLLM, self).__init__()
+        self.warmed_up = False
+
+    async def simple_response(self, *_, **__) -> LLMResponseEvent[Any]:
+        return LLMResponseEvent(text="Simple response", original=None)
+
+    async def on_warmup(self) -> bool:
+        return True
+
+    async def on_warmed_up(self, *_) -> None:
+        self.warmed_up = True
+
+
+@pytest.fixture()
+async def stream_edge_mock() -> MagicMock:
+    mock = MagicMock()
+    mock.events.return_value = EventManager()
+    return mock
+
+
+class TestAgentLauncher:
+    async def test_warmup(self, stream_edge_mock):
+        llm = DummyLLM()
+        tts = DummyTTS()
+
+        async def create_agent(**kwargs) -> Agent:
+            return Agent(
+                llm=llm,
+                tts=tts,
+                edge=stream_edge_mock,
+                agent_user=User(name="test"),
+            )
+
+        launcher = AgentLauncher(create_agent=create_agent)
+        await launcher.warmup()
+        assert llm.warmed_up
+
+    async def test_launch(self, stream_edge_mock):
+        llm = DummyLLM()
+        tts = DummyTTS()
+
+        async def create_agent(**kwargs) -> Agent:
+            return Agent(
+                llm=llm,
+                tts=tts,
+                edge=stream_edge_mock,
+                agent_user=User(name="test"),
+            )
+
+        launcher = AgentLauncher(create_agent=create_agent)
+        agent = await launcher.launch()
+        assert agent
+
+    async def test_idle_agents_stopped(self, stream_edge_mock):
+        llm = DummyLLM()
+        tts = DummyTTS()
+
+        async def create_agent(**kwargs) -> Agent:
+            return Agent(
+                llm=llm,
+                tts=tts,
+                edge=stream_edge_mock,
+                agent_user=User(name="test"),
+            )
+
+        launcher = AgentLauncher(
+            create_agent=create_agent,
+            agent_idle_timeout=1.0,
+            agent_idle_cleanup_interval=0.5,
+        )
+        with patch.object(Agent, "idle_for", return_value=10):
+            # Start the launcher internals
+            async with launcher:
+                # Launch a couple of idle agents
+                agent1 = await launcher.launch()
+                agent2 = await launcher.launch()
+                # Sleep 2s to let the launcher clean up the agents
+                await asyncio.sleep(2)
+
+        # The agents must be closed
+        assert agent1.closed
+        assert agent2.closed
+
+    async def test_idle_agents_alive_with_idle_timeout_zero(self, stream_edge_mock):
+        llm = DummyLLM()
+        tts = DummyTTS()
+
+        async def create_agent(**kwargs) -> Agent:
+            return Agent(
+                llm=llm,
+                tts=tts,
+                edge=stream_edge_mock,
+                agent_user=User(name="test"),
+            )
+
+        launcher = AgentLauncher(
+            create_agent=create_agent,
+            agent_idle_timeout=0,
+        )
+        with patch.object(Agent, "idle_for", return_value=10):
+            # Start the launcher internals
+            async with launcher:
+                # Launch a couple of idle agents
+                agent1 = await launcher.launch()
+                agent2 = await launcher.launch()
+                # Sleep 2s to let the launcher clean up the agents
+                await asyncio.sleep(2)
+
+        # The agents must not be closed because agent_idle_timeout=0
+        assert not agent1.closed
+        assert not agent2.closed
+
+    async def test_active_agents_alive(self, stream_edge_mock):
+        llm = DummyLLM()
+        tts = DummyTTS()
+
+        async def create_agent(**kwargs) -> Agent:
+            return Agent(
+                llm=llm,
+                tts=tts,
+                edge=stream_edge_mock,
+                agent_user=User(name="test"),
+            )
+
+        launcher = AgentLauncher(
+            create_agent=create_agent,
+            agent_idle_timeout=1.0,
+            agent_idle_cleanup_interval=0.5,
+        )
+        with patch.object(Agent, "idle_for", return_value=0):
+            # Start the launcher internals
+            async with launcher:
+                # Launch a couple of active agents (idle_for=0)
+                agent1 = await launcher.launch()
+                agent2 = await launcher.launch()
+                # Sleep 2s to let the launcher clean up the agents
+                await asyncio.sleep(2)
+
+        # The agents must be closed
+        assert not agent1.closed
+        assert not agent2.closed