mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-04 02:31:27 +00:00
Add compaction for vertex ai
This commit is contained in:
parent
1374b9dc02
commit
260cd5b49a
5 changed files with 152 additions and 242 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,
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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