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`
- `LLMThoughtTextFrame`
- `LLMThoughtEndFrame`
3. A mechanism for appending arbitrary context messages after a function call
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
3. A generic mechanism for recording LLM thoughts to context, used
specifically to support Anthropic, whose thought signatures are expected to
appear alongside the text of the thoughts within assistant context
messages. See:
- `LLMThoughtEndFrame.signature`
- `LLMAssistantAggregator` handling of the above field
- `AnthropicLLMAdapter` handling of `"thought"` context messages
5. Google-specific logic for inserting non-function-call-related thought
signatures into the context, to help maintain thinking continuity in a
chain of LLM calls. See:
4. Google-specific logic for inserting thought signatures into the context,
to help maintain thinking continuity in a chain of LLM calls. See:
- `GoogleLLMService` sending `LLMMessagesAppendFrame`s to add LLM-specific
`"non_fn_thought_signature"` messages to context
- `GeminiLLMAdapter` handling of `"non_fn_thought_signature"` messages
6. An expansion of `TranscriptProcessor` to process LLM thoughts in addition
`"thought_signature"` messages to context
- `GeminiLLMAdapter` handling of `"thought_signature"` messages
5. An expansion of `TranscriptProcessor` to process LLM thoughts in addition
to user and assistant utterances. See:
- `TranscriptProcessor(process_thoughts=True)` (defaults to `False`)
- `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.",
}
)
# Here are some example prompts conducive to demonstrating
# thinking (picked from Google and Anthropic docs).
# Replace the above with one of these example prompts to demonstrate
# thinking.
# These examples come from Gemini and Anthropic docs.
# messages.append(
# {
# "role": "user",

View File

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

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,
)