nanobot/tests/providers/test_enforce_role_alternation.py

"""Tests for LLMProvider._enforce_role_alternation."""

from nanobot.providers.base import _SYNTHETIC_USER_CONTENT, LLMProvider


class TestEnforceRoleAlternation:
    """Verify trailing-assistant removal and consecutive same-role merging."""

    def test_empty_messages(self):
        assert LLMProvider._enforce_role_alternation([]) == []

    def test_no_change_needed(self):
        msgs = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "Hi"},
            {"role": "assistant", "content": "Hello!"},
            {"role": "user", "content": "Bye"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert len(result) == 4
        assert result[-1]["role"] == "user"

    def test_trailing_assistant_removed(self):
        msgs = [
            {"role": "user", "content": "Hi"},
            {"role": "assistant", "content": "Hello!"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert len(result) == 1
        assert result[0]["role"] == "user"

    def test_multiple_trailing_assistants_removed(self):
        msgs = [
            {"role": "user", "content": "Hi"},
            {"role": "assistant", "content": "A"},
            {"role": "assistant", "content": "B"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert len(result) == 1
        assert result[0]["role"] == "user"

    def test_consecutive_user_messages_merged(self):
        msgs = [
            {"role": "user", "content": "Hello"},
            {"role": "user", "content": "How are you?"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert len(result) == 1
        assert "Hello" in result[0]["content"]
        assert "How are you?" in result[0]["content"]

    def test_consecutive_assistant_messages_merged(self):
        msgs = [
            {"role": "user", "content": "Hi"},
            {"role": "assistant", "content": "Hello!"},
            {"role": "assistant", "content": "How can I help?"},
            {"role": "user", "content": "Thanks"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert len(result) == 3
        assert "Hello!" in result[1]["content"]
        assert "How can I help?" in result[1]["content"]

    def test_system_messages_not_merged(self):
        msgs = [
            {"role": "system", "content": "System A"},
            {"role": "system", "content": "System B"},
            {"role": "user", "content": "Hi"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert len(result) == 3
        assert result[0]["content"] == "System A"
        assert result[1]["content"] == "System B"

    def test_tool_messages_not_merged(self):
        msgs = [
            {"role": "user", "content": "Hi"},
            {"role": "assistant", "content": None, "tool_calls": [{"id": "1"}]},
            {"role": "tool", "content": "result1", "tool_call_id": "1"},
            {"role": "tool", "content": "result2", "tool_call_id": "2"},
            {"role": "user", "content": "Next"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        tool_msgs = [m for m in result if m["role"] == "tool"]
        assert len(tool_msgs) == 2

    def test_consecutive_assistant_keeps_later_tool_call_message(self):
        msgs = [
            {"role": "user", "content": "Hi"},
            {"role": "assistant", "content": "Previous reply"},
            {"role": "assistant", "content": None, "tool_calls": [{"id": "1"}]},
            {"role": "tool", "content": "result1", "tool_call_id": "1"},
            {"role": "user", "content": "Next"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert result[1]["role"] == "assistant"
        assert result[1]["tool_calls"] == [{"id": "1"}]
        assert result[1]["content"] is None
        assert result[2]["role"] == "tool"

    def test_consecutive_assistant_does_not_overwrite_existing_tool_call_message(self):
        msgs = [
            {"role": "user", "content": "Hi"},
            {"role": "assistant", "content": None, "tool_calls": [{"id": "1"}]},
            {"role": "assistant", "content": "Later plain assistant"},
            {"role": "tool", "content": "result1", "tool_call_id": "1"},
            {"role": "user", "content": "Next"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert result[1]["role"] == "assistant"
        assert result[1]["tool_calls"] == [{"id": "1"}]
        assert result[1]["content"] is None
        assert result[2]["role"] == "tool"

    def test_consecutive_user_messages_preserve_text_before_multimodal_content(self):
        image = {
            "type": "image_url",
            "image_url": {"url": "data:image/png;base64,aW1hZ2U="},
        }
        msgs = [
            {"role": "user", "content": "Earlier unanswered question"},
            {
                "role": "user",
                "content": [image, {"type": "text", "text": "The error is here"}],
            },
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert result == [{
            "role": "user",
            "content": [
                {"type": "text", "text": "Earlier unanswered question"},
                image,
                {"type": "text", "text": "The error is here"},
            ],
        }]

    def test_consecutive_user_messages_preserve_multimodal_content_before_text(self):
        image = {
            "type": "image_url",
            "image_url": {"url": "data:image/png;base64,aW1hZ2U="},
        }
        msgs = [
            {
                "role": "user",
                "content": [image, {"type": "text", "text": "First question"}],
            },
            {"role": "user", "content": "Follow-up detail"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert result == [{
            "role": "user",
            "content": [
                image,
                {"type": "text", "text": "First question"},
                {"type": "text", "text": "Follow-up detail"},
            ],
        }]

    def test_original_messages_not_mutated(self):
        msgs = [
            {"role": "user", "content": "Hello"},
            {"role": "user", "content": "World"},
        ]
        original_first = dict(msgs[0])
        LLMProvider._enforce_role_alternation(msgs)
        assert msgs[0] == original_first
        assert len(msgs) == 2

    def test_trailing_assistant_recovered_as_user_when_only_system_remains(self):
        """A trailing assistant message must not be silently dropped.

        An externally supplied [system, assistant] sequence would otherwise leave
        only [system]. Most providers reject such requests, so the trailing
        assistant should be recovered as a user message instead.
        """
        msgs = [
            {"role": "system", "content": "You are helpful."},
            {"role": "assistant", "content": "Subagent completed successfully."},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert len(result) == 2
        assert result[0]["role"] == "system"
        assert result[1]["role"] == "user"
        assert "Subagent completed successfully." in result[1]["content"]

    def test_trailing_assistant_not_recovered_when_user_message_present(self):
        """Recovery should NOT happen when a user message already exists."""
        msgs = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "Hi"},
            {"role": "assistant", "content": "Hello!"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert len(result) == 2
        assert result[-1]["role"] == "user"

    def test_trailing_assistant_recovered_with_tool_result_preceding(self):
        """When only [system, tool, assistant] remains, recovery is not needed
        because tool messages are valid non-system content."""
        msgs = [
            {"role": "system", "content": "You are helpful."},
            {"role": "tool", "content": "result", "tool_call_id": "1"},
            {"role": "assistant", "content": "Done."},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert len(result) == 2
        assert result[-1]["role"] == "tool"

    def test_only_assistant_messages(self):
        msgs = [
            {"role": "assistant", "content": "A"},
            {"role": "assistant", "content": "B"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert result == []

    def test_realistic_conversation(self):
        msgs = [
            {"role": "system", "content": "You are helpful."},
            {"role": "user", "content": "What is 2+2?"},
            {"role": "assistant", "content": "4"},
            {"role": "user", "content": "And 3+3?"},
            {"role": "user", "content": "(please be quick)"},
            {"role": "assistant", "content": "6"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert len(result) == 4
        assert result[2]["role"] == "assistant"
        assert result[3]["role"] == "user"
        assert "And 3+3?" in result[3]["content"]
        assert "(please be quick)" in result[3]["content"]

    def test_leading_assistant_after_system_inserts_synthetic_user(self):
        """When the first non-system message is assistant (no tool_calls), a
        synthetic user message is inserted to prevent GLM error 1214."""
        msgs = [
            {"role": "system", "content": "sys"},
            {"role": "assistant", "content": "previous reply"},
            {"role": "tool", "tool_call_id": "tc_1", "content": "result"},
            {"role": "assistant", "content": "after tool"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        non_system = [m for m in result if m["role"] != "system"]
        assert non_system[0]["role"] == "user"
        assert non_system[0]["content"] == _SYNTHETIC_USER_CONTENT
        # The original assistant should follow.
        assert non_system[1]["role"] == "assistant"

    def test_leading_assistant_with_tool_calls_not_patched(self):
        """An assistant message with tool_calls at the start is left as-is
        because tool messages will follow and some providers accept this."""
        msgs = [
            {"role": "system", "content": "sys"},
            {"role": "assistant", "content": None, "tool_calls": [
                {"id": "tc_1", "type": "function", "function": {"name": "ls", "arguments": "{}"}}
            ]},
            {"role": "tool", "tool_call_id": "tc_1", "content": "result"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        non_system = [m for m in result if m["role"] != "system"]
        # The assistant has tool_calls so it should NOT be patched.
        assert non_system[0]["role"] == "assistant"
        assert non_system[0].get("tool_calls") is not None

    def test_user_after_system_not_patched(self):
        """Normal system→user sequence is not modified."""
        msgs = [
            {"role": "system", "content": "sys"},
            {"role": "user", "content": "hello"},
            {"role": "assistant", "content": "hi"},
        ]
        result = LLMProvider._enforce_role_alternation(msgs)
        assert result[1]["role"] == "user"
        assert result[1]["content"] == "hello"