diff --git a/ms_agent/llm/openai_llm.py b/ms_agent/llm/openai_llm.py index 730205e2a..ff4bf6b92 100644 --- a/ms_agent/llm/openai_llm.py +++ b/ms_agent/llm/openai_llm.py @@ -467,6 +467,10 @@ 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( @@ -474,7 +478,9 @@ def _merge_stream_message(self, pre_message_chunk: Optional[Message], 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 @@ -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 '' @@ -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) diff --git a/ms_agent/llm/transport/openai_compat.py b/ms_agent/llm/transport/openai_compat.py index 4dbe214ee..9bec7394b 100644 --- a/ms_agent/llm/transport/openai_compat.py +++ b/ms_agent/llm/transport/openai_compat.py @@ -585,6 +585,10 @@ 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( @@ -592,7 +596,9 @@ def _merge_stream_message(self, pre_message_chunk: Optional[Message], 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 @@ -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 '' @@ -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) diff --git a/ms_agent/llm/utils.py b/ms_agent/llm/utils.py index 09a04a444..7d619abae 100644 --- a/ms_agent/llm/utils.py +++ b/ms_agent/llm/utils.py @@ -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): @@ -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'], @@ -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 = [ diff --git a/tests/llm/test_tool_call_extra_content.py b/tests/llm/test_tool_call_extra_content.py new file mode 100644 index 000000000..fa40cb54d --- /dev/null +++ b/tests/llm/test_tool_call_extra_content.py @@ -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()