scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-27 06:13:07 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 72

Merge bd158c8108 into 06e3e32c88

Browse source
This commit is contained in:
Nate Armstrong 2026-10-20 14:56:16 +02:00 • committed by GitHub
commit 8d5e58c3df
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 172 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -1131,10 +1131,12 @@ class DBSpendUpdateWriter:
traceback.format_exc(),
)
daily_spend_base: Final = await self._shared_daily_spend_base_transaction(payload_copy, prisma_client)
try:
await self.add_spend_log_transaction_to_daily_user_transaction(
payload=payload_copy,
prisma_client=prisma_client,
daily_spend_base=daily_spend_base,
)
except Exception:
verbose_proxy_logger.debug(
@ -1146,6 +1148,7 @@ class DBSpendUpdateWriter:
await self.add_spend_log_transaction_to_daily_end_user_transaction(
payload=payload_copy,
prisma_client=prisma_client,
daily_spend_base=daily_spend_base,
)
except Exception:
verbose_proxy_logger.debug(
@ -1157,6 +1160,7 @@ class DBSpendUpdateWriter:
await self.add_spend_log_transaction_to_daily_agent_transaction(
payload=payload_copy,
prisma_client=prisma_client,
daily_spend_base=daily_spend_base,
)
except Exception:
verbose_proxy_logger.debug(
@ -1168,6 +1172,7 @@ class DBSpendUpdateWriter:
await self.add_spend_log_transaction_to_daily_team_transaction(
payload=payload_copy,
prisma_client=prisma_client,
daily_spend_base=daily_spend_base,
)
except Exception:
verbose_proxy_logger.debug(
@ -1187,6 +1192,7 @@ class DBSpendUpdateWriter:
payload=payload_copy,
org_id=org_id,
prisma_client=prisma_client,
daily_spend_base=daily_spend_base,
)
except Exception:
verbose_proxy_logger.debug(
@ -1198,6 +1204,7 @@ class DBSpendUpdateWriter:
await self.add_spend_log_transaction_to_daily_tag_transaction(
payload=payload_copy,
prisma_client=prisma_client,
daily_spend_base=daily_spend_base,
)
except Exception:
verbose_proxy_logger.debug(
@ -2761,9 +2768,10 @@ class DBSpendUpdateWriter:
async def _common_add_spend_log_transaction_to_daily_transaction(
self,
payload: dict | SpendLogsPayload,
payload: SpendLogsPayload,
prisma_client: PrismaClient,
type: Literal["user", "team", "org", "request_tags", "end_user", "agent"] = "user",
daily_spend_base: BaseDailySpendTransaction | None = None,
) -> BaseDailySpendTransaction | None:
entity: Final = "tag" if type == "request_tags" else type
identity_payload: Final = (
@ -2771,6 +2779,31 @@ class DBSpendUpdateWriter:
)
if not daily_spend_entity_ids(identity_payload, entity):
return None
if daily_spend_base is not None:
return daily_spend_base
return await self._daily_spend_base_transaction(payload, prisma_client)
async def _shared_daily_spend_base_transaction(
self,
payload: SpendLogsPayload,
prisma_client: PrismaClient | None,
) -> BaseDailySpendTransaction | None:
if prisma_client is None:
return None
try:
return await self._daily_spend_base_transaction(payload, prisma_client)
except Exception: # noqa: BLE001 # each daily rollup falls back to building its own row, as before
verbose_proxy_logger.debug(
"_batch_database_updates: _daily_spend_base_transaction failed: %s",
traceback.format_exc(),
)
return None
async def _daily_spend_base_transaction(
self,
payload: SpendLogsPayload,
prisma_client: PrismaClient,
) -> BaseDailySpendTransaction | None:
expected_keys: Final = ("startTime", "api_key")
if not all(key in payload for key in expected_keys):
verbose_proxy_logger.debug(
@ -2806,7 +2839,7 @@ class DBSpendUpdateWriter:
if isinstance(payload["startTime"], datetime):
start_time: Final = payload["startTime"].isoformat()
date = start_time.split("T")[0]
elif isinstance(payload["startTime"], str):
elif isinstance(payload["startTime"], str): # pyright: ignore[reportUnnecessaryIsInstance] # spend log payloads are not validated against the TypedDict at runtime
date = payload["startTime"].split("T")[0]
else:
verbose_proxy_logger.debug(
@ -2874,8 +2907,9 @@ class DBSpendUpdateWriter:
async def add_spend_log_transaction_to_daily_user_transaction(
self,
payload: dict | SpendLogsPayload,
payload: SpendLogsPayload,
prisma_client: PrismaClient | None = None,
daily_spend_base: BaseDailySpendTransaction | None = None,
):
"""
Add a spend log transaction to the `daily_spend_update_queue`
@ -2889,7 +2923,7 @@ class DBSpendUpdateWriter:
return
base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction(
payload, prisma_client, "user"
payload, prisma_client, "user", daily_spend_base
)
if base_daily_transaction is None:
return
@ -2903,13 +2937,14 @@ class DBSpendUpdateWriter:
self,
payload: SpendLogsPayload,
prisma_client: PrismaClient | None = None,
daily_spend_base: BaseDailySpendTransaction | None = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")
return
base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction(
payload, prisma_client, "team"
payload, prisma_client, "team", daily_spend_base
)
if base_daily_transaction is None:
return
@ -2927,6 +2962,7 @@ class DBSpendUpdateWriter:
payload: SpendLogsPayload,
prisma_client: PrismaClient | None = None,
org_id: str | None = None,
daily_spend_base: BaseDailySpendTransaction | None = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")
@ -2945,7 +2981,7 @@ class DBSpendUpdateWriter:
)
base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction(
payload_with_org, prisma_client, "org"
payload_with_org, prisma_client, "org", daily_spend_base
)
if base_daily_transaction is None:
return
@ -2959,6 +2995,7 @@ class DBSpendUpdateWriter:
self,
payload: SpendLogsPayload,
prisma_client: PrismaClient | None = None,
daily_spend_base: BaseDailySpendTransaction | None = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")
@ -2978,7 +3015,7 @@ class DBSpendUpdateWriter:
)
base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction(
payload_with_end_user_id, prisma_client, "end_user"
payload_with_end_user_id, prisma_client, "end_user", daily_spend_base
)
if base_daily_transaction is None:
return
@ -2992,6 +3029,7 @@ class DBSpendUpdateWriter:
self,
payload: SpendLogsPayload,
prisma_client: PrismaClient | None = None,
daily_spend_base: BaseDailySpendTransaction | None = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")
@ -3006,7 +3044,7 @@ class DBSpendUpdateWriter:
},
)
base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction(
payload_with_agent_id, prisma_client, "agent"
payload_with_agent_id, prisma_client, "agent", daily_spend_base
)
if base_daily_transaction is None:
return
@ -3019,13 +3057,14 @@ class DBSpendUpdateWriter:
self,
payload: SpendLogsPayload,
prisma_client: PrismaClient | None = None,
daily_spend_base: BaseDailySpendTransaction | None = None,
) -> None:
if prisma_client is None:
verbose_proxy_logger.debug("prisma_client is None. Skipping writing spend logs to db.")
return
base_daily_transaction: Final = await self._common_add_spend_log_transaction_to_daily_transaction(
payload, prisma_client, "request_tags"
payload, prisma_client, "request_tags", daily_spend_base
)
if base_daily_transaction is None:
return

124
tests/unit/proxy/db/test_db_spend_update_writer.py
View file

@ -1,5 +1,6 @@
import asyncio
import copy
import itertools
import json
import logging
import re
@ -35,6 +36,7 @@ from litellm.proxy.db.db_transaction_queue.window_spend_update_queue import (
build_window_spend_transaction,
)
from litellm.proxy.db.rollup_lock_timeout import ROLLUP_LOCK_TIMEOUT_SQL
from litellm.proxy.spend_tracking.savings import compute_savings_spend
from tests.unit.proxy.db.fake_prisma_engine import engine_call
@ -1440,6 +1442,128 @@ async def test_failed_project_enqueue_is_reported_and_does_not_drop_the_rest_of_
assert transactions["team_list_transactions"] == {"team-1": 0.25}
@pytest.mark.asyncio
async def test_batch_database_updates_prices_savings_once_for_every_daily_rollup():
db_writer: Final = DBSpendUpdateWriter()
prisma_client: Final = MagicMock()
prisma_client.get_request_status = MagicMock(return_value="success")
payload: Final = {
"request_id": "req-daily",
"user": "user-1",
"team_id": "team-1",
"end_user": "end-user-1",
"agent_id": "agent-1",
"request_tags": ["tag-a", "tag-b"],
"startTime": "2024-01-01T12:00:00",
"api_key": "key-1",
"model": "claude-sonnet-5",
"custom_llm_provider": "anthropic",
"model_group": "claude-sonnet-5",
"prompt_tokens": 1000,
"completion_tokens": 10,
"spend": 0.01,
"metadata": json.dumps({"usage_object": {"prompt_tokens": 1000, "cache_read_input_tokens": 800}}),
}
with patch("litellm.proxy.db.db_spend_update_writer.compute_savings_spend", wraps=compute_savings_spend) as savings:
await db_writer._batch_database_updates(
response_cost=0.01,
user_id="user-1",
hashed_token="key-1",
team_id="team-1",
org_id="org-1",
end_user_id="end-user-1",
prisma_client=prisma_client,
litellm_proxy_budget_name=None,
payload=payload,
)
assert savings.call_count == 1
queues: Final = {
"user_id": db_writer.daily_spend_update_queue,
"team_id": db_writer.daily_team_spend_update_queue,
"organization_id": db_writer.daily_org_spend_update_queue,
"end_user_id": db_writer.daily_end_user_spend_update_queue,
"agent_id": db_writer.daily_agent_spend_update_queue,
"tag": db_writer.daily_tag_spend_update_queue,
}
rows: Final = {
entity: tuple(
sorted((await queue.flush_and_get_aggregated_daily_spend_update_transactions()).values(), key=str)
)
for entity, queue in queues.items()
}
assert {entity: tuple(row[entity] for row in entity_rows) for entity, entity_rows in rows.items()} == {
"user_id": ("user-1",),
"team_id": ("team-1",),
"organization_id": ("org-1",),
"end_user_id": ("end-user-1",),
"agent_id": ("agent-1",),
"tag": ("tag-a", "tag-b"),
}
all_rows: Final = tuple(itertools.chain.from_iterable(rows.values()))
savings_columns: Final = {
(row["prompt_caching_savings_spend"], row["cache_read_input_tokens"], row["spend"]) for row in all_rows
}
assert len(savings_columns) == 1
assert next(iter(savings_columns))[0] > 0
@pytest.mark.asyncio
@pytest.mark.parametrize(
"overrides, dropped, has_prisma",
[
({"metadata": "{not json"}, (), True),
({}, (), False),
({}, ("api_key",), True),
({}, ("model",), True),
({}, ("model_group",), True),
],
)
async def test_batch_database_updates_skips_daily_rows_it_cannot_build_but_keeps_key_spend(
overrides: dict[str, str], dropped: tuple[str, ...], has_prisma: bool
):
db_writer: Final = DBSpendUpdateWriter()
prisma_client: Final = MagicMock() if has_prisma else None
if prisma_client is not None:
prisma_client.get_request_status = MagicMock(return_value="success")
complete_payload: Final = {
"request_id": "req-bad",
"user": "user-1",
"team_id": "team-1",
"startTime": "2024-01-01T12:00:00",
"api_key": "key-1",
"model": "gpt-4o-mini",
"custom_llm_provider": "openai",
"model_group": "gpt-4o-mini",
"prompt_tokens": 10,
"completion_tokens": 5,
"spend": 0.01,
"metadata": '{"usage_object": {}}',
**overrides,
}
payload: Final = {key: value for key, value in complete_payload.items() if key not in dropped}
await db_writer._batch_database_updates(
response_cost=0.01,
user_id="user-1",
hashed_token="key-1",
team_id="team-1",
org_id=None,
end_user_id=None,
prisma_client=prisma_client,
litellm_proxy_budget_name=None,
payload=payload,
)
assert await db_writer.daily_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() == {}
assert (
await db_writer.daily_team_spend_update_queue.flush_and_get_aggregated_daily_spend_update_transactions() == {}
)
transactions: Final = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
assert transactions["key_list_transactions"] == ({"key-1": 0.01} if has_prisma else {})
@pytest.mark.asyncio
async def test_add_spend_log_transaction_to_daily_tag_transaction_with_request_id():
"""