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

Merge branch 'litellm_fix-batch-spend-key-double-hash-bcae' of https://github.com/BerriAI/litellm into litellm_fix-batch-spend-key-double-hash-bcae

Browse source
This commit is contained in:
mateo-berri 2026-09-03 18:46:25 -07:00
commit b93cf2e20a
2 changed files with 50 additions and 4 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

6
litellm/proxy/spend_tracking/key_metadata_recovery.py
View file

@ -186,7 +186,7 @@ def _row_with_recovered_fields(
**row,
alias_field: meta.get("key_alias") or row.get(alias_field),
team_id_field: meta.get("team_id") or row.get(team_id_field),
user_email_field: meta.get("user_email") or row.get(user_email_field),
user_email_field: row.get(user_email_field) or meta.get("user_email"),
}
)
@ -211,9 +211,7 @@ async def fill_missing_api_key_aliases(
key
for row in rows
for key in (row.get(api_key_field),)
if isinstance(key, str)
and key
and (row.get(alias_field) in (None, "") or row.get(user_email_field) in (None, ""))
if isinstance(key, str) and key and row.get(alias_field) in (None, "")
)
if not missing_keys:
return tuple(rows)

48
tests/test_litellm/proxy/spend_tracking/test_key_metadata_recovery.py
View file

@ -170,3 +170,51 @@ async def test_fill_missing_api_key_aliases_leaves_rows_untouched_when_nothing_i
assert filled == rows
mock_prisma.db.query_raw.assert_not_called()
@pytest.mark.asyncio
async def test_fill_missing_api_key_aliases_keeps_spend_user_email_when_alias_is_missing():
double_hashed = hash_token("f" * 64)
mock_prisma = MagicMock()
mock_prisma.db.query_raw = _query_raw_by_table(
active_rows=[_digest_row(double_hashed, "team-key", "team-9", "key-owner")],
deleted_rows=[],
)
mock_prisma.db.litellm_usertable.find_many = AsyncMock(
return_value=[SimpleNamespace(user_id="key-owner", user_email="owner@example.com")]
)
rows = (
{
"api_key": double_hashed,
"api_key_alias": None,
"team_id": None,
"user_email": "spender@example.com",
"spend": 4.0,
},
)
filled = await fill_missing_api_key_aliases(mock_prisma, rows)
assert filled[0]["api_key_alias"] == "team-key"
assert filled[0]["team_id"] == "team-9"
assert filled[0]["user_email"] == "spender@example.com"
@pytest.mark.asyncio
async def test_fill_missing_api_key_aliases_skips_named_keys_that_have_no_email():
mock_prisma = MagicMock()
mock_prisma.db.query_raw = AsyncMock(return_value=[])
rows = (
{
"api_key": hash_token("g" * 64),
"api_key_alias": "service-key",
"team_id": "team-svc",
"user_email": None,
},
)
filled = await fill_missing_api_key_aliases(mock_prisma, rows)
assert filled == rows
mock_prisma.db.query_raw.assert_not_called()