remove grounding metadata commits
This commit is contained in:
@@ -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
|
||||
|
||||
|
||||
|
||||
@@ -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}")
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user