remove grounding metadata commits

This commit is contained in:
vipyne
2025-07-01 17:17:13 -05:00
parent f1c9f5040b
commit c7cbfe7a4f
3 changed files with 11 additions and 190 deletions

View File

@@ -12,7 +12,6 @@ import json
from enum import Enum from enum import Enum
from typing import List, Literal, Optional from typing import List, Literal, Optional
from loguru import logger
from PIL import Image from PIL import Image
from pydantic import BaseModel, Field from pydantic import BaseModel, Field
@@ -249,55 +248,6 @@ class Config(BaseModel):
setup: Setup setup: Setup
#
# Grounding metadata models
#
class SearchEntryPoint(BaseModel):
"""Represents the search entry point with rendered content for search suggestions."""
renderedContent: Optional[str] = None
class WebSource(BaseModel):
"""Represents a web source from grounding chunks."""
uri: Optional[str] = None
title: Optional[str] = None
class GroundingChunk(BaseModel):
"""Represents a grounding chunk containing web source information."""
web: Optional[WebSource] = None
class GroundingSegment(BaseModel):
"""Represents a segment of text that is grounded."""
startIndex: Optional[int] = None
endIndex: Optional[int] = None
text: Optional[str] = None
class GroundingSupport(BaseModel):
"""Represents support information for grounded text segments."""
segment: Optional[GroundingSegment] = None
groundingChunkIndices: Optional[List[int]] = None
confidenceScores: Optional[List[float]] = None
class GroundingMetadata(BaseModel):
"""Represents grounding metadata from Google Search."""
searchEntryPoint: Optional[SearchEntryPoint] = None
groundingChunks: Optional[List[GroundingChunk]] = None
groundingSupports: Optional[List[GroundingSupport]] = None
webSearchQueries: Optional[List[str]] = None
# #
# Server events # Server events
# #
@@ -389,7 +339,6 @@ class ServerContent(BaseModel):
turnComplete: Optional[bool] = None turnComplete: Optional[bool] = None
inputTranscription: Optional[BidiGenerateContentTranscription] = None inputTranscription: Optional[BidiGenerateContentTranscription] = None
outputTranscription: Optional[BidiGenerateContentTranscription] = None outputTranscription: Optional[BidiGenerateContentTranscription] = None
groundingMetadata: Optional[GroundingMetadata] = None
class FunctionCall(BaseModel): class FunctionCall(BaseModel):
@@ -482,7 +431,7 @@ class ServerEvent(BaseModel):
usageMetadata: Optional[UsageMetadata] = None usageMetadata: Optional[UsageMetadata] = None
def parse_server_event(message_str): def parse_server_event(str):
"""Parse a server event from JSON string. """Parse a server event from JSON string.
Args: Args:
@@ -491,26 +440,11 @@ def parse_server_event(message_str):
Returns: Returns:
ServerEvent instance if parsing succeeds, None otherwise. ServerEvent instance if parsing succeeds, None otherwise.
""" """
from loguru import logger # Import logger locally to avoid scoping issues
try: try:
evt_dict = json.loads(message_str) evt = json.loads(str)
return ServerEvent.model_validate(evt)
# Only log grounding metadata detection if truly needed for debugging
# In production, this could be removed entirely or moved to TRACE level
if "serverContent" in evt_dict:
server_content = evt_dict["serverContent"]
if "groundingMetadata" in server_content:
# Consider removing this log entirely for production
pass
evt = ServerEvent.model_validate(evt_dict)
return evt
except Exception as e: except Exception as e:
logger.error(f"Error parsing server event: {e}") print(f"Error parsing server event: {e}")
# Truncate raw message to avoid logging potentially sensitive or overly long data
truncated_message = message_str[:200] + "..." if len(message_str) > 200 else message_str
logger.error(f"Raw message (truncated): {truncated_message}")
return None return None

View File

@@ -31,8 +31,8 @@ class GeminiFileAPI:
api_key: Google AI API key api_key: Google AI API key
base_url: Base URL for the Gemini File API (default is the v1beta endpoint) base_url: Base URL for the Gemini File API (default is the v1beta endpoint)
""" """
self.api_key = api_key self._api_key = api_key
self.base_url = base_url self._base_url = base_url
# Upload URL uses the /upload/ path # Upload URL uses the /upload/ path
self.upload_base_url = "https://generativelanguage.googleapis.com/upload/v1beta/files" self.upload_base_url = "https://generativelanguage.googleapis.com/upload/v1beta/files"
@@ -76,7 +76,7 @@ class GeminiFileAPI:
logger.debug(f"Step 1: Getting upload URL from {self.upload_base_url}") logger.debug(f"Step 1: Getting upload URL from {self.upload_base_url}")
async with session.post( async with session.post(
f"{self.upload_base_url}?key={self.api_key}", headers=headers, json=metadata f"{self.upload_base_url}?key={self._api_key}", headers=headers, json=metadata
) as response: ) as response:
if response.status != 200: if response.status != 200:
error_text = await response.text() error_text = await response.text()
@@ -123,7 +123,7 @@ class GeminiFileAPI:
name = name.split("/")[-1] name = name.split("/")[-1]
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
async with session.get(f"{self.base_url}/{name}?key={self.api_key}") as response: async with session.get(f"{self._base_url}/{name}?key={self._api_key}") as response:
if response.status != 200: if response.status != 200:
error_text = await response.text() error_text = await response.text()
logger.error(f"Error getting file metadata: {error_text}") logger.error(f"Error getting file metadata: {error_text}")
@@ -144,13 +144,13 @@ class GeminiFileAPI:
Returns: Returns:
List of files and next page token if available List of files and next page token if available
""" """
params = {"key": self.api_key, "pageSize": page_size} params = {"key": self._api_key, "pageSize": page_size}
if page_token: if page_token:
params["pageToken"] = page_token params["pageToken"] = page_token
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
async with session.get(self.base_url, params=params) as response: async with session.get(self._base_url, params=params) as response:
if response.status != 200: if response.status != 200:
error_text = await response.text() error_text = await response.text()
logger.error(f"Error listing files: {error_text}") logger.error(f"Error listing files: {error_text}")
@@ -173,7 +173,7 @@ class GeminiFileAPI:
name = name.split("/")[-1] name = name.split("/")[-1]
async with aiohttp.ClientSession() as session: async with aiohttp.ClientSession() as session:
async with session.delete(f"{self.base_url}/{name}?key={self.api_key}") as response: async with session.delete(f"{self._base_url}/{name}?key={self._api_key}") as response:
if response.status != 200: if response.status != 200:
error_text = await response.text() error_text = await response.text()
logger.error(f"Error deleting file: {error_text}") logger.error(f"Error deleting file: {error_text}")

View File

