From 0c8a484394634895b3bf6dd9ab4b6572d8ed179f Mon Sep 17 00:00:00 2001 From: Sameer Kankute Date: Fri, 6 Feb 2026 12:51:55 +0530 Subject: [PATCH] Add compaction support for vertex ai --- litellm/llms/anthropic/chat/transformation.py | 3 + .../anthropic/transformation.py | 39 ++++++++ .../test_azure_anthropic_transformation.py | 94 +++++++++++++++++++ ...partner_models_anthropic_transformation.py | 70 ++++++++++++++ 4 files changed, 206 insertions(+) diff --git a/litellm/llms/anthropic/chat/transformation.py b/litellm/llms/anthropic/chat/transformation.py index 7c4901fc5f2..f4e08d2218c 100644 --- a/litellm/llms/anthropic/chat/transformation.py +++ b/litellm/llms/anthropic/chat/transformation.py @@ -865,6 +865,9 @@ class AnthropicConfig(AnthropicModelInfo, BaseConfig): ) elif param == "extra_headers": optional_params["extra_headers"] = value + elif param == "context_management" and isinstance(value, dict): + # Pass through Anthropic-specific context_management parameter + optional_params["context_management"] = value ## handle thinking tokens self.update_optional_params_with_thinking_tokens( diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py index 1df07f405e6..0b728d88e76 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/transformation.py @@ -51,6 +51,40 @@ class VertexAIAnthropicConfig(AnthropicConfig): def custom_llm_provider(self) -> Optional[str]: return "vertex_ai" + def _add_context_management_beta_headers( + self, beta_set: set, context_management: dict + ) -> None: + """ + Add context_management beta headers to the beta_set. + + - If any edit has type "compact_20260112", add compact-2026-01-12 header + - For all other edits, add context-management-2025-06-27 header + + Args: + beta_set: Set of beta headers to modify in-place + context_management: The context_management dict from optional_params + """ + from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES + + edits = context_management.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_set.add(ANTHROPIC_BETA_HEADER_VALUES.COMPACT_2026_01_12.value) + + # Add context management header if any other edits exist + if has_other: + beta_set.add(ANTHROPIC_BETA_HEADER_VALUES.CONTEXT_MANAGEMENT_2025_06_27.value) + def transform_request( self, model: str, @@ -86,6 +120,11 @@ class VertexAIAnthropicConfig(AnthropicConfig): beta_set = set(auto_betas) if tool_search_used: beta_set.add("tool-search-tool-2025-10-19") # Vertex requires this header for tool search + + # Add context_management beta headers (compact and/or context-management) + context_management = optional_params.get("context_management") + if context_management: + self._add_context_management_beta_headers(beta_set, context_management) if beta_set: data["anthropic_beta"] = list(beta_set) diff --git a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py index e43a899325f..f0f8a9d91bf 100644 --- a/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py +++ b/tests/test_litellm/llms/azure_ai/claude/test_azure_anthropic_transformation.py @@ -235,3 +235,97 @@ class TestAzureAnthropicConfig: assert result["max_tokens"] == 100 assert "messages" in result + def test_context_management_compact_beta_header(self): + """Test that context_management with compact adds the correct beta header for Azure AI""" + config = AzureAnthropicConfig() + + messages = [{"role": "user", "content": "Hello"}] + optional_params = { + "context_management": { + "edits": [ + { + "type": "compact_20260112" + } + ] + }, + "max_tokens": 100 + } + litellm_params = {"api_key": "test-key"} + headers = {"api-key": "test-key"} + + with patch( + "litellm.llms.azure.common_utils.BaseAzureLLM._base_validate_azure_environment" + ) as mock_validate: + mock_validate.return_value = {"api-key": "test-key"} + result = config.transform_request( + model="claude-opus-4-6", + messages=messages, + optional_params=optional_params, + litellm_params=litellm_params, + headers=headers, + ) + + # Verify context_management is included + assert "context_management" in result + assert result["context_management"]["edits"][0]["type"] == "compact_20260112" + + def test_context_management_compact_beta_header_in_headers(self): + """Test that compact beta header is added to headers for Azure AI""" + config = AzureAnthropicConfig() + + messages = [{"role": "user", "content": "Hello"}] + optional_params = { + "context_management": { + "edits": [ + { + "type": "compact_20260112" + } + ] + }, + "max_tokens": 100 + } + + # Test that the parent's update_headers_with_optional_anthropic_beta is called + # which should add the compact beta header + headers = {} + headers = config.update_headers_with_optional_anthropic_beta( + headers=headers, + optional_params=optional_params + ) + + # Verify compact beta header is present + assert "anthropic-beta" in headers + assert "compact-2026-01-12" in headers["anthropic-beta"] + + def test_context_management_mixed_edits_beta_headers(self): + """Test that context_management with both compact and other edits adds both beta headers""" + config = AzureAnthropicConfig() + + messages = [{"role": "user", "content": "Hello"}] + optional_params = { + "context_management": { + "edits": [ + { + "type": "compact_20260112" + }, + { + "type": "replace", + "message_id": "msg_123", + "content": "new content" + } + ] + }, + "max_tokens": 100 + } + + headers = {} + headers = config.update_headers_with_optional_anthropic_beta( + headers=headers, + optional_params=optional_params + ) + + # Verify both beta headers are present + assert "anthropic-beta" in headers + assert "compact-2026-01-12" in headers["anthropic-beta"] + assert "context-management-2025-06-27" in headers["anthropic-beta"] + diff --git a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py index 7b60a0a3369..9b29e796c46 100644 --- a/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_transformation.py @@ -74,6 +74,76 @@ def test_vertex_ai_anthropic_web_search_header_in_completion(): "anthropic-beta with web-search should not be present for non-Vertex requests" +def test_vertex_ai_anthropic_context_management_compact_beta_header(): + """Test that context_management with compact adds the correct beta header for Vertex AI""" + config = VertexAIAnthropicConfig() + + messages = [{"role": "user", "content": "Hello"}] + optional_params = { + "context_management": { + "edits": [ + { + "type": "compact_20260112" + } + ] + }, + "max_tokens": 100, + "is_vertex_request": True + } + + result = config.transform_request( + model="claude-opus-4-6", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={} + ) + + # Verify context_management is included + assert "context_management" in result + assert result["context_management"]["edits"][0]["type"] == "compact_20260112" + + # Verify compact beta header is in anthropic_beta field + assert "anthropic_beta" in result + assert "compact-2026-01-12" in result["anthropic_beta"] + + +def test_vertex_ai_anthropic_context_management_mixed_edits(): + """Test that context_management with both compact and other edits adds both beta headers""" + config = VertexAIAnthropicConfig() + + messages = [{"role": "user", "content": "Hello"}] + optional_params = { + "context_management": { + "edits": [ + { + "type": "compact_20260112" + }, + { + "type": "replace", + "message_id": "msg_123", + "content": "new content" + } + ] + }, + "max_tokens": 100, + "is_vertex_request": True + } + + result = config.transform_request( + model="claude-opus-4-6", + messages=messages, + optional_params=optional_params, + litellm_params={}, + headers={} + ) + + # Verify both beta headers are present + assert "anthropic_beta" in result + assert "compact-2026-01-12" in result["anthropic_beta"] + assert "context-management-2025-06-27" in result["anthropic_beta"] + + def test_vertex_ai_anthropic_structured_output_header_not_added(): """Test that structured output beta headers are NOT added for Vertex AI requests""" from litellm.llms.anthropic.chat.transformation import AnthropicConfig