mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-09 22:31:41 +00:00
Merge pull request #40176 from BerriAI/litellm_backport_1_100_x_spend_key_hash
chore(release): backport #39568 and #39572 to stable/1.100.x and cut 1.100.1
This commit is contained in:
commit
ecc04bf811
28 changed files with 922 additions and 58 deletions
|
|
@ -155,6 +155,7 @@ class CheckBatchCost:
|
|||
metadata: Dict[str, Any] = {
|
||||
"user_api_key_user_id": job.created_by,
|
||||
"user_api_key": api_key,
|
||||
"user_api_key_hash": api_key,
|
||||
"user_api_key_team_id": team_id,
|
||||
**(await self._get_user_info(batch_id, job.created_by)),
|
||||
}
|
||||
|
|
|
|||
|
|
@ -94,8 +94,18 @@ class LiteLLMDatabase:
|
|||
|
||||
try:
|
||||
db_response: Final = await client.db.query_raw(query, *params)
|
||||
# Convert the response to polars DataFrame with full schema inference
|
||||
# This prevents schema mismatch errors when data types vary across rows
|
||||
return pl.DataFrame(db_response, infer_schema_length=None)
|
||||
from litellm.proxy.spend_tracking.key_metadata_recovery import (
|
||||
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 []
|
||||
)
|
||||
recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows)
|
||||
return pl.DataFrame([dict(row) for row in recovered_rows], infer_schema_length=None)
|
||||
except Exception as e:
|
||||
raise Exception(f"Error retrieving usage data: {e}")
|
||||
|
|
|
|||
|
|
@ -96,7 +96,19 @@ class FocusLiteLLMDatabase:
|
|||
|
||||
try:
|
||||
db_response: Final = await client.db.query_raw(query, *query_params)
|
||||
return pl.DataFrame(db_response, infer_schema_length=None)
|
||||
from litellm.proxy.spend_tracking.key_metadata_recovery import (
|
||||
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 []
|
||||
)
|
||||
recovered_rows: Final = await fill_missing_api_key_aliases(client, usage_rows)
|
||||
return pl.DataFrame([dict(row) for row in recovered_rows], infer_schema_length=None)
|
||||
except Exception as exc:
|
||||
raise RuntimeError(f"Error retrieving usage data: {exc}") from exc
|
||||
|
||||
|
|
|
|||
|
|
@ -3151,6 +3151,17 @@
|
|||
}
|
||||
],
|
||||
"title": "Team Id"
|
||||
},
|
||||
"user_email": {
|
||||
"anyOf": [
|
||||
{
|
||||
"type": "string"
|
||||
},
|
||||
{
|
||||
"type": "null"
|
||||
}
|
||||
],
|
||||
"title": "User Email"
|
||||
}
|
||||
},
|
||||
"title": "KeyMetadata",
|
||||
|
|
|
|||
|
|
@ -88,8 +88,10 @@ class CallbackLogsReplayer:
|
|||
)
|
||||
|
||||
metadata: Final[dict[str, Any]] = payload.get("metadata") or {}
|
||||
user_api_key_hash: Final = metadata.get("user_api_key_hash")
|
||||
litellm_metadata: Final[dict[str, Any]] = {
|
||||
"user_api_key": metadata.get("user_api_key_hash"),
|
||||
"user_api_key": user_api_key_hash,
|
||||
"user_api_key_hash": user_api_key_hash,
|
||||
"user_api_key_alias": metadata.get("user_api_key_alias"),
|
||||
"user_api_key_user_id": metadata.get("user_api_key_user_id"),
|
||||
"user_api_key_team_id": metadata.get("user_api_key_team_id"),
|
||||
|
|
|
|||
|
|
@ -2,15 +2,19 @@ import asyncio
|
|||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from collections.abc import Set as AbstractSet
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from types import SimpleNamespace
|
||||
from types import MappingProxyType, SimpleNamespace
|
||||
from typing import TYPE_CHECKING, Final, Protocol
|
||||
|
||||
from fastapi import HTTPException, status
|
||||
from typing_extensions import TypedDict
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.constants import PTU_SENTINEL_API_KEY
|
||||
from litellm.proxy._types import CommonProxyErrors
|
||||
from litellm.proxy.spend_tracking.key_metadata_recovery import (
|
||||
attach_user_emails,
|
||||
recover_double_hashed_key_metadata,
|
||||
)
|
||||
from litellm.proxy.spend_tracking.ptu_feature_flag import is_ptu_cost_attribution_enabled
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.table_repositories import DeletedVerificationTokenRepository
|
||||
|
|
@ -111,8 +115,19 @@ class DailySpendRecord(Protocol):
|
|||
|
||||
|
||||
class _KeyMetadataDict(TypedDict, total=False):
|
||||
key_alias: str | None
|
||||
team_id: str | None
|
||||
key_alias: ReadOnly[str | None]
|
||||
team_id: ReadOnly[str | None]
|
||||
user_id: ReadOnly[str | None]
|
||||
user_email: ReadOnly[str | None]
|
||||
|
||||
|
||||
def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str) -> KeyMetadata:
|
||||
meta: Final = api_key_metadata.get(api_key, {})
|
||||
return KeyMetadata(
|
||||
key_alias=meta.get("key_alias"),
|
||||
team_id=meta.get("team_id"),
|
||||
user_email=meta.get("user_email"),
|
||||
)
|
||||
|
||||
|
||||
_WhereValue = str | dict[str, object]
|
||||
|
|
@ -283,10 +298,7 @@ def update_breakdown_metrics(
|
|||
if record.api_key not in breakdown.models[model_key].api_key_breakdown:
|
||||
breakdown.models[model_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
metadata=_key_metadata(api_key_metadata, record.api_key),
|
||||
)
|
||||
breakdown.models[model_key].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.models[model_key].api_key_breakdown[record.api_key].metrics,
|
||||
|
|
@ -310,10 +322,7 @@ def update_breakdown_metrics(
|
|||
if record.api_key not in breakdown.model_groups[model_group_key].api_key_breakdown:
|
||||
breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
metadata=_key_metadata(api_key_metadata, record.api_key),
|
||||
)
|
||||
breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.model_groups[model_group_key].api_key_breakdown[record.api_key].metrics,
|
||||
|
|
@ -335,10 +344,7 @@ def update_breakdown_metrics(
|
|||
breakdown.mcp_servers[record.mcp_namespaced_tool_name].api_key_breakdown[record.api_key] = (
|
||||
KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
metadata=_key_metadata(api_key_metadata, record.api_key),
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -363,10 +369,7 @@ def update_breakdown_metrics(
|
|||
if record.api_key not in breakdown.providers[provider].api_key_breakdown:
|
||||
breakdown.providers[provider].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
metadata=_key_metadata(api_key_metadata, record.api_key),
|
||||
)
|
||||
breakdown.providers[provider].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.providers[provider].api_key_breakdown[record.api_key].metrics,
|
||||
|
|
@ -388,10 +391,7 @@ def update_breakdown_metrics(
|
|||
if record.api_key not in breakdown.endpoints[record.endpoint].api_key_breakdown:
|
||||
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
metadata=_key_metadata(api_key_metadata, record.api_key),
|
||||
)
|
||||
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.endpoints[record.endpoint].api_key_breakdown[record.api_key].metrics,
|
||||
|
|
@ -403,10 +403,7 @@ def update_breakdown_metrics(
|
|||
if record.api_key not in breakdown.api_keys:
|
||||
breakdown.api_keys[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
), # Add any api_key-specific metadata here
|
||||
metadata=_key_metadata(api_key_metadata, record.api_key),
|
||||
)
|
||||
breakdown.api_keys[record.api_key].metrics = update_metrics(breakdown.api_keys[record.api_key].metrics, record)
|
||||
|
||||
|
|
@ -426,10 +423,7 @@ def update_breakdown_metrics(
|
|||
if record.api_key not in breakdown.entities[entity_value].api_key_breakdown:
|
||||
breakdown.entities[entity_value].api_key_breakdown[record.api_key] = KeyMetricWithMetadata(
|
||||
metrics=SpendMetrics(),
|
||||
metadata=KeyMetadata(
|
||||
key_alias=api_key_metadata.get(record.api_key, {}).get("key_alias", None),
|
||||
team_id=api_key_metadata.get(record.api_key, {}).get("team_id", None),
|
||||
),
|
||||
metadata=_key_metadata(api_key_metadata, record.api_key),
|
||||
)
|
||||
breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics = update_metrics(
|
||||
breakdown.entities[entity_value].api_key_breakdown[record.api_key].metrics,
|
||||
|
|
@ -442,17 +436,23 @@ def update_breakdown_metrics(
|
|||
async def get_api_key_metadata(
|
||||
prisma_client: PrismaClient,
|
||||
api_keys: AbstractSet[str],
|
||||
) -> dict[str, _KeyMetadataDict]:
|
||||
) -> Mapping[str, _KeyMetadataDict]:
|
||||
"""Get api key metadata, falling back to deleted keys table for keys not found in active table.
|
||||
|
||||
This ensures that key_alias and team_id are preserved in historical activity logs
|
||||
even after a key is deleted or regenerated.
|
||||
even after a key is deleted or regenerated. Also recovers aliases for api_key
|
||||
values that were double-hashed by the v1.99 spend-log provenance gate.
|
||||
"""
|
||||
key_records: Sequence[PrismaVerificationToken] = await VerificationTokenRepository(prisma_client).table.find_many(
|
||||
where={"token": {"in": list(api_keys)}}
|
||||
)
|
||||
result: Final[dict[str, _KeyMetadataDict]] = {
|
||||
k.token: {"key_alias": k.key_alias, "team_id": k.team_id} for k in key_records
|
||||
k.token: {
|
||||
"key_alias": k.key_alias,
|
||||
"team_id": k.team_id,
|
||||
"user_id": getattr(k, "user_id", None),
|
||||
}
|
||||
for k in key_records
|
||||
}
|
||||
|
||||
# For any keys not found in the active table, check the deleted keys table
|
||||
|
|
@ -471,6 +471,7 @@ async def get_api_key_metadata(
|
|||
result[k.token] = {
|
||||
"key_alias": k.key_alias,
|
||||
"team_id": k.team_id,
|
||||
"user_id": getattr(k, "user_id", None),
|
||||
}
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
|
|
@ -479,7 +480,13 @@ async def get_api_key_metadata(
|
|||
e,
|
||||
)
|
||||
|
||||
return result
|
||||
still_missing: Final = api_keys - frozenset(result)
|
||||
combined: Final = (
|
||||
result
|
||||
if not still_missing
|
||||
else MappingProxyType({**result, **(await recover_double_hashed_key_metadata(prisma_client, still_missing))})
|
||||
)
|
||||
return await attach_user_emails(prisma_client, combined)
|
||||
|
||||
|
||||
def _adjust_dates_for_timezone(
|
||||
|
|
@ -951,11 +958,6 @@ def _record_to_spend_metrics(record: _GroupingSetsRow) -> SpendMetrics:
|
|||
)
|
||||
|
||||
|
||||
def _key_metadata(api_key_metadata: Mapping[str, _KeyMetadataDict], api_key: str) -> KeyMetadata:
|
||||
meta: Final = api_key_metadata.get(api_key, {})
|
||||
return KeyMetadata(key_alias=meta.get("key_alias"), team_id=meta.get("team_id"))
|
||||
|
||||
|
||||
def _aggregate_grouping_sets_records_sync(
|
||||
*,
|
||||
records: Sequence[_GroupingSetsRow],
|
||||
|
|
|
|||
|
|
@ -15066,6 +15066,7 @@ async def async_queue_request(
|
|||
# extra_body); see above for the same guard upstream.
|
||||
data["metadata"] = {}
|
||||
data["metadata"]["user_api_key"] = user_api_key_dict.api_key
|
||||
data["metadata"]["user_api_key_hash"] = user_api_key_dict.api_key
|
||||
data["metadata"]["user_api_key_metadata"] = strip_callback_config(user_api_key_dict.metadata)
|
||||
_headers: Final = _safe_get_request_headers(request).copy()
|
||||
_headers.pop("authorization", None) # do not store the original `sk-..` api key in the db
|
||||
|
|
|
|||
236
litellm/proxy/spend_tracking/key_metadata_recovery.py
Normal file
236
litellm/proxy/spend_tracking/key_metadata_recovery.py
Normal file
|
|
@ -0,0 +1,236 @@
|
|||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from collections.abc import Set as AbstractSet
|
||||
from types import MappingProxyType
|
||||
from typing import Final, TypeVar
|
||||
|
||||
from pydantic import BaseModel, TypeAdapter
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.litellm_logging import is_valid_sha256_hash
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.user_repository import UserRepository
|
||||
|
||||
_T = TypeVar("_T")
|
||||
|
||||
_ACTIVE_TOKEN_DIGEST_SQL: Final = """
|
||||
SELECT encode(sha256(convert_to(token, 'UTF8')), 'hex') AS digest, key_alias, team_id, user_id
|
||||
FROM "LiteLLM_VerificationToken"
|
||||
WHERE encode(sha256(convert_to(token, 'UTF8')), 'hex') = ANY($1::text[])
|
||||
"""
|
||||
|
||||
_DELETED_TOKEN_DIGEST_SQL: Final = """
|
||||
SELECT DISTINCT ON (token)
|
||||
encode(sha256(convert_to(token, 'UTF8')), 'hex') AS digest, key_alias, team_id, user_id
|
||||
FROM "LiteLLM_DeletedVerificationToken"
|
||||
WHERE encode(sha256(convert_to(token, 'UTF8')), 'hex') = ANY($1::text[])
|
||||
ORDER BY token, deleted_at DESC
|
||||
"""
|
||||
|
||||
|
||||
class KeyMetadataDict(TypedDict, total=False):
|
||||
key_alias: ReadOnly[str | None]
|
||||
team_id: ReadOnly[str | None]
|
||||
user_id: ReadOnly[str | None]
|
||||
user_email: ReadOnly[str | None]
|
||||
|
||||
|
||||
class _TokenDigestRow(BaseModel):
|
||||
digest: str
|
||||
key_alias: str | None = None
|
||||
team_id: str | None = None
|
||||
user_id: str | None = None
|
||||
|
||||
|
||||
_TOKEN_DIGEST_ROWS: Final = TypeAdapter(tuple[_TokenDigestRow, ...])
|
||||
_EMPTY_KEY_METADATA: Final[Mapping[str, KeyMetadataDict]] = MappingProxyType({})
|
||||
_EMPTY_EMAILS: Final[Mapping[str, str]] = MappingProxyType({})
|
||||
|
||||
|
||||
async def _db_or_empty(
|
||||
load: Callable[[], Awaitable[_T]],
|
||||
warning: str,
|
||||
count: int,
|
||||
) -> _T | None:
|
||||
from prisma.errors import PrismaError
|
||||
|
||||
try:
|
||||
return await load()
|
||||
except PrismaError as e:
|
||||
verbose_proxy_logger.warning(warning, count, e)
|
||||
return None
|
||||
|
||||
|
||||
async def _reverse_hash_key_metadata(
|
||||
prisma_client: PrismaClient,
|
||||
sql: str,
|
||||
wanted: AbstractSet[str],
|
||||
*,
|
||||
warning: str,
|
||||
) -> Mapping[str, KeyMetadataDict]:
|
||||
rows: Final = await _db_or_empty(
|
||||
lambda: prisma_client.db.query_raw(sql, sorted(wanted)),
|
||||
warning,
|
||||
len(wanted),
|
||||
)
|
||||
if rows is None:
|
||||
return _EMPTY_KEY_METADATA
|
||||
return MappingProxyType(
|
||||
{
|
||||
row.digest: KeyMetadataDict(key_alias=row.key_alias, team_id=row.team_id, user_id=row.user_id)
|
||||
for row in _TOKEN_DIGEST_ROWS.validate_python(rows)
|
||||
if row.digest in wanted
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _emails_for_user_ids(
|
||||
prisma_client: PrismaClient,
|
||||
user_ids: AbstractSet[str],
|
||||
) -> Mapping[str, str]:
|
||||
if not user_ids:
|
||||
return _EMPTY_EMAILS
|
||||
users: Final = await _db_or_empty(
|
||||
lambda: UserRepository(prisma_client).table.find_many(
|
||||
where={"user_id": {"in": list(user_ids)}}, # mutable-ok: Prisma find_many where= is a dict
|
||||
),
|
||||
"Failed user_email recovery for %d user ids: %s",
|
||||
len(user_ids),
|
||||
)
|
||||
if users is None:
|
||||
return _EMPTY_EMAILS
|
||||
return MappingProxyType(
|
||||
{
|
||||
user.user_id: user.user_email
|
||||
for user in users
|
||||
if getattr(user, "user_id", None) and getattr(user, "user_email", None)
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _meta_with_email(meta: KeyMetadataDict, emails: Mapping[str, str]) -> KeyMetadataDict:
|
||||
if meta.get("user_email"):
|
||||
return meta
|
||||
user_id: Final = meta.get("user_id")
|
||||
if not isinstance(user_id, str) or user_id not in emails:
|
||||
return meta
|
||||
updated: Final[KeyMetadataDict] = {**meta, "user_email": emails[user_id]}
|
||||
return updated
|
||||
|
||||
|
||||
async def attach_user_emails(
|
||||
prisma_client: PrismaClient,
|
||||
recovered: Mapping[str, KeyMetadataDict],
|
||||
) -> Mapping[str, KeyMetadataDict]:
|
||||
needing_email: Final = frozenset(
|
||||
user_id
|
||||
for meta in recovered.values()
|
||||
for user_id in (meta.get("user_id"),)
|
||||
if isinstance(user_id, str) and user_id and not meta.get("user_email")
|
||||
)
|
||||
emails: Final = await _emails_for_user_ids(prisma_client, needing_email)
|
||||
if not emails:
|
||||
return recovered
|
||||
return MappingProxyType({api_key: _meta_with_email(meta, emails) for api_key, meta in recovered.items()})
|
||||
|
||||
|
||||
async def recover_double_hashed_key_metadata(
|
||||
prisma_client: PrismaClient,
|
||||
missing_keys: AbstractSet[str],
|
||||
) -> Mapping[str, KeyMetadataDict]:
|
||||
"""
|
||||
Recover key_alias/team_id/user_id for DailyUserSpend.api_key values that
|
||||
were double-hashed by the v1.99 spend-log provenance gate.
|
||||
|
||||
Those rows store hash(VerificationToken.token) instead of the token, so the
|
||||
exact join misses. Postgres hashes the token column itself, one pass over
|
||||
active keys and one over deleted keys, so no key row crosses the wire.
|
||||
"""
|
||||
sha_missing: Final = frozenset(key for key in missing_keys if is_valid_sha256_hash(key))
|
||||
if not sha_missing:
|
||||
return _EMPTY_KEY_METADATA
|
||||
|
||||
from_active: Final = await _reverse_hash_key_metadata(
|
||||
prisma_client,
|
||||
_ACTIVE_TOKEN_DIGEST_SQL,
|
||||
sha_missing,
|
||||
warning="Failed reverse-hash recovery against active keys for %d missing keys: %s",
|
||||
)
|
||||
still_missing: Final = sha_missing - frozenset(from_active)
|
||||
if not still_missing:
|
||||
return from_active
|
||||
from_deleted: Final = await _reverse_hash_key_metadata(
|
||||
prisma_client,
|
||||
_DELETED_TOKEN_DIGEST_SQL,
|
||||
still_missing,
|
||||
warning="Failed reverse-hash recovery against deleted keys for %d missing keys: %s",
|
||||
)
|
||||
return MappingProxyType({**from_active, **from_deleted})
|
||||
|
||||
|
||||
def _row_with_recovered_fields(
|
||||
row: Mapping[str, object],
|
||||
recovered: Mapping[str, KeyMetadataDict],
|
||||
*,
|
||||
api_key_field: str,
|
||||
alias_field: str,
|
||||
team_id_field: str,
|
||||
user_email_field: str,
|
||||
) -> Mapping[str, object]:
|
||||
api_key: Final = row.get(api_key_field)
|
||||
if not isinstance(api_key, str) or api_key not in recovered:
|
||||
return row
|
||||
meta: Final = recovered[api_key]
|
||||
return MappingProxyType(
|
||||
{
|
||||
**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: row.get(user_email_field) or meta.get("user_email"),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def fill_missing_api_key_aliases(
|
||||
prisma_client: PrismaClient,
|
||||
rows: Sequence[Mapping[str, object]],
|
||||
*,
|
||||
api_key_field: str = "api_key",
|
||||
alias_field: str = "api_key_alias",
|
||||
team_id_field: str = "team_id",
|
||||
user_email_field: str = "user_email",
|
||||
) -> tuple[Mapping[str, object], ...]:
|
||||
"""
|
||||
Fill null api_key_alias / team_id / user_email on export rows whose api_key
|
||||
was double-hashed.
|
||||
|
||||
Used by CloudZero and Focus, which join DailyUserSpend.api_key to
|
||||
VerificationToken.token and otherwise export null aliases for those rows.
|
||||
"""
|
||||
missing_keys: Final = frozenset(
|
||||
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, "")
|
||||
)
|
||||
if not missing_keys:
|
||||
return tuple(rows)
|
||||
|
||||
recovered: Final = await attach_user_emails(
|
||||
prisma_client,
|
||||
await recover_double_hashed_key_metadata(prisma_client, missing_keys),
|
||||
)
|
||||
if not recovered:
|
||||
return tuple(rows)
|
||||
|
||||
return tuple(
|
||||
_row_with_recovered_fields(
|
||||
row,
|
||||
recovered,
|
||||
api_key_field=api_key_field,
|
||||
alias_field=alias_field,
|
||||
team_id_field=team_id_field,
|
||||
user_email_field=user_email_field,
|
||||
)
|
||||
for row in rows
|
||||
)
|
||||
|
|
@ -14,6 +14,8 @@ from litellm.constants import (
|
|||
LITELLM_PROXY_MASTER_KEY_ALIAS,
|
||||
LITELLM_TRUNCATED_PAYLOAD_FIELD,
|
||||
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE,
|
||||
LITTELM_CLI_SERVICE_ACCOUNT_NAME,
|
||||
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
|
||||
REDACTED_BY_LITELM_STRING,
|
||||
)
|
||||
from litellm.constants import (
|
||||
|
|
@ -72,13 +74,18 @@ def _is_master_key(api_key: str | None, _master_key: str | None) -> bool:
|
|||
|
||||
|
||||
_HASHED_JWT_RE = re.compile(r"hashed-jwt-[a-fA-F0-9]{64}")
|
||||
_NON_SECRET_KEY_ALIASES: Final = frozenset(
|
||||
{
|
||||
LITELLM_PROXY_MASTER_KEY_ALIAS,
|
||||
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
|
||||
LITTELM_CLI_SERVICE_ACCOUNT_NAME,
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
def _is_non_secret_key_value(value: str) -> bool:
|
||||
return (
|
||||
value == LITELLM_PROXY_MASTER_KEY_ALIAS
|
||||
or is_valid_sha256_hash(value)
|
||||
or _HASHED_JWT_RE.fullmatch(value) is not None
|
||||
value in _NON_SECRET_KEY_ALIASES or is_valid_sha256_hash(value) or _HASHED_JWT_RE.fullmatch(value) is not None
|
||||
)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -43,6 +43,7 @@ class KeyMetadata(BaseModel):
|
|||
|
||||
key_alias: str | None = None
|
||||
team_id: str | None = None
|
||||
user_email: str | None = None
|
||||
|
||||
|
||||
class KeyMetricWithMetadata(MetricBase):
|
||||
|
|
|
|||
|
|
@ -1,6 +1,6 @@
|
|||
[project]
|
||||
name = "litellm"
|
||||
version = "1.100.0"
|
||||
version = "1.100.1"
|
||||
description = "Library to easily interface with LLM API providers"
|
||||
readme = "README.md"
|
||||
requires-python = ">=3.10, <3.15"
|
||||
|
|
@ -311,7 +311,7 @@ members = ["enterprise", "litellm-proxy-extras"]
|
|||
profile = "black"
|
||||
|
||||
[tool.commitizen]
|
||||
version = "1.100.0"
|
||||
version = "1.100.1"
|
||||
version_files = [
|
||||
"pyproject.toml:^version",
|
||||
]
|
||||
|
|
|
|||
|
|
@ -2445,6 +2445,7 @@ class TestBatchCostAttribution:
|
|||
metadata = await instance._build_creator_attribution_metadata(self._job(), "batch-1")
|
||||
|
||||
assert metadata["user_api_key"] == "hash-alice"
|
||||
assert metadata["user_api_key_hash"] == "hash-alice"
|
||||
assert metadata["user_api_key_user_id"] == "alice"
|
||||
assert metadata["user_api_key_team_id"] == "team-alpha"
|
||||
assert metadata["user_api_key_alias"] == "prod-key"
|
||||
|
|
@ -2553,6 +2554,48 @@ class TestBatchCostAttribution:
|
|||
|
||||
assert metadata["user_api_key_alias"] == "prod-key"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_metadata_provenance_keeps_spend_log_api_key_joinable(self):
|
||||
"""
|
||||
CheckBatchCost stores the VerificationToken hash on the managed object. The
|
||||
spend-log writer must receive matching user_api_key_hash provenance so it
|
||||
does not re-hash that value; otherwise DailyUserSpend.api_key no longer joins
|
||||
VerificationToken and Usage shows key-hash-... with a null alias/email.
|
||||
"""
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
|
||||
from litellm.proxy.utils import hash_token
|
||||
|
||||
token_hash = hash_token("sk-batch-creator-key")
|
||||
instance = self._instance(
|
||||
key_row=SimpleNamespace(key_alias="prod-key"),
|
||||
user_row=SimpleNamespace(user_email="alice@example.com", user_alias=None),
|
||||
)
|
||||
metadata = await instance._build_creator_attribution_metadata(
|
||||
self._job(api_key=token_hash), "batch-1"
|
||||
)
|
||||
|
||||
assert metadata["user_api_key"] == token_hash
|
||||
assert metadata["user_api_key_hash"] == token_hash
|
||||
|
||||
payload = get_logging_payload(
|
||||
kwargs={
|
||||
"model": "gpt-4o",
|
||||
"call_type": "aretrieve_batch",
|
||||
"litellm_params": {"metadata": metadata},
|
||||
},
|
||||
response_obj={
|
||||
"id": "batch_123",
|
||||
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
|
||||
},
|
||||
start_time=datetime.now(timezone.utc),
|
||||
end_time=datetime.now(timezone.utc),
|
||||
)
|
||||
assert payload["api_key"] == token_hash
|
||||
assert payload["api_key"] != hash_token(token_hash)
|
||||
|
||||
|
||||
class TestPollPageStarvation:
|
||||
"""LIT-5462 regression: a row that can never be costed used to keep its slot in the
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import hashlib
|
||||
from datetime import datetime
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
|
|
@ -74,6 +75,8 @@ class TestCloudZeroHourlyExport:
|
|||
fake_db = MagicMock()
|
||||
|
||||
async def query_raw_mock(query: str, *params):
|
||||
if "sha256(" 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 +149,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
|
||||
|
||||
|
|
@ -160,3 +166,26 @@ class TestCloudZeroHourlyExport:
|
|||
logger = CloudZeroLogger(api_key="test", connection_id="test")
|
||||
|
||||
await logger._hourly_usage_data_export()
|
||||
|
||||
|
||||
class TestLiteLLMDatabaseUsageData:
|
||||
@pytest.mark.asyncio
|
||||
async def test_builds_frame_from_rows_recovered_for_double_hashed_keys(self, monkeypatch: pytest.MonkeyPatch):
|
||||
double_hashed = hashlib.sha256(b"sk-hashed-token").hexdigest()
|
||||
joined_row = {"api_key": "sk-joined", "api_key_alias": "joined", "team_id": "team-0", "user_email": None, "spend": 0.1}
|
||||
dirty_row = {"api_key": double_hashed, "api_key_alias": None, "team_id": None, "user_email": None, "spend": 0.5}
|
||||
|
||||
async def query_raw(query: str, *params):
|
||||
if "sha256(" in query:
|
||||
return [{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": None}]
|
||||
return [joined_row, dirty_row]
|
||||
|
||||
fake_client = MagicMock()
|
||||
fake_client.db.query_raw = AsyncMock(side_effect=query_raw)
|
||||
db = LiteLLMDatabase()
|
||||
monkeypatch.setattr(db, "_ensure_prisma_client", lambda: fake_client)
|
||||
|
||||
result = await db.get_usage_data()
|
||||
|
||||
assert result["api_key_alias"].to_list() == ["joined", "batch-worker"]
|
||||
assert result["team_id"].to_list() == ["team-0", "team-1"]
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
"""Tests for FocusLiteLLMDatabase query construction."""
|
||||
|
||||
import hashlib
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock
|
||||
|
|
@ -87,3 +88,24 @@ async def test_should_join_organization_table(monkeypatch: pytest.MonkeyPatch):
|
|||
)
|
||||
assert "ot.organization_alias as organization_alias" in query_text
|
||||
assert 'LEFT JOIN "LiteLLM_OrganizationTable" ot' in query_text
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_should_build_frame_from_rows_recovered_for_double_hashed_keys(monkeypatch: pytest.MonkeyPatch):
|
||||
double_hashed = hashlib.sha256(b"sk-hashed-token").hexdigest()
|
||||
joined_row = {"api_key": "sk-joined", "api_key_alias": "joined", "team_id": "team-0", "user_email": None, "spend": 0.1}
|
||||
dirty_row = {"api_key": double_hashed, "api_key_alias": None, "team_id": None, "user_email": None, "spend": 0.5}
|
||||
|
||||
async def query_raw(query: str, *params):
|
||||
if "sha256(" in query:
|
||||
return [{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": None}]
|
||||
return [joined_row, dirty_row]
|
||||
|
||||
mock_client = SimpleNamespace(db=SimpleNamespace(query_raw=AsyncMock(side_effect=query_raw)))
|
||||
db = FocusLiteLLMDatabase()
|
||||
monkeypatch.setattr(db, "_ensure_prisma_client", lambda: mock_client)
|
||||
|
||||
result = await db.get_usage_data()
|
||||
|
||||
assert result["api_key_alias"].to_list() == ["joined", "batch-worker"]
|
||||
assert result["team_id"].to_list() == ["team-0", "team-1"]
|
||||
|
|
|
|||
|
|
@ -61,6 +61,7 @@ def test_build_logging_obj_seeds_model_call_details():
|
|||
# Metadata is mapped to the keys the cost-tracking callback reads.
|
||||
md = details["litellm_params"]["metadata"]
|
||||
assert md["user_api_key"] == "rust-gateway-test-key"
|
||||
assert md["user_api_key_hash"] == "rust-gateway-test-key"
|
||||
assert md["user_api_key_user_id"] == "user-cb-logs-test"
|
||||
assert md["user_api_key_team_id"] == "team-cb-logs-test"
|
||||
|
||||
|
|
|
|||
|
|
@ -454,6 +454,151 @@ async def test_get_api_key_metadata_regenerated_key_uses_most_recent_deleted_rec
|
|||
assert result["old-key-hash"]["team_id"] == "latest-team"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_api_key_metadata_recovers_double_hashed_key_via_reverse_hash():
|
||||
"""
|
||||
v1.99 spend logging re-hashed already-hashed api_key values when provenance was
|
||||
missing. Usage joins DailyUserSpend.api_key to VerificationToken.token, so those
|
||||
rows looked like key-hash-... with a null alias. Recovery asks Postgres for the
|
||||
key whose hashed token matches the dirty value and maps it back to its alias.
|
||||
"""
|
||||
from litellm.proxy.utils import hash_token
|
||||
|
||||
double_hashed = hash_token("a" * 64)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_usertable.find_many = AsyncMock(
|
||||
return_value=[SimpleNamespace(user_id="alice", user_email="alice@example.com")]
|
||||
)
|
||||
mock_prisma.db.query_raw = AsyncMock(
|
||||
return_value=[
|
||||
{"digest": double_hashed, "key_alias": "batch-worker", "team_id": "team-1", "user_id": "alice"}
|
||||
]
|
||||
)
|
||||
|
||||
result = await get_api_key_metadata(
|
||||
prisma_client=mock_prisma,
|
||||
api_keys={double_hashed},
|
||||
)
|
||||
|
||||
assert result[double_hashed]["key_alias"] == "batch-worker"
|
||||
assert result[double_hashed]["team_id"] == "team-1"
|
||||
assert result[double_hashed]["user_email"] == "alice@example.com"
|
||||
((digest_sql, digests),) = [call.args for call in mock_prisma.db.query_raw.call_args_list]
|
||||
assert '"LiteLLM_VerificationToken"' in digest_sql
|
||||
assert digests == [double_hashed]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_get_api_key_metadata_permanent_miss_never_pages_tokens_or_reads_spend_logs():
|
||||
"""A dirty key no table can explain costs two digest lookups, never a token page walk or a SpendLogs scan."""
|
||||
from litellm.proxy.utils import hash_token
|
||||
|
||||
double_hashed = hash_token("b" * 64)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.litellm_usertable.find_many = AsyncMock(return_value=[])
|
||||
mock_prisma.db.query_raw = AsyncMock(return_value=[])
|
||||
|
||||
result = await get_api_key_metadata(
|
||||
prisma_client=mock_prisma,
|
||||
api_keys={double_hashed},
|
||||
)
|
||||
|
||||
assert double_hashed not in result
|
||||
issued_sql = [call.args[0] for call in mock_prisma.db.query_raw.call_args_list]
|
||||
assert len(issued_sql) == 2
|
||||
assert not any("LiteLLM_SpendLogs" in sql for sql in issued_sql)
|
||||
token_lookups = (
|
||||
mock_prisma.db.litellm_verificationtoken.find_many.call_args_list
|
||||
+ mock_prisma.db.litellm_deletedverificationtoken.find_many.call_args_list
|
||||
)
|
||||
assert all("take" not in call.kwargs and "skip" not in call.kwargs for call in token_lookups)
|
||||
|
||||
|
||||
def test_key_metadata_includes_recovered_user_email():
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import _key_metadata
|
||||
|
||||
meta = _key_metadata(
|
||||
{
|
||||
"dirty-key": {
|
||||
"key_alias": "batch-worker",
|
||||
"team_id": "team-1",
|
||||
"user_email": "alice@example.com",
|
||||
}
|
||||
},
|
||||
"dirty-key",
|
||||
)
|
||||
|
||||
assert meta.key_alias == "batch-worker"
|
||||
assert meta.user_email == "alice@example.com"
|
||||
|
||||
|
||||
def test_update_breakdown_metrics_includes_user_email():
|
||||
from litellm.proxy.management_endpoints.common_daily_activity import update_breakdown_metrics
|
||||
from litellm.types.proxy.management_endpoints.common_daily_activity import BreakdownMetrics
|
||||
|
||||
breakdown = BreakdownMetrics()
|
||||
record = SimpleNamespace(
|
||||
api_key="dirty-key",
|
||||
model="gpt-4o-mini",
|
||||
model_group="grp",
|
||||
mcp_namespaced_tool_name="srv/tool",
|
||||
custom_llm_provider="openai",
|
||||
endpoint="/v1/chat/completions",
|
||||
spend=1.23,
|
||||
prompt_tokens=1,
|
||||
completion_tokens=1,
|
||||
cache_read_input_tokens=0,
|
||||
cache_creation_input_tokens=0,
|
||||
compression_saved_tokens=0,
|
||||
compression_savings_spend=0,
|
||||
prompt_caching_savings_spend=0,
|
||||
gateway_injected_caching_savings_spend=0,
|
||||
autorouter_savings_spend=0,
|
||||
total_tokens=2,
|
||||
api_requests=1,
|
||||
successful_requests=1,
|
||||
failed_requests=0,
|
||||
ptu_flat_cost=0.0,
|
||||
user_id="alice",
|
||||
)
|
||||
api_key_metadata = {
|
||||
"dirty-key": {
|
||||
"key_alias": "batch-worker",
|
||||
"team_id": "team-1",
|
||||
"user_email": "alice@example.com",
|
||||
}
|
||||
}
|
||||
|
||||
update_breakdown_metrics(
|
||||
breakdown,
|
||||
record,
|
||||
{},
|
||||
{},
|
||||
api_key_metadata,
|
||||
entity_id_field="user_id",
|
||||
)
|
||||
|
||||
expected = ("batch-worker", "alice@example.com")
|
||||
top = breakdown.api_keys["dirty-key"].metadata
|
||||
assert (top.key_alias, top.user_email) == expected
|
||||
assert (
|
||||
breakdown.models["gpt-4o-mini"].api_key_breakdown["dirty-key"].metadata.key_alias,
|
||||
breakdown.models["gpt-4o-mini"].api_key_breakdown["dirty-key"].metadata.user_email,
|
||||
) == expected
|
||||
assert (
|
||||
breakdown.providers["openai"].api_key_breakdown["dirty-key"].metadata.key_alias,
|
||||
breakdown.providers["openai"].api_key_breakdown["dirty-key"].metadata.user_email,
|
||||
) == expected
|
||||
assert (
|
||||
breakdown.entities["alice"].api_key_breakdown["dirty-key"].metadata.key_alias,
|
||||
breakdown.entities["alice"].api_key_breakdown["dirty-key"].metadata.user_email,
|
||||
) == expected
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_tag_daily_activity_metadata_totals_not_zero():
|
||||
"""Test that tag daily activity returns correct metadata totals.
|
||||
|
|
|
|||
|
|
@ -0,0 +1,220 @@
|
|||
from collections.abc import Sequence
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import pytest
|
||||
from prisma.errors import PrismaError
|
||||
|
||||
from litellm.proxy.spend_tracking.key_metadata_recovery import (
|
||||
fill_missing_api_key_aliases,
|
||||
recover_double_hashed_key_metadata,
|
||||
)
|
||||
from litellm.proxy.utils import hash_token
|
||||
|
||||
|
||||
def _digest_row(digest: str, key_alias: str, team_id: str | None, user_id: str | None) -> dict[str, str | None]:
|
||||
return {"digest": digest, "key_alias": key_alias, "team_id": team_id, "user_id": user_id}
|
||||
|
||||
|
||||
def _query_raw_by_table(
|
||||
active_rows: Sequence[dict[str, str | None]],
|
||||
deleted_rows: Sequence[dict[str, str | None]],
|
||||
) -> AsyncMock:
|
||||
async def query_raw(sql: str, *params: object) -> list[dict[str, str | None]]:
|
||||
if '"LiteLLM_VerificationToken"' in sql:
|
||||
return list(active_rows)
|
||||
if '"LiteLLM_DeletedVerificationToken"' in sql:
|
||||
return list(deleted_rows)
|
||||
raise AssertionError(f"unexpected query: {sql}")
|
||||
|
||||
return AsyncMock(side_effect=query_raw)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recover_double_hashed_key_metadata_via_active_token_digest():
|
||||
double_hashed = hash_token("a" * 64)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.query_raw = _query_raw_by_table(
|
||||
active_rows=[_digest_row(double_hashed, "batch-worker", "team-1", "alice")],
|
||||
deleted_rows=[],
|
||||
)
|
||||
|
||||
result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed})
|
||||
|
||||
assert result[double_hashed]["key_alias"] == "batch-worker"
|
||||
assert result[double_hashed]["team_id"] == "team-1"
|
||||
assert result[double_hashed]["user_id"] == "alice"
|
||||
((_, digests),) = [call.args for call in mock_prisma.db.query_raw.call_args_list]
|
||||
assert digests == [double_hashed]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recover_double_hashed_key_metadata_falls_back_to_deleted_tokens():
|
||||
double_hashed = hash_token("y" * 64)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.query_raw = _query_raw_by_table(
|
||||
active_rows=[],
|
||||
deleted_rows=[_digest_row(double_hashed, "deleted-key", "team-del", "erin")],
|
||||
)
|
||||
|
||||
result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed})
|
||||
|
||||
assert result[double_hashed]["key_alias"] == "deleted-key"
|
||||
assert result[double_hashed]["team_id"] == "team-del"
|
||||
assert result[double_hashed]["user_id"] == "erin"
|
||||
assert [call.args[1] for call in mock_prisma.db.query_raw.call_args_list] == [[double_hashed], [double_hashed]]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recover_only_asks_deleted_tokens_for_digests_active_keys_missed():
|
||||
found_active = hash_token("1" * 64)
|
||||
found_deleted = hash_token("2" * 64)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.query_raw = _query_raw_by_table(
|
||||
active_rows=[_digest_row(found_active, "active-key", None, None)],
|
||||
deleted_rows=[_digest_row(found_deleted, "deleted-key", None, None)],
|
||||
)
|
||||
|
||||
result = await recover_double_hashed_key_metadata(mock_prisma, {found_active, found_deleted})
|
||||
|
||||
assert result[found_active]["key_alias"] == "active-key"
|
||||
assert result[found_deleted]["key_alias"] == "deleted-key"
|
||||
assert [call.args[1] for call in mock_prisma.db.query_raw.call_args_list] == [
|
||||
sorted((found_active, found_deleted)),
|
||||
[found_deleted],
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recover_permanent_miss_costs_two_digest_lookups_and_no_table_walk():
|
||||
double_hashed = hash_token("b" * 64)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.query_raw = _query_raw_by_table(active_rows=[], deleted_rows=[])
|
||||
|
||||
result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed})
|
||||
|
||||
assert result == {}
|
||||
assert len(mock_prisma.db.query_raw.call_args_list) == 2
|
||||
mock_prisma.db.litellm_verificationtoken.find_many.assert_not_called()
|
||||
mock_prisma.db.litellm_deletedverificationtoken.find_many.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recover_skips_keys_that_are_not_sha256_digests():
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(return_value=[])
|
||||
|
||||
result = await recover_double_hashed_key_metadata(mock_prisma, {"sk-plain-key", "key-hash-short"})
|
||||
|
||||
assert result == {}
|
||||
mock_prisma.db.query_raw.assert_not_called()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_recover_returns_empty_when_digest_lookup_raises_prisma_error():
|
||||
double_hashed = hash_token("c" * 64)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(side_effect=PrismaError("db down"))
|
||||
|
||||
result = await recover_double_hashed_key_metadata(mock_prisma, {double_hashed})
|
||||
|
||||
assert result == {}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fill_missing_api_key_aliases_updates_null_alias_and_email_rows():
|
||||
double_hashed = hash_token("d" * 64)
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.query_raw = _query_raw_by_table(
|
||||
active_rows=[_digest_row(double_hashed, "recovered-alias", "team-9", "bob")],
|
||||
deleted_rows=[],
|
||||
)
|
||||
mock_prisma.db.litellm_usertable.find_many = AsyncMock(
|
||||
return_value=[SimpleNamespace(user_id="bob", user_email="bob@example.com")]
|
||||
)
|
||||
|
||||
rows = (
|
||||
{
|
||||
"api_key": double_hashed,
|
||||
"api_key_alias": None,
|
||||
"team_id": None,
|
||||
"user_email": None,
|
||||
"spend": 12.5,
|
||||
},
|
||||
{
|
||||
"api_key": "already-joined-token",
|
||||
"api_key_alias": "named-key",
|
||||
"team_id": "team-ok",
|
||||
"user_email": "other@example.com",
|
||||
"spend": 1.0,
|
||||
},
|
||||
)
|
||||
|
||||
filled = await fill_missing_api_key_aliases(mock_prisma, rows)
|
||||
|
||||
assert filled[0]["api_key_alias"] == "recovered-alias"
|
||||
assert filled[0]["team_id"] == "team-9"
|
||||
assert filled[0]["user_email"] == "bob@example.com"
|
||||
assert filled[0]["spend"] == 12.5
|
||||
assert filled[1]["api_key_alias"] == "named-key"
|
||||
assert mock_prisma.db.litellm_usertable.find_many.call_args.kwargs["where"] == {"user_id": {"in": ["bob"]}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_fill_missing_api_key_aliases_leaves_rows_untouched_when_nothing_is_missing():
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db.query_raw = AsyncMock(return_value=[])
|
||||
rows = ({"api_key": hash_token("e" * 64), "api_key_alias": "named", "user_email": "x@example.com"},)
|
||||
|
||||
filled = await fill_missing_api_key_aliases(mock_prisma, rows)
|
||||
|
||||
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()
|
||||
|
|
@ -12,9 +12,13 @@ import litellm
|
|||
from litellm.constants import (
|
||||
LITELLM_TRUNCATED_PAYLOAD_FIELD,
|
||||
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE,
|
||||
LITTELM_CLI_SERVICE_ACCOUNT_NAME,
|
||||
LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME,
|
||||
REDACTED_BY_LITELM_STRING,
|
||||
)
|
||||
from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup
|
||||
from litellm.proxy.spend_tracking.spend_tracking_utils import (
|
||||
_get_messages_for_spend_logs_payload,
|
||||
_get_proxy_server_request_for_spend_logs_payload,
|
||||
|
|
@ -2754,6 +2758,41 @@ def test_get_spend_logs_metadata_already_hashed_no_provenance_is_rehashed():
|
|||
assert meta["user_api_key"] == hash_token(already_hashed)
|
||||
|
||||
|
||||
def test_get_logging_payload_batch_attribution_keeps_verification_token_hash():
|
||||
"""
|
||||
Batch cost rebuilds metadata with the managed object's already-hashed api_key.
|
||||
That hash must land in SpendLogs.api_key unchanged so Usage/CloudZero can join
|
||||
LiteLLM_VerificationToken for api_key_alias and user_email. Regression: without
|
||||
user_api_key_hash provenance, v1.99+ re-hashed the token and broke the join.
|
||||
"""
|
||||
token_hash = hash_token("sk-batch-creator-key")
|
||||
kwargs = {
|
||||
"model": "gpt-4o",
|
||||
"call_type": "aretrieve_batch",
|
||||
"litellm_params": {
|
||||
"metadata": {
|
||||
"user_api_key": token_hash,
|
||||
"user_api_key_hash": token_hash,
|
||||
"user_api_key_alias": "batch-creator",
|
||||
"user_api_key_user_id": "alice",
|
||||
"user_api_key_team_id": "team-1",
|
||||
}
|
||||
},
|
||||
}
|
||||
payload = get_logging_payload(
|
||||
kwargs=kwargs,
|
||||
response_obj={"id": "batch_123", "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}},
|
||||
start_time=datetime.datetime.now(timezone.utc),
|
||||
end_time=datetime.datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
assert payload["api_key"] == token_hash
|
||||
assert payload["api_key"] != hash_token(token_hash)
|
||||
parsed_meta = json.loads(payload["metadata"])
|
||||
assert parsed_meta["user_api_key"] == token_hash
|
||||
assert parsed_meta["user_api_key_alias"] == "batch-creator"
|
||||
|
||||
|
||||
def test_get_spend_logs_metadata_provenance_bypass_requires_hash_match():
|
||||
already_hashed = hash_token("sk-some-key")
|
||||
different_hash = hash_token("sk-other-key")
|
||||
|
|
@ -3043,6 +3082,45 @@ def test_get_logging_payload_keeps_master_key_alias_readable():
|
|||
assert parsed_meta["user_api_key"] == LITELLM_PROXY_MASTER_KEY_ALIAS
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"service_account",
|
||||
[LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME, LITTELM_CLI_SERVICE_ACCOUNT_NAME],
|
||||
)
|
||||
def test_get_logging_payload_keeps_internal_service_account_key_readable(service_account: str):
|
||||
data = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata(
|
||||
data={"metadata": {}},
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key=service_account,
|
||||
team_id=service_account,
|
||||
key_alias=service_account,
|
||||
team_alias=service_account,
|
||||
),
|
||||
_metadata_variable_name="metadata",
|
||||
)
|
||||
kwargs = {
|
||||
"model": "openai/gpt-4.1",
|
||||
"messages": [{"role": "user", "content": "Hello"}],
|
||||
"call_type": "acompletion",
|
||||
"litellm_params": {"metadata": data["metadata"]},
|
||||
}
|
||||
payload = get_logging_payload(
|
||||
kwargs=kwargs,
|
||||
response_obj=Exception("error"),
|
||||
start_time=datetime.datetime.now(timezone.utc),
|
||||
end_time=datetime.datetime.now(timezone.utc),
|
||||
)
|
||||
|
||||
assert payload["api_key"] == service_account
|
||||
parsed_meta = json.loads(payload["metadata"])
|
||||
assert parsed_meta["user_api_key"] == service_account
|
||||
assert parsed_meta["user_api_key_alias"] == service_account
|
||||
|
||||
|
||||
def test_redact_logged_api_key_service_account_name_without_provenance_is_hashed():
|
||||
result = _redact_logged_api_key(LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME)
|
||||
assert result == hash_token(LITTELM_INTERNAL_HEALTH_SERVICE_ACCOUNT_NAME)
|
||||
|
||||
|
||||
@patch("litellm.proxy.proxy_server.master_key", None)
|
||||
@patch("litellm.proxy.proxy_server.general_settings", {})
|
||||
def test_get_logging_payload_hashes_bearer_prefixed_api_key():
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel";
|
||||
import { BreakdownMetrics, DailyData, KeyMetricWithMetadata, TagUsage } from "@/components/UsagePage/types";
|
||||
|
||||
export type ExtendedDailyData = DailyData & {
|
||||
|
|
@ -118,6 +119,7 @@ export const getTopAPIKeys = (results: ExtendedDailyData[], topKeysLimit: number
|
|||
metadata: {
|
||||
key_alias: metrics.metadata.key_alias,
|
||||
team_id: metrics.metadata.team_id || null,
|
||||
user_email: metrics.metadata.user_email,
|
||||
tags: tagDictionary[key] || [],
|
||||
},
|
||||
};
|
||||
|
|
@ -137,7 +139,7 @@ export const getTopAPIKeys = (results: ExtendedDailyData[], topKeysLimit: number
|
|||
return Object.entries(keySpend)
|
||||
.map(([api_key, metrics]) => ({
|
||||
api_key,
|
||||
key_alias: metrics.metadata.key_alias || "-", // Using truncated key as alias
|
||||
key_alias: keyActivityLabel(metrics.metadata),
|
||||
tags: metrics.metadata.tags || "-",
|
||||
spend: metrics.metrics.spend,
|
||||
}))
|
||||
|
|
|
|||
|
|
@ -44,6 +44,7 @@ import { Tag } from "@/components/tag_management/types";
|
|||
import UserAgentActivity from "@/components/user_agent_activity";
|
||||
import ViewUserSpend from "@/components/view_user_spend";
|
||||
import { usePaginatedDailyActivity } from "../hooks/usePaginatedDailyActivity";
|
||||
import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel";
|
||||
import { DailyData, KeyMetricWithMetadata, MetricWithMetadata } from "@/components/UsagePage/types";
|
||||
import { valueFormatterSpend } from "@/components/UsagePage/utils/value_formatters";
|
||||
import {
|
||||
|
|
@ -426,6 +427,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
|
|||
metadata: {
|
||||
key_alias: metrics.metadata.key_alias,
|
||||
team_id: null,
|
||||
user_email: metrics.metadata.user_email,
|
||||
tags: metrics.metadata.tags || [],
|
||||
},
|
||||
};
|
||||
|
|
@ -445,7 +447,7 @@ const UsagePage: React.FC<UsagePageProps> = ({ teams, organizations }) => {
|
|||
return Object.entries(keySpend)
|
||||
.map(([api_key, metrics]) => ({
|
||||
api_key,
|
||||
key_alias: metrics.metadata.key_alias || "-",
|
||||
key_alias: keyActivityLabel(metrics.metadata),
|
||||
tags: metrics.metadata.tags || [],
|
||||
spend: metrics.metrics.spend,
|
||||
}))
|
||||
|
|
|
|||
|
|
@ -1,6 +1,7 @@
|
|||
import { formatNumberWithCommas } from "@/utils/dataUtils";
|
||||
import type { DateRangePickerValue } from "@/components/shared/date_picker_types";
|
||||
import Papa from "papaparse";
|
||||
import { keyActivityLabel } from "@/components/UsagePage/keyActivityLabel";
|
||||
import type { EntityBreakdown, EntitySpendData, EntityType, ExportMetadata, ExportScope } from "./types";
|
||||
|
||||
const resolveEntityDisplay = (
|
||||
|
|
@ -186,7 +187,7 @@ export const generateDailyWithKeysData = (
|
|||
|
||||
// Iterate through each API key in the breakdown
|
||||
Object.entries(apiKeyBreakdown).forEach(([keyId, keyData]: [string, any]) => {
|
||||
const keyAlias = keyData?.metadata?.key_alias || null;
|
||||
const keyAlias = keyActivityLabel(keyData?.metadata, "") || null;
|
||||
|
||||
// Create unique key for aggregation: Date_EntityID_KeyID
|
||||
const uniqueKey = `${day.date}_${entityId}_${keyId}`;
|
||||
|
|
|
|||
|
|
@ -0,0 +1,15 @@
|
|||
import { keyActivityLabel } from "./keyActivityLabel";
|
||||
|
||||
describe("keyActivityLabel", () => {
|
||||
it("prefers key_alias", () => {
|
||||
expect(keyActivityLabel({ key_alias: "batch-worker", user_email: "alice@example.com" })).toBe("batch-worker");
|
||||
});
|
||||
|
||||
it("falls back to user_email when alias is missing", () => {
|
||||
expect(keyActivityLabel({ key_alias: null, user_email: "alice@example.com" })).toBe("alice@example.com");
|
||||
});
|
||||
|
||||
it("uses the fallback when both alias and email are missing", () => {
|
||||
expect(keyActivityLabel({ key_alias: null, user_email: null }, "key-hash-abc")).toBe("key-hash-abc");
|
||||
});
|
||||
});
|
||||
|
|
@ -0,0 +1,8 @@
|
|||
import type { KeyMetadata } from "./types";
|
||||
|
||||
export function keyActivityLabel(
|
||||
metadata: Pick<KeyMetadata, "key_alias" | "user_email"> | null | undefined,
|
||||
fallback = "-",
|
||||
): string {
|
||||
return metadata?.key_alias || metadata?.user_email || fallback;
|
||||
}
|
||||
|
|
@ -46,6 +46,7 @@ export interface KeyMetricWithMetadata {
|
|||
export interface KeyMetadata {
|
||||
key_alias: string | null;
|
||||
team_id: string | null;
|
||||
user_email?: string | null;
|
||||
tags?: { tag: string; usage: number }[];
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -101,7 +101,7 @@ const createMockDailyData = (
|
|||
});
|
||||
|
||||
const createMockKeyMetricWithMetadata = (
|
||||
metadata: { key_alias: string | null; team_id: string | null },
|
||||
metadata: { key_alias: string | null; team_id: string | null; user_email?: string | null },
|
||||
metrics: typeof EMPTY_SPEND_METRICS = EMPTY_SPEND_METRICS,
|
||||
): KeyMetricWithMetadata => ({
|
||||
metrics,
|
||||
|
|
@ -1450,6 +1450,17 @@ describe("formatKeyLabel", () => {
|
|||
expect(result).toBe("key-hash-actual-key (team: Test Team 1)");
|
||||
});
|
||||
|
||||
it("should use user_email when key_alias is null", () => {
|
||||
const modelData = createMockKeyMetricWithMetadata({
|
||||
key_alias: null,
|
||||
team_id: "team1",
|
||||
user_email: "alice@example.com",
|
||||
});
|
||||
|
||||
const result = formatKeyLabel(modelData, "actual-key", MOCK_TEAMS);
|
||||
expect(result).toBe("alice@example.com (team: Test Team 1)");
|
||||
});
|
||||
|
||||
it("should return key_alias with team_id when teams array is empty", () => {
|
||||
const modelData = createMockKeyMetricWithMetadata({
|
||||
key_alias: "my-key",
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ import { ChevronDown } from "lucide-react";
|
|||
import React, { useState } from "react";
|
||||
import { Team } from "./key_team_helpers/key_list";
|
||||
import KeyModelUsageView from "./UsagePage/components/KeyModelUsageView";
|
||||
import { keyActivityLabel } from "./UsagePage/keyActivityLabel";
|
||||
import { DailyData, KeyMetricWithMetadata, ModelActivityData, TopApiKeyData, TopModelData } from "./UsagePage/types";
|
||||
import { valueFormatter } from "./UsagePage/utils/value_formatters";
|
||||
|
||||
|
|
@ -433,7 +434,7 @@ export const ActivityMetrics: React.FC<ActivityMetricsProps> = ({ modelMetrics,
|
|||
|
||||
// Helper function to format key label
|
||||
export const formatKeyLabel = (modelData: KeyMetricWithMetadata, model: string, teams: Team[]): string => {
|
||||
const keyAlias = modelData.metadata.key_alias || `key-hash-${model}`;
|
||||
const keyAlias = keyActivityLabel(modelData.metadata, `key-hash-${model}`);
|
||||
const teamId = modelData.metadata.team_id;
|
||||
if (teamId) {
|
||||
const teamAlias = resolveTeamAliasFromTeamID(teamId, teams);
|
||||
|
|
@ -516,7 +517,7 @@ export const processActivityData = (
|
|||
if (!apiKeyBreakdown[apiKey]) {
|
||||
apiKeyBreakdown[apiKey] = {
|
||||
api_key: apiKey,
|
||||
key_alias: keyData.metadata.key_alias,
|
||||
key_alias: keyActivityLabel(keyData.metadata, "") || null,
|
||||
team_id: keyData.metadata.team_id,
|
||||
spend: 0,
|
||||
requests: 0,
|
||||
|
|
|
|||
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
2
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -27617,6 +27617,8 @@ export interface components {
|
|||
key_alias?: string | null;
|
||||
/** Team Id */
|
||||
team_id?: string | null;
|
||||
/** User Email */
|
||||
user_email?: string | null;
|
||||
};
|
||||
/**
|
||||
* KeyMetricWithMetadata
|
||||
|
|
|
|||
2
uv.lock
generated
2
uv.lock
generated
|
|
@ -4266,7 +4266,7 @@ wheels = [
|
|||
|
||||
[[package]]
|
||||
name = "litellm"
|
||||
version = "1.100.0"
|
||||
version = "1.100.1"
|
||||
source = { editable = "." }
|
||||
dependencies = [
|
||||
{ name = "aiohttp" },
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue