Add compaction for vertex ai

This commit is contained in:
Sameer Kankute 2026-02-09 10:18:19 +05:30
parent 1374b9dc02
commit 260cd5b49a
5 changed files with 152 additions and 242 deletions

View file

@ -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,
@ -117,8 +118,16 @@ class AmazonAnthropicClaudeConfig(AmazonInvokeConfig, AnthropicConfig):
if "opus-4" in model.lower() or "opus_4" in model.lower():
beta_set.add("tool-search-tool-2025-10-19")
if beta_set:
_anthropic_request["anthropic_beta"] = list(beta_set)
# Filter out beta headers that Bedrock Invoke doesn't support
# 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 filtered_beta_list:
_anthropic_request["anthropic_beta"] = filtered_beta_list
return _anthropic_request

View file

@ -67,6 +67,30 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
existing_beta = headers.get("anthropic-beta")
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):

View file

@ -902,66 +902,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,
@ -1082,66 +1022,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,

View file

@ -902,66 +902,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,
@ -1082,66 +1022,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,

View file

@ -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']}"