Merge pull request #907 from pipecat-ai/khk/gemini-20241221

Gemini unary API fixes and natural conversation demo
This commit is contained in:
Kwindla Hultman Kramer
2024-12-23 17:34:57 -08:00
committed by GitHub
3 changed files with 540 additions and 85 deletions

View File

@@ -10,6 +10,7 @@ import sys
import time
import aiohttp
import google.ai.generativelanguage as glm
from dotenv import load_dotenv
from loguru import logger
from runner import configure
@@ -20,6 +21,8 @@ from pipecat.frames.frames import (
EndFrame,
Frame,
InputAudioRawFrame,
LLMFullResponseEndFrame,
LLMFullResponseStartFrame,
LLMMessagesFrame,
StartFrame,
StartInterruptionFrame,
@@ -34,6 +37,7 @@ from pipecat.pipeline.parallel_pipeline import ParallelPipeline
from pipecat.pipeline.pipeline import Pipeline
from pipecat.pipeline.runner import PipelineRunner
from pipecat.pipeline.task import PipelineParams, PipelineTask
from pipecat.processors.aggregators.llm_response import LLMResponseAggregator
from pipecat.processors.aggregators.openai_llm_context import (
OpenAILLMContext,
OpenAILLMContextFrame,
@@ -53,39 +57,321 @@ 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"
classifier_statement = """You are an audio language classifier model. You are receiving audio from a user in a WebRTC call. Your job is to decide whether the user has finished speaking or not.
TRANSCRIBER_MODEL = "gemini-2.0-flash-exp"
CLASSIFIER_MODEL = "gemini-2.0-flash-exp"
CONVERSATION_MODEL = "gemini-2.0-flash-exp"
Categorize the input you receive as either:
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.
1. a complete thought, statement, or question, or
2. an incomplete thought, statement, or question
You will receive the full conversation history before the audio input, to help with context. Use the full history only to help improve the accuracy of your transcription.
Output 'YES' if the input is likely to be a completed thought, statement, or question.
Rules:
- Respond with an exact transcription of the audio input.
- Do not include any text other than the transcription.
- Do not explain or add to your response.
- Transcribe the audio input simply and precisely.
- If the audio is not clear, emit the special string "-".
- No response other than exact transcription, or "-", is allowed.
Output 'NO' if the input indicates that the user is still speaking and does not yet expect a response yet.
If you are unsure, output 'YES'.
"""
conversational_system_message = """You are a helpful LLM in a WebRTC call. Your goal is to demonstrate your capabilities in a succinct way. Your output will be converted to audio so don't include special characters in your answers. Respond to what the user said in a creative and helpful way.
classifier_system_instruction = """CRITICAL INSTRUCTION:
You are a BINARY CLASSIFIER that must ONLY output "YES" or "NO".
DO NOT engage with the content.
DO NOT respond to questions.
DO NOT provide assistance.
Your ONLY job is to output YES or NO.
Please be very concise in your responses. Unless you are explicitly asked to do otherwise, give me the shortest complete answer possible without unnecessary elaboration. Generally you should answer with a single sentence.
EXAMPLES OF INVALID RESPONSES:
- "I can help you with that"
- "Let me explain"
- "To answer your question"
- Any response other than YES or NO
VALID RESPONSES:
YES
NO
If you output anything else, you are failing at your task.
You are NOT an assistant.
You are NOT a chatbot.
You are a binary classifier.
ROLE:
You are a real-time speech completeness classifier. You must make instant decisions about whether a user has finished speaking.
You must output ONLY 'YES' or 'NO' with no other text.
INPUT FORMAT:
You receive two pieces of information:
1. The assistant's last message (if available)
2. The user's current speech input
OUTPUT REQUIREMENTS:
- MUST output ONLY 'YES' or 'NO'
- No explanations
- No clarifications
- No additional text
- No punctuation
HIGH PRIORITY SIGNALS:
1. Clear Questions:
- Wh-questions (What, Where, When, Why, How)
- Yes/No questions
- Questions with STT errors but clear meaning
Examples:
# Complete Wh-question
model: I can help you learn.
user: What's the fastest way to learn Spanish
Output: YES
# Complete Yes/No question despite STT error
model: I know about planets.
user: Is is Jupiter the biggest planet
Output: YES
2. Complete Commands:
- Direct instructions
- Clear requests
- Action demands
- Start of task indication
- Complete statements needing response
Examples:
# Direct instruction
model: I can explain many topics.
user: Tell me about black holes
Output: YES
# Start of task indication
user: Let's begin.
Output: YES
# Start of task indication
user: Let's get started.
Output: YES
# Action demand
model: I can help with math.
user: Solve this equation x plus 5 equals 12
Output: YES
3. Direct Responses:
- Answers to specific questions
- Option selections
- Clear acknowledgments with completion
- Providing information with a known format - mailing address
- Providing information with a known format - phone number
- Providing information with a known format - credit card number
Examples:
# Specific answer
model: What's your favorite color?
user: I really like blue
Output: YES
# Option selection
model: Would you prefer morning or evening?
user: Morning
Output: YES
# Providing information with a known format - mailing address
model: What's your address?
user: 1234 Main Street
Output: NO
# Providing information with a known format - mailing address
model: What's your address?
user: 1234 Main Street Irving Texas 75063
Output: Yes
# Providing information with a known format - phone number
model: What's your phone number?
user: 41086753
Output: NO
# Providing information with a known format - phone number
model: What's your phone number?
user: 4108675309
Output: Yes
# Providing information with a known format - phone number
model: What's your phone number?
user: 220
Output: No
# Providing information with a known format - credit card number
model: What's your credit card number?
user: 5556
Output: NO
# Providing information with a known format - phone number
model: What's your credit card number?
user: 5556710454680800
Output: Yes
model: What's your credit card number?
user: 414067
Output: NO
MEDIUM PRIORITY SIGNALS:
1. Speech Pattern Completions:
- Self-corrections reaching completion
- False starts with clear ending
- Topic changes with complete thought
- Mid-sentence completions
Examples:
# Self-correction reaching completion
model: What would you like to know?
user: Tell me about... no wait, explain how rainbows form
Output: YES
# Topic change with complete thought
model: The weather is nice today.
user: Actually can you tell me who invented the telephone
Output: YES
# Mid-sentence completion
model: Hello I'm ready.
user: What's the capital of? France
Output: YES
2. Context-Dependent Brief Responses:
- Acknowledgments (okay, sure, alright)
- Agreements (yes, yeah)
- Disagreements (no, nah)
- Confirmations (correct, exactly)
Examples:
# Acknowledgment
model: Should we talk about history?
user: Sure
Output: YES
# Disagreement with completion
model: Is that what you meant?
user: No not really
Output: YES
LOW PRIORITY SIGNALS:
1. STT Artifacts (Consider but don't over-weight):
- Repeated words
- Unusual punctuation
- Capitalization errors
- Word insertions/deletions
Examples:
# Word repetition but complete
model: I can help with that.
user: What what is the time right now
Output: YES
# Missing punctuation but complete
model: I can explain that.
user: Please tell me how computers work
Output: YES
2. Speech Features:
- Filler words (um, uh, like)
- Thinking pauses
- Word repetitions
- Brief hesitations
Examples:
# Filler words but complete
model: What would you like to know?
user: Um uh how do airplanes fly
Output: YES
# Thinking pause but incomplete
model: I can explain anything.
user: Well um I want to know about the
Output: NO
DECISION RULES:
1. Return YES if:
- ANY high priority signal shows clear completion
- Medium priority signals combine to show completion
- Meaning is clear despite low priority artifacts
2. Return NO if:
- No high priority signals present
- Thought clearly trails off
- Multiple incomplete indicators
- User appears mid-formulation
3. When uncertain:
- If you can understand the intent → YES
- If meaning is unclear → NO
- Always make a binary decision
- Never request clarification
Examples:
# Incomplete despite corrections
model: What would you like to know about?
user: Can you tell me about
Output: NO
# Complete despite multiple artifacts
model: I can help you learn.
user: How do you I mean what's the best way to learn programming
Output: YES
# Trailing off incomplete
model: I can explain anything.
user: I was wondering if you could tell me why
Output: NO
"""
conversation_system_instruction = """You are a helpful assistant participating in a voice converation.
Your goal is to demonstrate your capabilities in a succinct way. Your output will be converted to audio so don't include special characters in your answers. Respond to what the user said in a creative and helpful way.
If you know that a number string is a phone number from the context of the conversation, write it as a phone number. For example 210-333-4567.
If you know that a number string is a credit card number, write it as a credit card number. For example 4111-1111-1111-1111.
Please be very concise in your responses. Unless you are explicitly asked to do otherwise, give me shortest complete answer possible without unnecessary elaboration. Generally you should answer with a single sentence.
"""
class StatementJudgeAudioContextAccumulator(FrameProcessor):
def __init__(self, *, notifier: BaseNotifier, **kwargs):
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._notifier = notifier
self._audio_frames = []
self._audio_frames = []
self._start_secs = 0.2 # this should match VAD start_secs (hardcoding for now)
self._user_speaking = False
self._max_buffer_size_secs = 30
self._user_speaking_vad_state = False
self._user_speaking_utterance_state = False
async def reset(self):
self._audio_frames = []
self._user_speaking = False
self._user_speaking_vad_state = False
self._user_speaking_utterance_state = False
async def process_frame(self, frame: Frame, direction: FrameDirection):
await super().process_frame(frame, direction)
@@ -99,22 +385,33 @@ class StatementJudgeAudioContextAccumulator(FrameProcessor):
# but let's leave that as an exercise to the reader. :-)
return
if isinstance(frame, UserStartedSpeakingFrame):
self._user_speaking = True
self._user_speaking_vad_state = True
self._user_speaking_utterance_state = True
elif isinstance(frame, UserStoppedSpeakingFrame):
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}"
)
self._user_speaking = False
context = GoogleLLMContext()
context.set_messages([{"role": "system", "content": classifier_statement}])
context.add_audio_frames_message(audio_frames=self._audio_frames)
context.add_audio_frames_message(text="Audio follows", audio_frames=self._audio_frames)
await self.push_frame(OpenAILLMContextFrame(context=context))
elif isinstance(frame, InputAudioRawFrame):
if self._user_speaking:
self._audio_frames.append(frame)
# Append the audio frame to our buffer. Treat the buffer as a ring buffer, dropping the oldest
# frames as necessary.
# Use a small buffer size when an utterance is not in progress. Just big enough to backfill the start_secs.
# Use a larger buffer size when an utterance is in progress.
# Assume all audio frames have the same duration.
self._audio_frames.append(frame)
frame_duration = len(frame.audio) / 2 * frame.num_channels / frame.sample_rate
buffer_duration = frame_duration * len(self._audio_frames)
# logger.debug(f"!!! Frame duration: {frame_duration}")
if self._user_speaking_utterance_state:
while buffer_duration > self._max_buffer_size_secs:
self._audio_frames.pop(0)
buffer_duration -= frame_duration
else:
# Append the audio frame to our buffer. Treat the buffer as a ring buffer, dropping the oldest
# frames as necessary. Assume all audio frames have the same duration.
self._audio_frames.append(frame)
frame_duration = len(frame.audio) / 16 * frame.num_channels / frame.sample_rate
buffer_duration = frame_duration * len(self._audio_frames)
while buffer_duration > self._start_secs:
self._audio_frames.pop(0)
buffer_duration -= frame_duration
@@ -123,32 +420,143 @@ class StatementJudgeAudioContextAccumulator(FrameProcessor):
class CompletenessCheck(FrameProcessor):
def __init__(
self, notifier: BaseNotifier, audio_accumulator: StatementJudgeAudioContextAccumulator
):
"""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):
super().__init__()
self._notifier = notifier
self._audio_accumulator = audio_accumulator
self._idle_task = None
self._wakeup_time = 0
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:
self._idle_task.cancel()
await self.push_frame(UserStoppedSpeakingFrame())
await self._audio_accumulator.reset()
await self._notifier.notify()
elif isinstance(frame, TextFrame):
if frame.text.strip():
logger.debug(f"Completeness check NO - '{frame.text}'")
# start timer to wake up if necessary
if self._wakeup_time:
self._wakeup_time = time.time() + self.wait_time
else:
# 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())
async def _idle_task_handler(self):
try:
while time.time() < self._wakeup_time:
await asyncio.sleep(0.01)
# 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):
def __init__(self, notifier: BaseNotifier, **kwargs):
"""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,
context: OpenAILLMContext,
user_transcription_buffer: "UserAggregatorBuffer",
**kwargs,
):
super().__init__(**kwargs)
self._gate_open = False
self._frames_buffer = []
self._notifier = notifier
self._context = context
self._transcription_buffer = user_transcription_buffer
def close_gate(self):
self._gate_open = False
@@ -176,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
@@ -194,12 +609,22 @@ class OutputGate(FrameProcessor):
while True:
try:
await self._notifier.wait()
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:
await self.push_frame(frame, direction)
self._frames_buffer = []
except asyncio.CancelledError:
break
except Exception as e:
logger.error(f"OutputGate error: {e}")
raise e
break
async def main():
@@ -215,64 +640,63 @@ async def main():
vad_enabled=True,
vad_analyzer=SileroVADAnalyzer(),
vad_audio_passthrough=True,
audio_in_sample_rate=16000,
),
)
stt = DeepgramSTTService(api_key=os.getenv("DEEPGRAM_API_KEY"))
tts = CartesiaTTSService(
api_key=os.getenv("CARTESIA_API_KEY"),
voice_id="79a125e8-cd45-4c13-8a67-188112f4dd22", # British Lady
)
# This is the LLM that will be used to detect if the user has finished a
# statement. This doesn't really need to be an LLM, we could use NLP
# libraries for that, but we have the machinery to use an LLM, so we might as well!
statement_llm = GoogleLLMService(
model="gemini-1.5-flash-latest", api_key=os.getenv("GOOGLE_API_KEY")
# This is the LLM that will transcribe user speech.
tx_llm = GoogleLLMService(
name="Transcriber",
model=TRANSCRIBER_MODEL,
api_key=os.getenv("GOOGLE_API_KEY"),
temperature=0.0,
system_instruction=transcriber_system_instruction,
)
# This is the regular LLM.
llm = GoogleLLMService(model="gemini-1.5-flash-latest", api_key=os.getenv("GOOGLE_API_KEY"))
# This is the LLM that will classify user speech as complete or incomplete.
classifier_llm = GoogleLLMService(
name="Classifier",
model=CLASSIFIER_MODEL,
api_key=os.getenv("GOOGLE_API_KEY"),
temperature=0.0,
system_instruction=classifier_system_instruction,
)
messages = [
{
"role": "system",
"content": conversational_system_message,
},
]
# This is the regular LLM that responds conversationally.
conversation_llm = GoogleLLMService(
name="Conversation",
model=CONVERSATION_MODEL,
api_key=os.getenv("GOOGLE_API_KEY"),
system_instruction=conversation_system_instruction,
)
context = OpenAILLMContext(messages)
context_aggregator = llm.create_context_aggregator(context)
context = OpenAILLMContext()
context_aggregator = conversation_llm.create_context_aggregator(context)
# We have instructed the LLM to return 'YES' if it thinks the user
# completed a sentence. So, if it's 'YES' we will return true in this
# We have instructed the LLM to return 'True' if it thinks the user
# completed a sentence. So, if it's 'True' we will return true in this
# predicate which will wake up the notifier.
async def wake_check_filter(frame):
return frame.text == "YES"
return frame.text == "True"
# This is a notifier that we use to synchronize the two LLMs.
notifier = EventNotifier()
# This turns the LLM context into an inference request to classify the user's speech
# as complete or incomplete.
statement_judge_context_filter = StatementJudgeAudioContextAccumulator(notifier=notifier)
# statement_judge_context_filter = StatementJudgeAudioContextAccumulator(notifier=notifier)
audio_accumulater = AudioAccumulator()
# This sends a UserStoppedSpeakingFrame and triggers the notifier event
completeness_check = CompletenessCheck(
notifier=notifier, audio_accumulator=statement_judge_context_filter
notifier=notifier, audio_accumulator=audio_accumulater
)
# # Notify if the user hasn't said anything.
async def user_idle_notifier(frame):
await notifier.notify()
# Sometimes the LLM will fail detecting if a user has completed a
# sentence, this will wake up the notifier if that happens.
user_idle = UserIdleProcessor(callback=user_idle_notifier, timeout=5.0)
bot_output_gate = OutputGate(notifier=notifier)
async def block_user_stopped_speaking(frame):
return not isinstance(frame, UserStoppedSpeakingFrame)
@@ -284,9 +708,18 @@ async def main():
or isinstance(frame, StopInterruptionFrame)
)
conversation_audio_context_assembler = ConversationAudioContextAssembler(context=context)
user_aggregator_buffer = UserAggregatorBuffer()
bot_output_gate = OutputGate(
notifier=notifier, context=context, user_transcription_buffer=user_aggregator_buffer
)
pipeline = Pipeline(
[
transport.input(),
audio_accumulater,
ParallelPipeline(
[
# Pass everything except UserStoppedSpeaking to the elements after
@@ -294,24 +727,28 @@ async def main():
FunctionFilter(filter=block_user_stopped_speaking),
],
[
statement_judge_context_filter,
statement_llm,
completeness_check,
ParallelPipeline(
[
classifier_llm,
completeness_check,
],
[
tx_llm,
user_aggregator_buffer,
],
)
],
[
stt,
context_aggregator.user(),
# Block everything except OpenAILLMContextFrame and LLMMessagesFrame
FunctionFilter(filter=pass_only_llm_trigger_frames),
llm,
bot_output_gate, # Buffer all llm/tts output until notified.
conversation_audio_context_assembler,
conversation_llm,
bot_output_gate, # buffer output until notified, then flush frames and update context
# TempPrinter(),
],
),
tts,
user_idle,
transport.output(),
context_aggregator.assistant(),
]
],
)
task = PipelineTask(

View File

@@ -230,7 +230,6 @@ class GeminiMultimodalLiveLLMService(LLMService):
async def start(self, frame: StartFrame):
await super().start(frame)
await self._connect()
async def stop(self, frame: EndFrame):
await super().stop(frame)
@@ -434,8 +433,9 @@ class GeminiMultimodalLiveLLMService(LLMService):
async def _ws_send(self, message):
# logger.debug(f"Sending message to websocket: {message}")
try:
if self._websocket:
await self._websocket.send(json.dumps(message))
if not self._websocket:
await self._connect()
await self._websocket.send(json.dumps(message))
except Exception as e:
if self._disconnecting:
return

View File

@@ -593,6 +593,8 @@ class GoogleLLMService(LLMService):
model: str = "gemini-1.5-flash-latest",
params: InputParams = InputParams(),
system_instruction: Optional[str] = None,
tools: Optional[List[Dict[str, Any]]] = None,
tool_config: Optional[Dict[str, Any]] = None,
**kwargs,
):
super().__init__(**kwargs)
@@ -607,6 +609,8 @@ class GoogleLLMService(LLMService):
"top_p": params.top_p,
"extra": params.extra if isinstance(params.extra, dict) else {},
}
self._tools = tools
self._tool_config = tool_config
def can_generate_metrics(self) -> bool:
return True
@@ -625,7 +629,8 @@ class GoogleLLMService(LLMService):
try:
logger.debug(
f"Generating chat: {self._system_instruction} | {context.get_messages_for_logging()}"
# f"Generating chat: {self._system_instruction} | {context.get_messages_for_logging()}"
f"Generating chat: {context.get_messages_for_logging()}"
)
messages = context.messages
@@ -649,28 +654,41 @@ class GoogleLLMService(LLMService):
generation_config = GenerationConfig(**generation_params) if generation_params else None
await self.start_ttfb_metrics()
tools = context.tools if context.tools else []
tools = []
if context.tools:
tools = context.tools
elif self._tools:
tools = self._tools
tool_config = None
if self._tool_config:
tool_config = self._tool_config
response = await self._client.generate_content_async(
contents=messages, tools=tools, stream=True, generation_config=generation_config
contents=messages,
tools=tools,
stream=True,
generation_config=generation_config,
tool_config=tool_config,
)
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:
await self.push_frame(TextFrame(c.text))
elif c.function_call:
logger.debug(f"!!! Function call: {c.function_call}")
args = type(c.function_call).to_dict(c.function_call).get("args", {})
await self.call_function(
context=context,