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.remove(0)
|
||||||
logger.add(sys.stderr, level="DEBUG")
|
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
|
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.
|
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):
|
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):
|
def __init__(self, **kwargs):
|
||||||
super().__init__(**kwargs)
|
super().__init__(**kwargs)
|
||||||
self._audio_frames = []
|
self._audio_frames = []
|
||||||
@@ -376,14 +389,6 @@ class AudioAccumulator(FrameProcessor):
|
|||||||
self._user_speaking_utterance_state = True
|
self._user_speaking_utterance_state = True
|
||||||
|
|
||||||
elif isinstance(frame, UserStoppedSpeakingFrame):
|
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)
|
data = b"".join(frame.audio for frame in self._audio_frames)
|
||||||
logger.debug(
|
logger.debug(
|
||||||
f"Processing audio buffer seconds: ({len(self._audio_frames)}) ({len(data)}) {len(data) / 2 / 16000}"
|
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):
|
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
|
wait_time = 5.0
|
||||||
|
|
||||||
def __init__(self, notifier: BaseNotifier, audio_accumulator: AudioAccumulator, **kwargs):
|
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):
|
async def process_frame(self, frame: Frame, direction: FrameDirection):
|
||||||
await super().process_frame(frame, direction)
|
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")
|
logger.debug("Completeness check YES")
|
||||||
if self._idle_task:
|
if self._idle_task:
|
||||||
logger.debug(f"CompletenessCheck idle wait CANCEL")
|
|
||||||
self._idle_task.cancel()
|
self._idle_task.cancel()
|
||||||
self._idle_task = None
|
|
||||||
await self.push_frame(UserStoppedSpeakingFrame())
|
await self.push_frame(UserStoppedSpeakingFrame())
|
||||||
await self._audio_accumulator.reset()
|
await self._audio_accumulator.reset()
|
||||||
await self._notifier.notify()
|
await self._notifier.notify()
|
||||||
@@ -443,7 +455,7 @@ class CompletenessCheck(FrameProcessor):
|
|||||||
if self._wakeup_time:
|
if self._wakeup_time:
|
||||||
self._wakeup_time = time.time() + self.wait_time
|
self._wakeup_time = time.time() + self.wait_time
|
||||||
else:
|
else:
|
||||||
logger.debug("CompletenessCheck idle wait START")
|
# logger.debug("!!! CompletenessCheck idle wait START")
|
||||||
self._wakeup_time = time.time() + self.wait_time
|
self._wakeup_time = time.time() + self.wait_time
|
||||||
self._idle_task = self.get_event_loop().create_task(self._idle_task_handler())
|
self._idle_task = self.get_event_loop().create_task(self._idle_task_handler())
|
||||||
|
|
||||||
@@ -451,16 +463,87 @@ class CompletenessCheck(FrameProcessor):
|
|||||||
try:
|
try:
|
||||||
while time.time() < self._wakeup_time:
|
while time.time() < self._wakeup_time:
|
||||||
await asyncio.sleep(0.01)
|
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()
|
await self._notifier.notify()
|
||||||
except asyncio.CancelledError:
|
except asyncio.CancelledError:
|
||||||
|
# logger.debug(f"!!! CompletenessCheck idle wait CANCEL")
|
||||||
pass
|
pass
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
logger.error(f"CompletenessCheck idle wait error: {e}")
|
logger.error(f"CompletenessCheck idle wait error: {e}")
|
||||||
raise 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):
|
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__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
notifier: BaseNotifier,
|
notifier: BaseNotifier,
|
||||||
@@ -501,6 +584,13 @@ class OutputGate(FrameProcessor):
|
|||||||
await self.push_frame(frame, direction)
|
await self.push_frame(frame, direction)
|
||||||
return
|
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:
|
if self._gate_open:
|
||||||
await self.push_frame(frame, direction)
|
await self.push_frame(frame, direction)
|
||||||
return
|
return
|
||||||
@@ -517,16 +607,13 @@ class OutputGate(FrameProcessor):
|
|||||||
|
|
||||||
async def _gate_task_handler(self):
|
async def _gate_task_handler(self):
|
||||||
while True:
|
while True:
|
||||||
# logger.debug("!!! Waiting for notifier")
|
|
||||||
try:
|
try:
|
||||||
await self._notifier.wait()
|
await self._notifier.wait()
|
||||||
|
|
||||||
# logger.debug("!!! Notified")
|
transcription = await self._transcription_buffer.wait_for_transcription() or "-"
|
||||||
transcription = await self._transcription_buffer.wait_for_transcription()
|
self._context._messages.append(
|
||||||
|
glm.Content(role="user", parts=[glm.Part(text=transcription)])
|
||||||
last_message = self._context.messages[-1]
|
)
|
||||||
if last_message.role == "user":
|
|
||||||
last_message.parts = [glm.Part(text=transcription)]
|
|
||||||
|
|
||||||
self.open_gate()
|
self.open_gate()
|
||||||
for frame, direction in self._frames_buffer:
|
for frame, direction in self._frames_buffer:
|
||||||
@@ -540,54 +627,6 @@ class OutputGate(FrameProcessor):
|
|||||||
break
|
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 def main():
|
||||||
async with aiohttp.ClientSession() as session:
|
async with aiohttp.ClientSession() as session:
|
||||||
(room_url, _) = await configure(session)
|
(room_url, _) = await configure(session)
|
||||||
@@ -613,7 +652,7 @@ async def main():
|
|||||||
# This is the LLM that will transcribe user speech.
|
# This is the LLM that will transcribe user speech.
|
||||||
tx_llm = GoogleLLMService(
|
tx_llm = GoogleLLMService(
|
||||||
name="Transcriber",
|
name="Transcriber",
|
||||||
model="gemini-2.0-flash-exp",
|
model=TRANSCRIBER_MODEL,
|
||||||
api_key=os.getenv("GOOGLE_API_KEY"),
|
api_key=os.getenv("GOOGLE_API_KEY"),
|
||||||
temperature=0.0,
|
temperature=0.0,
|
||||||
system_instruction=transcriber_system_instruction,
|
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.
|
# This is the LLM that will classify user speech as complete or incomplete.
|
||||||
classifier_llm = GoogleLLMService(
|
classifier_llm = GoogleLLMService(
|
||||||
name="Classifier",
|
name="Classifier",
|
||||||
model="gemini-2.0-flash-exp",
|
model=CLASSIFIER_MODEL,
|
||||||
api_key=os.getenv("GOOGLE_API_KEY"),
|
api_key=os.getenv("GOOGLE_API_KEY"),
|
||||||
temperature=0.0,
|
temperature=0.0,
|
||||||
system_instruction=classifier_system_instruction,
|
system_instruction=classifier_system_instruction,
|
||||||
@@ -631,7 +670,7 @@ async def main():
|
|||||||
# This is the regular LLM that responds conversationally.
|
# This is the regular LLM that responds conversationally.
|
||||||
conversation_llm = GoogleLLMService(
|
conversation_llm = GoogleLLMService(
|
||||||
name="Conversation",
|
name="Conversation",
|
||||||
model="gemini-2.0-flash-exp",
|
model=CONVERSATION_MODEL,
|
||||||
api_key=os.getenv("GOOGLE_API_KEY"),
|
api_key=os.getenv("GOOGLE_API_KEY"),
|
||||||
system_instruction=conversation_system_instruction,
|
system_instruction=conversation_system_instruction,
|
||||||
)
|
)
|
||||||
|
|||||||
@@ -635,7 +635,7 @@ class GoogleLLMService(LLMService):
|
|||||||
|
|
||||||
messages = context.messages
|
messages = context.messages
|
||||||
if context.system_message and self._system_instruction != context.system_message:
|
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._system_instruction = context.system_message
|
||||||
self._create_client()
|
self._create_client()
|
||||||
|
|
||||||
@@ -673,15 +673,16 @@ class GoogleLLMService(LLMService):
|
|||||||
await self.stop_ttfb_metrics()
|
await self.stop_ttfb_metrics()
|
||||||
|
|
||||||
if response.usage_metadata:
|
if response.usage_metadata:
|
||||||
|
# Use only the prompt token count from the response object
|
||||||
prompt_tokens = response.usage_metadata.prompt_token_count
|
prompt_tokens = response.usage_metadata.prompt_token_count
|
||||||
completion_tokens = response.usage_metadata.candidates_token_count
|
total_tokens = prompt_tokens
|
||||||
total_tokens = response.usage_metadata.total_token_count
|
|
||||||
|
|
||||||
async for chunk in response:
|
async for chunk in response:
|
||||||
if chunk.usage_metadata:
|
if chunk.usage_metadata:
|
||||||
prompt_tokens += response.usage_metadata.prompt_token_count
|
# Use only the completion_tokens from the chunks. Prompt tokens are already counted and
|
||||||
completion_tokens += response.usage_metadata.candidates_token_count
|
# are repeated here.
|
||||||
total_tokens += response.usage_metadata.total_token_count
|
completion_tokens += chunk.usage_metadata.candidates_token_count
|
||||||
|
total_tokens += chunk.usage_metadata.candidates_token_count
|
||||||
try:
|
try:
|
||||||
for c in chunk.parts:
|
for c in chunk.parts:
|
||||||
if c.text:
|
if c.text:
|
||||||
|
|||||||
Reference in New Issue
Block a user