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

Author: Sameerlite

litellm/proxy/auth/user_api_key_auth.py

@@ -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,

tests/test_litellm/proxy/auth/test_user_api_key_auth.py

@@ -14,6 +14,7 @@
 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"