feature complete gemini audio, transcription, and phrase endpointing demo
This commit is contained in:
@@ -57,6 +57,14 @@ load_dotenv(override=True)
|
||||
logger.remove(0)
|
||||
logger.add(sys.stderr, level="DEBUG")
|
||||
|
||||
# TRANSCRIBER_MODEL = "gemini-1.5-flash-latest"
|
||||
# CLASSIFIER_MODEL = "gemini-1.5-flash-latest"
|
||||
# CONVERSATION_MODEL = "gemini-1.5-flash-latest"
|
||||
|
||||
TRANSCRIBER_MODEL = "gemini-2.0-flash-exp"
|
||||
CLASSIFIER_MODEL = "gemini-2.0-flash-exp"
|
||||
CONVERSATION_MODEL = "gemini-2.0-flash-exp"
|
||||
|
||||
transcriber_system_instruction = """You are an audio transcriber. You are receiving audio from a user. Your job is to
|
||||
transcribe the input audio to text exactly as it was said by the user.
|
||||
|
||||
@@ -347,6 +355,11 @@ Please be very concise in your responses. Unless you are explicitly asked to do
|
||||
|
||||
|
||||
class AudioAccumulator(FrameProcessor):
|
||||
"""Buffers user audio until the user stops speaking.
|
||||
|
||||
Always pushes a fresh context with a single audio message.
|
||||
"""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self._audio_frames = []
|
||||
@@ -376,14 +389,6 @@ class AudioAccumulator(FrameProcessor):
|
||||
self._user_speaking_utterance_state = True
|
||||
|
||||
elif isinstance(frame, UserStoppedSpeakingFrame):
|
||||
if self._audio_frames[-1]:
|
||||
fr = self._audio_frames[-1]
|
||||
frame_duration = len(fr.audio) / 2 * fr.num_channels / fr.sample_rate
|
||||
|
||||
logger.debug(
|
||||
f"!!! Frame duration: ({len(fr.audio)}) ({fr.num_channels}) ({fr.sample_rate}) {frame_duration}"
|
||||
)
|
||||
|
||||
data = b"".join(frame.audio for frame in self._audio_frames)
|
||||
logger.debug(
|
||||
f"Processing audio buffer seconds: ({len(self._audio_frames)}) ({len(data)}) {len(data) / 2 / 16000}"
|
||||
@@ -415,6 +420,12 @@ class AudioAccumulator(FrameProcessor):
|
||||
|
||||
|
||||
class CompletenessCheck(FrameProcessor):
|
||||
"""Checks the result of the classifier LLM to determine if the user has finished speaking.
|
||||
|
||||
Triggers the notifier if the user has finished speaking. Also triggers the notifier if an
|
||||
idle timeout is reached.
|
||||
"""
|
||||
|
||||
wait_time = 5.0
|
||||
|
||||
def __init__(self, notifier: BaseNotifier, audio_accumulator: AudioAccumulator, **kwargs):
|
||||
@@ -427,12 +438,13 @@ class CompletenessCheck(FrameProcessor):
|
||||
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
||||
await super().process_frame(frame, direction)
|
||||
|
||||
if isinstance(frame, TextFrame) and frame.text.startswith("YES"):
|
||||
if isinstance(frame, UserStartedSpeakingFrame):
|
||||
if self._idle_task:
|
||||
self._idle_task.cancel()
|
||||
elif isinstance(frame, TextFrame) and frame.text.startswith("YES"):
|
||||
logger.debug("Completeness check YES")
|
||||
if self._idle_task:
|
||||
logger.debug(f"CompletenessCheck idle wait CANCEL")
|
||||
self._idle_task.cancel()
|
||||
self._idle_task = None
|
||||
await self.push_frame(UserStoppedSpeakingFrame())
|
||||
await self._audio_accumulator.reset()
|
||||
await self._notifier.notify()
|
||||
@@ -443,7 +455,7 @@ class CompletenessCheck(FrameProcessor):
|
||||
if self._wakeup_time:
|
||||
self._wakeup_time = time.time() + self.wait_time
|
||||
else:
|
||||
logger.debug("CompletenessCheck idle wait START")
|
||||
# logger.debug("!!! CompletenessCheck idle wait START")
|
||||
self._wakeup_time = time.time() + self.wait_time
|
||||
self._idle_task = self.get_event_loop().create_task(self._idle_task_handler())
|
||||
|
||||
@@ -451,16 +463,87 @@ class CompletenessCheck(FrameProcessor):
|
||||
try:
|
||||
while time.time() < self._wakeup_time:
|
||||
await asyncio.sleep(0.01)
|
||||
logger.debug(f"CompletenessCheck idle wait OVER")
|
||||
# logger.debug(f"!!! CompletenessCheck idle wait OVER")
|
||||
await self._audio_accumulator.reset()
|
||||
await self._notifier.notify()
|
||||
except asyncio.CancelledError:
|
||||
# logger.debug(f"!!! CompletenessCheck idle wait CANCEL")
|
||||
pass
|
||||
except Exception as e:
|
||||
logger.error(f"CompletenessCheck idle wait error: {e}")
|
||||
raise e
|
||||
finally:
|
||||
# logger.debug(f"!!! CompletenessCheck idle wait FINALLY")
|
||||
self._wakeup_time = 0
|
||||
self._idle_task = None
|
||||
|
||||
|
||||
class UserAggregatorBuffer(LLMResponseAggregator):
|
||||
"""Buffers the output of the transcription LLM. Used by the bot output gate."""
|
||||
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(
|
||||
messages=None,
|
||||
role=None,
|
||||
start_frame=LLMFullResponseStartFrame,
|
||||
end_frame=LLMFullResponseEndFrame,
|
||||
accumulator_frame=TextFrame,
|
||||
handle_interruptions=True,
|
||||
expect_stripped_words=False,
|
||||
)
|
||||
self._transcription = ""
|
||||
|
||||
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
||||
await super().process_frame(frame, direction)
|
||||
# parent method pushes frames
|
||||
if isinstance(frame, UserStartedSpeakingFrame):
|
||||
self._transcription = ""
|
||||
|
||||
async def _push_aggregation(self):
|
||||
if self._aggregation:
|
||||
self._transcription = self._aggregation
|
||||
self._aggregation = ""
|
||||
|
||||
logger.debug(f"[Transcription] {self._transcription}")
|
||||
|
||||
async def wait_for_transcription(self):
|
||||
while not self._transcription:
|
||||
await asyncio.sleep(0.01)
|
||||
tx = self._transcription
|
||||
self._transcription = ""
|
||||
return tx
|
||||
|
||||
|
||||
class ConversationAudioContextAssembler(FrameProcessor):
|
||||
"""Takes the single-message context generated by the AudioAccumulator and adds it to the conversation LLM's context."""
|
||||
|
||||
def __init__(self, context: OpenAILLMContext, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self._context = context
|
||||
|
||||
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
||||
await super().process_frame(frame, direction)
|
||||
|
||||
# We must not block system frames.
|
||||
if isinstance(frame, SystemFrame):
|
||||
await self.push_frame(frame, direction)
|
||||
return
|
||||
|
||||
if isinstance(frame, OpenAILLMContextFrame):
|
||||
GoogleLLMContext.upgrade_to_google(self._context)
|
||||
last_message = frame.context.messages[-1]
|
||||
self._context._messages.append(last_message)
|
||||
await self.push_frame(OpenAILLMContextFrame(context=self._context))
|
||||
|
||||
|
||||
class OutputGate(FrameProcessor):
|
||||
"""Buffers output frames until the notifier is triggered.
|
||||
|
||||
When the notifier fires, waits until a transcription is ready, then:
|
||||
1. Replaces the last user audio message with the transcription.
|
||||
2. Flushes the frames buffer.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
notifier: BaseNotifier,
|
||||
@@ -501,6 +584,13 @@ class OutputGate(FrameProcessor):
|
||||
await self.push_frame(frame, direction)
|
||||
return
|
||||
|
||||
if isinstance(frame, LLMFullResponseStartFrame):
|
||||
# Remove the audio message from the context. We will never need it again.
|
||||
# If the completeness check fails, a new audio message will be appended to the context.
|
||||
# If the completeness check succeeds, our notifier will fire and we will append the
|
||||
# transcription to the context.
|
||||
self._context._messages.pop()
|
||||
|
||||
if self._gate_open:
|
||||
await self.push_frame(frame, direction)
|
||||
return
|
||||
@@ -517,16 +607,13 @@ class OutputGate(FrameProcessor):
|
||||
|
||||
async def _gate_task_handler(self):
|
||||
while True:
|
||||
# logger.debug("!!! Waiting for notifier")
|
||||
try:
|
||||
await self._notifier.wait()
|
||||
|
||||
# logger.debug("!!! Notified")
|
||||
transcription = await self._transcription_buffer.wait_for_transcription()
|
||||
|
||||
last_message = self._context.messages[-1]
|
||||
if last_message.role == "user":
|
||||
last_message.parts = [glm.Part(text=transcription)]
|
||||
transcription = await self._transcription_buffer.wait_for_transcription() or "-"
|
||||
self._context._messages.append(
|
||||
glm.Content(role="user", parts=[glm.Part(text=transcription)])
|
||||
)
|
||||
|
||||
self.open_gate()
|
||||
for frame, direction in self._frames_buffer:
|
||||
@@ -540,54 +627,6 @@ class OutputGate(FrameProcessor):
|
||||
break
|
||||
|
||||
|
||||
class ConversationAudioContextAssembler(FrameProcessor):
|
||||
def __init__(self, context: OpenAILLMContext, **kwargs):
|
||||
super().__init__(**kwargs)
|
||||
self._context = context
|
||||
|
||||
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
||||
await super().process_frame(frame, direction)
|
||||
|
||||
# We must not block system frames.
|
||||
if isinstance(frame, SystemFrame):
|
||||
await self.push_frame(frame, direction)
|
||||
return
|
||||
|
||||
if isinstance(frame, OpenAILLMContextFrame):
|
||||
GoogleLLMContext.upgrade_to_google(self._context)
|
||||
last_message = frame.context.messages[-1]
|
||||
self._context._messages.append(last_message)
|
||||
await self.push_frame(OpenAILLMContextFrame(context=self._context))
|
||||
|
||||
|
||||
class UserAggregatorBuffer(LLMResponseAggregator):
|
||||
def __init__(self, **kwargs):
|
||||
super().__init__(
|
||||
messages=None,
|
||||
role=None,
|
||||
start_frame=LLMFullResponseStartFrame,
|
||||
end_frame=LLMFullResponseEndFrame,
|
||||
accumulator_frame=TextFrame,
|
||||
handle_interruptions=True,
|
||||
expect_stripped_words=False,
|
||||
)
|
||||
self._transcription = ""
|
||||
|
||||
async def _push_aggregation(self):
|
||||
if self._aggregation:
|
||||
self._transcription = self._aggregation
|
||||
self._aggregation = ""
|
||||
|
||||
logger.debug(f"[Transcription] {self._transcription}")
|
||||
|
||||
async def wait_for_transcription(self):
|
||||
while not self._transcription:
|
||||
await asyncio.sleep(0.01)
|
||||
tx = self._transcription
|
||||
self._transcription = ""
|
||||
return tx
|
||||
|
||||
|
||||
async def main():
|
||||
async with aiohttp.ClientSession() as session:
|
||||
(room_url, _) = await configure(session)
|
||||
@@ -613,7 +652,7 @@ async def main():
|
||||
# This is the LLM that will transcribe user speech.
|
||||
tx_llm = GoogleLLMService(
|
||||
name="Transcriber",
|
||||
model="gemini-2.0-flash-exp",
|
||||
model=TRANSCRIBER_MODEL,
|
||||
api_key=os.getenv("GOOGLE_API_KEY"),
|
||||
temperature=0.0,
|
||||
system_instruction=transcriber_system_instruction,
|
||||
@@ -622,7 +661,7 @@ async def main():
|
||||
# This is the LLM that will classify user speech as complete or incomplete.
|
||||
classifier_llm = GoogleLLMService(
|
||||
name="Classifier",
|
||||
model="gemini-2.0-flash-exp",
|
||||
model=CLASSIFIER_MODEL,
|
||||
api_key=os.getenv("GOOGLE_API_KEY"),
|
||||
temperature=0.0,
|
||||
system_instruction=classifier_system_instruction,
|
||||
@@ -631,7 +670,7 @@ async def main():
|
||||
# This is the regular LLM that responds conversationally.
|
||||
conversation_llm = GoogleLLMService(
|
||||
name="Conversation",
|
||||
model="gemini-2.0-flash-exp",
|
||||
model=CONVERSATION_MODEL,
|
||||
api_key=os.getenv("GOOGLE_API_KEY"),
|
||||
system_instruction=conversation_system_instruction,
|
||||
)
|
||||
|
||||
@@ -635,7 +635,7 @@ class GoogleLLMService(LLMService):
|
||||
|
||||
messages = context.messages
|
||||
if context.system_message and self._system_instruction != context.system_message:
|
||||
# logger.debug(f"System instruction changed: {context.system_message}")
|
||||
logger.debug(f"System instruction changed: {context.system_message}")
|
||||
self._system_instruction = context.system_message
|
||||
self._create_client()
|
||||
|
||||
@@ -673,15 +673,16 @@ class GoogleLLMService(LLMService):
|
||||
await self.stop_ttfb_metrics()
|
||||
|
||||
if response.usage_metadata:
|
||||
# Use only the prompt token count from the response object
|
||||
prompt_tokens = response.usage_metadata.prompt_token_count
|
||||
completion_tokens = response.usage_metadata.candidates_token_count
|
||||
total_tokens = response.usage_metadata.total_token_count
|
||||
total_tokens = prompt_tokens
|
||||
|
||||
async for chunk in response:
|
||||
if chunk.usage_metadata:
|
||||
prompt_tokens += response.usage_metadata.prompt_token_count
|
||||
completion_tokens += response.usage_metadata.candidates_token_count
|
||||
total_tokens += response.usage_metadata.total_token_count
|
||||
# Use only the completion_tokens from the chunks. Prompt tokens are already counted and
|
||||
# are repeated here.
|
||||
completion_tokens += chunk.usage_metadata.candidates_token_count
|
||||
total_tokens += chunk.usage_metadata.candidates_token_count
|
||||
try:
|
||||
for c in chunk.parts:
|
||||
if c.text:
|
||||
|
||||
Reference in New Issue
Block a user