scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-21 00:21:49 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 74

Merge pull request #41354 from BerriAI/litellm_lit_3269_project_spend_tracking

Browse source 
fix(proxy): track project spend and enforce project budgets additively
This commit is contained in:
ryan-crabbe-berri 2026-09-18 17:17:28 -07:00 • committed by GitHub
commit 1073b9eff7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
23 changed files with 630 additions and 35 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/_types.py
View file

@ -5403,6 +5403,7 @@ class DBSpendUpdateTransactions(TypedDict):
team_member_list_transactions: dict[str, float] | None
org_list_transactions: dict[str, float] | None
org_member_list_transactions: ReadOnly[dict[str, float] | None]
project_list_transactions: ReadOnly[dict[str, float] | None]
tag_list_transactions: dict[str, float] | None
agent_list_transactions: dict[str, float] | None
model_access_group_list_transactions: ReadOnly[dict[str, float] | None]

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

@ -94,6 +94,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
model_access_group_registry_cache_key,
model_access_group_spend_counter_key,
object_permission_cache_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
tag_registry_cache_key,
team_membership_auth_cache_key,
@ -5680,16 +5682,22 @@ async def _project_max_budget_check(
if project_object.litellm_budget_table is not None:
max_budget = project_object.litellm_budget_table.max_budget
if (
max_budget is not None
and project_object.spend is not None
and math.isfinite(max_budget)
and project_object.spend > max_budget
):
if max_budget is None or max_budget <= 0 or not math.isfinite(max_budget):
return
from litellm.proxy.proxy_server import get_current_spend
project_spend: Final = await get_current_spend(
counter_key=project_spend_counter_key(project_object.project_id),
fallback_spend=project_object.spend or 0.0,
max_budget=max_budget,
)
if project_spend >= max_budget:
if valid_token:
call_info: Final = CallInfo(
token=valid_token.token,
spend=project_object.spend,
spend=project_spend,
max_budget=max_budget,
user_id=valid_token.user_id,
team_id=valid_token.team_id,
@ -5705,9 +5713,9 @@ async def _project_max_budget_check(
)
raise litellm.BudgetExceededError(
current_cost=project_object.spend,
current_cost=project_spend,
max_budget=max_budget,
message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_object.spend}, Max budget: {max_budget}",
message=f"Budget has been exceeded! Project={project_object.project_id} Current cost: {project_spend}, Max budget: {max_budget}",
entity_type=Litellm_EntityType.PROJECT.value,
entity_id=project_object.project_id,
)
@ -5757,10 +5765,6 @@ async def _project_soft_budget_check(
)
def _project_cache_key(project_id: str) -> str:
return f"project_id:{project_id}"
async def get_project_object(
project_id: str,
prisma_client: PrismaClient | None,
@ -5778,7 +5782,7 @@ async def get_project_object(
return None
# Check cache first
cache_key: Final = _project_cache_key(project_id)
cache_key: Final = project_cache_key(project_id)
deserialized_project: Final = await user_api_key_cache.async_get_cache(
key=cache_key,
model_type=LiteLLM_ProjectTableCachedObj,
@ -5820,7 +5824,7 @@ async def delete_cached_project_object(
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import evict_and_broadcast
await evict_and_broadcast(
cache_keys=(_project_cache_key(project_id),),
cache_keys=(project_cache_key(project_id),),
user_api_key_cache=user_api_key_cache,
)

24
litellm/proxy/common_utils/reset_budget_job.py
View file

@ -40,6 +40,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
end_user_cache_key,
model_access_group_cache_key,
model_access_group_spend_counter_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
)
from litellm.proxy.db.budget_window_spend_writer import roll_window_spend_row
@ -48,6 +50,7 @@ from litellm.proxy.db.exception_handler import call_with_db_reconnect_retry
from litellm.proxy.utils import PrismaClient, ProxyLogging
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.prisma_protocols import SpendLinkedTable
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import (
EndUserRepository,
ModelAccessGroupBudgetRepository,
@ -114,6 +117,11 @@ class _ModelAccessGroupRow(_BudgetLinkedRow, Protocol):
def access_group_name(self) -> str: ...
class _ProjectRow(_BudgetLinkedRow, Protocol):
@property
def project_id(self) -> str: ...
class _EndUserRow(_BudgetLinkedRow, Protocol):
@property
def user_id(self) -> str: ...
@ -184,6 +192,14 @@ def _model_access_group_cache_keys(row: _ModelAccessGroupRow) -> tuple[str, ...]
return (model_access_group_cache_key(row.access_group_name),)
def _project_counter_key(row: _ProjectRow) -> str:
return project_spend_counter_key(row.project_id)
def _project_cache_keys(row: _ProjectRow) -> tuple[str, ...]:
return (project_cache_key(row.project_id),)
def _enduser_counter_key(row: _EndUserRow) -> str:
return f"spend:end_user:{row.user_id}"
@ -754,6 +770,11 @@ class ResetBudgetJob:
where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE),
log_subject="model access groups",
)
projects: Final[tuple[_ProjectRow, ...]] = await self._fetch_linked_rows(
table=ProjectRepository(self.prisma_client).table,
where=_budget_link_where(budget_ids, _SPENT_ROWS_WHERE),
log_subject="projects",
)
rollover_caps: Final[Mapping[str, float]] = MappingProxyType(
{ # mutable-ok: MappingProxyType wraps a one-shot dict comprehension
b.budget_id: cap
@ -786,6 +807,7 @@ class ResetBudgetJob:
(_model_access_group_counter_key(row), _row_carried_spend(row, rollover_caps))
for row in model_access_groups
),
*((_project_counter_key(row), _row_carried_spend(row, rollover_caps)) for row in projects),
),
rollover_caps=rollover_caps,
cache_keys=(
@ -794,6 +816,7 @@ class ResetBudgetJob:
*(key for row in orgs for key in _org_cache_keys(row)),
*(key for row in tags for key in _tag_cache_keys(row)),
*(key for row in model_access_groups for key in _model_access_group_cache_keys(row)),
*(key for row in projects for key in _project_cache_keys(row)),
),
)
@ -820,6 +843,7 @@ class ResetBudgetJob:
_queue_budget_linked_resets(uow.organizations, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.tags, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.model_access_groups, cascade, extra=_SPENT_ROWS_WHERE)
_queue_budget_linked_resets(uow.projects, cascade, extra=_SPENT_ROWS_WHERE)
_queue_enduser_resets(uow.endusers, cascade)
for budget_id, budget_reset_at in cascade.budget_resets:
uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at)

8
litellm/proxy/common_utils/user_api_key_cache.py
View file

@ -325,6 +325,14 @@ def model_access_group_spend_counter_key(access_group_name: str) -> str:
return f"spend:model_access_group:{access_group_name}"
def project_cache_key(project_id: str) -> str:
return f"project_id:{project_id}"
def project_spend_counter_key(project_id: str) -> str:
return f"spend:project:{project_id}"
#: Cached under ``end_user_restricted_registry_cache_key`` when the restricted set exceeds
#: ``END_USER_RESTRICTED_REGISTRY_MAX_SIZE``: registry unusable, fall back to the per-id fetch.
END_USER_RESTRICTED_REGISTRY_OVERFLOW_SENTINEL: Final = "__end_user_restricted_registry_overflow__"

98
litellm/proxy/db/db_spend_update_writer.py
View file

@ -46,6 +46,7 @@ from litellm.proxy._types import (
SpendUpdateQueueItem,
ToolDiscoveryQueueItem,
)
from litellm.proxy.common_utils.user_api_key_cache import project_cache_key
from litellm.proxy.db.daily_spend_bulk_upsert import (
DAILY_SPEND_TABLES,
build_bulk_upsert,
@ -122,6 +123,7 @@ class _SpendBatch(Protocol):
litellm_teammembership: BatchTable
litellm_organizationtable: BatchTable
litellm_organizationmembership: BatchTable
litellm_projecttable: BatchTable
litellm_tagtable: BatchTable
litellm_agentstable: BatchTable
litellm_modelaccessgroupbudgettable: BatchTable
@ -300,6 +302,7 @@ class DBSpendUpdateWriter:
start_time: datetime,
end_time: datetime,
response_cost: float | None,
project_id: str | None = None,
) -> bool:
"""Record the request's spend, answering whether its cost still needs charging.
@ -382,6 +385,7 @@ class DBSpendUpdateWriter:
hashed_token=hashed_token,
team_id=team_id,
org_id=org_id,
project_id=project_id,
end_user_id=end_user_id,
prisma_client=prisma_client,
litellm_proxy_budget_name=litellm_proxy_budget_name,
@ -678,6 +682,7 @@ class DBSpendUpdateWriter:
litellm_proxy_budget_name: str | None,
payload: SpendLogsPayload,
request_model_access_groups: Sequence[str] = (),
project_id: str | None = None,
):
"""
Runs all 13 spend-update helpers sequentially inside a single asyncio task.
@ -741,6 +746,18 @@ class DBSpendUpdateWriter:
traceback.format_exc(),
)
try:
await self._update_project_db(
response_cost=response_cost,
project_id=project_id,
prisma_client=prisma_client,
)
except Exception: # noqa: BLE001 # a project enqueue failure must not skip the sibling spend writes
verbose_proxy_logger.debug(
"_batch_database_updates: _update_project_db failed: %s",
traceback.format_exc(),
)
try:
await self._update_tag_db(
response_cost=response_cost,
@ -1003,6 +1020,32 @@ class DBSpendUpdateWriter:
)
raise e
async def _update_project_db(
self,
response_cost: float | None,
project_id: str | None,
prisma_client: PrismaClient | None,
) -> None:
if project_id is None or prisma_client is None:
return
try:
await self.spend_update_queue.add_update(
update=SpendUpdateQueueItem(
entity_type=Litellm_EntityType.PROJECT,
entity_id=project_id,
response_cost=response_cost,
)
)
except Exception as e:
spend_log_error(
"Spend tracking - failed to enqueue project spend update. project_id=%s, response_cost=%s - %s",
project_id,
response_cost,
str(e),
exc=e,
)
raise e
async def _update_agent_db(
self,
response_cost: float | None,
@ -1240,18 +1283,19 @@ class DBSpendUpdateWriter:
if db_spend_update_transactions is not None:
verbose_proxy_logger.info(
"Spend tracking - committing spend updates from Redis to DB: "
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, tags=%d, "
"agents=%d, model_access_groups=%d",
len(db_spend_update_transactions.get("key_list_transactions") or {}),
len(db_spend_update_transactions.get("user_list_transactions") or {}),
len(db_spend_update_transactions.get("team_list_transactions") or {}),
len(db_spend_update_transactions.get("org_list_transactions") or {}),
len(db_spend_update_transactions.get("end_user_list_transactions") or {}),
len(db_spend_update_transactions.get("team_member_list_transactions") or {}),
len(db_spend_update_transactions.get("org_member_list_transactions") or {}),
len(db_spend_update_transactions.get("tag_list_transactions") or {}),
len(db_spend_update_transactions.get("agent_list_transactions") or {}),
len(db_spend_update_transactions.get("model_access_group_list_transactions") or {}),
"keys=%d, users=%d, teams=%d, orgs=%d, end_users=%d, team_members=%d, org_members=%d, "
"projects=%d, tags=%d, agents=%d, model_access_groups=%d",
len(db_spend_update_transactions.get("key_list_transactions") or ()),
len(db_spend_update_transactions.get("user_list_transactions") or ()),
len(db_spend_update_transactions.get("team_list_transactions") or ()),
len(db_spend_update_transactions.get("org_list_transactions") or ()),
len(db_spend_update_transactions.get("end_user_list_transactions") or ()),
len(db_spend_update_transactions.get("team_member_list_transactions") or ()),
len(db_spend_update_transactions.get("org_member_list_transactions") or ()),
len(db_spend_update_transactions.get("project_list_transactions") or ()),
len(db_spend_update_transactions.get("tag_list_transactions") or ()),
len(db_spend_update_transactions.get("agent_list_transactions") or ()),
len(db_spend_update_transactions.get("model_access_group_list_transactions") or ()),
)
await self._commit_spend_updates_to_db(
prisma_client=prisma_client,
@ -1797,6 +1841,22 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE PROJECT TABLE ###
project_list_transactions: Final = db_spend_update_transactions.get("project_list_transactions")
await DBSpendUpdateWriter._update_entity_spend_in_db(
entity_name="Project",
transactions=project_list_transactions,
table_accessor="litellm_projecttable",
where_field="project_id",
n_retry_times=n_retry_times,
prisma_client=prisma_client,
proxy_logging_obj=proxy_logging_obj,
)
await DBSpendUpdateWriter._invalidate_project_caches(
project_ids=tuple(project_list_transactions or ()),
proxy_logging_obj=proxy_logging_obj,
)
### UPDATE TAG TABLE ###
tag_list_transactions: Final = db_spend_update_transactions["tag_list_transactions"]
await DBSpendUpdateWriter._update_entity_spend_in_db(
@ -1835,11 +1895,23 @@ class DBSpendUpdateWriter:
proxy_logging_obj=proxy_logging_obj,
)
@staticmethod
async def _invalidate_project_caches(project_ids: Sequence[str], proxy_logging_obj: ProxyLogging | None) -> None:
if not project_ids or proxy_logging_obj is None:
return
user_api_key_cache: Final = proxy_logging_obj.call_details.get("user_api_key_cache")
if user_api_key_cache is None:
return
for project_id in project_ids:
await user_api_key_cache.async_delete_cache(key=project_cache_key(project_id))
@staticmethod
async def _update_entity_spend_in_db(
entity_name: str,
transactions: dict[str, float] | None,
table_accessor: Literal["litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable"],
table_accessor: Literal[
"litellm_tagtable", "litellm_agentstable", "litellm_modelaccessgroupbudgettable", "litellm_projecttable"
],
where_field: str,
n_retry_times: int,
prisma_client: PrismaClient,

7
litellm/proxy/db/db_transaction_queue/redis_update_buffer.py
View file

@ -70,6 +70,7 @@ _SpendTransactionField: TypeAlias = Literal[
"team_member_list_transactions",
"org_list_transactions",
"org_member_list_transactions",
"project_list_transactions",
"tag_list_transactions",
"agent_list_transactions",
"model_access_group_list_transactions",
@ -83,6 +84,7 @@ _SPEND_TRANSACTION_FIELDS: Final[tuple[_SpendTransactionField, ...]] = (
"team_member_list_transactions",
"org_list_transactions",
"org_member_list_transactions",
"project_list_transactions",
"tag_list_transactions",
"agent_list_transactions",
"model_access_group_list_transactions",
@ -418,6 +420,10 @@ class RedisUpdateBuffer:
Litellm_EntityType.ORGANIZATION_MEMBER,
db_spend_update_transactions.get("org_member_list_transactions"),
),
(
Litellm_EntityType.PROJECT,
db_spend_update_transactions.get("project_list_transactions"),
),
(
Litellm_EntityType.TAG,
db_spend_update_transactions.get("tag_list_transactions"),
@ -885,6 +891,7 @@ class RedisUpdateBuffer:
org_member_list_transactions=_merged_entity_transactions(
list_of_transactions, "org_member_list_transactions"
),
project_list_transactions=_merged_entity_transactions(list_of_transactions, "project_list_transactions"),
tag_list_transactions=_merged_entity_transactions(list_of_transactions, "tag_list_transactions"),
agent_list_transactions=_merged_entity_transactions(list_of_transactions, "agent_list_transactions"),
model_access_group_list_transactions=_merged_entity_transactions(

4
litellm/proxy/db/db_transaction_queue/spend_update_queue.py
View file

@ -138,6 +138,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
team_member_list_transactions={},
org_list_transactions={},
org_member_list_transactions={},
project_list_transactions={},
tag_list_transactions={},
agent_list_transactions={},
model_access_group_list_transactions={},
@ -152,6 +153,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions",
Litellm_EntityType.ORGANIZATION: "org_list_transactions",
Litellm_EntityType.ORGANIZATION_MEMBER: "org_member_list_transactions",
Litellm_EntityType.PROJECT: "project_list_transactions",
Litellm_EntityType.TAG: "tag_list_transactions",
Litellm_EntityType.AGENT: "agent_list_transactions",
Litellm_EntityType.MODEL_ACCESS_GROUP: "model_access_group_list_transactions",
@ -192,6 +194,8 @@ class SpendUpdateQueue(BaseUpdateQueue):
transactions_dict = db_spend_update_transactions["org_list_transactions"]
elif dict_key == "org_member_list_transactions":
transactions_dict = db_spend_update_transactions["org_member_list_transactions"]
elif dict_key == "project_list_transactions":
transactions_dict = db_spend_update_transactions["project_list_transactions"]
elif dict_key == "tag_list_transactions":
transactions_dict = db_spend_update_transactions["tag_list_transactions"]
elif dict_key == "agent_list_transactions":

5
litellm/proxy/db/spend_counter_reseed.py
View file

@ -26,6 +26,7 @@ from litellm.proxy._types import Litellm_EntityType
from litellm.proxy.db.db_lookup_gate import db_lookup_gate
from litellm.proxy.spend_tracking.spend_counter_batch import read_batched_spend_counter, record_spend_counter_value
from litellm.repositories.organization_repository import OrganizationRepository
from litellm.repositories.project_repository import ProjectRepository
from litellm.repositories.table_repositories import (
BudgetWindowSpendRepository,
EndUserRepository,
@ -77,6 +78,7 @@ class SpendCounterReseed:
spend:team_member:{uid}:{tid} -> LiteLLM_TeamMembership.spend
spend:user:{user_id} -> LiteLLM_UserTable.spend
spend:org:{org_id} -> LiteLLM_OrganizationTable.spend
spend:project:{project_id} -> LiteLLM_ProjectTable.spend
End-user and tag spend counters intentionally do not reseed here. Their
auth paths already load the corresponding objects via get_end_user_object()
@ -157,6 +159,9 @@ class SpendCounterReseed:
row = await OrganizationRepository(prisma_client).table.find_unique(
where={"organization_id": org_id}
)
elif counter_key.startswith("spend:project:"):
project_id: Final = counter_key[len("spend:project:") :]
row = await ProjectRepository(prisma_client).table.find_unique(where={"project_id": project_id})
else:
return None
except Exception:

12
litellm/proxy/hooks/proxy_track_cost_callback.py
View file

@ -267,6 +267,7 @@ class _ProxyDBLogger(CustomLogger):
start_time=actual_start_time,
end_time=datetime.now(),
org_id=user_api_key_dict.org_id,
project_id=user_api_key_dict.project_id,
)
@log_db_metrics
@ -318,6 +319,11 @@ class _ProxyDBLogger(CustomLogger):
user_id: Final = cast(str | None, metadata.get("user_api_key_user_id", None))
team_id: Final = cast(str | None, metadata.get("user_api_key_team_id", None))
org_id: Final = cast(str | None, metadata.get("user_api_key_org_id", None))
project_id: Final = (
project_id_value
if isinstance(project_id_value := metadata.get("user_api_key_project_id"), str)
else None
)
key_alias: Final = cast(str | None, metadata.get("user_api_key_alias", None))
end_user_max_budget: Final = metadata.get("user_api_end_user_max_budget", None)
sl_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object", None)
@ -368,6 +374,7 @@ class _ProxyDBLogger(CustomLogger):
budget_reservation=budget_reservation,
request_tags=tags,
model_access_groups=model_access_groups,
project_id=project_id,
)
if not charged:
return
@ -501,6 +508,8 @@ class _ProxyDBLogger(CustomLogger):
metadata["user_api_key_team_id"] = key_obj.team_id
if metadata.get("user_api_key_org_id") is None:
metadata["user_api_key_org_id"] = key_obj.org_id
if metadata.get("user_api_key_project_id") is None:
metadata["user_api_key_project_id"] = key_obj.project_id
except Exception:
verbose_proxy_logger.debug(
"Failed to enrich failure metadata with key info for api_key=%s",
@ -651,6 +660,7 @@ async def _update_database_and_spend_counters(
budget_reservation: dict | None,
request_tags: list[str] | None = None,
model_access_groups: Sequence[str] | None = None,
project_id: str | None = None,
) -> bool:
if budget_reservation is not None:
await _reconcile_budget_reservation_before_db_update(
@ -668,6 +678,7 @@ async def _update_database_and_spend_counters(
start_time=start_time,
end_time=end_time,
org_id=org_id,
project_id=project_id,
)
except Exception:
if budget_reservation is not None:
@ -698,6 +709,7 @@ async def _update_database_and_spend_counters(
tags=request_tags,
request_started_at=start_time,
model_access_groups=model_access_groups,
project_id=project_id,
)
except Exception:
if budget_reservation is not None:

30
litellm/proxy/proxy_server.py
View file

@ -438,6 +438,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
get_management_object_ttl,
model_access_group_cache_key,
model_access_group_spend_counter_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
)
from litellm.proxy.config_resolvers import SettingsStore, resolve_fields
@ -2836,6 +2838,7 @@ async def increment_spend_counters(
tags: list[str] | None = None,
request_started_at: datetime | None = None,
model_access_groups: Sequence[str] | None = None,
project_id: str | None = None,
):
"""
Atomically increment spend counters for budget enforcement.
@ -2857,6 +2860,7 @@ async def increment_spend_counters(
end_user_id=end_user_id,
tags=tags,
model_access_groups=model_access_groups,
project_id=project_id,
),
):
await _increment_spend_counters_batched(
@ -2870,6 +2874,7 @@ async def increment_spend_counters(
tags=tags,
request_started_at=request_started_at,
model_access_groups=model_access_groups,
project_id=project_id,
)
@ -2884,6 +2889,7 @@ async def _increment_spend_counters_batched(
tags: list[str] | None,
request_started_at: datetime | None,
model_access_groups: Sequence[str] | None,
project_id: str | None = None,
):
"""Runs inside one spend counter batch: the reservation reconcile and the warm checks share a single MGET."""
reserved_counter_keys: Final = await _reconcile_budget_reservation_for_counter_update(
@ -3084,6 +3090,13 @@ async def _increment_spend_counters_batched(
)
if org_id is not None
else None,
_prepare_project_spend_increment(
project_id=project_id,
response_cost=cost,
reserved_counter_keys=reserved_counter_keys,
)
if project_id is not None
else None,
)
if coro is not None
)
@ -3236,6 +3249,23 @@ async def _prepare_org_spend_increment(
return (pending,) if pending is not None else ()
async def _prepare_project_spend_increment(
project_id: str | None,
response_cost: float,
reserved_counter_keys: set[str],
) -> tuple[PendingSpendIncrement, ...]:
if project_id is None:
return ()
pending: Final = await _prepare_unreserved_spend_counter_increment(
counter_key=project_spend_counter_key(project_id),
source_cache_key=project_cache_key(project_id),
increment=response_cost,
reserved_counter_keys=reserved_counter_keys,
)
return (pending,) if pending is not None else ()
async def _prepare_unreserved_spend_counter_increment(
counter_key: str,
source_cache_key: str | list[str],

41
litellm/proxy/spend_tracking/budget_reservation.py
View file

@ -2,6 +2,7 @@ from __future__ import annotations
import asyncio
import json
import math
from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from datetime import datetime, timedelta, timezone
@ -29,6 +30,8 @@ from litellm.proxy.common_utils.user_api_key_cache import (
end_user_cache_key,
model_access_group_cache_key,
model_access_group_spend_counter_key,
project_cache_key,
project_spend_counter_key,
tag_cache_key,
team_membership_reservation_cache_key,
)
@ -62,6 +65,7 @@ _COUNTER_ENTITY_TYPES: Final[Mapping[str, str]] = {
"Tag": Litellm_EntityType.TAG.value,
"Model access group": Litellm_EntityType.MODEL_ACCESS_GROUP.value,
"Organization": Litellm_EntityType.ORGANIZATION.value,
"Project": Litellm_EntityType.PROJECT.value,
}
@ -542,6 +546,13 @@ async def _get_budget_counters(
if org_counter is not None:
counters.append(org_counter)
project_counter: Final = await _get_project_budget_counter(
valid_token=valid_token,
user_api_key_cache=user_api_key_cache,
)
if project_counter is not None:
counters.append(project_counter)
return counters
@ -757,6 +768,36 @@ async def _get_org_budget_counter(
)
async def _get_project_budget_counter(
valid_token: UserAPIKeyAuth,
user_api_key_cache: UserApiKeyCache,
) -> _BudgetCounter | None:
if valid_token.project_id is None:
return None
source_cache_key: Final = project_cache_key(valid_token.project_id)
project_object: Final = await user_api_key_cache.async_get_cache(key=source_cache_key)
if project_object is None:
return None
project_budget_table: Final = _get_value(project_object, "litellm_budget_table")
if project_budget_table is None:
return None
project_max_budget: Final = _to_float(_get_value(project_budget_table, "max_budget"))
if project_max_budget is None or project_max_budget <= 0 or not math.isfinite(project_max_budget):
return None
return _BudgetCounter(
counter_key=project_spend_counter_key(valid_token.project_id),
source_cache_key=source_cache_key,
max_budget=project_max_budget,
fallback_spend=_to_float(_get_value(project_object, "spend")) or 0.0,
entity_type="Project",
entity_id=valid_token.project_id,
)
def _get_budget_limit_counters(
entity_prefix: str,
entity_type: str,

11
litellm/proxy/spend_tracking/spend_counter_batch.py
View file

@ -12,7 +12,10 @@ from pydantic import TypeAdapter
from litellm._logging import verbose_proxy_logger
from litellm.caching.redis_cache import RedisCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import model_access_group_spend_counter_key
from litellm.proxy.common_utils.user_api_key_cache import (
model_access_group_spend_counter_key,
project_spend_counter_key,
)
_CounterValues: Final = TypeAdapter(dict[str, float | None])
_NO_VALUES: Final[Mapping[str, float | None]] = MappingProxyType({})
@ -154,6 +157,8 @@ def _iter_admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None)
yield f"spend:end_user:{end_user_id}"
if token.org_id is not None:
yield f"spend:org:{token.org_id}"
if token.project_id is not None:
yield project_spend_counter_key(token.project_id)
def admission_counter_keys(token: UserAPIKeyAuth, end_user_id: str | None) -> frozenset[str]:
@ -168,10 +173,12 @@ def post_call_counter_keys(
end_user_id: str | None,
tags: Sequence[object] | None,
model_access_groups: Sequence[object] | None,
project_id: str | None = None,
) -> frozenset[str]:
"""Every counter ``increment_spend_counters`` warm-checks, except budget windows which bind on read."""
entity_keys: Final = admission_counter_keys(
UserAPIKeyAuth(token=token, team_id=team_id, user_id=user_id, org_id=org_id), end_user_id
UserAPIKeyAuth(token=token, team_id=team_id, user_id=user_id, org_id=org_id, project_id=project_id),
end_user_id,
)
tag_keys: Final = frozenset(f"spend:tag:{tag}" for tag in tags or () if tag and isinstance(tag, str))
group_keys: Final = frozenset(

3
litellm/repositories/prisma_protocols.py
View file

@ -152,4 +152,7 @@ class PrismaBatch(Protocol):
@property
def litellm_modelaccessgroupbudgettable(self) -> BatchTable: ...
@property
def litellm_projecttable(self) -> BatchTable: ...
async def commit(self) -> None: ...

2
litellm/repositories/unit_of_work.py
View file

@ -109,6 +109,7 @@ class BudgetCascadeUnitOfWork:
organizations: LinkedSpendResetWrites
tags: LinkedSpendResetWrites
model_access_groups: LinkedSpendResetWrites
projects: LinkedSpendResetWrites
endusers: LinkedSpendResetWrites
budgets: BudgetWindowWrites
@ -135,6 +136,7 @@ async def budget_cascade_unit_of_work(
organizations=LinkedSpendResetWrites(table=batch.litellm_organizationtable),
tags=LinkedSpendResetWrites(table=batch.litellm_tagtable),
model_access_groups=LinkedSpendResetWrites(table=batch.litellm_modelaccessgroupbudgettable),
projects=LinkedSpendResetWrites(table=batch.litellm_projecttable),
endusers=LinkedSpendResetWrites(table=batch.litellm_endusertable),
budgets=BudgetWindowWrites(table=batch.litellm_budgettable),
)

64
tests/test_litellm/proxy/auth/test_auth_checks.py
View file

@ -7545,6 +7545,70 @@ async def test_project_allowlist_enforced_when_key_models_empty():
assert exc_info.value.code == "403"
def _project_with_budget(spend: float, max_budget: float):
from litellm.proxy._types import LiteLLM_BudgetTable, LiteLLM_ProjectTableCachedObj
return LiteLLM_ProjectTableCachedObj(
project_id="p-budget",
team_id="t-1",
budget_id="b-1",
spend=spend,
litellm_budget_table=LiteLLM_BudgetTable(budget_id="b-1", max_budget=max_budget),
)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"counter_spend, db_spend, max_budget, blocks",
[
pytest.param(5.0, 0.0, 5.0, True, id="counter-at-budget-blocks-despite-stale-db-row"),
pytest.param(4.99, 0.0, 5.0, False, id="counter-under-budget-admits"),
pytest.param(None, 5.0, 5.0, True, id="no-counter-falls-back-to-persisted-spend"),
pytest.param(None, 0.0, 5.0, False, id="no-counter-and-no-persisted-spend-admits"),
pytest.param(12.5, 12.5, 0.0, False, id="zero-budget-is-unbudgeted"),
pytest.param(12.5, 12.5, -1.0, False, id="negative-budget-is-unbudgeted"),
],
)
async def test_project_max_budget_check_blocks_only_when_live_spend_reaches_a_positive_budget(
counter_spend, db_spend, max_budget, blocks
):
from litellm.caching.dual_cache import DualCache
from litellm.proxy.auth.auth_checks import _project_max_budget_check
real_spend_counter_cache = DualCache()
if counter_spend is not None:
real_spend_counter_cache.in_memory_cache.set_cache(key="spend:project:p-budget", value=counter_spend)
valid_token = UserAPIKeyAuth(api_key="hashed-key", project_id="p-budget", team_id="t-1", user_id="u-1")
proxy_logging_obj = MagicMock()
proxy_logging_obj.budget_alerts = AsyncMock()
with patch( # test-quality-ok: injects a real DualCache for the module global, not a behavior mock
"litellm.proxy.proxy_server.spend_counter_cache", real_spend_counter_cache
):
if not blocks:
await _project_max_budget_check(
project_object=_project_with_budget(spend=db_spend, max_budget=max_budget),
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
await asyncio.sleep(0)
proxy_logging_obj.budget_alerts.assert_not_awaited()
return
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await _project_max_budget_check(
project_object=_project_with_budget(spend=db_spend, max_budget=max_budget),
valid_token=valid_token,
proxy_logging_obj=proxy_logging_obj,
)
await asyncio.sleep(0)
assert exc_info.value.entity_type == Litellm_EntityType.PROJECT.value
assert exc_info.value.entity_id == "p-budget"
assert exc_info.value.current_cost == 5.0
proxy_logging_obj.budget_alerts.assert_awaited_once()
assert proxy_logging_obj.budget_alerts.await_args.kwargs["type"] == "project_budget"
def test_is_user_proxy_admin_rejects_view_only_admin():
"""This predicate skips `non_proxy_admin_allowed_routes_check` entirely, so an
Admin Viewer answering True here would gain every write route. Read parity for

29
tests/test_litellm/proxy/common_utils/test_reset_budget_job.py
View file

@ -102,6 +102,7 @@ class MockBatcher:
self.litellm_organizationtable = _Table("org", self)
self.litellm_tagtable = _Table("tag", self)
self.litellm_modelaccessgroupbudgettable = _Table("model_access_group", self)
self.litellm_projecttable = _Table("project", self)
self.litellm_endusertable = _Table("enduser", self)
async def commit(self):
@ -117,6 +118,7 @@ class MockDB:
self.litellm_organizationtable = MockTable()
self.litellm_tagtable = MockTable()
self.litellm_modelaccessgroupbudgettable = MockTable()
self.litellm_projecttable = MockTable()
self.batch_calls: List[Dict[str, Any]] = []
self.batchers: List[MockBatcher] = []
@ -1575,13 +1577,19 @@ _INVALIDATION_CASES = [
"spend:model_access_group:gpt-4-group",
{"model_access_group:gpt-4-group"},
),
(
"litellm_projecttable",
type("Project", (), {"project_id": "proj-1"}),
"spend:project:proj-1",
{"project_id:proj-1"},
),
]
@pytest.mark.parametrize(
"table_attr, linked_row, counter_key, cache_keys",
_INVALIDATION_CASES,
ids=["team_membership", "key", "org", "tag", "model_access_group"],
ids=["team_membership", "key", "org", "tag", "model_access_group", "project"],
)
def test_budget_table_reset_invalidates_counters_and_management_cache(
reset_budget_job, mock_prisma_client, monkeypatch, table_attr, linked_row, counter_key, cache_keys
@ -1826,6 +1834,24 @@ def test_budget_table_reset_invalidates_every_access_group_not_just_the_first(
counter_cache.in_memory_cache.delete_cache.assert_any_call(key=f"spend:model_access_group:{name}")
def test_project_reset_zeroes_spend_on_due_tiers(reset_budget_job, mock_prisma_client, monkeypatch):
_make_counter_invalidation_job(monkeypatch)
mock_prisma_client.data["budget"] = [_budget_row(budget_id="budget-due", budget_duration="7d")]
mock_prisma_client.db.litellm_projecttable.set_find_many_results(
[type("Project", (), {"project_id": "proj-1", "spend": 12.0, "budget_id": "budget-due"})]
)
asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table())
expected_where = {"budget_id": {"in": ["budget-due"]}, "spend": {"gt": 0}}
assert mock_prisma_client.db.litellm_projecttable.find_many_calls == [{"where": expected_where}]
writes = _batch_writes(mock_prisma_client, "project", op="update_many")
assert len(writes) == 1
assert writes[0]["where"] == expected_where
assert writes[0]["data"] == {"spend": 0}
assert mock_prisma_client.db.batchers[0].committed is True
def test_budget_cascade_carries_access_group_overage_when_rollover_enabled(
rollover_enabled, reset_budget_job, mock_prisma_client, monkeypatch
):
@ -1971,6 +1997,7 @@ def test_budget_cascade_writes_land_in_a_single_transaction(reset_budget_job, mo
("org", "update_many"),
("tag", "update_many"),
("model_access_group", "update_many"),
("project", "update_many"),
("enduser", "update_many"),
("budget", "update_many"),
}

113
tests/test_litellm/proxy/db/test_db_spend_update_writer.py
View file

@ -1,6 +1,7 @@
import asyncio
import copy
import json
import logging
import re
@ -15,12 +16,14 @@ import pytest
from redis.exceptions import DataError
import litellm
from litellm.proxy._types import Litellm_EntityType
from litellm._logging import verbose_proxy_logger
from litellm.proxy._types import Litellm_EntityType, SpendUpdateQueueItem
from litellm.proxy.db.db_spend_update_writer import (
_TEAM_ADVISORY_LOCK_SQL,
_TEAM_MEMBER_SPEND_SQL,
DBSpendUpdateWriter,
)
from litellm.proxy.db.db_transaction_queue.spend_update_queue import SpendUpdateQueue
from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
build_window_spend_transaction,
)
@ -1146,6 +1149,114 @@ async def test_batch_database_updates_queues_org_member_spend_for_the_request_us
assert transactions["org_member_list_transactions"] == {"organization_id::org1::user_id::u1": 0.1}
@pytest.mark.asyncio
async def test_project_spend_is_persisted_to_project_table_and_project_cache_is_evicted():
db_writer: Final = DBSpendUpdateWriter()
await db_writer._batch_database_updates(
response_cost=0.25,
user_id="u1",
hashed_token="t1",
team_id="team-1",
org_id=None,
end_user_id=None,
prisma_client=MagicMock(),
litellm_proxy_budget_name=None,
payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.25},
project_id="proj-1",
)
await db_writer._batch_database_updates(
response_cost=0.5,
user_id="u1",
hashed_token="t1",
team_id="team-1",
org_id=None,
end_user_id=None,
prisma_client=MagicMock(),
litellm_proxy_budget_name=None,
payload={"request_id": "req-2", "model": "gpt-4o-mini", "spend": 0.5},
project_id="proj-1",
)
transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
assert transactions["project_list_transactions"] == {"proj-1": 0.75}
assert transactions["team_member_list_transactions"] == {"team_id::team-1::user_id::u1": 0.75}
mock_batcher: Final = MagicMock()
mock_prisma_client: Final = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=_good_tx(mock_batcher))
user_api_key_cache: Final = MagicMock()
user_api_key_cache.async_delete_cache = AsyncMock()
proxy_logging: Final = MagicMock()
proxy_logging.call_details = {"user_api_key_cache": user_api_key_cache}
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=proxy_logging,
db_spend_update_transactions=transactions,
)
mock_batcher.litellm_projecttable.update_many.assert_called_once_with(
where={"project_id": "proj-1"},
data={"spend": {"increment": 0.75}},
)
user_api_key_cache.async_delete_cache.assert_any_await(key="project_id:proj-1")
@pytest.mark.asyncio
async def test_batch_database_updates_without_project_id_touches_no_project_row():
db_writer: Final = DBSpendUpdateWriter()
await db_writer._batch_database_updates(
response_cost=0.1,
user_id="u1",
hashed_token="t1",
team_id=None,
org_id=None,
end_user_id=None,
prisma_client=MagicMock(),
litellm_proxy_budget_name=None,
payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.1},
)
transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
assert transactions["project_list_transactions"] == {}
@pytest.mark.asyncio
async def test_failed_project_enqueue_is_reported_and_does_not_drop_the_rest_of_the_batch(
caplog: pytest.LogCaptureFixture,
):
class _ProjectRejectingQueue(SpendUpdateQueue):
async def add_update(self, update: SpendUpdateQueueItem):
if update.get("entity_type") is Litellm_EntityType.PROJECT:
raise RuntimeError("project enqueue boom")
await super().add_update(update)
db_writer: Final = DBSpendUpdateWriter()
db_writer.spend_update_queue = _ProjectRejectingQueue()
with caplog.at_level(logging.ERROR, logger=verbose_proxy_logger.name):
await db_writer._batch_database_updates(
response_cost=0.25,
user_id="u1",
hashed_token="t1",
team_id="team-1",
org_id="org-1",
end_user_id=None,
prisma_client=MagicMock(),
litellm_proxy_budget_name=None,
payload={"request_id": "req-1", "model": "gpt-4o-mini", "spend": 0.25, "request_tags": ["tag-1"]},
project_id="proj-1",
)
assert any("proj-1" in record.getMessage() for record in caplog.records if record.levelno >= logging.ERROR)
transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
assert transactions["project_list_transactions"] == {}
assert transactions["tag_list_transactions"] == {"tag-1": 0.25}
assert transactions["key_list_transactions"] == {"t1": 0.25}
assert transactions["team_list_transactions"] == {"team-1": 0.25}
@pytest.mark.asyncio
async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id():
"""

17
tests/test_litellm/proxy/db/test_spend_counter_reseed.py
View file

@ -67,12 +67,14 @@ class _FakePrismaClient:
error: Exception | None = None,
end_user_row: SimpleNamespace | None = None,
end_user_error: Exception | None = None,
project_row: SimpleNamespace | None = None,
) -> None:
self.db = SimpleNamespace(
litellm_budgetwindowspend=_FakeFindUniqueTable(row=row, error=error),
litellm_spendlogs=_FakeSpendLogsTable(total=spend_logs_total),
litellm_endusertable=_FakeFindUniqueTable(row=end_user_row, error=end_user_error),
litellm_verificationtoken=_InFlightCountingTable(),
litellm_projecttable=_FakeFindUniqueTable(row=project_row),
)
@ -428,6 +430,21 @@ async def test_from_db_bounds_in_flight_prisma_requests_across_counter_keys():
assert prisma.db.litellm_verificationtoken.max_in_flight == PROXY_DB_LOOKUP_MAX_CONCURRENCY
@pytest.mark.asyncio
async def test_from_db_reseeds_project_counter_from_the_project_row():
prisma: Final = _FakePrismaClient(project_row=SimpleNamespace(project_id="proj-1", spend=7.25))
assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") == 7.25
assert prisma.db.litellm_projecttable.where_clauses == [{"project_id": "proj-1"}]
@pytest.mark.asyncio
async def test_from_db_returns_none_for_a_missing_project_row():
prisma: Final = _FakePrismaClient(project_row=None)
assert await SpendCounterReseed.from_db(prisma_client=prisma, counter_key="spend:project:proj-1") is None
@pytest.mark.asyncio
async def test_from_db_still_never_reads_the_end_user_row():
"""A cold end-user counter keeps seeding from the cached end-user object the auth

4
tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py
View file

@ -586,6 +586,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda
tags=["tag-a"],
request_started_at=start_time,
model_access_groups=("premium",),
project_id=None,
)
@ -1371,6 +1372,7 @@ async def test_enrich_failure_metadata_with_full_key_lookup():
mock_key_obj.user_id = "fetched-user-id"
mock_key_obj.team_id = "fetched-team-id"
mock_key_obj.org_id = "fetched-org-id"
mock_key_obj.project_id = "fetched-project-id"
mock_team_obj = MagicMock()
mock_team_obj.team_alias = "fetched-team-alias"
@ -1394,12 +1396,14 @@ async def test_enrich_failure_metadata_with_full_key_lookup():
"user_api_key_team_id": None,
"user_api_key_team_alias": None,
"user_api_key_org_id": None,
"user_api_key_project_id": None,
}
result = await _ProxyDBLogger._enrich_failure_metadata_with_key_info(metadata)
assert result["user_api_key_alias"] == "fetched-key-alias"
assert result["user_api_key_user_id"] == "fetched-user-id"
assert result["user_api_key_team_id"] == "fetched-team-id"
assert result["user_api_key_org_id"] == "fetched-org-id"
assert result["user_api_key_project_id"] == "fetched-project-id"
assert result["user_api_key_team_alias"] == "fetched-team-alias"

5
tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py
View file

@ -39,11 +39,10 @@ from litellm.proxy._types import (
from litellm.models.object_permission import LiteLLM_ObjectPermissionTable
from litellm.proxy.auth.auth_checks import (
_delete_cache_key_object,
_project_cache_key,
jwt_key_mapping_cache_key,
)
from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache, project_cache_key
from litellm.litellm_core_utils.duration_parser import duration_in_seconds
from litellm.proxy.management_endpoints.key_management_endpoints import (
_check_org_key_limits,
@ -19465,7 +19464,7 @@ async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read(
async def _cache_with_project(project_id: str, project_models: list[str]) -> UserApiKeyCache:
user_api_key_cache = UserApiKeyCache()
await user_api_key_cache.async_set_cache(
key=_project_cache_key(project_id),
key=project_cache_key(project_id),
value=LiteLLM_ProjectTableCachedObj(project_id=project_id, team_id="team-lit-5823", models=project_models),
model_type=LiteLLM_ProjectTableCachedObj,
)

1
tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py
View file

@ -1209,6 +1209,7 @@ async def test_api_key_preserved_through_failure_hook_to_database():
start_time,
end_time,
org_id,
project_id=None,
):
"""Mock update_database and capture the payload it creates"""
from litellm.proxy.spend_tracking.spend_tracking_utils import (

149
tests/test_litellm/proxy/test_budget_reservation.py
View file

@ -22,6 +22,7 @@ from litellm.proxy._types import (
LiteLLM_EndUserTable,
Litellm_EntityType,
LiteLLM_OrganizationTable,
LiteLLM_ProjectTableCachedObj,
LiteLLM_TagTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
@ -631,6 +632,154 @@ async def test_should_reserve_team_member_and_org_budget_counters(spend_counter_
await release_budget_reservation(reservation)
def _project_scoped_token() -> UserAPIKeyAuth:
return UserAPIKeyAuth(
token="key-project-scoped",
spend=0.0,
user_id="user-proj",
team_id="team-proj",
project_id="proj-1",
)
async def _seed_project_scoped_budgets(
key_cache: DualCache,
team_member_spend: float,
team_member_max_budget: float,
project_spend: float,
project_max_budget: float,
) -> None:
await key_cache.async_set_cache(
key="team_membership:user-proj:team-proj",
value=LiteLLM_TeamMembership(
user_id="user-proj",
team_id="team-proj",
spend=team_member_spend,
litellm_budget_table=LiteLLM_BudgetTable(max_budget=team_member_max_budget),
).model_dump(),
)
await key_cache.async_set_cache(
key="project_id:proj-1",
value=LiteLLM_ProjectTableCachedObj(
project_id="proj-1",
team_id="team-proj",
budget_id="project-budget-id",
spend=project_spend,
litellm_budget_table=LiteLLM_BudgetTable(max_budget=project_max_budget),
).model_dump(),
)
@pytest.mark.asyncio
async def test_should_reserve_project_and_team_member_counters_for_project_scoped_key(spend_counter_state):
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
await _seed_project_scoped_budgets(
key_cache,
team_member_spend=0.1,
team_member_max_budget=1.0,
project_spend=0.2,
project_max_budget=1.0,
)
estimated = estimate_request_max_cost(request_body=_request_body(), route="/chat/completions", llm_router=None)
assert estimated is not None and estimated > 0
reservation = await reserve_budget_for_request(
request_body=_request_body(),
route="/chat/completions",
llm_router=None,
valid_token=_project_scoped_token(),
team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None),
user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0),
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:team_member:user-proj:team-proj") == pytest.approx(
0.1 + estimated
)
assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") == pytest.approx(0.2 + estimated)
from litellm.proxy.proxy_server import increment_spend_counters
await increment_spend_counters(
token="key-project-scoped",
team_id="team-proj",
user_id="user-proj",
response_cost=0.05,
budget_reservation=reservation,
project_id="proj-1",
)
assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") == pytest.approx(0.25)
assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") == pytest.approx(0.15)
@pytest.mark.asyncio
async def test_exhausted_team_member_budget_still_blocks_project_scoped_key(spend_counter_state):
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
await _seed_project_scoped_budgets(
key_cache,
team_member_spend=1.0,
team_member_max_budget=1.0,
project_spend=0.0,
project_max_budget=100.0,
)
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await reserve_budget_for_request(
request_body=_request_body(),
route="/chat/completions",
llm_router=None,
valid_token=_project_scoped_token(),
team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None),
user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0),
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=proxy_logging_obj,
)
assert "TeamMember=user-proj:team-proj" in str(exc_info.value)
assert counter_cache.in_memory_cache.get_cache(key="spend:project:proj-1") in (None, pytest.approx(0.0))
@pytest.mark.asyncio
async def test_exhausted_project_budget_blocks_project_scoped_key(spend_counter_state):
counter_cache, key_cache = spend_counter_state
proxy_logging_obj = ProxyLogging(user_api_key_cache=key_cache)
await _seed_project_scoped_budgets(
key_cache,
team_member_spend=0.0,
team_member_max_budget=100.0,
project_spend=5.0,
project_max_budget=5.0,
)
with pytest.raises(litellm.BudgetExceededError) as exc_info:
await reserve_budget_for_request(
request_body=_request_body(),
route="/chat/completions",
llm_router=None,
valid_token=_project_scoped_token(),
team_object=LiteLLM_TeamTable(team_id="team-proj", spend=0.0, max_budget=None),
user_object=LiteLLM_UserTable(user_id="user-proj", spend=0.0),
prisma_client=None,
user_api_key_cache=key_cache,
proxy_logging_obj=proxy_logging_obj,
)
assert "Project=proj-1" in str(exc_info.value)
assert exc_info.value.entity_type == Litellm_EntityType.PROJECT.value
assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-proj:team-proj") in (
None,
pytest.approx(0.0),
)
@pytest.mark.asyncio
async def test_should_not_reserve_user_budget_counter_for_team_key(spend_counter_state):
"""The reservation path mirrors the read path: no personal user counter for a team key.

3
tests/test_litellm/repositories/test_unit_of_work.py
View file

@ -35,6 +35,7 @@ class FakeBatch:
self.litellm_organizationtable = FakeBatchTable("litellm_organizationtable", self.calls)
self.litellm_tagtable = FakeBatchTable("litellm_tagtable", self.calls)
self.litellm_modelaccessgroupbudgettable = FakeBatchTable("litellm_modelaccessgroupbudgettable", self.calls)
self.litellm_projecttable = FakeBatchTable("litellm_projecttable", self.calls)
self.litellm_endusertable = FakeBatchTable("litellm_endusertable", self.calls)
async def commit(self) -> None:
@ -94,6 +95,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch():
uow.organizations.queue_spend_zero(where=linked)
uow.tags.queue_spend_zero(where=linked)
uow.model_access_groups.queue_spend_zero(where=linked)
uow.projects.queue_spend_zero(where=linked)
uow.endusers.queue_spend_zero(where={"user_id": {"in": ["enduser-1"]}})
uow.budgets.queue_window_advance(budget_id="budget-1", budget_reset_at=reset_at)
assert batch.commit_count == 0
@ -105,6 +107,7 @@ async def test_budget_cascade_dependents_and_window_advance_share_one_batch():
("litellm_organizationtable.update_many", linked, {"spend": 0}),
("litellm_tagtable.update_many", linked, {"spend": 0}),
("litellm_modelaccessgroupbudgettable.update_many", linked, {"spend": 0}),
("litellm_projecttable.update_many", linked, {"spend": 0}),
("litellm_endusertable.update_many", {"user_id": {"in": ["enduser-1"]}}, {"spend": 0}),
("litellm_budgettable.update_many", {"budget_id": "budget-1"}, {"budget_reset_at": reset_at}),
]