Add compaction support for vertex ai

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(

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)

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"]
+

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