mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-14 09:03:16 +00:00
Merge 02723a28a6 into 9ffb82e139
This commit is contained in:
quietjules
GitHub
commit
0be6313b22
5 changed files with 315 additions and 127 deletions
|
|
@ -36373,6 +36373,11 @@
|
|||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Vertex Discovery Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
|
|
@ -36412,6 +36417,11 @@
|
|||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Vertex Discovery Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
|
|
@ -36451,6 +36461,11 @@
|
|||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Vertex Discovery Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
|
|
@ -36490,6 +36505,11 @@
|
|||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Vertex Discovery Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
|
|
@ -36529,6 +36549,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,
|
||||
|
|
@ -3093,7 +3094,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,
|
||||
):
|
||||
"""
|
||||
|
|
@ -3109,7 +3110,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 (
|
||||
|
|
@ -3123,18 +3124,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,
|
||||
|
|
@ -3246,6 +3235,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"],
|
||||
|
|
@ -3255,6 +3256,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.
|
||||
|
|
@ -3293,6 +3295,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,
|
||||
)
|
||||
|
||||
|
|
@ -3312,7 +3315,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