"""Reconciles Gemini Live translation sessions to room demand.""" from __future__ import annotations import asyncio import logging from livekit import rtc from gemini.config import ( NATIVE_LANG, PARTICIPANT_LANG_ATTR, RECONCILE_DEBOUNCE_SEC, SESSION_GRACE_SEC, gemini_lang, ) from gemini.session import GeminiSession logger = logging.getLogger("gemini.router") SessionKey = tuple[str, str] def participant_lang(participant: rtc.Participant) -> str | None: attrs = participant.attributes or {} raw = attrs.get(PARTICIPANT_LANG_ATTR) or attrs.get("language") if not raw or raw == NATIVE_LANG: return None return gemini_lang(raw) class TranslationRouter: def __init__(self, room: rtc.Room, gemini_api_key: str) -> None: self._room = room self._gemini_api_key = gemini_api_key self._speaker_tracks: dict[str, rtc.RemoteAudioTrack] = {} self._sessions: dict[SessionKey, GeminiSession] = {} self._grace_tasks: dict[SessionKey, asyncio.Task] = {} self._detached_tasks: set[asyncio.Task] = set() self._reconcile_handle: asyncio.TimerHandle | None = None self._reconcile_lock = asyncio.Lock() def start(self) -> None: room = self._room @room.on("participant_connected") def _on_conn(_: rtc.RemoteParticipant) -> None: self._schedule_reconcile() @room.on("participant_disconnected") def _on_disc(p: rtc.RemoteParticipant) -> None: self._on_participant_left(p.identity) self._schedule_reconcile() @room.on("participant_attributes_changed") def _on_attrs(_changed: dict[str, str], _p: rtc.Participant) -> None: self._schedule_reconcile() @room.on("track_subscribed") def _on_subscribed( track: rtc.Track, _pub: rtc.RemoteTrackPublication, participant: rtc.RemoteParticipant, ) -> None: if participant.kind == rtc.ParticipantKind.PARTICIPANT_KIND_AGENT: return if track.kind == rtc.TrackKind.KIND_AUDIO and isinstance(track, rtc.RemoteAudioTrack): self._speaker_tracks[participant.identity] = track self._schedule_reconcile() @room.on("track_unsubscribed") def _on_unsubscribed( track: rtc.Track, _pub: rtc.RemoteTrackPublication, participant: rtc.RemoteParticipant, ) -> None: if track.kind == rtc.TrackKind.KIND_AUDIO: self._speaker_tracks.pop(participant.identity, None) self._schedule_reconcile() @room.on("track_muted") def _on_muted(_pub: rtc.TrackPublication, _p: rtc.Participant) -> None: self._schedule_reconcile() @room.on("track_unmuted") def _on_unmuted(_pub: rtc.TrackPublication, _p: rtc.Participant) -> None: self._schedule_reconcile() for p in room.remote_participants.values(): if p.kind == rtc.ParticipantKind.PARTICIPANT_KIND_AGENT: continue for pub in p.track_publications.values(): if pub.track and pub.kind == rtc.TrackKind.KIND_AUDIO and isinstance(pub.track, rtc.RemoteAudioTrack): self._speaker_tracks[p.identity] = pub.track self._schedule_reconcile() async def aclose(self) -> None: if self._reconcile_handle: self._reconcile_handle.cancel() self._reconcile_handle = None for task in self._grace_tasks.values(): task.cancel() self._grace_tasks.clear() await asyncio.gather(*(s.aclose() for s in self._sessions.values()), return_exceptions=True) self._sessions.clear() def _schedule_reconcile(self) -> None: loop = asyncio.get_event_loop() if self._reconcile_handle is not None: self._reconcile_handle.cancel() self._reconcile_handle = loop.call_later( RECONCILE_DEBOUNCE_SEC, lambda: asyncio.create_task(self._reconcile()), ) async def _reconcile(self) -> None: async with self._reconcile_lock: desired = self._compute_desired_sessions() existing = set(self._sessions.keys()) for key in desired & set(self._grace_tasks.keys()): self._grace_tasks.pop(key).cancel() for key in existing - desired: if key not in self._grace_tasks: self._grace_tasks[key] = asyncio.create_task(self._grace_teardown(key)) for key in desired - existing: if key in self._grace_tasks: continue speaker_identity, target_lang = key track = self._speaker_tracks.get(speaker_identity) if track is None: continue session = GeminiSession( room=self._room, speaker_identity=speaker_identity, speaker_track=track, target_lang=target_lang, gemini_api_key=self._gemini_api_key, ) self._sessions[key] = session try: await session.start() except Exception: logger.exception("failed to start gemini session %s -> %s", speaker_identity, target_lang) self._sessions.pop(key, None) def _compute_desired_sessions(self) -> set[SessionKey]: target_langs = self._listener_target_langs() if not target_langs: return set() desired: set[SessionKey] = set() for speaker_identity, source_lang in self._active_speakers(): for tgt in target_langs: if tgt == source_lang: continue desired.add((speaker_identity, tgt)) return desired def _listener_target_langs(self) -> set[str]: langs: set[str] = set() for p in self._room.remote_participants.values(): if p.kind == rtc.ParticipantKind.PARTICIPANT_KIND_AGENT: continue lang = participant_lang(p) if lang: langs.add(lang) return langs def _active_speakers(self) -> list[tuple[str, str]]: out: list[tuple[str, str]] = [] for p in self._room.remote_participants.values(): if p.kind == rtc.ParticipantKind.PARTICIPANT_KIND_AGENT: continue lang = participant_lang(p) if not lang: continue if p.identity not in self._speaker_tracks: continue if not self._has_unmuted_mic(p): continue out.append((p.identity, lang)) return out def _has_unmuted_mic(self, p: rtc.RemoteParticipant) -> bool: for pub in p.track_publications.values(): if pub.kind == rtc.TrackKind.KIND_AUDIO and not pub.muted: return True return False async def _grace_teardown(self, key: SessionKey) -> None: try: await asyncio.sleep(SESSION_GRACE_SEC) except asyncio.CancelledError: return if key in self._sessions and key not in self._compute_desired_sessions(): session = self._sessions.pop(key) await session.aclose() self._grace_tasks.pop(key, None) def _on_participant_left(self, identity: str) -> None: self._speaker_tracks.pop(identity, None) for key in list(self._sessions.keys()): if key[0] == identity: session = self._sessions.pop(key) pending = self._grace_tasks.pop(key, None) if pending: pending.cancel() task = asyncio.create_task(session.aclose()) self._detached_tasks.add(task) task.add_done_callback(self._detached_tasks.discard)