scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 8

Merge 1825b41768 into cdb60af024

Browse source
This commit is contained in:
Yuzhong Zhang 2026-08-26 17:02:33 +00:00 • committed by GitHub
commit e1bad27203
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
4 changed files with 276 additions and 3 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -39,6 +39,7 @@ from litellm.proxy._types import (
SpendLogsPayload,
SpendUpdateQueueItem,
ToolDiscoveryQueueItem,
hash_token,
)
from litellm.proxy.db.daily_spend_bulk_upsert import (
DAILY_SPEND_TABLES,
@ -166,7 +167,7 @@ class DBSpendUpdateWriter:
litellm_proxy_budget_name,
prisma_client,
)
from litellm.proxy.utils import ProxyUpdateSpend, hash_token
from litellm.proxy.utils import ProxyUpdateSpend
try:
verbose_proxy_logger.debug(
@ -195,6 +196,10 @@ class DBSpendUpdateWriter:
end_time=end_time,
)
payload["spend"] = response_cost or 0.0
hashed_token = self._resolve_hashed_token_for_key_spend( # rebind-ok: SpendLogs api_key fallback
hashed_token=hashed_token,
payload_api_key=payload.get("api_key"),
)
if isinstance(payload["startTime"], datetime):
payload["startTime"] = payload["startTime"].isoformat()
if isinstance(payload["endTime"], datetime):
@ -583,6 +588,26 @@ class DBSpendUpdateWriter:
traceback.format_exc(),
)
@staticmethod
def _resolve_hashed_token_for_key_spend(
hashed_token: str | None,
payload_api_key: str | None,
) -> str | None:
"""
Use the same key identifier SpendLogs persist when `token` is missing.
SpendLogs can resolve api_key from standard_logging_object.metadata.user_api_key_hash
even when the update_database `token` argument is None. Key spend must use that
same identifier so LiteLLM_VerificationToken.spend tracks logged spend.
"""
if hashed_token:
return hashed_token
if not payload_api_key:
return hashed_token
if payload_api_key.startswith("sk-"):
return hash_token(payload_api_key)
return payload_api_key
async def _update_key_db(
self,
response_cost: float | None,
@ -590,7 +615,7 @@ class DBSpendUpdateWriter:
prisma_client: PrismaClient | None,
):
try:
if hashed_token is None or prisma_client is None:
if not hashed_token or prisma_client is None:
return
await self.spend_update_queue.add_update(

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

@ -257,7 +257,10 @@ class _ProxyDBLogger(CustomLogger):
)
if response_cost is not None:
user_api_key: Final = metadata.get("user_api_key", None)
sl_metadata: Final = sl_object.get("metadata") if sl_object is not None else None
user_api_key: Final = metadata.get("user_api_key", None) or (
sl_metadata.get("user_api_key_hash") if sl_metadata else None
)
if kwargs.get("cache_hit", False) is True:
response_cost = 0.0
verbose_proxy_logger.debug("Cache Hit: response_cost %s, for user_id %s", response_cost, user_id)

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

@ -2740,3 +2740,198 @@ async def test_insert_spend_log_asks_for_an_immediate_flush_on_responses_calls(
assert prisma.spend_log_transactions == [{"request_id": "req-1", "call_type": call_type}]
assert PrismaClient.spend_log_flush_requested.is_set() is expects_flush
PrismaClient.spend_log_flush_requested.clear()
def _mock_prisma_for_key_spend_commit():
mock_batcher = MagicMock()
mock_batcher.litellm_verificationtoken = MagicMock()
mock_batcher.litellm_verificationtoken.update_many = MagicMock()
mock_batcher.litellm_usertable = MagicMock()
mock_batcher.litellm_usertable.update_many = MagicMock()
mock_batcher.litellm_teamtable = MagicMock()
mock_batcher.litellm_teamtable.update_many = MagicMock()
mock_batcher.litellm_teammembership = MagicMock()
mock_batcher.litellm_teammembership.update_many = MagicMock()
mock_batcher.litellm_organizationtable = MagicMock()
mock_batcher.litellm_organizationtable.update_many = MagicMock()
mock_batcher.litellm_tagtable = MagicMock()
mock_batcher.litellm_tagtable.update_many = MagicMock()
mock_batcher.litellm_agentstable = MagicMock()
mock_batcher.litellm_agentstable.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_prisma_client._spend_log_transactions_lock = AsyncMock()
mock_prisma_client.spend_log_transactions = []
return mock_prisma_client, mock_batcher
def _stub_proxy_server_for_key_spend(monkeypatch, prisma_client):
mock_cache = MagicMock()
mock_cache.async_get_cache = AsyncMock(return_value=None)
monkeypatch.setattr("litellm.proxy.proxy_server.disable_spend_logs", False)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma_client)
monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", mock_cache)
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_budget_name", "test-budget")
monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {})
monkeypatch.setattr("litellm.proxy.proxy_server.master_key", None)
@pytest.mark.asyncio
async def test_logged_spend_updates_verification_token_spend_when_token_provided(monkeypatch):
"""
Successful logged spend must increment LiteLLM_VerificationToken.spend
for the hashed key written to SpendLogs.
"""
db_writer = DBSpendUpdateWriter()
hashed_token = "a" * 64
response_cost = 0.05
mock_prisma_client, mock_batcher = _mock_prisma_for_key_spend_commit()
_stub_proxy_server_for_key_spend(monkeypatch, mock_prisma_client)
await db_writer.update_database(
token=hashed_token,
user_id=None,
end_user_id=None,
start_time=datetime.now(),
end_time=datetime.now(),
team_id=None,
org_id=None,
completion_response=MagicMock(),
response_cost=response_cost,
kwargs={
"model": "gpt-4",
"custom_llm_provider": "openai",
"litellm_params": {"metadata": {"user_api_key": hashed_token}},
},
)
await asyncio.sleep(0)
aggregated = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
assert aggregated["key_list_transactions"][hashed_token] == response_cost
monkeypatch.setattr("litellm.proxy.utils._raise_failed_update_spend_exception", MagicMock())
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=MagicMock(),
db_spend_update_transactions=aggregated,
)
mock_batcher.litellm_verificationtoken.update_many.assert_called()
call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1]
assert call_kwargs["where"] == {"token": hashed_token}
assert call_kwargs["data"]["spend"] == {"increment": response_cost}
@pytest.mark.asyncio
async def test_logged_spend_updates_verification_token_spend_from_spend_log_api_key(monkeypatch):
"""
Fixes #37144: SpendLogs can persist api_key from standard_logging_object
even when update_database is called with token=None (metadata.user_api_key
missing, but user/team ids still trigger spend tracking).
LiteLLM_VerificationToken.spend must increment for that same api_key.
"""
db_writer = DBSpendUpdateWriter()
hashed_token = "b" * 64
response_cost = 0.123
mock_prisma_client, mock_batcher = _mock_prisma_for_key_spend_commit()
_stub_proxy_server_for_key_spend(monkeypatch, mock_prisma_client)
kwargs = {
"model": "custom_openai/test-model",
"custom_llm_provider": "openai",
"litellm_params": {
"metadata": {
"user_api_key_user_id": "user-1",
"user_api_key_team_id": "team-1",
}
},
"standard_logging_object": {
"response_cost": response_cost,
"prompt_tokens": 10,
"completion_tokens": 20,
"total_tokens": 30,
"model_map_information": None,
"hidden_params": {},
"metadata": {"user_api_key_hash": hashed_token},
},
}
await db_writer.update_database(
token=None,
user_id="user-1",
end_user_id=None,
start_time=datetime.now(),
end_time=datetime.now(),
team_id="team-1",
org_id=None,
completion_response=MagicMock(),
response_cost=response_cost,
kwargs=kwargs,
)
await asyncio.sleep(0)
spend_log_api_key = mock_prisma_client.spend_log_transactions[0]["api_key"]
assert spend_log_api_key == hashed_token
aggregated = await db_writer.spend_update_queue.flush_and_get_aggregated_db_spend_update_transactions()
assert hashed_token in aggregated["key_list_transactions"], (
"VerificationToken.spend was not queued for the SpendLogs api_key"
)
assert aggregated["key_list_transactions"][hashed_token] == response_cost
key_only_transactions = {
"user_list_transactions": {},
"end_user_list_transactions": {},
"key_list_transactions": aggregated["key_list_transactions"],
"team_list_transactions": {},
"team_member_list_transactions": {},
"org_list_transactions": {},
"tag_list_transactions": {},
"agent_list_transactions": {},
}
monkeypatch.setattr("litellm.proxy.utils._raise_failed_update_spend_exception", MagicMock())
await db_writer._commit_spend_updates_to_db(
prisma_client=mock_prisma_client,
n_retry_times=0,
proxy_logging_obj=MagicMock(),
db_spend_update_transactions=key_only_transactions,
)
mock_batcher.litellm_verificationtoken.update_many.assert_called()
call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1]
assert call_kwargs["where"] == {"token": hashed_token}
assert call_kwargs["data"]["spend"] == {"increment": response_cost}
def test_resolve_hashed_token_for_key_spend_uses_payload_when_token_missing():
hashed = "c" * 64
assert (
DBSpendUpdateWriter._resolve_hashed_token_for_key_spend(
hashed_token=None, payload_api_key=hashed
)
== hashed
)
assert (
DBSpendUpdateWriter._resolve_hashed_token_for_key_spend(
hashed_token="", payload_api_key=hashed
)
== hashed
)
assert (
DBSpendUpdateWriter._resolve_hashed_token_for_key_spend(
hashed_token="existing-token", payload_api_key=hashed
)
== "existing-token"
)

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

@ -1875,3 +1875,53 @@ async def test_track_cost_callback_logs_unauthenticated_pass_through_request(
assert mock_proxy_logging.db_spend_update_writer.update_database.await_count == (
1 if expect_spend_log else 0
)
@pytest.mark.asyncio
async def test_track_cost_callback_uses_standard_logging_key_hash_when_metadata_missing(monkeypatch):
"""
When metadata.user_api_key is missing but standard_logging_object has
user_api_key_hash, the cost callback must still pass that hash as token
so LiteLLM_VerificationToken.spend is updated (issue #37144).
"""
logger = _ProxyDBLogger()
hashed_token = "d" * 64
kwargs = {
"model": "custom_openai/test-model",
"litellm_params": {
"metadata": {
"user_api_key_user_id": "user-1",
"user_api_key_team_id": "team-1",
}
},
"standard_logging_object": {
"response_cost": 0.05,
"metadata": {"user_api_key_hash": hashed_token},
},
"stream": False,
}
mock_proxy_logging = MagicMock()
mock_proxy_logging.db_spend_update_writer = MagicMock()
mock_proxy_logging.db_spend_update_writer.update_database = AsyncMock()
mock_proxy_logging.slack_alerting_instance = MagicMock()
mock_proxy_logging.slack_alerting_instance.customer_spend_alert = AsyncMock()
mock_increment = AsyncMock()
monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging)
monkeypatch.setattr("litellm.proxy.proxy_server.increment_spend_counters", mock_increment)
monkeypatch.setattr("litellm.proxy.proxy_server.update_cache", AsyncMock())
await logger._PROXY_track_cost_callback(
kwargs=kwargs,
completion_response=MagicMock(),
start_time=datetime.now(),
end_time=datetime.now(),
)
mock_proxy_logging.db_spend_update_writer.update_database.assert_awaited_once()
call_kwargs = (
mock_proxy_logging.db_spend_update_writer.update_database.call_args[1]
)
assert call_kwargs["token"] == hashed_token
mock_increment.assert_awaited_once()
assert mock_increment.call_args[1]["token"] == hashed_token