Skip to content

fix(llm): keep Gemini thought signatures on replayed tool calls #1044

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ZhuLinsen wants to merge 1 commit into
base: main
Choose a base branch
from
+170 −10
Open

fix(llm): keep Gemini thought signatures on replayed tool calls #1044

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
19 changes: 14 additions & 5 deletions ms_agent/llm/openai_llm.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -467,14 +467,20 @@ def _merge_stream_message(self, pre_message_chunk: Optional[Message],
message.tool_calls[-1][
'tool_name'] = message_chunk.tool_calls[0][
'tool_name']
if message_chunk.tool_calls[0].get('extra_content'):
message.tool_calls[-1][
'extra_content'] = message_chunk.tool_calls[0][
'extra_content']
else:
message.tool_calls.append(
ToolCall(
id=message_chunk.tool_calls[0]['id'],
arguments=message_chunk.tool_calls[0]['arguments'],
type='function',
tool_name=message_chunk.tool_calls[0]['tool_name'],
index=message_chunk.tool_calls[0]['index']))
index=message_chunk.tool_calls[0]['index'],
extra_content=message_chunk.tool_calls[0].get(
'extra_content')))
else:
message.tool_calls = message_chunk.tool_calls
return message
Expand Down Expand Up @@ -566,8 +572,9 @@ def _stream_format_output_message(completion_chunk) -> Message:
index=tool_call.index,
type=tool_call.type,
arguments=tool_call.function.arguments,
tool_name=tool_call.function.name)
for tool_call in func
tool_name=tool_call.function.name,
extra_content=getattr(tool_call, 'extra_content',
None)) for tool_call in func
]
content = content or ''
reasoning_content = reasoning_content or ''
Expand Down Expand Up @@ -605,8 +612,10 @@ def _format_output_message(completion) -> Message:
index=getattr(tool_call, 'index', idx),
type=tool_call.type,
arguments=tool_call.function.arguments,
tool_name=tool_call.function.name) for idx, tool_call in
enumerate(completion.choices[0].message.tool_calls)
tool_name=tool_call.function.name,
extra_content=getattr(tool_call, 'extra_content', None))
for idx, tool_call in enumerate(
completion.choices[0].message.tool_calls)
]
usage = getattr(completion, 'usage', None)
cached, created = OpenAI._extract_cache_info(usage)
Expand Down
19 changes: 14 additions & 5 deletions ms_agent/llm/transport/openai_compat.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -585,14 +585,20 @@ def _merge_stream_message(self, pre_message_chunk: Optional[Message],
message.tool_calls[-1][
'tool_name'] = message_chunk.tool_calls[0][
'tool_name']
if message_chunk.tool_calls[0].get('extra_content'):
message.tool_calls[-1][
'extra_content'] = message_chunk.tool_calls[0][
'extra_content']
else:
message.tool_calls.append(
ToolCall(
id=message_chunk.tool_calls[0]['id'],
arguments=message_chunk.tool_calls[0]['arguments'],
type='function',
tool_name=message_chunk.tool_calls[0]['tool_name'],
index=message_chunk.tool_calls[0]['index']))
index=message_chunk.tool_calls[0]['index'],
extra_content=message_chunk.tool_calls[0].get(
'extra_content')))
else:
message.tool_calls = message_chunk.tool_calls
return message
Expand Down Expand Up @@ -679,8 +685,9 @@ def _stream_format_output_message(completion_chunk) -> Message:
index=tool_call.index,
type=tool_call.type,
arguments=tool_call.function.arguments,
tool_name=tool_call.function.name)
for tool_call in func
tool_name=tool_call.function.name,
extra_content=getattr(tool_call, 'extra_content',
None)) for tool_call in func
]
content = content or ''
reasoning_content = reasoning_content or ''
Expand Down Expand Up @@ -709,8 +716,10 @@ def _format_output_message(completion) -> Message:
index=getattr(tool_call, 'index', idx),
type=tool_call.type,
arguments=tool_call.function.arguments,
tool_name=tool_call.function.name) for idx, tool_call in
enumerate(completion.choices[0].message.tool_calls)
tool_name=tool_call.function.name,
extra_content=getattr(tool_call, 'extra_content', None))
for idx, tool_call in enumerate(
completion.choices[0].message.tool_calls)
]
usage = getattr(completion, 'usage', None)
cached, created = OpenAICompatTransport._extract_cache_info(usage)
Expand Down
6 changes: 6 additions & 0 deletions ms_agent/llm/utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,8 @@ class ToolCall(TypedDict, total=False):
type: str = 'function'
tool_name: str = ''
arguments: str = '{}'
# Provider data to replay verbatim, e.g. Gemini's thought signature.
extra_content: Optional[Dict[str, Any]] = None


