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 typing import List, Literal, Optional
from loguru import logger
from PIL import Image
from pydantic import BaseModel, Field
@@ -249,55 +248,6 @@ class Config(BaseModel):
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
#
@@ -389,7 +339,6 @@ class ServerContent(BaseModel):
turnComplete: Optional[bool] = None
inputTranscription: Optional[BidiGenerateContentTranscription] = None
outputTranscription: Optional[BidiGenerateContentTranscription] = None
groundingMetadata: Optional[GroundingMetadata] = None
class FunctionCall(BaseModel):
@@ -482,7 +431,7 @@ class ServerEvent(BaseModel):
usageMetadata: Optional[UsageMetadata] = None
def parse_server_event(message_str):
def parse_server_event(str):
"""Parse a server event from JSON string.
Args:
@@ -491,26 +440,11 @@ def parse_server_event(message_str):
Returns:
ServerEvent instance if parsing succeeds, None otherwise.
"""
from loguru import logger # Import logger locally to avoid scoping issues
try:
evt_dict = json.loads(message_str)
# 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
evt = json.loads(str)
return ServerEvent.model_validate(evt)
except Exception as e:
logger.error(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}")
print(f"Error parsing server event: {e}")
return None

View File

@@ -31,8 +31,8 @@ class GeminiFileAPI:
api_key: Google AI API key
base_url: Base URL for the Gemini File API (default is the v1beta endpoint)
"""
self.api_key = api_key
self.base_url = base_url
self._api_key = api_key
self._base_url = base_url
# Upload URL uses the /upload/ path
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}")
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:
if response.status != 200:
error_text = await response.text()
@@ -123,7 +123,7 @@ class GeminiFileAPI:
name = name.split("/")[-1]
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:
error_text = await response.text()
logger.error(f"Error getting file metadata: {error_text}")
@@ -144,13 +144,13 @@ class GeminiFileAPI:
Returns:
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:
params["pageToken"] = page_token
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:
error_text = await response.text()
logger.error(f"Error listing files: {error_text}")
@@ -173,7 +173,7 @@ class GeminiFileAPI:
name = name.split("/")[-1]
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:
error_text = await response.text()
logger.error(f"Error deleting file: {error_text}")

View File

@@ -564,10 +564,6 @@ class GeminiMultimodalLiveLLMService(LLMService):
# Initialize the File API client
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:
"""Check if the service can generate usage metrics.
@@ -917,23 +913,12 @@ class GeminiMultimodalLiveLLMService(LLMService):
await self._handle_evt_input_transcription(evt)
elif evt.serverContent and evt.serverContent.outputTranscription:
await self._handle_evt_output_transcription(evt)
elif evt.serverContent and evt.serverContent.groundingMetadata:
await self._handle_evt_grounding_metadata(evt)
elif evt.toolCall:
await self._handle_evt_tool_call(evt)
elif False: # !!! todo: error events?
await self._handle_evt_error(evt)
# errors are fatal, so exit the receive loop
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())
self._bot_text_buffer += text
self._search_result_buffer += text # Also accumulate for grounding
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
if not inline_data:
return
@@ -1166,20 +1145,6 @@ class GeminiMultimodalLiveLLMService(LLMService):
self._llm_output_buffer = ""
# 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
# is produced.
if not text:
@@ -1257,13 +1222,6 @@ class GeminiMultimodalLiveLLMService(LLMService):
# Collect text for tracing
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(TTSTextFrame(text=text))
@@ -1286,77 +1244,6 @@ class GeminiMultimodalLiveLLMService(LLMService):
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(
self,
context: OpenAILLMContext,