From 3cdabff323538df780ce0dc736c22d2f57fd323c Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 18 Mar 2026 15:45:09 +0530 Subject: [PATCH] Fix greptile review --- .../prompt_templates/common_utils.py | 38 ++-------------- litellm/main.py | 4 -- tests/llm_translation/test_prompt_factory.py | 44 +++++++++++++++++-- 3 files changed, 44 insertions(+), 42 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 8791c769af..aa2e07234d 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/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], diff --git a/litellm/main.py b/litellm/main.py index cb5a92caf8..81319bc432 100644 --- a/litellm/main.py +++ b/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, diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index 2eed9fa212..12efb47e06 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/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