scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-14 09:03:16 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 214

fix(proxy): track project spend so project budgets and alerts enforce

Browse source
This commit is contained in:
Devin AI 2026-07-18 23:26:54 +00:00
parent 366ec6f487
commit 227a99cf6b
8 changed files with 265 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
litellm/proxy/_types.py
View file

@ -4505,6 +4505,7 @@ class DBSpendUpdateTransactions(TypedDict):
team_list_transactions: Optional[Dict[str, float]]
team_member_list_transactions: Optional[Dict[str, float]]
org_list_transactions: Optional[Dict[str, float]]
project_list_transactions: Dict[str, float] | None
tag_list_transactions: Optional[Dict[str, float]]
agent_list_transactions: Optional[Dict[str, float]]

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

@ -131,6 +131,7 @@ class DBSpendUpdateWriter:
start_time: Optional[datetime],
end_time: Optional[datetime],
response_cost: Optional[float],
project_id: str | None = None,
):
from litellm.proxy.proxy_server import (
disable_spend_logs,
@ -191,6 +192,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,
@ -210,12 +212,13 @@ class DBSpendUpdateWriter:
spend_log_error(
"Spend tracking - update_database failed. Spend log insertion or daily transaction enqueue "
"may not have completed for this request. "
"response_cost=%s, token=%s, user_id=%s, team_id=%s, org_id=%s, end_user_id=%s",
"response_cost=%s, token=%s, user_id=%s, team_id=%s, org_id=%s, project_id=%s, end_user_id=%s",
response_cost,
token,
user_id,
team_id,
org_id,
project_id,
end_user_id,
)
@ -325,6 +328,7 @@ class DBSpendUpdateWriter:
prisma_client: Optional[PrismaClient],
litellm_proxy_budget_name: Optional[str],
payload: SpendLogsPayload,
project_id: str | None = None,
):
"""
Runs all 11 spend-update helpers sequentially inside a single asyncio task.
@ -387,6 +391,12 @@ class DBSpendUpdateWriter:
traceback.format_exc(),
)
await self._update_project_db(
response_cost=response_cost,
project_id=project_id,
prisma_client=prisma_client,
)
try:
await self._update_tag_db(
response_cost=response_cost,
@ -631,6 +641,35 @@ 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:
verbose_proxy_logger.debug(
"track_cost_callback: project_id is None or prisma_client is None. Not tracking spend for project"
)
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: # noqa: BLE001 # background spend enqueue must not break request flow or block sibling helpers
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,
)
async def _update_agent_db(
self,
response_cost: Optional[float],
@ -1278,6 +1317,18 @@ class DBSpendUpdateWriter:
e=e, start_time=start_time, proxy_logging_obj=proxy_logging_obj
)
### UPDATE PROJECT TABLE ###
project_list_transactions = 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,
)
### UPDATE TAG TABLE ###
tag_list_transactions = db_spend_update_transactions["tag_list_transactions"]
await DBSpendUpdateWriter._update_entity_spend_in_db(

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

@ -339,6 +339,10 @@ class RedisUpdateBuffer:
Litellm_EntityType.ORGANIZATION,
db_spend_update_transactions.get("org_list_transactions"),
),
(
Litellm_EntityType.PROJECT,
db_spend_update_transactions.get("project_list_transactions"),
),
(
Litellm_EntityType.TAG,
db_spend_update_transactions.get("tag_list_transactions"),
@ -717,6 +721,7 @@ class RedisUpdateBuffer:
team_list_transactions={},
team_member_list_transactions={},
org_list_transactions={},
project_list_transactions={},
tag_list_transactions={},
agent_list_transactions={},
)
@ -729,6 +734,7 @@ class RedisUpdateBuffer:
"team_list_transactions",
"team_member_list_transactions",
"org_list_transactions",
"project_list_transactions",
"tag_list_transactions",
"agent_list_transactions",
]

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

@ -138,6 +138,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
team_list_transactions={},
team_member_list_transactions={},
org_list_transactions={},
project_list_transactions={},
tag_list_transactions={},
agent_list_transactions={},
)
@ -150,6 +151,7 @@ class SpendUpdateQueue(BaseUpdateQueue):
Litellm_EntityType.TEAM: "team_list_transactions",
Litellm_EntityType.TEAM_MEMBER: "team_member_list_transactions",
Litellm_EntityType.ORGANIZATION: "org_list_transactions",
Litellm_EntityType.PROJECT: "project_list_transactions",
Litellm_EntityType.TAG: "tag_list_transactions",
Litellm_EntityType.AGENT: "agent_list_transactions",
}
@ -187,6 +189,8 @@ class SpendUpdateQueue(BaseUpdateQueue):
transactions_dict = db_spend_update_transactions["team_member_list_transactions"]
elif dict_key == "org_list_transactions":
transactions_dict = db_spend_update_transactions["org_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/hooks/proxy_track_cost_callback.py
View file

@ -174,6 +174,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
@ -208,6 +209,7 @@ class _ProxyDBLogger(CustomLogger):
user_id = cast(Optional[str], metadata.get("user_api_key_user_id", None))
team_id = cast(Optional[str], metadata.get("user_api_key_team_id", None))
org_id = cast(Optional[str], metadata.get("user_api_key_org_id", None))
project_id = cast("str | None", metadata.get("user_api_key_project_id", None)) # cast-ok: untyped metadata
key_alias = cast(Optional[str], metadata.get("user_api_key_alias", None))
end_user_max_budget = metadata.get("user_api_end_user_max_budget", None)
sl_object: Optional[StandardLoggingPayload] = kwargs.get("standard_logging_object", None)
@ -245,6 +247,7 @@ class _ProxyDBLogger(CustomLogger):
end_user_id=end_user_id,
team_id=team_id,
org_id=org_id,
project_id=project_id,
kwargs=kwargs,
completion_response=completion_response,
start_time=start_time,
@ -488,6 +491,7 @@ async def _update_database_and_spend_counters(
response_cost: float,
budget_reservation: Optional[dict],
request_tags: Optional[List[str]] = None,
project_id: str | None = None,
) -> None:
try:
await proxy_logging_obj.db_spend_update_writer.update_database(
@ -501,6 +505,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:

25
tests/test_litellm/proxy/db/db_transaction_queue/test_spend_update_queue.py
View file

@ -34,6 +34,31 @@ async def test_add_update(spend_queue):
assert spend_queue.update_queue.qsize() == 1
@pytest.mark.asyncio
async def test_project_entity_type_aggregates_into_project_transactions(spend_queue):
# PROJECT updates must aggregate into project_list_transactions so project spend is committed.
await spend_queue.add_update(
{
"entity_type": Litellm_EntityType.PROJECT,
"entity_id": "project123",
"response_cost": 0.5,
}
)
await spend_queue.add_update(
{
"entity_type": Litellm_EntityType.PROJECT,
"entity_id": "project123",
"response_cost": 0.25,
}
)
aggregated = (
await spend_queue.flush_and_get_aggregated_db_spend_update_transactions()
)
assert aggregated["project_list_transactions"]["project123"] == 0.75
@pytest.mark.asyncio
async def test_missing_response_cost(spend_queue):
# Test with missing response_cost - should default to 0

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

@ -627,6 +627,161 @@ async def test_update_agent_db_skips_when_prisma_client_none():
writer.spend_update_queue.add_update.assert_not_called()
@pytest.mark.asyncio
async def test_update_project_db_enqueues_project_spend():
"""
Test that _update_project_db enqueues a SpendUpdateQueueItem with entity_type=PROJECT.
"""
from litellm.proxy._types import Litellm_EntityType
writer = DBSpendUpdateWriter()
mock_prisma = MagicMock()
project_id = "project-123"
response_cost = 0.1
writer.spend_update_queue.add_update = AsyncMock()
await writer._update_project_db(
response_cost=response_cost,
project_id=project_id,
prisma_client=mock_prisma,
)
writer.spend_update_queue.add_update.assert_called_once()
call_args = writer.spend_update_queue.add_update.call_args[1]
assert call_args["update"]["entity_type"] == Litellm_EntityType.PROJECT
assert call_args["update"]["entity_id"] == project_id
assert call_args["update"]["response_cost"] == response_cost
@pytest.mark.asyncio
async def test_update_project_db_skips_when_project_id_none():
"""_update_project_db does not enqueue when project_id is None."""
writer = DBSpendUpdateWriter()
mock_prisma = MagicMock()
writer.spend_update_queue.add_update = AsyncMock()
await writer._update_project_db(
response_cost=0.05,
project_id=None,
prisma_client=mock_prisma,
)
writer.spend_update_queue.add_update.assert_not_called()
@pytest.mark.asyncio
async def test_update_project_db_skips_when_prisma_client_none():
"""_update_project_db does not enqueue when prisma_client is None."""
writer = DBSpendUpdateWriter()
writer.spend_update_queue.add_update = AsyncMock()
await writer._update_project_db(
response_cost=0.05,
project_id="project-456",
prisma_client=None,
)
writer.spend_update_queue.add_update.assert_not_called()
@pytest.mark.asyncio
async def test_batch_database_updates_forwards_project_id_to_project_helper():
"""
_batch_database_updates must route the request's project_id into _update_project_db
so project spend is attributed alongside user/key/team/org.
"""
db_writer = DBSpendUpdateWriter()
db_writer._update_user_db = AsyncMock()
db_writer._update_key_db = AsyncMock()
db_writer._update_team_db = AsyncMock()
db_writer._update_org_db = AsyncMock()
db_writer._update_project_db = AsyncMock()
db_writer._update_tag_db = AsyncMock()
db_writer._update_agent_db = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_end_user_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_agent_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_team_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_org_transaction = AsyncMock()
db_writer.add_spend_log_transaction_to_daily_tag_transaction = AsyncMock()
await db_writer._batch_database_updates(
response_cost=0.1,
user_id="u1",
hashed_token="t1",
team_id="team1",
org_id="org1",
project_id="project1",
end_user_id="eu1",
prisma_client=MagicMock(),
litellm_proxy_budget_name="budget",
payload={"key": "value"},
)
db_writer._update_project_db.assert_awaited_once()
call_args = db_writer._update_project_db.call_args[1]
assert call_args["project_id"] == "project1"
assert call_args["response_cost"] == 0.1
@pytest.mark.asyncio
async def test_commit_spend_updates_to_db_increments_project_spend():
"""
Test that _commit_spend_updates_to_db calls litellm_projecttable.update_many
with spend increment when project_list_transactions is present.
"""
db_writer = DBSpendUpdateWriter()
mock_batcher = MagicMock()
mock_batcher.litellm_projecttable = MagicMock()
mock_batcher.litellm_projecttable.update_many = MagicMock()
mock_transaction = AsyncMock()
mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction)
mock_transaction.__aexit__ = AsyncMock(return_value=False)
mock_transaction.batch_ = MagicMock(
return_value=AsyncMock(
__aenter__=AsyncMock(return_value=mock_batcher),
__aexit__=AsyncMock(return_value=False),
)
)
mock_prisma_client = MagicMock()
mock_prisma_client.db = MagicMock()
mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction)
mock_proxy_logging = MagicMock()
project_id = "project-789"
response_cost = 0.25
db_spend_update_transactions = {
"user_list_transactions": {},
"end_user_list_transactions": {},
"key_list_transactions": {},
"team_list_transactions": {},
"team_member_list_transactions": {},
"org_list_transactions": {},
"project_list_transactions": {project_id: response_cost},
"tag_list_transactions": {},
"agent_list_transactions": {},
}
with patch("litellm.proxy.utils._raise_failed_update_spend_exception"):
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=mock_proxy_logging,
db_spend_update_transactions=db_spend_update_transactions,
)
mock_batcher.litellm_projecttable.update_many.assert_called_once()
call_kwargs = mock_batcher.litellm_projecttable.update_many.call_args[1]
assert call_kwargs["where"] == {"project_id": project_id}
assert call_kwargs["data"] == {"spend": {"increment": response_cost}}
@pytest.mark.asyncio
async def test_commit_spend_updates_to_db_increments_agent_spend():
"""
@ -1644,6 +1799,15 @@ async def test_commit_spend_updates_uses_pipeline():
["agent_a", "agent_b", "agent_c"],
id="agent",
),
pytest.param(
"project_list_transactions",
{"project_c": 0.1, "project_a": 0.2, "project_b": 0.3},
"litellm_projecttable",
"update_many",
"project_id",
["project_a", "project_b", "project_c"],
id="project",
),
],
)
@pytest.mark.asyncio
@ -1693,6 +1857,7 @@ async def test_commit_spend_updates_iterates_in_sorted_order(
"team_list_transactions": {},
"team_member_list_transactions": {},
"org_list_transactions": {},
"project_list_transactions": {},
"tag_list_transactions": {},
"agent_list_transactions": {},
}

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

@ -32,6 +32,7 @@ async def test_async_post_call_failure_hook():
user_id="test_user_id",
team_id="test_team_id",
org_id="test_org_id",
project_id="test_project_id",
team_alias="test_team_alias",
end_user_id="test_end_user_id",
)
@ -70,6 +71,7 @@ async def test_async_post_call_failure_hook():
assert call_args["end_user_id"] == "test_end_user_id"
assert call_args["team_id"] == "test_team_id"
assert call_args["org_id"] == "test_org_id"
assert call_args["project_id"] == "test_project_id"
assert call_args["completion_response"] == original_exception
# Check that metadata was properly updated
@ -452,6 +454,7 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda
end_user_id="test_end_user_id",
team_id="test_team_id",
org_id="test_org_id",
project_id="test_project_id",
kwargs={},
completion_response=None,
start_time=datetime.now(),
@ -462,6 +465,10 @@ async def test_update_database_and_spend_counters_updates_counters_after_db_upda
)
proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once()
assert (
proxy_logging_obj.db_spend_update_writer.update_database.call_args[1]["project_id"]
== "test_project_id"
)
increment_spend_counters.assert_awaited_once_with(
token="test_api_key",
team_id="test_team_id",