fix(proxy): extract model from vertex ai passthrough url pattern (#18097)

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
This commit is contained in:
Kris Xia
2025-12-17 11:10:14 +05:30
committed by GitHub
parent 089d1eb08b
commit be2f429087
2 changed files with 61 additions and 0 deletions
+8
View File
@@ -616,6 +616,14 @@ def get_model_from_request(
if match:
model = match.group(1)
# If still not found, extract from Vertex AI passthrough route
# Pattern: /vertex_ai/.../models/{model_id}:*
# Example: /vertex_ai/v1/.../models/gemini-1.5-pro:generateContent
if model is None and "/vertex" in route.lower():
vertex_match = re.search(r"/models/([^/:]+)", route)
if vertex_match:
model = vertex_match.group(1)
return model
+53
View File
@@ -311,3 +311,56 @@ def test_get_internal_user_header_from_mapping_no_internal_returns_none():
single_mapping = {"header_name": "X-Only-Customer", "litellm_user_role": "customer"}
result = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(single_mapping)
assert result is None
@pytest.mark.parametrize(
"request_data, route, expected_model",
[
# Vertex AI passthrough URL patterns
(
{},
"/vertex_ai/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent",
"gemini-1.5-pro"
),
(
{},
"/vertex_ai/v1beta1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.0-pro:streamGenerateContent",
"gemini-1.0-pro"
),
(
{},
"/vertex_ai/v1/projects/my-project/locations/asia-southeast1/publishers/google/models/gemini-2.0-flash:generateContent",
"gemini-2.0-flash"
),
# Model without method suffix (no colon) - should still extract
(
{},
"/vertex_ai/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-pro",
"gemini-pro" # Should match even without colon
),
# Request body model takes precedence over URL
(
{"model": "gpt-4o"},
"/vertex_ai/v1/projects/my-project/locations/us-central1/publishers/google/models/gemini-1.5-pro:generateContent",
"gpt-4o"
),
# Non-vertex route should not extract from vertex pattern
(
{},
"/openai/v1/chat/completions",
None
),
# Azure deployment pattern should still work
(
{},
"/openai/deployments/my-deployment/chat/completions",
"my-deployment"
),
],
)
def test_get_model_from_request_vertex_ai_passthrough(request_data, route, expected_model):
"""Test that get_model_from_request correctly extracts Vertex AI model from URL"""
from litellm.proxy.auth.auth_utils import get_model_from_request
model = get_model_from_request(request_data, route)
assert model == expected_model