Refactor Gemini tracing to more closely match OpenAI Realtime, add TTFB metrics

This commit is contained in:
Mark Backman
2025-05-29 23:39:13 -04:00
parent dd1f7d0875
commit 8602557985
4 changed files with 125 additions and 76 deletions

View File

@@ -12,7 +12,7 @@ from loguru import logger
from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
from pipecat.adapters.schemas.function_schema import FunctionSchema
from pipecat.adapters.schemas.tools_schema import AdapterType, ToolsSchema
from pipecat.adapters.schemas.tools_schema import ToolsSchema
from pipecat.audio.vad.silero import SileroVADAnalyzer
from pipecat.frames.frames import TTSSpeakFrame
from pipecat.pipeline.pipeline import Pipeline
@@ -21,12 +21,10 @@ from pipecat.pipeline.task import PipelineParams, PipelineTask
from pipecat.processors.aggregators.openai_llm_context import OpenAILLMContext
from pipecat.services.cartesia.tts import CartesiaTTSService
from pipecat.services.deepgram.stt import DeepgramSTTService
from pipecat.services.gemini_multimodal_live.gemini import GeminiMultimodalLiveLLMService
from pipecat.services.llm_service import FunctionCallParams
from pipecat.services.openai.llm import OpenAILLMService
from pipecat.transports.base_transport import BaseTransport, TransportParams
from pipecat.transports.services.daily import DailyParams
from pipecat.utils.time import time_now_iso8601
from pipecat.utils.tracing.setup import setup_tracing
load_dotenv(override=True)
@@ -47,6 +45,11 @@ if IS_TRACING_ENABLED:
logger.info("OpenTelemetry tracing initialized")
async def fetch_weather_from_api(params: FunctionCallParams):
await params.llm.push_frame(TTSSpeakFrame("Let me check on that."))
await params.result_callback({"conditions": "nice", "temperature": "75"})
# We store functions so objects (e.g. SileroVADAnalyzer) don't get
# instantiated. The function will be called when the desired transport gets
# selected.
@@ -69,29 +72,24 @@ transport_params = {
}
async def fetch_weather_from_api(params: FunctionCallParams):
temperature = 75 if params.arguments["format"] == "fahrenheit" else 24
await params.result_callback(
{
"conditions": "nice",
"temperature": temperature,
"format": params.arguments["format"],
"timestamp": time_now_iso8601(),
}
)
system_instruction = """
You are a helpful assistant who can answer questions and use tools.
You have a tool called "get_current_weather" that can be used to get the current weather. If the user asks
for the weather, call this function.
"""
async def run_example(transport: BaseTransport, _: argparse.Namespace, handle_sigint: bool):
logger.info(f"Starting bot")
stt = DeepgramSTTService(api_key=os.getenv("DEEPGRAM_API_KEY"))
tts = CartesiaTTSService(
api_key=os.getenv("CARTESIA_API_KEY"),
voice_id="71a7ad14-091c-4e8e-a314-022ece01c121", # British Reading Lady
)
llm = OpenAILLMService(
api_key=os.getenv("OPENAI_API_KEY"), params=OpenAILLMService.InputParams(temperature=0.5)
)
# You can also register a function_name of None to get all functions
# sent to the same callback with an additional function_name parameter.
llm.register_function("get_current_weather", fetch_weather_from_api)
weather_function = FunctionSchema(
name="get_current_weather",
description="Get the current weather",
@@ -108,29 +106,25 @@ async def run_example(transport: BaseTransport, _: argparse.Namespace, handle_si
},
required=["location", "format"],
)
search_tool = {"google_search": {}}
tools = ToolsSchema(
standard_tools=[weather_function], custom_tools={AdapterType.GEMINI: [search_tool]}
)
tools = ToolsSchema(standard_tools=[weather_function])
llm = GeminiMultimodalLiveLLMService(
api_key=os.getenv("GOOGLE_API_KEY"),
system_instruction=system_instruction,
tools=tools,
)
messages = [
{
"role": "system",
"content": "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.",
},
]
llm.register_function("get_current_weather", fetch_weather_from_api)
context = OpenAILLMContext(
[{"role": "user", "content": "Say hello."}],
)
context = OpenAILLMContext(messages, tools)
context_aggregator = llm.create_context_aggregator(context)
pipeline = Pipeline(
[
transport.input(),
stt,
context_aggregator.user(),
llm,
tts,
transport.output(),
context_aggregator.assistant(),
]

View File

@@ -60,7 +60,7 @@ from pipecat.services.openai.llm import (
from pipecat.transcriptions.language import Language
from pipecat.utils.string import match_endofsentence
from pipecat.utils.time import time_now_iso8601
from pipecat.utils.tracing.service_decorators import traced_multimodal_llm, traced_stt, traced_tts
from pipecat.utils.tracing.service_decorators import traced_gemini_live, traced_stt, traced_tts
from . import events
@@ -473,6 +473,7 @@ class GeminiMultimodalLiveLLMService(LLMService):
async def _handle_user_stopped_speaking(self, frame):
self._user_is_speaking = False
self._user_audio_buffer = bytearray()
await self.start_ttfb_metrics()
if self._needs_turn_complete_message:
self._needs_turn_complete_message = False
evt = events.ClientContentMessage.model_validate(
@@ -754,6 +755,8 @@ class GeminiMultimodalLiveLLMService(LLMService):
logger.debug(f"Creating initial response: {messages}")
await self.start_ttfb_metrics()
evt = events.ClientContentMessage.model_validate(
{
"clientContent": {
@@ -795,6 +798,8 @@ class GeminiMultimodalLiveLLMService(LLMService):
return
logger.debug(f"Creating response: {messages}")
await self.start_ttfb_metrics()
evt = events.ClientContentMessage.model_validate(
{
"clientContent": {
@@ -805,7 +810,7 @@ class GeminiMultimodalLiveLLMService(LLMService):
)
await self.send_client_event(evt)
@traced_multimodal_llm(operation="tool_result")
@traced_gemini_live(operation="llm_tool_result")
async def _tool_result(self, tool_result_message):
# For now we're shoving the name into the tool_call_id field, so this
# will work until we revisit that.
@@ -830,7 +835,7 @@ class GeminiMultimodalLiveLLMService(LLMService):
await self._websocket.send(response_message)
# await self._websocket.send(json.dumps({"clientContent": {"turnComplete": True}}))
@traced_multimodal_llm(operation="setup")
@traced_gemini_live(operation="llm_setup")
async def _handle_evt_setup_complete(self, evt):
# If this is our first context frame, run the LLM
self._api_session_ready = True
@@ -844,6 +849,8 @@ class GeminiMultimodalLiveLLMService(LLMService):
if not part:
return
await self.stop_ttfb_metrics()
# part.text is added when `modalities` is set to TEXT; otherwise, it's None
text = part.text
if text:
@@ -877,7 +884,7 @@ class GeminiMultimodalLiveLLMService(LLMService):
)
await self.push_frame(frame)
@traced_multimodal_llm(operation="tool_call")
@traced_gemini_live(operation="llm_tool_call")
async def _handle_evt_tool_call(self, evt):
function_calls = evt.toolCall.functionCalls
if not function_calls:
@@ -892,24 +899,28 @@ class GeminiMultimodalLiveLLMService(LLMService):
arguments=call.args,
)
@traced_tts
async def _handle_bot_transcription(self, text: str):
"""Handle a transcription result with tracing."""
pass
@traced_gemini_live(operation="llm_response")
async def _handle_evt_turn_complete(self, evt):
self._bot_is_speaking = False
text = self._bot_text_buffer
# Determine output and modality for tracing
if text:
# TEXT modality
await self._handle_bot_transcription(text)
output_text = text
output_modality = "TEXT"
else:
# AUDIO modality
await self._handle_bot_transcription(self._llm_output_buffer)
output_text = self._llm_output_buffer
output_modality = "AUDIO"
# Trace the complete LLM response (this will be handled by the decorator)
# The decorator will extract the output text and usage metadata from the event
self._bot_text_buffer = ""
self._llm_output_buffer = ""
# Only push the TTSStoppedFrame the bot is outputting audio
# Only push the TTSStoppedFrame if the bot is outputting audio
# when text is found, modalities is set to TEXT and no audio
# is produced.
if not text:
@@ -990,7 +1001,6 @@ class GeminiMultimodalLiveLLMService(LLMService):
await self.push_frame(LLMTextFrame(text=text))
await self.push_frame(TTSTextFrame(text=text))
@traced_multimodal_llm(operation="usage_metadata")
async def _handle_evt_usage_metadata(self, evt):
if not evt.usageMetadata:
return

View File

@@ -258,7 +258,7 @@ def add_llm_span_attributes(
span.set_attribute(key, value)
def add_multimodal_llm_span_attributes(
def add_gemini_live_span_attributes(
span: "Span",
service_name: str,
model: str,

View File

@@ -24,8 +24,8 @@ if TYPE_CHECKING:
from opentelemetry import trace
from pipecat.utils.tracing.service_attributes import (
add_gemini_live_span_attributes,
add_llm_span_attributes,
add_multimodal_llm_span_attributes,
add_openai_realtime_span_attributes,
add_stt_span_attributes,
add_tts_span_attributes,
@@ -481,17 +481,14 @@ def traced_llm(func: Optional[Callable] = None, *, name: Optional[str] = None) -
return decorator
def traced_multimodal_llm(operation: str) -> Callable:
def traced_gemini_live(operation: str) -> Callable:
"""Traces Gemini Live service methods with operation-specific attributes.
This decorator automatically captures relevant information based on the operation type:
- setup_complete: Configuration, tools definitions, and system instructions
- model_turn: Text and audio output
- tool_call: Function call information
- tool_result: Function execution results
- input_transcription: User transcription
- output_transcription: Assistant transcription
- usage_metadata: Token usage metrics
- llm_setup: Configuration, tools definitions, and system instructions
- llm_tool_call: Function call information
- llm_tool_result: Function execution results
- llm_response: Complete LLM response with usage and output
Args:
operation: The operation name (matches the event type being handled)
@@ -542,7 +539,7 @@ def traced_multimodal_llm(operation: str) -> Callable:
# Operation-specific attribute collection
operation_attrs = {}
if operation == "setup":
if operation == "llm_setup":
# Capture detailed tool information
tools = getattr(self, "_tools", None)
if tools:
@@ -627,7 +624,7 @@ def traced_multimodal_llm(operation: str) -> Callable:
f"Error extracting context system instructions: {e}"
)
elif operation == "tool_call" and args:
elif operation == "llm_tool_call" and args:
# Extract tool call information
evt = args[0] if args else None
if evt and hasattr(evt, "toolCall") and evt.toolCall.functionCalls:
@@ -654,7 +651,7 @@ def traced_multimodal_llm(operation: str) -> Callable:
except Exception:
operation_attrs["tool.arguments"] = str(call.args)[:1000]
elif operation == "tool_result" and args:
elif operation == "llm_tool_result" and args:
# Extract tool result information
tool_result_message = args[0] if args else None
if tool_result_message and isinstance(tool_result_message, dict):
@@ -698,11 +695,13 @@ def traced_multimodal_llm(operation: str) -> Callable:
)
operation_attrs["tool.result_status"] = "processing_error"
elif operation == "usage_metadata" and args:
# Token usage will be handled by the original start_llm_usage_metrics method
elif operation == "llm_response" and args:
# Extract usage and response metadata from turn complete event
evt = args[0] if args else None
if evt and hasattr(evt, "usageMetadata") and evt.usageMetadata:
usage = evt.usageMetadata
# Token usage - basic attributes for span visibility
if hasattr(usage, "promptTokenCount"):
operation_attrs["tokens.prompt"] = usage.promptTokenCount or 0
if hasattr(usage, "responseTokenCount"):
@@ -712,8 +711,29 @@ def traced_multimodal_llm(operation: str) -> Callable:
if hasattr(usage, "totalTokenCount"):
operation_attrs["tokens.total"] = usage.totalTokenCount or 0
# Get output text and modality from service state
text = getattr(self, "_bot_text_buffer", "")
audio_text = getattr(self, "_llm_output_buffer", "")
if text:
# TEXT modality
operation_attrs["output"] = text
operation_attrs["output_modality"] = "TEXT"
elif audio_text:
# AUDIO modality
operation_attrs["output"] = audio_text
operation_attrs["output_modality"] = "AUDIO"
# Add turn completion status
if (
evt
and hasattr(evt, "serverContent")
and evt.serverContent.turnComplete
):
operation_attrs["turn_complete"] = True
# Add all attributes to the span
add_multimodal_llm_span_attributes(
add_gemini_live_span_attributes(
span=current_span,
service_name=service_class_name,
model=model_name,
@@ -725,10 +745,8 @@ def traced_multimodal_llm(operation: str) -> Callable:
**operation_attrs,
)
# For usage_metadata operation, also handle token usage metrics
if operation == "usage_metadata" and hasattr(
self, "start_llm_usage_metrics"
):
# For llm_response operation, also handle token usage metrics
if operation == "llm_response" and hasattr(self, "start_llm_usage_metrics"):
evt = args[0] if args else None
if evt and hasattr(evt, "usageMetadata") and evt.usageMetadata:
usage = evt.usageMetadata
@@ -742,6 +760,11 @@ def traced_multimodal_llm(operation: str) -> Callable:
)
_add_token_usage_to_span(current_span, tokens)
# Capture TTFB metric if available
ttfb_ms = getattr(getattr(self, "_metrics", None), "ttfb_ms", None)
if ttfb_ms is not None:
current_span.set_attribute("metrics.ttfb_ms", ttfb_ms)
# Run the original function
result = await func(self, *args, **kwargs)
@@ -766,11 +789,9 @@ def traced_openai_realtime(operation: str) -> Callable:
"""Traces OpenAI Realtime service methods with operation-specific attributes.
This decorator automatically captures relevant information based on the operation type:
- setup: Session configuration and tools
- transcription_completed: User transcription
- response_create: Context and input messages
- response_done: Usage metadata and output
- function_call: Function call information
- llm_setup: Session configuration and tools
- llm_request: Context and input messages
- llm_response: Usage metadata, output, and function calls
Args:
operation: The operation name (matches the event type being handled)
@@ -890,8 +911,10 @@ def traced_openai_realtime(operation: str) -> Callable:
if hasattr(response, "output") and response.output:
operation_attrs["response.output_items"] = len(response.output)
# Extract assistant transcript for output field
# Extract assistant transcript and function calls
assistant_transcript = ""
function_calls = []
for item in response.output:
if (
hasattr(item, "content")
@@ -906,9 +929,31 @@ def traced_openai_realtime(operation: str) -> Callable:
):
assistant_transcript += content.transcript + " "
elif hasattr(item, "type") and item.type == "function_call":
function_call_info = {
"name": getattr(item, "name", "unknown"),
"call_id": getattr(item, "call_id", "unknown"),
}
if hasattr(item, "arguments"):
args_str = item.arguments
if len(args_str) > 500:
args_str = args_str[:500] + "..."
function_call_info["arguments"] = args_str
function_calls.append(function_call_info)
if assistant_transcript.strip():
operation_attrs["output"] = assistant_transcript.strip()
if function_calls:
operation_attrs["function_calls"] = function_calls
operation_attrs["function_calls.count"] = len(
function_calls
)
all_names = [call["name"] for call in function_calls]
operation_attrs["function_calls.all_names"] = ",".join(
all_names
)
# Add all attributes to the span
add_openai_realtime_span_attributes(
span=current_span,