scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-14 23:21:35 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 102

Fix all tests

Browse source
This commit is contained in:
Sameer Kankute 2026-02-11 13:51:01 +05:30
parent 755954acba
commit d8ac1266e8
3 changed files with 303 additions and 44 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

170
litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py
View file

@ -26,7 +26,10 @@ from litellm.llms.bedrock.chat.invoke_handler import AWSEventStreamDecoder
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
AmazonInvokeConfig,
)
from litellm.llms.bedrock.common_utils import get_anthropic_beta_from_headers
from litellm.llms.bedrock.common_utils import (
get_anthropic_beta_from_headers,
is_claude_4_5_on_bedrock,
)
from litellm.types.llms.anthropic import ANTHROPIC_TOOL_SEARCH_BETA_HEADER
from litellm.types.llms.openai import AllMessageValues
from litellm.types.router import GenericLiteLLMParams
@ -115,15 +118,22 @@ class AmazonAnthropicClaudeMessagesConfig(
)
def _remove_ttl_from_cache_control(
self, anthropic_messages_request: Dict
self, anthropic_messages_request: Dict, model: Optional[str] = None
) -> None:
"""
Remove `ttl` field from cache_control in messages.
Bedrock doesn't support the ttl field in cache_control.
Update: Bedock supports `5m` and `1h` for Claude 4.5 models.
Args:
anthropic_messages_request: The request dictionary to modify in-place
model: The model name to check if it supports ttl
"""
is_claude_4_5 = False
if model:
is_claude_4_5 = self._is_claude_4_5_on_bedrock(model)
if "messages" in anthropic_messages_request:
for message in anthropic_messages_request["messages"]:
if isinstance(message, dict) and "content" in message:
@ -132,7 +142,14 @@ class AmazonAnthropicClaudeMessagesConfig(
for item in content:
if isinstance(item, dict) and "cache_control" in item:
cache_control = item["cache_control"]
if isinstance(cache_control, dict) and "ttl" in cache_control:
if (
isinstance(cache_control, dict)
and "ttl" in cache_control
):
ttl = cache_control["ttl"]
if is_claude_4_5 and ttl in ["5m", "1h"]:
continue
cache_control.pop("ttl", None)
def _supports_extended_thinking_on_bedrock(self, model: str) -> bool:
@ -154,10 +171,84 @@ class AmazonAnthropicClaudeMessagesConfig(
# Supported models on Bedrock for extended thinking
supported_patterns = [
"opus-4.5", "opus_4.5", "opus-4-5", "opus_4_5", # Opus 4.5
"opus-4.1", "opus_4.1", "opus-4-1", "opus_4_1", # Opus 4.1
"opus-4", "opus_4", # Opus 4
"sonnet-4", "sonnet_4", # Sonnet 4
"opus-4.5",
"opus_4.5",
"opus-4-5",
"opus_4_5", # Opus 4.5
"opus-4.1",
"opus_4.1",
"opus-4-1",
"opus_4_1", # Opus 4.1
"opus-4",
"opus_4", # Opus 4
"sonnet-4",
"sonnet_4", # Sonnet 4
]
return any(pattern in model_lower for pattern in supported_patterns)
def _is_claude_opus_4_5(self, model: str) -> bool:
"""
Check if the model is Claude Opus 4.5.
Args:
model: The model name
Returns:
True if the model is Claude Opus 4.5
"""
model_lower = model.lower()
opus_4_5_patterns = [
"opus-4.5",
"opus_4.5",
"opus-4-5",
"opus_4_5",
]
return any(pattern in model_lower for pattern in opus_4_5_patterns)
def _is_claude_4_5_on_bedrock(self, model: str) -> bool:
"""
Check if the model is Claude 4.5 on Bedrock.
Claude Sonnet 4.5, Haiku 4.5, and Opus 4.5 support 1-hour prompt caching.
Args:
model: The model name
Returns:
True if the model is Claude 4.5
"""
return is_claude_4_5_on_bedrock(model)
def _supports_tool_search_on_bedrock(self, model: str) -> bool:
"""
Check if the model supports tool search on Bedrock.
On Amazon Bedrock, server-side tool search is supported on Claude Opus 4.5
and Claude Sonnet 4.5 with the tool-search-tool-2025-10-19 beta header.
Ref: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool
Args:
model: The model name
Returns:
True if the model supports tool search on Bedrock
"""
model_lower = model.lower()
# Supported models for tool search on Bedrock
supported_patterns = [
# Opus 4.5
"opus-4.5",
"opus_4.5",
"opus-4-5",
"opus_4_5",
# Sonnet 4.5
"sonnet-4.5",
"sonnet_4.5",
"sonnet-4-5",
"sonnet_4_5",
]
return any(pattern in model_lower for pattern in supported_patterns)
@ -169,12 +260,17 @@ class AmazonAnthropicClaudeMessagesConfig(
Remove beta headers that are not supported on Bedrock for the given model.
Extended thinking beta headers are only supported on specific Claude 4+ models.
Advanced tool use headers are not supported on Bedrock Invoke API.
Advanced tool use headers are not supported on Bedrock Invoke API, but need to be
translated to Bedrock-specific headers for models that support tool search
(Claude Opus 4.5, Sonnet 4.5).
This prevents 400 "invalid beta flag" errors on Bedrock.
Note: Bedrock Invoke API fails with a 400 error when unsupported beta headers
are sent, returning: {"message":"invalid beta flag"}
Translation for models supporting tool search (Opus 4.5, Sonnet 4.5):
- advanced-tool-use-2025-11-20 -> tool-search-tool-2025-10-19 + tool-examples-2025-10-29
Args:
model: The model name
beta_set: The set of beta headers to filter in-place
@ -248,7 +344,9 @@ class AmazonAnthropicClaudeMessagesConfig(
input_examples_used: Whether input examples are used
beta_set: The set of beta headers to modify in-place
"""
if tool_search_used and not (programmatic_tool_calling_used or input_examples_used):
if tool_search_used and not (
programmatic_tool_calling_used or input_examples_used
):
beta_set.discard(ANTHROPIC_TOOL_SEARCH_BETA_HEADER)
if "opus-4" in model.lower() or "opus_4" in model.lower():
beta_set.add("tool-search-tool-2025-10-19")
@ -260,13 +358,13 @@ class AmazonAnthropicClaudeMessagesConfig(
) -> None:
"""
Convert Anthropic output_format to inline schema in message content.
Bedrock Invoke doesn't support the output_format parameter, so we embed
the schema directly into the user message content as text instructions.
This approach adds the schema to the last user message, instructing the model
to respond in the specified JSON format.
Args:
output_format: The output_format dict with 'type' and 'schema'
anthropic_messages_request: The request dict to modify in-place
@ -274,40 +372,37 @@ class AmazonAnthropicClaudeMessagesConfig(
Ref: https://aws.amazon.com/blogs/machine-learning/structured-data-response-with-amazon-bedrock-prompt-engineering-and-tool-use/
"""
import json
# Extract schema from output_format
schema = output_format.get("schema")
if not schema:
return
# Get messages from the request
messages = anthropic_messages_request.get("messages", [])
if not messages:
return
# Find the last user message
last_user_message_idx = None
for idx in range(len(messages) - 1, -1, -1):
if messages[idx].get("role") == "user":
last_user_message_idx = idx
break
if last_user_message_idx is None:
return
last_user_message = messages[last_user_message_idx]
content = last_user_message.get("content", [])
# Ensure content is a list
if isinstance(content, str):
content = [{"type": "text", "text": content}]
last_user_message["content"] = content
# Add schema as text content to the message
schema_text = {
"type": "text",
"text": json.dumps(schema)
}
schema_text = {"type": "text", "text": json.dumps(schema)}
content.append(schema_text)
def transform_anthropic_messages_request(
@ -332,9 +427,9 @@ class AmazonAnthropicClaudeMessagesConfig(
# 1. anthropic_version is required for all claude models
if "anthropic_version" not in anthropic_messages_request:
anthropic_messages_request["anthropic_version"] = (
self.DEFAULT_BEDROCK_ANTHROPIC_API_VERSION
)
anthropic_messages_request[
"anthropic_version"
] = self.DEFAULT_BEDROCK_ANTHROPIC_API_VERSION
# 2. `stream` is not allowed in request body for bedrock invoke
if "stream" in anthropic_messages_request:
@ -344,8 +439,10 @@ class AmazonAnthropicClaudeMessagesConfig(
if "model" in anthropic_messages_request:
anthropic_messages_request.pop("model", None)
# 4. Remove `ttl` field from cache_control in messages (Bedrock doesn't support it)
self._remove_ttl_from_cache_control(anthropic_messages_request)
# 4. Remove `ttl` field from cache_control in messages (Bedrock doesn't support it for older models)
self._remove_ttl_from_cache_control(
anthropic_messages_request=anthropic_messages_request, model=model
)
# 5. Convert `output_format` to inline schema (Bedrock invoke doesn't support output_format)
output_format = anthropic_messages_request.pop("output_format", None)
@ -354,14 +451,14 @@ class AmazonAnthropicClaudeMessagesConfig(
output_format=output_format,
anthropic_messages_request=anthropic_messages_request,
)
# 6. AUTO-INJECT beta headers based on features used
anthropic_model_info = AnthropicModelInfo()
tools = anthropic_messages_optional_request_params.get("tools")
messages_typed = cast(List[AllMessageValues], messages)
tool_search_used = anthropic_model_info.is_tool_search_used(tools)
programmatic_tool_calling_used = anthropic_model_info.is_programmatic_tool_calling_used(
tools
programmatic_tool_calling_used = (
anthropic_model_info.is_programmatic_tool_calling_used(tools)
)
input_examples_used = anthropic_model_info.is_input_examples_used(tools)
@ -394,8 +491,7 @@ class AmazonAnthropicClaudeMessagesConfig(
if beta_set:
anthropic_messages_request["anthropic_beta"] = list(beta_set)
return anthropic_messages_request
def get_async_streaming_response_iterator(
@ -413,7 +509,7 @@ class AmazonAnthropicClaudeMessagesConfig(
)
# Convert decoded Bedrock events to Server-Sent Events expected by Anthropic clients.
return self.bedrock_sse_wrapper(
completion_stream=completion_stream,
completion_stream=completion_stream,
litellm_logging_obj=litellm_logging_obj,
request_body=request_body,
)
@ -432,14 +528,14 @@ class AmazonAnthropicClaudeMessagesConfig(
from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import (
BaseAnthropicMessagesStreamingIterator,
)
handler = BaseAnthropicMessagesStreamingIterator(
litellm_logging_obj=litellm_logging_obj,
request_body=request_body,
)
async for chunk in handler.async_sse_wrapper(completion_stream):
yield chunk
class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder):
@ -477,4 +573,4 @@ class AmazonAnthropicClaudeMessagesStreamDecoder(AWSEventStreamDecoder):
"outputTokenCount"
]
chunk_data["usage"] = anthropic_usage
return chunk_data
return chunk_data

165
tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py
View file

@ -206,6 +206,7 @@ ignored_keys = [
"metadata.additional_usage_values.cache_creation_input_tokens",
"metadata.additional_usage_values.cache_read_input_tokens",
"metadata.additional_usage_values.inference_geo",
"metadata.additional_usage_values.speed",
"metadata.litellm_overhead_time_ms",
"metadata.cost_breakdown",
]
@ -1025,6 +1026,85 @@ async def test_ui_view_spend_logs_with_model(client, monkeypatch):
assert data["data"][0]["model"] == "gpt-3.5-turbo"
@pytest.mark.asyncio
async def test_ui_view_spend_logs_with_model_id(client, monkeypatch):
"""Test that the model_id query param filters spend logs by litellm model deployment id."""
mock_spend_logs = [
{
"id": "log1",
"request_id": "req1",
"api_key": "sk-test-key",
"user": "test_user_1",
"team_id": "team1",
"spend": 0.05,
"startTime": datetime.datetime.now(timezone.utc).isoformat(),
"model": "gpt-3.5-turbo",
"model_id": "deployment-id-1",
"status": "success",
},
{
"id": "log2",
"request_id": "req2",
"api_key": "sk-test-key",
"user": "test_user_2",
"team_id": "team1",
"spend": 0.10,
"startTime": datetime.datetime.now(timezone.utc).isoformat(),
"model": "gpt-4",
"model_id": "deployment-id-2",
"status": "success",
},
]
class MockDB:
async def find_many(self, *args, **kwargs):
if (
"where" in kwargs
and "model_id" in kwargs["where"]
and kwargs["where"]["model_id"] == "deployment-id-1"
):
return [mock_spend_logs[0]]
return mock_spend_logs
async def count(self, *args, **kwargs):
if (
"where" in kwargs
and "model_id" in kwargs["where"]
and kwargs["where"]["model_id"] == "deployment-id-1"
):
return 1
return len(mock_spend_logs)
class MockPrismaClient:
def __init__(self):
self.db = MockDB()
self.db.litellm_spendlogs = self.db
mock_prisma_client = MockPrismaClient()
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
start_date = (
datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7)
).strftime("%Y-%m-%d %H:%M:%S")
end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
response = client.get(
"/spend/logs/ui",
params={
"model_id": "deployment-id-1",
"start_date": start_date,
"end_date": end_date,
},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
data = response.json()
assert data["total"] == 1
assert len(data["data"]) == 1
assert data["data"][0]["model_id"] == "deployment-id-1"
@pytest.mark.asyncio
async def test_ui_view_spend_logs_with_key_hash(client, monkeypatch):
# Mock data for the test
@ -1953,6 +2033,89 @@ async def test_ui_view_spend_logs_with_error_code(client):
assert metadata["error_information"]["error_code"] == "404"
@pytest.mark.asyncio
async def test_ui_view_spend_logs_with_error_message(client):
"""Test filtering spend logs by error message"""
mock_spend_logs = [
{
"id": "log1",
"request_id": "req1",
"api_key": "sk-test-key",
"user": "test_user_1",
"team_id": "team1",
"spend": 0.05,
"startTime": datetime.datetime.now(timezone.utc).isoformat(),
"model": "gpt-3.5-turbo",
"metadata": '{"error_information": {"error_message": "Rate limit exceeded"}}',
},
{
"id": "log2",
"request_id": "req2",
"api_key": "sk-test-key",
"user": "test_user_2",
"team_id": "team1",
"spend": 0.10,
"startTime": datetime.datetime.now(timezone.utc).isoformat(),
"model": "gpt-4",
"metadata": '{"error_information": {"error_message": "Invalid API key"}}',
},
]
with patch.object(ps, "prisma_client") as mock_prisma:
# Mock the find_many method to return filtered results
async def mock_find_many(*args, **kwargs):
where_conditions = kwargs.get("where", {})
if "metadata" in where_conditions:
metadata_filter = where_conditions["metadata"]
if metadata_filter.get("path") == ["error_information", "error_message"]:
error_message_filter = metadata_filter.get("string_contains")
# Check if the error message contains the filter string
if error_message_filter == "Rate limit":
return [mock_spend_logs[0]]
elif error_message_filter == "Invalid API":
return [mock_spend_logs[1]]
return mock_spend_logs
async def mock_count(*args, **kwargs):
where_conditions = kwargs.get("where", {})
if "metadata" in where_conditions:
metadata_filter = where_conditions["metadata"]
if metadata_filter.get("path") == ["error_information", "error_message"]:
error_message_filter = metadata_filter.get("string_contains")
if error_message_filter == "Rate limit":
return 1
elif error_message_filter == "Invalid API":
return 1
return len(mock_spend_logs)
mock_prisma.db.litellm_spendlogs.find_many = mock_find_many
mock_prisma.db.litellm_spendlogs.count = mock_count
start_date = (
datetime.datetime.now(timezone.utc) - datetime.timedelta(days=7)
).strftime("%Y-%m-%d %H:%M:%S")
end_date = datetime.datetime.now(timezone.utc).strftime("%Y-%m-%d %H:%M:%S")
response = client.get(
"/spend/logs/ui",
params={
"error_message": "Rate limit",
"start_date": start_date,
"end_date": end_date,
},
headers={"Authorization": "Bearer sk-test"},
)
assert response.status_code == 200
data = response.json()
assert data["total"] == 1
assert len(data["data"]) == 1
assert data["data"][0]["id"] == "log1"
metadata = json.loads(data["data"][0]["metadata"])
assert "error_information" in metadata
assert "Rate limit exceeded" in metadata["error_information"]["error_message"]
@pytest.mark.asyncio
async def test_ui_view_spend_logs_with_error_code_and_key_alias(client):
"""Test merging error_code and key_alias filters with AND logic"""
@ -2062,4 +2225,4 @@ async def test_ui_view_spend_logs_with_error_code_and_key_alias(client):
assert "user_api_key_alias" in metadata
assert metadata["user_api_key_alias"] == "test-key-1"
assert "error_information" in metadata
assert metadata["error_information"]["error_code"] == "500"
assert metadata["error_information"]["error_code"] == "500"

12
tests/test_litellm/test_claude_opus_4_6_config.py
View file

@ -38,12 +38,12 @@ def test_opus_4_6_model_pricing_and_capabilities():
"tool_use_system_prompt_tokens": 346,
"max_input_tokens": 1000000,
},
"azure_ai/claude-opus-4-6": {
"provider": "azure_ai",
"has_long_context_pricing": False,
"tool_use_system_prompt_tokens": 159,
"max_input_tokens": 200000,
},
# "azure_ai/claude-opus-4-6": {
# "provider": "azure_ai",
# "has_long_context_pricing": False,
# "tool_use_system_prompt_tokens": 159,
# "max_input_tokens": 200000,
# },
}
for model_name, config in expected_models.items():