mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(proxy): support slashes in google generateContent model names (#19737)
* fix(proxy): support slashes in google route params * fix(proxy): extract google model ids with slashes * test(proxy): cover google model ids with slashes
This commit is contained in:
parent
51d744c9d0
commit
43642391c5
5 changed files with 60 additions and 9 deletions
|
|
@ -425,12 +425,12 @@ class LiteLLMRoutes(enum.Enum):
|
|||
]
|
||||
|
||||
google_routes = [
|
||||
"/v1beta/models/{model_name}:countTokens",
|
||||
"/v1beta/models/{model_name}:generateContent",
|
||||
"/v1beta/models/{model_name}:streamGenerateContent",
|
||||
"/models/{model_name}:countTokens",
|
||||
"/models/{model_name}:generateContent",
|
||||
"/models/{model_name}:streamGenerateContent",
|
||||
"/v1beta/models/{model_name:path}:countTokens",
|
||||
"/v1beta/models/{model_name:path}:generateContent",
|
||||
"/v1beta/models/{model_name:path}:streamGenerateContent",
|
||||
"/models/{model_name:path}:countTokens",
|
||||
"/models/{model_name:path}:generateContent",
|
||||
"/models/{model_name:path}:streamGenerateContent",
|
||||
# Google Interactions API
|
||||
"/interactions",
|
||||
"/v1beta/interactions",
|
||||
|
|
|
|||
|
|
@ -758,11 +758,27 @@ def get_model_from_request(
|
|||
if match:
|
||||
model = match.group(1)
|
||||
|
||||
# If still not found, extract model from Google generateContent-style routes.
|
||||
# These routes put the model in the path and allow "/" inside the model id.
|
||||
# Examples:
|
||||
# - /v1beta/models/gemini-2.0-flash:generateContent
|
||||
# - /v1beta/models/bedrock/claude-sonnet-3.7:generateContent
|
||||
# - /models/custom/ns/model:streamGenerateContent
|
||||
if model is None and not route.lower().startswith("/vertex"):
|
||||
google_match = re.search(r"/(?:v1beta|beta)/models/([^:]+):", route)
|
||||
if google_match:
|
||||
model = google_match.group(1)
|
||||
|
||||
if model is None and not route.lower().startswith("/vertex"):
|
||||
google_match = re.search(r"^/models/([^:]+):", route)
|
||||
if google_match:
|
||||
model = google_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 model is None and route.lower().startswith("/vertex"):
|
||||
vertex_match = re.search(r"/models/([^:]+)", route)
|
||||
if vertex_match:
|
||||
model = vertex_match.group(1)
|
||||
|
||||
|
|
|
|||
|
|
@ -392,7 +392,15 @@ class RouteChecks:
|
|||
# Ensure route is a string before attempting regex matching
|
||||
if not isinstance(route, str):
|
||||
return False
|
||||
pattern = re.sub(r"\{[^}]+\}", r"[^/]+", pattern)
|
||||
|
||||
def _placeholder_to_regex(match: re.Match) -> str:
|
||||
placeholder = match.group(0).strip("{}")
|
||||
if placeholder.endswith(":path"):
|
||||
# allow "/" in the placeholder value, but don't eat the route suffix after ":"
|
||||
return r"[^:]+"
|
||||
return r"[^/]+"
|
||||
|
||||
pattern = re.sub(r"\{[^}]+\}", _placeholder_to_regex, pattern)
|
||||
# Anchor the pattern to match the entire string
|
||||
pattern = f"^{pattern}$"
|
||||
if re.match(pattern, route):
|
||||
|
|
|
|||
|
|
@ -8,6 +8,7 @@ from litellm.proxy._types import UserAPIKeyAuth
|
|||
from litellm.proxy.auth.auth_utils import (
|
||||
_get_customer_id_from_standard_headers,
|
||||
get_end_user_id_from_request_body,
|
||||
get_model_from_request,
|
||||
get_key_model_rpm_limit,
|
||||
get_key_model_tpm_limit,
|
||||
)
|
||||
|
|
@ -186,3 +187,25 @@ class TestGetEndUserIdFromRequestBodyWithStandardHeaders:
|
|||
request_body=request_body, request_headers=headers
|
||||
)
|
||||
assert result == "body-user"
|
||||
|
||||
|
||||
def test_get_model_from_request_supports_google_model_names_with_slashes():
|
||||
assert (
|
||||
get_model_from_request(
|
||||
request_data={},
|
||||
route="/v1beta/models/bedrock/claude-sonnet-3.7:generateContent",
|
||||
)
|
||||
== "bedrock/claude-sonnet-3.7"
|
||||
)
|
||||
assert (
|
||||
get_model_from_request(
|
||||
request_data={},
|
||||
route="/models/hosted_vllm/gpt-oss-20b:generateContent",
|
||||
)
|
||||
== "hosted_vllm/gpt-oss-20b"
|
||||
)
|
||||
|
||||
|
||||
def test_get_model_from_request_vertex_passthrough_still_works():
|
||||
route = "/vertex_ai/v1/projects/p/locations/l/publishers/google/models/gemini-1.5-pro:generateContent"
|
||||
assert get_model_from_request(request_data={}, route=route) == "gemini-1.5-pro"
|
||||
|
|
|
|||
|
|
@ -161,9 +161,11 @@ def test_virtual_key_llm_api_route_includes_passthrough_prefix(route):
|
|||
[
|
||||
"/v1beta/models/gemini-2.5-flash:countTokens",
|
||||
"/v1beta/models/gemini-2.0-flash:generateContent",
|
||||
"/v1beta/models/bedrock/claude-sonnet-3.7:generateContent",
|
||||
"/v1beta/models/gemini-1.5-pro:streamGenerateContent",
|
||||
"/models/gemini-2.5-flash:countTokens",
|
||||
"/models/gemini-2.0-flash:generateContent",
|
||||
"/models/bedrock/claude-sonnet-3.7:generateContent",
|
||||
"/models/gemini-1.5-pro:streamGenerateContent",
|
||||
],
|
||||
)
|
||||
|
|
@ -187,9 +189,11 @@ def test_virtual_key_llm_api_routes_allows_google_routes(route):
|
|||
"/v1beta/models/google-gemini-2-5-pro-code-reviewer-k8s:generateContent",
|
||||
"/v1beta/models/gemini-2.5-flash-exp:countTokens",
|
||||
"/v1beta/models/custom-model-name-123:streamGenerateContent",
|
||||
"/v1beta/models/bedrock/claude-sonnet-3.7:generateContent",
|
||||
"/models/google-gemini-2-5-pro-code-reviewer-k8s:generateContent",
|
||||
"/models/gemini-2.5-flash-exp:countTokens",
|
||||
"/models/custom-model-name-123:streamGenerateContent",
|
||||
"/models/bedrock/claude-sonnet-3.7:generateContent",
|
||||
],
|
||||
)
|
||||
def test_google_routes_with_dynamic_model_names_recognized_as_llm_api_route(route):
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue