From 22fc08d602598f5b5cbe2293ccbd146bb748e622 Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Wed, 18 Mar 2026 16:59:16 +0530 Subject: [PATCH] fix(prompting): revert _insert_assistant_continue_message to adjacent-check logic Restore backward-compatible behavior: only insert assistant_continue between directly adjacent user messages, not across tool-call chains. The _counts_for_alternation skip logic was a silent behavioral change for [user, assistant(tc), tool, user] sequences. Made-with: Cursor --- .../prompt_templates/common_utils.py | 42 +++++-------------- tests/llm_translation/test_prompt_factory.py | 3 +- 2 files changed, 12 insertions(+), 33 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index 8713d0283e..eb3755b71f 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -326,45 +326,25 @@ def _insert_assistant_continue_message( ) -> List[AllMessageValues]: """ Add assistant continuation messages between consecutive user messages. - Skips tool messages and assistant messages with tool calls in the - alternation check, matching strict templates like llama.cpp. - Args: - messages: List of message dictionaries - assistant_continue_message: Optional custom assistant message - ensure_alternating_roles: Whether to enforce alternating roles - - Returns: - Modified list of messages with inserted assistant messages + Only checks directly adjacent messages to preserve backward compatibility. """ if not ensure_alternating_roles or len(messages) <= 1: return messages + continue_message = assistant_continue_message or DEFAULT_ASSISTANT_CONTINUE_MESSAGE - insert_before_indexes = set() - - for i, message in enumerate(messages): - if message.get("role") != "user": - continue - - next_counted_index = i + 1 - while next_counted_index < len(messages) and not _counts_for_alternation( - messages[next_counted_index] - ): - next_counted_index += 1 - - if ( - next_counted_index < len(messages) - and messages[next_counted_index].get("role") == "user" - ): - # Insert before the next counted user turn. - # This avoids splitting assistant tool-call -> tool chains. - insert_before_indexes.add(next_counted_index) modified_messages: List[AllMessageValues] = [] - for idx, message in enumerate(messages): - if idx in insert_before_indexes: + for i, message in enumerate(messages): + if ( + i < len(messages) - 1 + and message.get("role") == "user" + and messages[i + 1].get("role") == "user" + ): + modified_messages.append(message) modified_messages.append(continue_message) - modified_messages.append(message) + else: + modified_messages.append(message) return modified_messages diff --git a/tests/llm_translation/test_prompt_factory.py b/tests/llm_translation/test_prompt_factory.py index 3a9f267e6a..fe46c24a29 100644 --- a/tests/llm_translation/test_prompt_factory.py +++ b/tests/llm_translation/test_prompt_factory.py @@ -860,6 +860,7 @@ def test_ensure_alternating_roles_three_consecutive_assistants(): def test_ensure_alternating_roles_does_not_split_tool_call_chain(): + """Tool-call chains [user, assistant(tc), tool, user] are preserved as-is.""" messages = [ {"role": "user", "content": "Search for X"}, { @@ -898,7 +899,6 @@ def test_ensure_alternating_roles_does_not_split_tool_call_chain(): ], }, {"role": "tool", "tool_call_id": "c1", "content": "results"}, - {"role": "assistant", "content": "Please continue."}, {"role": "user", "content": "Thanks, now do Y"}, ] @@ -945,7 +945,6 @@ def test_ensure_alternating_roles_trailing_tool_call_assistant(): } ], }, - {"role": "assistant", "content": "Please continue."}, {"role": "user", "content": "Please continue."}, ]