diff --git a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py index cfd5b5d3980..eefeda190ca 100644 --- a/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py +++ b/litellm/llms/vertex_ai/vertex_ai_partner_models/anthropic/experimental_pass_through/transformation.py @@ -1,4 +1,4 @@ -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from typing import Final from litellm.llms.anthropic.common_utils import AnthropicModelInfo @@ -26,34 +26,92 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert def should_strip_billing_metadata(self) -> bool: return True + @staticmethod + def _vertex_auth_parameters( + litellm_params: dict[str, object], # mutable-ok: delegates to the existing Vertex parameter readers + ) -> tuple[str | None, str | None, str | None]: + return ( + VertexBase.safe_get_vertex_ai_project(litellm_params), + VertexBase.safe_get_vertex_ai_location(litellm_params), + VertexBase.safe_get_vertex_ai_credentials(litellm_params), + ) + def validate_anthropic_messages_environment( self, - headers: dict, + headers: dict[str, str], model: str, messages: Sequence[object], - optional_params: dict, - litellm_params: dict, + optional_params: dict[str, object], + litellm_params: dict[str, object], api_key: str | None = None, api_base: str | None = None, - ) -> tuple[dict, str | None]: + ) -> tuple[dict[str, str], str | None]: """ OPTIONAL Validate the environment for the request """ - # Work on a local copy — router shallow-copies litellm_params so the caller's - # headers dict may be the shared deployment extra_headers object. - 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) + vertex_ai_project, vertex_ai_location, vertex_credentials = self._vertex_auth_parameters(litellm_params) access_token, project_id = self._ensure_access_token( credentials=vertex_credentials, project_id=vertex_ai_project, custom_llm_provider="vertex_ai", ) - headers["Authorization"] = f"Bearer {access_token}" + return self._finalize_vertex_messages_environment( + headers, + model, + messages, + optional_params, + api_base, + access_token, + project_id, + vertex_ai_project, + vertex_ai_location, + ) + + async def avalidate_anthropic_messages_environment( + self, + headers: dict[str, str], # mutable-ok: mirrors the provider validation contract + model: str, + messages: Sequence[object], + optional_params: dict[str, object], # mutable-ok: mirrors the provider validation contract + litellm_params: dict[str, object], # mutable-ok: mirrors the provider validation contract + api_key: str | None = None, + api_base: str | None = None, + ) -> tuple[dict[str, str], str | None]: # mutable-ok: the caller consumes request headers + vertex_ai_project, vertex_ai_location, vertex_credentials = self._vertex_auth_parameters(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_messages_environment( + headers, + model, + messages, + optional_params, + api_base, + access_token, + project_id, + vertex_ai_project, + vertex_ai_location, + ) + + def _finalize_vertex_messages_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[object], + optional_params: dict, # mutable-ok: preserves the existing provider optional-parameter contract + api_base: str | None, + access_token: str, + project_id: str, + vertex_ai_project: str | None, + vertex_ai_location: str | None, + ) -> tuple[dict[str, str], str | None]: # mutable-ok: the caller consumes request headers + # Router shallow-copies litellm_params, so headers may be shared between requests. + request_headers: Final = dict(headers) + request_headers["Authorization"] = f"Bearer {access_token}" api_base = self.get_complete_vertex_url( custom_api_base=api_base, @@ -65,14 +123,14 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert model=model, ) - headers["content-type"] = "application/json" + request_headers["content-type"] = "application/json" # Add beta headers for Vertex AI tools: Final = optional_params.get("tools", []) beta_values: Final[set[str]] = set() # Get existing beta headers if any - existing_beta: Final = headers.get("anthropic-beta") + existing_beta: Final = request_headers.get("anthropic-beta") if existing_beta: beta_values.update(b.strip() for b in existing_beta.split(",")) @@ -117,9 +175,9 @@ class VertexAIPartnerModelsAnthropicMessagesConfig(AnthropicMessagesConfig, Vert beta_values.add(ANTHROPIC_BETA_HEADER_VALUES.PER_TURN_CONTROL_2026_07_01.value) if beta_values: - headers["anthropic-beta"] = ",".join(beta_values) + request_headers["anthropic-beta"] = ",".join(beta_values) - return headers, api_base + return request_headers, api_base def get_complete_url( self, diff --git a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py index f4920505722..ff8b75af4d4 100644 --- a/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py +++ b/tests/unit/llms/anthropic/pass_through/messages/test_anthropic_experimental_pass_through_messages_handler.py @@ -6,6 +6,7 @@ import uuid from typing import Any, Dict, Final, List import httpx +import responses import respx import pytest from fastapi.testclient import TestClient @@ -1709,16 +1710,20 @@ async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_bet @pytest.mark.asyncio @pytest.mark.parametrize("client_headers", _CLIENT_BETA_HEADERS) async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_beta_to_vertex( - local_beta_headers_config, client_headers -): + local_beta_headers_config: None, client_headers: dict[str, str] +) -> None: """Vertex rawPredict takes the beta as the `anthropic-beta` header and 400s on `safeguards` without it, so the beta rides along with the field.""" from litellm.llms.anthropic.pass_through.messages import handler - from litellm.llms.vertex_ai.vertex_llm_base import VertexBase safeguards, safeguard_results = _claude_code_auto_mode_request() captured: dict[str, object] = {} - with patch.object(VertexBase, "_ensure_access_token", return_value=("test-token", "test-project")): + with responses.RequestsMock(assert_all_requests_are_fired=False) as oauth: + oauth.add( + responses.POST, + "https://oauth2.googleapis.com/token", + json={"access_token": "test-token", "expires_in": 3600}, + ) response = await handler.anthropic_messages( max_tokens=16, messages=[{"role": "user", "content": "hi"}], @@ -1726,7 +1731,14 @@ async def test_anthropic_messages_forwards_safeguards_and_dangerous_tool_use_bet custom_llm_provider="vertex_ai", vertex_project="test-project", vertex_location="global", - vertex_credentials="{}", + vertex_credentials=json.dumps( + { + "type": "authorized_user", + "client_id": "synthetic-client", + "client_secret": "synthetic-secret", + "refresh_token": "synthetic-refresh", + } + ), client=_upstream_answering_with(safeguard_results, captured), safeguards=safeguards, extra_headers=client_headers, diff --git a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py index ae9b9684f2d..136c444c9f1 100644 --- a/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py +++ b/tests/unit/llms/vertex_ai/vertex_ai_partner_models/anthropic/test_vertex_ai_partner_models_anthropic_messages_config.py @@ -2,18 +2,98 @@ import copy import json import os from collections.abc import Mapping +from datetime import datetime, timedelta +from queue import SimpleQueue +from threading import get_ident from typing import Final from unittest.mock import MagicMock, patch import pytest +import requests +import responses from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.experimental_pass_through.transformation import ( VertexAIPartnerModelsAnthropicMessagesConfig, ) from litellm.llms.vertex_ai.vertex_ai_partner_models.main import VertexAIPartnerModels +from litellm.types.llms.anthropic import ANTHROPIC_BETA_HEADER_VALUES, ANTHROPIC_HOSTED_TOOLS from litellm.types.router import GenericLiteLLMParams +@pytest.mark.parametrize("stream", (False, True)) +@pytest.mark.parametrize("infer_project", (False, True), ids=("explicit-project", "credential-project")) +@pytest.mark.asyncio +async def test_messages_auth_load_and_refresh_leave_the_event_loop(stream: bool, infer_project: bool) -> None: + config: Final = VertexAIPartnerModelsAnthropicMessagesConfig() + loop_thread: Final = get_ident() + token_threads: Final[SimpleQueue[int]] = SimpleQueue() + now: Final = datetime(2025, 1, 1) + model: Final = "claude-sonnet-5" + headers: Final = {"Authorization": "Bearer incoming", "anthropic-beta": "retained-beta", "x-request-id": "req-1"} + optional_params: Final = { + "stream": stream, + "tools": [{"type": ANTHROPIC_HOSTED_TOOLS.WEB_SEARCH.value, "name": "web_search"}], + } + litellm_params: Final = { + "vertex_ai_project": None if infer_project else "auth-project", + "vertex_ai_location": "global", + "vertex_credentials": json.dumps( + { + "type": "authorized_user", + "client_id": "synthetic-client", + "client_secret": "synthetic-secret", + "refresh_token": "synthetic-refresh", + "quota_project_id": "auth-project", + } + ), + } + + async def validate() -> tuple[Mapping[str, object], str | None]: + return await config.avalidate_anthropic_messages_environment( + headers=headers, model=model, messages=[], optional_params=optional_params, litellm_params=litellm_params + ) + + with patch("google.auth._helpers.utcnow", return_value=now) as clock, responses.RequestsMock() as oauth: + + def token_response(_request: requests.PreparedRequest) -> tuple[int, dict[str, str], str]: + token_threads.put(get_ident()) + return 200, {}, json.dumps({"access_token": f"token-{len(oauth.calls) + 1}", "expires_in": 3600}) + + oauth.add_callback(responses.POST, "https://oauth2.googleapis.com/token", callback=token_response) + cold_headers, cold_url = await validate() + assert token_threads.get_nowait() != loop_thread + assert cold_headers["Authorization"] == "Bearer token-1" + assert await validate() == (cold_headers, cold_url) + assert len(oauth.calls) == 1 + + clock.return_value = now + timedelta(hours=2) + refreshed_headers, refreshed_url = await validate() + assert token_threads.get_nowait() != loop_thread + assert refreshed_headers["Authorization"] == "Bearer token-2" + assert refreshed_url == cold_url + assert len(oauth.calls) == 2 + assert config.validate_anthropic_messages_environment( + headers=headers, model=model, messages=[], optional_params=optional_params, litellm_params=litellm_params + ) == (refreshed_headers, refreshed_url) + assert len(oauth.calls) == 2 + + endpoint: Final = "streamRawPredict" if stream else "rawPredict" + assert cold_url == ( + "https://aiplatform.googleapis.com/v1/projects/auth-project/locations/global/" + f"publishers/anthropic/models/{model}:{endpoint}" + ) + assert refreshed_headers == { + "Authorization": "Bearer token-2", + "anthropic-beta": cold_headers["anthropic-beta"], + "x-request-id": "req-1", + "content-type": "application/json", + } + beta_header: Final = cold_headers["anthropic-beta"] + assert isinstance(beta_header, str) + assert set(beta_header.split(",")) == {"retained-beta", ANTHROPIC_BETA_HEADER_VALUES.WEB_SEARCH_2025_03_05.value} + assert headers == {"Authorization": "Bearer incoming", "anthropic-beta": "retained-beta", "x-request-id": "req-1"} + + def test_validate_environment_uses_vertex_ai_location(): config = VertexAIPartnerModelsAnthropicMessagesConfig() headers = {}