mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-22 00:31:44 +00:00
fix(proxy): track project spend and enforce project budgets additively
Project-scoped keys never wrote spend to LiteLLM_ProjectTable, so /project/info stayed at 0 and project budgets could not block. Wire the PROJECT entity through the spend queue, redis buffer, and db writer, reserve and increment a spend:project counter, reseed it from the project row, reset project spend in the budget cascade, and read the live counter in the project max budget check. Team member budgets keep gating project-scoped keys alongside the project budget Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
67cb0bc089
commit
1a749d84bd
22 changed files with 583 additions and 21 deletions
|
|
@ -5261,6 +5261,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]
|
||||
|
|
|
|||
|
|
@ -92,6 +92,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,
|
||||
|
|
@ -5586,16 +5588,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 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,
|
||||
|
|
@ -5611,9 +5619,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,
|
||||
)
|
||||
|
|
@ -5663,10 +5671,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,
|
||||
|
|
@ -5684,7 +5688,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,
|
||||
|
|
@ -5726,7 +5730,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,
|
||||
)
|
||||
|
||||
|
|
|
|||
|
|
@ -41,6 +41,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
|
||||
|
|
@ -49,6 +51,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,
|
||||
|
|
@ -115,6 +118,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: ...
|
||||
|
|
@ -185,6 +193,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}"
|
||||
|
||||
|
|
@ -661,6 +677,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
|
||||
|
|
@ -695,6 +716,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),
|
||||
*((_enduser_counter_key(row), _enduser_carried_spend(row, rollover_caps)) for row in endusers),
|
||||
),
|
||||
rollover_caps=rollover_caps,
|
||||
|
|
@ -704,6 +726,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)),
|
||||
*(key for row in endusers for key in _enduser_cache_keys(row)),
|
||||
),
|
||||
)
|
||||
|
|
@ -731,6 +754,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)
|
||||
|
|
|
|||
|
|
@ -306,6 +306,16 @@ 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:
|
||||
"""Cache key one project row is stored under; shared by auth, spend tracking and the spend writer."""
|
||||
return f"project_id:{project_id}"
|
||||
|
||||
|
||||
def project_spend_counter_key(project_id: str) -> str:
|
||||
"""Spend counter key for one project; the reservation, cost callback, auth and reseed paths all read it."""
|
||||
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__"
|
||||
|
|
|
|||
|
|
@ -44,6 +44,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,
|
||||
|
|
@ -116,6 +117,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
|
||||
|
|
@ -254,6 +256,7 @@ class DBSpendUpdateWriter:
|
|||
start_time: datetime | None,
|
||||
end_time: datetime | None,
|
||||
response_cost: float | None,
|
||||
project_id: str | None = None,
|
||||
) -> bool:
|
||||
"""Record the request's spend, answering whether its cost still needs charging.
|
||||
|
||||
|
|
@ -335,6 +338,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,
|
||||
|
|
@ -631,6 +635,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.
|
||||
|
|
@ -694,6 +699,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:
|
||||
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,
|
||||
|
|
@ -956,6 +973,33 @@ class DBSpendUpdateWriter:
|
|||
)
|
||||
raise e
|
||||
|
||||
async def _update_project_db(
|
||||
self,
|
||||
response_cost: float | None,
|
||||
project_id: str | None,
|
||||
prisma_client: PrismaClient | None,
|
||||
):
|
||||
try:
|
||||
if project_id is None or prisma_client is None:
|
||||
return
|
||||
|
||||
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,
|
||||
|
|
@ -1193,8 +1237,8 @@ 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",
|
||||
"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 {}),
|
||||
|
|
@ -1202,6 +1246,7 @@ class DBSpendUpdateWriter:
|
|||
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 {}),
|
||||
|
|
@ -1762,6 +1807,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(
|
||||
|
|
@ -1800,11 +1861,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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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":
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,7 @@ 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 = cast(str | None, metadata.get("user_api_key_project_id", 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 +370,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
|
||||
|
|
@ -651,6 +654,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 +672,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 +703,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:
|
||||
|
|
|
|||
|
|
@ -417,6 +417,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 resolve_fields
|
||||
|
|
@ -2780,6 +2782,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.
|
||||
|
|
@ -2801,6 +2804,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(
|
||||
|
|
@ -2814,6 +2818,7 @@ async def increment_spend_counters(
|
|||
tags=tags,
|
||||
request_started_at=request_started_at,
|
||||
model_access_groups=model_access_groups,
|
||||
project_id=project_id,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -2828,6 +2833,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(
|
||||
|
|
@ -3028,6 +3034,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
|
||||
)
|
||||
|
|
@ -3180,6 +3193,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],
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
||||
|
|
@ -751,6 +762,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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -152,4 +152,7 @@ class PrismaBatch(Protocol):
|
|||
@property
|
||||
def litellm_modelaccessgroupbudgettable(self) -> BatchTable: ...
|
||||
|
||||
@property
|
||||
def litellm_projecttable(self) -> BatchTable: ...
|
||||
|
||||
async def commit(self) -> None: ...
|
||||
|
|
|
|||
|
|
@ -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),
|
||||
)
|
||||
|
|
|
|||
|
|
@ -7138,6 +7138,67 @@ 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, blocks",
|
||||
[
|
||||
pytest.param(5.0, 0.0, True, id="counter-at-budget-blocks-despite-stale-db-row"),
|
||||
pytest.param(4.99, 0.0, False, id="counter-under-budget-admits"),
|
||||
pytest.param(None, 5.0, True, id="no-counter-falls-back-to-persisted-spend"),
|
||||
pytest.param(None, 0.0, False, id="no-counter-and-no-persisted-spend-admits"),
|
||||
],
|
||||
)
|
||||
async def test_project_max_budget_check_reads_live_spend_counter(counter_spend, db_spend, blocks):
|
||||
"""LIT-3269: project budget enforcement must read the cross-pod
|
||||
``spend:project:{id}`` counter first and only fall back to the cached row's
|
||||
spend, matching key/team/org checks. The boundary is inclusive (>=)."""
|
||||
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=5.0),
|
||||
valid_token=valid_token,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
)
|
||||
return
|
||||
with pytest.raises(litellm.BudgetExceededError) as exc_info:
|
||||
await _project_max_budget_check(
|
||||
project_object=_project_with_budget(spend=db_spend, max_budget=5.0),
|
||||
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
|
||||
|
|
|
|||
|
|
@ -78,6 +78,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):
|
||||
|
|
@ -93,6 +94,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] = []
|
||||
|
||||
|
|
@ -1507,13 +1509,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
|
||||
|
|
@ -1657,6 +1665,25 @@ 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):
|
||||
"""A project linked to an expiring budget tier has its spend zeroed in the same cascade transaction."""
|
||||
_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
|
||||
):
|
||||
|
|
@ -1802,6 +1829,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"),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -1060,6 +1060,82 @@ 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():
|
||||
"""Regression for LIT-3269: a request made with a project-scoped key must
|
||||
increment LiteLLM_ProjectTable.spend, otherwise /project/info stays at 0
|
||||
and the project budget never blocks. The cached project row is evicted so
|
||||
the next auth check reads the fresh spend."""
|
||||
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_add_spend_log_transaction_to_daily_tag_transaction_with_request_id():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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,23 @@ 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():
|
||||
"""LIT-3269: a cold ``spend:project:{id}`` counter seeds from LiteLLM_ProjectTable.spend,
|
||||
so a fresh pod enforces the project budget against persisted spend rather than 0."""
|
||||
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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1187,6 +1187,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 (
|
||||
|
|
|
|||
|
|
@ -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,161 @@ 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):
|
||||
"""LIT-3269: a key carrying user_id, team_id and project_id reserves against
|
||||
both the team member counter and the project counter; neither replaces the
|
||||
other. After the call the project counter reflects the real cost once, not
|
||||
the reservation plus the post-call increment."""
|
||||
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):
|
||||
"""LIT-3269: the project budget is additive. A project with plenty of
|
||||
headroom must not let a key through once its team member budget is spent."""
|
||||
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):
|
||||
"""LIT-3269: with team member headroom left, the project budget alone blocks the key."""
|
||||
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.
|
||||
|
|
|
|||
|
|
@ -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}),
|
||||
]
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue