scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(auth): allow video retrieval after budget exhaustion (#51227)

Browse source 
* fix(auth): allow paid video retrieval after budget exhaustion

* test(auth): check video retrieval boundary on registered route

* fix(auth): exempt video retrieval in post-custom-auth checks

* style(auth): format custom video budget exemption

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
This commit is contained in:
joshua-berri 2026-11-11 17:10:01 -08:00 • committed by GitHub
parent 6116adeb74
commit a5d04ec9ac
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 485 additions and 6 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/proxy/auth/auth_checks.py
View file

@ -1081,6 +1081,15 @@ BUDGET_ENFORCED_SIDE_EFFECT_ROUTES: Final = frozenset(
}
)
VIDEO_RETRIEVAL_ROUTES: Final = frozenset(
{
"/videos/{video_id}",
"/v1/videos/{video_id}",
"/videos/{video_id}/content",
"/v1/videos/{video_id}/content",
}
)
MCP_DISCOVERY_ROUTES: Final = frozenset({"/mcp-rest/tools/list", "/v1/mcp/tools"})
MCP_ZERO_SPEND_JSONRPC_METHODS: Final = frozenset({"initialize", "notifications/initialized", "ping", "tools/list"})
@ -1097,6 +1106,10 @@ def is_mcp_discovery_request(route: str, request_body: Mapping[str, object]) ->
return isinstance(method, str) and method in MCP_ZERO_SPEND_JSONRPC_METHODS
def is_video_retrieval_request(route: str, http_method: str | None) -> bool:
return http_method == "GET" and RouteChecks.check_route_access(route=route, allowed_routes=VIDEO_RETRIEVAL_ROUTES)
class ModelMaxBudgetHolder(Protocol):
@property
def model_max_budget(self) -> Mapping[str, object] | None: ...
@ -1126,8 +1139,9 @@ def request_skips_budget_checks(
llm_router: Router | None,
*,
request_body: Mapping[str, object] | None = None,
http_method: str | None = None,
) -> bool:
if route_skips_budget_checks(route=route):
if route_skips_budget_checks(route=route) or is_video_retrieval_request(route=route, http_method=http_method):
return True
if request_body is not None and is_mcp_discovery_request(route=route, request_body=request_body):
return True
@ -1189,6 +1203,7 @@ async def common_checks(
model=_model,
llm_router=llm_router,
request_body=_typed_request_body(request_body),
http_method=RouteChecks.get_request_method(request),
)
)

13
litellm/proxy/auth/user_api_key_auth.py
View file

@ -85,6 +85,7 @@ from litellm.proxy.auth.auth_checks import ( # noqa: F401 # legacy module expo
is_model_cost_zero,
is_user_proxy_admin,
is_valid_fallback_model,
is_video_retrieval_request,
jwt_key_mapping_cache_key,
key_access_group_grants_model,
key_model_aliases_for_auth_check,
@ -2680,7 +2681,11 @@ async def user_api_key_auth_builder(
team_id=valid_token.team_id,
)
skip_budget_checks = request_skips_budget_checks(
route=route, model=model, llm_router=llm_router, request_body=request_data
route=route,
model=model,
llm_router=llm_router,
request_body=request_data,
http_method=RouteChecks.get_request_method(request),
)
# Fetch project object for JWT path if project_id is set
@ -3198,6 +3203,7 @@ async def validate_resolved_virtual_key( # noqa: C901 # Preserve ordering of e
model=await budget_served_model(model, request_data, request, llm_router, valid_token),
llm_router=llm_router,
request_body=request_data,
http_method=RouteChecks.get_request_method(request),
)
# Check 3. Check if user is in their team budget
@ -4299,6 +4305,7 @@ async def _should_skip_budget_checks(
model=(model if sub_call else await budget_served_model(model, request_data, request, llm_router, valid_token)),
llm_router=llm_router,
request_body=request_data,
http_method=RouteChecks.get_request_method(request),
)
@ -4976,7 +4983,9 @@ async def _run_post_custom_auth_checks(
# spend some other model accrued. The JWT and virtual-key paths already skip
# every budget check for these; this path did not, so the same request could
# be refused under custom auth and served under the other two.
skip_budget_checks: Final = (
skip_budget_checks: Final = is_video_retrieval_request(
route=route, http_method=RouteChecks.get_request_method(request)
) or (
is_model_cost_zero(
model=await budget_served_model(current_model, request_data, request, llm_router, valid_token),
llm_router=llm_router,

7
litellm/proxy/common_request_processing.py
View file

@ -111,6 +111,7 @@ from litellm.proxy.auth.auth_checks import (
tag_max_budget_check_for_tags,
)
from litellm.proxy.auth.auth_utils import check_response_size_is_safe, get_request_route
from litellm.proxy.auth.route_checks import RouteChecks
from litellm.proxy.bug_report_config import build_proxy_bug_report
from litellm.proxy.common_utils.callback_utils import (
get_logging_caching_headers,
@ -1048,11 +1049,14 @@ async def _enforce_guardrail_added_tag_budgets(
llm_router: Router | None,
user_api_key_dict: UserAPIKeyAuth,
proxy_logging_obj: ProxyLogging,
http_method: str | None = None,
) -> None:
added_tags: Final = tuple(
tag for tag in get_tags_from_request_body(request_body=data) if tag not in tags_before_guardrails
)
if not added_tags or request_skips_budget_checks(route=route, model=_request_model(data), llm_router=llm_router):
if not added_tags or request_skips_budget_checks(
route=route, model=_request_model(data), llm_router=llm_router, http_method=http_method
):
return
from litellm.proxy.proxy_server import prisma_client, user_api_key_cache
@ -2864,6 +2868,7 @@ class ProxyBaseLLMRequestProcessing:
llm_router=llm_router,
user_api_key_dict=user_api_key_dict,
proxy_logging_obj=proxy_logging_obj,
http_method=RouteChecks.get_request_method(request),
)
if route_type == "aget_responses":
attach_post_call_pipelines_to_retrieval(

67
tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py
View file

@ -7286,6 +7286,45 @@ async def test_model_discovery_route_bypasses_user_budget():
assert result is True
@pytest.mark.parametrize(
"method, route, expect_blocked",
[
("GET", "/v1/videos/video_abc", False),
("GET", "/videos/video_abc/content", False),
("POST", "/v1/videos", True),
("POST", "/v1/videos/video_abc/remix", True),
("GET", "/v1/videos", True),
],
)
@pytest.mark.asyncio
async def test_exhausted_team_budget_spares_only_video_retrieval(method: str, route: str, expect_blocked: bool) -> None:
from litellm.proxy.auth.auth_checks import common_checks
request: Final = MagicMock()
request.method = method
async def run() -> bool:
return await common_checks(
request_body={},
team_object=LiteLLM_TeamTable(team_id="test-team", spend=0.4, max_budget=0.3),
user_object=None,
end_user_object=None,
global_proxy_spend=None,
general_settings={},
route=route,
llm_router=None,
proxy_logging_obj=AsyncMock(),
valid_token=UserAPIKeyAuth(token="test-token", team_id="test-team"),
request=request,
)
if expect_blocked:
with pytest.raises(litellm.BudgetExceededError):
await run()
else:
assert await run() is True
@pytest.mark.parametrize("route", ["/health/services", "/auto_router/test_routing"])
@pytest.mark.asyncio
async def test_side_effectful_info_route_still_enforces_budget(route: str) -> None:
@ -11202,6 +11241,34 @@ def test_request_skips_budget_checks_for_non_spending_requests(
)
@pytest.mark.parametrize(
"route, http_method, expected",
[
("/v1/videos/video_abc", "GET", True),
("/videos/video_abc", "GET", True),
("/v1/videos/video_abc/content", "GET", True),
("/videos/video_abc/content", "GET", True),
("/v1/videos/video_abc", "POST", False),
("/v1/videos/video_abc/content", None, False),
("/v1/videos", "GET", False),
("/v1/videos", "POST", False),
("/v1/videos/video_abc/remix", "POST", False),
("/v1/videos/edits", "POST", False),
("/v1/videos/video_abc/remix", "GET", False),
("/v1/chat/completions", "GET", False),
],
)
def test_request_skips_budget_checks_only_for_paid_video_retrieval(
route: str, http_method: str | None, expected: bool, free_and_paid_router: litellm.Router
) -> None:
assert (
request_skips_budget_checks(
route=route, model="paid-model", llm_router=free_and_paid_router, http_method=http_method
)
is expected
)
def _agent_model_ceiling_resolver(
models: frozenset[str] | None,
) -> tuple[CeilingResolver, list[str]]:

93
tests/unit/proxy/auth/test_custom_auth_end_user_budget.py
View file

@ -1,3 +1,4 @@
from datetime import datetime, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
@ -6,15 +7,24 @@ from starlette.requests import Request
from starlette.routing import Route
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import (
LiteLLM_BudgetTable,
LiteLLM_EndUserTable,
Litellm_EntityType,
ProxyErrorTypes,
ProxyException,
UserAPIKeyAuth,
)
from litellm.proxy.auth.user_api_key_auth import (
_run_post_custom_auth_checks,
update_valid_token_with_end_user_params,
)
from litellm.proxy.hooks.model_max_budget_limiter import (
PROXY_VirtualKeyModelMaxBudgetLimiter,
model_budget_spend_cache_key,
)
from litellm.types.videos.utils import encode_video_id_with_provider
@pytest.mark.asyncio
@ -513,3 +523,86 @@ async def test_custom_auth_judges_the_zero_cost_waiver_by_the_key_alias_target(
result: Final = await run()
assert result.token == "test_token"
limiter.is_end_user_within_model_budget.assert_not_awaited()
_VIDEO_ID: Final = encode_video_id_with_provider(video_id="video_raw", provider="openai", model_id="sora-2")
async def _post_custom_auth_video_request(
valid_token: UserAPIKeyAuth, method: str, template: str, request_data: dict[str, object]
) -> UserAPIKeyAuth:
limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
await limiter.dual_cache.async_set_cache(
key=model_budget_spend_cache_key(
entity_type=Litellm_EntityType.KEY, entity_id=valid_token.token, budget_model="sora-2", budget_duration="1d"
),
value=0.4,
)
path: Final = template.format(video_id=_VIDEO_ID)
http_request: Final = Request(
{
"type": "http",
"method": method,
"path": path,
"headers": [],
"query_string": b"",
"route": Route(template, endpoint=lambda request: None),
}
)
with (
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.model_max_budget_limiter", limiter),
patch("litellm.proxy.proxy_server.general_settings", {}),
):
return await _run_post_custom_auth_checks(
valid_token=valid_token,
request=http_request,
request_data=request_data,
route=path,
parent_otel_span=None,
)
def _exhausted_model_budget_token() -> UserAPIKeyAuth:
return UserAPIKeyAuth(token="test_token", model_max_budget={"sora-2": {"max_budget": 0.3, "budget_duration": "1d"}})
@pytest.mark.asyncio
@pytest.mark.parametrize(
"template",
["/v1/videos/{video_id}", "/videos/{video_id}", "/v1/videos/{video_id}/content", "/videos/{video_id}/content"],
)
async def test_custom_auth_exhausted_model_budget_still_retrieves_an_already_paid_video(template: str) -> None:
result: Final = await _post_custom_auth_video_request(
_exhausted_model_budget_token(), "GET", template, {"video_id": _VIDEO_ID}
)
assert result.token == "test_token"
@pytest.mark.asyncio
@pytest.mark.parametrize(
"method, template, request_data",
[
("POST", "/v1/videos", {"model": "sora-2"}),
("POST", "/v1/videos/{video_id}/remix", {"video_id": _VIDEO_ID}),
("GET", "/v1/videos", {"model": "sora-2"}),
],
)
async def test_custom_auth_exhausted_model_budget_still_blocks_video_requests_other_than_retrieval(
method: str, template: str, request_data: dict[str, object]
) -> None:
with pytest.raises(litellm.BudgetExceededError):
await _post_custom_auth_video_request(_exhausted_model_budget_token(), method, template, request_data)
@pytest.mark.asyncio
async def test_custom_auth_video_retrieval_still_rejects_an_expired_token() -> None:
expired_token: Final = _exhausted_model_budget_token().model_copy(
update={"expires": datetime(2000, 1, 1, tzinfo=timezone.utc)}
)
with pytest.raises(ProxyException) as exc_info:
await _post_custom_auth_video_request(expired_token, "GET", "/v1/videos/{video_id}", {"video_id": _VIDEO_ID})
assert exc_info.value.type == ProxyErrorTypes.expired_key

109
tests/unit/proxy/auth/test_user_api_key_auth.py
View file

@ -2360,6 +2360,115 @@ async def test_jwt_path_skips_user_model_budget_only_for_zero_spend_mcp_methods(
assert user_model_budget_check.await_count == (1 if expect_user_model_budget_check else 0)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"method, path, expect_rejected",
[
("GET", "/v1/videos/{video_id}", False),
("GET", "/videos/{video_id}/content", False),
("POST", "/v1/videos", True),
("POST", "/v1/videos/{video_id}/remix", True),
],
)
async def test_jwt_path_spares_only_video_retrieval_from_an_exhausted_user_model_budget(
monkeypatch: pytest.MonkeyPatch, method: str, path: str, expect_rejected: bool
) -> None:
from litellm.caching.dual_cache import DualCache
from litellm.proxy._types import JWTAuthBuilderResult, LiteLLM_JWTAuth, Litellm_EntityType, ProxyErrorTypes
from litellm.proxy.auth import user_api_key_auth as user_api_key_auth_module
from litellm.proxy.auth.handle_jwt import JWTAuthManager, JWTHandler
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth_builder
from litellm.proxy.hooks.model_max_budget_limiter import (
PROXY_VirtualKeyModelMaxBudgetLimiter,
model_budget_spend_cache_key,
)
from litellm.types.videos.utils import encode_video_id_with_provider
video_id: Final = encode_video_id_with_provider(video_id="video_raw", provider="openai", model_id="sora-2")
jwt_result: Final = JWTAuthBuilderResult(
is_proxy_admin=False,
team_object=None,
user_object=LiteLLM_UserTable(
user_id="jwt-user",
model_max_budget={"sora-2": {"budget_limit": 0.3, "time_period": "1d"}},
),
end_user_object=None,
org_object=None,
token="fake.jwt.token",
team_id=None,
user_id="jwt-user",
user_email=None,
end_user_id=None,
org_id=None,
team_membership=None,
jwt_claims={},
agent_id=None,
)
class _VerifiedJWTAuthManager(JWTAuthManager):
@staticmethod
async def auth_builder(
api_key: object,
jwt_handler: object,
request_data: object,
general_settings: object,
route: object,
prisma_client: object,
user_api_key_cache: object,
parent_otel_span: object,
proxy_logging_obj: object,
request_headers: object = None,
request_method: object = None,
) -> JWTAuthBuilderResult:
return jwt_result
limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
await limiter.dual_cache.async_set_cache(
key=model_budget_spend_cache_key(
entity_type=Litellm_EntityType.USER, entity_id="jwt-user", budget_model="sora-2", budget_duration="1d"
),
value=0.4,
)
jwt_handler: Final = JWTHandler()
jwt_handler.update_environment(
prisma_client=None,
user_api_key_cache=DualCache(),
litellm_jwtauth=LiteLLM_JWTAuth(),
)
mock_request: Final = MagicMock()
mock_request.url.path = path.format(video_id=video_id)
mock_request.method = method
mock_request.headers = {"authorization": "Bearer fake.jwt.token"}
mock_request.query_params = {}
request_data: Final[dict[str, object]] = {"model": "sora-2"} if path == "/v1/videos" else {"video_id": video_id}
monkeypatch.setattr(litellm.proxy.proxy_server, "general_settings", {"enable_jwt_auth": True})
monkeypatch.setattr(litellm.proxy.proxy_server, "premium_user", True)
monkeypatch.setattr(litellm.proxy.proxy_server, "jwt_handler", jwt_handler)
monkeypatch.setattr(litellm.proxy.proxy_server, "llm_router", None)
monkeypatch.setattr(litellm.proxy.proxy_server, "model_max_budget_limiter", limiter)
monkeypatch.setattr(user_api_key_auth_module, "JWTAuthManager", _VerifiedJWTAuthManager)
async def authenticate() -> str | None:
result: Final = await user_api_key_auth_builder(
request=mock_request,
api_key="Bearer fake.jwt.token",
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data=request_data,
)
return result.user_id
if not expect_rejected:
assert await authenticate() == "jwt-user"
return
with pytest.raises(ProxyException) as exc_info:
await authenticate()
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
def _qa_websocket(messages):
async def receive():
return {"type": "websocket.connect"}

142
tests/unit/proxy/auth/test_user_api_key_auth_request_flow.py
View file

@ -13,8 +13,8 @@ from functools import partial
from io import BytesIO
from pathlib import Path
from textwrap import dedent
from types import SimpleNamespace
from typing import Final
from types import MappingProxyType, SimpleNamespace
from typing import Final, Literal
from unittest.mock import ANY, AsyncMock, MagicMock, patch
@ -98,6 +98,7 @@ from litellm.proxy.auth.user_api_key_auth import (
)
from litellm.proxy.pass_through_endpoints import pass_through_endpoints as pass_through
from litellm.proxy.spend_tracking.carried_budget_state import carried_budget_metadata
from litellm.types.videos.utils import encode_video_id_with_provider
from tests.unit.proxy.db.fake_prisma_engine import engine_call
from tests._master_key import MASTER_KEY
@ -13956,6 +13957,143 @@ async def test_exhausted_key_budget_spares_only_zero_spend_mcp_discovery(
assert result.spend == 5.0
_VIDEO_ID: Final = encode_video_id_with_provider(video_id="video_raw", provider="openai", model_id="sora-2")
_VIDEO_RETRIEVAL_ROUTES: Final = (
f"/v1/videos/{_VIDEO_ID}",
f"/videos/{_VIDEO_ID}",
f"/v1/videos/{_VIDEO_ID}/content",
f"/videos/{_VIDEO_ID}/content",
)
_EXPIRED_BEFORE_ANY_TEST_RUN: Final = datetime(2000, 1, 1, tzinfo=timezone.utc)
_VIDEO_SPENDING_REQUESTS: Final = (
("POST", "/v1/videos", MappingProxyType({"model": "sora-2"})),
("POST", f"/v1/videos/{_VIDEO_ID}/remix", MappingProxyType({"video_id": _VIDEO_ID})),
("POST", "/v1/videos/edits", MappingProxyType({"model": "sora-2"})),
("POST", f"/v1/videos/{_VIDEO_ID}", MappingProxyType({"video_id": _VIDEO_ID})),
("GET", "/v1/videos", MappingProxyType({"model": "sora-2"})),
)
async def _authenticate_video_request(
key: UserAPIKeyAuth, api_key: str, method: str, route: str, request_data: dict[str, object]
) -> UserAPIKeyAuth:
user_api_key_cache: Final = DualCache()
await cache_key_object(
hashed_token=key.token or "",
user_api_key_obj=key,
user_api_key_cache=user_api_key_cache,
proxy_logging_obj=None,
)
model_budget_limiter: Final = PROXY_VirtualKeyModelMaxBudgetLimiter(dual_cache=DualCache())
await model_budget_limiter.dual_cache.async_set_cache(
key=model_budget_spend_cache_key(
entity_type=Litellm_EntityType.KEY, entity_id=key.token, budget_model="sora-2", budget_duration="1d"
),
value=0.4,
)
request: Final = Request(
{
"type": "http",
"method": method,
"path": route,
"headers": [(b"authorization", f"Bearer {api_key}".encode())],
"query_string": b"",
}
)
proxy_logging_obj: Final = MagicMock()
proxy_logging_obj.budget_alerts = AsyncMock()
proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None)
with (
patch("litellm.proxy.proxy_server.general_settings", {}),
patch("litellm.proxy.proxy_server.master_key", "sk-master"),
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.llm_router", None),
patch("litellm.proxy.proxy_server.llm_model_list", []),
patch("litellm.proxy.proxy_server.user_api_key_cache", user_api_key_cache),
patch("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj),
patch("litellm.proxy.proxy_server.model_max_budget_limiter", model_budget_limiter),
patch(
"litellm.proxy.auth.user_api_key_auth.virtual_key_max_budget_alert_check",
new_callable=AsyncMock,
),
):
resolved: Final = await user_api_key_auth_builder(
request=request,
api_key=f"Bearer {api_key}",
azure_api_key_header="",
anthropic_api_key_header=None,
google_ai_studio_api_key_header=None,
azure_apim_header=None,
request_data=request_data,
)
await _authorize_authenticated_request(resolved, request, request_data, route, api_key)
return resolved
def _exhausted_video_key(budget: Literal["max_budget", "model_max_budget"]) -> tuple[UserAPIKeyAuth, str]:
from litellm.proxy.utils import hash_token
api_key: Final = f"sk-exhausted-video-{budget}"
if budget == "max_budget":
return UserAPIKeyAuth(token=hash_token(api_key), max_budget=0.3, spend=0.4), api_key
return (
UserAPIKeyAuth(
token=hash_token(api_key), model_max_budget={"sora-2": {"max_budget": 0.3, "budget_duration": "1d"}}
),
api_key,
)
@pytest.mark.asyncio
@pytest.mark.parametrize("budget", ["max_budget", "model_max_budget"])
@pytest.mark.parametrize("route", _VIDEO_RETRIEVAL_ROUTES)
async def test_exhausted_key_budget_still_retrieves_an_already_paid_video(
budget: Literal["max_budget", "model_max_budget"], route: str
) -> None:
key, api_key = _exhausted_video_key(budget)
result: Final = await _authenticate_video_request(key, api_key, "GET", route, {"video_id": _VIDEO_ID})
assert result.token == key.token
@pytest.mark.asyncio
@pytest.mark.parametrize("budget", ["max_budget", "model_max_budget"])
@pytest.mark.parametrize("method, route, request_data", _VIDEO_SPENDING_REQUESTS)
async def test_exhausted_key_budget_still_blocks_video_requests_other_than_retrieval(
budget: Literal["max_budget", "model_max_budget"], method: str, route: str, request_data: Mapping[str, object]
) -> None:
key, api_key = _exhausted_video_key(budget)
with pytest.raises(ProxyException) as exc_info:
await _authenticate_video_request(key, api_key, method, route, dict(request_data))
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
@pytest.mark.asyncio
@pytest.mark.parametrize(
"overrides, expected_type",
[
({"expires": _EXPIRED_BEFORE_ANY_TEST_RUN}, ProxyErrorTypes.expired_key),
({"blocked": True}, ProxyErrorTypes.auth_error),
({"models": ["other-model"]}, ProxyErrorTypes.key_model_access_denied),
],
)
@pytest.mark.parametrize("route", (_VIDEO_RETRIEVAL_ROUTES[0], _VIDEO_RETRIEVAL_ROUTES[2]))
async def test_video_retrieval_on_exhausted_key_keeps_key_access_controls(
overrides: dict[str, object], expected_type: ProxyErrorTypes, route: str
) -> None:
exhausted_key, api_key = _exhausted_video_key("max_budget")
with pytest.raises(ProxyException) as exc_info:
await _authenticate_video_request(
exhausted_key.model_copy(update=overrides), api_key, "GET", route, {"video_id": _VIDEO_ID}
)
assert exc_info.value.type == expected_type
@pytest.mark.asyncio
@pytest.mark.parametrize(
"requested_model, key_aliases, team_model_aliases, expect_blocked",

43
tests/unit/proxy/test_common_request_processing.py
View file

@ -1020,6 +1020,49 @@ class TestProxyBaseLLMRequestProcessing:
)
tag_budget_check.assert_not_awaited()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"http_method, route, route_type, expect_rejected",
[
("GET", "/v1/videos/video_abc", "avideo_status", False),
("GET", "/videos/video_abc/content", "avideo_content", False),
("POST", "/v1/videos", "avideo_generation", True),
("POST", "/v1/videos/video_abc/remix", "avideo_remix", True),
("GET", "/v1/videos", "avideo_list", True),
],
)
async def test_common_processing_pre_call_logic_spares_only_video_retrieval_from_exhausted_guardrail_tags(
self, monkeypatch: pytest.MonkeyPatch, http_method: str, route: str, route_type: str, expect_rejected: bool
) -> None:
processing_obj, mock_request, mock_proxy_logging_obj, mock_proxy_config, tag_budget_check = (
self._guardrail_tag_budget_harness(
monkeypatch,
request_body={"model": "paid-model", "metadata": {"tags": ["existing-tag"]}},
guardrail_tags=["guardrail-tag"],
route=route,
)
)
tag_budget_check.side_effect = litellm.BudgetExceededError(current_cost=0.4, max_budget=0.3)
mock_request.method = http_method
async def pre_call() -> None:
await processing_obj.common_processing_pre_call_logic(
request=mock_request,
general_settings={},
user_api_key_dict=ProxyUserAPIKeyAuth(api_key="sk-test"),
proxy_logging_obj=mock_proxy_logging_obj,
proxy_config=mock_proxy_config,
route_type=route_type,
llm_router=self._router_with_free_and_paid_models(),
)
if not expect_rejected:
await pre_call()
return
with pytest.raises(ProxyException) as exc_info:
await pre_call()
assert exc_info.value.type == ProxyErrorTypes.budget_exceeded
@pytest.mark.asyncio
async def test_enforce_guardrail_added_tag_budgets_raises_budget_exceeded_proxy_exception(self, monkeypatch):
from litellm.proxy.common_request_processing import _enforce_guardrail_added_tag_budgets