scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-07 08:26:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 89

fix(proxy): apply default_vertex_config location before building the Vertex passthrough base URL

Browse source 
Routes without /projects/<project>/locations/<location>/ built the upstream host from the URL's
still-empty location and 500ed even with default_vertex_config set. Build the base URL once after
the configured project and location are applied, drop the hook that re-derived it afterwards, and
answer 400 with a fix-it message when no location is available at all.

Resolves LIT-6905
This commit is contained in:
mateo-berri 2026-09-03 15:27:47 -07:00
parent 8699998c9e
commit 57da95a77c
3 changed files with 147 additions and 67 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

42
litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py
View file

@ -1659,6 +1659,12 @@ async def azure_proxy_route(
from abc import ABC, abstractmethod
_VERTEX_LOCATION_REQUIRED_DETAIL: Final = (
"No Vertex AI location for this request. Include /projects/<project>/locations/<location>/ in the "
"route, set vertex_location in default_vertex_config (or DEFAULT_VERTEXAI_LOCATION), or add the "
"model to model_list with use_in_pass_through: true."
)
class BaseVertexAIPassThroughHandler(ABC):
@staticmethod
@ -1666,29 +1672,18 @@ class BaseVertexAIPassThroughHandler(ABC):
def get_default_base_target_url(vertex_location: str | None) -> str:
pass
@staticmethod
@abstractmethod
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
pass
class VertexAIDiscoveryPassThroughHandler(BaseVertexAIPassThroughHandler):
@staticmethod
def get_default_base_target_url(vertex_location: str | None) -> str:
return "https://discoveryengine.googleapis.com/"
@staticmethod
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
return base_target_url
class VertexAIPassThroughHandler(BaseVertexAIPassThroughHandler):
@staticmethod
def get_default_base_target_url(vertex_location: str | None) -> str:
return get_vertex_base_url(vertex_location)
@staticmethod
def update_base_target_url_with_credential_location(base_target_url: str, vertex_location: str | None) -> str:
if vertex_location is None:
raise HTTPException(status_code=400, detail=_VERTEX_LOCATION_REQUIRED_DETAIL)
return get_vertex_base_url(vertex_location)
@ -1911,10 +1906,8 @@ async def _prepare_vertex_auth_headers(
router_credentials: LiteLLM_ManagedVectorStore | None,
vertex_project: str | None,
vertex_location: str | None,
base_target_url: str | None,
get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler,
user_api_key_dict: UserAPIKeyAuth,
) -> tuple[Mapping[str, str], str | None, bool, str | None, str | None]:
) -> tuple[Mapping[str, str], bool, str | None, str | None]:
"""
Prepare authentication headers for Vertex AI pass-through requests.
@ -1924,15 +1917,12 @@ async def _prepare_vertex_auth_headers(
router_credentials: Optional vector store credentials from registry
vertex_project: Vertex project ID
vertex_location: Vertex location
base_target_url: Base URL for the Vertex AI service
get_vertex_pass_through_handler: Handler for the specific Vertex AI service
user_api_key_dict: The caller's resolved authentication, so only the secret that
authenticated them is stripped on the credential-less branch
Returns:
tuple containing:
- headers: dict - Authentication headers to use
- base_target_url: str | None - Updated base target URL
- headers_passed_through: bool - Whether headers were passed through from request
- vertex_project: str | None - Updated vertex project ID
- vertex_location: str | None - Updated vertex location
@ -1985,14 +1975,8 @@ async def _prepare_vertex_auth_headers(
# Add the Authorization header with vendor credentials
headers["Authorization"] = f"Bearer {auth_header}"
if base_target_url is not None:
base_target_url = get_vertex_pass_through_handler.update_base_target_url_with_credential_location(
base_target_url, vertex_location
)
return (
headers,
base_target_url,
headers_passed_through,
vertex_project,
vertex_location,
@ -2085,12 +2069,9 @@ async def _base_vertex_proxy_route(
location=vertex_location,
)
base_target_url = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
# Prepare authentication headers
(
headers,
base_target_url,
headers_passed_through,
vertex_project,
vertex_location,
@ -2100,13 +2081,10 @@ async def _base_vertex_proxy_route(
router_credentials=router_credentials,
vertex_project=vertex_project,
vertex_location=vertex_location,
base_target_url=base_target_url,
get_vertex_pass_through_handler=get_vertex_pass_through_handler,
user_api_key_dict=user_api_key_dict,
)
if base_target_url is None:
base_target_url = get_vertex_base_url(vertex_location)
base_target_url: Final = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
request_route: Final = encoded_endpoint
verbose_proxy_logger.debug("request_route %s", request_route)

147
tests/test_litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py
View file

@ -319,9 +319,6 @@ class TestVertexAIPassThroughHandler:
mock_handler.get_default_base_target_url.return_value = (
f"https://{test_location}-aiplatform.googleapis.com/"
)
mock_handler.update_base_target_url_with_credential_location = Mock(
return_value=f"https://{test_location}-aiplatform.googleapis.com/"
)
mock_get_handler.return_value = mock_handler
# Mock create_pass_through_route to return a function that returns a mock response
@ -427,9 +424,6 @@ class TestVertexAIPassThroughHandler:
mock_handler.get_default_base_target_url.return_value = (
"https://aiplatform.googleapis.com/"
)
mock_handler.update_base_target_url_with_credential_location = Mock(
return_value="https://aiplatform.googleapis.com/"
)
mock_get_handler.return_value = mock_handler
# Mock create_pass_through_route to return a function that returns a mock response
@ -530,9 +524,6 @@ class TestVertexAIPassThroughHandler:
mock_handler.get_default_base_target_url.return_value = (
f"https://{default_location}-aiplatform.googleapis.com/"
)
mock_handler.update_base_target_url_with_credential_location = Mock(
return_value=f"https://{default_location}-aiplatform.googleapis.com/"
)
mock_get_handler.return_value = mock_handler
# Mock create_pass_through_route to return a function that returns a mock response
@ -1308,9 +1299,6 @@ class TestVertexAIDiscoveryPassThroughHandler:
mock_handler.get_default_base_target_url.return_value = (
"https://discoveryengine.googleapis.com"
)
mock_handler.update_base_target_url_with_credential_location = Mock(
return_value="https://discoveryengine.googleapis.com"
)
mock_get_handler.return_value = mock_handler
# Mock create_pass_through_route to return a function that returns a mock response
@ -3650,7 +3638,6 @@ class TestVertexRawPredictStreamingClassification:
base_url = "https://us-east5-aiplatform.googleapis.com/"
mock_handler = Mock()
mock_handler.get_default_base_target_url.return_value = base_url
mock_handler.update_base_target_url_with_credential_location = Mock(return_value=base_url)
module = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints"
with (
@ -4234,6 +4221,140 @@ class TestVertexCredentiallessPassthroughVirtualKeyLeak:
assert "sk-master-1234" not in " ".join(f"{name}:{value}" for name, value in forwarded.items())
class TestVertexPassthroughDefaultLocationOnShortRoutes:
"""Regression coverage for LIT-6905.
``default_vertex_config`` carries the project and location, yet a route that
omits ``/projects/<project>/locations/<location>/`` built the upstream base URL
from the still-unresolved URL location and 500ed with ``vertex_location is
required``. The base URL must be built after the configured location is
applied, and a request with no location anywhere must fail with a clean 400
that says where a location can come from, never a 500.
"""
PROJECT = "test-project"
SHORT_ROUTE = "publishers/google/models/gemini-2.5-flash:generateContent"
async def _forward(
self,
monkeypatch,
endpoint: str,
default_config: dict | None,
headers: list[tuple[bytes, bytes]],
) -> tuple[HTTPException | None, dict]:
from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import (
PassthroughEndpointRouter,
)
async def receive():
return {"type": "http.request", "body": b"{}", "more_body": False}
request = Request(
{
"type": "http",
"method": "POST",
"path": f"/vertex_ai/{endpoint}",
"headers": headers,
"query_string": b"",
},
receive=receive,
)
captured: dict = {}
def fake_create_pass_through_route(**kwargs):
captured.update(kwargs)
return AsyncMock(return_value={"status": "success"})
router = PassthroughEndpointRouter()
if default_config is not None:
router.set_default_vertex_config(dict(default_config))
module = "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints"
monkeypatch.setattr(f"{module}.passthrough_endpoint_router", router)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None)
mock_credentials = Mock()
mock_credentials.token = "test-token"
caller: Final = UserAPIKeyAuth(api_key="test-key")
raised: HTTPException | None = None
with (
mock.patch( # test-quality-ok: the route mints its Google token through its own VertexBase, nothing injects the credential loader
"litellm.llms.vertex_ai.vertex_llm_base.VertexBase.load_auth",
return_value=(mock_credentials, self.PROJECT),
),
mock.patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route),
mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=caller)),
):
try:
await vertex_proxy_route(
endpoint=endpoint,
request=request,
fastapi_response=Response(),
user_api_key_dict=caller,
)
except HTTPException as exc:
raised = exc
return raised, captured
@pytest.mark.asyncio
@pytest.mark.parametrize(
("endpoint", "location", "expected_target"),
[
(
SHORT_ROUTE,
"global",
"https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/" + SHORT_ROUTE,
),
(
f"v1/{SHORT_ROUTE}",
"global",
"https://aiplatform.googleapis.com/v1/projects/test-project/locations/global/" + SHORT_ROUTE,
),
(
f"v1beta1/{SHORT_ROUTE}",
"global",
"https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/global/" + SHORT_ROUTE,
),
(
SHORT_ROUTE,
"us-central1",
"https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/"
+ SHORT_ROUTE,
),
],
)
async def test_default_vertex_config_location_fills_routes_without_project_and_location(
self, monkeypatch, endpoint, location, expected_target
):
raised, captured = await self._forward(
monkeypatch,
endpoint,
{"vertex_project": self.PROJECT, "vertex_location": location, "vertex_credentials": "test-creds"},
[(b"content-type", b"application/json"), (b"authorization", b"Bearer test-key")],
)
assert raised is None
assert str(captured["target"]) == expected_target
assert captured["custom_headers"]["Authorization"] == "Bearer test-token"
@pytest.mark.asyncio
@pytest.mark.parametrize(
("default_config", "headers"),
[
(None, [(b"content-type", b"application/json"), (b"authorization", b"Bearer ya29.byo-google-oauth")]),
(
{"vertex_project": PROJECT, "vertex_credentials": "test-creds"},
[(b"content-type", b"application/json"), (b"authorization", b"Bearer test-key")],
),
],
)
async def test_no_location_anywhere_is_a_400_not_a_500(self, monkeypatch, default_config, headers):
raised, captured = await self._forward(monkeypatch, self.SHORT_ROUTE, default_config, headers)
assert not captured, "a request with no location must never reach the upstream forwarder"
assert raised is not None
assert raised.status_code == 400
assert "/projects/<project>/locations/<location>/" in str(raised.detail)
assert "default_vertex_config" in str(raised.detail)
class TestGetAzureAISearchIndexFromEndpoint:
"""The operable index is only the segment right after ``indexes``.

25
tests/test_litellm/proxy/pass_through_endpoints/test_vertex_passthrough_load_balancing.py
View file

@ -4,6 +4,7 @@ import pytest
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import (
VertexAIPassThroughHandler,
_base_vertex_proxy_route,
_upstream_headers_for_vertex_route,
)
@ -20,6 +21,7 @@ async def test_vertex_passthrough_load_balancing():
mock_request = MagicMock()
mock_response = MagicMock()
mock_handler = MagicMock()
mock_handler.get_default_base_target_url.return_value = "https://test.url"
# Mock the router
mock_router = MagicMock()
@ -68,7 +70,6 @@ async def test_vertex_passthrough_load_balancing():
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
mock_prep_headers.return_value = (
{},
"https://test.url",
False,
"test-project-lb",
"us-central1-lb",
@ -290,12 +291,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header():
mock_vertex_credentials.vertex_location = "us-central1"
mock_vertex_credentials.vertex_credentials = "test-credentials"
# Create mock handler
mock_handler = MagicMock()
mock_handler.update_base_target_url_with_credential_location.return_value = (
"https://us-central1-aiplatform.googleapis.com"
)
with (
patch.object(
VertexBase,
@ -313,7 +308,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header():
# Call the function
(
headers,
base_target_url,
headers_passed_through,
vertex_project,
vertex_location,
@ -323,8 +317,6 @@ async def test_vertex_passthrough_forwards_anthropic_beta_header():
router_credentials=None,
vertex_project="test-project",
vertex_location="us-central1",
base_target_url="https://us-central1-aiplatform.googleapis.com",
get_vertex_pass_through_handler=mock_handler,
user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"),
)
@ -394,7 +386,6 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
"content-type": "application/json",
"Authorization": "Bearer vertex-access-token",
},
"https://aiplatform.googleapis.com",
False,
"test-project",
"global",
@ -406,7 +397,7 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}",
request=MagicMock(),
fastapi_response=MagicMock(),
get_vertex_pass_through_handler=MagicMock(),
get_vertex_pass_through_handler=VertexAIPassThroughHandler(),
)
upstream_headers = mock_create_route.call_args.kwargs["custom_headers"]
@ -473,12 +464,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token():
mock_vertex_credentials.vertex_location = "us-central1"
mock_vertex_credentials.vertex_credentials = "test-credentials"
# Create mock handler
mock_handler = MagicMock()
mock_handler.update_base_target_url_with_credential_location.return_value = (
"https://us-central1-aiplatform.googleapis.com"
)
with (
patch.object(
VertexBase,
@ -495,7 +480,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token():
(
headers,
_base_target_url,
_headers_passed_through,
_vertex_project,
_vertex_location,
@ -505,8 +489,6 @@ async def test_vertex_passthrough_does_not_forward_litellm_auth_token():
router_credentials=None,
vertex_project="test-project",
vertex_location="us-central1",
base_target_url="https://us-central1-aiplatform.googleapis.com",
get_vertex_pass_through_handler=mock_handler,
user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"),
)
@ -742,7 +724,6 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url():
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
mock_prep_headers.return_value = (
{},
"https://global-aiplatform.googleapis.com",
False,
"nv-gcpllmgwit-20250411173346",
"global",