Merge 146028d907 into 8d6b8dab3d

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

@@ -52,6 +52,71 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert
             project_id=vertex_ai_project,
             custom_llm_provider="vertex_ai",
         )
+        return self._finalize_vertex_anthropic_messages_environment(
+            headers=headers,
+            model=model,
+            messages=messages,
+            optional_params=optional_params,
+            api_base=api_base,
+            access_token=access_token,
+            project_id=project_id,
+            vertex_ai_project=vertex_ai_project,
+            vertex_ai_location=vertex_ai_location,
+        )
+
+    async def avalidate_anthropic_messages_environment(
+        self,
+        headers: dict,  # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
+        model: str,
+        messages: list[Any],  # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
+        optional_params: dict,  # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
+        litellm_params: dict,  # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
+        api_key: str | None = None,
+        api_base: str | None = None,
+    ) -> tuple[dict, str | None]:  # mutable-ok: mirrors the sync validate_anthropic_messages_environment contract
+        """Async override so /v1/messages does not block the event loop on Google auth.
+
+        AnthropicMessagesConfig.avalidate delegates to a subclass's sync
+        validate when sync is overridden. That would call
+        ``_ensure_access_token`` (blocking ``requests`` / GCE metadata) on the
+        loop thread. Use ``_ensure_access_token_async`` instead, matching the
+        chat/completions Vertex path.
+        """
+        headers = dict(headers)
+        vertex_ai_project: Final = VertexBase.safe_get_vertex_ai_project(litellm_params)
+        vertex_ai_location: Final = VertexBase.safe_get_vertex_ai_location(litellm_params)
+
+        vertex_credentials: Final = VertexBase.safe_get_vertex_ai_credentials(litellm_params)
+        access_token, project_id = await self._ensure_access_token_async(
+            credentials=vertex_credentials,
+            project_id=vertex_ai_project,
+            custom_llm_provider="vertex_ai",
+        )
+        return self._finalize_vertex_anthropic_messages_environment(
+            headers=headers,
+            model=model,
+            messages=messages,
+            optional_params=optional_params,
+            api_base=api_base,
+            access_token=access_token,
+            project_id=project_id,
+            vertex_ai_project=vertex_ai_project,
+            vertex_ai_location=vertex_ai_location,
+        )
+
+    def _finalize_vertex_anthropic_messages_environment(
+        self,
+        headers: dict,
+        model: str,
+        messages: list[Any],
+        optional_params: dict,
+        api_base: str | None,
+        access_token: str,
+        project_id: str | None,
+        vertex_ai_project: str | None,
+        vertex_ai_location: str | None,
+    ) -> tuple[dict, str | None]:
+        """Shared header/URL building after the access token is resolved."""
         headers["Authorization"] = f"Bearer {access_token}"
 
         api_base = self.get_complete_vertex_url(

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

@@ -3,7 +3,7 @@ import json
 import os
 from collections.abc import Mapping
 from typing import Final
-from unittest.mock import MagicMock, patch
+from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 
@@ -823,3 +823,98 @@ def test_vertex_claude_4_8_plus_cost_map_entries_carry_mid_conversation_system_f
         and info.get("supports_mid_conversation_system") is not True
     ]
     assert missing == []
+
+
+@pytest.mark.asyncio
+async def test_avalidate_uses_async_access_token_not_sync():
+    """/v1/messages must not call blocking _ensure_access_token on the event loop.
+
+    AnthropicMessagesConfig.avalidate would otherwise delegate to the Vertex
+    sync validate (which does), freezing the proxy worker on slow GCE metadata.
+    """
+    config = VertexAIPartnerModelsAnthropicMessagesConfig()
+    litellm_params = {
+        "vertex_ai_project": "test-project",
+        "vertex_ai_location": "us-east5",
+        "vertex_credentials": "{}",
+    }
+    optional_params = {"tools": [{"type": "web_search_20250305", "name": "web_search", "max_uses": 1}]}
+
+    with (
+        patch.object(
+            config,
+            "_ensure_access_token_async",
+            new_callable=AsyncMock,
+            return_value=("async-token", "test-project"),
+        ) as mock_async,
+        patch.object(
+            config,
+            "_ensure_access_token",
+            side_effect=AssertionError("sync _ensure_access_token must not run on async path"),
+        ) as mock_sync,
+        patch.object(config, "get_complete_vertex_url", return_value="https://mock-url") as mock_get_url,
+    ):
+        updated_headers, api_base = await config.avalidate_anthropic_messages_environment(
+            headers={},
+            model="claude-sonnet-4-5",
+            messages=[{"role": "user", "content": "hi"}],
+            optional_params=optional_params,
+            litellm_params=litellm_params,
+            api_base=None,
+        )
+
+    mock_async.assert_awaited_once()
+    mock_sync.assert_not_called()
+    assert mock_async.await_args.kwargs["custom_llm_provider"] == "vertex_ai"
+    assert mock_async.await_args.kwargs["project_id"] == "test-project"
+    assert updated_headers["Authorization"] == "Bearer async-token"
+    assert updated_headers["content-type"] == "application/json"
+    assert "web-search-2025-03-05" in updated_headers["anthropic-beta"]
+    assert api_base == "https://mock-url"
+    assert mock_get_url.call_args.kwargs["vertex_location"] == "us-east5"
+
+
+@pytest.mark.asyncio
+async def test_avalidate_matches_sync_headers_and_url():
+    """Async and sync finalize paths should produce the same headers/api_base."""
+    config = VertexAIPartnerModelsAnthropicMessagesConfig()
+    litellm_params = {
+        "vertex_ai_project": "proj-a",
+        "vertex_ai_location": "europe-west1",
+        "vertex_credentials": "{}",
+    }
+    optional_params = {
+        "context_management": {"edits": [{"type": "clear_thinking_20251015"}]},
+        "tools": [],
+    }
+    messages = [{"role": "user", "content": "hi"}]
+
+    with (
+        patch.object(config, "_ensure_access_token", return_value=("tok", "proj-a")),
+        patch.object(
+            config,
+            "_ensure_access_token_async",
+            new_callable=AsyncMock,
+            return_value=("tok", "proj-a"),
+        ),
+        patch.object(config, "get_complete_vertex_url", return_value="https://eu.mock"),
+    ):
+        sync_headers, sync_base = config.validate_anthropic_messages_environment(
+            headers={"anthropic-beta": "existing-beta"},
+            model="claude-3-sonnet",
+            messages=messages,
+            optional_params=optional_params,
+            litellm_params=dict(litellm_params),
+            api_base=None,
+        )
+        async_headers, async_base = await config.avalidate_anthropic_messages_environment(
+            headers={"anthropic-beta": "existing-beta"},
+            model="claude-3-sonnet",
+            messages=messages,
+            optional_params=optional_params,
+            litellm_params=dict(litellm_params),
+            api_base=None,
+        )
+
+    assert async_headers == sync_headers
+    assert async_base == sync_base