mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-06 08:16:43 +00:00
fix(spend): keep CloudZero export and spend-log snapshots compatible with email recovery
Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
This commit is contained in:
parent
7e4032cfcc
commit
d3c839147e
4 changed files with 18 additions and 2 deletions
|
|
@ -98,9 +98,14 @@ class LiteLLMDatabase:
|
|||
fill_missing_api_key_aliases,
|
||||
)
|
||||
|
||||
usage_rows: Final = (
|
||||
db_response.to_dicts()
|
||||
if isinstance(db_response, pl.DataFrame)
|
||||
else db_response if isinstance(db_response, list) else []
|
||||
)
|
||||
# v1.99 double-hashed DailyUserSpend.api_key values miss the
|
||||
# VerificationToken join above; recover alias/team for those rows.
|
||||
recovered_rows: Final = await fill_missing_api_key_aliases(client, db_response)
|
||||
recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows)
|
||||
return pl.DataFrame(tuple(recovered_rows), infer_schema_length=None)
|
||||
except Exception as e:
|
||||
raise Exception(f"Error retrieving usage data: {e}")
|
||||
|
|
|
|||
|
|
@ -100,9 +100,14 @@ class FocusLiteLLMDatabase:
|
|||
fill_missing_api_key_aliases,
|
||||
)
|
||||
|
||||
usage_rows: Final = (
|
||||
db_response.to_dicts()
|
||||
if isinstance(db_response, pl.DataFrame)
|
||||
else db_response if isinstance(db_response, list) else []
|
||||
)
|
||||
# v1.99 double-hashed DailyUserSpend.api_key values miss the
|
||||
# VerificationToken join above; recover alias/team for those rows.
|
||||
recovered_rows: Final = await fill_missing_api_key_aliases(client, db_response)
|
||||
recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows)
|
||||
return pl.DataFrame(tuple(recovered_rows), infer_schema_length=None)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Error retrieving usage data: {exc}") from exc
|
||||
|
|
|
|||
|
|
@ -74,6 +74,8 @@ class TestCloudZeroHourlyExport:
|
|||
fake_db = MagicMock()
|
||||
|
||||
async def query_raw_mock(query: str, *params):
|
||||
if "LiteLLM_SpendLogs" in query:
|
||||
return []
|
||||
start_time_utc = params[0] if len(params) > 0 else None
|
||||
end_time_utc = params[1] if len(params) > 1 else None
|
||||
limit = params[2] if len(params) > 2 else None
|
||||
|
|
@ -146,6 +148,9 @@ class TestCloudZeroHourlyExport:
|
|||
return joined
|
||||
|
||||
fake_db.query_raw = AsyncMock(side_effect=query_raw_mock)
|
||||
fake_db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
fake_db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
|
||||
fake_db.litellm_usertable.find_many = AsyncMock(return_value=[])
|
||||
fake_client.db = fake_db
|
||||
mock_prisma_client_getter.return_value = fake_client
|
||||
|
||||
|
|
|
|||
|
|
@ -483,6 +483,7 @@ ignored_keys = [
|
|||
"metadata.user_api_key_project_alias",
|
||||
"metadata.user_api_key_org_id",
|
||||
"metadata.user_api_key_user_id",
|
||||
"metadata.user_api_key_user_email",
|
||||
"metadata.user_api_key_team_alias",
|
||||
"metadata.spend_logs_metadata",
|
||||
"metadata.requester_ip_address",
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue