wehub-resource-sync
388 lines
12 KiB
Python
388 lines
12 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import contextvars
|
|
import os
|
|
import signal
|
|
import threading
|
|
from types import FrameType
|
|
from typing import TYPE_CHECKING, Any, Literal
|
|
|
|
from ..job import JobExecutorType
|
|
from ..log import logger
|
|
from ..voice import AgentSession, io
|
|
from ..voice.transcription import TranscriptSynchronizer
|
|
from ..worker import AgentServer, WorkerOptions
|
|
from . import proto
|
|
from .log import setup_logging
|
|
|
|
if TYPE_CHECKING:
|
|
from ..voice.remote_session import TcpSessionTransport
|
|
from .tcp_console import TcpAudioInput, TcpAudioOutput
|
|
|
|
HANDLED_SIGNALS = (
|
|
signal.SIGINT,
|
|
signal.SIGTERM,
|
|
)
|
|
|
|
|
|
class _ExitCli(BaseException):
|
|
pass
|
|
|
|
|
|
ConsoleMode = Literal["text", "audio"]
|
|
|
|
|
|
class AgentsConsole:
|
|
"""Minimal console stub for TCP console mode (Go CLI handles the TUI)."""
|
|
|
|
_instance: AgentsConsole | None = None
|
|
|
|
@classmethod
|
|
def get_instance(cls) -> AgentsConsole:
|
|
if cls._instance is None:
|
|
cls._instance = cls()
|
|
return cls._instance
|
|
|
|
def __init__(self) -> None:
|
|
import datetime
|
|
import pathlib
|
|
|
|
self._lock = threading.Lock()
|
|
self._io_acquired = False
|
|
self._io_acquired_event = threading.Event()
|
|
self._enabled = False
|
|
self._record = False
|
|
self._console_mode: ConsoleMode = "audio"
|
|
self._tcp_transport: TcpSessionTransport | None = None
|
|
self._tcp_audio_input: TcpAudioInput | None = None
|
|
self._tcp_audio_output: TcpAudioOutput | None = None
|
|
self._session_directory = pathlib.Path(
|
|
"console-recordings",
|
|
f"session-{datetime.datetime.now().strftime('%m-%d-%H%M%S')}",
|
|
)
|
|
|
|
def acquire_io(self, *, loop: asyncio.AbstractEventLoop, session: AgentSession) -> None:
|
|
with self._lock:
|
|
if self._io_acquired:
|
|
raise RuntimeError("the ConsoleIO was already acquired by another session")
|
|
|
|
if asyncio.get_running_loop() != loop:
|
|
raise RuntimeError(
|
|
"the ConsoleIO must be acquired in the same asyncio loop as the session"
|
|
)
|
|
|
|
self._io_acquired = True
|
|
self._io_loop = loop
|
|
self._io_context = contextvars.copy_context()
|
|
|
|
assert self._tcp_transport is not None
|
|
assert self._tcp_audio_input is not None
|
|
assert self._tcp_audio_output is not None
|
|
self._io_audio_input = self._tcp_audio_input
|
|
self._io_audio_output = self._tcp_audio_output
|
|
|
|
self._io_transcription_sync = TranscriptSynchronizer(
|
|
next_in_chain_audio=self._io_audio_output,
|
|
next_in_chain_text=None,
|
|
)
|
|
self._io_acquired_event.set()
|
|
self._io_session = session
|
|
|
|
if session:
|
|
self._update_sess_io(
|
|
session,
|
|
self.console_mode,
|
|
self._io_audio_input,
|
|
self._io_transcription_sync.audio_output,
|
|
self._io_transcription_sync.text_output,
|
|
)
|
|
|
|
@property
|
|
def enabled(self) -> bool:
|
|
return self._enabled
|
|
|
|
@enabled.setter
|
|
def enabled(self, val: bool) -> None:
|
|
self._enabled = val
|
|
|
|
@property
|
|
def record(self) -> bool:
|
|
return self._record
|
|
|
|
@record.setter
|
|
def record(self, val: bool) -> None:
|
|
self._record = val
|
|
|
|
@property
|
|
def session_directory(self) -> Any:
|
|
return self._session_directory
|
|
|
|
@property
|
|
def io_acquired(self) -> bool:
|
|
with self._lock:
|
|
return self._io_acquired
|
|
|
|
@property
|
|
def io_session(self) -> AgentSession:
|
|
if not self._io_acquired:
|
|
raise RuntimeError("AgentsConsole is not acquired")
|
|
return self._io_session
|
|
|
|
@property
|
|
def io_loop(self) -> asyncio.AbstractEventLoop:
|
|
if not self._io_acquired:
|
|
raise RuntimeError("AgentsConsole is not acquired")
|
|
return self._io_loop
|
|
|
|
@property
|
|
def io_context(self) -> contextvars.Context:
|
|
if not self._io_acquired:
|
|
raise RuntimeError("AgentsConsole is not acquired")
|
|
return self._io_context
|
|
|
|
def wait_for_io_acquisition(self) -> None:
|
|
self._io_acquired_event.wait()
|
|
|
|
@property
|
|
def console_mode(self) -> ConsoleMode:
|
|
return self._console_mode
|
|
|
|
@console_mode.setter
|
|
def console_mode(self, mode: ConsoleMode) -> None:
|
|
with self._lock:
|
|
self._console_mode = mode
|
|
|
|
if not self._io_acquired:
|
|
return
|
|
|
|
self.io_loop.call_soon_threadsafe(
|
|
self._update_sess_io,
|
|
self.io_session,
|
|
mode,
|
|
self._io_audio_input,
|
|
self._io_transcription_sync.audio_output,
|
|
self._io_transcription_sync.text_output,
|
|
)
|
|
|
|
def _update_sess_io(
|
|
self,
|
|
sess: AgentSession,
|
|
mode: ConsoleMode,
|
|
audio_input: io.AudioInput,
|
|
audio_output: io.AudioOutput,
|
|
text_output: io.TextOutput,
|
|
) -> None:
|
|
if asyncio.get_running_loop() != self.io_loop:
|
|
raise RuntimeError("_update_sess_io must be executed on the io_loop")
|
|
|
|
with self._lock:
|
|
if not self._io_acquired:
|
|
return
|
|
|
|
if self._io_session != sess or self._console_mode != mode:
|
|
return
|
|
|
|
if mode == "text":
|
|
sess.input.audio = None
|
|
sess.output.audio = None
|
|
sess.output.transcription = None
|
|
else:
|
|
sess.input.audio = audio_input
|
|
sess.output.audio = audio_output
|
|
sess.output.transcription = text_output
|
|
|
|
|
|
def _run_tcp_console(*, server: AgentServer, connect_addr: str, record: bool = False) -> None:
|
|
"""Run console in TCP mode — connects to the Go CLI's TCP server."""
|
|
from ..voice.remote_session import TcpSessionTransport
|
|
from .tcp_console import TcpAudioInput, TcpAudioOutput
|
|
|
|
host, port_str = connect_addr.rsplit(":", 1)
|
|
port = int(port_str)
|
|
|
|
setup_logging("DEBUG", devmode=True, console=True)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
|
|
tcp_audio_input: TcpAudioInput | None = None
|
|
|
|
async def _tcp_main() -> None:
|
|
nonlocal tcp_audio_input
|
|
transport = TcpSessionTransport(host, port)
|
|
|
|
server._job_executor_type = JobExecutorType.THREAD
|
|
|
|
console_inst = AgentsConsole.get_instance()
|
|
console_inst.enabled = True
|
|
console_inst.record = record
|
|
console_inst._tcp_transport = transport
|
|
tcp_audio_input = TcpAudioInput()
|
|
console_inst._tcp_audio_input = tcp_audio_input
|
|
console_inst._tcp_audio_output = TcpAudioOutput(transport)
|
|
|
|
@server.once("worker_started")
|
|
def _simulate_job() -> None:
|
|
asyncio.run_coroutine_threadsafe(
|
|
server.simulate_job("console-room", agent_identity="console", fake_job=True),
|
|
loop,
|
|
)
|
|
|
|
try:
|
|
await server.run(devmode=True, unregistered=True)
|
|
finally:
|
|
await transport.close()
|
|
|
|
exit_triggered = False
|
|
|
|
async def _graceful_shutdown() -> None:
|
|
if tcp_audio_input is not None:
|
|
tcp_audio_input.close()
|
|
await server.aclose()
|
|
|
|
def _handle_exit(sig: int, frame: FrameType | None) -> None:
|
|
nonlocal exit_triggered
|
|
if exit_triggered:
|
|
os.killpg(os.getpgid(0), signal.SIGKILL)
|
|
exit_triggered = True
|
|
asyncio.run_coroutine_threadsafe(_graceful_shutdown(), loop)
|
|
|
|
for sig in HANDLED_SIGNALS:
|
|
signal.signal(sig, _handle_exit)
|
|
|
|
try:
|
|
loop.run_until_complete(_tcp_main())
|
|
finally:
|
|
for sig in HANDLED_SIGNALS:
|
|
signal.signal(sig, lambda *_: os._exit(1))
|
|
|
|
try:
|
|
tasks = asyncio.all_tasks(loop)
|
|
for task in tasks:
|
|
task.cancel()
|
|
|
|
loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
try:
|
|
loop.run_until_complete(loop.shutdown_asyncgens())
|
|
loop.run_until_complete(loop.shutdown_default_executor())
|
|
except Exception:
|
|
pass
|
|
loop.close()
|
|
|
|
|
|
def _run_worker(server: AgentServer, args: proto.CliArgs) -> None:
|
|
kwargs: dict = {}
|
|
if args.url:
|
|
kwargs["ws_url"] = args.url
|
|
if args.api_key:
|
|
kwargs["api_key"] = args.api_key
|
|
if args.api_secret:
|
|
kwargs["api_secret"] = args.api_secret
|
|
if kwargs:
|
|
server.update_options(**kwargs)
|
|
|
|
if args.simulation:
|
|
server._simulation = True
|
|
|
|
if args.cli_addr and not args.dev:
|
|
raise ValueError("--cli-addr requires --dev")
|
|
|
|
devmode = args.dev
|
|
colored_logs = devmode or args.log_format == "colored"
|
|
|
|
exit_raised = False
|
|
|
|
def _handle_exit(sig: int, frame: FrameType | None) -> None:
|
|
nonlocal exit_raised
|
|
if exit_raised:
|
|
os._exit(1)
|
|
exit_raised = True
|
|
raise _ExitCli()
|
|
|
|
for sig in HANDLED_SIGNALS:
|
|
signal.signal(sig, _handle_exit)
|
|
|
|
setup_logging(args.log_level, devmode=colored_logs, console=False, compact=args.simulation)
|
|
|
|
loop = asyncio.new_event_loop()
|
|
asyncio.set_event_loop(loop)
|
|
|
|
loop.slow_callback_duration = 0.1 # 100ms
|
|
|
|
async def _worker_run(worker: AgentServer) -> None:
|
|
try:
|
|
await server.run(devmode=devmode, unregistered=False)
|
|
except Exception:
|
|
logger.exception("worker failed")
|
|
|
|
watch_client = None
|
|
if args.cli_addr:
|
|
from .watcher import WatchClient
|
|
|
|
watch_client = WatchClient(server, args.cli_addr, loop=loop)
|
|
watch_client.start()
|
|
|
|
try:
|
|
main_task = loop.create_task(_worker_run(server), name="worker_main_task_cli")
|
|
try:
|
|
loop.run_until_complete(main_task)
|
|
except _ExitCli:
|
|
pass
|
|
|
|
# Second Ctrl+C force-exits.
|
|
def _force_exit(sig: int, frame: FrameType | None) -> None:
|
|
logger.warning("exiting forcefully", extra={"signal": sig})
|
|
os._exit(1)
|
|
|
|
for sig in HANDLED_SIGNALS:
|
|
signal.signal(sig, _force_exit)
|
|
|
|
try:
|
|
if not devmode:
|
|
try:
|
|
loop.run_until_complete(server.drain())
|
|
except asyncio.TimeoutError:
|
|
logger.warning("drain timed out, forcing shutdown")
|
|
loop.run_until_complete(server.aclose())
|
|
|
|
if watch_client:
|
|
loop.run_until_complete(watch_client.aclose())
|
|
except _ExitCli:
|
|
pass # stray from first signal — ignore
|
|
finally:
|
|
# Re-enable force exit for the final cleanup phase
|
|
for sig in HANDLED_SIGNALS:
|
|
signal.signal(sig, lambda *_: os._exit(1))
|
|
|
|
try:
|
|
tasks = asyncio.all_tasks(loop)
|
|
for task in tasks:
|
|
task.cancel()
|
|
|
|
loop.run_until_complete(asyncio.gather(*tasks, return_exceptions=True))
|
|
except Exception:
|
|
pass
|
|
finally:
|
|
try:
|
|
loop.run_until_complete(loop.shutdown_asyncgens())
|
|
loop.run_until_complete(loop.shutdown_default_executor())
|
|
except Exception:
|
|
pass
|
|
loop.close()
|
|
|
|
|
|
def run_app(server: AgentServer | WorkerOptions) -> None:
|
|
"""Run the agent via the (deprecated) rich Python CLI.
|
|
|
|
This is the default entry used by ``python myagent.py <command>``. The rich
|
|
CLI lives in ``._legacy`` and is being phased out in favor of the LiveKit CLI
|
|
(``lk agent ...``) and the thin interface in ``livekit.agents.__main__``.
|
|
"""
|
|
from . import _legacy
|
|
|
|
_legacy.run_app(server)
|