mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-28 06:23:07 +00:00
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.
This commit is contained in:
mateo-berri
parent
f0f046db12
commit
7ca83f9807
5 changed files with 122 additions and 112 deletions
|
|
@ -34564,6 +34564,11 @@
|
|||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Vertex Discovery Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
|
|
@ -34603,6 +34608,11 @@
|
|||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Vertex Discovery Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
|
|
@ -34642,6 +34652,11 @@
|
|||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Vertex Discovery Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
|
|
@ -34681,6 +34696,11 @@
|
|||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Vertex Discovery Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
|
|
@ -34720,6 +34740,11 @@
|
|||
"description": "Validation Error"
|
||||
}
|
||||
},
|
||||
"security": [
|
||||
{
|
||||
"APIKeyHeader": []
|
||||
}
|
||||
],
|
||||
"summary": "Vertex Discovery Proxy Route",
|
||||
"tags": [
|
||||
"llm_passthrough"
|
||||
|
|
|
|||
|
|
@ -105,7 +105,6 @@ from litellm.proxy.common_utils.resource_ownership import is_proxy_admin
|
|||
from litellm.proxy.common_utils.sse_keepalive import (
|
||||
wrap_passthrough_sse_bytes_with_keepalive_pings,
|
||||
)
|
||||
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,
|
||||
|
|
@ -2942,7 +2941,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,
|
||||
):
|
||||
"""
|
||||
|
|
@ -2958,7 +2957,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: The caller the route already authenticated, or None to authenticate here
|
||||
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 (
|
||||
|
|
@ -2972,15 +2971,10 @@ async def _base_vertex_proxy_route(
|
|||
encoded_endpoint = httpx.URL(endpoint).path
|
||||
verbose_proxy_logger.debug("requested endpoint %s", endpoint)
|
||||
headers: Mapping[str, str] = {}
|
||||
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 router_credentials is not None:
|
||||
await assert_user_can_access_vector_store(
|
||||
vector_store=router_credentials,
|
||||
user_api_key_dict=caller,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
vertex_project: str | None = get_vertex_project_id_from_url(endpoint)
|
||||
|
|
@ -3032,7 +3026,7 @@ async def _base_vertex_proxy_route(
|
|||
router_credentials=router_credentials,
|
||||
vertex_project=vertex_project,
|
||||
vertex_location=vertex_location,
|
||||
user_api_key_dict=caller,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
)
|
||||
|
||||
base_target_url: Final = get_vertex_pass_through_handler.get_default_base_target_url(vertex_location)
|
||||
|
|
@ -3078,7 +3072,7 @@ async def _base_vertex_proxy_route(
|
|||
received_value: Final = await endpoint_func(
|
||||
request,
|
||||
fastapi_response,
|
||||
caller,
|
||||
user_api_key_dict,
|
||||
)
|
||||
except ProxyException as e:
|
||||
if headers_passed_through:
|
||||
|
|
@ -3088,6 +3082,16 @@ 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 or "", custom_litellm_key_header=litellm_api_key
|
||||
)
|
||||
|
||||
|
||||
@router.api_route(
|
||||
"/vertex_ai/discovery/{endpoint:path}",
|
||||
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
|
||||
|
|
@ -3097,6 +3101,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.
|
||||
|
|
@ -3135,20 +3140,11 @@ 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,
|
||||
)
|
||||
|
||||
|
||||
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 or "", custom_litellm_key_header=litellm_api_key
|
||||
)
|
||||
|
||||
|
||||
@router.api_route(
|
||||
"/vertex-ai/{endpoint:path}",
|
||||
methods=["GET", "POST", "PUT", "DELETE", "PATCH"],
|
||||
|
|
|
|||
|
|
@ -324,12 +324,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,
|
||||
|
|
@ -340,8 +334,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()
|
||||
|
|
@ -426,12 +418,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,
|
||||
|
|
@ -442,8 +428,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()
|
||||
|
|
@ -1238,12 +1222,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,
|
||||
|
|
@ -1254,8 +1232,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()
|
||||
|
|
@ -1271,6 +1247,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
|
||||
|
|
@ -1285,36 +1262,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
|
||||
|
|
@ -4375,8 +4345,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(
|
||||
|
|
@ -8163,6 +8131,10 @@ class TestVertexPassthroughAuthenticatesOncePerRequest:
|
|||
"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,
|
||||
|
|
@ -8211,17 +8183,7 @@ class TestVertexPassthroughAuthenticatesOncePerRequest:
|
|||
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={
|
||||
"candidates": [
|
||||
{"content": {"role": "model", "parts": [{"text": "hello"}]}, "finishReason": "STOP"}
|
||||
],
|
||||
"usageMetadata": {"promptTokenCount": 1, "candidatesTokenCount": 1, "totalTokenCount": 2},
|
||||
},
|
||||
)
|
||||
)
|
||||
vertex: Final = upstream.post(vertex_url).mock(return_value=httpx.Response(200, json=self.GEMINI_ANSWER))
|
||||
|
||||
response, recorder = self._post(
|
||||
monkeypatch,
|
||||
|
|
@ -8258,6 +8220,51 @@ class TestVertexPassthroughAuthenticatesOncePerRequest:
|
|||
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:
|
||||
|
|
|
|||
|
|
@ -66,10 +66,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
|
||||
|
|
@ -83,7 +79,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(
|
||||
|
|
@ -92,6 +87,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
|
||||
|
|
@ -376,10 +372,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,
|
||||
|
|
@ -398,7 +390,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}",
|
||||
|
|
@ -406,6 +397,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)
|
||||
|
|
@ -826,11 +814,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()
|
||||
|
||||
|
|
@ -840,6 +823,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"
|
||||
|
|
|
|||
|
|
@ -441,19 +441,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
|
||||
|
||||
|
|
@ -541,6 +538,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