Add groundingMetadata events.py

This commit is contained in:
getchannel
2025-05-30 18:07:09 -04:00
committed by GitHub
parent 43c6f1f5cd
commit 8070e156d8

View File

@@ -11,6 +11,7 @@ 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
@@ -138,6 +139,49 @@ 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
# #
@@ -178,6 +222,7 @@ class ServerContent(BaseModel):
interrupted: Optional[bool] = None interrupted: Optional[bool] = None
turnComplete: Optional[bool] = None turnComplete: Optional[bool] = None
outputTranscription: Optional[BidiGenerateContentTranscription] = None outputTranscription: Optional[BidiGenerateContentTranscription] = None
groundingMetadata: Optional[GroundingMetadata] = None
class FunctionCall(BaseModel): class FunctionCall(BaseModel):
@@ -196,12 +241,27 @@ class ServerEvent(BaseModel):
toolCall: Optional[ToolCall] = None toolCall: Optional[ToolCall] = None
def parse_server_event(str): def parse_server_event(message_str):
from loguru import logger # Import logger locally to avoid scoping issues
try: try:
evt = json.loads(str) evt_dict = json.loads(message_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:
print(f"Error parsing server event: {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}")
return None return None