mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-11 08:38:24 +00:00
fix(vertex_ai): authenticate Vertex AI pass-through requests once (#48942)
* test(vertex_ai): run auth once per Vertex pass-through request Drive /vertex_ai, /vertex-ai and /vertex_ai/discovery requests through the real proxy app with a counting custom_auth, and assert each request authenticates once and is served as the caller that auth returned. The generate-content routes fail today because the route authenticates a second time. * 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. * fix(vertex): read the caller's key only from x-litellm-api-key or Authorization on Vertex routes The single auth check on /vertex_ai and /vertex-ai now reads the same two headers main's effective check did, so a key sent only in another key header gets 401 again instead of authenticating. * fix(vertex): authenticate the discovery route through the same Vertex key check Both Vertex pass-through routes now resolve the caller once through _authenticate_vertex_caller and hand it to _base_vertex_proxy_route, which no longer re-runs auth. x-litellm-api-key accepts the same scheme prefixes on /vertex_ai/discovery as on /vertex_ai and every other proxy route. * fix(proxy): size a Gemini request's budget hold from its prompt and maxOutputTokens A Gemini generateContent body's budget hold now counts its contents and systemInstruction and reads generationConfig.maxOutputTokens. Before, it held the model's whole context window, about $0.85 on gemini-3.8-flash, so a key with a small budget got 422s on back-to-back calls it could afford Also removes get_litellm_virtual_key, which nothing calls since the Vertex routes stopped using it * fix(proxy): count every Gemini system part and candidateCount in the budget hold A Gemini generateContent body with several systemInstruction text parts lost every part after the first when LiteLLM translated it to a chat request. That dropped the later system text for non-Gemini models behind /v1beta/models/<model>:generateContent, and made the budget hold count only the first part. The adapter now keeps every text part, as a plain string when there is one part and as text blocks when there are several. The output hold now also multiplies by generationConfig.candidateCount, the same way chat bodies use n, so a request asking for several candidates holds budget for all of them. * fix(proxy): keep the context-window hold for Gemini bodies the prompt count cannot see A Gemini generateContent body with media, cachedContent, other roles, or non-function tools keeps the full context-window hold, since the chat translation drops or undercounts those parts. The hold also reads the snake_case generation_config output fields, and several systemInstruction text parts reach non-Gemini models joined into one system string * fix(proxy): read only camelCase Gemini output fields in the budget hold /v1beta forwards only generationConfig.maxOutputTokens to non-Gemini models, so a snake_case cap the model never receives must not shrink the hold. Snake_case bodies now get the default output hold. * fix(proxy): count snake_case candidate_count in the Gemini budget hold * fix(vertex_ai): skip malformed system parts and read Gemini output fields only on Gemini bodies A systemInstruction part whose text is not a string is skipped, as main did, instead of failing the request. The budget hold reads generationConfig output and candidate fields, and counts contents, only when the body has no messages, prompt, input, query or documents field. An output cap of 0 or below counts as no cap. Test recorders keep tuples instead of mutating lists. * fix(proxy): take the budget hold's request shape from the route On generateContent routes the hold counts only contents and reads its output cap and multiplier from generationConfig, so stray fields of other APIs can no longer hide the prompt. A max_tokens of 0 holds no output again, while negative values and a Gemini maxOutputTokens of 0 or below hold the default output * fix(proxy): hold one candidate for a count too large to price and read system parts like content parts * fix(proxy): keep the spend already counted when a budget hold is far over the budget * fix(proxy): count candidateCount in the budget hold only for models that generate the candidates A non-Gemini model reached through /v1beta generateContent never receives candidateCount, so a huge count held the whole budget for the length of the call while the model produced one answer * fix(proxy): hold a candidate count only when every deployment generates it, and drop n when every deployment drops it * fix(proxy): count candidateCount when any deployment generates it and always count n The stricter rules from the previous commit let a key spend past max_budget through a Gemini name no deployment serves, an allowed_openai_params n, and a router fallback. Restore the earlier rules and add regression tests for those three shapes. * fix(proxy): cap the candidateCount hold at main's default output hold A candidateCount like 1e300 sent to a mixed Gemini and Claude group, or through a key alias from a Gemini name to Claude, held a whole key while Claude served the call and answered a strict key 422. On generateContent routes the counted candidates now raise the output hold at most to the default output hold main used for the model, so no count makes a call hold more than on main. Chat n is unchanged. * fix(proxy): hold every Gemini candidate up to the merge base's whole hold The previous cap stopped the candidateCount output hold at the default 16384 tokens. Because the hold now counts the Gemini prompt instead of the whole input window, a strict key admitted a call like 4 candidates of 8000 tokens and then spent past max_budget, where the merge base refused it. Every candidate is now held at its output cap, like chat n, and the whole hold is clamped at what the merge base held for the same call (the model's input window plus the default output), never below one candidate's hold. A huge count still holds no more than the merge base did. Chat n is unchanged * fix(proxy): let a Gemini output cap that /v1beta drops only raise the hold The Vertex and Gemini pass-throughs forward every spelling of the output cap, but /v1beta passes only generationConfig.maxOutputTokens to other models. The hold read only that spelling, so a snake_case cap above the default output held less than the pass-through could spend. A cap in any other spelling now raises the hold to it and never lowers it * fix(proxy): hold Gemini candidates of a model with no context window up to the default output A Gemini model the cost map has no context window for, priced in its deployment, held one candidate for any candidateCount, so parallel calls could pass their budget holds. Its candidates are held again up to its counted prompt plus the default output hold * fix(proxy): keep main's Gemini hold unless the request caps its output An uncapped generateContent call, or one capped only by a snake_case spelling, now holds what main holds (the context window plus the default output). A call with generationConfig.maxOutputTokens holds its counted prompt plus cap times candidateCount, never clamped to the deployment's max_output_tokens and never more than main's hold. A model with no context window keeps main's no-hold, and candidateCount always counts on generateContent routes. Also keeps get_litellm_virtual_key in pass_through_endpoints/common_utils.py, so an out-of-tree importer does not break * fix(proxy): keep main's budget hold on generateContent routes The PR now leaves the generateContent budget estimate as main has it. A live risk check showed the counted Gemini hold under-reserves (the prompt counter undercounts Gemini bodies, tiered and image output rates are not read, and request fields the handler merges later are not counted), so that estimate moves to a follow-up. The single auth pass on the Vertex pass-through and the systemInstruction adapter fix stay * fix(vertex_ai): keep main's Gemini adapter and pass a missing Authorization as no key The Gemini adapter change for later systemInstruction parts goes back to main, since it is not the double-auth fix this PR is for. The Vertex auth dependency now hands a missing Authorization header to the auth check as no key, like main's first auth pass, so custom auth and the 401 text see what they see on main --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
quietjules
GitHub
parent
5103d8a0ee
commit
210e2a075a
5 changed files with 315 additions and 127 deletions
|
|
@ -36405,6 +36405,11 @@
|
|||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Vertex Discovery Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
|
|
@ -36444,6 +36449,11 @@
|
|||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Vertex Discovery Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
|
|
@ -36483,6 +36493,11 @@
|
|||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Vertex Discovery Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
|
|
@ -36522,6 +36537,11 @@
|
|||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Vertex Discovery Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
|
|
@ -36561,6 +36581,11 @@
|
|||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Vertex Discovery Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
|
|
|
|||
|
|
@ -23,7 +23,7 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias
|
|||
|
||||
import httpx
|
||||
import openai
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response, WebSocket
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, Response, Security, WebSocket
|
||||
from fastapi.responses import StreamingResponse
|
||||
from pydantic import ConfigDict, TypeAdapter, ValidationError
|
||||
from starlette.websockets import WebSocketState
|
||||
|
|
@ -97,6 +97,8 @@ from litellm.proxy.auth.handle_jwt import JWTHandler
|
|||
from litellm.proxy.auth.route_checks import RouteChecks
|
||||
from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy module exports
|
||||
_get_bearer_token, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export
|
||||
api_key_header,
|
||||
custom_litellm_key_header,
|
||||
get_bearer_token,
|
||||
is_no_auth_dev_mode,
|
||||
user_api_key_auth,
|
||||
|
|
@ -133,7 +135,6 @@ from litellm.proxy.common_utils.token_count_log import (
|
|||
logged_provider_token_count,
|
||||
rejected_token_count,
|
||||
)
|
||||
from litellm.proxy.pass_through_endpoints.common_utils import get_litellm_virtual_key
|
||||
from litellm.proxy.pass_through_endpoints.injected_stream_usage import (
|
||||
stream_usage_relay,
|
||||
without_injected_usage_frames,
|
||||
|
|
@ -3097,7 +3098,7 @@ async def _base_vertex_proxy_route(
|
|||
fastapi_response: Response,
|
||||
get_vertex_pass_through_handler: BaseVertexAIPassThroughHandler,
|
||||
general_settings: Mapping[str, object],
|
||||
user_api_key_dict: UserAPIKeyAuth | None = None,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
router_credentials: LiteLLM_ManagedVectorStore | None = None,
|
||||
):
|
||||
"""
|
||||
|
|
@ -3113,7 +3114,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 authenticated
|
||||
router_credentials: Optional vector store credentials from registry (LiteLLM_ManagedVectorStore)
|
||||
"""
|
||||
from litellm.llms.vertex_ai.common_utils import (
|
||||
|
|
@ -3127,18 +3128,6 @@ 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,
|
||||
)
|
||||
|
||||
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,
|
||||
|
|
@ -3250,6 +3239,18 @@ async def _base_vertex_proxy_route(
|
|||
return received_value
|
||||
|
||||
|
||||
async def _authenticate_vertex_caller(
|
||||
request: Request,
|
||||
authorization: Annotated[str | None, Security(api_key_header)],
|
||||
litellm_api_key: Annotated[str | None, Security(custom_litellm_key_header)],
|
||||
) -> UserAPIKeyAuth:
|
||||
return await user_api_key_auth(
|
||||
request=request,
|
||||
api_key=authorization, # pyright: ignore[reportArgumentType] # its own Security default is None when unset
|
||||
custom_litellm_key_header=litellm_api_key,
|
||||
)
|
||||
|
||||
|
||||
@router.api_route(
|
||||
"/vertex_ai/discovery/{endpoint:path}",
|
||||
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
|
||||
|
|
@ -3259,6 +3260,7 @@ async def vertex_discovery_proxy_route(
|
|||
endpoint: str,
|
||||
request: Request,
|
||||
fastapi_response: Response,
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(_authenticate_vertex_caller)],
|
||||
):
|
||||
"""
|
||||
Call any vertex discovery endpoint using the proxy.
|
||||
|
|
@ -3297,6 +3299,7 @@ async def vertex_discovery_proxy_route(
|
|||
fastapi_response=fastapi_response,
|
||||
get_vertex_pass_through_handler=discovery_handler,
|
||||
general_settings=_proxy_general_settings(),
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
router_credentials=vector_store_credentials,
|
||||
)
|
||||
|
||||
|
|
@ -3316,7 +3319,7 @@ async def vertex_proxy_route(
|
|||
endpoint: str,
|
||||
request: Request,
|
||||
fastapi_response: Response,
|
||||
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
|
||||
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(_authenticate_vertex_caller)],
|
||||
):
|
||||
"""
|
||||
Call LiteLLM proxy via Vertex AI SDK.
|
||||
|
|
|
|||
|
|
@ -296,12 +296,6 @@ 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.get_litellm_virtual_key"
|
||||
) as mock_get_virtual_key,
|
||||
mock.patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth"
|
||||
) as mock_user_auth,
|
||||
mock.patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler"
|
||||
) as mock_get_handler,
|
||||
|
|
@ -312,8 +306,6 @@ class TestVertexAIPassThroughHandler:
|
|||
|
||||
# Setup mocks
|
||||
mock_load_auth.return_value = (mock_credentials, test_project)
|
||||
mock_get_virtual_key.return_value = "Bearer test-key"
|
||||
mock_user_auth.return_value = {"api_key": "test-key"}
|
||||
|
||||
# Mock the vertex handler
|
||||
mock_handler = Mock()
|
||||
|
|
@ -398,12 +390,6 @@ 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.get_litellm_virtual_key"
|
||||
) as mock_get_virtual_key,
|
||||
mock.patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth"
|
||||
) as mock_user_auth,
|
||||
mock.patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler"
|
||||
) as mock_get_handler,
|
||||
|
|
@ -414,8 +400,6 @@ class TestVertexAIPassThroughHandler:
|
|||
|
||||
# Setup mocks
|
||||
mock_load_auth.return_value = (mock_credentials, test_project)
|
||||
mock_get_virtual_key.return_value = "Bearer test-key"
|
||||
mock_user_auth.return_value = {"api_key": "test-key"}
|
||||
|
||||
# Mock the vertex handler for global location
|
||||
mock_handler = Mock()
|
||||
|
|
@ -500,17 +484,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()
|
||||
|
|
@ -528,6 +507,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()
|
||||
|
|
@ -602,20 +582,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
|
||||
|
|
@ -627,8 +603,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
|
||||
|
|
@ -637,26 +614,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):
|
||||
"""
|
||||
|
|
@ -1222,12 +1194,6 @@ class TestVertexAIDiscoveryPassThroughHandler:
|
|||
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.get_litellm_virtual_key"
|
||||
) as mock_get_virtual_key,
|
||||
mock.patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth"
|
||||
) as mock_user_auth,
|
||||
mock.patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.get_vertex_pass_through_handler"
|
||||
) as mock_get_handler,
|
||||
|
|
@ -1238,8 +1204,6 @@ class TestVertexAIDiscoveryPassThroughHandler:
|
|||
|
||||
# Setup mocks
|
||||
mock_load_auth.return_value = (mock_credentials, test_project)
|
||||
mock_get_virtual_key.return_value = "Bearer test-key"
|
||||
mock_user_auth.return_value = {"api_key": "test-key"}
|
||||
|
||||
# Mock the discovery handler
|
||||
mock_handler = Mock()
|
||||
|
|
@ -1255,6 +1219,7 @@ class TestVertexAIDiscoveryPassThroughHandler:
|
|||
endpoint=endpoint,
|
||||
request=mock_request,
|
||||
fastapi_response=mock_response,
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="test-key"),
|
||||
)
|
||||
|
||||
# Verify create_pass_through_route was called with correct arguments
|
||||
|
|
@ -1269,36 +1234,29 @@ class TestVertexAIDiscoveryPassThroughHandler:
|
|||
@pytest.mark.asyncio
|
||||
async def test_vertex_discovery_proxy_route_api_key_auth(self):
|
||||
"""
|
||||
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
|
||||
With no Vertex credential configured, the caller's virtual key must not be forwarded to
|
||||
Google, so the request fails with a clean 401 (LIT-5997).
|
||||
"""
|
||||
# Mock dependencies
|
||||
mock_request = Mock()
|
||||
mock_request.headers = {"x-litellm-api-key": "test-key-123"}
|
||||
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_discovery_proxy_route(
|
||||
endpoint="v1/projects/test-project/locations/us-central1/dataStores/default/servingConfigs/default:search",
|
||||
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_discovery_proxy_route(
|
||||
endpoint="v1/projects/test-project/locations/us-central1/dataStores/default/servingConfigs/default:search",
|
||||
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()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -4545,8 +4503,6 @@ class TestVertexRawPredictStreamingClassification:
|
|||
return_value=(mock_credentials, "test-project"),
|
||||
),
|
||||
mock.patch(f"{module}.create_pass_through_route", side_effect=fake_create_pass_through_route),
|
||||
mock.patch(f"{module}.get_litellm_virtual_key", return_value="Bearer test-key"),
|
||||
mock.patch(f"{module}.user_api_key_auth", new=AsyncMock(return_value=UserAPIKeyAuth(api_key="test-key"))),
|
||||
mock.patch(f"{module}.get_vertex_pass_through_handler", return_value=mock_handler),
|
||||
):
|
||||
await vertex_proxy_route(
|
||||
|
|
@ -8703,6 +8659,228 @@ class TestAzureSpeechRawBodyThroughRealAuth:
|
|||
assert "Invalid JSON payload" in response.text
|
||||
|
||||
|
||||
class _CountingCustomAuth:
|
||||
def __init__(self, only_accepts: str | None = None) -> None:
|
||||
self.only_accepts: Final = only_accepts
|
||||
self.api_keys_seen: tuple[str | None, ...] = ()
|
||||
|
||||
async def __call__(self, request: Request, api_key: str | None) -> UserAPIKeyAuth:
|
||||
self.api_keys_seen = (*self.api_keys_seen, api_key)
|
||||
if self.only_accepts is not None and api_key != self.only_accepts:
|
||||
raise ValueError("unknown caller")
|
||||
return UserAPIKeyAuth(user_id=f"custom-auth-caller-{len(self.api_keys_seen)}")
|
||||
|
||||
|
||||
class _ServedCallerRecorder(CustomLogger):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.served_user_ids: tuple[str | None, ...] = ()
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self, user_api_key_dict: UserAPIKeyAuth, cache: DualCache, data: dict, call_type: CallTypesLiteral
|
||||
) -> dict:
|
||||
self.served_user_ids = (*self.served_user_ids, user_api_key_dict.user_id)
|
||||
return data
|
||||
|
||||
|
||||
class TestVertexPassthroughAuthenticatesOncePerRequest:
|
||||
"""Each Vertex pass-through request runs the proxy's auth once and is served as the caller that auth resolved"""
|
||||
|
||||
PROJECT: Final = "vertex-proj"
|
||||
VIRTUAL_KEY: Final = "sk-vertex-caller"
|
||||
GENERATE_CONTENT: Final = (
|
||||
"v1/projects/vertex-proj/locations/us-central1/publishers/google/models/gemini-3.8-flash:generateContent"
|
||||
)
|
||||
DISCOVERY_SEARCH: Final = (
|
||||
"v1/projects/vertex-proj/locations/global/collections/default_collection/engines/support-engine"
|
||||
"/servingConfigs/default_search:search"
|
||||
)
|
||||
GEMINI_ANSWER: Final = {
|
||||
"candidates": [{"content": {"role": "model", "parts": [{"text": "hello"}]}, "finishReason": "STOP"}],
|
||||
"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2},
|
||||
}
|
||||
|
||||
def _post(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
path: str,
|
||||
headers: Mapping[str, str],
|
||||
custom_auth: _CountingCustomAuth,
|
||||
credential_locations: tuple[str, ...],
|
||||
) -> tuple[httpx.Response, _ServedCallerRecorder]:
|
||||
from litellm.proxy.pass_through_endpoints.passthrough_endpoint_router import (
|
||||
PassthroughEndpointRouter,
|
||||
)
|
||||
from litellm.proxy.proxy_server import app
|
||||
|
||||
router: Final = PassthroughEndpointRouter()
|
||||
for location in credential_locations:
|
||||
router.add_vertex_credentials(
|
||||
project_id=self.PROJECT, location=location, vertex_credentials="vertex-service-account"
|
||||
)
|
||||
recorder: Final = _ServedCallerRecorder()
|
||||
monkeypatch.setattr(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.passthrough_endpoint_router", router
|
||||
)
|
||||
monkeypatch.delitem(app.dependency_overrides, user_api_key_auth, raising=False)
|
||||
monkeypatch.setattr(litellm, "callbacks", [recorder])
|
||||
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
|
||||
litellm.in_memory_llm_clients_cache.flush_cache()
|
||||
with (
|
||||
patch.multiple( # test-quality-ok: the real user_api_key_auth reads proxy_server module globals (custom auth, master_key, caches) that have no injection seam
|
||||
"litellm.proxy.proxy_server", **{**_azure_speech_real_auth_attrs(), "user_custom_auth": custom_auth}
|
||||
),
|
||||
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(token="google-access-token"), self.PROJECT),
|
||||
),
|
||||
):
|
||||
response: Final = TestClient(app).post(
|
||||
path, json={"contents": [{"role": "user", "parts": [{"text": "hi"}]}]}, headers=dict(headers)
|
||||
)
|
||||
return response, recorder
|
||||
|
||||
@pytest.mark.parametrize("route_prefix", ["/vertex_ai", "/vertex-ai"])
|
||||
def test_generate_content_runs_custom_auth_once_and_is_served_as_its_caller(
|
||||
self, monkeypatch: pytest.MonkeyPatch, route_prefix: str
|
||||
) -> None:
|
||||
custom_auth: Final = _CountingCustomAuth()
|
||||
vertex_url: Final = f"https://us-central1-aiplatform.googleapis.com/{self.GENERATE_CONTENT}"
|
||||
with respx.mock(assert_all_called=True) as upstream:
|
||||
vertex: Final = upstream.post(vertex_url).mock(return_value=httpx.Response(200, json=self.GEMINI_ANSWER))
|
||||
|
||||
response, recorder = self._post(
|
||||
monkeypatch,
|
||||
f"{route_prefix}/{self.GENERATE_CONTENT}",
|
||||
{"Authorization": f"Bearer {self.VIRTUAL_KEY}"},
|
||||
custom_auth,
|
||||
("us-central1",),
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,)
|
||||
assert recorder.served_user_ids == ("custom-auth-caller-1",)
|
||||
assert vertex.call_count == 1
|
||||
|
||||
def test_discovery_search_runs_custom_auth_once_and_is_served_as_its_caller(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
custom_auth: Final = _CountingCustomAuth()
|
||||
with respx.mock(assert_all_called=True) as upstream:
|
||||
discovery: Final = upstream.post(f"https://discoveryengine.googleapis.com/{self.DISCOVERY_SEARCH}").mock(
|
||||
return_value=httpx.Response(200, json={"results": []})
|
||||
)
|
||||
|
||||
response, recorder = self._post(
|
||||
monkeypatch,
|
||||
f"/vertex_ai/discovery/{self.DISCOVERY_SEARCH}",
|
||||
{"Authorization": f"Bearer {self.VIRTUAL_KEY}"},
|
||||
custom_auth,
|
||||
("global",),
|
||||
)
|
||||
|
||||
assert (response.status_code, response.json()) == (200, {"results": []})
|
||||
assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,)
|
||||
assert recorder.served_user_ids == ("custom-auth-caller-1",)
|
||||
assert discovery.call_count == 1
|
||||
|
||||
@pytest.mark.parametrize("scheme", ["bearer", "Basic"])
|
||||
@pytest.mark.parametrize(
|
||||
"route, upstream_url, credential_location, answer",
|
||||
[
|
||||
(
|
||||
f"/vertex_ai/{GENERATE_CONTENT}",
|
||||
f"https://us-central1-aiplatform.googleapis.com/{GENERATE_CONTENT}",
|
||||
"us-central1",
|
||||
GEMINI_ANSWER,
|
||||
),
|
||||
(
|
||||
f"/vertex_ai/discovery/{DISCOVERY_SEARCH}",
|
||||
f"https://discoveryengine.googleapis.com/{DISCOVERY_SEARCH}",
|
||||
"global",
|
||||
{"results": []},
|
||||
),
|
||||
],
|
||||
ids=["generate_content", "discovery_search"],
|
||||
)
|
||||
def test_x_litellm_api_key_with_any_auth_scheme_authenticates_once_on_every_vertex_route(
|
||||
self,
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
scheme: str,
|
||||
route: str,
|
||||
upstream_url: str,
|
||||
credential_location: str,
|
||||
answer: Mapping[str, object],
|
||||
) -> None:
|
||||
custom_auth: Final = _CountingCustomAuth(only_accepts=self.VIRTUAL_KEY)
|
||||
with respx.mock(assert_all_called=True) as upstream:
|
||||
google: Final = upstream.post(upstream_url).mock(return_value=httpx.Response(200, json=answer))
|
||||
|
||||
response, recorder = self._post(
|
||||
monkeypatch,
|
||||
route,
|
||||
{"x-litellm-api-key": f"{scheme} {self.VIRTUAL_KEY}"},
|
||||
custom_auth,
|
||||
(credential_location,),
|
||||
)
|
||||
|
||||
assert response.status_code == 200, response.text
|
||||
assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,)
|
||||
assert recorder.served_user_ids == ("custom-auth-caller-1",)
|
||||
assert google.call_count == 1
|
||||
|
||||
def test_x_litellm_api_key_authenticates_once_and_is_never_forwarded_without_a_vertex_credential(
|
||||
self, monkeypatch: pytest.MonkeyPatch
|
||||
) -> None:
|
||||
custom_auth: Final = _CountingCustomAuth()
|
||||
with respx.mock(assert_all_called=False) as upstream:
|
||||
google: Final = upstream.route().mock(return_value=httpx.Response(200))
|
||||
|
||||
response, _ = self._post(
|
||||
monkeypatch,
|
||||
f"/vertex_ai/{self.GENERATE_CONTENT}",
|
||||
{"x-litellm-api-key": self.VIRTUAL_KEY},
|
||||
custom_auth,
|
||||
(),
|
||||
)
|
||||
|
||||
assert response.status_code == 401, response.text
|
||||
assert custom_auth.api_keys_seen == (self.VIRTUAL_KEY,)
|
||||
assert not google.called
|
||||
|
||||
@pytest.mark.parametrize("header", ["x-goog-api-key", "x-api-key", "API-Key", "Ocp-Apim-Subscription-Key"])
|
||||
def test_key_outside_litellm_headers_is_rejected_and_never_reaches_google_without_a_vertex_credential(
|
||||
self, monkeypatch: pytest.MonkeyPatch, header: str
|
||||
) -> None:
|
||||
custom_auth: Final = _CountingCustomAuth(only_accepts=self.VIRTUAL_KEY)
|
||||
with respx.mock(assert_all_called=False) as upstream:
|
||||
google: Final = upstream.route().mock(return_value=httpx.Response(200))
|
||||
|
||||
response, recorder = self._post(
|
||||
monkeypatch, f"/vertex_ai/{self.GENERATE_CONTENT}", {header: self.VIRTUAL_KEY}, custom_auth, ()
|
||||
)
|
||||
|
||||
assert response.status_code == 401, response.text
|
||||
assert self.VIRTUAL_KEY not in custom_auth.api_keys_seen
|
||||
assert recorder.served_user_ids == ()
|
||||
assert not google.called
|
||||
|
||||
|
||||
@pytest.mark.parametrize("path", [f"/vertex_ai/{GENERATE_CONTENT}", f"/vertex_ai/discovery/{DISCOVERY_SEARCH}"])
|
||||
def test_request_without_a_key_hands_custom_auth_none(
|
||||
self, monkeypatch: pytest.MonkeyPatch, path: str
|
||||
) -> None:
|
||||
custom_auth: Final = _CountingCustomAuth(only_accepts=self.VIRTUAL_KEY)
|
||||
with respx.mock(assert_all_called=False) as upstream:
|
||||
google: Final = upstream.route().mock(return_value=httpx.Response(200))
|
||||
|
||||
response, recorder = self._post(monkeypatch, path, {}, custom_auth, ("us-central1", "global"))
|
||||
|
||||
assert response.status_code == 401, response.text
|
||||
assert custom_auth.api_keys_seen == (None,)
|
||||
assert recorder.served_user_ids == ()
|
||||
assert not google.called
|
||||
|
||||
class TestTypeSafePassthroughRoute:
|
||||
@staticmethod
|
||||
def _request(body: object, query_params: Mapping[str, str] | None = None) -> MagicMock:
|
||||
|
|
|
|||
|
|
@ -67,10 +67,6 @@ async def test_vertex_passthrough_load_balancing():
|
|||
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,
|
||||
):
|
||||
|
||||
# Setup additional mocks to avoid side effects
|
||||
|
|
@ -84,7 +80,6 @@ async def test_vertex_passthrough_load_balancing():
|
|||
|
||||
mock_endpoint_func = AsyncMock()
|
||||
mock_create_route.return_value = mock_endpoint_func
|
||||
mock_auth.return_value = {}
|
||||
|
||||
# Execute
|
||||
await _base_vertex_proxy_route(
|
||||
|
|
@ -93,6 +88,7 @@ async def test_vertex_passthrough_load_balancing():
|
|||
fastapi_response=mock_response,
|
||||
get_vertex_pass_through_handler=mock_handler,
|
||||
general_settings={},
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
# Verify
|
||||
|
|
@ -375,10 +371,6 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
|
|||
patch( # test-quality-ok: the upstream call is captured here, the route offers no injection point
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route"
|
||||
) as mock_create_route,
|
||||
patch( # test-quality-ok: the route calls auth directly rather than through Depends
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new_callable=AsyncMock,
|
||||
) as mock_auth,
|
||||
patch( # test-quality-ok: the route reads the request body for this, a MagicMock request has none
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.is_streaming_request_fn",
|
||||
new_callable=AsyncMock,
|
||||
|
|
@ -397,7 +389,6 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
|
|||
"global",
|
||||
)
|
||||
mock_create_route.return_value = AsyncMock()
|
||||
mock_auth.return_value = UserAPIKeyAuth(api_key="sk-litellm-secret-key")
|
||||
|
||||
await _base_vertex_proxy_route(
|
||||
endpoint=f"{VERTEX_ANTHROPIC_MODELS_PREFIX}{model_segment}",
|
||||
|
|
@ -405,6 +396,7 @@ async def test_vertex_passthrough_drops_anthropic_beta_only_on_count_tokens(
|
|||
fastapi_response=MagicMock(),
|
||||
get_vertex_pass_through_handler=VertexAIPassThroughHandler(),
|
||||
general_settings={},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-litellm-secret-key"),
|
||||
)
|
||||
|
||||
upstream_headers = mock_create_route.call_args.kwargs["custom_headers"]
|
||||
|
|
@ -723,10 +715,6 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url():
|
|||
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()
|
||||
|
|
@ -738,7 +726,6 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url():
|
|||
)
|
||||
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"
|
||||
|
|
@ -750,6 +737,7 @@ async def test_vertex_passthrough_custom_model_name_replaced_in_url():
|
|||
fastapi_response=mock_response,
|
||||
get_vertex_pass_through_handler=mock_handler,
|
||||
general_settings={},
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
# Verify the router was called with the custom model name (extracted from URL)
|
||||
|
|
@ -825,11 +813,6 @@ async def test_vertex_passthrough_attributes_the_call_to_the_resolved_deployment
|
|||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.create_pass_through_route",
|
||||
return_value=relay_returning_logging_kwargs,
|
||||
),
|
||||
patch( # test-quality-ok: the route calls auth directly rather than through Depends
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new_callable=AsyncMock,
|
||||
return_value=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
),
|
||||
):
|
||||
mock_pt_router.get_vertex_credentials.return_value = MagicMock()
|
||||
|
||||
|
|
@ -839,6 +822,7 @@ async def test_vertex_passthrough_attributes_the_call_to_the_resolved_deployment
|
|||
fastapi_response=MagicMock(),
|
||||
get_vertex_pass_through_handler=mock_handler,
|
||||
general_settings={},
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"),
|
||||
)
|
||||
|
||||
assert logging_kwargs["litellm_params"]["metadata"]["model_info"]["id"] == "vertex-gemini-38-flash-dep"
|
||||
|
|
|
|||
|
|
@ -442,19 +442,16 @@ async def test_vertex_discovery_denies_other_team_vector_store_credentials():
|
|||
"team_id": "team-b",
|
||||
}
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints.user_api_key_auth",
|
||||
new=AsyncMock(return_value=UserAPIKeyAuth(team_id="team-a")),
|
||||
):
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _base_vertex_proxy_route(
|
||||
endpoint="projects/p/locations/us-central1/dataStores/vs_other_team",
|
||||
request=request,
|
||||
fastapi_response=Response(),
|
||||
get_vertex_pass_through_handler=MagicMock(),
|
||||
general_settings={},
|
||||
router_credentials=vector_store_credentials,
|
||||
)
|
||||
with pytest.raises(HTTPException) as exc_info:
|
||||
await _base_vertex_proxy_route(
|
||||
endpoint="projects/p/locations/us-central1/dataStores/vs_other_team",
|
||||
request=request,
|
||||
fastapi_response=Response(),
|
||||
get_vertex_pass_through_handler=MagicMock(),
|
||||
general_settings={},
|
||||
user_api_key_dict=UserAPIKeyAuth(team_id="team-a"),
|
||||
router_credentials=vector_store_credentials,
|
||||
)
|
||||
|
||||
assert exc_info.value.status_code == 403
|
||||
|
||||
|
|
@ -542,6 +539,7 @@ async def test_vertex_discovery_allows_unregistered_provider_native_datastore_id
|
|||
endpoint="projects/p/locations/us-central1/dataStores/vs_unknown",
|
||||
request=request,
|
||||
fastapi_response=Response(),
|
||||
user_api_key_dict=UserAPIKeyAuth(),
|
||||
)
|
||||
|
||||
assert response == {"ok": True}
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue