Merge 147d44c582 into 452254963e

litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py

@@ -24,6 +24,11 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
     def should_strip_billing_metadata(self) -> bool:
         return True
 
+    @staticmethod
+    def _has_vertex_predict_endpoint(api_base: str) -> bool:
+        api_base_without_query = api_base.split("?", 1)[0]
+        return api_base_without_query.endswith((":rawPredict", ":streamRawPredict"))
+
     def validate_anthropic_messages_environment(
         self,
         headers: dict,
@@ -53,15 +58,18 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
         )
         headers["Authorization"] = f"Bearer {access_token}"
 
-        api_base = self.get_complete_vertex_url(
-            custom_api_base=api_base,
-            vertex_location=vertex_ai_location,
-            vertex_project=vertex_ai_project,
-            project_id=project_id or "",
-            partner=VertexPartnerProvider.claude,
-            stream=optional_params.get("stream", False),
-            model=model,
-        )
+        if api_base is None or not self._has_vertex_predict_endpoint(api_base):
+            # Normalize Vertex model URLs, but preserve fully-qualified endpoints
+            # provided by callers (for example ...:rawPredict).
+            api_base = self.get_complete_vertex_url(
+                custom_api_base=api_base,
+                vertex_location=vertex_ai_location,
+                vertex_project=vertex_ai_project,
+                project_id=project_id or "",
+                partner=VertexPartnerProvider.claude,
+                stream=optional_params.get("stream", False),
+                model=model,
+            )
 
         headers["content-type"] = "application/json"
 

tests/test_litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py

@@ -23,12 +23,8 @@ def test_validate_environment_uses_vertex_ai_location():
     optional_params = {}
 
     with (
-        patch.object(
-            config, "_ensure_access_token", return_value=("token", "test-project")
-        ),
-        patch.object(
-            config, "get_complete_vertex_url", return_value="https://mock-url"
-        ) as mock_get_url,
+        patch.object(config, "_ensure_access_token", return_value=("token", "test-project")),
+        patch.object(config, "get_complete_vertex_url", return_value="https://mock-url") as mock_get_url,
     ):
         config.validate_anthropic_messages_environment(
             headers=headers,
@@ -51,17 +47,11 @@ def test_web_search_header_added_for_messages_endpoint():
         "vertex_credentials": "{}",
     }
     # Include web search tool in optional_params
-    optional_params = {
-        "tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 5}]
-    }
+    optional_params = {"tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 5}]}
 
     with (
-        patch.object(
-            config, "_ensure_access_token", return_value=("token", "test-project")
-        ),
-        patch.object(
-            config, "get_complete_vertex_url", return_value="https://mock-url"
-        ),
+        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,
@@ -73,12 +63,10 @@ def test_web_search_header_added_for_messages_endpoint():
         )
 
         # Assert that the anthropic-beta header with web-search is present
-        assert (
-            "anthropic-beta" in updated_headers
-        ), "anthropic-beta header should be present"
-        assert (
-            updated_headers["anthropic-beta"] == "web-search-2025-03-05"
-        ), f"anthropic-beta should be 'web-search-2025-03-05', got: {updated_headers['anthropic-beta']}"
+        assert "anthropic-beta" in updated_headers, "anthropic-beta header should be present"
+        assert updated_headers["anthropic-beta"] == "web-search-2025-03-05", (
+            f"anthropic-beta should be 'web-search-2025-03-05', got: {updated_headers['anthropic-beta']}"
+        )
 
 
 def test_web_search_header_not_added_without_tool():
@@ -94,12 +82,8 @@ def test_web_search_header_not_added_without_tool():
     optional_params = {}
 
     with (
-        patch.object(
-            config, "_ensure_access_token", return_value=("token", "test-project")
-        ),
-        patch.object(
-            config, "get_complete_vertex_url", return_value="https://mock-url"
-        ),
+        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,
@@ -111,9 +95,9 @@ 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"
+        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():
@@ -129,12 +113,8 @@ def test_compact_context_management_header_added():
     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"
-        ),
+        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,
@@ -146,12 +126,10 @@ def test_compact_context_management_header_added():
         )
 
         # 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']}"
+        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():
@@ -167,12 +145,8 @@ def test_context_management_header_added_for_other_edits():
     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"
-        ),
+        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,
@@ -184,12 +158,10 @@ def test_context_management_header_added_for_other_edits():
         )
 
         # 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']}"
+        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():
@@ -202,19 +174,11 @@ def test_both_compact_and_context_management_headers_added():
         "vertex_credentials": "{}",
     }
     # Include context_management with both compact and other edit types
-    optional_params = {
-        "context_management": {
-            "edits": [{"type": "compact_20260112"}, {"type": "some_other_type"}]
-        }
-    }
+    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"
-        ),
+        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,
@@ -226,15 +190,13 @@ def test_both_compact_and_context_management_headers_added():
         )
 
         # 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']}"
+        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']}"
+        )
 
 
 def test_validate_environment_always_refreshes_token_ignoring_stale_bearer():
@@ -248,12 +210,8 @@ def test_validate_environment_always_refreshes_token_ignoring_stale_bearer():
     }
 
     with (
-        patch.object(
-            config, "_ensure_access_token", return_value=("fresh-token", "test-project")
-        ) as mock_ensure,
-        patch.object(
-            config, "get_complete_vertex_url", return_value="https://mock-vertex-url"
-        ),
+        patch.object(config, "_ensure_access_token", return_value=("fresh-token", "test-project")) as mock_ensure,
+        patch.object(config, "get_complete_vertex_url", return_value="https://mock-vertex-url"),
     ):
         updated_headers, api_base = config.validate_anthropic_messages_environment(
             headers=headers,
@@ -286,9 +244,7 @@ def test_validate_environment_appends_stream_raw_predict_with_custom_api_base():
             "get_complete_vertex_url",
             wraps=config.get_complete_vertex_url,
         ) as spy_get_url,
-        patch.object(
-            config, "_ensure_access_token", return_value=("token", "test-project")
-        ),
+        patch.object(config, "_ensure_access_token", return_value=("token", "test-project")),
     ):
         _, api_base = config.validate_anthropic_messages_environment(
             headers={},
@@ -318,9 +274,7 @@ def test_validate_environment_appends_raw_predict_with_custom_api_base():
             "get_complete_vertex_url",
             wraps=config.get_complete_vertex_url,
         ) as spy_get_url,
-        patch.object(
-            config, "_ensure_access_token", return_value=("token", "test-project")
-        ),
+        patch.object(config, "_ensure_access_token", return_value=("token", "test-project")),
     ):
         _, api_base = config.validate_anthropic_messages_environment(
             headers={},
@@ -336,6 +290,70 @@ def test_validate_environment_appends_raw_predict_with_custom_api_base():
     assert api_base.endswith(":rawPredict")
 
 
+def test_validate_environment_with_custom_api_base_appends_streaming_suffix():
+    """Ensure custom Vertex Anthropic api_base values still get the streaming endpoint suffix."""
+    config = VertexAIPartnerModelsAnthropicMessagesConfig()
+    headers = {"Authorization": "Bearer existing-token"}
+    custom_api_base = (
+        "https://aiplatform.us.rep.googleapis.com/v1/projects/test-project/"
+        "locations/us/publishers/anthropic/models/claude-sonnet-4-5@20250929"
+    )
+
+    with patch.object(config, "_ensure_access_token", return_value=("fresh-token", "test-project")):
+        updated_headers, api_base = config.validate_anthropic_messages_environment(
+            headers=headers,
+            model="claude-sonnet-4-5@20250929",
+            messages=[],
+            optional_params={"stream": True},
+            litellm_params={
+                "vertex_ai_project": "test-project",
+                "vertex_ai_location": "us",
+            },
+            api_base=custom_api_base,
+        )
+
+    assert api_base == f"{custom_api_base}:streamRawPredict?alt=sse"
+    assert updated_headers["Authorization"] == "Bearer fresh-token"
+
+
+@pytest.mark.parametrize(
+    "custom_api_base,stream",
+    [
+        (
+            "https://aiplatform.us.rep.googleapis.com/v1/projects/test-project/"
+            "locations/us/publishers/anthropic/models/claude-sonnet-4-5@20250929"
+            ":streamRawPredict?alt=sse",
+            True,
+        ),
+        (
+            "https://aiplatform.us.rep.googleapis.com/v1/projects/test-project/"
+            "locations/us/publishers/anthropic/models/claude-sonnet-4-5@20250929"
+            ":rawPredict",
+            False,
+        ),
+    ],
+)
+def test_validate_environment_with_fully_qualified_custom_api_base_is_preserved(custom_api_base: str, stream: bool):
+    config = VertexAIPartnerModelsAnthropicMessagesConfig()
+    headers = {"Authorization": "Bearer existing-token"}
+
+    with patch.object(config, "_ensure_access_token", return_value=("fresh-token", "test-project")):
+        updated_headers, api_base = config.validate_anthropic_messages_environment(
+            headers=headers,
+            model="claude-sonnet-4-5@20250929",
+            messages=[],
+            optional_params={"stream": stream},
+            litellm_params={
+                "vertex_ai_project": "test-project",
+                "vertex_ai_location": "us",
+            },
+            api_base=custom_api_base,
+        )
+
+    assert api_base == custom_api_base
+    assert updated_headers["Authorization"] == "Bearer fresh-token"
+
+
 def test_transform_anthropic_messages_request_removes_scope_from_cache_control():
     """Ensure scope field is removed from cache_control for Vertex AI (not supported)."""
     config = VertexAIPartnerModelsAnthropicMessagesConfig()
@@ -447,20 +465,14 @@ def test_validate_environment_does_not_mutate_caller_headers():
     caller_headers: dict = {}
 
     with (
-        patch.object(
-            config, "_ensure_access_token", return_value=("token", "test-project")
-        ),
-        patch.object(
-            config, "get_complete_vertex_url", return_value="https://mock-url"
-        ),
+        patch.object(config, "_ensure_access_token", return_value=("token", "test-project")),
+        patch.object(config, "get_complete_vertex_url", return_value="https://mock-url"),
     ):
         config.validate_anthropic_messages_environment(
             headers=caller_headers,
             model="claude-sonnet-4",
             messages=[],
-            optional_params={
-                "tools": [{"type": "web_search_20250305", "name": "web_search"}]
-            },
+            optional_params={"tools": [{"type": "web_search_20250305", "name": "web_search"}]},
             litellm_params={
                 "vertex_ai_project": "p",
                 "vertex_ai_location": "us-central1",
@@ -468,9 +480,7 @@ def test_validate_environment_does_not_mutate_caller_headers():
             api_base=None,
         )
 
-    assert (
-        caller_headers == {}
-    ), "validate_anthropic_messages_environment must not mutate the caller's headers dict"
+    assert caller_headers == {}, "validate_anthropic_messages_environment must not mutate the caller's headers dict"
 
 
 def test_vertex_claude_completion_does_not_mutate_shared_extra_headers():
@@ -483,12 +493,8 @@ def test_vertex_claude_completion_does_not_mutate_shared_extra_headers():
     mock_response = MagicMock()
 
     with (
-        patch.object(
-            handler, "_ensure_access_token", return_value=("ya29.fresh", "proj")
-        ),
-        patch.object(
-            handler, "get_complete_vertex_url", return_value="https://mock-url"
-        ),
+        patch.object(handler, "_ensure_access_token", return_value=("ya29.fresh", "proj")),
+        patch.object(handler, "get_complete_vertex_url", return_value="https://mock-url"),
         patch(
             "litellm.llms.anthropic.chat.AnthropicChatCompletion.completion",
             return_value=mock_response,
@@ -509,9 +515,7 @@ def test_vertex_claude_completion_does_not_mutate_shared_extra_headers():
             litellm_params={},
         )
 
-    assert (
-        shared_extra_headers == {}
-    ), "extra_headers must not be mutated by completion()"
+    assert shared_extra_headers == {}, "extra_headers must not be mutated by completion()"
 
 
 
@@ -541,9 +545,7 @@ def test_messages_thinking_shape_follows_exact_vertex_entry_flag(local_model_cos
     assert result.get("thinking") == {"type": "adaptive", "display": "summarized"}
     assert result.get("output_config") == {"effort": "medium"}
 
-    monkeypatch.setitem(
-        litellm.model_cost["vertex_ai/claude-opus-4-8"], "supports_adaptive_thinking", False
-    )
+    monkeypatch.setitem(litellm.model_cost["vertex_ai/claude-opus-4-8"], "supports_adaptive_thinking", False)
     litellm.get_model_info.cache_clear()
     assert litellm.model_cost["claude-opus-4-8"]["supports_adaptive_thinking"] is True
 
@@ -614,9 +616,7 @@ class TestVertexAnthropicMidConversationSystem:
             {"role": "assistant", "content": "reading"},
             {"role": "user", "content": "continue"},
         ]
-        result = _vertex_transform(
-            "claude-sonnet-4-6", messages, system=[{"type": "text", "text": "Base."}]
-        )
+        result = _vertex_transform("claude-sonnet-4-6", messages, system=[{"type": "text", "text": "Base."}])
         assert result["messages"] == [
             {"role": "user", "content": "read the file"},
             {
@@ -660,9 +660,7 @@ def test_vertex_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_f
 
     import litellm
 
-    cost_map_path = os.path.join(
-        os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json"
-    )
+    cost_map_path = os.path.join(os.path.dirname(litellm.__file__), "model_prices_and_context_window_backup.json")
     with open(cost_map_path) as f:
         cost_map = json.load(f)
     rules = cost_map["fallback_generalizations"]["rules"]