mirror of
https://github.com/tiennm99/litellm.git
synced 2026-07-12 03:06:22 +00:00
harden end-user and tag budget reservations
This commit is contained in:
user
@@ -655,13 +655,7 @@ async def common_checks( # noqa: PLR0915
|
||||
end_user_object is not None
|
||||
and end_user_object.litellm_budget_table is not None
|
||||
):
|
||||
end_user_budget = end_user_object.litellm_budget_table.max_budget
|
||||
if end_user_budget is not None and end_user_object.spend > end_user_budget:
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=end_user_object.spend,
|
||||
max_budget=end_user_budget,
|
||||
message=f"ExceededBudget: End User={end_user_object.user_id} over budget. Spend={end_user_object.spend}, Budget={end_user_budget}",
|
||||
)
|
||||
await _check_end_user_budget(end_user_obj=end_user_object, route=route)
|
||||
|
||||
_enforce_user_param_check(general_settings, request, request_body, route)
|
||||
_reject_clientside_metadata_tags_check(general_settings, request_body, route)
|
||||
@@ -1005,7 +999,7 @@ async def _apply_default_budget_to_end_user(
|
||||
return end_user_obj
|
||||
|
||||
|
||||
def _check_end_user_budget(
|
||||
async def _check_end_user_budget(
|
||||
end_user_obj: LiteLLM_EndUserTable,
|
||||
route: str,
|
||||
) -> None:
|
||||
@@ -1026,11 +1020,20 @@ def _check_end_user_budget(
|
||||
return
|
||||
|
||||
end_user_budget = end_user_obj.litellm_budget_table.max_budget
|
||||
if end_user_budget is not None and end_user_obj.spend > end_user_budget:
|
||||
if end_user_budget is None:
|
||||
return
|
||||
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
end_user_spend = await get_current_spend(
|
||||
counter_key=f"spend:end_user:{end_user_obj.user_id}",
|
||||
fallback_spend=end_user_obj.spend or 0.0,
|
||||
)
|
||||
if end_user_spend > end_user_budget:
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=end_user_obj.spend,
|
||||
current_cost=end_user_spend,
|
||||
max_budget=end_user_budget,
|
||||
message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_obj.spend}, Budget={end_user_budget}",
|
||||
message=f"ExceededBudget: End User={end_user_obj.user_id} over budget. Spend={end_user_spend}, Budget={end_user_budget}",
|
||||
)
|
||||
|
||||
|
||||
@@ -1082,7 +1085,7 @@ async def get_end_user_object(
|
||||
)
|
||||
|
||||
# Check budget limits
|
||||
_check_end_user_budget(end_user_obj=return_obj, route=route)
|
||||
await _check_end_user_budget(end_user_obj=return_obj, route=route)
|
||||
|
||||
return return_obj
|
||||
|
||||
@@ -1113,7 +1116,7 @@ async def get_end_user_object(
|
||||
)
|
||||
|
||||
# Check budget limits
|
||||
_check_end_user_budget(end_user_obj=_response, route=route)
|
||||
await _check_end_user_budget(end_user_obj=_response, route=route)
|
||||
|
||||
return _response
|
||||
|
||||
@@ -3823,13 +3826,19 @@ async def _tag_max_budget_check(
|
||||
if (
|
||||
tag_object.litellm_budget_table is not None
|
||||
and tag_object.litellm_budget_table.max_budget is not None
|
||||
and tag_object.spend is not None
|
||||
and tag_object.spend > tag_object.litellm_budget_table.max_budget
|
||||
):
|
||||
from litellm.proxy.proxy_server import get_current_spend
|
||||
|
||||
tag_spend = await get_current_spend(
|
||||
counter_key=f"spend:tag:{tag_name}",
|
||||
fallback_spend=tag_object.spend or 0.0,
|
||||
)
|
||||
if tag_spend <= tag_object.litellm_budget_table.max_budget:
|
||||
continue
|
||||
raise litellm.BudgetExceededError(
|
||||
current_cost=tag_object.spend,
|
||||
current_cost=tag_spend,
|
||||
max_budget=tag_object.litellm_budget_table.max_budget,
|
||||
message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_object.spend}, Max budget: {tag_object.litellm_budget_table.max_budget}",
|
||||
message=f"Budget has been exceeded! Tag={tag_name} Current cost: {tag_spend}, Max budget: {tag_object.litellm_budget_table.max_budget}",
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -1869,6 +1869,8 @@ async def _run_centralized_common_checks(
|
||||
llm_router=llm_router,
|
||||
team_object=team_object,
|
||||
user_object=user_object,
|
||||
end_user_id=end_user_id,
|
||||
end_user_object=end_user_object,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
@@ -1893,6 +1895,8 @@ async def _reserve_budget_after_common_checks(
|
||||
user_api_key_cache: DualCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
skip_budget_checks: bool,
|
||||
end_user_id: Optional[str] = None,
|
||||
end_user_object: Optional[LiteLLM_EndUserTable] = None,
|
||||
) -> None:
|
||||
user_api_key_auth_obj.budget_reservation = None
|
||||
if skip_budget_checks:
|
||||
@@ -1912,6 +1916,8 @@ async def _reserve_budget_after_common_checks(
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
end_user_id=end_user_id,
|
||||
end_user_object=end_user_object,
|
||||
)
|
||||
|
||||
|
||||
|
||||
@@ -35,6 +35,8 @@ class SpendCounterReseed:
|
||||
spend:team:{team_id} -> LiteLLM_TeamTable.spend
|
||||
spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend
|
||||
spend:user:{user_id} -> LiteLLM_UserTable.spend
|
||||
spend:end_user:{end_user_id} -> LiteLLM_EndUserTable.spend
|
||||
spend:tag:{tag_name} -> LiteLLM_TagTable.spend
|
||||
spend:org:{org_id} -> LiteLLM_OrganizationTable.spend
|
||||
"""
|
||||
|
||||
@@ -98,6 +100,16 @@ class SpendCounterReseed:
|
||||
row = await prisma_client.db.litellm_usertable.find_unique(
|
||||
where={"user_id": user_id}
|
||||
)
|
||||
elif counter_key.startswith("spend:end_user:"):
|
||||
end_user_id = counter_key[len("spend:end_user:") :]
|
||||
row = await prisma_client.db.litellm_endusertable.find_unique(
|
||||
where={"user_id": end_user_id}
|
||||
)
|
||||
elif counter_key.startswith("spend:tag:"):
|
||||
tag_name = counter_key[len("spend:tag:") :]
|
||||
row = await prisma_client.db.litellm_tagtable.find_unique(
|
||||
where={"tag_name": tag_name}
|
||||
)
|
||||
elif counter_key.startswith("spend:org:"):
|
||||
org_id = counter_key[len("spend:org:") :]
|
||||
row = await prisma_client.db.litellm_organizationtable.find_unique(
|
||||
|
||||
@@ -54,6 +54,8 @@ async def reserve_budget_for_request(
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
end_user_id: Optional[str] = None,
|
||||
end_user_object: Optional[Any] = None,
|
||||
) -> Optional[dict]:
|
||||
if valid_token is None or not RouteChecks.is_llm_api_route(route=route):
|
||||
return None
|
||||
@@ -63,12 +65,15 @@ async def reserve_budget_for_request(
|
||||
return None
|
||||
|
||||
counters = await _get_budget_counters(
|
||||
request_body=request_body,
|
||||
valid_token=valid_token,
|
||||
team_object=team_object,
|
||||
user_object=user_object,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
end_user_id=end_user_id,
|
||||
end_user_object=end_user_object,
|
||||
)
|
||||
if not counters:
|
||||
return None
|
||||
@@ -166,12 +171,15 @@ async def invalidate_budget_reservation_counters(
|
||||
|
||||
|
||||
async def _get_budget_counters(
|
||||
request_body: dict,
|
||||
valid_token: UserAPIKeyAuth,
|
||||
team_object: Optional[LiteLLM_TeamTable],
|
||||
user_object: Optional[LiteLLM_UserTable],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
end_user_id: Optional[str] = None,
|
||||
end_user_object: Optional[Any] = None,
|
||||
) -> List[_BudgetCounter]:
|
||||
counters: List[_BudgetCounter] = []
|
||||
|
||||
@@ -236,6 +244,25 @@ async def _get_budget_counters(
|
||||
)
|
||||
)
|
||||
|
||||
end_user_counter = await _get_end_user_budget_counter(
|
||||
valid_token=valid_token,
|
||||
end_user_id=end_user_id,
|
||||
end_user_object=end_user_object,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
if end_user_counter is not None:
|
||||
counters.append(end_user_counter)
|
||||
|
||||
counters.extend(
|
||||
await _get_tag_budget_counters(
|
||||
request_body=request_body,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
)
|
||||
|
||||
team_member_counter = await _get_team_member_budget_counter(
|
||||
valid_token=valid_token,
|
||||
team_object=team_object,
|
||||
@@ -256,6 +283,217 @@ async def _get_budget_counters(
|
||||
return counters
|
||||
|
||||
|
||||
async def _get_end_user_budget_counter(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
end_user_id: Optional[str],
|
||||
end_user_object: Optional[Any],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
) -> Optional[_BudgetCounter]:
|
||||
end_user_id = end_user_id or valid_token.end_user_id
|
||||
if end_user_id is None:
|
||||
return None
|
||||
|
||||
source_cache_key = f"end_user_id:{end_user_id}"
|
||||
end_user_obj = end_user_object or (
|
||||
await _get_end_user_from_cache_or_db(
|
||||
end_user_id=end_user_id,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
)
|
||||
|
||||
max_budget = _to_float(valid_token.end_user_max_budget)
|
||||
fallback_spend = 0.0
|
||||
if end_user_obj is not None:
|
||||
fallback_spend = _to_float(_get_value(end_user_obj, "spend")) or 0.0
|
||||
if max_budget is None:
|
||||
budget_table = _get_value(end_user_obj, "litellm_budget_table")
|
||||
max_budget = _to_float(_get_value(budget_table, "max_budget"))
|
||||
|
||||
if max_budget is None:
|
||||
max_budget = await _get_default_end_user_max_budget(
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
)
|
||||
|
||||
if max_budget is None or max_budget <= 0:
|
||||
return None
|
||||
|
||||
return _BudgetCounter(
|
||||
counter_key=f"spend:end_user:{end_user_id}",
|
||||
source_cache_key=source_cache_key,
|
||||
max_budget=max_budget,
|
||||
fallback_spend=fallback_spend,
|
||||
entity_type="EndUser",
|
||||
entity_id=end_user_id,
|
||||
)
|
||||
|
||||
|
||||
async def _get_end_user_from_cache_or_db(
|
||||
end_user_id: str,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
) -> Optional[Any]:
|
||||
cache_key = f"end_user_id:{end_user_id}"
|
||||
cached_end_user = await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
if cached_end_user is not None:
|
||||
return cached_end_user
|
||||
|
||||
if prisma_client is None:
|
||||
return None
|
||||
|
||||
try:
|
||||
row = await prisma_client.db.litellm_endusertable.find_unique(
|
||||
where={"user_id": end_user_id},
|
||||
include={"litellm_budget_table": True, "object_permission": True},
|
||||
)
|
||||
except Exception:
|
||||
verbose_proxy_logger.debug(
|
||||
"Unable to fetch end-user budget for reservation", exc_info=True
|
||||
)
|
||||
return None
|
||||
|
||||
if row is None:
|
||||
return None
|
||||
|
||||
row_dict = _object_to_dict(row)
|
||||
if row_dict:
|
||||
await user_api_key_cache.async_set_cache(key=cache_key, value=row_dict)
|
||||
return row_dict
|
||||
return row
|
||||
|
||||
|
||||
async def _get_default_end_user_max_budget(
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
) -> Optional[float]:
|
||||
if litellm.max_end_user_budget_id is None or prisma_client is None:
|
||||
return None
|
||||
|
||||
cache_key = f"default_end_user_budget:{litellm.max_end_user_budget_id}"
|
||||
cached_budget = await user_api_key_cache.async_get_cache(key=cache_key)
|
||||
max_budget = _to_float(_get_value(cached_budget, "max_budget"))
|
||||
if max_budget is not None:
|
||||
return max_budget
|
||||
|
||||
try:
|
||||
budget_record = await prisma_client.db.litellm_budgettable.find_unique(
|
||||
where={"budget_id": litellm.max_end_user_budget_id}
|
||||
)
|
||||
except Exception:
|
||||
verbose_proxy_logger.debug(
|
||||
"Unable to fetch default end-user budget for reservation", exc_info=True
|
||||
)
|
||||
return None
|
||||
|
||||
if budget_record is None:
|
||||
return None
|
||||
|
||||
budget_dict = _object_to_dict(budget_record)
|
||||
if budget_dict:
|
||||
await user_api_key_cache.async_set_cache(key=cache_key, value=budget_dict)
|
||||
return _to_float(_get_value(budget_record, "max_budget"))
|
||||
|
||||
|
||||
async def _get_tag_budget_counters(
|
||||
request_body: dict,
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> List[_BudgetCounter]:
|
||||
from litellm.proxy.common_utils.http_parsing_utils import get_tags_from_request_body
|
||||
|
||||
tag_names = _dedupe_tags(get_tags_from_request_body(request_body=request_body))
|
||||
if not tag_names:
|
||||
return []
|
||||
|
||||
tag_objects = await _get_tag_objects_for_reservation(
|
||||
tag_names=tag_names,
|
||||
prisma_client=prisma_client,
|
||||
user_api_key_cache=user_api_key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
counters: List[_BudgetCounter] = []
|
||||
for tag_name in tag_names:
|
||||
tag_object = tag_objects.get(tag_name)
|
||||
if tag_object is None:
|
||||
continue
|
||||
budget_table = _get_value(tag_object, "litellm_budget_table")
|
||||
max_budget = _to_float(_get_value(budget_table, "max_budget"))
|
||||
if max_budget is None or max_budget <= 0:
|
||||
continue
|
||||
counters.append(
|
||||
_BudgetCounter(
|
||||
counter_key=f"spend:tag:{tag_name}",
|
||||
source_cache_key=f"tag:{tag_name}",
|
||||
max_budget=max_budget,
|
||||
fallback_spend=_to_float(_get_value(tag_object, "spend")) or 0.0,
|
||||
entity_type="Tag",
|
||||
entity_id=tag_name,
|
||||
)
|
||||
)
|
||||
return counters
|
||||
|
||||
|
||||
async def _get_tag_objects_for_reservation(
|
||||
tag_names: List[str],
|
||||
prisma_client: Optional[PrismaClient],
|
||||
user_api_key_cache: DualCache,
|
||||
proxy_logging_obj: ProxyLogging,
|
||||
) -> Dict[str, Any]:
|
||||
tag_objects: Dict[str, Any] = {}
|
||||
uncached_tags: List[str] = []
|
||||
|
||||
for tag_name in tag_names:
|
||||
cached_tag = await user_api_key_cache.async_get_cache(key=f"tag:{tag_name}")
|
||||
if cached_tag is not None:
|
||||
tag_objects[tag_name] = cached_tag
|
||||
else:
|
||||
uncached_tags.append(tag_name)
|
||||
|
||||
if not uncached_tags or prisma_client is None:
|
||||
return tag_objects
|
||||
|
||||
try:
|
||||
db_tags = await prisma_client.db.litellm_tagtable.find_many(
|
||||
where={"tag_name": {"in": uncached_tags}},
|
||||
include={"litellm_budget_table": True},
|
||||
)
|
||||
except Exception:
|
||||
verbose_proxy_logger.debug(
|
||||
"Unable to fetch tag budgets for reservation", exc_info=True
|
||||
)
|
||||
return tag_objects
|
||||
|
||||
for db_tag in db_tags:
|
||||
tag_name = _get_value(db_tag, "tag_name")
|
||||
if not isinstance(tag_name, str):
|
||||
continue
|
||||
row_dict = _object_to_dict(db_tag)
|
||||
if row_dict:
|
||||
await user_api_key_cache.async_set_cache(
|
||||
key=f"tag:{tag_name}", value=row_dict
|
||||
)
|
||||
tag_objects[tag_name] = row_dict
|
||||
else:
|
||||
tag_objects[tag_name] = db_tag
|
||||
|
||||
return tag_objects
|
||||
|
||||
|
||||
def _dedupe_tags(tags: List[str]) -> List[str]:
|
||||
seen = set()
|
||||
deduped_tags = []
|
||||
for tag in tags:
|
||||
if tag in seen:
|
||||
continue
|
||||
seen.add(tag)
|
||||
deduped_tags.append(tag)
|
||||
return deduped_tags
|
||||
|
||||
|
||||
async def _get_team_member_budget_counter(
|
||||
valid_token: UserAPIKeyAuth,
|
||||
team_object: Optional[LiteLLM_TeamTable],
|
||||
@@ -727,3 +965,15 @@ def _get_value(obj: Any, key: str) -> Any:
|
||||
if isinstance(obj, dict):
|
||||
return obj.get(key)
|
||||
return getattr(obj, key, None)
|
||||
|
||||
|
||||
def _object_to_dict(obj: Any) -> dict:
|
||||
if isinstance(obj, dict):
|
||||
return obj
|
||||
if hasattr(obj, "model_dump"):
|
||||
value = obj.model_dump()
|
||||
return value if isinstance(value, dict) else {}
|
||||
if hasattr(obj, "dict"):
|
||||
value = obj.dict()
|
||||
return value if isinstance(value, dict) else {}
|
||||
return {}
|
||||
|
||||
@@ -17,7 +17,10 @@ import litellm
|
||||
from litellm.proxy._types import (
|
||||
CallInfo,
|
||||
Litellm_EntityType,
|
||||
LiteLLM_BudgetTable,
|
||||
LiteLLM_EndUserTable,
|
||||
LiteLLM_ObjectPermissionTable,
|
||||
LiteLLM_TagTable,
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_UserTable,
|
||||
LitellmUserRoles,
|
||||
@@ -29,10 +32,12 @@ from litellm.proxy._types import (
|
||||
from litellm.proxy.auth.auth_checks import (
|
||||
ExperimentalUIJWTToken,
|
||||
_can_object_call_vector_stores,
|
||||
_check_end_user_budget,
|
||||
_check_team_member_budget,
|
||||
_get_fuzzy_user_object,
|
||||
_get_team_db_check,
|
||||
_log_budget_lookup_failure,
|
||||
_tag_max_budget_check,
|
||||
_team_max_budget_check,
|
||||
_virtual_key_max_budget_alert_check,
|
||||
_virtual_key_max_budget_check,
|
||||
@@ -1962,6 +1967,67 @@ async def test_team_budget_check_reads_from_spend_counter():
|
||||
assert exc_info.value.current_cost == 1.5
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_end_user_budget_check_reads_from_spend_counter():
|
||||
"""End-user budget check should use get_current_spend when counter exists."""
|
||||
end_user_object = LiteLLM_EndUserTable(
|
||||
user_id="customer-1",
|
||||
blocked=False,
|
||||
spend=0.0,
|
||||
litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0),
|
||||
)
|
||||
|
||||
async def mock_get_current_spend(counter_key, fallback_spend):
|
||||
if counter_key == "spend:end_user:customer-1":
|
||||
return 1.5
|
||||
return fallback_spend
|
||||
|
||||
with patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend):
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _check_end_user_budget(
|
||||
end_user_obj=end_user_object,
|
||||
route="/chat/completions",
|
||||
)
|
||||
assert exc_info.value.current_cost == 1.5
|
||||
assert exc_info.value.max_budget == 1.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tag_budget_check_reads_from_spend_counter():
|
||||
"""Tag budget check should use get_current_spend when counter exists."""
|
||||
from litellm.proxy.utils import ProxyLogging
|
||||
|
||||
tag_object = LiteLLM_TagTable(
|
||||
tag_name="paid-tag",
|
||||
spend=0.0,
|
||||
litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0),
|
||||
)
|
||||
|
||||
async def mock_get_current_spend(counter_key, fallback_spend):
|
||||
if counter_key == "spend:tag:paid-tag":
|
||||
return 1.5
|
||||
return fallback_spend
|
||||
|
||||
with (
|
||||
patch("litellm.proxy.proxy_server.get_current_spend", mock_get_current_spend),
|
||||
patch(
|
||||
"litellm.proxy.auth.auth_checks.get_tag_objects_batch",
|
||||
new_callable=AsyncMock,
|
||||
return_value={"paid-tag": tag_object},
|
||||
),
|
||||
):
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _tag_max_budget_check(
|
||||
request_body={"metadata": {"tags": ["paid-tag"]}},
|
||||
prisma_client=MagicMock(),
|
||||
user_api_key_cache=MagicMock(),
|
||||
proxy_logging_obj=ProxyLogging(user_api_key_cache=None),
|
||||
valid_token=UserAPIKeyAuth(token="test-token"),
|
||||
)
|
||||
assert exc_info.value.current_cost == 1.5
|
||||
assert exc_info.value.max_budget == 1.0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_member_budget_check_reads_from_spend_counter():
|
||||
"""Team member budget check should use get_current_spend when counter exists."""
|
||||
|
||||
@@ -13,6 +13,8 @@ import litellm.proxy.proxy_server
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_JWTAuth,
|
||||
LiteLLM_BudgetTable,
|
||||
LiteLLM_EndUserTable,
|
||||
LiteLLM_UserTable,
|
||||
LitellmUserRoles,
|
||||
ProxyErrorTypes,
|
||||
@@ -2150,6 +2152,80 @@ async def test_centralized_common_checks_propagates_end_user_budget_error():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_centralized_common_checks_reserves_request_end_user_budget():
|
||||
"""Regression: reservation runs before user_api_key_auth() copies the
|
||||
request end-user onto the token, so centralized checks must pass the
|
||||
locally extracted end_user_id/end_user_object into reservation."""
|
||||
import litellm.proxy.proxy_server as _proxy_server_mod
|
||||
from fastapi import Request
|
||||
from starlette.datastructures import URL
|
||||
|
||||
token = UserAPIKeyAuth(api_key="sk-test", user_id="u")
|
||||
request = Request(scope={"type": "http", "headers": []})
|
||||
request._url = URL(url="/chat/completions")
|
||||
request_data = {
|
||||
"model": "gpt-4o",
|
||||
"messages": [{"role": "user", "content": "hello"}],
|
||||
"user": "alice",
|
||||
}
|
||||
end_user_object = LiteLLM_EndUserTable(
|
||||
user_id="alice",
|
||||
blocked=False,
|
||||
spend=0.0,
|
||||
litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0),
|
||||
)
|
||||
|
||||
attrs = _proxy_attrs_for_centralized_checks(user_custom_auth=None)
|
||||
counter_cache = DualCache()
|
||||
attrs["spend_counter_cache"] = counter_cache
|
||||
originals = {a: getattr(_proxy_server_mod, a, None) for a in attrs}
|
||||
try:
|
||||
for k, v in attrs.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
with (
|
||||
patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.get_end_user_object",
|
||||
new_callable=AsyncMock,
|
||||
return_value=end_user_object,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.auth.user_api_key_auth.common_checks",
|
||||
new_callable=AsyncMock,
|
||||
),
|
||||
patch(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
|
||||
return_value=0.6,
|
||||
),
|
||||
):
|
||||
assert token.end_user_id is None
|
||||
|
||||
await _run_centralized_common_checks(
|
||||
user_api_key_auth_obj=token,
|
||||
request=request,
|
||||
request_data=request_data,
|
||||
route="/chat/completions",
|
||||
)
|
||||
|
||||
finally:
|
||||
for k, v in originals.items():
|
||||
setattr(_proxy_server_mod, k, v)
|
||||
|
||||
assert token.end_user_id is None
|
||||
assert token.budget_reservation is not None
|
||||
assert token.budget_reservation["entries"] == [
|
||||
{
|
||||
"counter_key": "spend:end_user:alice",
|
||||
"entity_type": "EndUser",
|
||||
"entity_id": "alice",
|
||||
"applied_adjustment": 0.0,
|
||||
}
|
||||
]
|
||||
assert counter_cache.in_memory_cache.get_cache(
|
||||
key="spend:end_user:alice"
|
||||
) == pytest.approx(0.6)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_centralized_common_checks_short_circuits_when_master_key_unset():
|
||||
"""master_key=None is no-auth dev mode — admin-only routes and
|
||||
|
||||
@@ -6,7 +6,9 @@ import litellm
|
||||
from litellm.caching.dual_cache import DualCache
|
||||
from litellm.proxy._types import (
|
||||
LiteLLM_BudgetTable,
|
||||
LiteLLM_EndUserTable,
|
||||
LiteLLM_OrganizationTable,
|
||||
LiteLLM_TagTable,
|
||||
LiteLLM_TeamMembership,
|
||||
LiteLLM_TeamTable,
|
||||
LiteLLM_UserTable,
|
||||
@@ -118,6 +120,171 @@ async def test_should_prevent_second_key_reservation_over_budget(
|
||||
await release_budget_reservation(reservation)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_prevent_second_end_user_reservation_over_budget(
|
||||
spend_counter_state,
|
||||
):
|
||||
counter_cache, key_cache = spend_counter_state
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
|
||||
valid_token = UserAPIKeyAuth(
|
||||
token="key-budget-end-user",
|
||||
end_user_id="end-user-budget-race",
|
||||
end_user_max_budget=1.0,
|
||||
)
|
||||
await key_cache.async_set_cache(
|
||||
key="end_user_id:end-user-budget-race",
|
||||
value=LiteLLM_EndUserTable(
|
||||
user_id="end-user-budget-race",
|
||||
blocked=False,
|
||||
spend=0.0,
|
||||
litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0),
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
|
||||
return_value=0.6,
|
||||
):
|
||||
reservation = await reserve_budget_for_request(
|
||||
request_body=_request_body(),
|
||||
route="/chat/completions",
|
||||
llm_router=None,
|
||||
valid_token=valid_token,
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
assert reservation is not None
|
||||
assert counter_cache.in_memory_cache.get_cache(
|
||||
key="spend:end_user:end-user-budget-race"
|
||||
) == pytest.approx(0.6)
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await reserve_budget_for_request(
|
||||
request_body=_request_body(),
|
||||
route="/chat/completions",
|
||||
llm_router=None,
|
||||
valid_token=valid_token,
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
assert counter_cache.in_memory_cache.get_cache(
|
||||
key="spend:end_user:end-user-budget-race"
|
||||
) == pytest.approx(0.6)
|
||||
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
await increment_spend_counters(
|
||||
token=None,
|
||||
team_id=None,
|
||||
user_id=None,
|
||||
response_cost=0.2,
|
||||
budget_reservation=reservation,
|
||||
)
|
||||
|
||||
assert counter_cache.in_memory_cache.get_cache(
|
||||
key="spend:end_user:end-user-budget-race"
|
||||
) == pytest.approx(0.2)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_prevent_second_tag_reservation_over_budget(
|
||||
spend_counter_state,
|
||||
):
|
||||
counter_cache, key_cache = spend_counter_state
|
||||
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
|
||||
valid_token = UserAPIKeyAuth(token="key-budget-tag")
|
||||
request_body = _request_body()
|
||||
request_body["metadata"] = {
|
||||
"tags": ["tag-budget-race", "tag-without-budget", "tag-budget-race"]
|
||||
}
|
||||
await key_cache.async_set_cache(
|
||||
key="tag:tag-budget-race",
|
||||
value=LiteLLM_TagTable(
|
||||
tag_name="tag-budget-race",
|
||||
spend=0.0,
|
||||
budget_id="tag-budget-id",
|
||||
litellm_budget_table=LiteLLM_BudgetTable(max_budget=1.0),
|
||||
).model_dump(),
|
||||
)
|
||||
await key_cache.async_set_cache(
|
||||
key="tag:tag-without-budget",
|
||||
value=LiteLLM_TagTable(
|
||||
tag_name="tag-without-budget",
|
||||
spend=0.0,
|
||||
).model_dump(),
|
||||
)
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.spend_tracking.budget_reservation.estimate_request_max_cost",
|
||||
return_value=0.6,
|
||||
):
|
||||
reservation = await reserve_budget_for_request(
|
||||
request_body=request_body,
|
||||
route="/chat/completions",
|
||||
llm_router=None,
|
||||
valid_token=valid_token,
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
assert reservation is not None
|
||||
assert reservation["entries"] == [
|
||||
{
|
||||
"counter_key": "spend:tag:tag-budget-race",
|
||||
"entity_type": "Tag",
|
||||
"entity_id": "tag-budget-race",
|
||||
"applied_adjustment": 0.0,
|
||||
}
|
||||
]
|
||||
assert counter_cache.in_memory_cache.get_cache(
|
||||
key="spend:tag:tag-budget-race"
|
||||
) == pytest.approx(0.6)
|
||||
assert (
|
||||
counter_cache.in_memory_cache.get_cache(key="spend:tag:tag-without-budget")
|
||||
is None
|
||||
)
|
||||
|
||||
with pytest.raises(litellm.BudgetExceededError):
|
||||
await reserve_budget_for_request(
|
||||
request_body=request_body,
|
||||
route="/chat/completions",
|
||||
llm_router=None,
|
||||
valid_token=valid_token,
|
||||
team_object=None,
|
||||
user_object=None,
|
||||
prisma_client=None,
|
||||
user_api_key_cache=key_cache,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
|
||||
assert counter_cache.in_memory_cache.get_cache(
|
||||
key="spend:tag:tag-budget-race"
|
||||
) == pytest.approx(0.6)
|
||||
|
||||
from litellm.proxy.proxy_server import increment_spend_counters
|
||||
|
||||
await increment_spend_counters(
|
||||
token=None,
|
||||
team_id=None,
|
||||
user_id=None,
|
||||
response_cost=0.2,
|
||||
budget_reservation=reservation,
|
||||
)
|
||||
|
||||
assert counter_cache.in_memory_cache.get_cache(
|
||||
key="spend:tag:tag-budget-race"
|
||||
) == pytest.approx(0.2)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_reserve_team_member_and_org_budget_counters(spend_counter_state):
|
||||
counter_cache, key_cache = spend_counter_state
|
||||
|
||||
@@ -5058,11 +5058,19 @@ async def test_reseed_spend_from_db_user_and_org_prefixes():
|
||||
|
||||
user_row = MagicMock()
|
||||
user_row.spend = 17.0
|
||||
end_user_row = MagicMock()
|
||||
end_user_row.spend = 21.0
|
||||
tag_row = MagicMock()
|
||||
tag_row.spend = 8.0
|
||||
org_row = MagicMock()
|
||||
org_row.spend = 305.0
|
||||
|
||||
fake_prisma = MagicMock()
|
||||
fake_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row)
|
||||
fake_prisma.db.litellm_endusertable.find_unique = AsyncMock(
|
||||
return_value=end_user_row
|
||||
)
|
||||
fake_prisma.db.litellm_tagtable.find_unique = AsyncMock(return_value=tag_row)
|
||||
fake_prisma.db.litellm_organizationtable.find_unique = AsyncMock(
|
||||
return_value=org_row
|
||||
)
|
||||
@@ -5072,6 +5080,19 @@ async def test_reseed_spend_from_db_user_and_org_prefixes():
|
||||
where={"user_id": "alice"}
|
||||
)
|
||||
|
||||
assert (
|
||||
await SpendCounterReseed.from_db(fake_prisma, "spend:end_user:customer-1")
|
||||
== 21.0
|
||||
)
|
||||
fake_prisma.db.litellm_endusertable.find_unique.assert_awaited_once_with(
|
||||
where={"user_id": "customer-1"}
|
||||
)
|
||||
|
||||
assert await SpendCounterReseed.from_db(fake_prisma, "spend:tag:paid-tag") == 8.0
|
||||
fake_prisma.db.litellm_tagtable.find_unique.assert_awaited_once_with(
|
||||
where={"tag_name": "paid-tag"}
|
||||
)
|
||||
|
||||
assert await SpendCounterReseed.from_db(fake_prisma, "spend:org:acme") == 305.0
|
||||
fake_prisma.db.litellm_organizationtable.find_unique.assert_awaited_once_with(
|
||||
where={"organization_id": "acme"}
|
||||
|
||||
Reference in New Issue
Block a user