Skip to content

Commit da27392

Browse files
committed
fix(proxy): extract model from vertex ai passthrough url pattern
extract model id from vertex ai passthrough routes that follow the pattern: /vertex_ai/*/models/{model_id}:* the model extraction now handles vertex ai routes by regex matching the model segment from the url path, which allows proper model identification for authentication and authorization in proxy pass-through endpoints. adds comprehensive test coverage for vertex ai model extraction including: - various vertex api versions (v1, v1beta1) - different locations (us-central1, asia-southeast1) - model names with special suffixes (gemini-1.5-pro, gemini-2.0-flash) - precedence verification (request body model over url) - non-vertex route isolation
1 parent 5e7bb0c commit da27392

File tree

3 files changed

+446
-0
lines changed

3 files changed

+446
-0
lines changed

litellm/proxy/auth/auth_utils.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -616,6 +616,14 @@ def get_model_from_request(
616616
if match:
617617
model = match.group(1)
618618

619+
# If still not found, extract from Vertex AI passthrough route
620+
# Pattern: /vertex_ai/.../models/{model_id}:*
621+
# Example: /vertex_ai/v1/.../models/gemini-1.5-pro:generateContent
622+
if model is None and "/vertex" in route.lower():
623+
vertex_match = re.search(r"/models/([^/:]+)", route)
624+
if vertex_match:
625+
model = vertex_match.group(1)
626+
619627
return model
620628

621629

tests/local_testing/test_auth_utils.py

Lines changed: 53 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -311,3 +311,56 @@ def test_get_internal_user_header_from_mapping_no_internal_returns_none():
311311
single_mapping = {"header_name": "X-Only-Customer", "litellm_user_role": "customer"}
312312
result = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(single_mapping)
313313
assert result is None
314+
315+
316+
@pytest.mark.parametrize(
317+
"request_data, route, expected_model",
318+
[
319+
# Vertex AI passthrough URL patterns
320+
(
321+
{},
322+
"/vertex_ai/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent",
323+
"gemini-1.5-pro"
324+
),
325+
(
326+
{},
327+
"/vertex_ai/v1beta1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.0-pro:streamGenerateContent",
328+
"gemini-1.0-pro"
329+
),
330+
(
331+
{},
332+
"/vertex_ai/v1/projects/my-project/locations/asia-southeast1/publishers/google/models/gemini-2.0-flash:generateContent",
333+
"gemini-2.0-flash"
334+
),
335+
# Model without method suffix (no colon) - should still extract
336+
(
337+
{},
338+
"/vertex_ai/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-pro",
339+
"gemini-pro" # Should match even without colon
340+
),
341+
# Request body model takes precedence over URL
342+
(
343+
{"model": "gpt-4o"},
344+
"/vertex_ai/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent",
345+
"gpt-4o"
346+
),
347+
# Non-vertex route should not extract from vertex pattern
348+
(
349+
{},
350+
"/openai/v1/chat/completions",
351+
None
352+
),
353+
# Azure deployment pattern should still work
354+
(
355+
{},
356+
"/openai/deployments/my-deployment/chat/completions",
357+
"my-deployment"
358+
),
359+
],
360+
)
361+
def test_get_model_from_request_vertex_ai_passthrough(request_data, route, expected_model):
362+
"""Test that get_model_from_request correctly extracts Vertex AI model from URL"""
363+
from litellm.proxy.auth.auth_utils import get_model_from_request
364+
365+
model = get_model_from_request(request_data, route)
366+
assert model == expected_model

0 commit comments

Comments
 (0)