"""One Gemini Live session per (speaker, target language) pair.""" from __future__ import annotations import asyncio import base64 import contextlib import json import logging import random import websockets from livekit import rtc from gemini.audio import iter_pcm_for_gemini, make_audio_source, push_pcm_to_source from gemini.config import ( GEMINI_INPUT_SAMPLE_RATE, GEMINI_MAX_FAILURES_BEFORE_LONG_BACKOFF, GEMINI_MODEL, GEMINI_RECONNECT_BACKOFF_SEC, ) logger = logging.getLogger("gemini.session") GEMINI_WS_URL = ( "wss://generativelanguage.googleapis.com/ws/" "google.ai.generativelanguage.v1beta.GenerativeService.BidiGenerateContent" ) class GeminiSession: def __init__( self, *, room: rtc.Room, speaker_identity: str, speaker_track: rtc.RemoteAudioTrack, target_lang: str, gemini_api_key: str, ) -> None: self._room = room self._speaker_identity = speaker_identity self._speaker_track = speaker_track self._target_lang = target_lang self._gemini_api_key = gemini_api_key self._audio_source = make_audio_source() self._local_track: rtc.LocalAudioTrack | None = None self._track_sid: str | None = None self._consecutive_failures = 0 self._tasks: list[asyncio.Task] = [] self._closed = asyncio.Event() async def start(self) -> None: track_name = f"tx:{self._speaker_identity}:{self._target_lang}" self._local_track = rtc.LocalAudioTrack.create_audio_track(track_name, self._audio_source) publish_opts = rtc.TrackPublishOptions(source=rtc.TrackSource.SOURCE_MICROPHONE) pub = await self._room.local_participant.publish_track(self._local_track, publish_opts) self._track_sid = pub.sid logger.info( "started gemini track sid=%s name=%s for %s -> %s", self._track_sid, track_name, self._speaker_identity, self._target_lang, ) self._tasks.append(asyncio.create_task(self._run(), name=f"gemini/{track_name}")) async def aclose(self) -> None: if self._closed.is_set(): return self._closed.set() for task in self._tasks: task.cancel() for task in self._tasks: with contextlib.suppress(asyncio.CancelledError, Exception): await task self._tasks.clear() if self._track_sid: with contextlib.suppress(Exception): await self._room.local_participant.unpublish_track(self._track_sid) with contextlib.suppress(Exception): await self._audio_source.aclose() async def _run(self) -> None: while not self._closed.is_set(): try: await self._connect_and_pump() return except asyncio.CancelledError: raise except Exception as exc: self._consecutive_failures += 1 idx = min(self._consecutive_failures - 1, len(GEMINI_RECONNECT_BACKOFF_SEC) - 1) delay = GEMINI_RECONNECT_BACKOFF_SEC[idx] + random.uniform(0, GEMINI_RECONNECT_BACKOFF_SEC[idx] * 0.2) if self._consecutive_failures >= GEMINI_MAX_FAILURES_BEFORE_LONG_BACKOFF: logger.error( "Gemini session %s -> %s failed %d times", self._speaker_identity, self._target_lang, self._consecutive_failures, ) logger.warning( "Gemini session error (%s -> %s): %s; backoff %.2fs", self._speaker_identity, self._target_lang, exc, delay, ) try: await asyncio.wait_for(self._closed.wait(), timeout=delay) return except asyncio.TimeoutError: pass async def _connect_and_pump(self) -> None: url = f"{GEMINI_WS_URL}?key={self._gemini_api_key}" async with websockets.connect(url, max_size=2**22, ping_interval=20, ping_timeout=20) as ws: await ws.send(json.dumps(self._build_setup_payload())) setup_complete = asyncio.Event() send_task = asyncio.create_task(self._pump_input(ws, setup_complete)) recv_task = asyncio.create_task(self._pump_output(ws, setup_complete)) done, pending = await asyncio.wait({send_task, recv_task}, return_when=asyncio.FIRST_EXCEPTION) for task in pending: task.cancel() for task in done: exc = task.exception() if exc is not None: raise exc def _build_setup_payload(self) -> dict: return { "setup": { "model": f"models/{GEMINI_MODEL}", "outputAudioTranscription": {}, "generationConfig": { "responseModalities": ["AUDIO"], "translationConfig": { "targetLanguageCode": self._target_lang, "echoTargetLanguage": False, }, }, "realtimeInputConfig": { "automaticActivityDetection": {"disabled": False}, }, } } async def _pump_input(self, ws: websockets.WebSocketClientProtocol, setup_complete: asyncio.Event) -> None: await setup_complete.wait() mime = f"audio/pcm;rate={GEMINI_INPUT_SAMPLE_RATE}" async for pcm in iter_pcm_for_gemini(self._speaker_track): if self._closed.is_set(): return await ws.send( json.dumps( { "realtimeInput": { "audio": { "mimeType": mime, "data": base64.b64encode(pcm).decode("ascii"), } } } ) ) async def _pump_output(self, ws: websockets.WebSocketClientProtocol, setup_complete: asyncio.Event) -> None: async for raw in ws: if self._closed.is_set(): return try: msg = json.loads(raw) except json.JSONDecodeError: continue if msg.get("setupComplete") is not None: self._consecutive_failures = 0 setup_complete.set() continue sc = msg.get("serverContent") if not sc: continue model_turn = sc.get("modelTurn") if model_turn is not None: for part in model_turn.get("parts", []) or []: inline = part.get("inlineData") if inline and inline.get("data"): await push_pcm_to_source(self._audio_source, base64.b64decode(inline["data"])) ot = sc.get("outputTranscription") if ot and ot.get("text"): await self._publish_transcript(ot["text"], final=False) if sc.get("turnComplete"): await self._publish_transcript("", final=True) async def _publish_transcript(self, text: str, *, final: bool) -> None: if not text and not final: return try: writer = await self._room.local_participant.stream_text( topic="lk.translation", sender_identity=self._speaker_identity, attributes={ "target_lang": self._target_lang, "source_identity": self._speaker_identity, "final": "true" if final else "false", }, ) if text: await writer.write(text) await writer.aclose() except Exception as exc: logger.debug("text-stream publish failed: %s", exc)