diff --git a/litellm/proxy/management_endpoints/common_daily_activity.py b/litellm/proxy/management_endpoints/common_daily_activity.py index 833f463621a..ce6a97708ab 100644 --- a/litellm/proxy/management_endpoints/common_daily_activity.py +++ b/litellm/proxy/management_endpoints/common_daily_activity.py @@ -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], diff --git a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py index f7ad01c3bba..b9c0d953086 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py +++ b/tests/test_litellm/proxy/management_endpoints/test_common_daily_activity.py @@ -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.