Skip to content

fix: strip embedded thought signature from LiteLLM tool call ids - #6743

Open
benclarkeio wants to merge 1 commit into
google:mainfrom
benclarkeio:fix/thought-signature-in-litellm-tool-call-ids
Open

fix: strip embedded thought signature from LiteLLM tool call ids#6743
benclarkeio wants to merge 1 commit into
google:mainfrom
benclarkeio:fix/thought-signature-in-litellm-tool-call-ids

Conversation

@benclarkeio

Copy link
Copy Markdown
Contributor

Fixes #6742

Problem

LiteLLM embeds a Gemini thought_signature inside the tool call id, separated by __thought__ (see _THOUGHT_SIGNATURE_SEPARATOR). _message_to_generate_content_response already extracts that signature onto part.thought_signature, but then assigns the raw id to the function call:

thought_signature = _extract_thought_signature_from_tool_call(tool_call)
part = types.Part.from_function_call(...)
function_call.id = tool_call.id      # still "call_abc__thought__AY89a18..."
if thought_signature:
    part.thought_signature = thought_signature

So every consumer of function_call_id gets several hundred characters of base64 glued onto the real id:

call_abc__thought__AY89a18qpllj6mkpxPjhFEcbBbtEsqy4Ia4Eam2lD_NsoZ...

It leaks into logs, into any UI that displays a tool call id, and into anything that derives a name from one.

Fix

Split the separator off before assigning the id. The signature is already preserved on the part, so nothing is lost.

Why the round trip still works

  • _content_to_message_param re-attaches the signature to the outgoing tool call from part.thought_signature, via both provider_specific_fields and extra_content.google.thought_signature.
  • _extract_thought_signature_from_tool_call checks those two locations before falling back to the id-embedded form, which exists only for providers that drop the other channels.
  • Both sides of the call/response pairing (the assistant tool_calls entry and the tool message's tool_call_id) are generated from ADK's own stored parts, so stripping consistently keeps them matched.

Testing

Added test_message_to_generate_content_response_strips_signature_from_id, asserting the id is split and the signature still lands on the part. tests/unittests/models/test_litellm.py passes in full (387 passed). Formatted with pyink per the repo config.

Note

This replaces #6599, which was opened from the same work but carried a co-author trailer on its commit. Rebased onto current main (the surrounding code now goes through a local function_call variable) and reopened clean.

LiteLLM embeds a Gemini `thought_signature` in the tool call id, separated
by `__thought__`. `_message_to_generate_content_response` already lifts that
signature onto `part.thought_signature`, but then assigned the raw id to
`part.function_call.id`, so every consumer of `function_call_id` saw a few
hundred characters of base64 appended to the real id.

- Split the separator off the id before assigning it to the function call
- Leave `thought_signature` extraction unchanged, so nothing is lost

The round trip is unaffected: `_content_to_message_param` re-attaches the
signature to the outgoing tool call from `part.thought_signature` via
`provider_specific_fields` and `extra_content.google.thought_signature`, and
`_extract_thought_signature_from_tool_call` reads both of those before it
falls back to the id-embedded form. Both sides of the call/response pairing
are generated from ADK's own stored parts, so ids stay matched.

Fixes google#6742

@varunbiluri varunbiluri left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The ID should only be shortened when the suffix was actually accepted as the embedded thought signature. As written, every provider-issued ID containing __thought__ is truncated, even when _decode_thought_signature() rejects the suffix (or when the signature was sourced from extra_content / provider_specific_fields). Tool-call IDs are otherwise opaque, so an ID such as job__thought__not-base64 becomes job, which can break response pairing or collide with another call. Please have the extraction path return the cleaned ID only when it successfully decodes the ID-embedded form (or otherwise gate this assignment on successful decoding specifically from that fallback), and add coverage showing a non-signature ID containing the separator is preserved.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

LiteLLM tool call ids carry the embedded Gemini thought signature

3 participants