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

@@ -10,31 +10,19 @@
- `LLMThoughtStartFrame` - `LLMThoughtStartFrame`
- `LLMThoughtTextFrame` - `LLMThoughtTextFrame`
- `LLMThoughtEndFrame` - `LLMThoughtEndFrame`
3. A mechanism for appending arbitrary context messages after a function call 3. A generic mechanism for recording LLM thoughts to context, used
message, used specifically to support Google's function-call-related
"thought signatures", which are necessary to ensure thinking continuity
between function calls in a chain (where the model thinks, makes a function
call, thinks some more, etc.). See:
- `append_extra_context_messages` field in `FunctionInProgressFrame` and
helper types
- `GoogleLLMService` leveraging the new mechanism to add a Google-specific
`"fn_thought_signature"` message
- `LLMAssistantAggregator` handling of `append_extra_context_messages`
- `GeminiLLMAdapter` handling of `"fn_thought_signature"` messages
4. A generic mechanism for recording LLM thoughts to context, used
specifically to support Anthropic, whose thought signatures are expected to specifically to support Anthropic, whose thought signatures are expected to
appear alongside the text of the thoughts within assistant context appear alongside the text of the thoughts within assistant context
messages. See: messages. See:
- `LLMThoughtEndFrame.signature` - `LLMThoughtEndFrame.signature`
- `LLMAssistantAggregator` handling of the above field - `LLMAssistantAggregator` handling of the above field
- `AnthropicLLMAdapter` handling of `"thought"` context messages - `AnthropicLLMAdapter` handling of `"thought"` context messages
5. Google-specific logic for inserting non-function-call-related thought 4. Google-specific logic for inserting thought signatures into the context,
signatures into the context, to help maintain thinking continuity in a to help maintain thinking continuity in a chain of LLM calls. See:
chain of LLM calls. See:
- `GoogleLLMService` sending `LLMMessagesAppendFrame`s to add LLM-specific - `GoogleLLMService` sending `LLMMessagesAppendFrame`s to add LLM-specific
`"non_fn_thought_signature"` messages to context `"thought_signature"` messages to context
- `GeminiLLMAdapter` handling of `"non_fn_thought_signature"` messages - `GeminiLLMAdapter` handling of `"thought_signature"` messages
6. An expansion of `TranscriptProcessor` to process LLM thoughts in addition 5. An expansion of `TranscriptProcessor` to process LLM thoughts in addition
to user and assistant utterances. See: to user and assistant utterances. See:
- `TranscriptProcessor(process_thoughts=True)` (defaults to `False`) - `TranscriptProcessor(process_thoughts=True)` (defaults to `False`)
- `ThoughtTranscriptionMessage`, which is now also emitted with the - `ThoughtTranscriptionMessage`, which is now also emitted with the

View File

@@ -123,8 +123,9 @@ async def run_bot(transport: BaseTransport, runner_args: RunnerArguments):
"content": "Say hello briefly.", "content": "Say hello briefly.",
} }
) )
# Here are some example prompts conducive to demonstrating # Replace the above with one of these example prompts to demonstrate
# thinking (picked from Google and Anthropic docs). # thinking.
# These examples come from Gemini and Anthropic docs.
# messages.append( # messages.append(
# { # {
# "role": "user", # "role": "user",

View File

@@ -149,8 +149,8 @@ async def run_bot(transport: BaseTransport, runner_args: RunnerArguments):
"content": "Say hello briefly.", "content": "Say hello briefly.",
} }
) )
# Here is an example prompt conducive to demonstrating thinking and # Replace the above with one of these example prompts to demonstrate
# function calling. # thinking and function calling.
# This example comes from Gemini docs. # This example comes from Gemini docs.
# messages.append( # messages.append(
# { # {

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_thought_signatures = [] thought_signature_dicts = []
# 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:
@@ -218,29 +218,11 @@ class GeminiLLMAdapter(BaseLLMAdapter[GeminiLLMInvocationParams]):
# special way, or a message already in Google format that we can # special way, or a message already in Google format that we can
# use directly # use directly
if isinstance(message, LLMSpecificMessage): if isinstance(message, LLMSpecificMessage):
# Special handling for function-call-related thought signature
# messages
if ( if (
isinstance(message.message, dict) isinstance(message.message, dict)
and message.message.get("type") == "fn_thought_signature" and message.message.get("type") == "thought_signature"
and (thought_signature := message.message.get("signature"))
): ):
self._apply_function_thought_signature_to_messages( thought_signature_dicts.append(message.message)
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}
)
continue continue
# Fall back to assuming that the message is already in Google # 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: if result.tool_call_id_to_name_mapping:
tool_call_id_to_name_mapping.update(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 # Apply thought signatures to the corresponding messages
# messages self._apply_thought_signatures_to_messages(thought_signature_dicts, 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(
@@ -447,136 +428,135 @@ class GeminiLLMAdapter(BaseLLMAdapter[GeminiLLMInvocationParams]):
tool_call_id_to_name_mapping=tool_call_id_to_name_mapping, tool_call_id_to_name_mapping=tool_call_id_to_name_mapping,
) )
def _apply_function_thought_signature_to_messages( def _apply_thought_signatures_to_messages(
self, thought_signature: bytes, tool_call_id: str, messages: List[Content] self, thought_signature_dicts: List[dict], messages: List[Content]
) -> None: ) -> 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: Args:
thought_signature: The thought signature bytes to apply. thought_signature_dicts: A list of dicts containing:
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:
- "signature": a thought signature - "signature": a thought signature
- "bookmark": a bookmark to identify the message part to apply the signature to. - "bookmark": a bookmark to identify the message part to apply the signature to.
The bookmark may contain either: The bookmark may contain one of:
- "text" - "function_call" (a function call ID string)
- "inline_data" - "text" (a text string)
messages: List of messages to search through. - "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 return
# For debugging, print out thought signatures and their bookmarks # For debugging, print out thought signatures and their bookmarks
logger.trace(f"Thought signatures to apply: {len(thought_signatures)}") logger.debug(f"Thought signatures to apply: {len(thought_signature_dicts)}")
for ts in thought_signatures: for ts in thought_signature_dicts:
bookmark = ts.get("bookmark") 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"] text = bookmark["text"]
log_display_text = f"{text[:50]}..." if len(text) > 50 else 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"): 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 # Get all assistant messages
non_fn_assistant_messages = [] assistant_messages = [
for message in messages: message
if not isinstance(message, Content) or not message.parts: for message in messages
continue if isinstance(message, Content) and message.role == "model"
# 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)
# 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) # Thought signatures are already in message order.
message_start_index = 0 # Track where to start searching for the next match thought_signatures_applied = 0
for thought_signature_dict in thought_signatures: 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") signature = thought_signature_dict.get("signature")
bookmark = thought_signature_dict.get("bookmark") bookmark = thought_signature_dict.get("bookmark")
if not signature: if not signature or not bookmark:
continue continue
# Search through remaining non-function assistant messages for a match # Search through remaining assistant messages for a match
for i in range(message_start_index, len(non_fn_assistant_messages)): for i in range(message_start_index, len(assistant_messages)):
message = non_fn_assistant_messages[i] message = assistant_messages[i]
if not message.parts: if not message.parts:
continue continue
# We're assuming that the thought signature always applies to the last part
last_part = message.parts[-1] 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 # If the bookmark matches the part...
# - is a prefix of that text (in case spoken text was truncated due to interruption) if self._thought_signature_bookmark_matches_part(bookmark, last_part):
# - is prefixed by that text (in case bookmark represents just first chunk of multi-chunk text) # Apply the thought signature
if bookmark_text := bookmark.get("text"): last_part.thought_signature = signature
if hasattr(last_part, "text") and last_part.text: thought_signatures_applied += 1
# 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
# Check if signed part has inline_data and last message part has matching inline_data # Update the start index and stop searching for a match
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:
message_start_index = i + 1 message_start_index = i + 1
break 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. tool_call_id: A unique identifier for the function call.
arguments: The arguments to pass to the function. arguments: The arguments to pass to the function.
context: The LLM context when the function call was made. 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 function_name: str
tool_call_id: str tool_call_id: str
arguments: Mapping[str, Any] arguments: Mapping[str, Any]
context: Any context: Any
append_extra_context_messages: Optional[List["LLMContextMessage"]] = None
@dataclass @dataclass
@@ -1745,16 +1741,12 @@ class FunctionCallInProgressFrame(ControlFrame, UninterruptibleFrame):
tool_call_id: Unique identifier for this function call. tool_call_id: Unique identifier for this function call.
arguments: Arguments passed to the function. arguments: Arguments passed to the function.
cancel_on_interruption: Whether to cancel this call if interrupted. 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 function_name: str
tool_call_id: str tool_call_id: str
arguments: Any arguments: Any
cancel_on_interruption: bool = False cancel_on_interruption: bool = False
append_extra_context_messages: Optional[List["LLMContextMessage"]] = None
@dataclass @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 self._function_calls_in_progress[frame.tool_call_id] = frame
async def _handle_function_call_result(self, frame: FunctionCallResultFrame): async def _handle_function_call_result(self, frame: FunctionCallResultFrame):

View File

@@ -932,7 +932,7 @@ class GoogleLLMService(LLMService):
reasoning_tokens = 0 reasoning_tokens = 0
grounding_metadata = None grounding_metadata = None
search_result = "" accumulated_text = ""
try: try:
# Generate content using either OpenAILLMContext or universal LLMContext # Generate content using either OpenAILLMContext or universal LLMContext
@@ -943,7 +943,6 @@ 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()
@@ -966,6 +965,7 @@ class GoogleLLMService(LLMService):
for candidate in chunk.candidates: for candidate in chunk.candidates:
if candidate.content and candidate.content.parts: if candidate.content and candidate.content.parts:
for part in candidate.content.parts: for part in candidate.content.parts:
function_call_id = None
if part.text: if part.text:
if part.thought: if part.thought:
# Gemini emits fully-formed thoughts rather # Gemini emits fully-formed thoughts rather
@@ -975,29 +975,20 @@ class GoogleLLMService(LLMService):
await self.push_frame(LLMThoughtTextFrame(part.text)) await self.push_frame(LLMThoughtTextFrame(part.text))
await self.push_frame(LLMThoughtEndFrame()) await self.push_frame(LLMThoughtEndFrame())
else: else:
search_result += part.text accumulated_text += part.text
await self.push_frame(LLMTextFrame(part.text)) await self.push_frame(LLMTextFrame(part.text))
elif part.function_call: elif part.function_call:
function_call = part.function_call function_call = part.function_call
id = function_call.id or str(uuid.uuid4()) function_call_id = function_call.id or str(uuid.uuid4())
logger.debug(f"Function call: {function_call.name}:{id}") logger.debug(
f"Function call: {function_call.name}:{function_call_id}"
)
function_calls.append( function_calls.append(
FunctionCallFromLLM( FunctionCallFromLLM(
context=context, context=context,
tool_call_id=id, tool_call_id=function_call_id,
function_name=function_call.name, function_name=function_call.name,
arguments=function_call.args or {}, 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: elif part.inline_data and part.inline_data.data:
@@ -1007,14 +998,14 @@ class GoogleLLMService(LLMService):
) )
await self.push_frame(frame) await self.push_frame(frame)
# With Gemini 3 Pro (and, contrary to Google's # Handle Gemini thought signatures.
# 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.
# #
# They should always be included in the last # - Gemini 2.5: they appear on function_call Parts,
# response Part. (*) # 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, # (*) Since we're using the streaming API, though,
# where text Parts may be split across multiple # where text Parts may be split across multiple
@@ -1022,34 +1013,44 @@ class GoogleLLMService(LLMService):
# signatures may actually appear with the first # signatures may actually appear with the first
# chunk (Gemini 2.5) or in a trailing empty-text # chunk (Gemini 2.5) or in a trailing empty-text
# chunk (Gemini 3 Pro). # 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 # Save a "bookmark" for the signature, so we
# can later stick it in the right place in # can later be sure we've put it in the right
# context when sending it back to the LLM to # place in context when sending the context
# continue the conversation. # back to the LLM to continue the conversation.
bookmark = {} bookmark = {}
if part.inline_data and part.inline_data.data: if part.function_call:
bookmark["inline_data"] = {"inline_data": part.inline_data} 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: elif part.text is not None:
# Account for Gemini 3 Pro trailing # Account for Gemini 3 Pro trailing
# empty-text chunk by using search_result, # empty-text chunk by using all the text
# which accumulates all text so far. # seen so far in this response's chunks.
bookmark["text"] = search_result bookmark["text"] = accumulated_text
await self.push_frame( else:
LLMMessagesAppendFrame( logger.warning("Thought signature found on unhandled Part type")
[ if bookmark:
self.get_llm_adapter().create_llm_specific_message( await self.push_frame(
{ LLMMessagesAppendFrame(
"type": "non_fn_thought_signature", [
"signature": part.thought_signature, self.get_llm_adapter().create_llm_specific_message(
"bookmark": bookmark, {
} "type": "thought_signature",
) "signature": part.thought_signature,
] "bookmark": bookmark,
}
)
]
)
) )
)
previous_part = part
if ( if (
candidate.grounding_metadata candidate.grounding_metadata
@@ -1098,7 +1099,7 @@ class GoogleLLMService(LLMService):
finally: finally:
if grounding_metadata and isinstance(grounding_metadata, dict): if grounding_metadata and isinstance(grounding_metadata, dict):
llm_search_frame = LLMSearchResponseFrame( llm_search_frame = LLMSearchResponseFrame(
search_result=search_result, search_result=accumulated_text,
origins=grounding_metadata["origins"], origins=grounding_metadata["origins"],
rendered_content=grounding_metadata["rendered_content"], rendered_content=grounding_metadata["rendered_content"],
) )

View File

@@ -132,9 +132,6 @@ class FunctionCallRunnerItem:
tool_call_id: A unique identifier for the function call. tool_call_id: A unique identifier for the function call.
arguments: The arguments for the function. arguments: The arguments for the function.
context: The LLM context. 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. run_llm: Optional flag to control LLM execution after function call.
""" """
@@ -143,7 +140,6 @@ class FunctionCallRunnerItem:
tool_call_id: str tool_call_id: str
arguments: Mapping[str, Any] arguments: Mapping[str, Any]
context: OpenAILLMContext | LLMContext context: OpenAILLMContext | LLMContext
append_extra_context_messages: Optional[List[LLMContextMessage]] = None
run_llm: Optional[bool] = None run_llm: Optional[bool] = None
@@ -465,7 +461,6 @@ class LLMService(AIService):
tool_call_id=function_call.tool_call_id, tool_call_id=function_call.tool_call_id,
arguments=function_call.arguments, arguments=function_call.arguments,
context=function_call.context, 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, function_name=runner_item.function_name,
tool_call_id=runner_item.tool_call_id, tool_call_id=runner_item.tool_call_id,
arguments=runner_item.arguments, arguments=runner_item.arguments,
append_extra_context_messages=runner_item.append_extra_context_messages,
cancel_on_interruption=item.cancel_on_interruption, cancel_on_interruption=item.cancel_on_interruption,
) )