mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-08 22:21:35 +00:00
Fix Token Spend is under budget for passthrough (#15805)
This commit is contained in:
parent
abe67df95e
commit
2ab2d15efc
2 changed files with 129 additions and 1 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue