ultimatepos/Modules/Communication/agent/gemini/audio.py

50 lines
1.4 KiB
Python

"""Audio frame plumbing for Gemini Live Translate."""
from __future__ import annotations
import logging
from collections.abc import AsyncIterator
from livekit import rtc
from gemini.config import AUDIO_CHANNELS, GEMINI_INPUT_SAMPLE_RATE, GEMINI_OUTPUT_SAMPLE_RATE
logger = logging.getLogger("gemini.audio")
async def iter_pcm_for_gemini(track: rtc.RemoteAudioTrack) -> AsyncIterator[bytes]:
stream = rtc.AudioStream(
track,
sample_rate=GEMINI_INPUT_SAMPLE_RATE,
num_channels=AUDIO_CHANNELS,
)
try:
async for ev in stream:
yield bytes(ev.frame.data)
finally:
await stream.aclose()
def make_audio_source() -> rtc.AudioSource:
return rtc.AudioSource(GEMINI_OUTPUT_SAMPLE_RATE, AUDIO_CHANNELS)
async def push_pcm_to_source(source: rtc.AudioSource, pcm_bytes: bytes) -> None:
import array
samples = array.array("h")
samples.frombytes(pcm_bytes)
frame = rtc.AudioFrame(
data=samples.tobytes(),
sample_rate=GEMINI_OUTPUT_SAMPLE_RATE,
num_channels=AUDIO_CHANNELS,
samples_per_channel=len(samples),
)
try:
await source.capture_frame(frame)
except Exception as exc:
if "closed" in str(exc).lower() or "invalidstate" in str(exc).lower():
logger.debug("AudioSource closed mid-capture; dropping frame")
else:
raise