218 lines
7.9 KiB
Python
218 lines
7.9 KiB
Python
"""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)
|