diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index b48b8d81494..29f216fd450 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -1730,6 +1730,16 @@ def get_vertex_ai_allowed_incoming_headers(request: Request) -> dict: return headers +def _is_vertex_anthropic_count_tokens_route(endpoint: str) -> bool: + return endpoint.rsplit("/", 1)[-1].split(":", 1)[0] == "count-tokens" + + +def _upstream_headers_for_vertex_route(endpoint: str, headers: Mapping[str, str]) -> Mapping[str, str]: + if not _is_vertex_anthropic_count_tokens_route(endpoint): + return headers + return MappingProxyType({name: value for name, value in headers.items() if name.lower() != "anthropic-beta"}) + + def get_vertex_pass_through_handler( call_type: Literal["discovery", "aiplatform"], # noqa: UP037 # ruff reports quoted Literal values here ) -> BaseVertexAIPassThroughHandler: @@ -2128,7 +2138,7 @@ async def _base_vertex_proxy_route( endpoint_func: Final = create_pass_through_route( endpoint=endpoint, target=target, - custom_headers=headers, + custom_headers=_upstream_headers_for_vertex_route(endpoint, headers), is_streaming_request=is_streaming_request, ) # dynamically construct pass-through endpoint based on incoming path diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py index 961479c0393..6735f2a3780 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py @@ -5,6 +5,7 @@ import pytest from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( _base_vertex_proxy_route, + _upstream_headers_for_vertex_route, ) from litellm.types.router import DeploymentTypedDict @@ -348,6 +349,93 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header(): assert headers_passed_through is False +VERTEX_ANTHROPIC_MODELS_PREFIX = "v1/projects/test-project/locations/global/publishers/anthropic/models/" + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("model_segment", "expects_anthropic_beta"), + [ + ("count-tokens:rawPredict", False), + ("claude-sonnet-4-6:streamRawPredict", True), + ], +) +async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens( + model_segment: str, expects_anthropic_beta: bool +): + with ( + patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it + "litellm.proxy.proxy_server.llm_router", None + ), + patch( # test-quality-ok: the route reads this proxy global at call time, nothing injects it + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router" + ) as mock_pt_router, + patch( # test-quality-ok: the route offers no injection point for its header preparation + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers", + new_callable=AsyncMock, + ) as mock_prep_headers, + patch( # test-quality-ok: the upstream call is captured here, the route offers no injection point + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" + ) as mock_create_route, + patch( # test-quality-ok: the route calls auth directly rather than through Depends + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", + new_callable=AsyncMock, + ) as mock_auth, + patch( # test-quality-ok: the route reads the request body for this, a MagicMock request has none + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn", + new_callable=AsyncMock, + return_value=False, + ), + ): + mock_pt_router.get_vertex_credentials.return_value = MagicMock() + mock_prep_headers.return_value = ( + { + "anthropic-beta": "tool-search-tool-2025-10-19,web-search-2025-03-05", + "content-type": "application/json", + "Authorization": "Bearer vertex-access-token", + }, + "https://aiplatform.googleapis.com", + False, + "test-project", + "global", + ) + mock_create_route.return_value = AsyncMock() + mock_auth.return_value = UserAPIKeyAuth(api_key="sk-litellm-secret-key") + + await _base_vertex_proxy_route( + endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}", + request=MagicMock(), + fastapi_response=MagicMock(), + get_vertex_pass_through_handler=MagicMock(), + ) + + upstream_headers = mock_create_route.call_args.kwargs["custom_headers"] + assert ("anthropic-beta" in upstream_headers) is expects_anthropic_beta + assert upstream_headers["Authorization"] == "Bearer vertex-access-token" + assert upstream_headers["content-type"] == "application/json" + + +def test_upstream_headers_for_vertex_route_filters_anthropic_beta_by_route(): + headers = { + "Anthropic-Beta": "effort-2025-11-24", + "content-type": "application/json", + "Authorization": "Bearer vertex-access-token", + } + + count_tokens_headers = _upstream_headers_for_vertex_route( + f"{VERTEX_ANTHROPIC_MODELS_PREFIX}count-tokens:rawPredict", headers + ) + model_headers = _upstream_headers_for_vertex_route( + f"{VERTEX_ANTHROPIC_MODELS_PREFIX}claude-sonnet-4-6:rawPredict", headers + ) + + assert dict(count_tokens_headers) == { + "content-type": "application/json", + "Authorization": "Bearer vertex-access-token", + } + assert dict(model_headers) == headers + + @pytest.mark.asyncio async def test_vertex_passthrough_does_not_forward_litellm_auth_token(): """