chore: import upstream snapshot with attribution

tests/processing/test_message_processing.py

@@ -0,0 +1,344 @@
+"""
+Tests for message processing optimizations.
+"""
+
+from typing import cast
+
+from openai.types.chat import ChatCompletionMessageParam
+
+from instructor.utils import (
+    merge_consecutive_messages,
+    get_message_content,
+    transform_to_gemini_prompt,
+    update_gemini_kwargs,
+    combine_system_messages,
+    extract_system_messages,
+    SystemMessage,
+)
+
+
+class TestMergeConsecutiveMessages:
+    """Test the merge_consecutive_messages function."""
+
+    def test_empty_messages(self):
+        """Test merging empty messages list."""
+        result = merge_consecutive_messages([])
+        assert result == []
+
+    def test_single_message(self):
+        """Test merging a single message."""
+        messages = [{"role": "user", "content": "Hello"}]
+        result = merge_consecutive_messages(messages)
+        assert result == messages
+
+    def test_consecutive_same_role(self):
+        """Test merging consecutive messages with the same role."""
+        messages = [
+            {"role": "user", "content": "Hello"},
+            {"role": "user", "content": "World"},
+        ]
+        result = merge_consecutive_messages(messages)
+        assert len(result) == 1
+        assert result[0]["role"] == "user"
+        assert "Hello" in result[0]["content"]
+        assert "World" in result[0]["content"]
+
+    def test_alternating_roles(self):
+        """Test merging messages with alternating roles."""
+        messages = [
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi there"},
+            {"role": "user", "content": "How are you?"},
+        ]
+        result = merge_consecutive_messages(messages)
+        assert len(result) == 3
+        assert result[0]["role"] == "user"
+        assert result[1]["role"] == "assistant"
+        assert result[2]["role"] == "user"
+
+    def test_mixed_content_types(self):
+        """Test merging messages with mixed content types."""
+        messages = [
+            {"role": "user", "content": "Hello"},
+            {"role": "user", "content": [{"type": "text", "text": "World"}]},
+        ]
+        result = merge_consecutive_messages(messages)
+        assert len(result) == 1
+        assert result[0]["role"] == "user"
+        assert isinstance(result[0]["content"], list)
+        assert len(result[0]["content"]) == 2
+
+    def test_multiple_consecutive(self):
+        """Test merging multiple consecutive messages."""
+        messages = [
+            {"role": "user", "content": "Hello"},
+            {"role": "user", "content": "World"},
+            {"role": "assistant", "content": "Hi there"},
+            {"role": "assistant", "content": "How can I help?"},
+            {"role": "user", "content": "I need help"},
+        ]
+        result = merge_consecutive_messages(messages)
+        assert len(result) == 3
+        assert result[0]["role"] == "user"
+        assert "Hello" in result[0]["content"]
+        assert "World" in result[0]["content"]
+        assert result[1]["role"] == "assistant"
+        assert "Hi there" in result[1]["content"]
+        assert "How can I help?" in result[1]["content"]
+        assert result[2]["role"] == "user"
+        assert "I need help" in result[2]["content"]
+
+
+class TestGetMessageContent:
+    """Test the get_message_content function."""
+
+    def test_string_content(self):
+        """Test getting content from a message with string content."""
+        message = cast(ChatCompletionMessageParam, {"role": "user", "content": "Hello"})
+        result = get_message_content(message)
+        assert result == ["Hello"]
+
+    def test_list_content(self):
+        """Test getting content from a message with list content."""
+        message = cast(
+            ChatCompletionMessageParam,
+            {"role": "user", "content": [{"type": "text", "text": "Hello"}]},
+        )
+        result = get_message_content(message)
+        assert result == [{"type": "text", "text": "Hello"}]
+
+    def test_empty_content(self):
+        """Test getting content from a message with empty content."""
+        message = cast(ChatCompletionMessageParam, {"role": "user", "content": ""})
+        result = get_message_content(message)
+        assert result == [""]
+
+    def test_none_content(self):
+        """Test getting content from a message with None content."""
+        message = cast(ChatCompletionMessageParam, {"role": "user", "content": None})
+        result = get_message_content(message)
+        assert result == [""]
+
+    def test_missing_content(self):
+        """Test getting content from a message with missing content."""
+        message = cast(ChatCompletionMessageParam, {"role": "user"})
+        result = get_message_content(message)
+        assert result == [""]
+
+    def test_empty_message(self):
+        """Test getting content from an empty message."""
+        message = cast(ChatCompletionMessageParam, {})
+        result = get_message_content(message)
+        assert result == [""]
+
+
+class TestTransformToGeminiPrompt:
+    """Test the transform_to_gemini_prompt function."""
+
+    def test_empty_messages(self):
+        """Test transforming empty messages."""
+        result = transform_to_gemini_prompt([])
+        assert result == []
+
+    def test_user_message(self):
+        """Test transforming a user message."""
+        messages = [{"role": "user", "content": "Hello"}]
+        result = transform_to_gemini_prompt(messages)
+        assert len(result) == 1
+        assert result[0]["role"] == "user"
+        assert result[0]["parts"] == ["Hello"]
+
+    def test_assistant_message(self):
+        """Test transforming an assistant message."""
+        messages = [{"role": "assistant", "content": "Hello"}]
+        result = transform_to_gemini_prompt(messages)
+        assert len(result) == 1
+        assert result[0]["role"] == "model"
+        assert result[0]["parts"] == ["Hello"]
+
+    def test_system_message(self):
+        """Test transforming a system message."""
+        messages = [{"role": "system", "content": "You are an AI assistant"}]
+        result = transform_to_gemini_prompt(messages)
+        assert len(result) == 1
+        assert result[0]["role"] == "user"
+        assert "*You are an AI assistant*" in result[0]["parts"][0]
+
+    def test_full_conversation(self):
+        """Test transforming a full conversation."""
+        messages = [
+            {"role": "system", "content": "You are an AI assistant"},
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi there"},
+            {"role": "user", "content": "How are you?"},
+        ]
+        result = transform_to_gemini_prompt(messages)
+        assert len(result) == 3
+        assert result[0]["role"] == "user"
+        assert "*You are an AI assistant*" in result[0]["parts"][0]
+        assert "Hello" in result[0]["parts"][1]
+        assert result[1]["role"] == "model"
+        assert result[1]["parts"] == ["Hi there"]
+        assert result[2]["role"] == "user"
+        assert result[2]["parts"] == ["How are you?"]
+
+    def test_multiple_system_messages(self):
+        """Test transforming multiple system messages."""
+        messages = [
+            {"role": "system", "content": "You are an AI assistant"},
+            {"role": "system", "content": "Be helpful and concise"},
+            {"role": "user", "content": "Hello"},
+        ]
+        result = transform_to_gemini_prompt(messages)
+        assert len(result) == 1
+        assert result[0]["role"] == "user"
+        assert any("You are an AI assistant" in part for part in result[0]["parts"])
+        assert any("Be helpful and concise" in part for part in result[0]["parts"])
+        assert any("Hello" in part for part in result[0]["parts"])
+
+
+class TestUpdateGeminiKwargs:
+    """Test the update_gemini_kwargs function."""
+
+    def test_transform_messages(self):
+        """Test transforming messages to Gemini format."""
+        kwargs = {"messages": [{"role": "user", "content": "Hello"}]}
+        result = update_gemini_kwargs(kwargs)
+        assert "contents" in result
+        assert "messages" not in result
+        assert len(result["contents"]) == 1
+        assert result["contents"][0]["role"] == "user"
+
+    def test_generation_config(self):
+        """Test updating generation config."""
+        kwargs = {
+            "messages": [{"role": "user", "content": "Hello"}],
+            "generation_config": {
+                "max_tokens": 100,
+                "temperature": 0.7,
+                "n": 3,
+                "top_p": 0.9,
+                "stop": ["END"],
+            },
+        }
+        result = update_gemini_kwargs(kwargs)
+        assert "generation_config" in result
+        assert "max_output_tokens" in result["generation_config"]
+        assert "candidate_count" in result["generation_config"]
+        assert "stop_sequences" in result["generation_config"]
+        assert "max_tokens" not in result["generation_config"]
+        assert "n" not in result["generation_config"]
+        assert "stop" not in result["generation_config"]
+
+    def test_safety_settings(self):
+        """Test setting safety settings."""
+        kwargs = {
+            "messages": [{"role": "user", "content": "Hello"}],
+        }
+        result = update_gemini_kwargs(kwargs)
+        assert "safety_settings" in result
+        assert len(result["safety_settings"]) >= 3  # At least 3 safety settings
+
+    def test_existing_safety_settings(self):
+        """Test respecting existing safety settings."""
+        from google.genai.types import HarmCategory, HarmBlockThreshold
+
+        kwargs = {
+            "messages": [{"role": "user", "content": "Hello"}],
+            "safety_settings": {
+                HarmCategory.HARM_CATEGORY_HATE_SPEECH: HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE
+            },
+        }
+        result = update_gemini_kwargs(kwargs)
+        assert (
+            result["safety_settings"][HarmCategory.HARM_CATEGORY_HATE_SPEECH]
+            == HarmBlockThreshold.BLOCK_MEDIUM_AND_ABOVE
+        )
+
+
+class TestSystemMessages:
+    """Test the system message utility functions."""
+
+    def test_combine_system_messages_strings(self):
+        """Test combining two string system messages."""
+        existing = "You are an AI assistant"
+        new = "Be helpful"
+        result = combine_system_messages(existing, new)
+        assert result == "You are an AI assistant\n\nBe helpful"
+
+    def test_combine_system_messages_lists(self):
+        """Test combining two list system messages."""
+        existing = [SystemMessage(type="text", text="You are an AI assistant")]
+        new = [SystemMessage(type="text", text="Be helpful")]
+        result = combine_system_messages(existing, new)
+        assert len(result) == 2
+        assert result[0]["text"] == "You are an AI assistant"
+        assert result[1]["text"] == "Be helpful"
+
+    def test_combine_system_messages_mixed(self):
+        """Test combining mixed system message types."""
+        existing = "You are an AI assistant"
+        new = [SystemMessage(type="text", text="Be helpful")]
+        result = combine_system_messages(existing, new)
+        assert len(result) == 2
+        assert result[0]["text"] == "You are an AI assistant"
+        assert result[1]["text"] == "Be helpful"
+
+    def test_combine_system_messages_none(self):
+        """Test combining None with a system message."""
+        existing = None
+        new = "Be helpful"
+        result = combine_system_messages(existing, new)
+        assert result == "Be helpful"
+
+    def test_extract_system_messages_empty(self):
+        """Test extracting system messages from an empty list."""
+        messages = []
+        result = extract_system_messages(messages)
+        assert result == []
+
+    def test_extract_system_messages_no_system(self):
+        """Test extracting system messages when there are none."""
+        messages = [
+            {"role": "user", "content": "Hello"},
+            {"role": "assistant", "content": "Hi there"},
+        ]
+        result = extract_system_messages(messages)
+        assert result == []
+
+    def test_extract_system_messages_string(self):
+        """Test extracting string system messages."""
+        messages = [
+            {"role": "system", "content": "You are an AI assistant"},
+            {"role": "user", "content": "Hello"},
+        ]
+        result = extract_system_messages(messages)
+        assert len(result) == 1
+        assert result[0]["type"] == "text"
+        assert result[0]["text"] == "You are an AI assistant"
+
+    def test_extract_system_messages_list(self):
+        """Test extracting list system messages."""
+        messages = [
+            {
+                "role": "system",
+                "content": [{"type": "text", "text": "You are an AI assistant"}],
+            },
+            {"role": "user", "content": "Hello"},
+        ]
+        result = extract_system_messages(messages)
+        assert len(result) == 1
+        assert result[0]["type"] == "text"
+        assert result[0]["text"] == "You are an AI assistant"
+
+    def test_extract_system_messages_multiple(self):
+        """Test extracting multiple system messages."""
+        messages = [
+            {"role": "system", "content": "You are an AI assistant"},
+            {"role": "system", "content": "Be helpful"},
+            {"role": "user", "content": "Hello"},
+        ]
+        result = extract_system_messages(messages)
+        assert len(result) == 2
+        assert result[0]["text"] == "You are an AI assistant"
+        assert result[1]["text"] == "Be helpful"