diff --git a/litellm/llms/base_llm/base_utils.py b/litellm/llms/base_llm/base_utils.py index 1ad3bc4b4aa..0e4fa8c8eeb 100644 --- a/litellm/llms/base_llm/base_utils.py +++ b/litellm/llms/base_llm/base_utils.py @@ -312,14 +312,17 @@ def _leading_system_block_length(messages: Sequence[AllMessageValues]) -> int: def _closing_instruction_block_start(messages: Sequence[AllMessageValues], leading_length: int) -> int: - return next( + last_conversation_index: Final = next( ( - index + 1 + index for index in range(len(messages) - 1, leading_length - 1, -1) if messages[index]["role"] not in _INSTRUCTION_ROLES ), - leading_length, + None, ) + if last_conversation_index is None or messages[last_conversation_index]["role"] != "assistant": + return len(messages) + return last_conversation_index + 1 def _move_later_developer_messages_up(messages: Sequence[AllMessageValues]) -> tuple[AllMessageValues, ...]: @@ -347,9 +350,10 @@ def hoist_developer_messages_into_leading_system_message( Translate `developer` role to `system` role for OpenAI-compatible backends whose chat template allows a single system message and only at the start: developer messages that arrive after the first user turn move into the leading system - block, except a developer message that closes the conversation, which stays in - place so the request does not end on the assistant's turn. Each run of - consecutive system messages is then folded into one message in a single pass. + block, except a developer message that closes the conversation right after an + assistant turn, which stays in place so the request does not end on the + assistant's turn. Each run of consecutive system messages is then folded into + one message in a single pass. """ translated: Final = tuple(map(_as_system_message, _move_later_developer_messages_up(messages))) return tuple(_merged_system_runs(translated)) diff --git a/tests/test_litellm/llms/base_llm/test_base_utils.py b/tests/test_litellm/llms/base_llm/test_base_utils.py index d124b35f5a4..9e78d787d8e 100644 --- a/tests/test_litellm/llms/base_llm/test_base_utils.py +++ b/tests/test_litellm/llms/base_llm/test_base_utils.py @@ -275,12 +275,15 @@ class TestHoistDeveloperMessagesIntoLeadingSystemMessage: {"role": "system", "content": "Reply with the single word PONG and nothing else."}, ] - def test_only_the_closing_developer_run_stays_in_place_and_is_folded_into_one_message(self): + def test_only_the_closing_developer_run_after_an_assistant_turn_stays_in_place_and_is_folded_into_one_message( + self, + ): messages = [ {"role": "system", "content": "Base"}, {"role": "user", "content": "Turn 1"}, {"role": "developer", "content": "Update A"}, {"role": "user", "content": "Turn 2"}, + {"role": "assistant", "content": "Reply 2"}, {"role": "developer", "content": "Closing B"}, {"role": "developer", "content": "Closing C"}, ] @@ -288,9 +291,43 @@ class TestHoistDeveloperMessagesIntoLeadingSystemMessage: {"role": "system", "content": "Base\n\nUpdate A"}, {"role": "user", "content": "Turn 1"}, {"role": "user", "content": "Turn 2"}, + {"role": "assistant", "content": "Reply 2"}, {"role": "system", "content": "Closing B\n\nClosing C"}, ] + def test_developer_message_that_closes_the_conversation_after_a_user_turn_is_hoisted(self): + messages = [ + {"role": "user", "content": "Hi there"}, + {"role": "developer", "content": "Reply with exactly one word: the capital of France"}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "Reply with exactly one word: the capital of France"}, + {"role": "user", "content": "Hi there"}, + ] + + def test_developer_message_that_closes_the_conversation_after_a_tool_result_is_hoisted(self): + messages = [ + {"role": "system", "content": "Base"}, + {"role": "user", "content": "Look it up"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "Paris"}, + {"role": "developer", "content": "Answer with exactly one word."}, + ] + assert list(hoist_developer_messages_into_leading_system_message(messages)) == [ + {"role": "system", "content": "Base\n\nAnswer with exactly one word."}, + {"role": "user", "content": "Look it up"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "Paris"}, + ] + def test_hoisted_block_content_developer_message_merges_as_blocks_after_string_instructions(self): messages = [ {"role": "system", "content": "You are Codex"},