diff --git a/litellm/proxy/management_endpoints/daily_activity_routes.py b/litellm/proxy/management_endpoints/daily_activity_routes.py index 1a3445aae52..1abda1aae33 100644 --- a/litellm/proxy/management_endpoints/daily_activity_routes.py +++ b/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 diff --git a/litellm/repositories/daily_activity_sql.py b/litellm/repositories/daily_activity_sql.py index 4bb52767de6..5393fc42100 100644 --- a/litellm/repositories/daily_activity_sql.py +++ b/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 diff --git a/litellm/types/repositories/daily_activity.py b/litellm/types/repositories/daily_activity.py index b013ce43113..7baf1600c03 100644 --- a/litellm/types/repositories/daily_activity.py +++ b/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) diff --git a/tests/e2e/coverage_registry/quota_management.yaml b/tests/e2e/coverage_registry/quota_management.yaml index cbaf46130d6..75dd883742f 100644 --- a/tests/e2e/coverage_registry/quota_management.yaml +++ b/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"} diff --git a/tests/e2e/models.py b/tests/e2e/models.py index f76e5cec8e1..4ae54de9df5 100644 --- a/tests/e2e/models.py +++ b/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): diff --git a/tests/e2e/quota_management/spend_tracking/test_usage_export_metadata_e2e.py b/tests/e2e/quota_management/spend_tracking/test_usage_export_metadata_e2e.py new file mode 100644 index 00000000000..e70eba87689 --- /dev/null +++ b/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}" + ) diff --git a/tests/integration/spend/_daily_activity_fixtures.py b/tests/integration/spend/_daily_activity_fixtures.py index 8b4f35365ef..15574d88e82 100644 --- a/tests/integration/spend/_daily_activity_fixtures.py +++ b/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( diff --git a/tests/integration/spend/test_daily_activity_repository.py b/tests/integration/spend/test_daily_activity_repository.py index 1ed85d4a3bc..4c16c5c951b 100644 --- a/tests/integration/spend/test_daily_activity_repository.py +++ b/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), diff --git a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py b/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py index e487c00f25d..f39dedb3dcb 100644 --- a/tests/unit/proxy/management_endpoints/test_daily_activity_routes.py +++ b/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( diff --git a/tests/unit/repositories/test_daily_activity_sql.py b/tests/unit/repositories/test_daily_activity_sql.py index f8375b93dad..aed120c85c5 100644 --- a/tests/unit/repositories/test_daily_activity_sql.py +++ b/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(