@@ -564,10 +564,6 @@ class GeminiMultimodalLiveLLMService(LLMService):
# Initialize the File API client # Initialize the File API client
self.file_api = GeminiFileAPI(api_key=api_key, base_url=file_api_base_url) self.file_api = GeminiFileAPI(api_key=api_key, base_url=file_api_base_url)
# Grounding metadata tracking
self._search_result_buffer = ""
self._accumulated_grounding_metadata = None
def can_generate_metrics(self) -> bool: def can_generate_metrics(self) -> bool:
"""Check if the service can generate usage metrics. """Check if the service can generate usage metrics.
@@ -917,23 +913,12 @@ class GeminiMultimodalLiveLLMService(LLMService):
await self._handle_evt_input_transcription(evt) await self._handle_evt_input_transcription(evt)
elif evt.serverContent and evt.serverContent.outputTranscription: elif evt.serverContent and evt.serverContent.outputTranscription:
await self._handle_evt_output_transcription(evt) await self._handle_evt_output_transcription(evt)
elif evt.serverContent and evt.serverContent.groundingMetadata:
await self._handle_evt_grounding_metadata(evt)
elif evt.toolCall: elif evt.toolCall:
await self._handle_evt_tool_call(evt) await self._handle_evt_tool_call(evt)
elif False: # !!! todo: error events? elif False: # !!! todo: error events?
await self._handle_evt_error(evt) await self._handle_evt_error(evt)
# errors are fatal, so exit the receive loop # errors are fatal, so exit the receive loop
return return
else:
# Log unhandled events that might contain grounding metadata
logger.warning(f"Received unhandled server event type: {evt}")
pass
async def _transcribe_audio_handler(self):
while True:
audio = await self._transcribe_audio_queue.get()
await self._handle_transcribe_user_audio(audio, self._context)
# #
# #
@@ -1092,14 +1077,8 @@ class GeminiMultimodalLiveLLMService(LLMService):
await self.push_frame(LLMFullResponseStartFrame()) await self.push_frame(LLMFullResponseStartFrame())
self._bot_text_buffer += text self._bot_text_buffer += text
self._search_result_buffer += text # Also accumulate for grounding
await self.push_frame(LLMTextFrame(text=text)) await self.push_frame(LLMTextFrame(text=text))
# Check for grounding metadata in server content
if evt.serverContent and evt.serverContent.groundingMetadata:
self._accumulated_grounding_metadata = evt.serverContent.groundingMetadata
logger.debug("Grounding metadata detected in model turn.")
inline_data = part.inlineData inline_data = part.inlineData
if not inline_data: if not inline_data:
return return
@@ -1166,20 +1145,6 @@ class GeminiMultimodalLiveLLMService(LLMService):
self._llm_output_buffer = "" self._llm_output_buffer = ""
# Only push the TTSStoppedFrame if the bot is outputting audio # Only push the TTSStoppedFrame if the bot is outputting audio
# Process grounding metadata if we have accumulated any
if self._accumulated_grounding_metadata:
logger.debug("Processing grounding metadata...")
await self._process_grounding_metadata(
self._accumulated_grounding_metadata, self._search_result_buffer
)
else:
logger.debug("No grounding metadata to process")
# Reset grounding tracking for next response
self._search_result_buffer = ""
self._accumulated_grounding_metadata = None
# Only push the TTSStoppedFrame the bot is outputting audio
# when text is found, modalities is set to TEXT and no audio # when text is found, modalities is set to TEXT and no audio
# is produced. # is produced.
if not text: if not text:
@@ -1257,13 +1222,6 @@ class GeminiMultimodalLiveLLMService(LLMService):
# Collect text for tracing # Collect text for tracing
self._llm_output_buffer += text self._llm_output_buffer += text
# Accumulate text for grounding as well
self._search_result_buffer += text
# Check for grounding metadata in server content
if evt.serverContent and evt.serverContent.groundingMetadata:
self._accumulated_grounding_metadata = evt.serverContent.groundingMetadata
await self.push_frame(LLMTextFrame(text=text)) await self.push_frame(LLMTextFrame(text=text))
await self.push_frame(TTSTextFrame(text=text)) await self.push_frame(TTSTextFrame(text=text))
@@ -1286,77 +1244,6 @@ class GeminiMultimodalLiveLLMService(LLMService):
await self.start_llm_usage_metrics(tokens) await self.start_llm_usage_metrics(tokens)
async def _handle_evt_grounding_metadata(self, evt):
"""Handle dedicated grounding metadata events."""
logger.debug("Received dedicated grounding metadata event.")
if evt.serverContent and evt.serverContent.groundingMetadata:
grounding_metadata = evt.serverContent.groundingMetadata
logger.debug(
f"Grounding data: {len(grounding_metadata.groundingChunks or [])} chunks, {len(grounding_metadata.groundingSupports or [])} supports"
)
# Process the grounding metadata immediately
await self._process_grounding_metadata(grounding_metadata, self._search_result_buffer)
async def _process_grounding_metadata(
self, grounding_metadata: events.GroundingMetadata, search_result: str = ""
):
"""Process grounding metadata and emit LLMSearchResponseFrame."""
logger.debug(
f"Processing grounding metadata. Search result text length: {len(search_result)}"
)
if not grounding_metadata:
logger.warning("No grounding metadata provided to _process_grounding_metadata")
return
# logger.debug(f"Processing grounding metadata: {grounding_metadata}") # Too verbose for PR
# Extract rendered content for search suggestions
rendered_content = None
if (
grounding_metadata.searchEntryPoint
and grounding_metadata.searchEntryPoint.renderedContent
):
rendered_content = grounding_metadata.searchEntryPoint.renderedContent
# Convert grounding chunks and supports to LLMSearchOrigin format
origins = []
if grounding_metadata.groundingChunks and grounding_metadata.groundingSupports:
# Create a mapping of chunk indices to origins
chunk_to_origin = {}
for index, chunk in enumerate(grounding_metadata.groundingChunks):
if chunk.web:
origin = LLMSearchOrigin(
site_uri=chunk.web.uri, site_title=chunk.web.title, results=[]
)
chunk_to_origin[index] = origin
origins.append(origin)
# Add grounding support results to the appropriate origins
for support in grounding_metadata.groundingSupports:
if support.segment and support.groundingChunkIndices:
text = support.segment.text or ""
confidence_scores = support.confidenceScores or []
# Add this result to all origins referenced by this support
for chunk_index in support.groundingChunkIndices:
if chunk_index in chunk_to_origin:
result = LLMSearchResult(text=text, confidence=confidence_scores)
chunk_to_origin[chunk_index].results.append(result)
# Create and push the search response frame
search_frame = LLMSearchResponseFrame(
search_result=search_result, origins=origins, rendered_content=rendered_content
)
logger.debug(
f"Emitting LLMSearchResponseFrame with {len(origins)} origins, rendered_content available: {rendered_content is not None}"
)
await self.push_frame(search_frame)
def create_context_aggregator( def create_context_aggregator(
self, self,
context: OpenAILLMContext, context: OpenAILLMContext,