Add groundingMetadata events.py
This commit is contained in:
@@ -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
|
||||||
|
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user