From 2d5c7f809d5b451443818c02cd1036b8f516f67d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Sat, 22 Mar 2025 11:49:51 -0700 Subject: [PATCH] fix(llm_passthrough_endpoints.py): raise verbose error if credentials not found on proxy --- .../llm_passthrough_endpoints.py | 25 ++++-- .../test_llm_pass_through_endpoints.py | 85 ++++++++++++++++++- 2 files changed, 100 insertions(+), 10 deletions(-) diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 0fae1e6f0b..c4d96b67f6 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -463,9 +463,11 @@ async def vertex_proxy_route( location=vertex_location, ) + headers_passed_through = False # Use headers from the incoming request if no vertex credentials are found if vertex_credentials is None or vertex_credentials.vertex_project is None: headers = dict(request.headers) or {} + headers_passed_through = True verbose_proxy_logger.debug( "default_vertex_config not set, incoming request headers %s", headers ) @@ -516,8 +518,6 @@ async def vertex_proxy_route( vertex_location=vertex_location, vertex_project=vertex_project, ) - # base_url = httpx.URL(base_target_url) - # updated_url = base_url.copy_with(path=encoded_endpoint) verbose_proxy_logger.debug("updated url %s", updated_url) @@ -534,12 +534,21 @@ async def vertex_proxy_route( target=target, custom_headers=headers, ) # dynamically construct pass-through endpoint based on incoming path - received_value = await endpoint_func( - request, - fastapi_response, - user_api_key_dict, - stream=is_streaming_request, # type: ignore - ) + + try: + received_value = await endpoint_func( + request, + fastapi_response, + user_api_key_dict, + stream=is_streaming_request, # type: ignore + ) + except Exception as e: + if headers_passed_through: + raise Exception( + f"No credentials found on proxy for this request. Headers were passed through directly but request failed with error: {str(e)}" + ) + else: + raise e return received_value diff --git a/tests/litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 6283c8ebce..74a3dd45c8 100644 --- a/tests/litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/litellm/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -188,7 +188,7 @@ class TestVertexAIPassThroughHandler: Case 2: User set default credentials, no exact passthrough credentials - confirm default credentials used. - Case 3: No default credentials, incorrect project/base passed - confirm no credentials used. + Case 3: No default credentials, no mapped credentials - request passed through directly. """ @pytest.mark.asyncio @@ -225,7 +225,10 @@ class TestVertexAIPassThroughHandler: "type": "http", "method": "POST", "path": endpoint, - "headers": {}, + "headers": [ + (b"Authorization", b"Bearer test-creds"), + (b"Content-Type", b"application/json"), + ], } ) @@ -338,3 +341,81 @@ class TestVertexAIPassThroughHandler: target=f"https://{default_location}-aiplatform.googleapis.com/v1/projects/{default_project}/locations/{default_location}/publishers/google/models/gemini-1.5-flash:generateContent", custom_headers={"Authorization": f"Bearer {default_credentials}"}, ) + + @pytest.mark.asyncio + async def test_vertex_passthrough_with_no_default_credentials(self, monkeypatch): + """ + Test that when no default credentials are set, the request fails + """ + """ + Test that when passthrough credentials are set, they are correctly used in the request + """ + from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import ( + PassthroughEndpointRouter, + ) + + vertex_project = "my-project" + vertex_location = "us-central1" + vertex_credentials = "test-creds" + + test_project = "test-project" + test_location = "test-location" + test_token = "test-creds" + + pass_through_router = PassthroughEndpointRouter() + + pass_through_router.add_vertex_credentials( + project_id=vertex_project, + location=vertex_location, + vertex_credentials=vertex_credentials, + ) + + monkeypatch.setattr( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router", + pass_through_router, + ) + + endpoint = f"/v1/projects/{test_project}/locations/{test_location}/publishers/google/models/gemini-1.5-flash:generateContent" + + # Mock request + mock_request = Request( + scope={ + "type": "http", + "method": "POST", + "path": endpoint, + "headers": [ + (b"authorization", b"Bearer test-creds"), + ], + } + ) + + # Mock response + mock_response = Response() + + with mock.patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.vertex_llm_base._ensure_access_token_async" + ) as mock_ensure_token, mock.patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.vertex_llm_base._get_token_and_url" + ) as mock_get_token, mock.patch( + "litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route" + ) as mock_create_route: + mock_ensure_token.return_value = ("test-auth-header", test_project) + mock_get_token.return_value = (test_token, "") + + # Call the route + try: + await vertex_proxy_route( + endpoint=endpoint, + request=mock_request, + fastapi_response=mock_response, + ) + except Exception as e: + traceback.print_exc() + print(f"Error: {e}") + + # Verify create_pass_through_route was called with correct arguments + mock_create_route.assert_called_once_with( + endpoint=endpoint, + target=f"https://{test_location}-aiplatform.googleapis.com/v1/projects/{test_project}/locations/{test_location}/publishers/google/models/gemini-1.5-flash:generateContent", + custom_headers={"authorization": f"Bearer {test_token}"}, + )