mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-07 08:26:10 +00:00
fix(usage): include user_email on daily activity key breakdowns
Co-authored-by: Mateo Wang <mateo-berri@users.noreply.github.com>
This commit is contained in:
Cursor Agent
parent
eeb8b4f7f3
commit
e109d89c20
2 changed files with 81 additions and 39 deletions
|
|
@ -115,12 +115,21 @@ 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]
|
||||
|
||||
|
||||
|
|
@ -289,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,
|
||||
|
|
@ -316,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,
|
||||
|
|
@ -341,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),
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -369,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,
|
||||
|
|
@ -394,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,
|
||||
|
|
@ -409,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)
|
||||
|
||||
|
|
@ -432,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,
|
||||
|
|
@ -970,15 +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"),
|
||||
user_email=meta.get("user_email"),
|
||||
)
|
||||
|
||||
|
||||
def _aggregate_grouping_sets_records_sync(
|
||||
*,
|
||||
records: Sequence[_GroupingSetsRow],
|
||||
|
|
|
|||
|
|
@ -546,6 +546,69 @@ def test_key_metadata_includes_recovered_user_email():
|
|||
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.
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue