mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
Add compaction for vertex ai
This commit is contained in:
parent
4f96a3b126
commit
8cd8a01d5a
5 changed files with 149 additions and 260 deletions
|
|
@ -2,6 +2,7 @@ from typing import TYPE_CHECKING, Any, List, Optional
|
|||
|
||||
import httpx
|
||||
|
||||
from litellm.anthropic_beta_headers_manager import filter_and_transform_beta_headers
|
||||
from litellm.llms.anthropic.chat.transformation import AnthropicConfig
|
||||
from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import (
|
||||
AmazonInvokeConfig,
|
||||
|
|
@ -133,27 +134,15 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
|
|||
beta_set.add("tool-search-tool-2025-10-19")
|
||||
|
||||
# Filter out beta headers that Bedrock Invoke doesn't support
|
||||
# AWS Bedrock only supports a specific whitelist of beta flags
|
||||
# Reference: https://docs.aws.amazon.com/bedrock/latest/userguide/model-parameters-anthropic-claude-messages-request-response.html
|
||||
BEDROCK_SUPPORTED_BETAS = {
|
||||
"computer-use-2024-10-22", # Legacy computer use
|
||||
"computer-use-2025-01-24", # Current computer use (Claude 3.7 Sonnet)
|
||||
"token-efficient-tools-2025-02-19", # Tool use (Claude 3.7+ and Claude 4+)
|
||||
"interleaved-thinking-2025-05-14", # Interleaved thinking (Claude 4+)
|
||||
"output-128k-2025-02-19", # 128K output tokens (Claude 3.7 Sonnet)
|
||||
"dev-full-thinking-2025-05-14", # Developer mode for raw thinking (Claude 4+)
|
||||
"context-1m-2025-08-07", # 1 million tokens (Claude Sonnet 4)
|
||||
"context-management-2025-06-27", # Context management (Claude Sonnet/Haiku 4.5)
|
||||
"effort-2025-11-24", # Effort parameter (Claude Opus 4.5)
|
||||
"tool-search-tool-2025-10-19", # Tool search (Claude Opus 4.5)
|
||||
"tool-examples-2025-10-29", # Tool use examples (Claude Opus 4.5)
|
||||
}
|
||||
|
||||
# Only keep beta headers that Bedrock supports
|
||||
beta_set = {beta for beta in beta_set if beta in BEDROCK_SUPPORTED_BETAS}
|
||||
# Uses centralized configuration from anthropic_beta_headers_config.json
|
||||
beta_list = list(beta_set)
|
||||
filtered_beta_list = filter_and_transform_beta_headers(
|
||||
beta_headers=beta_list,
|
||||
provider="bedrock",
|
||||
)
|
||||
|
||||
if beta_set:
|
||||
_anthropic_request["anthropic_beta"] = list(beta_set)
|
||||
if filtered_beta_list:
|
||||
_anthropic_request["anthropic_beta"] = filtered_beta_list
|
||||
|
||||
return _anthropic_request
|
||||
|
||||
|
|
|
|||
|
|
@ -68,6 +68,29 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
|
|||
if existing_beta:
|
||||
beta_values.update(b.strip() for b in existing_beta.split(","))
|
||||
|
||||
# Check for context management
|
||||
context_management_param = optional_params.get("context_management")
|
||||
if context_management_param is not None:
|
||||
# Check edits array for compact_20260112 type
|
||||
edits = context_management_param.get("edits", [])
|
||||
has_compact = False
|
||||
has_other = False
|
||||
|
||||
for edit in edits:
|
||||
edit_type = edit.get("type", "")
|
||||
if edit_type == "compact_20260112":
|
||||
has_compact = True
|
||||
else:
|
||||
has_other = True
|
||||
|
||||
# Add compact header if any compact edits exist
|
||||
if has_compact:
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value)
|
||||
|
||||
# Add context management header if any other edits exist
|
||||
if has_other:
|
||||
beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value)
|
||||
|
||||
# Check for web search tool
|
||||
for tool in tools:
|
||||
if isinstance(tool, dict) and tool.get("type", "").startswith(ANTHROPIC_HOSTED_TOOLS.WEB_SEARCH.value):
|
||||
|
|
|
|||
|
|
@ -993,66 +993,6 @@
|
|||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
},
|
||||
"anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 1.25e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 1e-06,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 1e-05,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.5e-05,
|
||||
"output_cost_per_token_above_200k_tokens": 3.75e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.01,
|
||||
"search_context_size_low": 0.01,
|
||||
"search_context_size_medium": 0.01
|
||||
},
|
||||
"supports_assistant_prefill": false,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
},
|
||||
"global.anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 1.25e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 1e-06,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 1e-05,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.5e-05,
|
||||
"output_cost_per_token_above_200k_tokens": 3.75e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.01,
|
||||
"search_context_size_low": 0.01,
|
||||
"search_context_size_medium": 0.01
|
||||
},
|
||||
"supports_assistant_prefill": false,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
},
|
||||
"global.anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 1.25e-05,
|
||||
|
|
@ -1143,66 +1083,6 @@
|
|||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
},
|
||||
"eu.anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 1.375e-05,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 1.1e-06,
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 1.1e-05,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.75e-05,
|
||||
"output_cost_per_token_above_200k_tokens": 4.125e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.01,
|
||||
"search_context_size_low": 0.01,
|
||||
"search_context_size_medium": 0.01
|
||||
},
|
||||
"supports_assistant_prefill": false,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
},
|
||||
"apac.anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 1.375e-05,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 1.1e-06,
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 1.1e-05,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.75e-05,
|
||||
"output_cost_per_token_above_200k_tokens": 4.125e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.01,
|
||||
"search_context_size_low": 0.01,
|
||||
"search_context_size_medium": 0.01
|
||||
},
|
||||
"supports_assistant_prefill": false,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
},
|
||||
"apac.anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 1.375e-05,
|
||||
|
|
|
|||
|
|
@ -993,66 +993,6 @@
|
|||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
},
|
||||
"anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 1.25e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 1e-06,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 1e-05,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.5e-05,
|
||||
"output_cost_per_token_above_200k_tokens": 3.75e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.01,
|
||||
"search_context_size_low": 0.01,
|
||||
"search_context_size_medium": 0.01
|
||||
},
|
||||
"supports_assistant_prefill": false,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
},
|
||||
"global.anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 1.25e-05,
|
||||
"cache_read_input_token_cost": 5e-07,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 1e-06,
|
||||
"input_cost_per_token": 5e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 1e-05,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.5e-05,
|
||||
"output_cost_per_token_above_200k_tokens": 3.75e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.01,
|
||||
"search_context_size_low": 0.01,
|
||||
"search_context_size_medium": 0.01
|
||||
},
|
||||
"supports_assistant_prefill": false,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
},
|
||||
"global.anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.25e-06,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 1.25e-05,
|
||||
|
|
@ -1143,66 +1083,6 @@
|
|||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
},
|
||||
"eu.anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 1.375e-05,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 1.1e-06,
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 1.1e-05,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.75e-05,
|
||||
"output_cost_per_token_above_200k_tokens": 4.125e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.01,
|
||||
"search_context_size_low": 0.01,
|
||||
"search_context_size_medium": 0.01
|
||||
},
|
||||
"supports_assistant_prefill": false,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
},
|
||||
"apac.anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 1.375e-05,
|
||||
"cache_read_input_token_cost": 5.5e-07,
|
||||
"cache_read_input_token_cost_above_200k_tokens": 1.1e-06,
|
||||
"input_cost_per_token": 5.5e-06,
|
||||
"input_cost_per_token_above_200k_tokens": 1.1e-05,
|
||||
"litellm_provider": "bedrock_converse",
|
||||
"max_input_tokens": 1000000,
|
||||
"max_output_tokens": 128000,
|
||||
"max_tokens": 128000,
|
||||
"mode": "chat",
|
||||
"output_cost_per_token": 2.75e-05,
|
||||
"output_cost_per_token_above_200k_tokens": 4.125e-05,
|
||||
"search_context_cost_per_query": {
|
||||
"search_context_size_high": 0.01,
|
||||
"search_context_size_low": 0.01,
|
||||
"search_context_size_medium": 0.01
|
||||
},
|
||||
"supports_assistant_prefill": false,
|
||||
"supports_computer_use": true,
|
||||
"supports_function_calling": true,
|
||||
"supports_pdf_input": true,
|
||||
"supports_prompt_caching": true,
|
||||
"supports_reasoning": true,
|
||||
"supports_response_schema": true,
|
||||
"supports_tool_choice": true,
|
||||
"supports_vision": true,
|
||||
"tool_use_system_prompt_tokens": 346
|
||||
},
|
||||
"apac.anthropic.claude-opus-4-6-v1": {
|
||||
"cache_creation_input_token_cost": 6.875e-06,
|
||||
"cache_creation_input_token_cost_above_200k_tokens": 1.375e-05,
|
||||
|
|
|
|||
|
|
@ -98,3 +98,120 @@ def test_web_search_header_not_added_without_tool():
|
|||
# Assert that the anthropic-beta header is NOT present when no web search tool
|
||||
assert "anthropic-beta" not in updated_headers, \
|
||||
"anthropic-beta header should not be present without web search tool"
|
||||
|
||||
|
||||
def test_compact_context_management_header_added():
|
||||
"""Test that compact-2026-01-12 beta header is added when context_management with compact_20260112 is used"""
|
||||
config = VertexAIPartnerModelsAnthropicMessagesConfig()
|
||||
headers = {}
|
||||
litellm_params = {
|
||||
"vertex_ai_project": "test-project",
|
||||
"vertex_ai_location": "us-central1",
|
||||
"vertex_credentials": "{}",
|
||||
}
|
||||
# Include context_management with compact_20260112
|
||||
optional_params = {
|
||||
"context_management": {
|
||||
"edits": [
|
||||
{"type": "compact_20260112"}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
config, "_ensure_access_token", return_value=("token", "test-project")
|
||||
), patch.object(
|
||||
config, "get_complete_vertex_url", return_value="https://mock-url"
|
||||
):
|
||||
updated_headers, api_base = config.validate_anthropic_messages_environment(
|
||||
headers=headers,
|
||||
model="claude-vertex-ai-opus-4-6",
|
||||
messages=[],
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
api_base=None,
|
||||
)
|
||||
|
||||
# Assert that the anthropic-beta header with compact-2026-01-12 is present
|
||||
assert "anthropic-beta" in updated_headers, "anthropic-beta header should be present"
|
||||
assert "compact-2026-01-12" in updated_headers["anthropic-beta"], \
|
||||
f"anthropic-beta should contain 'compact-2026-01-12', got: {updated_headers['anthropic-beta']}"
|
||||
|
||||
|
||||
def test_context_management_header_added_for_other_edits():
|
||||
"""Test that context-management-2025-06-27 beta header is added for non-compact edits"""
|
||||
config = VertexAIPartnerModelsAnthropicMessagesConfig()
|
||||
headers = {}
|
||||
litellm_params = {
|
||||
"vertex_ai_project": "test-project",
|
||||
"vertex_ai_location": "us-central1",
|
||||
"vertex_credentials": "{}",
|
||||
}
|
||||
# Include context_management with other edit types
|
||||
optional_params = {
|
||||
"context_management": {
|
||||
"edits": [
|
||||
{"type": "some_other_type"}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
config, "_ensure_access_token", return_value=("token", "test-project")
|
||||
), patch.object(
|
||||
config, "get_complete_vertex_url", return_value="https://mock-url"
|
||||
):
|
||||
updated_headers, api_base = config.validate_anthropic_messages_environment(
|
||||
headers=headers,
|
||||
model="claude-vertex-ai-opus-4-6",
|
||||
messages=[],
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
api_base=None,
|
||||
)
|
||||
|
||||
# Assert that the anthropic-beta header with context-management-2025-06-27 is present
|
||||
assert "anthropic-beta" in updated_headers, "anthropic-beta header should be present"
|
||||
assert "context-management-2025-06-27" in updated_headers["anthropic-beta"], \
|
||||
f"anthropic-beta should contain 'context-management-2025-06-27', got: {updated_headers['anthropic-beta']}"
|
||||
|
||||
|
||||
def test_both_compact_and_context_management_headers_added():
|
||||
"""Test that both compact and context-management beta headers are added when both edit types are present"""
|
||||
config = VertexAIPartnerModelsAnthropicMessagesConfig()
|
||||
headers = {}
|
||||
litellm_params = {
|
||||
"vertex_ai_project": "test-project",
|
||||
"vertex_ai_location": "us-central1",
|
||||
"vertex_credentials": "{}",
|
||||
}
|
||||
# Include context_management with both compact and other edit types
|
||||
optional_params = {
|
||||
"context_management": {
|
||||
"edits": [
|
||||
{"type": "compact_20260112"},
|
||||
{"type": "some_other_type"}
|
||||
]
|
||||
}
|
||||
}
|
||||
|
||||
with patch.object(
|
||||
config, "_ensure_access_token", return_value=("token", "test-project")
|
||||
), patch.object(
|
||||
config, "get_complete_vertex_url", return_value="https://mock-url"
|
||||
):
|
||||
updated_headers, api_base = config.validate_anthropic_messages_environment(
|
||||
headers=headers,
|
||||
model="claude-vertex-ai-opus-4-6",
|
||||
messages=[],
|
||||
optional_params=optional_params,
|
||||
litellm_params=litellm_params,
|
||||
api_base=None,
|
||||
)
|
||||
|
||||
# Assert that both beta headers are present
|
||||
assert "anthropic-beta" in updated_headers, "anthropic-beta header should be present"
|
||||
assert "compact-2026-01-12" in updated_headers["anthropic-beta"], \
|
||||
f"anthropic-beta should contain 'compact-2026-01-12', got: {updated_headers['anthropic-beta']}"
|
||||
assert "context-management-2025-06-27" in updated_headers["anthropic-beta"], \
|
||||
f"anthropic-beta should contain 'context-management-2025-06-27', got: {updated_headers['anthropic-beta']}"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue