feat(usage-export): include key and user metadata objects in daily activity export rows

litellm/proxy/management_endpoints/daily_activity_routes.py

@@ -221,6 +221,8 @@ def _search_rows(keys: Sequence[str], response: SpendAnalyticsPaginatedResponse)
 
 
 def _csv_cell(value: object) -> object:
+    if isinstance(value, Mapping):
+        return json.dumps(value, separators=(",", ":"))
     if isinstance(value, str) and value.startswith(("=", "+", "-", "@", "\t", "\r")):
         return f"'{value}"
     return value

litellm/repositories/daily_activity_sql.py

@@ -520,12 +520,23 @@ def build_export_sql(
     )
 
 
+_EMPTY_METADATA_FIELDS: Final = ", '{}'::jsonb AS key_metadata, '{}'::jsonb AS user_metadata"
+_USER_METADATA_FIELD: Final = (
+    "COALESCE(MIN(CASE WHEN jsonb_typeof(u.metadata) = 'object' THEN u.metadata::text END), '{}')::jsonb "
+    "AS user_metadata"
+)
+_KEY_METADATA_FIELD: Final = (
+    "COALESCE(MIN(CASE WHEN jsonb_typeof(COALESCE(vt.metadata, dvt.metadata)) = 'object' "
+    "THEN COALESCE(vt.metadata, dvt.metadata)::text END), '{}')::jsonb AS key_metadata"
+)
+
+
 def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str, ...]]:
     if export_type is ExportType.DAILY:
         return (
             "''",
             "NULL::text",
-            "NULL::text AS key_alias, NULL::text AS user_id, NULL::text AS user_email",
+            "NULL::text AS key_alias, NULL::text AS user_id, NULL::text AS user_email" + _EMPTY_METADATA_FIELDS,
             (),
         )
     if export_type is ExportType.DAILY_WITH_KEYS:
@@ -534,11 +545,14 @@ def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str,
             _API_KEY_EXPR,
             "MAX(CASE WHEN scoped.api_key = '' THEN 'Unattributed' "
             "ELSE COALESCE(vt.key_alias, dvt.key_alias) END) AS key_alias, "
-            "MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, MAX(u.user_email) AS user_email",
+            "MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, MAX(u.user_email) AS user_email, "
+            + _KEY_METADATA_FIELD
+            + ", "
+            + _USER_METADATA_FIELD,
             (
                 "LEFT JOIN \"LiteLLM_VerificationToken\" vt ON vt.token = scoped.api_key AND scoped.api_key <> ''",
                 """LEFT JOIN LATERAL (
-                    SELECT key_alias, user_id
+                    SELECT key_alias, user_id, metadata
                     FROM "LiteLLM_DeletedVerificationToken"
                     WHERE token = scoped.api_key
                     ORDER BY deleted_at DESC
@@ -551,7 +565,7 @@ def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str,
         return (
             _MODEL_EXPR,
             "NULL::text",
-            "NULL::text AS key_alias, NULL::text AS user_id, NULL::text AS user_email",
+            "NULL::text AS key_alias, NULL::text AS user_id, NULL::text AS user_email" + _EMPTY_METADATA_FIELDS,
             (),
         )
     if export_type is ExportType.DAILY_WITH_USERS:
@@ -559,11 +573,11 @@ def _export_grouping(export_type: ExportType) -> tuple[str, str, str, tuple[str,
             "COALESCE(vt.user_id, dvt.user_id, '')",
             "NULL::text",
             "NULL::text AS key_alias, MAX(COALESCE(vt.user_id, dvt.user_id)) AS user_id, "
-            "MAX(u.user_email) AS user_email",
+            "MAX(u.user_email) AS user_email, '{}'::jsonb AS key_metadata, " + _USER_METADATA_FIELD,
             (
                 'LEFT JOIN "LiteLLM_VerificationToken" vt ON vt.token = scoped.api_key',
                 """LEFT JOIN LATERAL (
-                    SELECT key_alias, user_id
+                    SELECT key_alias, user_id, metadata
                     FROM "LiteLLM_DeletedVerificationToken"
                     WHERE token = scoped.api_key
                     ORDER BY deleted_at DESC

litellm/types/repositories/daily_activity.py

@@ -126,6 +126,8 @@ class ExportRow:
     failed_requests: int
     cache_read_input_tokens: int
     cache_creation_input_tokens: int
+    key_metadata: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({}))
+    user_metadata: Mapping[str, object] = field(default_factory=lambda: MappingProxyType({}))
 
 
 @dataclass(frozen=True, slots=True)

tests/e2e/coverage_registry/quota_management.yaml

@@ -40,6 +40,7 @@
 - {id: quota_management.budget.fallback.routes_to_fallback, module: quota_management, tier: P1, behavior: budget, variant: fallback, assertions: [routes_to_fallback], exercised_on: [messages], source: "proxy/hooks/model_max_budget_limiter.py", rationale: "budget_fallbacks reroute to the fallback model once the primary's budget is exhausted"}
 - {id: quota_management.budget.spend_counter.reseed_matches_db, module: quota_management, tier: P2, behavior: budget, variant: spend_counter, assertions: [reseed_matches_db], exercised_on: [chat_completions], source: "proxy/spend_tracking/budget_reservation.py", rationale: "Concurrent cold-counter reseeds keep the enforcement counter equal to DB spend (#26829)"}
 - {id: quota_management.spend_tracking.chat_completions.logs_cost, module: quota_management, tier: P0, behavior: spend_tracking, variant: chat_completions, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "A paid chat call writes a nonzero spend row"}
+- {id: quota_management.spend_tracking.usage_export.carries_key_and_user_metadata, module: quota_management, tier: P1, behavior: spend_tracking, variant: usage_export, assertions: [carries_key_and_user_metadata], exercised_on: [chat_completions], source: "repositories/daily_activity_sql.py / LIT-8524", rationale: "The daily_with_keys usage export row carries the key's and the owning user's stored metadata objects"}
 - {id: quota_management.spend_tracking.stream.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: stream, assertions: [logs_cost], exercised_on: [chat_completions], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Streaming responses aggregate token counts into a spend row"}
 - {id: quota_management.spend_tracking.messages_bridge.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: messages_bridge, assertions: [logs_cost], exercised_on: [messages], source: "llms/anthropic/messages/responses_adapters/handler.py", rationale: "A streaming /v1/messages request served by an openai-provider model is bridged through the anthropic-messages -> Responses adapter and must aggregate the consumed SSE stream into one spend row with nonzero cost and token counts, attributed to custom_llm_provider openai under call_type anthropic_messages"}
 - {id: quota_management.spend_tracking.embeddings.logs_cost, module: quota_management, tier: P1, behavior: spend_tracking, variant: embeddings, assertions: [logs_cost], exercised_on: [embeddings], source: "proxy/spend_tracking/spend_tracking_utils.py", rationale: "Embedding calls write nonzero spend rows"}

tests/e2e/models.py

@@ -1677,6 +1677,7 @@ class UserNewBody(BaseModel):
     user_id: str | None = None
     auto_create_key: bool | None = None
     max_budget: float | None = None
+    metadata: dict[str, str] | None = None
 
 
 class UserNewResponse(BaseModel):

tests/e2e/quota_management/spend_tracking/test_usage_export_metadata_e2e.py

@@ -0,0 +1,152 @@
+"""The usage export carries the stored key and user metadata objects.
+
+A customer charging usage back by cost center keeps that cost center in the
+virtual key's metadata and the owner's user metadata. Before LIT-8524 the
+Usage page export only carried the key alias, user id and email, so every
+exported row needed a /key/info and /user/info round trip to recover the
+metadata. The export now joins both objects onto each daily_with_keys row.
+"""
+
+from __future__ import annotations
+
+from datetime import datetime, timedelta, timezone
+from typing import Final
+
+import pytest
+from e2e_config import provider_edge_base, unique_marker
+from e2e_metadata import Domain, Mode, Provider, Subject, meta
+from lifecycle import ResourceManager
+from models import (
+    ChatBody,
+    ChatMessage,
+    KeyGenerateBody,
+    KeyMetadata,
+    LiteLLMParamsBody,
+    UserNewBody,
+    UserNewResponse,
+)
+from pydantic import BaseModel, TypeAdapter
+from spend_e2e_client import SpendClient, unwrap
+from spend_reconciliation import BACKEND, INPUT_RATE, OUTPUT_RATE
+
+pytestmark = pytest.mark.e2e
+
+EXPORT_ROUTE: Final = "/user/daily/activity/export"
+
+
+class UsageExportParams(BaseModel):
+    start_date: str
+    end_date: str
+    api_key: str
+    export_type: str = "daily_with_keys"
+    format: str = "json"
+
+
+class UsageExportKeyMetadata(BaseModel):
+    tag: str | None = None
+
+
+class UsageExportUserMetadata(BaseModel):
+    department: str | None = None
+
+
+class UsageExportRow(BaseModel):
+    api_key: str | None = None
+    user_id: str | None = None
+    key_metadata: UsageExportKeyMetadata
+    user_metadata: UsageExportUserMetadata
+
+
+_EXPORT_ROWS: Final = TypeAdapter(list[UsageExportRow])
+
+
+class TestUsageExportMetadata:
+    @pytest.mark.replayable
+    @pytest.mark.covers("quota_management.spend_tracking.usage_export.carries_key_and_user_metadata")
+    @meta(
+        Subject(
+            domain=Domain.SPEND_BUDGETS,
+            providers=(Provider.OPENAI,),
+            models=(BACKEND,),
+            mode=Mode.NONSTREAM,
+        )
+    )
+    def test_daily_with_keys_export_row_carries_key_and_user_metadata(
+        self, client: SpendClient, resources: ResourceManager
+    ) -> None:
+        started: Final = datetime.now(timezone.utc)
+        marker: Final = unique_marker()
+        base: Final = provider_edge_base("openai")
+        model: Final = f"e2e-usage-export-{marker}"
+        model_id: Final = client.proxy.create_model(
+            model,
+            LiteLLMParamsBody(
+                model=BACKEND,
+                api_key="os.environ/OPENAI_API_KEY",
+                api_base=None if base is None else f"{base}/v1",
+                input_cost_per_token=INPUT_RATE,
+                output_cost_per_token=OUTPUT_RATE,
+            ),
+        )
+        resources.defer(lambda: client.proxy.delete_model(model_id))
+        user_id: Final = unwrap(
+            client.proxy.transport.post(
+                "/user/new",
+                headers=client.proxy.transport.master,
+                json=UserNewBody(
+                    user_email=f"e2e-usage-export-{marker}@example.com",
+                    user_role="internal_user",
+                    user_id=f"e2e-usage-export-{marker}",
+                    metadata={"department": f"research-{marker}"},
+                ),
+                response_type=UserNewResponse,
+            )
+        ).user_id
+        resources.defer(lambda: client.delete_user(user_id))
+        record: Final = client.generate_key_record(
+            KeyGenerateBody(
+                user_id=user_id,
+                models=[model],
+                key_alias=f"e2e-usage-export-{marker}",
+                metadata=KeyMetadata(tag=f"cost-center-{marker}"),
+            )
+        )
+        resources.defer(lambda: client.proxy.delete_key(record.key))
+        assert record.token, "/key/generate answered without the key's token hash"
+        token: Final = record.token
+
+        response: Final = unwrap(
+            client.proxy.chat(
+                record.key,
+                ChatBody(
+                    model=model,
+                    messages=[ChatMessage(role="user", content=f"Reply with one word. {marker}")],
+                    max_completion_tokens=32,
+                ),
+            )
+        )
+        assert response.id, "successful response must have an ID"
+        end: Final = datetime.now(timezone.utc) + timedelta(days=1)
+        activity: Final = client.poll_daily_activity_for_key(
+            token, start=started - timedelta(days=1), end=end, min_requests=1
+        )
+        assert activity is not None, f"/user/daily/activity never listed key {token} under api_keys"
+
+        exported: Final = client.proxy.transport.probe(
+            EXPORT_ROUTE,
+            params=UsageExportParams(
+                start_date=(started - timedelta(days=1)).strftime("%Y-%m-%d"),
+                end_date=end.strftime("%Y-%m-%d"),
+                api_key=token,
+            ),
+        )
+        assert exported.status_code == 200, f"{EXPORT_ROUTE} answered {exported.status_code}: {exported.body}"
+        rows: Final = _EXPORT_ROWS.validate_json(exported.body)
+        key_rows: Final = tuple(row for row in rows if row.api_key == token)
+        assert key_rows, f"{EXPORT_ROUTE} returned no daily_with_keys row for {token}: {exported.body}"
+        for row in key_rows:
+            assert row.user_id == user_id, f"export row for {token} is owned by {row.user_id}, not {user_id}"
+            assert row.key_metadata.tag == f"cost-center-{marker}", f"key metadata not exported: {row.key_metadata}"
+            assert row.user_metadata.department == f"research-{marker}", (
+                f"user metadata not exported: {row.user_metadata}"
+            )

tests/integration/spend/_daily_activity_fixtures.py

@@ -198,8 +198,10 @@ def seed_daily_activity_fixture(connection: psycopg.Connection, *, schema: str,
             ),
         )
         cursor.execute(
-            sql.SQL("INSERT INTO {} (user_id, user_email, models) VALUES (%s, %s, %s)").format(user_table),
-            ("user-1", "user@example.com", []),
+            sql.SQL("INSERT INTO {} (user_id, user_email, models, metadata) VALUES (%s, %s, %s, %s)").format(
+                user_table
+            ),
+            ("user-1", "user@example.com", [], '{"department": "research"}'),
         )
         cursor.execute(
             sql.SQL("INSERT INTO {} (team_id, team_alias, admins, members, models) VALUES (%s, %s, %s, %s, %s)").format(

tests/integration/spend/test_daily_activity_repository.py

@@ -236,6 +236,11 @@ async def test_repository_queries_and_exports_seeded_daily_activity(monkeypatch:
             "user-1",
             "user@example.com",
         )
+        assert deleted_key_export.key_metadata == {"tags": ["archived"]}
+        assert deleted_key_export.user_metadata == {"department": "research"}
+        active_key_export: Final = next(row for row in exports if row.api_key == "key-a")
+        assert active_key_export.key_metadata == {"tags": ["blue", "gold"]}
+        assert active_key_export.user_metadata == {"department": "research"}
         user_exports: Final = tuple(
             [row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY_WITH_USERS)]
         )
@@ -245,11 +250,14 @@ async def test_repository_queries_and_exports_seeded_daily_activity(monkeypatch:
             "user@example.com",
             273.0,
         )
+        assert user_exports[0].user_metadata == {"department": "research"}
+        assert user_exports[0].key_metadata == {}
         daily_export: Final = tuple(
             [row async for row in repository.export_rows(team_scope, export_type=ExportType.DAILY)]
         )
         assert len(daily_export) == 1
         assert daily_export[0].spend == 1273.0
+        assert (daily_export[0].key_metadata, daily_export[0].user_metadata) == ({}, {})
         assert daily_export[0].flat_cost == 42.0
 
         metadata: Final = await repository.key_metadata(frozenset(("key-a", "key-target")), None)
@@ -684,9 +692,7 @@ async def test_user_page_ranks_users_and_folds_null_ids() -> None:
         repository: Final = _repository(database)
         full_page: Final = await repository.user_page(scope, offset=0, limit=50)
         empty_page: Final = await repository.user_page(scope, offset=full_page.total_users + 1, limit=50)
-        scoped_page: Final = await repository.user_page(
-            replace(scope, entity_ids=("user-alpha",)), offset=0, limit=50
-        )
+        scoped_page: Final = await repository.user_page(replace(scope, entity_ids=("user-alpha",)), offset=0, limit=50)
 
         assert tuple((row.user_id, row.spend) for row in full_page.rows) == (
             (None, 45.0),

tests/unit/proxy/management_endpoints/test_daily_activity_routes.py

@@ -1,5 +1,6 @@
 import csv
 import io
+import json
 from collections.abc import AsyncIterator, Iterator, Mapping, Sequence
 from dataclasses import dataclass, fields
 from itertools import chain
@@ -390,6 +391,8 @@ class _FakeRepository:
                 failed_requests=row.failed_requests,
                 cache_read_input_tokens=row.cache_read_input_tokens,
                 cache_creation_input_tokens=row.cache_creation_input_tokens,
+                key_metadata={"cost_center": "=finance", "tier": 2},
+                user_metadata={"department": "research"},
             )
 
 
@@ -1160,6 +1163,9 @@ def test_export_routes_stream_csv_and_preserve_row_counts(
     assert len(records) == 7
     assert records[1][2] == "'=entity"
     assert records[1][4] == "'+key"
+    columns: Final = dict(zip(records[0], records[1], strict=True))
+    assert json.loads(columns["key_metadata"]) == {"cost_center": "=finance", "tier": 2}
+    assert json.loads(columns["user_metadata"]) == {"department": "research"}
 
 
 @pytest.mark.parametrize(("prefix", "query_name", "entity_id"), _ENTITY_CASES)
@@ -1180,6 +1186,8 @@ def test_export_routes_stream_json_arrays(
     records: Final = response.json()
     assert isinstance(records, list) and len(records) == 6, response.text
     assert records[0]["entity_alias"] == "=entity"
+    assert records[0]["key_metadata"] == {"cost_center": "=finance", "tier": 2}
+    assert records[0]["user_metadata"] == {"department": "research"}
 
 
 def test_export_first_row_error_returns_json_error_before_streaming(

tests/unit/repositories/test_daily_activity_sql.py

@@ -451,6 +451,34 @@ def test_export_uses_latest_deleted_key_metadata(export_type: ExportType) -> Non
     assert "COALESCE(vt.user_id, dvt.user_id)" in query.sql
 
 
+@pytest.mark.parametrize(
+    ("export_type", "key_metadata", "user_metadata"),
+    [
+        (
+            ExportType.DAILY_WITH_KEYS,
+            "THEN COALESCE(vt.metadata, dvt.metadata)::text END), '{}')::jsonb AS key_metadata",
+            "THEN u.metadata::text END), '{}')::jsonb AS user_metadata",
+        ),
+        (
+            ExportType.DAILY_WITH_USERS,
+            "'{}'::jsonb AS key_metadata",
+            "THEN u.metadata::text END), '{}')::jsonb AS user_metadata",
+        ),
+        (ExportType.DAILY, "'{}'::jsonb AS key_metadata", "'{}'::jsonb AS user_metadata"),
+        (ExportType.DAILY_WITH_MODELS, "'{}'::jsonb AS key_metadata", "'{}'::jsonb AS user_metadata"),
+    ],
+)
+def test_export_selects_stored_key_and_user_metadata_objects(
+    export_type: ExportType, key_metadata: str, user_metadata: str
+) -> None:
+    query = build_export_sql(_scope(entity_ids=None), export_type=export_type, after=None, batch_size=2)
+
+    assert key_metadata in query.sql
+    assert user_metadata in query.sql
+    if export_type is ExportType.DAILY_WITH_KEYS:
+        assert "SELECT key_alias, user_id, metadata" in query.sql
+
+
 @pytest.mark.parametrize("export_type", tuple(ExportType))
 def test_export_without_cursor_omits_cursor_predicate_and_parameters(export_type: ExportType) -> None:
     query = build_export_sql(