Add generate content in llm route

This commit is contained in:
Sameer Kankute
2025-12-24 12:16:24 +05:30
parent ecd628b4ab
commit 4c4c9cb353
2 changed files with 87 additions and 1 deletions
+19 -1
View File
@@ -293,6 +293,9 @@ class RouteChecks:
if route in LiteLLMRoutes.anthropic_routes.value:
return True
if route in LiteLLMRoutes.google_routes.value:
return True
if RouteChecks.check_route_access(
route=route, allowed_routes=LiteLLMRoutes.mcp_routes.value
@@ -315,13 +318,28 @@ class RouteChecks:
):
return True
# Check for Google routes with placeholders like "/v1beta/models/{model_name}:generateContent"
for google_route in LiteLLMRoutes.google_routes.value:
if "{" in google_route:
if RouteChecks._route_matches_pattern(
route=route, pattern=google_route
):
return True
# Check for Anthropic routes with placeholders
for anthropic_route in LiteLLMRoutes.anthropic_routes.value:
if "{" in anthropic_route:
if RouteChecks._route_matches_pattern(
route=route, pattern=anthropic_route
):
return True
if RouteChecks._is_azure_openai_route(route=route):
return True
for _llm_passthrough_route in LiteLLMRoutes.mapped_pass_through_routes.value:
if _llm_passthrough_route in route:
return True
return False
@staticmethod
@@ -181,6 +181,74 @@ def test_virtual_key_llm_api_routes_allows_google_routes(route):
assert result is True
@pytest.mark.parametrize(
"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",
"/models/google-gemini-2-5-pro-code-reviewer-k8s:generateContent",
"/models/gemini-2.5-flash-exp:countTokens",
"/models/custom-model-name-123:streamGenerateContent",
],
)
def test_google_routes_with_dynamic_model_names_recognized_as_llm_api_route(route):
"""
Test that Google routes with dynamic model names (including custom names) are recognized as LLM API routes.
This test verifies the fix for the issue where routes like:
/v1beta/models/google-gemini-2-5-pro-code-reviewer-k8s:generateContent
were incorrectly classified as "custom admin only route" instead of LLM API routes.
The fix adds pattern matching for Google routes with placeholders like {model_name}.
"""
# Test that the route is recognized as an LLM API route
assert RouteChecks.is_llm_api_route(route) is True
def test_google_routes_with_dynamic_model_names_accessible_to_internal_users():
"""
Test that internal users can access Google routes with dynamic model names.
This ensures that routes like /v1beta/models/{model_name}:generateContent
are properly accessible to internal users and not blocked as admin-only routes.
"""
# Create an internal user object
user_obj = LiteLLM_UserTable(
user_id="test_user",
user_email="test@example.com",
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
# Create an internal user API key auth
valid_token = UserAPIKeyAuth(
user_id="test_user",
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
# Create a mock request
request = MagicMock(spec=Request)
request.query_params = {}
# Test that calling Google route with dynamic model name does NOT raise an exception
try:
RouteChecks.non_proxy_admin_allowed_routes_check(
user_obj=user_obj,
_user_role=LitellmUserRoles.INTERNAL_USER.value,
route="/v1beta/models/google-gemini-2-5-pro-code-reviewer-k8s:generateContent",
request=request,
valid_token=valid_token,
request_data={"contents": [{"parts": [{"text": "test"}]}]},
)
# If no exception is raised, the test passes
except Exception as e:
pytest.fail(
f"Internal user should be able to access Google generateContent route. Got error: {str(e)}"
)
def test_virtual_key_allowed_routes_with_multiple_litellm_routes_member_names():
"""Test that virtual key works with multiple LiteLLMRoutes member names in allowed_routes"""