mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
fix(vertex_ai): authenticate Vertex pass-through requests once
vertex_proxy_route already resolves the caller through Depends(user_api_key_auth), but _base_vertex_proxy_route ignored it and authenticated the request again. A custom_auth function ran twice per request, and a key with max_budget reserved its estimated cost twice, so the unreleased first reservation exhausted the budget after one or two calls. The base route now uses the caller it is given and authenticates once itself only when none is passed, which is the discovery route's case.
This commit is contained in:
mateo-berri
parent
2b9ed94aa3
commit
60e16040fa
2 changed files with 26 additions and 45 deletions
|
|
@ -2956,7 +2956,7 @@ async def _base_vertex_proxy_route(
|
|||
fastapi_response: FastAPI response object
|
||||
get_vertex_pass_through_handler: Handler for the specific Vertex AI service
|
||||
general_settings: Proxy general settings, read per request for the anthropic-beta values to drop
|
||||
user_api_key_dict: User API key authentication dict
|
||||
user_api_key_dict: The caller the route already authenticated, or None to authenticate here
|
||||
router_credentials: Optional vector store credentials from registry (LiteLLM_ManagedVectorStore)
|
||||
"""
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
|
|
@ -2970,22 +2970,15 @@ async def _base_vertex_proxy_route(
|
|||
encoded_endpoint = httpx.URL(endpoint).path
|
||||
verbose_proxy_logger.debug("requested endpoint %s", endpoint)
|
||||
headers: Mapping[str, str] = {}
|
||||
api_key_to_use = get_litellm_virtual_key(request=request)
|
||||
user_api_key_dict = await user_api_key_auth(
|
||||
request=request,
|
||||
api_key=api_key_to_use,
|
||||
caller: Final = (
|
||||
user_api_key_dict
|
||||
if user_api_key_dict is not None
|
||||
else await user_api_key_auth(request=request, api_key=get_litellm_virtual_key(request=request))
|
||||
)
|
||||
|
||||
if user_api_key_dict is None:
|
||||
api_key_to_use = get_litellm_virtual_key(request=request)
|
||||
user_api_key_dict = await user_api_key_auth(
|
||||
request=request,
|
||||
api_key=api_key_to_use,
|
||||
)
|
||||
if router_credentials is not None:
|
||||
await assert_user_can_access_vector_store(
|
||||
vector_store=router_credentials,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
user_api_key_dict=caller,
|
||||
)
|
||||
|
||||
vertex_project: str | None = get_vertex_project_id_from_url(endpoint)
|
||||
|
|
@ -3037,7 +3030,7 @@ async def _base_vertex_proxy_route(
|
|||
router_credentials=router_credentials,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
user_api_key_dict=caller,
|
||||
)
|
||||
|
||||
base_target_url: Final = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
|
||||
|
|
@ -3083,7 +3076,7 @@ async def _base_vertex_proxy_route(
|
|||
received_value: Final = await endpoint_func(
|
||||
request,
|
||||
fastapi_response,
|
||||
user_api_key_dict,
|
||||
caller,
|
||||
)
|
||||
except ProxyException as e:
|
||||
if headers_passed_through:
|
||||
|
|
|
|||
|
|
@ -528,17 +528,12 @@ class TestVertexAIPassThroughHandler:
|
|||
mock.patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler"
|
||||
) as mock_get_handler,
|
||||
mock.patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_auth,
|
||||
):
|
||||
# Mock credentials object with necessary attributes
|
||||
mock_credentials = Mock()
|
||||
mock_credentials.token = default_credentials
|
||||
|
||||
mock_load_auth.return_value = (mock_credentials, default_project)
|
||||
mock_auth.return_value = MagicMock()
|
||||
|
||||
# Mock the vertex handler
|
||||
mock_handler = Mock()
|
||||
|
|
@ -556,6 +551,7 @@ class TestVertexAIPassThroughHandler:
|
|||
endpoint=endpoint,
|
||||
request=mock_request,
|
||||
fastapi_response=mock_response,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
|
||||
)
|
||||
except Exception as e:
|
||||
traceback.print_exc()
|
||||
|
|
@ -630,20 +626,16 @@ class TestVertexAIPassThroughHandler:
|
|||
mock.patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
|
||||
) as mock_create_route,
|
||||
mock.patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_auth,
|
||||
):
|
||||
mock_ensure_token.return_value = ("test-auth-header", test_project)
|
||||
mock_get_token.return_value = (test_token, "")
|
||||
mock_auth.return_value = UserAPIKeyAuth(api_key="sk-test-creds")
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await vertex_proxy_route(
|
||||
endpoint=endpoint,
|
||||
request=mock_request,
|
||||
fastapi_response=mock_response,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test-creds"),
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
|
|
@ -655,8 +647,9 @@ class TestVertexAIPassThroughHandler:
|
|||
Critical
|
||||
|
||||
This is how Vertex AI JS SDK will Auth to Litellm Proxy: the virtual key
|
||||
arrives in x-litellm-api-key and must reach user_api_key_auth. With no Vertex
|
||||
credential configured, that virtual key must not be forwarded to Google, so
|
||||
arrives in x-litellm-api-key and authenticates the caller (covered through the
|
||||
real auth dependency in TestVertexPassthroughAuthenticatesOncePerRequest). With no
|
||||
Vertex credential configured, that virtual key must not be forwarded to Google, so
|
||||
the request fails with a clean 401 (LIT-5997).
|
||||
"""
|
||||
# Mock dependencies
|
||||
|
|
@ -665,26 +658,21 @@ class TestVertexAIPassThroughHandler:
|
|||
mock_request.method = "POST"
|
||||
mock_response = Mock()
|
||||
|
||||
with patch("litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth") as mock_auth:
|
||||
mock_auth.return_value = {"api_key": "test-key-123"}
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
|
||||
) as mock_pass_through:
|
||||
mock_pass_through.return_value = AsyncMock(return_value={"status": "success"})
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
|
||||
) as mock_pass_through:
|
||||
mock_pass_through.return_value = AsyncMock(return_value={"status": "success"})
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await vertex_proxy_route(
|
||||
endpoint="v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent",
|
||||
request=mock_request,
|
||||
fastapi_response=mock_response,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test-key-123"),
|
||||
)
|
||||
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await vertex_proxy_route(
|
||||
endpoint="v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent",
|
||||
request=mock_request,
|
||||
fastapi_response=mock_response,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 401
|
||||
mock_pass_through.assert_not_called()
|
||||
mock_auth.assert_called_once()
|
||||
call_args = mock_auth.call_args[1]
|
||||
assert call_args["api_key"] == "Bearer test-key-123"
|
||||
assert exc_info.value.status_code == 401
|
||||
mock_pass_through.assert_not_called()
|
||||
|
||||
def test_vertex_passthrough_handler_multimodal_embedding_response(self):
|
||||
"""
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue