Clean up logic related to applying Gemini thought signatures to context messages

This commit is contained in:
Paul Kompfner
2025-12-12 10:50:16 -05:00
parent 22288648e6
commit 64471d65f8
8 changed files with 170 additions and 218 deletions

View File

@@ -209,7 +209,7 @@ class GeminiLLMAdapter(BaseLLMAdapter[GeminiLLMInvocationParams]):
system_instruction = None
messages = []
tool_call_id_to_name_mapping = {}
non_fn_thought_signatures = []
thought_signature_dicts = []
# Process each message, converting to Google format as needed
for message in universal_context_messages:
@@ -218,29 +218,11 @@ class GeminiLLMAdapter(BaseLLMAdapter[GeminiLLMInvocationParams]):
# special way, or a message already in Google format that we can
# use directly
if isinstance(message, LLMSpecificMessage):
# Special handling for function-call-related thought signature
# messages
if (
isinstance(message.message, dict)
and message.message.get("type") == "fn_thought_signature"
and (thought_signature := message.message.get("signature"))
and message.message.get("type") == "thought_signature"
):
self._apply_function_thought_signature_to_messages(
thought_signature, message.message.get("tool_call_id"), messages
)
continue
# Special handling for non-function-call-related thought-
# signature-containing messages
if (
isinstance(message.message, dict)
and message.message.get("type") == "non_fn_thought_signature"
and (thought_signature := message.message.get("signature"))
and (bookmark := message.message.get("bookmark"))
):
non_fn_thought_signatures.append(
{"signature": thought_signature, "bookmark": bookmark}
)
thought_signature_dicts.append(message.message)
continue
# Fall back to assuming that the message is already in Google
@@ -268,9 +250,8 @@ class GeminiLLMAdapter(BaseLLMAdapter[GeminiLLMInvocationParams]):
if result.tool_call_id_to_name_mapping:
tool_call_id_to_name_mapping.update(result.tool_call_id_to_name_mapping)
# Apply non-function-call-related thought signatures to the appropriate
# messages
self._apply_non_function_thought_signatures_to_messages(non_fn_thought_signatures, messages)
# Apply thought signatures to the corresponding messages
self._apply_thought_signatures_to_messages(thought_signature_dicts, messages)
# Check if we only have function-related messages (no regular text)
has_regular_messages = any(
@@ -447,136 +428,135 @@ class GeminiLLMAdapter(BaseLLMAdapter[GeminiLLMInvocationParams]):
tool_call_id_to_name_mapping=tool_call_id_to_name_mapping,
)
def _apply_function_thought_signature_to_messages(
self, thought_signature: bytes, tool_call_id: str, messages: List[Content]
def _apply_thought_signatures_to_messages(
self, thought_signature_dicts: List[dict], messages: List[Content]
) -> None:
"""Apply a function-related thought signature to the corresponding function call message.
"""Apply thought signatures to corresponding assistant messages.
See GoogleLLMService for more details about thought signatures.
Args:
thought_signature: The thought signature bytes to apply.
tool_call_id: ID of the tool call message to find and modify.
messages: List of messages to search through.
"""
# Search backwards through messages to find the matching function call
for message in reversed(messages):
if not isinstance(message, Content) or not message.parts:
continue
# Find the specific part with the matching function call
for part in message.parts:
if (
hasattr(part, "function_call")
and part.function_call
and part.function_call.id == tool_call_id
):
part.thought_signature = thought_signature
break
else:
# Continue outer loop if inner loop didn't break
continue
# Break outer loop if inner loop broke (found match)
break
def _apply_non_function_thought_signatures_to_messages(
self, thought_signatures: List[dict], messages: List[Content]
) -> None:
"""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
functions are involved in the conversation) outputs thought signatures
at the end of assistant responses.
Args:
thought_signatures: A list of dicts containing:
thought_signature_dicts: 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.
The bookmark may contain one of:
- "function_call" (a function call ID string)
- "text" (a text string)
- "inline_data" (a Blob)
The list of thought signature dicts is in order.
messages: List of messages to apply the thought signatures to.
"""
if not thought_signatures:
if not thought_signature_dicts:
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:
logger.debug(f"Thought signatures to apply: {len(thought_signature_dicts)}")
for ts in thought_signature_dicts:
bookmark = ts.get("bookmark")
if bookmark.get("text"):
if bookmark.get("function_call"):
logger.trace(f" - To function call: {bookmark['function_call']}")
elif 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}")
logger.trace(f" - To text: {log_display_text}")
elif bookmark.get("inline_data"):
logger.trace(f" - At inline data")
logger.trace(f" - To inline data")
# Find all assistant (model) messages that aren't function calls
non_fn_assistant_messages = []
for message in messages:
if not isinstance(message, Content) or not message.parts:
continue
# Check if this is a model message without function calls
if message.role == "model":
has_function_call = any(
hasattr(part, "function_call") and part.function_call for part in message.parts
)
if not has_function_call:
non_fn_assistant_messages.append(message)
# Get all assistant messages
assistant_messages = [
message
for message in messages
if isinstance(message, Content) and message.role == "model"
]
# Apply thought signatures to the corresponding assistant messages
# 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
for thought_signature_dict in thought_signatures:
# Apply thought signatures to the corresponding assistant messages.
# Thought signatures are already in message order.
thought_signatures_applied = 0
message_start_index = 0 # Track where to start searching for the next matching message.
for thought_signature_dict in thought_signature_dicts:
signature = thought_signature_dict.get("signature")
bookmark = thought_signature_dict.get("bookmark")
if not signature:
if not signature or not bookmark:
continue
# Search through remaining non-function assistant messages for a match
for i in range(message_start_index, len(non_fn_assistant_messages)):
message = non_fn_assistant_messages[i]
# Search through remaining assistant messages for a match
for i in range(message_start_index, len(assistant_messages)):
message = assistant_messages[i]
if not message.parts:
continue
# We're assuming that the thought signature always applies to the last part
last_part = message.parts[-1]
matched = False
# 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 prefixed by that text (in case bookmark represents just first chunk of multi-chunk text)
if bookmark_text := bookmark.get("text"):
if hasattr(last_part, "text") and last_part.text:
# Normalize whitespace for comparison
signed_text = " ".join(bookmark_text.split())
last_text = " ".join(last_part.text.split())
if (
last_text == signed_text
or signed_text.startswith(last_text)
or last_text.startswith(signed_text)
):
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
# If the bookmark matches the part...
if self._thought_signature_bookmark_matches_part(bookmark, last_part):
# Apply the thought signature
last_part.thought_signature = signature
thought_signatures_applied += 1
# Check if signed part has inline_data and last message part has matching inline_data
elif inline_data := bookmark.get("inline_data"):
if (
hasattr(last_part, "inline_data")
and last_part.inline_data
and last_part.inline_data.data == inline_data.data
):
logger.trace(
f"Applying thought signature to part with matching inline_data"
)
last_part.thought_signature = signature
matched = True
# If we found a match, update start index and stop searching for this signed part
if matched:
# Update the start index and stop searching for a match
message_start_index = i + 1
break
# For debugging, print out how many thought signatures were applied
logger.debug(f"Applied {thought_signatures_applied} thought signatures.")
def _thought_signature_bookmark_matches_part(self, bookmark: dict, part: Part) -> bool:
if function_call_bookmark := bookmark.get("function_call"):
return self._thought_signature_function_call_bookmark_matches_part(
function_call_bookmark, part
)
elif text_bookmark := bookmark.get("text"):
return self._thought_signature_text_bookmark_matches_part(text_bookmark, part)
elif inline_data := bookmark.get("inline_data"):
return self._thought_signature_inline_data_bookmark_matches_part(inline_data, part)
else:
logger.warning(f"Unknown thought signature bookmark type: {bookmark}")
return False
def _thought_signature_function_call_bookmark_matches_part(
self, bookmark_function_call_id: str, part: Part
) -> bool:
if (
hasattr(part, "function_call")
and part.function_call
and part.function_call.id == bookmark_function_call_id
):
logger.trace(f"Thought signature function call match: {bookmark_function_call_id}")
return True
return False
def _thought_signature_text_bookmark_matches_part(self, bookmark_text: str, part: Part) -> bool:
if hasattr(part, "text") and part.text:
# Normalize whitespace for comparison
bookmark_text = " ".join(bookmark_text.split())
part_text = " ".join(part.text.split())
# Check that either:
# - the part text is the same as the bookmark text
# - a prefix of the bookmark text (in case the part text was truncated due to interruption)
# - the bookmark text is a prefix of the part text (in case the bookmark represents just first chunk of multi-chunk text)
if (
part_text == bookmark_text
or bookmark_text.startswith(part_text)
or part_text.startswith(bookmark_text)
):
log_display_text = f"{part.text[:50]}..." if len(part.text) > 50 else part.text
logger.trace(f"Thought signature text match: {log_display_text}")
return True
return False
def _thought_signature_inline_data_bookmark_matches_part(
self, bookmark_inline_data: Blob, part: Part
) -> bool:
if (
hasattr(part, "inline_data")
and part.inline_data
and part.inline_data.data == bookmark_inline_data.data
):
logger.trace(f"Thought signature inline data match")
return True
return False

View File

@@ -1197,16 +1197,12 @@ class FunctionCallFromLLM:
tool_call_id: A unique identifier for the function call.
arguments: The arguments to pass to the function.
context: The LLM context when the function call was made.
append_extra_context_messages: Optional extra messages to append to the
context after the function call message. Used to add Google
function-call-related thought signatures to the context.
"""
function_name: str
tool_call_id: str
arguments: Mapping[str, Any]
context: Any
append_extra_context_messages: Optional[List["LLMContextMessage"]] = None
@dataclass
@@ -1745,16 +1741,12 @@ class FunctionCallInProgressFrame(ControlFrame, UninterruptibleFrame):
tool_call_id: Unique identifier for this function call.
arguments: Arguments passed to the function.
cancel_on_interruption: Whether to cancel this call if interrupted.
append_extra_context_messages: Optional extra messages to append to the
context after the function call message. Used to add Google
function-call-related thought signatures to the context.
"""
function_name: str
tool_call_id: str
arguments: Any
cancel_on_interruption: bool = False
append_extra_context_messages: Optional[List["LLMContextMessage"]] = None
@dataclass

View File

@@ -740,10 +740,6 @@ class LLMAssistantAggregator(LLMContextAggregator):
}
)
# Append to context any specified extra context messages
if frame.append_extra_context_messages:
self._context.add_messages(frame.append_extra_context_messages)
self._function_calls_in_progress[frame.tool_call_id] = frame
async def _handle_function_call_result(self, frame: FunctionCallResultFrame):

View File

@@ -932,7 +932,7 @@ class GoogleLLMService(LLMService):
reasoning_tokens = 0
grounding_metadata = None
search_result = ""
accumulated_text = ""
try:
# Generate content using either OpenAILLMContext or universal LLMContext
@@ -943,7 +943,6 @@ class GoogleLLMService(LLMService):
)
function_calls = []
previous_part = None
async for chunk in response:
# Stop TTFB metrics after the first chunk
await self.stop_ttfb_metrics()
@@ -966,6 +965,7 @@ class GoogleLLMService(LLMService):
for candidate in chunk.candidates:
if candidate.content and candidate.content.parts:
for part in candidate.content.parts:
function_call_id = None
if part.text:
if part.thought:
# Gemini emits fully-formed thoughts rather
@@ -975,29 +975,20 @@ class GoogleLLMService(LLMService):
await self.push_frame(LLMThoughtTextFrame(part.text))
await self.push_frame(LLMThoughtEndFrame())
else:
search_result += part.text
accumulated_text += part.text
await self.push_frame(LLMTextFrame(part.text))
elif part.function_call:
function_call = part.function_call
id = function_call.id or str(uuid.uuid4())
logger.debug(f"Function call: {function_call.name}:{id}")
function_call_id = function_call.id or str(uuid.uuid4())
logger.debug(
f"Function call: {function_call.name}:{function_call_id}"
)
function_calls.append(
FunctionCallFromLLM(
context=context,
tool_call_id=id,
tool_call_id=function_call_id,
function_name=function_call.name,
arguments=function_call.args or {},
append_extra_context_messages=[
self.get_llm_adapter().create_llm_specific_message(
{
"type": "fn_thought_signature",
"signature": part.thought_signature,
"tool_call_id": id,
}
)
]
if part.thought_signature
else None,
)
)
elif part.inline_data and part.inline_data.data:
@@ -1007,14 +998,14 @@ class GoogleLLMService(LLMService):
)
await self.push_frame(frame)
# With Gemini 3 Pro (and, contrary to Google's
# docs, other models models, too, especially when
# functions are involved in the conversation),
# thought signatures can be associated with any
# kind of Part, not just function calls.
# Handle Gemini thought signatures.
#
# They should always be included in the last
# response Part. (*)
# - Gemini 2.5: they appear on function_call Parts,
# and then (surprisingly) on the last(*) Part of
# model responses following the first function_call
# in a conversation.
# - Gemini 3 Pro: they appear on the last(*) Part
# of model responses, regardless of Part type.
#
# (*) Since we're using the streaming API, though,
# where text Parts may be split across multiple
@@ -1022,34 +1013,44 @@ class GoogleLLMService(LLMService):
# 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:
# 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.
# can later be sure we've put it in the right
# place in context when sending the context
# 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}
if part.function_call:
bookmark["function_call"] = function_call_id
elif part.inline_data and part.inline_data.data:
# NOTE: missing feature: we don't store
# inline_data messages (like generated
# images) in context today, so this thought
# signature is not fully supported yet.
# (A conversation with
# "gemini-3-pro-image-preview" doesn't work
# today due to the missing context.)
bookmark["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(
LLMMessagesAppendFrame(
[
self.get_llm_adapter().create_llm_specific_message(
{
"type": "non_fn_thought_signature",
"signature": part.thought_signature,
"bookmark": bookmark,
}
)
]
# empty-text chunk by using all the text
# seen so far in this response's chunks.
bookmark["text"] = accumulated_text
else:
logger.warning("Thought signature found on unhandled Part type")
if bookmark:
await self.push_frame(
LLMMessagesAppendFrame(
[
self.get_llm_adapter().create_llm_specific_message(
{
"type": "thought_signature",
"signature": part.thought_signature,
"bookmark": bookmark,
}
)
]
)
)
)
previous_part = part
if (
candidate.grounding_metadata
@@ -1098,7 +1099,7 @@ class GoogleLLMService(LLMService):
finally:
if grounding_metadata and isinstance(grounding_metadata, dict):
llm_search_frame = LLMSearchResponseFrame(
search_result=search_result,
search_result=accumulated_text,
origins=grounding_metadata["origins"],
rendered_content=grounding_metadata["rendered_content"],
)

View File

@@ -132,9 +132,6 @@ class FunctionCallRunnerItem:
tool_call_id: A unique identifier for the function call.
arguments: The arguments for the function.
context: The LLM context.
append_extra_context_messages: Optional extra messages to append to the
context after the function call message. Used to add Google
function-call-related thought signatures to the context.
run_llm: Optional flag to control LLM execution after function call.
"""
@@ -143,7 +140,6 @@ class FunctionCallRunnerItem:
tool_call_id: str
arguments: Mapping[str, Any]
context: OpenAILLMContext | LLMContext
append_extra_context_messages: Optional[List[LLMContextMessage]] = None
run_llm: Optional[bool] = None
@@ -465,7 +461,6 @@ class LLMService(AIService):
tool_call_id=function_call.tool_call_id,
arguments=function_call.arguments,
context=function_call.context,
append_extra_context_messages=function_call.append_extra_context_messages,
)
)
@@ -590,7 +585,6 @@ class LLMService(AIService):
function_name=runner_item.function_name,
tool_call_id=runner_item.tool_call_id,
arguments=runner_item.arguments,
append_extra_context_messages=runner_item.append_extra_context_messages,
cancel_on_interruption=item.cancel_on_interruption,
)