Merge 81e675210c into b4a8bc715f

litellm/llms/vertex_ai/gemini/transformation.py

@@ -44,6 +44,7 @@ from litellm.types.files import (
 from litellm.types.llms.openai import (
     AllMessageValues,
     ChatCompletionAssistantMessage,
+    ChatCompletionAssistantToolCall,
     ChatCompletionTextObject,
     ChatCompletionUserMessage,
 )
@@ -844,6 +845,11 @@ def gemini_convert_messages_with_history(
                     )  # add a blank text, to ensure Gemini doesn't fail the request.
                 contents.append(ContentType(role="user", parts=user_content))
             assistant_content = []
+            # Tool calls of every assistant message merged into this model turn.
+            # A tool result may answer a call from an earlier message of the run
+            # (e.g. [function_call, reasoning, function_call] from the Responses
+            # API), so names are recovered from all of them, not only the last.
+            turn_tool_calls: tuple[ChatCompletionAssistantToolCall, ...] = ()
             ## MERGE CONSECUTIVE ASSISTANT CONTENT ##
             while msg_i < len(messages) and messages[msg_i]["role"] == "assistant":
                 if isinstance(messages[msg_i], BaseModel):
@@ -954,6 +960,13 @@ def gemini_convert_messages_with_history(
                     # corresponding tool call for tool response message".
                     if assistant_msg.get("tool_calls") or assistant_msg.get("function_call") is not None:
                         last_message_with_tool_calls = assistant_msg
+                        if assistant_msg.get("tool_calls"):
+                            prior: Final = turn_tool_calls
+                            turn_tool_calls = (*prior, *assistant_msg["tool_calls"])
+                            if prior:
+                                last_message_with_tool_calls = ChatCompletionAssistantMessage(
+                                    **{**assistant_msg, "tool_calls": [*turn_tool_calls]}
+                                )
 
                 ## HANDLE SERVER-SIDE TOOL INVOCATIONS (context circulation)
                 _psf = assistant_msg.get("provider_specific_fields")

tests/unit/llms/vertex_ai/gemini/test_tool_call_followed_by_text_assistant.py

@@ -55,3 +55,99 @@ def test_tool_result_matches_tool_call_with_text_assistant_in_between():
     ]
     assert function_responses, f"expected a functionResponse part, got: {contents}"
     assert function_responses[0]["name"] == "shell"
+
+
+def _function_response_names(contents):
+    return [
+        part["function_response"]["name"]
+        for content in contents
+        for part in content["parts"]
+        if isinstance(part, dict) and part.get("function_response")
+    ]
+
+
+def test_tool_results_match_tool_calls_split_across_assistant_messages():
+    """Each tool result is matched against every tool call of the merged model turn.
+
+    Gemini streams `[functionCall, thought, functionCall]` when it reasons between two
+    calls. Replayed as chat messages this becomes two assistant messages with one tool
+    call each; the result for the first call used to be looked up only in the second
+    message and failed with "Missing corresponding tool call" (#49063).
+    """
+    messages = [
+        {"role": "user", "content": "Weather in Paris and the time in Rome?"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "call_a",
+                    "type": "function",
+                    "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'},
+                }
+            ],
+        },
+        {
+            "role": "assistant",
+            "content": None,
+            "reasoning_content": "Now the time in Rome.",
+            "tool_calls": [
+                {
+                    "id": "call_b",
+                    "type": "function",
+                    "function": {"name": "get_time", "arguments": '{"city": "Rome"}'},
+                }
+            ],
+        },
+        {"role": "tool", "tool_call_id": "call_a", "content": "sunny"},
+        {"role": "tool", "tool_call_id": "call_b", "content": "noon"},
+    ]
+
+    contents = gemini_convert_messages_with_history(messages=messages)
+
+    function_calls = [
+        part["function_call"]["name"]
+        for content in contents
+        for part in content["parts"]
+        if isinstance(part, dict) and part.get("function_call")
+    ]
+    assert function_calls == ["get_weather", "get_time"]
+    assert _function_response_names(contents) == ["get_weather", "get_time"]
+    # The original assistant messages are not modified.
+    assert [tc["id"] for tc in messages[2]["tool_calls"]] == ["call_b"]
+
+
+def test_responses_input_with_reasoning_between_function_calls_converts_for_gemini():
+    from litellm.responses.litellm_completion_transformation.transformation import (
+        LiteLLMCompletionResponsesConfig,
+    )
+
+    chat_messages = LiteLLMCompletionResponsesConfig._transform_response_input_param_to_chat_completion_message(
+        input=[
+            {"role": "user", "content": "Weather in Paris and Rome?"},
+            {
+                "type": "function_call",
+                "call_id": "call_a",
+                "name": "get_weather",
+                "arguments": '{"city": "Paris"}',
+            },
+            {
+                "type": "reasoning",
+                "id": "rs_1",
+                "content": [{"type": "output_text", "text": "Now Rome too."}],
+            },
+            {
+                "type": "function_call",
+                "call_id": "call_b",
+                "name": "get_weather",
+                "arguments": '{"city": "Rome"}',
+            },
+            {"type": "function_call_output", "call_id": "call_a", "output": "sunny"},
+            {"type": "function_call_output", "call_id": "call_b", "output": "cloudy"},
+        ],
+        replay_reasoning=True,
+    )
+
+    contents = gemini_convert_messages_with_history(messages=chat_messages)
+
+    assert _function_response_names(contents) == ["get_weather", "get_weather"]