diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 840f64cfed..a798a9fe9f 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -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}", ) diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 0d60b36d35..b21c648398 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -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, ) diff --git a/litellm/proxy/db/spend_counter_reseed.py b/litellm/proxy/db/spend_counter_reseed.py index b4dc4a4d0a..1a2f64d988 100644 --- a/litellm/proxy/db/spend_counter_reseed.py +++ b/litellm/proxy/db/spend_counter_reseed.py @@ -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( diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 18ef6dff0e..7bb9c28522 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -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 {} diff --git a/tests/test_litellm/proxy/auth/test_auth_checks.py b/tests/test_litellm/proxy/auth/test_auth_checks.py index 676a32c202..c50278a492 100644 --- a/tests/test_litellm/proxy/auth/test_auth_checks.py +++ b/tests/test_litellm/proxy/auth/test_auth_checks.py @@ -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.""" diff --git a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py index bcd56cf872..96a900ffbf 100644 --- a/tests/test_litellm/proxy/auth/test_user_api_key_auth.py +++ b/tests/test_litellm/proxy/auth/test_user_api_key_auth.py @@ -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 diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index c14582638c..57dc894989 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -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 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index 8f3297780e..ed03f3431c 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -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"}