From d12ce3cd5de9f4f760249b85a7bc619208ccc393 Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Wed, 28 Jan 2026 16:54:01 -0800 Subject: [PATCH] [Fix] VertexAI Pass through - fix regression that caused vertex ai passthroughs to stop working for router models (#19967) * fix(vertex_ai): replace custom model names with actual Vertex AI model names in passthrough URLs (#19948) When the passthrough URL already contains project and location, the code was skipping the deployment lookup and forwarding the URL as-is to Vertex AI. For custom model names like gcp/google/gemini-2.5-flash, Vertex AI returned 404 because it only knows the actual model name (gemini-2.5-flash). The fix makes the deployment lookup always run, so the custom model name gets replaced with the actual Vertex AI model name before forwarding. * add _resolve_vertex_model_from_router * fix: get_llm_provider * Potential fix for code scanning alert no. 4020: Clear-text logging of sensitive information Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --------- Co-authored-by: michelligabriele Co-authored-by: Copilot Autofix powered by AI <62310815+github-advanced-security[bot]@users.noreply.github.com> --- .../llm_passthrough_endpoints.py | 122 +++++++++++++++--- .../test_vertex_passthrough_load_balancing.py | 72 +++++++++++ 2 files changed, 175 insertions(+), 19 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index b079e16151..c11dcd79bc 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -16,6 +16,7 @@ from fastapi.responses import StreamingResponse from starlette.websockets import WebSocketState import litellm +from litellm import get_llm_provider from litellm._logging import verbose_proxy_logger from litellm.constants import ( ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS, @@ -1052,6 +1053,89 @@ async def bedrock_proxy_route( return received_value +def _resolve_vertex_model_from_router( + model_id: str, + llm_router: Optional[litellm.Router], + encoded_endpoint: str, + endpoint: str, + vertex_project: Optional[str], + vertex_location: Optional[str], +) -> Tuple[str, str, Optional[str], Optional[str]]: + """ + Resolve Vertex AI model configuration from router. + + Args: + model_id: The model ID extracted from the URL (e.g., "gcp/google/gemini-2.5-flash") + llm_router: The LiteLLM router instance + encoded_endpoint: The encoded endpoint path + endpoint: The original endpoint path + vertex_project: Current vertex project (may be from URL) + vertex_location: Current vertex location (may be from URL) + + Returns: + Tuple of (encoded_endpoint, endpoint, vertex_project, vertex_location) + with resolved values from router config + """ + if not llm_router: + return encoded_endpoint, endpoint, vertex_project, vertex_location + + try: + deployment = llm_router.get_available_deployment_for_pass_through(model=model_id) + if not deployment: + return encoded_endpoint, endpoint, vertex_project, vertex_location + + litellm_params = deployment.get("litellm_params", {}) + + # Always override with router config values (they take precedence over URL values) + config_vertex_project = litellm_params.get("vertex_project") + config_vertex_location = litellm_params.get("vertex_location") + if config_vertex_project: + vertex_project = config_vertex_project + if config_vertex_location: + vertex_location = config_vertex_location + + # Get the actual Vertex AI model name by stripping the provider prefix + # e.g., "vertex_ai/gemini-2.0-flash-exp" -> "gemini-2.0-flash-exp" + model_from_config = litellm_params.get("model", "") + if model_from_config: + + # get_llm_provider returns (model, custom_llm_provider, dynamic_api_key, api_base) + # For "vertex_ai/gemini-2.0-flash-exp" it returns: + # model="gemini-2.0-flash-exp", custom_llm_provider="vertex_ai" + actual_model, custom_llm_provider, _, _ = get_llm_provider( + model=model_from_config + ) + + # Log only non-sensitive information (model names and provider), never API keys or secrets. + safe_actual_model = actual_model + safe_custom_llm_provider = custom_llm_provider + verbose_proxy_logger.debug( + "get_llm_provider returned: actual_model=%s, custom_llm_provider=%s, model_id=%s", + safe_actual_model, + safe_custom_llm_provider, + model_id, + ) + + if actual_model and model_id != actual_model: + verbose_proxy_logger.debug( + "Resolved router model '%s' to '%s' (provider=%s) with project=%s, location=%s", + model_id, + actual_model, + custom_llm_provider, + vertex_project, + vertex_location, + ) + encoded_endpoint = encoded_endpoint.replace(model_id, actual_model) + endpoint = endpoint.replace(model_id, actual_model) + + except Exception as e: + verbose_proxy_logger.debug( + f"Error resolving vertex model from router for model {model_id}: {e}" + ) + + return encoded_endpoint, endpoint, vertex_project, vertex_location + + def _is_bedrock_agent_runtime_route(endpoint: str) -> bool: """ Return True, if the endpoint should be routed to the `bedrock-agent-runtime` endpoint. @@ -1512,8 +1596,11 @@ async def _prepare_vertex_auth_headers( if router_credentials is not None: vertex_credentials_str = None elif vertex_credentials is not None: - vertex_project = vertex_credentials.vertex_project - vertex_location = vertex_credentials.vertex_location + # Only override vertex_project and vertex_location if they're not already set from router config + if vertex_project is None: + vertex_project = vertex_credentials.vertex_project + if vertex_location is None: + vertex_location = vertex_credentials.vertex_location vertex_credentials_str = vertex_credentials.vertex_credentials else: raise ValueError("No vertex credentials found") @@ -1583,6 +1670,7 @@ async def _base_vertex_proxy_route( get_vertex_model_id_from_url, get_vertex_project_id_from_url, ) + from litellm.proxy.proxy_server import llm_router encoded_endpoint = httpx.URL(endpoint).path verbose_proxy_logger.debug("requested endpoint %s", endpoint) @@ -1610,24 +1698,20 @@ async def _base_vertex_proxy_route( vertex_location=vertex_location, ) - if vertex_project is None or vertex_location is None: - # Check if model is in router config - model_id = get_vertex_model_id_from_url(endpoint) - if model_id: - from litellm.proxy.proxy_server import llm_router + # Check if model is in router config - always do this to resolve custom model names + model_id = get_vertex_model_id_from_url(endpoint) + if model_id: - if llm_router: - try: - # Use the dedicated pass-through deployment selection method to automatically filter use_in_pass_through=True - deployment = llm_router.get_available_deployment_for_pass_through(model=model_id) - if deployment: - litellm_params = deployment.get("litellm_params", {}) - vertex_project = litellm_params.get("vertex_project") - vertex_location = litellm_params.get("vertex_location") - except Exception as e: - verbose_proxy_logger.debug( - f"Error getting available deployment for model {model_id}: {e}" - ) + if llm_router: + # Resolve model configuration from router + encoded_endpoint, endpoint, vertex_project, vertex_location = _resolve_vertex_model_from_router( + model_id=model_id, + llm_router=llm_router, + encoded_endpoint=encoded_endpoint, + endpoint=endpoint, + vertex_project=vertex_project, + vertex_location=vertex_location, + ) vertex_credentials = passthrough_endpoint_router.get_vertex_credentials( project_id=vertex_project, 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 28b3ba0a17..d2fdb157c8 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 @@ -447,3 +447,75 @@ def test_forward_headers_from_request_x_pass_prefix(): assert "x-pass-anthropic-beta" not in result assert "x-pass-custom-header" not in result + +@pytest.mark.asyncio +async def test_vertex_passthrough_custom_model_name_replaced_in_url(): + """ + Test that when a passthrough URL contains a custom model_name (e.g., gcp/google/gemini-3-pro), + the URL is rewritten to use the actual Vertex AI model name (e.g., gemini-3-pro) + before being forwarded to Vertex AI. + + This prevents 404 errors from Vertex AI when custom model names are used in the config. + + Config example: + model_name: gcp/google/gemini-3-pro + litellm_params: + model: vertex_ai/gemini-3-pro + vertex_project: "my-project" + vertex_location: "global" + use_in_pass_through: true + """ + mock_request = MagicMock() + mock_response = MagicMock() + mock_handler = MagicMock() + + # Deployment with custom model_name but real vertex model + mock_deployment = { + "litellm_params": { + "model": "vertex_ai/gemini-3-pro", + "vertex_project": "nv-gcpllmgwit-20250411173346", + "vertex_location": "global", + "use_in_pass_through": True, + } + } + mock_router = MagicMock() + mock_router.get_available_deployment_for_pass_through.return_value = mock_deployment + + # The URL contains project/location AND a custom model name with slashes + test_endpoint = "v1/projects/nv-gcpllmgwit-20250411173346/locations/global/publishers/google/models/gcp/google/gemini-3-pro:generateContent" + + with patch("litellm.proxy.proxy_server.llm_router", mock_router), \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router") as mock_pt_router, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints._prepare_vertex_auth_headers", new_callable=AsyncMock) as mock_prep_headers, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route") as mock_create_route, \ + patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth", new_callable=AsyncMock) as mock_auth: + + mock_pt_router.get_vertex_credentials.return_value = MagicMock() + mock_prep_headers.return_value = ({}, "https://global-aiplatform.googleapis.com", False, "nv-gcpllmgwit-20250411173346", "global") + mock_endpoint_func = AsyncMock() + mock_create_route.return_value = mock_endpoint_func + mock_auth.return_value = {} + + mock_handler.get_default_base_target_url.return_value = "https://global-aiplatform.googleapis.com" + + await _base_vertex_proxy_route( + endpoint=test_endpoint, + request=mock_request, + fastapi_response=mock_response, + get_vertex_pass_through_handler=mock_handler, + ) + + # Verify the router was called with the custom model name (extracted from URL) + mock_router.get_available_deployment_for_pass_through.assert_called_once_with( + model="gcp/google/gemini-3-pro" + ) + + # Verify the target URL passed to create_pass_through_route contains + # the REAL Vertex AI model name, not the custom one + create_route_call = mock_create_route.call_args + target_url = create_route_call.kwargs.get("target", "") + assert "gcp/google/gemini-3-pro" not in target_url, \ + f"Custom model name should have been replaced in target URL. Got: {target_url}" + assert "gemini-3-pro" in target_url, \ + f"Actual Vertex AI model name should be in target URL. Got: {target_url}" +