Better handle Gemini non-function thought signatures

This commit is contained in:
Paul Kompfner
2025-12-08 15:56:40 -05:00
parent 1249ee3de3
commit 229ff794d6
2 changed files with 83 additions and 30 deletions

View File

@@ -209,7 +209,7 @@ class GeminiLLMAdapter(BaseLLMAdapter[GeminiLLMInvocationParams]):
system_instruction = None system_instruction = None
messages = [] messages = []
tool_call_id_to_name_mapping = {} tool_call_id_to_name_mapping = {}
non_fn_signed_parts = [] non_fn_thought_signatures = []
# Process each message, converting to Google format as needed # Process each message, converting to Google format as needed
for message in universal_context_messages: for message in universal_context_messages:
@@ -230,16 +230,17 @@ class GeminiLLMAdapter(BaseLLMAdapter[GeminiLLMInvocationParams]):
) )
continue continue
# Special handling for non-function-call-related thought # Special handling for non-function-call-related thought-
# signature messages (Gemini 3 Pro mainly, but possibly others, # signature-containing messages
# too, especially when functions are involved in the
# conversation)
if ( if (
isinstance(message.message, dict) isinstance(message.message, dict)
and message.message.get("type") == "non_fn_thought_signature" and message.message.get("type") == "non_fn_thought_signature"
and (signed_part := message.message.get("signed_part")) and (thought_signature := message.message.get("signature"))
and (bookmark := message.message.get("bookmark"))
): ):
non_fn_signed_parts.append(signed_part) non_fn_thought_signatures.append(
{"signature": thought_signature, "bookmark": bookmark}
)
continue continue
# Fall back to assuming that the message is already in Google # Fall back to assuming that the message is already in Google
@@ -269,7 +270,7 @@ class GeminiLLMAdapter(BaseLLMAdapter[GeminiLLMInvocationParams]):
# Apply non-function-call-related thought signatures to the appropriate # Apply non-function-call-related thought signatures to the appropriate
# messages # messages
self._apply_non_function_thought_signatures_to_messages(non_fn_signed_parts, messages) self._apply_non_function_thought_signatures_to_messages(non_fn_thought_signatures, messages)
# Check if we only have function-related messages (no regular text) # Check if we only have function-related messages (no regular text)
has_regular_messages = any( has_regular_messages = any(
@@ -476,21 +477,37 @@ class GeminiLLMAdapter(BaseLLMAdapter[GeminiLLMInvocationParams]):
break break
def _apply_non_function_thought_signatures_to_messages( def _apply_non_function_thought_signatures_to_messages(
self, signed_parts: List[Part], messages: List[Content] self, thought_signatures: List[dict], messages: List[Content]
) -> None: ) -> None:
"""Apply (optional, but recommended) non-function-call-related thought signatures to the last part of corresponding non-function-call assistant messages. """Apply (optional, but recommended) non-function-call-related thought signatures to the last part of corresponding non-function-call assistant messages.
Gemini 3 Pro (and, somewhat surprisingly, other models, too, when Gemini 3 Pro (and, somewhat surprisingly, other models, too, when
functions are involved in the conversation) outputs a thought signature functions are involved in the conversation) outputs thought signatures
at the end of assistant responses. at the end of assistant responses.
Args: Args:
signed_parts: A list of signed received Parts containing thought signatures to apply. thought_signatures: A list of dicts containing:
- "signature": a thought signature
- "bookmark": a bookmark to identify the message part to apply the signature to.
The bookmark may contain either:
- "text"
- "inline_data"
messages: List of messages to search through. messages: List of messages to search through.
""" """
if not signed_parts: if not thought_signatures:
return return
# For debugging, print out thought signatures and their bookmarks
logger.trace(f"Thought signatures to apply: {len(thought_signatures)}")
for ts in thought_signatures:
bookmark = ts.get("bookmark")
if bookmark.get("text"):
text = bookmark["text"]
log_display_text = f"{text[:50]}..." if len(text) > 50 else text
logger.trace(f" - At text: {log_display_text}")
elif bookmark.get("inline_data"):
logger.trace(f" - At inline data")
# Find all assistant (model) messages that aren't function calls # Find all assistant (model) messages that aren't function calls
non_fn_assistant_messages = [] non_fn_assistant_messages = []
for message in messages: for message in messages:
@@ -507,9 +524,10 @@ class GeminiLLMAdapter(BaseLLMAdapter[GeminiLLMInvocationParams]):
# Apply thought signatures to the corresponding assistant messages # Apply thought signatures to the corresponding assistant messages
# Match them using content heuristics, maintaining order (messages without signatures are skipped) # Match them using content heuristics, maintaining order (messages without signatures are skipped)
message_start_index = 0 # Track where to start searching for the next match message_start_index = 0 # Track where to start searching for the next match
for signed_part in signed_parts: for thought_signature_dict in thought_signatures:
thought_signature = getattr(signed_part, "thought_signature", None) signature = thought_signature_dict.get("signature")
if not thought_signature: bookmark = thought_signature_dict.get("bookmark")
if not signature:
continue continue
# Search through remaining non-function assistant messages for a match # Search through remaining non-function assistant messages for a match
@@ -521,31 +539,41 @@ class GeminiLLMAdapter(BaseLLMAdapter[GeminiLLMInvocationParams]):
last_part = message.parts[-1] last_part = message.parts[-1]
matched = False matched = False
# Check if signed part has text and last message part text has the same text or # If it's a text bookmark, check that the last message part text has the same text or
# - is a prefix of that text (in case spoken text was truncated due to interruption) # - is a prefix of that text (in case spoken text was truncated due to interruption)
# - is prefixed by that text (in case signed part was not the end of the assistant response... # - is prefixed by that text (in case bookmark represents just first chunk of multi-chunk text)
# which is NOT supposed to happen, according to Google's docs, but seems to, for long responses...) if bookmark_text := bookmark.get("text"):
if hasattr(signed_part, "text") and signed_part.text:
if hasattr(last_part, "text") and last_part.text: if hasattr(last_part, "text") and last_part.text:
# Normalize whitespace for comparison # Normalize whitespace for comparison
signed_text = " ".join(signed_part.text.split()) signed_text = " ".join(bookmark_text.split())
last_text = " ".join(last_part.text.split()) last_text = " ".join(last_part.text.split())
if ( if (
last_text == signed_text last_text == signed_text
or signed_text.startswith(last_text) or signed_text.startswith(last_text)
or last_text.startswith(signed_text) or last_text.startswith(signed_text)
): ):
last_part.thought_signature = thought_signature log_display_text = (
f"{last_part.text[:50]}..."
if len(last_part.text) > 50
else last_part.text
)
logger.trace(
f"Applying thought signature to part with matching text: {log_display_text}"
)
last_part.thought_signature = signature
matched = True matched = True
# Check if signed part has inline_data and last message part has matching inline_data # Check if signed part has inline_data and last message part has matching inline_data
elif hasattr(signed_part, "inline_data") and signed_part.inline_data: elif inline_data := bookmark.get("inline_data"):
if ( if (
hasattr(last_part, "inline_data") hasattr(last_part, "inline_data")
and last_part.inline_data and last_part.inline_data
and last_part.inline_data.data == signed_part.inline_data.data and last_part.inline_data.data == inline_data.data
): ):
last_part.thought_signature = thought_signature logger.trace(
f"Applying thought signature to part with matching inline_data"
)
last_part.thought_signature = signature
matched = True matched = True
# If we found a match, update start index and stop searching for this signed part # If we found a match, update start index and stop searching for this signed part

View File

@@ -942,6 +942,7 @@ class GoogleLLMService(LLMService):
) )
function_calls = [] function_calls = []
previous_part = None
async for chunk in response: async for chunk in response:
# Stop TTFB metrics after the first chunk # Stop TTFB metrics after the first chunk
await self.stop_ttfb_metrics() await self.stop_ttfb_metrics()
@@ -1005,26 +1006,50 @@ class GoogleLLMService(LLMService):
) )
await self.push_frame(frame) await self.push_frame(frame)
# With Gemini 3 Pro (and, somewhat surprisingly, # With Gemini 3 Pro (and, contrary to Google's
# other models models, too, especially when # docs, other models models, too, especially when
# functions are involved in the conversation), # functions are involved in the conversation),
# thought signatures can be included in any kind of # thought signatures can be associated with any
# part, not just function calls. It will come in # kind of Part, not just function calls.
# the last part of a response. #
# They should always be included in the last
# response Part. (*)
#
# (*) Since we're using the streaming API, though,
# where text Parts may be split across multiple
# chunks (each represented by a Part, confusingly),
# signatures may actually appear with the first
# chunk (Gemini 2.5) or in a trailing empty-text
# chunk (Gemini 3 Pro).
if part.thought_signature and not part.function_call: if part.thought_signature and not part.function_call:
# Save a "bookmark" for the signature, so we
# can later stick it in the right place in
# context when sending it back to the LLM to
# continue the conversation.
bookmark = {}
if part.inline_data and part.inline_data.data:
bookmark["inline_data"] = {"inline_data": part.inline_data}
elif part.text is not None:
# Account for Gemini 3 Pro trailing
# empty-text chunk by using search_result,
# which accumulates all text so far.
bookmark["text"] = search_result
await self.push_frame( await self.push_frame(
LLMMessagesAppendFrame( LLMMessagesAppendFrame(
[ [
self.get_llm_adapter().create_llm_specific_message( self.get_llm_adapter().create_llm_specific_message(
{ {
"type": "non_fn_thought_signature", "type": "non_fn_thought_signature",
"signed_part": part, "signature": part.thought_signature,
"bookmark": bookmark,
} }
) )
] ]
) )
) )
previous_part = part
if ( if (
candidate.grounding_metadata candidate.grounding_metadata
and candidate.grounding_metadata.grounding_chunks and candidate.grounding_metadata.grounding_chunks