class Tool(TypedDict, total=False):
Expand Down Expand Up @@ -143,6 +145,7 @@ def to_dict_clean(self):
json.loads(tool_call['arguments'])
except Exception:
tool_call['arguments'] = '{}'
extra_content = tool_call.get('extra_content')
raw_dict['tool_calls'][idx] = {
'id': tool_call['id'],
'type': tool_call['type'],
Expand All @@ -151,6 +154,9 @@ def to_dict_clean(self):
'arguments': tool_call['arguments'],
}
}
if extra_content:
raw_dict['tool_calls'][idx][
'extra_content'] = extra_content
required = ['content', 'role']
# Never send UI-only fields to model providers.
rm = [
Expand Down
136 changes: 136 additions & 0 deletions tests/llm/test_tool_call_extra_content.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,136 @@
"""Gemini 3 returns a thought signature on each tool call as
``extra_content.google.thought_signature`` and answers 400 when the next
request replays the call without it, so the field must survive parsing, stream
merging and serialization in both OpenAI-compatible clients."""
import unittest

from openai.types.chat import ChatCompletion, ChatCompletionChunk
from omegaconf import OmegaConf

from ms_agent.llm.openai_llm import OpenAI
from ms_agent.llm.transport.openai_compat import OpenAICompatTransport
from ms_agent.llm.utils import Message

EXTRA_CONTENT = {'google': {'thought_signature': 'sig=='}}


def _completion(extra=EXTRA_CONTENT):
tool_call = {
'id': 'call_1',
'type': 'function',
'function': {
'name': 'read_file',
'arguments': '{"path": "a.txt"}',
},
}
if extra:
tool_call['extra_content'] = extra
return ChatCompletion.model_validate({
'id': 'c1',
'object': 'chat.completion',
'created': 0,
'model': 'm',
'usage': {
'prompt_tokens': 1,
'completion_tokens': 1,
'total_tokens': 2,
},
'choices': [{
'index': 0,
'finish_reason': 'tool_calls',
'message': {
'role': 'assistant',
'content': None,
'tool_calls': [tool_call],
},
}],
})


def _chunk(index, call_id=None, name=None, arguments='', extra=None):
tool_call = {
'index': index,
'id': call_id,
'type': 'function' if call_id else None,
'function': {
'name': name,
'arguments': arguments,
},
}
if extra:
tool_call['extra_content'] = extra
return ChatCompletionChunk.model_validate({
'id': 'c1',
'object': 'chat.completion.chunk',
'created': 0,
'model': 'm',
'choices': [{
'index': 0,
'delta': {
'tool_calls': [tool_call],
},
}],
})


def _merge_stream(client, chunks):
merged = None
for chunk in chunks:
merged = client._merge_stream_message(
client, merged, client._stream_format_output_message(chunk))
return merged


class TestToolCallExtraContent(unittest.TestCase):
clients = (OpenAICompatTransport, OpenAI)

def test_non_stream_round_trip(self):
for client in self.clients:
message = client._format_output_message(_completion())
wire = message.to_dict_clean()['tool_calls'][0]
self.assertEqual(wire['extra_content'], EXTRA_CONTENT)

def test_stream_signature_on_any_chunk_of_any_call(self):
"""The signature may arrive on the opening chunk or a later one, and
a parallel call must keep its own."""
other = {'google': {'thought_signature': 'other=='}}
chunks = [
_chunk(0, 'call_1', 'read_file', '{"path"', EXTRA_CONTENT),
_chunk(0, arguments=': "a.txt"}'),
_chunk(1, 'call_2', 'read_file', '{"path"'),
_chunk(1, arguments=': "b.txt"}', extra=other),
]
for client in self.clients:
wire = _merge_stream(client, chunks).to_dict_clean()['tool_calls']
self.assertEqual(wire[0]['extra_content'], EXTRA_CONTENT)
self.assertEqual(wire[0]['function']['arguments'],
'{"path": "a.txt"}')
self.assertEqual(wire[1]['extra_content'], other)

def test_other_providers_send_no_extra_field(self):
"""Strict endpoints can reject unknown tool-call fields."""
for client in self.clients:
message = client._format_output_message(_completion(extra=None))
wire = message.to_dict_clean()['tool_calls'][0]
self.assertNotIn('extra_content', wire)

def test_next_request_replays_signature(self):
conf = OmegaConf.create({
'llm': {
'model': 'm',
'openai_base_url': 'http://x',
'openai_api_key': 'dummy',
},
'generation_config': {},
})
assistant = OpenAI._format_output_message(_completion())
tool = Message(role='tool', content='ok', tool_call_id='call_1')
request = OpenAI(conf)._format_input_message([
Message(role='user', content='read a.txt'), assistant, tool
])
sent = [m for m in request if m['role'] == 'assistant'][0]
self.assertEqual(sent['tool_calls'][0]['extra_content'], EXTRA_CONTENT)


if __name__ == '__main__':
unittest.main()