diff --git a/litellm/llms/vertex_ai/gemini/transformation.py b/litellm/llms/vertex_ai/gemini/transformation.py index 87bd4843822..afdadb46730 100644 --- a/litellm/llms/vertex_ai/gemini/transformation.py +++ b/litellm/llms/vertex_ai/gemini/transformation.py @@ -307,6 +307,7 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 contents: List[ContentType] = [] last_message_with_tool_calls = None + message_by_tool_call_id: Dict[str, ChatCompletionAssistantMessage] = {} msg_i = 0 tool_call_responses = [] @@ -527,9 +528,10 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 assistant_content.append(_part) ## HANDLE ASSISTANT FUNCTION CALL + tool_calls = assistant_msg.get("tool_calls") + function_call = assistant_msg.get("function_call") if ( - assistant_msg.get("tool_calls", []) is not None - or assistant_msg.get("function_call") is not None + tool_calls or function_call is not None ): # support assistant tool invoke conversion gemini_tool_call_parts = convert_to_gemini_tool_call_invoke( assistant_msg, model=model @@ -542,6 +544,11 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 excluded_keys=["thoughtSignature"], ): assistant_content.append(gemini_tool_call_part) + if tool_calls: + for tool_call in tool_calls: + tool_call_id = tool_call.get("id") + if tool_call_id: + message_by_tool_call_id[tool_call_id] = assistant_msg last_message_with_tool_calls = assistant_msg ## HANDLE SERVER-SIDE TOOL INVOCATIONS (context circulation) @@ -590,8 +597,17 @@ def _gemini_convert_messages_with_history( # noqa: PLR0915 msg_i < len(messages) and messages[msg_i]["role"] in tool_call_message_roles ): + tool_result_message = messages[msg_i] + tool_call_id = tool_result_message.get("tool_call_id") + message_with_tool_call = ( + message_by_tool_call_id.get( + tool_call_id, last_message_with_tool_calls + ) + if tool_call_id is not None + else last_message_with_tool_calls + ) _part = convert_to_gemini_tool_call_result( - messages[msg_i], last_message_with_tool_calls # type: ignore + tool_result_message, message_with_tool_call # type: ignore ) msg_i += 1 # Handle both single part and list of parts (for Computer Use with images) diff --git a/tests/litellm/llms/vertex_ai/gemini/test_transformation.py b/tests/litellm/llms/vertex_ai/gemini/test_transformation.py index 963e2d273a7..d4616adfa15 100644 --- a/tests/litellm/llms/vertex_ai/gemini/test_transformation.py +++ b/tests/litellm/llms/vertex_ai/gemini/test_transformation.py @@ -10,8 +10,6 @@ from litellm.llms.vertex_ai.gemini import transformation from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( VertexGeminiConfig, ) -from litellm.types.llms import openai -from litellm.types import completion from litellm.types.llms.vertex_ai import RequestBody @@ -95,6 +93,65 @@ async def test__transform_request_body_metadata(): } +def test_gemini_tool_result_resolves_after_text_only_assistant_message(): + messages = [ + {"role": "user", "content": "Read the file before answering."}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_read_file", + "type": "function", + "function": { + "name": "read_file", + "arguments": '{"path": "/tmp/example.txt"}', + }, + } + ], + }, + { + "role": "assistant", + "content": "I will inspect the file and then summarize it.", + "tool_calls": [], + }, + { + "role": "tool", + "tool_call_id": "call_read_file", + "content": '{"contents": "hello from the file"}', + }, + ] + + contents = transformation._gemini_convert_messages_with_history(messages) + + assert contents == [ + {"parts": [{"text": "Read the file before answering."}], "role": "user"}, + { + "parts": [ + { + "function_call": { + "name": "read_file", + "args": {"path": "/tmp/example.txt"}, + } + }, + {"text": "I will inspect the file and then summarize it."}, + ], + "role": "model", + }, + { + "parts": [ + { + "function_response": { + "name": "read_file", + "response": {"contents": "hello from the file"}, + } + } + ], + "role": "user", + }, + ] + + @pytest.mark.asyncio async def test__transform_request_body_labels_and_metadata(): """