Fix Token Spend is under budget for passthrough (#15805)

This commit is contained in:
Sameer Kankute 2025-10-22 23:25:22 +05:30 committed by GitHub
parent abe67df95e
commit 2ab2d15efc
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 129 additions and 1 deletions

View file

@ -998,7 +998,7 @@ async def _user_api_key_auth_builder( # noqa: PLR0915
)
# Check 4. Token Spend is under budget
if route in LiteLLMRoutes.llm_api_routes.value:
if RouteChecks.is_llm_api_route(route=route):
await _virtual_key_max_budget_check(
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,

View file

@ -14,6 +14,7 @@ from unittest.mock import MagicMock
import pytest
from litellm.proxy.auth.user_api_key_auth import get_api_key
from litellm.proxy.auth.route_checks import RouteChecks
def test_get_api_key():
@ -56,3 +57,130 @@ def test_get_api_key_with_custom_litellm_key_header(
route="",
request=MagicMock(),
) == (api_key, passed_in_key)
def test_route_checks_is_llm_api_route():
"""Test RouteChecks.is_llm_api_route() correctly identifies LLM API routes including passthrough endpoints"""
# Test OpenAI routes
openai_routes = [
"/v1/chat/completions",
"/chat/completions",
"/v1/completions",
"/completions",
"/v1/embeddings",
"/embeddings",
"/v1/images/generations",
"/images/generations",
"/v1/audio/transcriptions",
"/audio/transcriptions",
"/v1/audio/speech",
"/audio/speech",
"/v1/moderations",
"/moderations",
"/v1/models",
"/models",
"/v1/rerank",
"/rerank",
"/v1/realtime",
"/realtime",
]
for route in openai_routes:
assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route"
# Test Anthropic routes
anthropic_routes = [
"/v1/messages",
"/v1/messages/count_tokens",
]
for route in anthropic_routes:
assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route"
# Test passthrough routes (this is the key improvement over the old route checking)
passthrough_routes = [
"/bedrock/v1/chat/completions",
"/vertex-ai/v1/chat/completions",
"/vertex_ai/v1/chat/completions",
"/cohere/v1/chat/completions",
"/gemini/v1/chat/completions",
"/anthropic/v1/messages",
"/langfuse/v1/chat/completions",
"/azure/v1/chat/completions",
"/openai/v1/chat/completions",
"/assemblyai/v1/transcript",
"/eu.assemblyai/v1/transcript",
"/vllm/v1/chat/completions",
"/mistral/v1/chat/completions",
]
for route in passthrough_routes:
assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route"
# Test MCP routes
mcp_routes = [
"/mcp",
"/mcp/",
"/mcp/test",
]
for route in mcp_routes:
assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route"
# Test routes with placeholders
placeholder_routes = [
"/v1/threads/thread_49EIN5QF32s4mH20M7GFKdlZ",
"/threads/thread_49EIN5QF32s4mH20M7GFKdlZ",
"/v1/assistants/assistant_123",
"/assistants/assistant_123",
"/v1/files/file_123",
"/files/file_123",
"/v1/batches/batch_123",
"/batches/batch_123",
]
for route in placeholder_routes:
assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route"
# Test Azure OpenAI routes
azure_routes = [
"/openai/deployments/gpt-4/chat/completions",
"/openai/deployments/gpt-3.5-turbo/completions",
"/engines/gpt-4/chat/completions",
"/engines/gpt-3.5-turbo/completions",
]
for route in azure_routes:
assert RouteChecks.is_llm_api_route(route=route), f"Route {route} should be identified as LLM API route"
# Test non-LLM routes (should return False)
non_llm_routes = [
"/health",
"/metrics",
"/key/list",
"/team/list",
"/user/list",
"/config",
"/routes",
"/",
"/admin/settings",
"/logs",
"/debug",
"/test",
]
for route in non_llm_routes:
assert not RouteChecks.is_llm_api_route(route=route), f"Route {route} should NOT be identified as LLM API route"
# Test invalid inputs
invalid_inputs = [
None,
123,
[],
{},
"",
]
for invalid_input in invalid_inputs:
assert not RouteChecks.is_llm_api_route(route=invalid_input), f"Invalid input {invalid_input} should return False"