[draft] async text mode

examples/voice_agents/basic_agent.py

@@ -10,6 +10,7 @@
     JobProcess,
     MetricsCollectedEvent,
     RunContext,
+    TextMessageContext,
     cli,
     metrics,
     room_io,
@@ -40,6 +41,7 @@ async def on_enter(self):
         # when the agent is added to the session, it'll generate a reply
         # according to its instructions
         # Keep it uninterruptible so the client has time to calibrate AEC (Acoustic Echo Cancellation).
+        logger.debug("on_enter called")
         self.session.generate_reply(allow_interruptions=False)
 
     # all functions annotated with @function_tool will be passed to the LLM when this
@@ -74,6 +76,15 @@ def prewarm(proc: JobProcess):
 server.setup_fnc = prewarm
 
 
+@server.sms_handler()
+async def sms_handler(ctx: TextMessageContext):
+    logger.info(f"SMS received: {ctx.text}")
+
+    session = await ctx.rehydrated_session()
+    result = await session.run(user_input=ctx.text)
+    await ctx.send_result(result)
+
+
 @server.rtc_session()
 async def entrypoint(ctx: JobContext):
     # each log entry will include these fields

livekit-agents/livekit/agents/__init__.py

@@ -34,6 +34,7 @@
     JobExecutorType,
     JobProcess,
     JobRequest,
+    TextMessageContext,
     get_job_context,
 )
 from .llm.chat_context import (
@@ -121,6 +122,7 @@ def __getattr__(name: str) -> typing.Any:
     "JobProcess",
     "JobContext",
     "JobRequest",
+    "TextMessageContext",
     "get_job_context",
     "JobExecutorType",
     "AutoSubscribe",

livekit-agents/livekit/agents/cli/cli.py

@@ -10,14 +10,16 @@
 import logging
 import os
 import pathlib
+import pickle
 import re
 import signal
 import sys
+import tempfile
 import textwrap
 import threading
 import time
 import traceback
-from collections.abc import Iterator
+from collections.abc import Awaitable, Iterator
 from contextlib import contextmanager
 from types import FrameType
 from typing import TYPE_CHECKING, Annotated, Any, Callable, Literal, Optional, Union
@@ -37,7 +39,7 @@
 from livekit import rtc
 
 from .._exceptions import CLIError
-from ..job import JobExecutorType
+from ..job import JobExecutorType, TextMessageContext
 from ..log import logger
 from ..plugin import Plugin
 from ..utils import aio
@@ -985,7 +987,12 @@
         yield update
 
 
-def _text_mode(c: AgentsConsole) -> None:
+def _text_mode(
+    c: AgentsConsole,
+    *,
+    sms_handler: Callable[[TextMessageContext], Awaitable[None]] | None,
+    sess_dump_file: str | None,
+) -> None:
     def _key_read(ch: str) -> None:
         if ch == key.CTRL_T:
             raise _ToggleMode()
@@ -1006,8 +1013,32 @@
 
         def _generate_with_context(text: str, result_fut: asyncio.Future[list[RunEvent]]) -> None:
             async def _generate(text: str) -> list[RunEvent]:
-                sess = await c.io_session.run(user_input=text)  # type: ignore
-                return sess.events.copy()
+                if sms_handler is not None:
+                    assert sess_dump_file
+
+                    sess_data: bytes | None = None
+                    if os.path.exists(sess_dump_file):
+                        with open(sess_dump_file, "rb") as f:
+                            sess_data = f.read()
+                    text_context = TextMessageContext(
+                        text=text, session=c.io_session, sess_data=sess_data
+                    )
+
+                    await sms_handler(text_context)
+                    result = text_context.result
+                    if result is None:
+                        logger.warning("result is not set from the sms handler")
+                        return []
+
+                    # serialize the state of the session
+                    state = c.io_session.get_state()
+                    with open(sess_dump_file, "wb") as f:
+                        f.write(pickle.dumps(state))
+                    logger.debug("session state serialized")
+                else:
+                    result = await c.io_session.run(user_input=text)  # type: ignore
+
+                return result.events.copy()
 
             def _done_callback(task: asyncio.Task[list[RunEvent]]) -> None:
                 if exception := task.exception():
@@ -1176,6 +1207,8 @@
     output_device: str | None,
     mode: ConsoleMode,
     record: bool,
+    simulate_sms: bool,
+    sess_dump_file: str | None,
 ) -> None:
     c = AgentsConsole.get_instance()
     c.console_mode = mode
@@ -1228,14 +1261,26 @@
         console_worker = _ConsoleWorker(server=server, shutdown_cb=_on_worker_shutdown)
         console_worker.start()
 
+        temp_file: str | None = None
+        sms_handler: Callable[[TextMessageContext], Awaitable[None]] | None = None
+        if simulate_sms:
+            assert server._sms_handler_fnc is not None, (
+                "sms_handler is required when simulating SMS"
+            )
+            sms_handler = server._sms_handler_fnc
+            if not sess_dump_file:
+                sess_dump_file = temp_file = tempfile.mktemp(
+                    suffix=".pkl", prefix="lk-agents-sess_"
+                )
+
         # TODO: wait for a session request the agents console context before showing any of the mode
         try:
             c.wait_for_io_acquisition()
 
             while True:
                 try:
                     if c.console_mode == "text":
-                        _text_mode(c)
+                        _text_mode(c, sms_handler=sms_handler, sess_dump_file=sess_dump_file)
                     elif c.console_mode == "audio":
                         _audio_mode(c, input_device=input_device, output_device=output_device)
 
@@ -1247,6 +1292,8 @@
         finally:
             console_worker.shutdown()
             console_worker.join()
+            if temp_file and pathlib.Path(temp_file).exists():
+                pathlib.Path(temp_file).unlink()
 
     except CLIError as e:
         c.print(" ")
@@ -1370,6 +1417,13 @@
             bool,
             typer.Option(help="Whether to start the console in text mode"),
         ] = False,
+        sms: Annotated[
+            bool, typer.Option(help="Whether to simulate an SMS received event")
+        ] = False,
+        sess_dump_file: Annotated[
+            Optional[str],  # noqa: UP007
+            typer.Option(help="Path to the serialized AgentSession data file"),
+        ] = None,
         record: Annotated[bool, typer.Option(help="Whether to record the AgentSession")] = False,
     ) -> None:
         """
@@ -1389,8 +1443,10 @@
             server=server,
             input_device=input_device,
             output_device=output_device,
-            mode="text" if text else "audio",
+            mode="text" if text or sms else "audio",
             record=record,
+            simulate_sms=sms,
+            sess_dump_file=sess_dump_file,
         )
 
     @app.command()

livekit-agents/livekit/agents/job.py

@@ -22,6 +22,7 @@
 import json
 import logging
 import multiprocessing as mp
+import pickle
 import tempfile
 from collections.abc import Coroutine
 from dataclasses import dataclass
@@ -51,6 +52,7 @@
     from .ipc.inference_executor import InferenceExecutor
     from .voice.agent_session import AgentSession
     from .voice.report import SessionReport
+    from .voice.run_result import RunResult
 
 
 def get_job_context() -> JobContext:
@@ -129,6 +131,35 @@ def filter(self, record: logging.LogRecord) -> bool:
         return True
 
 
+class TextMessageContext:
+    def __init__(self, *, text: str, session: AgentSession, sess_data: bytes | None) -> None:
+        self._text = text
+        self._session = session
+        self._sess_data = sess_data
+        self._result: RunResult | None = None
+
+    async def send_result(self, result: RunResult) -> None:
+        self._result = result
+
+    async def rehydrated_session(self) -> AgentSession:
+        # TODO: it should be AgentSession(...).rehydrate(context)
+        if not self._sess_data:
+            return self._session
+
+        state = pickle.loads(self._sess_data)
+        await self._session.rehydrate(state)
+        logger.debug("session rehydrated")
+        return self._session
+
+    @property
+    def text(self) -> str:
+        return self._text
+
+    @property
+    def result(self) -> RunResult | None:
+        return self._result
+
+
 class JobContext:
     _PARTICIPANT_ENTRYPOINT_CALLBACK = Callable[
         ["JobContext", rtc.RemoteParticipant], Coroutine[None, None, None]

livekit-agents/livekit/agents/voice/agent.py

@@ -93,6 +93,7 @@ def __init__(
 
         self._mcp_servers = mcp_servers
         self._activity: AgentActivity | None = None
+        self._rehydrated = False
 
     @property
     def id(self) -> str:
@@ -359,6 +360,35 @@ def _get_activity_or_raise(self) -> AgentActivity:
 
         return self._activity
 
+    def __getstate__(self) -> dict[str, Any]:
+        tool_ctx = llm.ToolContext(self.tools)
+        return {
+            "tools": list(tool_ctx.function_tools.keys()),
+            "chat_ctx": self.chat_ctx.to_dict(),
+        }
+
+    def __setstate__(self, state: dict[str, Any]) -> None:
+        try:
+            self.__init__()
+        except TypeError as e:
+            logger.error("Agent rehydration requires a zero-argument constructor")
+            raise e
+
+        tool_ctx = llm.ToolContext(self.tools)
+        valid_tools: list[llm.FunctionTool | llm.RawFunctionTool] = []
+        for name in state["tools"]:
+            if name in tool_ctx.function_tools:
+                valid_tools.append(tool_ctx.function_tools[name])
+            else:
+                logger.warning("tool not found when unpickling", extra={"missing_tool": name})
+
+        # TODO: serialize the ongoing speech tasks?
+        # TODO: support AgentTask
+
+        self._tools = valid_tools
+        self._chat_ctx = llm.ChatContext.from_dict(state["chat_ctx"])
+        self._rehydrated = True  # skip on_enter when rehydrating
+
     class default:
         @staticmethod
         async def stt_node(

livekit-agents/livekit/agents/voice/agent_activity.py

@@ -474,10 +474,11 @@ async def _traceable_on_enter() -> None:
                     finally:
                         _OnEnterContextVar.reset(tk)
 
-                self._on_enter_task = task = self._create_speech_task(
-                    _traceable_on_enter(), name="AgentTask_on_enter"
-                )
-                _set_activity_task_info(task, inline_task=True)
+                if not self._agent._rehydrated:
+                    self._on_enter_task = task = self._create_speech_task(
+                        _traceable_on_enter(), name="AgentTask_on_enter"
+                    )
+                    _set_activity_task_info(task, inline_task=True)
             finally:
                 start_span.end()
 

livekit-agents/livekit/agents/voice/agent_session.py

@@ -10,6 +10,7 @@
 from types import TracebackType
 from typing import (
     TYPE_CHECKING,
+    Any,
     Generic,
     Literal,
     Optional,
@@ -838,6 +839,32 @@ async def _aclose_impl(
     async def aclose(self) -> None:
         await self._aclose_impl(reason=CloseReason.USER_INITIATED)
 
+    def get_state(self) -> dict[str, Any]:
+        tool_ctx = llm.ToolContext(self.tools)
+        return {
+            "tools": list(tool_ctx.function_tools.keys()),
+            "chat_ctx": self._chat_ctx.to_dict(),
+            "agent": self._agent,
+        }
+
+    async def rehydrate(self, state: dict[str, Any]) -> None:
+        tool_ctx = llm.ToolContext(self.tools)
+        valid_tools: list[llm.FunctionTool | llm.RawFunctionTool] = []
+        for name in state["tools"]:
+            if name in tool_ctx.function_tools:
+                valid_tools.append(tool_ctx.function_tools[name])
+            else:
+                logger.warning("tool not found when unpickling", extra={"missing_tool": name})
+
+        self._tools = valid_tools
+        self._chat_ctx = llm.ChatContext.from_dict(state["chat_ctx"])
+
+        if self._started:
+            # only allow rehydrate session that not started yet?
+            await self._update_activity(state["agent"])
+        else:
+            await self.start(agent=state["agent"])
+
     def update_options(
         self,
         *,
@@ -1076,8 +1103,7 @@ async def _update_activity(
                 await self._activity.resume()
 
         # move it outside the lock to allow calling _update_activity in on_enter of a new agent
-        if wait_on_enter:
-            assert self._activity._on_enter_task is not None
+        if wait_on_enter and self._activity._on_enter_task:
             await asyncio.shield(self._activity._on_enter_task)
 
     @utils.log_exceptions(logger=logger)