wehub-resource-sync
97e91a83f3
Ruff / Ruff (push) Waiting to run
Test / Core Tests (push) Waiting to run
Test / Offline Coverage Tests (Python 3.10) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.11) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.12) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.13) (push) Waiting to run
Test / Offline Coverage Tests (Python 3.9) (push) Waiting to run
Test / Full Coverage (Python 3.11) (push) Waiting to run
Test / Core Provider Tests (OpenAI) (push) Blocked by required conditions
Test / Core Provider Tests (Anthropic) (push) Blocked by required conditions
Test / Core Provider Tests (Google) (push) Blocked by required conditions
Test / Core Provider Tests (Other) (push) Blocked by required conditions
Test / Anthropic Tests (push) Blocked by required conditions
Test / Gemini Tests (push) Blocked by required conditions
Test / Google GenAI Tests (push) Blocked by required conditions
Test / Vertex AI Tests (push) Blocked by required conditions
Test / OpenAI Tests (push) Blocked by required conditions
Test / Writer Tests (push) Blocked by required conditions
Test / Auto Client Tests (push) Blocked by required conditions
ty / type-check (push) Waiting to run
345 lines
13 KiB
Python
345 lines
13 KiB
Python
"""
|
|
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"
|