219 lines
8.0 KiB
Python
219 lines
8.0 KiB
Python
"""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)
|