Fix greptile review

litellm/litellm_core_utils/prompt_templates/common_utils.py

@@ -297,6 +297,7 @@ def _insert_user_continue_message(
     i = 1
     while i < len(result_messages):
         curr_message = result_messages[i]
+        inserted_continue_message = False
         if (
             curr_message["role"] == "assistant"
             and _counts_for_alternation(curr_message)
@@ -308,9 +309,10 @@ def _insert_user_continue_message(
                     if previous_message["role"] == "assistant":
                         result_messages.insert(i, continue_message)
                         i += 2
+                        inserted_continue_message = True
                     break
                 j -= 1
-        if i < len(result_messages):
+        if not inserted_continue_message:
             i += 1
 
     # Handle final message
@@ -378,40 +380,6 @@ def _insert_assistant_continue_message(
     return modified_messages
 
 
-def strip_tool_messages_for_alternating_roles(
-    messages: List[AllMessageValues],
-) -> List[AllMessageValues]:
-    """
-    Prepare history for strict user/assistant-only chat templates.
-
-    - Drop tool/function role messages
-    - Drop assistant tool-dispatch turns with no content
-    - Keep assistant content turns but remove tool metadata fields
-    """
-    cleaned_messages: List[AllMessageValues] = []
-
-    for message in messages:
-        role = message.get("role")
-        if role in ("tool", "function"):
-            continue
-
-        if role == "assistant":
-            assistant_message = message.copy()
-            assistant_message.pop("tool_calls", None)
-            assistant_message.pop("function_call", None)
-            assistant_message.pop("tool_call_id", None)
-
-            if assistant_message.get("content") is None:
-                continue
-
-            cleaned_messages.append(assistant_message)
-            continue
-
-        cleaned_messages.append(message)
-
-    return cleaned_messages
-
-
 def get_completion_messages(
     messages: List[AllMessageValues],
     assistant_continue_message: Optional[ChatCompletionAssistantMessage],

litellm/main.py

@@ -166,7 +166,6 @@ from .litellm_core_utils.fallback_utils import (
 from .litellm_core_utils.prompt_templates.common_utils import (
     add_system_prompt_to_messages,
     get_completion_messages,
-    strip_tool_messages_for_alternating_roles,
     update_messages_with_model_file_ids,
 )
 from .litellm_core_utils.prompt_templates.factory import (
@@ -1299,9 +1298,6 @@ def completion(  # type: ignore # noqa: PLR0915
     prompt_variables = cast(Optional[dict], kwargs.get("prompt_variables", None))
     litellm_system_prompt = kwargs.get("litellm_system_prompt", None)
     ### COPY MESSAGES ### - related issue https://github.com/BerriAI/litellm/discussions/4489
-    if ensure_alternating_roles:
-        messages = strip_tool_messages_for_alternating_roles(messages=messages)
-
     messages = get_completion_messages(
         messages=messages,
         ensure_alternating_roles=ensure_alternating_roles or False,

tests/llm_translation/test_prompt_factory.py

@@ -25,7 +25,6 @@ from litellm.litellm_core_utils.prompt_templates.factory import (
 )
 from litellm.litellm_core_utils.prompt_templates.common_utils import (
     get_completion_messages,
-    strip_tool_messages_for_alternating_roles,
 )
 from litellm.llms.vertex_ai.gemini.transformation import (
     _gemini_convert_messages_with_history,
@@ -800,8 +799,6 @@ def test_ensure_alternating_roles_with_tool_calls():
         {"role": "user", "content": "What about next week?"},
     ]
 
-    messages = strip_tool_messages_for_alternating_roles(messages)
-
     transformed_messages = get_completion_messages(
         messages=messages,
         assistant_continue_message=None,
@@ -811,6 +808,21 @@ def test_ensure_alternating_roles_with_tool_calls():
 
     assert transformed_messages == [
         {"role": "user", "content": "What's the weather?"},
+        {
+            "role": "assistant",
+            "content": None,
+            "tool_calls": [
+                {
+                    "id": "call_123",
+                    "type": "function",
+                    "function": {
+                        "name": "get_weather",
+                        "arguments": '{"location": "NYC"}',
+                    },
+                }
+            ],
+        },
+        {"role": "tool", "tool_call_id": "call_123", "content": "72F, sunny"},
         {"role": "assistant", "content": "It's 72F and sunny in NYC."},
         {"role": "user", "content": "What about tomorrow?"},
         {"role": "assistant", "content": "Please continue."},
@@ -820,6 +832,32 @@ def test_ensure_alternating_roles_with_tool_calls():
     ]
 
 
+def test_ensure_alternating_roles_three_consecutive_assistants():
+    messages = [
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "A1"},
+        {"role": "assistant", "content": "A2"},
+        {"role": "assistant", "content": "A3"},
+    ]
+
+    transformed_messages = get_completion_messages(
+        messages=messages,
+        assistant_continue_message=None,
+        user_continue_message=None,
+        ensure_alternating_roles=True,
+    )
+
+    assert transformed_messages == [
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "A1"},
+        {"role": "user", "content": "Please continue."},
+        {"role": "assistant", "content": "A2"},
+        {"role": "user", "content": "Please continue."},
+        {"role": "assistant", "content": "A3"},
+        {"role": "user", "content": "Please continue."},
+    ]
+
+
 def test_alternating_roles_e2e():
     from litellm.llms.custom_httpx.http_handler import HTTPHandler
     import json