fix(llm_passthrough_endpoints.py): raise verbose error if credentials not found on proxy

This commit is contained in:
Krrish Dholakia 2025-03-22 11:49:51 -07:00
parent ecf2243b69
commit 2d5c7f809d
2 changed files with 100 additions and 10 deletions

View file

@ -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
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

View file

@ -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}"},
)