scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-09 03:18:44 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 108

refactor(budgets): declare route dependencies with Annotated instead of argument defaults

Browse source
This commit is contained in:
ryan-crabbe-berri 2026-08-29 14:10:58 -07:00
parent 06d0665f50
commit 6b2e7f8a1f
2 changed files with 17 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

22
litellm/proxy/management_endpoints/model_access_group_management_endpoints.py
View file

@ -15,7 +15,7 @@ Endpoints here:
import json
from collections.abc import Mapping, Sequence
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Protocol
from typing import TYPE_CHECKING, Annotated, Any, Final, Protocol
from fastapi import APIRouter, Depends, HTTPException
from typing_extensions import ReadOnly, TypedDict
@ -548,7 +548,7 @@ async def get_all_access_groups_from_db(
)
async def create_model_group(
data: NewModelGroupRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Create a new access group containing multiple model names.
@ -693,7 +693,7 @@ async def create_model_group(
response_model=ListAccessGroupsResponse,
)
async def list_access_groups(
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
List all access groups.
@ -743,7 +743,7 @@ async def list_access_groups(
)
async def get_access_group_info(
access_group: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Get information about a specific access group.
@ -808,7 +808,7 @@ async def get_access_group_info(
async def update_access_group(
access_group: str,
data: UpdateModelGroupRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
):
"""
Update an access group's model names.
@ -961,8 +961,8 @@ async def update_access_group(
)
async def delete_access_group(
access_group: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
auth_cache: UserApiKeyCache = Depends(_auth_cache),
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)],
):
"""
Delete an access group.
@ -1071,7 +1071,6 @@ async def delete_access_group(
)
async def get_access_group_budget(
access_group: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
) -> AccessGroupBudgetResponse:
"""
Get the shared budget of an access group, and the spend drawn against it.
@ -1108,8 +1107,8 @@ async def get_access_group_budget(
async def set_access_group_budget(
access_group: str,
data: AccessGroupBudgetRequest,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
auth_cache: UserApiKeyCache = Depends(_auth_cache),
user_api_key_dict: Annotated[UserAPIKeyAuth, Depends(user_api_key_auth)],
auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)],
) -> AccessGroupBudgetResponse:
"""
Set or replace the shared budget of an access group. Idempotent.
@ -1185,8 +1184,7 @@ async def set_access_group_budget(
)
async def delete_access_group_budget(
access_group: str,
user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth),
auth_cache: UserApiKeyCache = Depends(_auth_cache),
auth_cache: Annotated[UserApiKeyCache, Depends(_auth_cache)],
) -> DeleteAccessGroupBudgetResponse:
"""
Clear the shared budget of an access group, leaving the group itself in place.

22
tests/test_litellm/proxy/management_endpoints/test_access_group_management.py
View file

@ -955,7 +955,7 @@ async def test_get_access_group_budget_returns_the_budget_and_the_shared_spend()
_seed_budget(prisma, "prod-models", spend=42.5, max_budget=100.0, budget_duration="30d")
with _proxy(prisma):
response = await get_access_group_budget(access_group="prod-models", user_api_key_dict=_admin())
response = await get_access_group_budget(access_group="prod-models")
assert response.access_group == "prod-models"
assert response.spend == 42.5
@ -974,7 +974,7 @@ async def test_get_access_group_budget_on_a_budgetless_group_is_200_not_404():
prisma = _FakePrismaClient([], deployments=[_deployment()])
with _proxy(prisma):
response = await get_access_group_budget(access_group="prod-models", user_api_key_dict=_admin())
response = await get_access_group_budget(access_group="prod-models")
assert response.spend == 0.0
assert response.budget is None
@ -998,16 +998,14 @@ async def test_access_group_budget_routes_404_on_an_unknown_group():
admin = _admin()
calls = (
lambda: get_access_group_budget(access_group="ghost-group", user_api_key_dict=admin),
lambda: get_access_group_budget(access_group="ghost-group"),
lambda: set_access_group_budget(
access_group="ghost-group",
data=AccessGroupBudgetRequest(max_budget=1.0),
user_api_key_dict=admin,
auth_cache=cache,
),
lambda: delete_access_group_budget(
access_group="ghost-group", user_api_key_dict=admin, auth_cache=cache
),
lambda: delete_access_group_budget(access_group="ghost-group", auth_cache=cache),
)
with _proxy(prisma):
@ -1033,9 +1031,7 @@ async def test_delete_access_group_budget_drops_the_row_and_spares_the_shared_bu
cache = _FakeAuthCache()
with _proxy(prisma):
response = await delete_access_group_budget(
access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache
)
response = await delete_access_group_budget(access_group="prod-models", auth_cache=cache)
assert response.budget_deleted is True
assert prisma.access_group_budget_table.rows == {}
@ -1056,9 +1052,7 @@ async def test_delete_access_group_budget_on_a_budgetless_group_still_evicts():
cache = _FakeAuthCache(journal)
with _proxy(prisma):
response = await delete_access_group_budget(
access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache
)
response = await delete_access_group_budget(access_group="prod-models", auth_cache=cache)
assert response.budget_deleted is False
assert prisma.budget_table.deleted_ids == []
@ -1160,9 +1154,7 @@ async def test_delete_access_group_budget_evicts_both_auth_cache_keys():
cache = _FakeAuthCache(journal)
with _proxy(prisma):
await delete_access_group_budget(
access_group="prod-models", user_api_key_dict=_admin(), auth_cache=cache
)
await delete_access_group_budget(access_group="prod-models", auth_cache=cache)
_assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models")