Better handle Gemini non-function thought signatures
This commit is contained in:
@@ -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
|
||||||
|
|||||||
@@ -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
|
||||||
|
|||||||
Reference in New Issue
Block a user