From dc5650eedaafe2e215cb9b9ddf67fc937399cb57 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 10 Sep 2025 15:47:18 -0700 Subject: [PATCH] Revert "fix: remove anthropic-beta header for Vertex AI requests with prompt caching" (#14421) --- litellm/llms/anthropic/common_utils.py | 11 +- .../test_vertex_ai_prompt_caching_fix.py | 134 ------------------ .../llms/vertex_ai/test_vertex.py | 2 +- 3 files changed, 3 insertions(+), 144 deletions(-) delete mode 100644 tests/litellm/llms/anthropic/test_vertex_ai_prompt_caching_fix.py diff --git a/litellm/llms/anthropic/common_utils.py b/litellm/llms/anthropic/common_utils.py index 06ebb5079d9..68b5341e954 100644 --- a/litellm/llms/anthropic/common_utils.py +++ b/litellm/llms/anthropic/common_utils.py @@ -107,10 +107,8 @@ class AnthropicModelInfo(BaseLLMModelInfo): user_anthropic_beta_headers: Optional[List[str]] = None, ) -> dict: betas = set() - # Note: prompt-caching-2024-07-31 header is no longer required for prompt caching - # as per current Anthropic documentation. It's now generally available. - # if prompt_caching_set: - # betas.add("prompt-caching-2024-07-31") + if prompt_caching_set: + betas.add("prompt-caching-2024-07-31") if computer_tool_used: betas.add("computer-use-2024-10-22") # if pdf_used: @@ -178,11 +176,6 @@ class AnthropicModelInfo(BaseLLMModelInfo): mcp_server_used=mcp_server_used, ) - # For Vertex AI requests, remove any user-provided anthropic-beta headers - # since Vertex AI rejects them and they're no longer required for prompt caching - if optional_params.get("is_vertex_request", False): - headers = {k: v for k, v in headers.items() if k != "anthropic-beta"} - headers = {**headers, **anthropic_headers} return headers diff --git a/tests/litellm/llms/anthropic/test_vertex_ai_prompt_caching_fix.py b/tests/litellm/llms/anthropic/test_vertex_ai_prompt_caching_fix.py deleted file mode 100644 index 67c37bc6f26..00000000000 --- a/tests/litellm/llms/anthropic/test_vertex_ai_prompt_caching_fix.py +++ /dev/null @@ -1,134 +0,0 @@ -""" -Test file for Vertex AI prompt caching fix. - -This test verifies that: -1. The anthropic-beta header is removed for Vertex AI requests -2. Regular Anthropic requests still work correctly -3. Prompt caching detection logic is preserved -""" - -import unittest -from unittest.mock import patch, MagicMock -from litellm.llms.anthropic.common_utils import AnthropicModelInfo - - -class TestVertexAIPromptCachingFix(unittest.TestCase): - """Test cases for the Vertex AI prompt caching fix.""" - - def setUp(self): - """Set up test fixtures.""" - self.model_info = AnthropicModelInfo() - - def _is_cache_control_set(self, messages): - """Helper method to test cache control detection.""" - for message in messages: - if message.get("cache_control", None) is not None: - return True - _message_content = message.get("content") - if _message_content is not None and isinstance(_message_content, list): - for content in _message_content: - if "cache_control" in content: - return True - return False - - def test_vertex_ai_removes_anthropic_beta_header(self): - """Test that anthropic-beta header is removed for Vertex AI requests.""" - # Mock the get_anthropic_headers method - with patch.object(self.model_info, 'get_anthropic_headers') as mock_get_headers: - # Set up the mock to return headers with anthropic-beta - mock_get_headers.return_value = { - 'anthropic-beta': 'prompt-caching-2024-07-31', - 'content-type': 'application/json' - } - - # Test with Vertex AI request - optional_params = {'is_vertex_request': True} - headers = self.model_info.get_anthropic_headers( - model='vertex/claude-3-5-sonnet-20240620', - messages=[], - optional_params=optional_params - ) - - # Verify anthropic-beta header is removed - self.assertNotIn('anthropic-beta', headers) - - def test_regular_anthropic_preserves_anthropic_beta_header(self): - """Test that regular Anthropic requests preserve anthropic-beta header.""" - # Mock the get_anthropic_headers method - with patch.object(self.model_info, 'get_anthropic_headers') as mock_get_headers: - # Set up the mock to return headers with anthropic-beta - mock_get_headers.return_value = { - 'anthropic-beta': 'prompt-caching-2024-07-31', - 'content-type': 'application/json' - } - - # Test with regular Anthropic request (not Vertex AI) - optional_params = {'is_vertex_request': False} - headers = self.model_info.get_anthropic_headers( - model='claude-3-5-sonnet-20240620', - messages=[], - optional_params=optional_params - ) - - # Verify anthropic-beta header is preserved - self.assertIn('anthropic-beta', headers) - - def test_prompt_caching_detection_still_works(self): - """Test that prompt caching detection logic still works.""" - # Test messages with cache_control - messages_with_cache = [ - { - 'role': 'user', - 'content': [ - { - 'type': 'text', - 'text': 'Test prompt', - 'cache_control': {'type': 'ephemeral'} - } - ] - } - ] - - # Test messages without cache_control - messages_without_cache = [ - { - 'role': 'user', - 'content': [ - { - 'type': 'text', - 'text': 'Regular prompt' - } - ] - } - ] - - # Test cache detection - cache_detected_with = self._is_cache_control_set(messages_with_cache) - cache_detected_without = self._is_cache_control_set(messages_without_cache) - - self.assertTrue(cache_detected_with) - self.assertFalse(cache_detected_without) - - def test_vertex_ai_without_user_beta_header(self): - """Test Vertex AI request when no user-provided beta header exists.""" - # Mock the get_anthropic_headers method - with patch.object(self.model_info, 'get_anthropic_headers') as mock_get_headers: - # Set up the mock to return headers without anthropic-beta - mock_get_headers.return_value = { - 'content-type': 'application/json' - } - - # Test with Vertex AI request - optional_params = {'is_vertex_request': True} - headers = self.model_info.get_anthropic_headers( - model='vertex/claude-3-5-sonnet-20240620', - messages=[], - optional_params=optional_params - ) - - # Verify no anthropic-beta header is present - self.assertNotIn('anthropic-beta', headers) - - -if __name__ == '__main__': - unittest.main() \ No newline at end of file diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex.py b/tests/test_litellm/llms/vertex_ai/test_vertex.py index 02bd622016b..7e683d1f54e 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex.py @@ -10,7 +10,7 @@ import litellm.litellm_core_utils.prompt_templates import litellm.litellm_core_utils.prompt_templates.factory load_dotenv() -from unittest.mock import MagicMock, patch +from unittest.mock import MagicMock sys.path.insert( 0, os.path.abspath("../..")