mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-08 03:08:45 +00:00
344 lines
11 KiB
Python
344 lines
11 KiB
Python
"""
|
|
Test message sanitization for Anthropic API when modify_params=True
|
|
|
|
Tests three cases:
|
|
A. Missing tool_result for tool_use (orphaned tool calls)
|
|
B. Orphaned tool_result without matching tool_use
|
|
C. Empty text content
|
|
"""
|
|
|
|
import pytest
|
|
import sys
|
|
import os
|
|
|
|
# Add the parent directory to the path so we can import litellm
|
|
sys.path.insert(
|
|
0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../.."))
|
|
)
|
|
|
|
import litellm
|
|
from litellm.litellm_core_utils.prompt_templates.factory import (
|
|
sanitize_messages_for_tool_calling,
|
|
anthropic_messages_pt,
|
|
)
|
|
|
|
|
|
class TestMessageSanitization:
|
|
"""Test message sanitization for tool calling scenarios"""
|
|
|
|
def setup_method(self):
|
|
"""Setup for each test"""
|
|
# Save original modify_params value
|
|
self.original_modify_params = litellm.modify_params
|
|
litellm.modify_params = True
|
|
|
|
def teardown_method(self):
|
|
"""Cleanup after each test"""
|
|
# Restore original modify_params value
|
|
litellm.modify_params = self.original_modify_params
|
|
|
|
def test_case_a_orphaned_tool_call_single(self):
|
|
"""
|
|
Test Case A: Assistant message with tool_calls but no tool result
|
|
Should add a dummy tool result message
|
|
"""
|
|
messages = [
|
|
{"role": "user", "content": "What is the weather in Nashik?"},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "toolu_01Kus2cC3ydjBW7UK4GJqBP4",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"arguments": '{"location": "Nashik, India"}',
|
|
},
|
|
}
|
|
],
|
|
},
|
|
]
|
|
|
|
sanitized = sanitize_messages_for_tool_calling(messages)
|
|
|
|
# Should have 3 messages: user, assistant, and dummy tool result
|
|
assert len(sanitized) == 3
|
|
assert sanitized[0]["role"] == "user"
|
|
assert sanitized[1]["role"] == "assistant"
|
|
assert sanitized[2]["role"] == "tool"
|
|
assert sanitized[2]["tool_call_id"] == "toolu_01Kus2cC3ydjBW7UK4GJqBP4"
|
|
assert (
|
|
"skipped" in sanitized[2]["content"].lower()
|
|
or "interrupted" in sanitized[2]["content"].lower()
|
|
)
|
|
assert "get_weather" in sanitized[2]["content"]
|
|
|
|
def test_case_a_orphaned_tool_call_multiple(self):
|
|
"""
|
|
Test Case A: Assistant message with multiple tool_calls, some missing results
|
|
"""
|
|
messages = [
|
|
{"role": "user", "content": "Get weather for Nashik and Mumbai"},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"arguments": '{"location": "Nashik"}',
|
|
},
|
|
},
|
|
{
|
|
"id": "call_2",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"arguments": '{"location": "Mumbai"}',
|
|
},
|
|
},
|
|
],
|
|
},
|
|
{
|
|
"role": "tool",
|
|
"tool_call_id": "call_1",
|
|
"content": "Weather in Nashik: 25°C",
|
|
},
|
|
]
|
|
|
|
sanitized = sanitize_messages_for_tool_calling(messages)
|
|
|
|
# Should have 4 messages: user, assistant, tool result for call_1, dummy for call_2
|
|
assert len(sanitized) == 4
|
|
assert sanitized[0]["role"] == "user"
|
|
assert sanitized[1]["role"] == "assistant"
|
|
assert (
|
|
sanitized[2]["tool_call_id"] == "call_1"
|
|
) # Original tool result (first in tool_calls)
|
|
assert (
|
|
sanitized[3]["tool_call_id"] == "call_2"
|
|
) # Dummy added for missing call_2
|
|
|
|
def test_case_b_orphaned_tool_result(self):
|
|
"""
|
|
Test Case B: Tool result without matching tool_call in previous assistant message
|
|
Should remove the orphaned tool result
|
|
"""
|
|
messages = [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi there!"},
|
|
{
|
|
"role": "tool",
|
|
"tool_call_id": "nonexistent_id",
|
|
"content": "Some result",
|
|
},
|
|
]
|
|
|
|
sanitized = sanitize_messages_for_tool_calling(messages)
|
|
|
|
# Should have only 2 messages, orphaned tool result removed
|
|
assert len(sanitized) == 2
|
|
assert sanitized[0]["role"] == "user"
|
|
assert sanitized[1]["role"] == "assistant"
|
|
|
|
def test_case_b_valid_tool_result_preserved(self):
|
|
"""
|
|
Test Case B: Valid tool result with matching tool_call should be preserved
|
|
"""
|
|
messages = [
|
|
{"role": "user", "content": "What's the weather?"},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_123",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"arguments": '{"location": "Boston"}',
|
|
},
|
|
}
|
|
],
|
|
},
|
|
{"role": "tool", "tool_call_id": "call_123", "content": "Weather: 20°C"},
|
|
]
|
|
|
|
sanitized = sanitize_messages_for_tool_calling(messages)
|
|
|
|
# All messages should be preserved
|
|
assert len(sanitized) == 3
|
|
assert sanitized[2]["role"] == "tool"
|
|
assert sanitized[2]["tool_call_id"] == "call_123"
|
|
|
|
def test_case_c_empty_text_content_user(self):
|
|
"""
|
|
Test Case C: Empty text content in user message
|
|
Should replace with placeholder
|
|
"""
|
|
messages = [
|
|
{"role": "user", "content": ""},
|
|
{"role": "assistant", "content": "Hello!"},
|
|
]
|
|
|
|
sanitized = sanitize_messages_for_tool_calling(messages)
|
|
|
|
assert len(sanitized) == 2
|
|
assert sanitized[0]["role"] == "user"
|
|
assert (
|
|
sanitized[0]["content"]
|
|
== "[System: Empty message content sanitised to satisfy protocol]"
|
|
)
|
|
|
|
def test_case_c_whitespace_only_content(self):
|
|
"""
|
|
Test Case C: Whitespace-only content
|
|
Should replace with placeholder
|
|
"""
|
|
messages = [
|
|
{"role": "user", "content": " \n \t "},
|
|
{"role": "assistant", "content": " "},
|
|
]
|
|
|
|
sanitized = sanitize_messages_for_tool_calling(messages)
|
|
|
|
assert len(sanitized) == 2
|
|
assert (
|
|
sanitized[0]["content"]
|
|
== "[System: Empty message content sanitised to satisfy protocol]"
|
|
)
|
|
assert (
|
|
sanitized[1]["content"]
|
|
== "[System: Empty message content sanitised to satisfy protocol]"
|
|
)
|
|
|
|
def test_case_c_valid_content_preserved(self):
|
|
"""
|
|
Test Case C: Valid non-empty content should be preserved
|
|
"""
|
|
messages = [
|
|
{"role": "user", "content": "Hello"},
|
|
{"role": "assistant", "content": "Hi there!"},
|
|
]
|
|
|
|
sanitized = sanitize_messages_for_tool_calling(messages)
|
|
|
|
assert len(sanitized) == 2
|
|
assert sanitized[0]["content"] == "Hello"
|
|
assert sanitized[1]["content"] == "Hi there!"
|
|
|
|
def test_combined_cases(self):
|
|
"""
|
|
Test combination of multiple cases
|
|
"""
|
|
messages = [
|
|
{"role": "user", "content": "Get weather"},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"arguments": '{"location": "NYC"}',
|
|
},
|
|
}
|
|
],
|
|
},
|
|
# Missing tool result for call_1
|
|
{"role": "user", "content": ""}, # Empty content
|
|
{"role": "assistant", "content": "Response"},
|
|
{
|
|
"role": "tool",
|
|
"tool_call_id": "orphaned_id", # Orphaned tool result
|
|
"content": "Some data",
|
|
},
|
|
]
|
|
|
|
sanitized = sanitize_messages_for_tool_calling(messages)
|
|
|
|
# Should have: user, assistant, dummy tool result, user (sanitized), assistant
|
|
# Orphaned tool result should be removed
|
|
assert len(sanitized) == 5
|
|
assert sanitized[0]["role"] == "user"
|
|
assert sanitized[1]["role"] == "assistant"
|
|
assert sanitized[2]["role"] == "tool"
|
|
assert sanitized[2]["tool_call_id"] == "call_1" # Dummy added
|
|
assert sanitized[3]["role"] == "user"
|
|
assert (
|
|
sanitized[3]["content"]
|
|
== "[System: Empty message content sanitised to satisfy protocol]"
|
|
)
|
|
assert sanitized[4]["role"] == "assistant"
|
|
|
|
def test_modify_params_false_no_sanitization(self):
|
|
"""
|
|
Test that sanitization is skipped when modify_params=False
|
|
"""
|
|
litellm.modify_params = False
|
|
|
|
messages = [
|
|
{"role": "user", "content": ""},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "call_1",
|
|
"type": "function",
|
|
"function": {"name": "get_weather", "arguments": "{}"},
|
|
}
|
|
],
|
|
},
|
|
]
|
|
|
|
sanitized = sanitize_messages_for_tool_calling(messages)
|
|
|
|
# Messages should be unchanged
|
|
assert len(sanitized) == 2
|
|
assert sanitized[0]["content"] == ""
|
|
assert len(sanitized[1].get("tool_calls", [])) == 1
|
|
|
|
def test_anthropic_messages_pt_integration(self):
|
|
"""
|
|
Test that sanitization is integrated into anthropic_messages_pt
|
|
"""
|
|
litellm.modify_params = True
|
|
|
|
messages = [
|
|
{"role": "user", "content": "What is the weather in Nashik?"},
|
|
{
|
|
"role": "assistant",
|
|
"content": None,
|
|
"tool_calls": [
|
|
{
|
|
"id": "toolu_01Kus2cC3ydjBW7UK4GJqBP4",
|
|
"type": "function",
|
|
"function": {
|
|
"name": "get_weather",
|
|
"arguments": '{"location": "Nashik, India"}',
|
|
},
|
|
}
|
|
],
|
|
},
|
|
]
|
|
|
|
# This should not raise an error and should add dummy tool result
|
|
result = anthropic_messages_pt(
|
|
messages=messages, model="claude-sonnet-4-5", llm_provider="anthropic"
|
|
)
|
|
|
|
# Should have at least 2 messages (user and assistant)
|
|
# The tool result will be merged into user content
|
|
assert len(result) >= 2
|
|
assert result[0]["role"] == "user"
|
|
assert result[1]["role"] == "assistant"
|
|
|
|
|
|
if __name__ == "__main__":
|
|
pytest.main([__file__, "-v"])
|