diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index e202c163aff..0faadbd8729 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -1080,10 +1080,7 @@ def model_max_budget_of(holder: ModelMaxBudgetHolder) -> Mapping[str, object] | def route_skips_budget_checks(route: str) -> bool: return route not in BUDGET_ENFORCED_SIDE_EFFECT_ROUTES and ( - route in MODEL_DISCOVERY_ROUTES - or not ( - RouteChecks.is_llm_api_route(route=route) or RouteChecks.is_provider_prefixed_batch_or_file_route(route) - ) + route in MODEL_DISCOVERY_ROUTES or not RouteChecks.is_llm_handler_route(route) ) diff --git a/litellm/proxy/auth/route_checks.py b/litellm/proxy/auth/route_checks.py index 5c60c31f03e..41491d57867 100644 --- a/litellm/proxy/auth/route_checks.py +++ b/litellm/proxy/auth/route_checks.py @@ -113,8 +113,10 @@ _TOKEN_COUNTING_ACTION: Final = "countTokens" class RouteChecks: @staticmethod - def is_provider_prefixed_batch_or_file_route(route: str) -> bool: - return _PROVIDER_PREFIXED_BATCH_OR_FILE_ROUTE.fullmatch(route) is not None + def is_llm_handler_route(route: str) -> bool: + return RouteChecks.is_llm_api_route(route=route) or ( + _PROVIDER_PREFIXED_BATCH_OR_FILE_ROUTE.fullmatch(route) is not None + ) @staticmethod def check_route_enabled(route: str) -> None: diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 67410e867d6..f887c04af79 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -3144,7 +3144,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e ) # Check 5. Token Spend is under budget - if RouteChecks.is_llm_api_route(route=route): + if RouteChecks.is_llm_handler_route(route): await virtual_key_max_budget_check( valid_token=valid_token, proxy_logging_obj=proxy_logging_obj, diff --git a/litellm/proxy/common_utils/request_tag_mappings.py b/litellm/proxy/common_utils/request_tag_mappings.py index 0e7b0401f4e..3f263403305 100644 --- a/litellm/proxy/common_utils/request_tag_mappings.py +++ b/litellm/proxy/common_utils/request_tag_mappings.py @@ -20,9 +20,8 @@ def is_baggage_mapped_route(route: str) -> bool: return ( ( - RouteChecks.is_llm_api_route(route) + RouteChecks.is_llm_handler_route(route) or InitPassThroughEndpointHelpers.is_registered_pass_through_route(route) - or RouteChecks.is_provider_prefixed_batch_or_file_route(route) ) and not RouteChecks.is_unbilled_route(route) and not RouteChecks.check_route_access(route=route, allowed_routes=LiteLLMRoutes.mcp_inference_routes.value) diff --git a/tests/integration/spend/test_baggage_tag_mappings.py b/tests/integration/spend/test_baggage_tag_mappings.py index c497b28a02b..4342bf4ceae 100644 --- a/tests/integration/spend/test_baggage_tag_mappings.py +++ b/tests/integration/spend/test_baggage_tag_mappings.py @@ -625,6 +625,26 @@ def test_spent_mapped_tag_budget_rejects_a_provider_prefixed_batch_before_the_pr assert mapped.provider.drain() == () +@pytest.mark.parametrize("route", ["/litellm_proxy/v1/batches", "/hosted_vllm/v1/batches"]) +def test_spent_key_budget_rejects_a_provider_prefixed_batch_before_the_provider(mapped: Rig, route: str) -> None: + entry: Final = uuid.uuid4().hex + baggage: Final = f"entrypoint={entry},tenant={entry}" + with mapped.proxy.scenario() as scenario: + key: Final = scenario.key(allowed_routes=["/v1/chat/completions", route], max_budget=0.0001) + + def batch() -> httpx.Response: + return mapped.proxy.request("POST", route, _batch_body(mapped), key=key, headers={"baggage": baggage}) + + first: Final = _chat(mapped, key, baggage) + assert first.status_code == 200, first.text + eventually(lambda: batch().status_code, lambda status: status == 422, seconds=70) + mapped.provider.drain() + blocked: Final = batch() + assert blocked.status_code == 422, blocked.text + assert "budget" in blocked.text.lower(), blocked.text + assert mapped.provider.drain() == () + + def test_require_budgeted_tag_rejects_an_unbudgeted_mapped_tag_until_the_tag_has_a_budget(mapped: Rig) -> None: value: Final = uuid.uuid4().hex tag: Final = mapped.budgeted_tag(value)