diff --git a/.github/pr-assets/customer-usage-detail/before-main-6685525.png b/.github/pr-assets/customer-usage-detail/before-main-6685525.png new file mode 100644 index 00000000000..563c231cfbe Binary files /dev/null and b/.github/pr-assets/customer-usage-detail/before-main-6685525.png differ diff --git a/.github/pr-assets/customer-usage-detail/final-overview.png b/.github/pr-assets/customer-usage-detail/final-overview.png new file mode 100644 index 00000000000..fb9438653a5 Binary files /dev/null and b/.github/pr-assets/customer-usage-detail/final-overview.png differ diff --git a/.github/pr-assets/customer-usage-detail/final-report.png b/.github/pr-assets/customer-usage-detail/final-report.png new file mode 100644 index 00000000000..7118992dfdc Binary files /dev/null and b/.github/pr-assets/customer-usage-detail/final-report.png differ diff --git a/litellm/accounting/application.py b/litellm/accounting/application.py index c911f4e35ce..58c060b81c6 100644 --- a/litellm/accounting/application.py +++ b/litellm/accounting/application.py @@ -1,3 +1,5 @@ +from hashlib import sha256 +from json import dumps from typing import TYPE_CHECKING, Final from litellm.accounting.delivery_projection import entity_transactions, project_deliveries @@ -30,6 +32,7 @@ async def apply_deliveries( await writer.apply_accounting_entities( prisma_client, proxy_logging_obj, transaction, entity_transactions(contributions) ) + await _apply_customer_usage(transaction, contributions) for table in sorted(DAILY_SPEND_TABLES.values(), key=lambda item: item.name): rows = tuple( {**dict(item.key), **dict(item.values), **dict(item.increments)} @@ -60,6 +63,45 @@ async def apply_deliveries( return contributions +async def _apply_customer_usage(transaction: LedgerTransaction, contributions: tuple[Contribution, ...]) -> None: + for item in contributions: + if item.table == "LiteLLM_DailyCustomerMemberSpend": + await _upsert_customer_daily(transaction, item) + if item.table == "LiteLLM_CustomerOrganizationTable": + await transaction.litellm_customerorganizationtable.update_many( + where=dict(item.key), data={"spend": {"increment": dict(item.increments)["spend"]}} + ) + if item.table == "LiteLLM_CustomerTeamTable": + await transaction.litellm_customerteamtable.update_many( + where=dict(item.key), data={"spend": {"increment": dict(item.increments)["spend"]}} + ) + if item.table == "LiteLLM_CustomerMemberTable": + await transaction.litellm_customermembertable.update_many( + where=dict(item.key), data={"spend": {"increment": dict(item.increments)["spend"]}} + ) + + +async def _upsert_customer_daily(transaction: LedgerTransaction, contribution: Contribution) -> None: + row_id: Final = _customer_daily_row_id(contribution.key) + await transaction.litellm_dailycustomermemberspend.upsert( + where={"id": row_id}, + data={ + "create": { + "id": row_id, + **dict(contribution.key), + **dict(contribution.values), + **dict(contribution.increments), + }, + "update": {name: {"increment": value} for name, value in contribution.increments}, + }, + ) + + +def _customer_daily_row_id(key: tuple[tuple[str, str], ...]) -> str: + encoded: Final = dumps(sorted(key), separators=(",", ":")).encode("utf-8") + return sha256(encoded).hexdigest() + + def _window(item: Contribution) -> WindowSpendTransaction: key: Final = dict(item.key) values: Final = dict(item.values) diff --git a/litellm/accounting/projection.py b/litellm/accounting/projection.py index d1e461efca9..2c11372d7cb 100644 --- a/litellm/accounting/projection.py +++ b/litellm/accounting/projection.py @@ -14,6 +14,7 @@ from litellm.constants import ( MODEL_INSIGHTS_TASK_TAG_PREFIX, ) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import GuardrailCostByUnitEntry +from litellm.proxy.common_utils.customer_usage import CUSTOMER_USAGE_METADATA_KEY, CustomerUsageAttribution from litellm.proxy.db.daily_spend_bulk_upsert import DAILY_SPEND_TABLES, daily_spend_entity_ids from litellm.proxy.db.model_insights_tasks import load_model_insight_tasks from litellm.proxy.db.model_usage_rollup import date_from_start_time @@ -202,10 +203,53 @@ def project(row: Mapping[str, object]) -> tuple[Contribution, ...]: *_window_contributions(row, metadata), *_model_contributions(row, metadata), *_guardrail_contributions(row, metadata), + *_customer_usage_contributions(row, metadata), ) return tuple(_capped(contribution) for contribution in contributions if _indexable(contribution)) +def _customer_usage_contributions( + row: Mapping[str, object], metadata: Mapping[str, object] +) -> tuple[Contribution, ...]: + serialized: Final = metadata.get(CUSTOMER_USAGE_METADATA_KEY) + date: Final = _date(row) + if not isinstance(serialized, str) or date is None: + return () + try: + attribution: Final = CustomerUsageAttribution.model_validate_json(serialized) + except ValidationError: + return () + counters: Final = _daily_counters(row, metadata) + identities: Final = ( + ("LiteLLM_CustomerOrganizationTable", "customer_org_id", attribution.customer_org_id), + ("LiteLLM_CustomerTeamTable", "customer_team_id", attribution.customer_team_id), + ("LiteLLM_CustomerMemberTable", "customer_member_id", attribution.customer_member_id), + ) + return ( + Contribution( + "LiteLLM_DailyCustomerMemberSpend", + ( + ("date", date), + ("customer_org_id", attribution.customer_org_id), + ("customer_team_id", attribution.customer_team_id or ""), + ("customer_member_id", attribution.customer_member_id or ""), + ("model", _text(row.get("model"))), + ("model_group", _text(row.get("model_group"))), + ("custom_llm_provider", _text(row.get("custom_llm_provider"))), + ("api_key", _text(row.get("api_key"))), + ("mcp_namespaced_tool_name", _text(row.get("mcp_namespaced_tool_name"))), + ("endpoint", _text(ROUTE_ENDPOINT_MAPPING.get(_text(row.get("call_type"))))), + ), + counters, + ), + *( + _spend(table, ((column, identity),), float(_number(row.get("spend")))) + for table, column, identity in identities + if identity is not None + ), + ) + + def _indexable(contribution: Contribution) -> bool: texts: Final = (value for _, value in (*contribution.key, *contribution.values) if value is not None) return all(len(text.encode("utf-8", "surrogatepass")) <= _MAX_INDEXED_TEXT_BYTES for text in texts) diff --git a/litellm/accounting/snapshot.py b/litellm/accounting/snapshot.py index 874b90110c1..b86f37c3f5e 100644 --- a/litellm/accounting/snapshot.py +++ b/litellm/accounting/snapshot.py @@ -31,6 +31,7 @@ _METADATA_FIELDS: Final = frozenset( "compression_savings", "internal_call_origin", "billed_retry_attempt", + "user_api_key_customer_usage", } ) diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index d1ed9f7b62b..2c9d64d50dd 100644 --- a/litellm/integrations/otel/model/spans.py +++ b/litellm/integrations/otel/model/spans.py @@ -308,6 +308,7 @@ _PRISMA_MODELS: Final[frozenset[str]] = frozenset( "LiteLLM_DailyUserSpend", "LiteLLM_DailyGlobalSpend", "LiteLLM_DailyOrganizationSpend", + "LiteLLM_DailyCustomerMemberSpend", "LiteLLM_DailyEndUserSpend", "LiteLLM_DailyAgentSpend", "LiteLLM_DailyTeamSpend", diff --git a/litellm/proxy-extras/litellm_proxy_extras/migrations/20261113000000_customer_daily_usage/migration.sql b/litellm/proxy-extras/litellm_proxy_extras/migrations/20261113000000_customer_daily_usage/migration.sql new file mode 100644 index 00000000000..f2746f41854 --- /dev/null +++ b/litellm/proxy-extras/litellm_proxy_extras/migrations/20261113000000_customer_daily_usage/migration.sql @@ -0,0 +1,47 @@ +SET LOCAL lock_timeout = '10s'; + +CREATE TABLE IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend" ( + "id" TEXT NOT NULL, + "date" TEXT NOT NULL, + "customer_org_id" TEXT NOT NULL, + "customer_team_id" TEXT NOT NULL DEFAULT '', + "customer_member_id" TEXT NOT NULL DEFAULT '', + "model" TEXT NOT NULL DEFAULT '', + "model_group" TEXT NOT NULL DEFAULT '', + "custom_llm_provider" TEXT NOT NULL DEFAULT '', + "mcp_namespaced_tool_name" TEXT NOT NULL DEFAULT '', + "endpoint" TEXT NOT NULL DEFAULT '', + "api_key" TEXT NOT NULL DEFAULT '', + "prompt_tokens" BIGINT NOT NULL DEFAULT 0, + "completion_tokens" BIGINT NOT NULL DEFAULT 0, + "cache_read_input_tokens" BIGINT NOT NULL DEFAULT 0, + "cache_creation_input_tokens" BIGINT NOT NULL DEFAULT 0, + "compression_saved_tokens" BIGINT NOT NULL DEFAULT 0, + "api_requests" BIGINT NOT NULL DEFAULT 0, + "successful_requests" BIGINT NOT NULL DEFAULT 0, + "failed_requests" BIGINT NOT NULL DEFAULT 0, + "total_response_time_ms" BIGINT NOT NULL DEFAULT 0, + "timed_requests" BIGINT NOT NULL DEFAULT 0, + "spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "compression_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "prompt_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "gateway_injected_caching_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "autorouter_savings_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "guardrail_spend" DOUBLE PRECISION NOT NULL DEFAULT 0.0, + "created_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + "updated_at" TIMESTAMP(3) NOT NULL DEFAULT CURRENT_TIMESTAMP, + CONSTRAINT "LiteLLM_DailyCustomerMemberSpend_pkey" PRIMARY KEY ("id") +); + +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_date_idx" + ON "LiteLLM_DailyCustomerMemberSpend"("date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_customer_org_id_date_idx" + ON "LiteLLM_DailyCustomerMemberSpend"("customer_org_id", "date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_customer_team_id_date_idx" + ON "LiteLLM_DailyCustomerMemberSpend"("customer_team_id", "date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_customer_member_id_date_idx" + ON "LiteLLM_DailyCustomerMemberSpend"("customer_member_id", "date"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_api_key_idx" + ON "LiteLLM_DailyCustomerMemberSpend"("api_key"); +CREATE INDEX IF NOT EXISTS "LiteLLM_DailyCustomerMemberSpend_model_idx" + ON "LiteLLM_DailyCustomerMemberSpend"("model"); diff --git a/litellm/proxy-extras/litellm_proxy_extras/schema.prisma b/litellm/proxy-extras/litellm_proxy_extras/schema.prisma index 2c0246bb94e..bac27daa5de 100644 --- a/litellm/proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm/proxy-extras/litellm_proxy_extras/schema.prisma @@ -1035,6 +1035,45 @@ model LiteLLM_DailyOrganizationSpend { } // Track daily end user (customer) spend metrics per model and key +model LiteLLM_DailyCustomerMemberSpend { + id String @id + date String + customer_org_id String + customer_team_id String @default("") + customer_member_id String @default("") + model String @default("") + model_group String @default("") + custom_llm_provider String @default("") + mcp_namespaced_tool_name String @default("") + endpoint String @default("") + api_key String @default("") + spend Float @default(0.0) + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + compression_saved_tokens BigInt @default(0) + compression_savings_spend Float @default(0.0) + prompt_caching_savings_spend Float @default(0.0) + gateway_injected_caching_savings_spend Float @default(0.0) + autorouter_savings_spend Float @default(0.0) + guardrail_spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@index([date]) + @@index([customer_org_id, date]) + @@index([customer_team_id, date]) + @@index([customer_member_id, date]) + @@index([api_key]) + @@index([model]) +} + model LiteLLM_DailyEndUserSpend { id String @id @default(uuid()) end_user_id String? diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3ee3ea92d7b..7a791ff7876 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -726,12 +726,15 @@ class LiteLLMRoutes(enum.Enum): "/jwt/key/mapping/info", "/management/v1/customer_organizations", "/management/v1/customer_organizations/{customer_org_id}", + "/management/v1/customer_organizations/daily/activity", "/management/v1/customer_organizations/{customer_org_id}/default_budgets", "/management/v1/customer_organizations/{customer_org_id}/default_budgets/{budget_id}", "/management/v1/customer_teams", "/management/v1/customer_teams/{customer_team_id}", + "/management/v1/customer_teams/daily/activity", "/management/v1/customer_members", "/management/v1/customer_members/{customer_member_id}", + "/management/v1/customer_members/daily/activity", ] + key_management_routes + mcp_management_routes @@ -3855,6 +3858,7 @@ class SpendLogsRouterMetadata(TypedDict): class SpendLogsMetadata(TypedDict): + user_api_key_customer_usage: ReadOnly[NotRequired[str | None]] actor_agent_id: ReadOnly[NotRequired[str | None]] target_agent_id: ReadOnly[NotRequired[str | None]] billing_agent_id: ReadOnly[NotRequired[str | None]] diff --git a/litellm/proxy/auth/login_utils.py b/litellm/proxy/auth/login_utils.py index 7d413121b56..16b59657d80 100644 --- a/litellm/proxy/auth/login_utils.py +++ b/litellm/proxy/auth/login_utils.py @@ -354,6 +354,7 @@ async def _sign_in( response = await generate_key_helper_fn( llm_router=None, request_type="key", + actor=None, **{ "user_role": LitellmUserRoles.PROXY_ADMIN, "duration": LITELLM_UI_SESSION_DURATION, diff --git a/litellm/proxy/common_utils/customer_usage.py b/litellm/proxy/common_utils/customer_usage.py new file mode 100644 index 00000000000..f356fc0f2bf --- /dev/null +++ b/litellm/proxy/common_utils/customer_usage.py @@ -0,0 +1,187 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from enum import Enum +from typing import Annotated, Final, Protocol, TypeAlias + +from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError + +from litellm.repositories.prisma_protocols import TableActions + +CustomerIdentifier: TypeAlias = Annotated[str, Field(strict=True, min_length=1, max_length=256)] +_IDENTITY_FIELDS: Final = ("customer_org_id", "customer_team_id", "customer_member_id") +CUSTOMER_USAGE_METADATA_KEY: Final = "user_api_key_customer_usage" +_FIELDS: Final = TypeAdapter(Mapping[str, object]) + + +class CustomerUsageInput(BaseModel): + model_config = ConfigDict(frozen=True) + + customer_org_id: CustomerIdentifier | None = None + customer_team_id: CustomerIdentifier | None = None + customer_member_id: CustomerIdentifier | None = None + + +class CustomerUsageAttribution(BaseModel): + model_config = ConfigDict(frozen=True, extra="forbid") + + customer_org_id: CustomerIdentifier + customer_team_id: CustomerIdentifier | None = None + customer_member_id: CustomerIdentifier | None = None + + +class _Team(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + parent_org_id: str + + +class _Member(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + parent_team_id: str + + +class _KeyMetadata(BaseModel): + model_config = ConfigDict(frozen=True, from_attributes=True) + + metadata: object = None + + +class CustomerUsageDatabase(Protocol): + @property + def litellm_customerorganizationtable(self) -> TableActions[object]: ... + + @property + def litellm_customerteamtable(self) -> TableActions[object]: ... + + @property + def litellm_customermembertable(self) -> TableActions[object]: ... + + +class CustomerUsageError(str, Enum): + INVALID = "Customer usage metadata contains invalid identifiers." + MISMATCH = "Customer usage identifiers must belong to the same organization and team." + FORBIDDEN = "Customer usage must match the customer organization assigned to this API key." + UNBOUND = "Assign a customer organization to the API key before tracking customer usage." + NOT_FOUND = "The customer usage organization, team, or member does not exist." + UNAVAILABLE = "Customer usage tracking requires a database and spend accounting ledger." + KEY_BINDING_FORBIDDEN = "Only proxy admins can change a key's customer organization, team, or member binding." + + +@dataclass(frozen=True, slots=True) +class CustomerUsageFailure: + error: CustomerUsageError + + +def normalize_customer_key_metadata( + metadata: Mapping[str, object] | None, + existing_metadata: Mapping[str, object] | None, + *, + administrator: bool, +) -> Mapping[str, object] | CustomerUsageFailure | None: + if administrator: + return metadata + stored: Final = existing_metadata or {} + requested: Final = metadata or {} + if any(name in requested and requested[name] != stored.get(name) for name in _IDENTITY_FIELDS): + return CustomerUsageFailure(CustomerUsageError.KEY_BINDING_FORBIDDEN) + bindings: Final = {name: stored[name] for name in _IDENTITY_FIELDS if stored.get(name) is not None} + if not bindings: + return metadata + return {**requested, **bindings} + + +def _input(metadata: Mapping[str, object], binding: Mapping[str, object]) -> CustomerUsageInput | CustomerUsageFailure: + if any( + name in metadata and name in binding and metadata[name] != binding[name] + for name in _IDENTITY_FIELDS + if binding.get(name) is not None + ): + return CustomerUsageFailure(CustomerUsageError.FORBIDDEN) + try: + return CustomerUsageInput.model_validate( + {**metadata, **{name: binding[name] for name in _IDENTITY_FIELDS if binding.get(name) is not None}} + ) + except ValidationError: + return CustomerUsageFailure(CustomerUsageError.INVALID) + + +def _hierarchy_failure(error: CustomerUsageError, *, administrator: bool) -> CustomerUsageFailure: + if administrator: + return CustomerUsageFailure(error) + return CustomerUsageFailure(CustomerUsageError.FORBIDDEN) + + +async def resolve_customer_usage( + *, + metadata: Mapping[str, object], + key_metadata: Mapping[str, object], + administrator: bool, + database: CustomerUsageDatabase | None, + ledger_enabled: bool, +) -> CustomerUsageAttribution | CustomerUsageFailure | None: + if not customer_usage_requested(metadata, key_metadata): + return None + candidate: Final = _input(metadata, key_metadata) + if isinstance(candidate, CustomerUsageFailure): + return candidate + if not administrator and key_metadata.get("customer_org_id") is None: + return CustomerUsageFailure(CustomerUsageError.UNBOUND) + if database is None or not ledger_enabled: + return CustomerUsageFailure(CustomerUsageError.UNAVAILABLE) + member: Final = ( + await database.litellm_customermembertable.find_unique( + where={"customer_member_id": candidate.customer_member_id} + ) + if candidate.customer_member_id is not None + else None + ) + if candidate.customer_member_id is not None and member is None: + return _hierarchy_failure(CustomerUsageError.NOT_FOUND, administrator=administrator) + member_team: Final = _Member.model_validate(member).parent_team_id if member is not None else None + if candidate.customer_team_id is not None and member_team is not None and candidate.customer_team_id != member_team: + return _hierarchy_failure(CustomerUsageError.MISMATCH, administrator=administrator) + team_id: Final = candidate.customer_team_id or member_team + team: Final = ( + await database.litellm_customerteamtable.find_unique(where={"customer_team_id": team_id}) + if team_id is not None + else None + ) + if team_id is not None and team is None: + return _hierarchy_failure(CustomerUsageError.NOT_FOUND, administrator=administrator) + team_org: Final = _Team.model_validate(team).parent_org_id if team is not None else None + if candidate.customer_org_id is not None and team_org is not None and candidate.customer_org_id != team_org: + return _hierarchy_failure(CustomerUsageError.MISMATCH, administrator=administrator) + organization_id: Final = candidate.customer_org_id or team_org + if organization_id is None: + return CustomerUsageFailure(CustomerUsageError.INVALID) + organization: Final = await database.litellm_customerorganizationtable.find_unique( + where={"customer_org_id": organization_id} + ) + if organization is None: + return _hierarchy_failure(CustomerUsageError.NOT_FOUND, administrator=administrator) + return CustomerUsageAttribution( + customer_org_id=organization_id, + customer_team_id=team_id, + customer_member_id=candidate.customer_member_id, + ) + + +def customer_usage_requested(metadata: Mapping[str, object], key_metadata: Mapping[str, object]) -> bool: + return any(metadata.get(name) is not None or key_metadata.get(name) is not None for name in _IDENTITY_FIELDS) + + +def customer_usage_fields(value: object) -> Mapping[str, object]: + return _FIELDS.validate_python(value) if isinstance(value, Mapping) else {} + + +def customer_usage_key_fields(key: object) -> Mapping[str, object]: + return customer_usage_fields(_KeyMetadata.model_validate(key).metadata) + + +def customer_usage_http_status(error: CustomerUsageError) -> int: + if error == CustomerUsageError.UNAVAILABLE: + return 503 + if error in (CustomerUsageError.UNBOUND, CustomerUsageError.FORBIDDEN, CustomerUsageError.KEY_BINDING_FORBIDDEN): + return 403 + return 400 diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 8abc91a72ed..5ba324e7154 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -59,6 +59,7 @@ from litellm.litellm_core_utils.core_helpers import ( iter_items, keyword_arguments, normalize_drop_params, + require_attributes, ) from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -93,6 +94,16 @@ from litellm.proxy.common_utils.callback_utils import ( key_logging_entries, strip_callback_config, ) +from litellm.proxy.common_utils.customer_usage import ( + CUSTOMER_USAGE_METADATA_KEY, + CustomerUsageDatabase, + CustomerUsageFailure, + customer_usage_fields, + customer_usage_http_status, + customer_usage_key_fields, + customer_usage_requested, + resolve_customer_usage, +) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export detach_request_parsed_body, @@ -2712,7 +2723,7 @@ async def add_litellm_data_to_request( """ - from litellm.proxy.proxy_server import llm_router, premium_user + from litellm.proxy.proxy_server import llm_router, premium_user, prisma_client from litellm.types.proxy.litellm_pre_call_utils import RedactedDict, SecretFields # Strip internal-only keys from user input before the proxy sets its own. @@ -2989,6 +3000,33 @@ async def add_litellm_data_to_request( if _metadata_variable_name == "metadata": data.pop("litellm_metadata", None) + _customer_metadata: Final = customer_usage_fields(_REQUEST_FIELDS.validate_python(data)[_metadata_variable_name]) + _customer_key_metadata: Final = customer_usage_key_fields(user_api_key_dict) + _customer_usage: Final = await resolve_customer_usage( + metadata=_customer_metadata, + key_metadata=_customer_key_metadata, + administrator=user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN, + database=require_attributes( + prisma_client.db, + CustomerUsageDatabase, + "litellm_customerorganizationtable", + "litellm_customerteamtable", + "litellm_customermembertable", + ) + if prisma_client is not None and customer_usage_requested(_customer_metadata, _customer_key_metadata) + else None, + ledger_enabled=general_settings_view is None + or general_settings_view.get("use_spend_logs_ledger", True) is not False, + ) + if isinstance(_customer_usage, CustomerUsageFailure): + raise HTTPException( + status_code=customer_usage_http_status(_customer_usage.error), detail=_customer_usage.error.value + ) + if _customer_usage is not None: + data[_metadata_variable_name][CUSTOMER_USAGE_METADATA_KEY] = ( # rebind-ok: data is an out-param + _customer_usage.model_dump_json() + ) + data = LiteLLMProxyRequestSetup.add_user_api_key_auth_to_request_metadata( data=data, user_api_key_dict=user_api_key_dict, diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index f973fadb53c..a2c32cdb008 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -79,9 +79,11 @@ from litellm.proxy.management_endpoints.common_utils import ( # noqa: F401 # l from litellm.proxy.management_endpoints.key_management_endpoints import ( # noqa: F401 # legacy module exports _check_permissions_caller_permission, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export check_permissions_caller_permission, + customer_key_actor, generate_key_helper_fn, prepare_metadata_fields, raise_on_logging_credential_change, + validated_customer_key_metadata, ) from litellm.proxy.management_helpers.bulk_user_deletion import delete_users_from_current_teams from litellm.proxy.management_helpers.object_permission_utils import ( # noqa: F401 # legacy module exports @@ -560,17 +562,15 @@ async def new_user( }' ``` """ + actor: Final = customer_key_actor(user_api_key_dict) + if data.auto_create_key and actor is not None: + validated_customer_key_metadata(data.metadata, None, actor) try: from litellm.proxy.proxy_server import _license_check, prisma_client if prisma_client is None: raise HTTPException(status_code=400, detail=CommonProxyErrors.db_not_connected_error.value) - if prisma_client is None: - raise HTTPException( - status_code=500, - detail=CommonProxyErrors.db_not_connected_error.value, - ) validate_budget_duration(data.budget_duration) # Check for duplicate user_id or email @@ -629,7 +629,12 @@ async def new_user( teams = check_if_default_team_set() organization_ids: Final = cast(list[str] | None, data_json.pop("organizations", None)) - response: Final = await generate_key_helper_fn(request_type="user", **data_json, llm_router=None) + response: Final = await generate_key_helper_fn( + request_type="user", + **data_json, + llm_router=None, + actor=actor, + ) created_user_id: Final = cast(str | None, response.get("user_id", None)) if created_user_id is not None: forget_missing_user(created_user_id) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 3bd17c4b5d8..b6f9dd0ef55 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -82,6 +82,12 @@ from litellm.proxy.common_utils.config_sync_pubsub import ( coordination_redis_cache, publish_config_change, ) +from litellm.proxy.common_utils.customer_usage import ( + CustomerUsageFailure, + customer_usage_fields, + customer_usage_http_status, + normalize_customer_key_metadata, +) from litellm.proxy.common_utils.openai_error_payload import openai_error_param, proxy_exception_for from litellm.proxy.common_utils.rbac_utils import check_org_admin_can_generate_keys from litellm.proxy.common_utils.timezone_utils import get_budget_reset_time, keeps_due_budget_reset @@ -1479,6 +1485,25 @@ async def _stored_user_role(prisma_client: PrismaClient | None, user_id: str | N return user_row.user_role if user_row is not None else None +def validated_customer_key_metadata( + metadata: object, + existing_metadata: object, + actor: UserAPIKeyAuth | None, +) -> Mapping[str, object] | None: + binding: Final = normalize_customer_key_metadata( + customer_usage_fields(metadata) if metadata is not None else None, + customer_usage_fields(existing_metadata) if existing_metadata is not None else None, + administrator=actor is not None and actor.user_role == LitellmUserRoles.PROXY_ADMIN.value, + ) + if isinstance(binding, CustomerUsageFailure): + raise HTTPException(status_code=customer_usage_http_status(binding.error), detail=binding.error.value) + return binding + + +def customer_key_actor(value: object) -> UserAPIKeyAuth | None: + return value if isinstance(value, UserAPIKeyAuth) else None + + async def _common_key_generation_helper( data: GenerateKeyRequest, user_api_key_dict: UserAPIKeyAuth, @@ -1493,6 +1518,8 @@ async def _common_key_generation_helper( prisma_client, ) + validated_customer_key_metadata(data.metadata, None, user_api_key_dict) + acting_as_team_admin: Final = _team_admin_requests_member_key( data=data, team_table=team_table, @@ -1857,7 +1884,7 @@ async def _common_key_generation_helper( try: response = await generate_key_helper_fn( - request_type="key", **data_json, table_name="key", llm_router=llm_router + request_type="key", **data_json, table_name="key", llm_router=llm_router, actor=user_api_key_dict ) except KeyProjectBindingError: if prisma_client is not None: @@ -3062,6 +3089,7 @@ async def prepare_key_update_data( *, prisma_client: PrismaClient | None = None, llm_router: Router | None = None, + user_api_key_dict: UserAPIKeyAuth | None = None, ): if data.router_settings is not None or ( "router_settings" not in data.model_fields_set @@ -3080,6 +3108,13 @@ async def prepare_key_update_data( llm_router=llm_router, ) data_json: Final[dict] = data.model_dump(exclude_unset=True) + if "metadata" in data_json: + customer_binding: Final = validated_customer_key_metadata( + data.metadata, + _StoredKeyMetadata.model_validate(existing_key_row, from_attributes=True).metadata, + user_api_key_dict, + ) + data_json["metadata"] = dict(customer_binding) if customer_binding is not None else None data_json.pop("key", None) data_json.pop("new_key", None) data_json.pop("grace_period", None) # Request-only param, not a DB column @@ -3477,7 +3512,11 @@ async def _process_single_key_update( # Prepare update data raise_on_logging_credential_change(key_request.metadata, existing_key_row, user_api_key_dict) non_default_values = await prepare_key_update_data( - data=key_request, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router + data=key_request, + existing_key_row=existing_key_row, + prisma_client=prisma_client, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, ) await _enforce_custom_key_policy( @@ -4231,7 +4270,11 @@ async def update_key_fn( _enforce_upperbound_key_params(data, fill_defaults=False) raise_on_logging_credential_change(data.metadata, existing_key_row, user_api_key_dict) non_default_values: Final = await prepare_key_update_data( - data=data, existing_key_row=existing_key_row, prisma_client=prisma_client, llm_router=llm_router + data=data, + existing_key_row=existing_key_row, + prisma_client=prisma_client, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, ) # Only validate key_alias format if it's actually being changed @@ -5384,7 +5427,9 @@ async def generate_key_helper_fn( budget_limits: list | None = None, # multiple concurrent budget windows *, llm_router: Router | None = None, + actor: UserAPIKeyAuth | None = None, ): + customer_metadata: Final = validated_customer_key_metadata(_DECODED_JSON.validate_python(metadata), None, actor) from litellm.proxy.proxy_server import premium_user, prisma_client if prisma_client is None: @@ -5436,7 +5481,7 @@ async def generate_key_helper_fn( router_settings_json: Final = safe_dumps(router_settings) if router_settings is not None else safe_dumps({}) metadata_json: Final = metadata_json_with_limits( - metadata, + customer_metadata, model_rpm_limit=model_rpm_limit, model_tpm_limit=model_tpm_limit, mcp_rpm_limit=mcp_rpm_limit, @@ -6419,7 +6464,11 @@ async def _execute_virtual_key_regeneration( _enforce_upperbound_key_params(data, fill_defaults=False) raise_on_logging_credential_change(data.metadata, key_in_db, user_api_key_dict) non_default_values = await prepare_key_update_data( - data=data, existing_key_row=key_in_db, prisma_client=prisma_client, llm_router=llm_router + data=data, + existing_key_row=key_in_db, + prisma_client=prisma_client, + llm_router=llm_router, + user_api_key_dict=user_api_key_dict, ) # Only validate key_alias format if it's actually being changed new_key_alias: Final = non_default_values.get("key_alias") diff --git a/litellm/proxy/management_endpoints/management_v1/__init__.py b/litellm/proxy/management_endpoints/management_v1/__init__.py index c42a786db0d..c4483ef7ff1 100644 --- a/litellm/proxy/management_endpoints/management_v1/__init__.py +++ b/litellm/proxy/management_endpoints/management_v1/__init__.py @@ -8,6 +8,9 @@ from litellm.proxy.management.users.endpoints import router as current_user_rout from litellm.proxy.management_endpoints.management_v1.budgets import ( router as budgets_router, ) +from litellm.proxy.management_endpoints.management_v1.customer_usage import ( + router as customer_usage_router, +) from litellm.proxy.management_endpoints.management_v1.customers import ( router as customers_router, ) @@ -24,6 +27,7 @@ from litellm.proxy.management_endpoints.management_v1.users import ( router: Final = APIRouter() router.include_router(budgets_router) router.include_router(customers_router) +router.include_router(customer_usage_router) router.include_router(spend_logs_router) router.include_router(teams_router) router.include_router(users_router) diff --git a/litellm/proxy/management_endpoints/management_v1/customer_usage.py b/litellm/proxy/management_endpoints/management_v1/customer_usage.py new file mode 100644 index 00000000000..4f0a7fc154e --- /dev/null +++ b/litellm/proxy/management_endpoints/management_v1/customer_usage.py @@ -0,0 +1,220 @@ +from collections.abc import Sequence +from enum import Enum +from functools import partial +from typing import Annotated, Final, Protocol + +from fastapi import APIRouter, Depends +from pydantic import TypeAdapter + +from litellm.litellm_core_utils.core_helpers import require_attributes +from litellm.proxy.list_api.common import reject_unknown_query_params +from litellm.proxy.management_endpoints.common_daily_activity import DailySpendRecord, get_daily_activity +from litellm.proxy.management_endpoints.management_v1.common import MANAGEMENT_V1_PREFIX +from litellm.proxy.management_endpoints.management_v1.customers import ( + raise_public, + reject_duplicate_list_params, + require_customer_admin, +) +from litellm.proxy.management_helpers.customer_hierarchy import ErrorCode, Failure, PrismaRecord +from litellm.proxy.management_helpers.utils import ( + management_endpoint_wrapper, # pyright: ignore[reportUnknownVariableType] # legacy decorator is untyped +) +from litellm.proxy.utils import PrismaClient +from litellm.repositories.chunked_in import find_many_in +from litellm.repositories.prisma_protocols import DatabaseClient, TableActions +from litellm.types.proxy.management_endpoints.common_daily_activity import SpendAnalyticsPaginatedResponse + + +class CustomerEntity(str, Enum): + ORGANIZATION = "customer_org_id" + TEAM = "customer_team_id" + MEMBER = "customer_member_id" + + +class CustomerMetadataDatabase(Protocol): + @property + def litellm_customerorganizationtable(self) -> TableActions[PrismaRecord]: ... + + @property + def litellm_customerteamtable(self) -> TableActions[PrismaRecord]: ... + + @property + def litellm_customermembertable(self) -> TableActions[PrismaRecord]: ... + + +_TEXT: Final = TypeAdapter(str) +_OPTIONAL_TEXT: Final[TypeAdapter[str | None]] = TypeAdapter(str | None) + + +def get_customer_usage_prisma() -> PrismaClient: + from litellm.proxy.proxy_server import prisma_client + + if prisma_client is None: + raise_public(Failure(ErrorCode.DATABASE_UNAVAILABLE)) + return prisma_client + + +def metadata_table(prisma_client: PrismaClient, entity: CustomerEntity) -> TableActions[PrismaRecord]: + database: Final = require_attributes(prisma_client, DatabaseClient, "db").db + tables: Final = require_attributes( + database, + CustomerMetadataDatabase, + "litellm_customerorganizationtable", + "litellm_customerteamtable", + "litellm_customermembertable", + ) + if entity is CustomerEntity.ORGANIZATION: + return tables.litellm_customerorganizationtable + if entity is CustomerEntity.TEAM: + return tables.litellm_customerteamtable + return tables.litellm_customermembertable + + +async def page_aliases( + table: TableActions[PrismaRecord], entity: CustomerEntity, records: Sequence[DailySpendRecord] +) -> dict[str, dict[str, object]]: # mutable-ok: get_daily_activity requires dictionary-returning metadata callbacks + identifiers: Final = tuple(_OPTIONAL_TEXT.validate_python(getattr(row, entity.value, None)) for row in records) + rows: Final = await find_many_in(table, entity.value, tuple(value for value in identifiers if value)) + values: Final = tuple(row.model_dump() for row in rows) + return { + _TEXT.validate_python(row[entity.value]): { + "alias": _OPTIONAL_TEXT.validate_python(row.get("alias")) or _TEXT.validate_python(row[entity.value]) + } + for row in values + } + + +async def customer_activity( + prisma_client: PrismaClient, + entity: CustomerEntity, + entity_ids: str | None, + excluded_ids: str | None, + start_date: str | None, + end_date: str | None, + model: str | None, + api_key: str | None, + page: int, + page_size: int, +) -> SpendAnalyticsPaginatedResponse: + return await get_daily_activity( + prisma_client=prisma_client, + table_name="litellm_dailycustomermemberspend", + entity_id_field=entity.value, + entity_id=entity_ids.split(",") if entity_ids else None, + entity_metadata_field=None, + exclude_entity_ids=excluded_ids.split(",") if excluded_ids else None, + start_date=start_date, + end_date=end_date, + model=model, + api_key=api_key, + page=page, + page_size=page_size, + resolve_entity_metadata=partial(page_aliases, metadata_table(prisma_client, entity), entity), + ) + + +router: Final = APIRouter( + prefix=MANAGEMENT_V1_PREFIX, + tags=["customer usage"], + dependencies=( + Depends(require_customer_admin), + Depends(reject_unknown_query_params), + Depends(reject_duplicate_list_params), + ), +) + + +@router.get( + "/customer_organizations/daily/activity", + response_model=SpendAnalyticsPaginatedResponse, + description="Get paginated daily customer organization activity with model, provider, key and organization breakdowns. " + "Filters accept comma-separated organization IDs. Proxy admins only.", +) +@management_endpoint_wrapper +async def customer_organization_daily_activity( + prisma_client: Annotated[PrismaClient, Depends(get_customer_usage_prisma)], + customer_org_ids: str | None = None, + start_date: str | None = None, + end_date: str | None = None, + model: str | None = None, + api_key: str | None = None, + page: int = 1, + page_size: int = 10, + exclude_customer_org_ids: str | None = None, +) -> SpendAnalyticsPaginatedResponse: + return await customer_activity( + prisma_client, + CustomerEntity.ORGANIZATION, + customer_org_ids, + exclude_customer_org_ids, + start_date, + end_date, + model, + api_key, + page, + page_size, + ) + + +@router.get( + "/customer_teams/daily/activity", + response_model=SpendAnalyticsPaginatedResponse, + description="Get paginated daily customer team activity with model, provider, key and team breakdowns. " + "Filters accept comma-separated team IDs. Proxy admins only.", +) +@management_endpoint_wrapper +async def customer_team_daily_activity( + prisma_client: Annotated[PrismaClient, Depends(get_customer_usage_prisma)], + customer_team_ids: str | None = None, + start_date: str | None = None, + end_date: str | None = None, + model: str | None = None, + api_key: str | None = None, + page: int = 1, + page_size: int = 10, + exclude_customer_team_ids: str | None = None, +) -> SpendAnalyticsPaginatedResponse: + return await customer_activity( + prisma_client, + CustomerEntity.TEAM, + customer_team_ids, + exclude_customer_team_ids, + start_date, + end_date, + model, + api_key, + page, + page_size, + ) + + +@router.get( + "/customer_members/daily/activity", + response_model=SpendAnalyticsPaginatedResponse, + description="Get paginated daily customer member activity with model, provider, key and member breakdowns. " + "Filters accept comma-separated member IDs. Proxy admins only.", +) +@management_endpoint_wrapper +async def customer_member_daily_activity( + prisma_client: Annotated[PrismaClient, Depends(get_customer_usage_prisma)], + customer_member_ids: str | None = None, + start_date: str | None = None, + end_date: str | None = None, + model: str | None = None, + api_key: str | None = None, + page: int = 1, + page_size: int = 10, + exclude_customer_member_ids: str | None = None, +) -> SpendAnalyticsPaginatedResponse: + return await customer_activity( + prisma_client, + CustomerEntity.MEMBER, + customer_member_ids, + exclude_customer_member_ids, + start_date, + end_date, + model, + api_key, + page, + page_size, + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b37eede1745..f9d4d5704d8 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -19970,6 +19970,7 @@ async def _generate_onboarding_ui_session_token(user_obj: _UserTableRow) -> str: response: Final = await generate_key_helper_fn( llm_router=llm_router, request_type="key", + actor=None, **{ "user_role": user_obj.user_role, "duration": LITELLM_UI_SESSION_DURATION, diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 2c0246bb94e..bac27daa5de 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -1035,6 +1035,45 @@ model LiteLLM_DailyOrganizationSpend { } // Track daily end user (customer) spend metrics per model and key +model LiteLLM_DailyCustomerMemberSpend { + id String @id + date String + customer_org_id String + customer_team_id String @default("") + customer_member_id String @default("") + model String @default("") + model_group String @default("") + custom_llm_provider String @default("") + mcp_namespaced_tool_name String @default("") + endpoint String @default("") + api_key String @default("") + spend Float @default(0.0) + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + compression_saved_tokens BigInt @default(0) + compression_savings_spend Float @default(0.0) + prompt_caching_savings_spend Float @default(0.0) + gateway_injected_caching_savings_spend Float @default(0.0) + autorouter_savings_spend Float @default(0.0) + guardrail_spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@index([date]) + @@index([customer_org_id, date]) + @@index([customer_team_id, date]) + @@index([customer_member_id, date]) + @@index([api_key]) + @@index([model]) +} + model LiteLLM_DailyEndUserSpend { id String @id @default(uuid()) end_user_id String? diff --git a/litellm/repositories/accounting_delivery_repository.py b/litellm/repositories/accounting_delivery_repository.py index 624c2b8851a..e172e863b63 100644 --- a/litellm/repositories/accounting_delivery_repository.py +++ b/litellm/repositories/accounting_delivery_repository.py @@ -14,6 +14,18 @@ if TYPE_CHECKING: class LedgerTransaction(Protocol): + @property + def litellm_dailycustomermemberspend(self) -> TableActions[object]: ... + + @property + def litellm_customerorganizationtable(self) -> TableActions[object]: ... + + @property + def litellm_customerteamtable(self) -> TableActions[object]: ... + + @property + def litellm_customermembertable(self) -> TableActions[object]: ... + @property def litellm_dailyguardrailusageunits(self) -> TableActions[object]: ... diff --git a/litellm/repositories/daily_activity_repository.py b/litellm/repositories/daily_activity_repository.py index 9aa0950785b..d8b4d1eb8d5 100644 --- a/litellm/repositories/daily_activity_repository.py +++ b/litellm/repositories/daily_activity_repository.py @@ -96,6 +96,9 @@ class _DailyActivityDatabase(Protocol): @property def litellm_dailyenduserspend(self) -> TableActions[DailyActivityRow]: ... + @property + def litellm_dailycustomermemberspend(self) -> TableActions[DailyActivityRow]: ... + @property def litellm_dailyagentspend(self) -> TableActions[DailyActivityRow]: ... @@ -178,6 +181,8 @@ def _daily_rows_table( return prisma_client.db.litellm_dailyorganizationspend if table is DailyActivityTable.CUSTOMER: return prisma_client.db.litellm_dailyenduserspend + if table is DailyActivityTable.CUSTOMER_MEMBER: + return prisma_client.db.litellm_dailycustomermemberspend if table is DailyActivityTable.AGENT: return prisma_client.db.litellm_dailyagentspend assert_never(table) diff --git a/litellm/repositories/daily_activity_sql.py b/litellm/repositories/daily_activity_sql.py index dc3ca19bb30..4b575761e19 100644 --- a/litellm/repositories/daily_activity_sql.py +++ b/litellm/repositories/daily_activity_sql.py @@ -43,6 +43,7 @@ PRISMA_TO_PG_TABLE: Final[Mapping[DailyActivityTable, str]] = MappingProxyType( DailyActivityTable.TAG: "LiteLLM_DailyTagSpend", DailyActivityTable.ORGANIZATION: "LiteLLM_DailyOrganizationSpend", DailyActivityTable.CUSTOMER: "LiteLLM_DailyEndUserSpend", + DailyActivityTable.CUSTOMER_MEMBER: "LiteLLM_DailyCustomerMemberSpend", DailyActivityTable.AGENT: "LiteLLM_DailyAgentSpend", } ) diff --git a/litellm/schema.prisma b/litellm/schema.prisma index 2c0246bb94e..bac27daa5de 100644 --- a/litellm/schema.prisma +++ b/litellm/schema.prisma @@ -1035,6 +1035,45 @@ model LiteLLM_DailyOrganizationSpend { } // Track daily end user (customer) spend metrics per model and key +model LiteLLM_DailyCustomerMemberSpend { + id String @id + date String + customer_org_id String + customer_team_id String @default("") + customer_member_id String @default("") + model String @default("") + model_group String @default("") + custom_llm_provider String @default("") + mcp_namespaced_tool_name String @default("") + endpoint String @default("") + api_key String @default("") + spend Float @default(0.0) + prompt_tokens BigInt @default(0) + completion_tokens BigInt @default(0) + cache_read_input_tokens BigInt @default(0) + cache_creation_input_tokens BigInt @default(0) + compression_saved_tokens BigInt @default(0) + compression_savings_spend Float @default(0.0) + prompt_caching_savings_spend Float @default(0.0) + gateway_injected_caching_savings_spend Float @default(0.0) + autorouter_savings_spend Float @default(0.0) + guardrail_spend Float @default(0.0) + api_requests BigInt @default(0) + successful_requests BigInt @default(0) + failed_requests BigInt @default(0) + total_response_time_ms BigInt @default(0) + timed_requests BigInt @default(0) + created_at DateTime @default(now()) + updated_at DateTime @updatedAt + + @@index([date]) + @@index([customer_org_id, date]) + @@index([customer_team_id, date]) + @@index([customer_member_id, date]) + @@index([api_key]) + @@index([model]) +} + model LiteLLM_DailyEndUserSpend { id String @id @default(uuid()) end_user_id String? diff --git a/litellm/types/repositories/daily_activity.py b/litellm/types/repositories/daily_activity.py index 8d68d92a336..babf4a6c8fd 100644 --- a/litellm/types/repositories/daily_activity.py +++ b/litellm/types/repositories/daily_activity.py @@ -12,6 +12,7 @@ class DailyActivityTable(str, Enum): TAG = "litellm_dailytagspend" ORGANIZATION = "litellm_dailyorganizationspend" CUSTOMER = "litellm_dailyenduserspend" + CUSTOMER_MEMBER = "litellm_dailycustomermemberspend" AGENT = "litellm_dailyagentspend" @@ -22,6 +23,7 @@ _ENTITY_FIELDS: Mapping[DailyActivityTable, frozenset[str]] = MappingProxyType( DailyActivityTable.TAG: frozenset(("tag",)), DailyActivityTable.ORGANIZATION: frozenset(("organization_id",)), DailyActivityTable.CUSTOMER: frozenset(("end_user_id",)), + DailyActivityTable.CUSTOMER_MEMBER: frozenset(("customer_org_id", "customer_team_id", "customer_member_id")), DailyActivityTable.AGENT: frozenset(("agent_id",)), } ) diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerMemberSpend.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerMemberSpend.tsx new file mode 100644 index 00000000000..dfcd9470764 --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerMemberSpend.tsx @@ -0,0 +1,166 @@ +"use client"; + +import { useMemo } from "react"; +import { parseAsString, parseAsInteger, useQueryStates } from "nuqs"; +import { Button } from "@/components/ui/button"; +import { useTranslation } from "react-i18next"; +import { Search } from "lucide-react"; +import { formatCompact, formatUsd } from "@/app/(dashboard)/usage/_components/components/overview/overviewData"; +import { Input } from "@/components/ui/input"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { cn } from "@/lib/cva.config"; +import type { CustomerHierarchyUsage, CustomerUsageEntity } from "./customerUsageApi"; + +const MEMBER_TABLE_PARSERS = { memberSearch: parseAsString.withDefault(""), memberPage: parseAsInteger.withDefault(1) }; +const PAGE_SIZE = 20; +const memberStatus = (member: CustomerUsageEntity) => { + if (member.budgetLimit === null) return "customerPortal.analytics.noCap" as const; + if (member.spend > member.budgetLimit) return "customerPortal.analytics.overBudget" as const; + if (member.spend === member.budgetLimit) return "customerPortal.analytics.atBudget" as const; + return "customerPortal.analytics.withinBudget" as const; +}; + +export default function CustomerMemberSpend({ + data, + onMemberClick, +}: { + data: CustomerHierarchyUsage; + onMemberClick: (id: string) => void; +}) { + const { t } = useTranslation(); + const [table, setTable] = useQueryStates(MEMBER_TABLE_PARSERS); + const search = table.memberSearch; + const usage = useMemo(() => new Map(data.members.map((member) => [member.id, member])), [data.members]); + const members = useMemo( + () => + data.memberBudgets + .filter((member) => member.label.toLocaleLowerCase().includes(search.toLocaleLowerCase())) + .map((member) => ({ ...member, rangeUsage: usage.get(member.id) })) + .sort((a, b) => (b.rangeUsage?.spend ?? 0) - (a.rangeUsage?.spend ?? 0) || a.label.localeCompare(b.label)), + [data.memberBudgets, search, usage], + ); + + const lastPage = Math.max(1, Math.ceil(members.length / PAGE_SIZE)); + const page = Math.max(1, Math.min(table.memberPage, lastPage)); + const start = (page - 1) * PAGE_SIZE; + return ( +
+
+
+

{t("customerPortal.analytics.memberSpend")}

+

{t("customerPortal.analytics.memberSpendHint")}

+
+
+
+
+ + + + {t("customerPortal.analytics.member")} + {t("customerPortal.analytics.spend")} + {t("customerPortal.analytics.tokens")} + {t("customerPortal.analytics.currentBudget")} + {t("customerPortal.analytics.status")} + + + + {members.slice(start, start + PAGE_SIZE).map((member) => { + const atLimit = member.budgetLimit !== null && member.spend >= member.budgetLimit; + return ( + + +
+ +
+ + {member.id} +
+
+
+ {formatUsd(member.rangeUsage?.spend ?? 0)} + + {formatCompact(member.rangeUsage?.tokens ?? 0)} + + + {member.budgetLimit === null + ? t("customerPortal.analytics.noBudget") + : `${formatUsd(member.spend)} / ${formatUsd(member.budgetLimit)}`} + + + + {t(memberStatus(member))} + + +
+ ); + })} + {members.length === 0 && ( + + + {t(search ? "customerPortal.analytics.noMatches" : "customerPortal.analytics.noMembers")} + + + )} +
+
+ {members.length > PAGE_SIZE && ( +
+ + {t("customerPortal.analytics.pagination", { + from: start + 1, + to: Math.min(start + PAGE_SIZE, members.length), + total: members.length, + })} + +
+ + +
+
+ )} +
+ ); +} diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerUsageHeader.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerUsageHeader.tsx new file mode 100644 index 00000000000..618aeb0b80b --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerUsageHeader.tsx @@ -0,0 +1,112 @@ +import type { ParseKeys } from "i18next"; +import { ChevronRight } from "lucide-react"; +import { useTranslation } from "react-i18next"; +import { formatUsd } from "@/app/(dashboard)/usage/_components/components/overview/overviewData"; +import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; +import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; +import { uiHref } from "@/utils/uiHref"; +import type { CustomerHierarchyUsage, CustomerUsageEntity, CustomerUsageScope } from "./customerUsageApi"; + +interface CustomerUsageHeaderProps { + scope: CustomerUsageScope; + data: CustomerHierarchyUsage | undefined; + dateValue: DateRangePickerValue; + preset: string | null; + onOrganizationClick: () => void; + onTeamClick: () => void; + onDateRangeChange: (value: DateRangePickerValue, preset: string | null) => void; +} + +const scopeLabel = (scope: CustomerUsageScope): ParseKeys => { + if (scope.member) return "customerPortal.analytics.memberUsage"; + if (scope.team) return "customerPortal.analytics.teamUsage"; + return "customerPortal.analytics.organizationUsage"; +}; + +function CustomerUsageBreadcrumbs({ + scope, + data, + onOrganizationClick, + onTeamClick, +}: Pick) { + const { t } = useTranslation(); + return ( + + ); +} + +function CustomerCurrentSpend({ entity }: { entity: CustomerUsageEntity }) { + const { t } = useTranslation(); + return ( +

+ {t("customerPortal.analytics.currentSpend")}{" "} + {formatUsd(entity.spend)} + {entity.budgetLimit === null + ? ` ยท ${t("customerPortal.analytics.noBudget")}` + : ` / ${formatUsd(entity.budgetLimit)}`} +

+ ); +} + +export default function CustomerUsageHeader({ + scope, + data, + dateValue, + preset, + onOrganizationClick, + onTeamClick, + onDateRangeChange, +}: CustomerUsageHeaderProps) { + const { t } = useTranslation(); + const entity = data?.member ?? data?.team ?? data?.organization; + const label = entity?.label ?? scope.member ?? scope.team ?? scope.org; + return ( +
+ +
+
+

{t(scopeLabel(scope))}

+

{label}

+

{t("customerPortal.analytics.subtitle")}

+ {entity && } +
+ +
+
+ ); +} diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerUsageRankings.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerUsageRankings.tsx new file mode 100644 index 00000000000..aa881cb997f --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerUsageRankings.tsx @@ -0,0 +1,200 @@ +"use client"; + +import type { ParseKeys } from "i18next"; +import { useTranslation } from "react-i18next"; +import { ChevronRight, Search } from "lucide-react"; +import { + type UsageMetric, + formatCompact, + formatUsd, +} from "@/app/(dashboard)/usage/_components/components/overview/overviewData"; +import { Input } from "@/components/ui/input"; +import { Button } from "@/components/ui/button"; +import { Table, TableBody, TableCell, TableHead, TableHeader, TableRow } from "@/components/ui/table"; +import { cn } from "@/lib/cva.config"; +import type { CustomerUsageRow } from "./customerUsageData"; +import { ModelMark } from "@/app/(dashboard)/usage/_components/components/overview/BreakdownChart"; +import { colorForModel } from "@/app/(dashboard)/usage/_components/components/overview/overviewData"; +import { providerForModel } from "@/app/(dashboard)/usage/_components/components/overview/modelProvider"; + +export const CUSTOMER_USAGE_DIMENSIONS = ["models", "api_keys", "teams"] as const; +export type CustomerUsageDimension = (typeof CUSTOMER_USAGE_DIMENSIONS)[number]; + +const DIMENSION_LABELS: Record = { + models: "customerPortal.analytics.models", + api_keys: "customerPortal.analytics.keys", + teams: "customerPortal.analytics.teams", +}; +const PAGE_SIZE = 10; + +interface CustomerUsageRankingsProps { + rows: readonly CustomerUsageRow[]; + dimension: CustomerUsageDimension; + metric: UsageMetric; + search: string; + page: number; + colorByModel: ReadonlyMap; + onSearchChange: (search: string) => void; + onPageChange: (page: number) => void; + onDrillDown: (dimension: "teams", id: string) => void; +} + +export default function CustomerUsageRankings({ + rows, + dimension, + metric, + search, + page, + colorByModel, + onSearchChange, + onPageChange, + onDrillDown, +}: CustomerUsageRankingsProps) { + const { t } = useTranslation(); + const filtered = rows.filter((row) => row.label.toLocaleLowerCase().includes(search.toLocaleLowerCase())); + const lastPage = Math.max(1, Math.ceil(filtered.length / PAGE_SIZE)); + const currentPage = Math.min(Math.max(1, page), lastPage); + const start = (currentPage - 1) * PAGE_SIZE; + const canDrillDown = dimension === "teams"; + + return ( +
+
+
+
+
+
+

{t("customerPortal.analytics.rankingHint")}

+ + + + {t("customerPortal.analytics.rank")} + {t(DIMENSION_LABELS[dimension])} + {t("customerPortal.analytics.share")} + + {t("customerPortal.analytics.spend")} + + + {t("customerPortal.analytics.tokens")} + + + {t("customerPortal.analytics.requests")} + + + + + {filtered.slice(start, start + PAGE_SIZE).map((row, index) => ( + + {start + index + 1} + + {canDrillDown && row.id ? ( + + ) : ( + + {dimension === "models" && ( + + )} + + {row.label} + {dimension === "models" && providerForModel(row.id) && ( + + {t("usage.breakdown.byProvider", { provider: providerForModel(row.id) })} + + )} + + + )} + + +
+ + + + {formatUsd(row.spend)} + + + {formatCompact(row.tokens)} + + + {formatCompact(row.requests)} + + + ))} + {filtered.length === 0 && ( + + + {t(search ? "customerPortal.analytics.noMatches" : "customerPortal.analytics.noBreakdown")} + + + )} + +
+ {filtered.length > PAGE_SIZE && ( +
+ + {t("customerPortal.analytics.pagination", { + from: start + 1, + to: Math.min(start + PAGE_SIZE, filtered.length), + total: filtered.length, + })} + +
+ + +
+
+ )} +
+
+ ); +} diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerUsageTrend.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerUsageTrend.tsx new file mode 100644 index 00000000000..0ba88ea4d6d --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/CustomerUsageTrend.tsx @@ -0,0 +1,152 @@ +"use client"; + +import { useMemo } from "react"; +import { useTranslation } from "react-i18next"; +import { + bucketSeries, + bucketTotals, + formatMetricValue, + labelForDate, + type UsageMetric, + type Granularity, + type OverviewTotals, +} from "@/app/(dashboard)/usage/_components/components/overview/overviewData"; +import { ChartSkeleton, Segmented } from "@/app/(dashboard)/usage/_components/components/overview/Primitives"; +import { StackedUsageChart } from "@/components/shared/charts"; +import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; +import type { CustomerHierarchyUsage } from "./customerUsageApi"; +import { customerUsageSeries } from "./customerUsageChartData"; +import { CUSTOMER_USAGE_DIMENSIONS, type CustomerUsageDimension } from "./CustomerUsageRankings"; +import { ModelMark } from "@/app/(dashboard)/usage/_components/components/overview/BreakdownChart"; + +const DIMENSION_LABELS = { + models: "customerPortal.analytics.models", + api_keys: "customerPortal.analytics.keys", + teams: "customerPortal.analytics.teams", +} as const; + +interface CustomerUsageTrendProps { + data: CustomerHierarchyUsage | undefined; + dimension: CustomerUsageDimension; + onDimensionChange: (dimension: CustomerUsageDimension) => void; + totals: OverviewTotals; + metric: UsageMetric; + bucket: Granularity; + loading: boolean; + colorByModel: ReadonlyMap; + onMetricChange: (metric: UsageMetric) => void; + onBucketChange: (bucket: Granularity) => void; +} + +export default function CustomerUsageTrend({ + data, + dimension, + onDimensionChange, + totals, + metric, + bucket, + loading, + colorByModel, + onMetricChange, + onBucketChange, +}: CustomerUsageTrendProps) { + const { t } = useTranslation(); + const dailySeries = useMemo( + () => + data + ? customerUsageSeries(data, dimension, metric, colorByModel) + : { data: [], keys: [], labels: [], colors: [] }, + [data, dimension, metric, colorByModel], + ); + const series = useMemo(() => bucketSeries(dailySeries, bucket), [dailySeries, bucket]); + const bucketValues = useMemo(() => bucketTotals(series), [series]); + const empty = totals.requests === 0 && totals.tokens === 0 && totals.spend === 0; + return ( +
+
+
+

{t("customerPortal.analytics.trend")}

+

+ {t("customerPortal.analytics.trendHint", { dimension: t(DIMENSION_LABELS[dimension]).toLocaleLowerCase() })} +

+
+
+ + + label={t("customerPortal.analytics.metric")} + value={metric} + options={[ + { value: "spend", label: t("customerPortal.analytics.spend") }, + { value: "tokens", label: t("customerPortal.analytics.tokens") }, + { value: "requests", label: t("customerPortal.analytics.requests") }, + ]} + onChange={(metric) => onMetricChange(metric)} + /> + +
+ {loading && ( +
+ +
+ )} + {!loading && empty && ( +
+ {t("customerPortal.analytics.empty")} +
+ )} + {!loading && !empty && ( + <> + labelForDate(series, date)} + format={(value) => formatMetricValue(value, metric)} + totalFor={(date) => bucketValues.get(date)} + className="h-72 sm:h-80" + /> +
+ {series.labels.map((label, index) => ( + + {dimension === "models" ? ( + + ) : ( + + ))} +
+ + )} +
+ ); +} diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageApi.test.ts b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageApi.test.ts new file mode 100644 index 00000000000..72443540bab --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageApi.test.ts @@ -0,0 +1,336 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { registerBaseUrlGetter } from "@/lib/http/runtime"; +import { EMPTY_DAILY_ACTIVITY_METADATA } from "@/components/UsagePage/dailyActivityApi"; +import type { DailyData, SpendMetrics } from "@/components/UsagePage/types"; +import { customerHierarchyUsageCall } from "./customerUsageApi"; + +const audit = { + created_at: "2026-11-01T00:00:00Z", + created_by: "admin", + updated_at: "2026-11-01T00:00:00Z", + updated_by: "admin", + spend: 40, +}; +const organization = { + ...audit, + customer_org_id: "org-1", + alias: "Customer organization", + budget: { max_budget: 100 }, + budget_reset_at: "2026-12-01T00:00:00Z", + default_budgets: [], +}; +const teams = [ + { ...audit, customer_team_id: "team-a", parent_org_id: "org-1", alias: "Product", budget: { max_budget: 50 } }, + { ...audit, customer_team_id: "team-b", parent_org_id: "org-1", alias: "Operations", budget: null }, +]; +const members = [ + { + ...audit, + customer_member_id: "member-a", + parent_team_id: "team-a", + alias: "Member One", + effective_budget: { max_budget: 10 }, + }, + { + ...audit, + customer_member_id: "member-b", + parent_team_id: "team-b", + alias: "Member Two", + effective_budget: { max_budget: 20 }, + }, +]; +const metrics = (spend: number, tokens: number, requests: number): SpendMetrics => ({ + spend, + prompt_tokens: tokens, + completion_tokens: 0, + total_tokens: tokens, + api_requests: requests, + successful_requests: requests, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, +}); +const records = [ + { + organization: "org-1", + team: "team-a", + member: "member-a", + model: "model-a", + key: "hash-a", + alias: "Assistant key", + metrics: metrics(3, 100, 8), + }, + { + organization: "org-1", + team: "team-b", + member: "member-b", + model: "model-b", + key: "hash-b", + alias: "Search key", + metrics: metrics(1, 500, 2), + }, +]; +type UsageRecord = (typeof records)[number]; +const daily = (record: UsageRecord, entityId: string): DailyData => { + const key = { metrics: record.metrics, metadata: { key_alias: record.alias, team_id: null } }; + const entity = { metrics: record.metrics, metadata: {}, api_key_breakdown: { [record.key]: key } }; + return { + date: "2026-11-01", + metrics: record.metrics, + breakdown: { + models: { [record.model]: entity }, + model_groups: { [record.model]: entity }, + providers: {}, + mcp_servers: {}, + api_keys: { [record.key]: key }, + entities: { [entityId]: entity }, + }, + }; +}; +const json = (body: unknown, status = 200) => + new Response(JSON.stringify(body), { status, headers: { "Content-Type": "application/json" } }); +const requestUrl = (input: Parameters[0]) => + new URL(input instanceof Request ? input.url : String(input), "http://localhost"); +const fetchMock = vi.fn(); +const urls = () => fetchMock.mock.calls.map(([input]) => requestUrl(input)); +const dailyUrls = () => urls().filter((url) => url.pathname.endsWith("/daily/activity")); +const request = { + scope: { org: "org-1", team: null, member: null }, + accessToken: "test-session", + startDate: "2026-11-01", + endDate: "2026-11-03", +}; +const list = (data: T[], page = 1, totalPages = 1) => ({ + data, + meta: { total_count: data.length, page, page_size: 100, total_pages: totalPages }, +}); + +const activityDimension = (resource: string | undefined): "organization" | "team" | "member" => { + if (resource === "customer_organizations") return "organization"; + if (resource === "customer_teams") return "team"; + return "member"; +}; +const filterFor = (resource: string | undefined) => + ({ + organization: "customer_org_ids", + team: "customer_team_ids", + member: "customer_member_ids", + })[activityDimension(resource)]; +const activityResponse = (url: URL) => { + const collection = url.pathname.split("/").at(-3); + const dimension = activityDimension(collection); + const filter = filterFor(collection); + const ids = url.searchParams.get(filter)?.split(",") ?? []; + const selected = records.filter((row) => ids.includes(row[dimension])); + const page = Number(url.searchParams.get("page")); + const current = selected.slice(page - 1, page); + const values = current[0]?.metrics ?? metrics(0, 0, 0); + return json({ + results: current.map((row) => daily(row, row[dimension])), + metadata: { + ...EMPTY_DAILY_ACTIVITY_METADATA, + total_spend: values.spend, + total_tokens: values.total_tokens, + total_api_requests: values.api_requests, + page, + total_pages: selected.length, + has_more: page < selected.length, + }, + }); +}; + +const serve = (input: Parameters[0]): Response => { + const url = requestUrl(input); + if (url.pathname.endsWith("/daily/activity")) return activityResponse(url); + if (url.pathname === "/management/v1/customer_organizations/org-1") return json(organization); + if (url.pathname === "/management/v1/customer_teams") { + const page = Number(url.searchParams.get("page")); + return json(list([teams[page - 1]], page, 2)); + } + if (url.pathname === "/management/v1/customer_members") + return json(list(members.filter((member) => member.parent_team_id === url.searchParams.get("parent_team_id")))); + const team = teams.find((team) => url.pathname.endsWith(`/customer_teams/${team.customer_team_id}`)); + if (team) return json(team); + const member = members.find((member) => url.pathname.endsWith(`/customer_members/${member.customer_member_id}`)); + if (member) return json(member); + throw new Error(`Unexpected usage request ${url.pathname}`); +}; + +beforeEach(() => { + registerBaseUrlGetter(() => "http://localhost"); + fetchMock.mockReset(); + fetchMock.mockImplementation((input) => Promise.resolve(serve(input))); + vi.stubGlobal("fetch", fetchMock); +}); +afterEach(() => vi.unstubAllGlobals()); + +describe("customerHierarchyUsageCall", () => { + it("should read exact daily activity routes, all pages, and scoped descendant lists", async () => { + fetchMock.mockImplementation((input) => { + const url = requestUrl(input); + if (url.pathname === "/management/v1/customer_members" && url.searchParams.get("parent_team_id") === "team-a") { + const page = Number(url.searchParams.get("page")); + const member = page === 1 ? members[0] : { ...members[0], customer_member_id: "member-c" }; + return Promise.resolve(json(list([member], page, 2))); + } + return Promise.resolve(serve(input)); + }); + const result = await customerHierarchyUsageCall(request); + expect(result.metadata.total_spend).toBe(4); + expect(result.metadata.total_tokens).toBe(600); + expect(result.metadata.total_api_requests).toBe(10); + expect(result.results).toHaveLength(1); + expect(result.results[0].metrics.spend).toBe(4); + expect(result.results[0].breakdown.model_groups["model-b"].metrics.total_tokens).toBe(500); + expect(result.api_keys.map((row) => row.label).sort()).toEqual(["Assistant key", "Search key"]); + expect(result.teams.map((row) => [row.id, row.spend])).toEqual([ + ["team-a", 3], + ["team-b", 1], + ]); + expect(result.members.map((row) => [row.id, row.tokens])).toEqual([ + ["member-a", 100], + ["member-b", 500], + ]); + expect(result.teamResults).toHaveLength(1); + expect(result.teamResults[0].metrics.spend).toBe(4); + expect(result.teamResults[0].breakdown.entities["team-a"]).toMatchObject({ + metadata: { alias: "Product" }, + metrics: { spend: 3 }, + }); + expect(result.memberResults).toHaveLength(1); + expect(result.memberResults[0].breakdown.entities["member-b"].metrics.total_tokens).toBe(500); + expect(result.memberBudgets[0]).toMatchObject({ id: "member-a", spend: 40, budgetLimit: 10 }); + for (const url of dailyUrls()) { + expect(url.pathname).toMatch(/^\/management\/v1\/customer_(organizations|teams|members)\/daily\/activity$/); + expect(url.searchParams.get("start_date")).toBe(request.startDate); + expect(url.searchParams.get("end_date")).toBe(request.endDate); + expect(url.searchParams.get("page_size")).toBe("100"); + expect([...url.searchParams.keys()].filter((name) => name.endsWith("_ids"))).toHaveLength(1); + } + const organizationPages = dailyUrls().filter((url) => url.searchParams.get("customer_org_ids") === "org-1"); + expect(organizationPages.map((url) => url.searchParams.get("page"))).toEqual(["1", "2"]); + expect( + urls() + .filter((url) => url.pathname === "/management/v1/customer_teams") + .map((url) => [url.searchParams.get("parent_org_id"), url.searchParams.get("page")]), + ).toEqual([ + ["org-1", "1"], + ["org-1", "2"], + ]); + expect( + dailyUrls() + .find((url) => url.searchParams.has("customer_team_ids")) + ?.searchParams.get("customer_team_ids"), + ).toBe("team-a,team-b"); + expect( + dailyUrls() + .find((url) => url.searchParams.has("customer_member_ids")) + ?.searchParams.get("customer_member_ids"), + ).toBe("member-a,member-c,member-b"); + expect( + urls().some( + (url) => + url.pathname === "/management/v1/customer_members" && + url.searchParams.get("parent_team_id") === "team-a" && + url.searchParams.get("page") === "2", + ), + ).toBe(true); + }); + + it("should infer member ancestry and retain live spend and effective budget separately from range totals", async () => { + const result = await customerHierarchyUsageCall({ + ...request, + scope: { org: "org-1", team: null, member: "member-a" }, + }); + expect(result.team?.id).toBe("team-a"); + expect(result.member).toMatchObject({ id: "member-a", spend: 40, budgetLimit: 10 }); + expect(result.organization).toMatchObject({ spend: 40, budgetLimit: 100, budgetResetAt: "2026-12-01T00:00:00Z" }); + expect(result.metadata.total_spend).toBe(3); + expect(result.teams).toEqual([{ id: "team-a", label: "Product", spend: 3, tokens: 100, requests: 8 }]); + expect(Object.keys(result.teamResults[0].breakdown.entities)).toEqual(["team-a"]); + expect(result.teamResults[0].breakdown.entities["team-a"].metrics.spend).toBe(3); + expect(result.memberResults[0].breakdown.entities["member-a"].metrics.total_tokens).toBe(100); + expect(dailyUrls()).toHaveLength(1); + expect(dailyUrls()[0].searchParams.get("customer_member_ids")).toBe("member-a"); + expect(urls().some((url) => url.pathname === "/management/v1/customer_members")).toBe(false); + }); + + it("should reject mismatched ancestry before requesting any usage", async () => { + fetchMock.mockImplementation((input) => + Promise.resolve( + requestUrl(input).pathname.endsWith("/customer_teams/team-a") + ? json({ ...teams[0], parent_org_id: "another-org" }) + : serve(input), + ), + ); + await expect( + customerHierarchyUsageCall({ ...request, scope: { org: "org-1", team: "team-a", member: null } }), + ).rejects.toThrow("does not belong"); + expect(dailyUrls()).toHaveLength(0); + }); + + it("should never send an empty descendant filter and should retain unassigned usage", async () => { + fetchMock.mockImplementation((input) => + Promise.resolve(requestUrl(input).pathname === "/management/v1/customer_teams" ? json(list([])) : serve(input)), + ); + const result = await customerHierarchyUsageCall(request); + expect(dailyUrls().every((url) => url.searchParams.get("customer_org_ids") === "org-1")).toBe(true); + expect(result.teams).toEqual([ + { id: "", label: "Unassigned or removed team", spend: 4, tokens: 600, requests: 10 }, + ]); + expect(result.members).toEqual([ + { id: "", label: "Unassigned or removed member", spend: 4, tokens: 600, requests: 10 }, + ]); + }); + + it("should fail a snapshot when a later page fails instead of presenting partial totals", async () => { + fetchMock.mockImplementation((input) => { + const url = requestUrl(input); + return Promise.resolve( + url.searchParams.get("customer_org_ids") === "org-1" && url.searchParams.get("page") === "2" + ? json({ error: "Second page unavailable" }, 503) + : serve(input), + ); + }); + await expect(customerHierarchyUsageCall(request)).rejects.toThrow("Second page unavailable"); + }); + + it("should reject identifiers the comma-separated API cannot represent instead of broadening filters", async () => { + await expect( + customerHierarchyUsageCall({ ...request, scope: { org: "org-1,org-2", team: null, member: null } }), + ).rejects.toThrow(); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it("should chunk descendant filters and keep every request phase within four concurrent reads", async () => { + const manyTeams = Array.from({ length: 205 }, (_, index) => ({ ...teams[0], customer_team_id: `team-${index}` })); + let active = 0; + let maximumActive = 0; + fetchMock.mockImplementation(async (input) => { + active++; + maximumActive = Math.max(maximumActive, active); + await Promise.resolve(); + const url = requestUrl(input); + active--; + if (url.pathname === "/management/v1/customer_teams") { + const page = Number(url.searchParams.get("page")); + return json(list(manyTeams.slice((page - 1) * 100, page * 100), page, 3)); + } + if (url.pathname === "/management/v1/customer_members") return json(list([])); + if (url.pathname === "/management/v1/customer_teams/daily/activity") { + return json({ results: [], metadata: EMPTY_DAILY_ACTIVITY_METADATA }); + } + return serve(input); + }); + + await customerHierarchyUsageCall(request); + + expect(maximumActive).toBeLessThanOrEqual(4); + const teamFilters = dailyUrls() + .filter((url) => url.searchParams.has("customer_team_ids")) + .map((url) => url.searchParams.get("customer_team_ids")!.split(",")); + expect(teamFilters.map((ids) => ids.length)).toEqual([100, 100, 5]); + expect(new Set(teamFilters.flat()).size).toBe(manyTeams.length); + expect(dailyUrls().some((url) => url.searchParams.has("customer_member_ids"))).toBe(false); + }); +}); diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageApi.ts b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageApi.ts new file mode 100644 index 00000000000..3bfeeaebef5 --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageApi.ts @@ -0,0 +1,265 @@ +import { apiClient } from "@/components/networking"; +import { type DailyActivityAggregatedResponse } from "@/components/UsagePage/dailyActivityApi"; +import type { components } from "@/lib/http/schema"; +import i18n from "@/i18n/i18n"; +import { + mergeCustomerActivityPages, + customerEntityRankings, + customerKeyRankings, + customerEntityDailyData, + type CustomerUsageRow, +} from "./customerUsageData"; + +type Organization = components["schemas"]["Organization"]; +type Team = components["schemas"]["Team"]; +type Member = components["schemas"]["CustomerMember"]; +type Collection = "customer_organizations" | "customer_teams" | "customer_members"; +type UsageRow = Omit; + +export interface CustomerUsageEntity { + id: string; + label: string; + spend: number; + budgetLimit: number | null; + budgetResetAt: string | null; +} + +export interface CustomerHierarchyUsage extends ReturnType { + organization: CustomerUsageEntity; + team: CustomerUsageEntity | null; + member: CustomerUsageEntity | null; + teams: UsageRow[]; + members: UsageRow[]; + api_keys: UsageRow[]; + teamResults: ReturnType["results"]; + memberResults: ReturnType["results"]; + memberBudgets: CustomerUsageEntity[]; +} + +export interface CustomerUsageScope { + org: string; + team: string | null; + member: string | null; +} + +interface CustomerUsageRequest { + scope: CustomerUsageScope; + accessToken: string; + startDate: string; + endDate: string; + signal?: AbortSignal; +} + +interface ResolvedScope { + organization: Organization; + team: Team | null; + member: Member | null; +} + +const PREFIX = "/management/v1"; +const PAGE_SIZE = 100; +const MAX_PAGES = 1_000; +const MAX_CONCURRENT = 4; +const ID_FILTER: Record = { + customer_organizations: "customer_org_ids", + customer_teams: "customer_team_ids", + customer_members: "customer_member_ids", +}; + +const assertIdentifiers = (ids: readonly string[]) => { + if (ids.some((id) => !id || id.includes(","))) throw new Error(i18n.t("customerPortal.analytics.invalidIdentifier")); +}; + +const assertScope = (valid: boolean) => { + if (!valid) throw new Error(i18n.t("customerPortal.analytics.scopeMismatch")); +}; + +const batches = async (items: readonly T[], task: (item: T) => Promise): Promise => { + const results: R[] = []; + for (let start = 0; start < items.length; start += MAX_CONCURRENT) { + results.push(...(await Promise.all(items.slice(start, start + MAX_CONCURRENT).map(task)))); + } + return results; +}; + +const info = (resource: Collection, id: string, request: CustomerUsageRequest) => + apiClient.get(`${PREFIX}/${resource}/${encodeURIComponent(id)}`, { + accessToken: request.accessToken, + signal: request.signal, + }); + +const resolveScope = async (request: CustomerUsageRequest): Promise => { + const { scope } = request; + assertIdentifiers([scope.org, ...[scope.team, scope.member].filter((id): id is string => id !== null)]); + const [organization, requestedTeam, member] = await Promise.all([ + info("customer_organizations", scope.org, request), + scope.team ? info("customer_teams", scope.team, request) : null, + scope.member ? info("customer_members", scope.member, request) : null, + ]); + const team = requestedTeam ?? (member ? await info("customer_teams", member.parent_team_id, request) : null); + assertScope(organization.customer_org_id === scope.org); + assertScope(!team || team.parent_org_id === scope.org); + assertScope(!scope.team || team?.customer_team_id === scope.team); + assertScope(!member || member.customer_member_id === scope.member); + assertScope(!member || member.parent_team_id === team?.customer_team_id); + return { organization, team, member }; +}; + +const listAll = async ( + resource: Collection, + parent: Record, + request: CustomerUsageRequest, +): Promise => { + const rows: T[] = []; + for (let page = 1; page <= MAX_PAGES; page++) { + const result = await apiClient.get<{ data: T[]; meta: components["schemas"]["ListMeta"] }>( + `${PREFIX}/${resource}`, + { + accessToken: request.accessToken, + signal: request.signal, + query: { ...parent, page, page_size: PAGE_SIZE }, + }, + ); + rows.push(...result.data); + if (page >= result.meta.total_pages) return rows; + } + throw new Error(i18n.t("customerPortal.analytics.incompletePages")); +}; + +const descendants = async (scope: ResolvedScope, request: CustomerUsageRequest) => { + const teams = scope.team + ? [scope.team] + : await listAll("customer_teams", { parent_org_id: request.scope.org }, request); + if (teams.some((team) => team.parent_org_id !== request.scope.org)) + throw new Error(i18n.t("customerPortal.analytics.scopeMismatch")); + const memberLists = scope.member + ? [[scope.member]] + : await batches(teams, async (team) => { + const members = await listAll("customer_members", { parent_team_id: team.customer_team_id }, request); + if (members.some((member) => member.parent_team_id !== team.customer_team_id)) + throw new Error(i18n.t("customerPortal.analytics.scopeMismatch")); + return members; + }); + return { + teams: [...new Map(teams.map((team) => [team.customer_team_id, team])).values()], + members: [...new Map(memberLists.flat().map((member) => [member.customer_member_id, member])).values()], + }; +}; + +const activityPages = async (resource: Collection, ids: readonly string[], request: CustomerUsageRequest) => { + assertIdentifiers(ids); + const pages: DailyActivityAggregatedResponse[] = []; + for (let page = 1; page <= MAX_PAGES; page++) { + const result = await apiClient.get(`${PREFIX}/${resource}/daily/activity`, { + accessToken: request.accessToken, + signal: request.signal, + query: { + start_date: request.startDate, + end_date: request.endDate, + [ID_FILTER[resource]]: ids.join(","), + page, + page_size: PAGE_SIZE, + }, + }); + if (!result.metadata) throw new Error(i18n.t("customerPortal.analytics.incompletePages")); + pages.push(result); + if (!result.metadata.has_more && page >= result.metadata.total_pages) return pages; + } + throw new Error(i18n.t("customerPortal.analytics.incompletePages")); +}; + +interface ActivityJob { + dimension: "scope" | "teams" | "members"; + resource: Collection; + ids: readonly string[]; +} + +const chunkJobs = (resource: Collection, dimension: ActivityJob["dimension"], ids: readonly string[]): ActivityJob[] => + Array.from({ length: Math.ceil(ids.length / PAGE_SIZE) }, (_, index) => ({ + resource, + dimension, + ids: ids.slice(index * PAGE_SIZE, (index + 1) * PAGE_SIZE), + })); + +const entity = (id: string, record: Organization | Team | Member): CustomerUsageEntity => ({ + id, + label: record.alias || id, + spend: record.spend, + budgetLimit: ("effective_budget" in record ? record.effective_budget : record.budget)?.max_budget ?? null, + budgetResetAt: record.budget_reset_at ?? null, +}); + +const selectedResource = (scope: ResolvedScope): Collection => { + if (scope.member) return "customer_members"; + if (scope.team) return "customer_teams"; + return "customer_organizations"; +}; + +export const customerHierarchyUsageCall = async (request: CustomerUsageRequest): Promise => { + const selected = await resolveScope(request); + const children = await descendants(selected, request); + const resource = selectedResource(selected); + const selectedId = selected.member?.customer_member_id ?? selected.team?.customer_team_id ?? request.scope.org; + const jobs: ActivityJob[] = [ + { dimension: "scope", resource, ids: [selectedId] }, + ...(!selected.team + ? chunkJobs( + "customer_teams", + "teams", + children.teams.map((team) => team.customer_team_id), + ) + : []), + ...(!selected.member + ? chunkJobs( + "customer_members", + "members", + children.members.map((member) => member.customer_member_id), + ) + : []), + ]; + const activity = await batches(jobs, async (job) => ({ + dimension: job.dimension, + pages: await activityPages(job.resource, job.ids, request), + })); + const primary = mergeCustomerActivityPages( + activity.filter((result) => result.dimension === "scope").flatMap((result) => result.pages), + ); + const teamDays = selected.team + ? primary.results + : mergeCustomerActivityPages( + activity.filter((result) => result.dimension === "teams").flatMap((result) => result.pages), + ).results; + const memberDays = selected.member + ? primary.results + : mergeCustomerActivityPages( + activity.filter((result) => result.dimension === "members").flatMap((result) => result.pages), + ).results; + const teamLabels = new Map( + children.teams.map((team) => [team.customer_team_id, team.alias || team.customer_team_id]), + ); + const memberLabels = new Map( + children.members.map((member) => [member.customer_member_id, member.alias || member.customer_member_id]), + ); + return { + ...primary, + organization: entity(selected.organization.customer_org_id, selected.organization), + team: selected.team ? entity(selected.team.customer_team_id, selected.team) : null, + member: selected.member ? entity(selected.member.customer_member_id, selected.member) : null, + teams: customerEntityRankings(teamDays, teamLabels, primary.metadata, { + unassignedLabel: i18n.t("customerPortal.analytics.unassignedTeam"), + onlyEntityId: selected.member ? selected.team?.customer_team_id : undefined, + }), + members: customerEntityRankings(memberDays, memberLabels, primary.metadata, { + unassignedLabel: i18n.t("customerPortal.analytics.unassignedMember"), + }), + api_keys: customerKeyRankings(primary.results, i18n.t("customerPortal.analytics.unnamedKey")), + teamResults: customerEntityDailyData( + primary.results, + teamDays, + teamLabels, + selected.member ? selected.team?.customer_team_id : undefined, + ), + memberResults: customerEntityDailyData(primary.results, memberDays, memberLabels), + memberBudgets: children.members.map((member) => entity(member.customer_member_id, member)), + }; +}; diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageChartData.test.ts b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageChartData.test.ts new file mode 100644 index 00000000000..2d2a46178e4 --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageChartData.test.ts @@ -0,0 +1,185 @@ +import { describe, expect, it } from "vitest"; +import { + bucketTotals, + OTHER_COLOR, + type Series, + type UsageMetric, +} from "@/app/(dashboard)/usage/_components/components/overview/overviewData"; +import { stackedUsageColor } from "@/components/shared/charts"; +import { EMPTY_DAILY_ACTIVITY_METADATA } from "@/components/UsagePage/dailyActivityApi"; +import type { DailyData, MetricWithMetadata, SpendMetrics } from "@/components/UsagePage/types"; +import i18n from "@/i18n/i18n"; +import type { CustomerHierarchyUsage } from "./customerUsageApi"; +import { customerUsageSeries } from "./customerUsageChartData"; +import type { CustomerUsageDimension } from "./CustomerUsageRankings"; + +const metrics = (spend: number, tokens: number, requests: number): SpendMetrics => ({ + spend, + prompt_tokens: tokens, + completion_tokens: 0, + total_tokens: tokens, + api_requests: requests, + successful_requests: requests, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, +}); + +const entry = (spend: number, tokens: number, requests: number): MetricWithMetadata => ({ + metrics: metrics(spend, tokens, requests), + metadata: {}, + api_key_breakdown: {}, +}); + +const day = (date: string, total: SpendMetrics, breakdown: Partial = {}): DailyData => ({ + date, + metrics: total, + breakdown: { models: {}, model_groups: {}, providers: {}, mcp_servers: {}, api_keys: {}, entities: {}, ...breakdown }, +}); + +const snapshot = (): CustomerHierarchyUsage => ({ + organization: { id: "org", label: "Organization", spend: 0, budgetLimit: null, budgetResetAt: null }, + team: null, + member: null, + memberBudgets: [], + metadata: { ...EMPTY_DAILY_ACTIVITY_METADATA }, + results: [ + day("2026-11-01", metrics(100, 1_000, 100), { + models: { deployed: entry(100, 1_000, 100) }, + model_groups: { alpha: entry(60, 300, 20), beta: entry(40, 700, 80) }, + api_keys: { + "hash-a": { metrics: metrics(30, 800, 70), metadata: { key_alias: "Support", team_id: null } }, + "hash-b": { metrics: metrics(70, 200, 30), metadata: { key_alias: "Search", team_id: null } }, + }, + }), + day("2026-11-02", metrics(60, 900, 30)), + day("2026-11-03", metrics(50, 500, 40)), + ], + teamResults: [ + day("2026-11-01", metrics(75, 800, 60), { + entities: { "team-a": entry(25, 100, 40), "team-b": entry(50, 700, 20) }, + }), + day("2026-11-02", metrics(10, 90, 5), { entities: { "team-a": entry(10, 90, 5) } }), + day("2026-11-09", metrics(999, 999, 999), { entities: { "team-a": entry(999, 999, 999) } }), + ], + memberResults: [ + day("2026-11-01", metrics(80, 700, 80), { + entities: { "member-a": entry(20, 200, 30), "member-b": entry(60, 500, 50) }, + }), + ], + teams: [ + { id: "team-a", label: "Product", spend: 35, tokens: 190, requests: 45 }, + { id: "team-b", label: "Operations", spend: 50, tokens: 700, requests: 20 }, + ], + members: [ + { id: "member-a", label: "Member A", spend: 20, tokens: 200, requests: 30 }, + { id: "member-b", label: "Member B", spend: 60, tokens: 500, requests: 50 }, + ], + api_keys: [ + { id: "hash-a", label: "Support", spend: 30, tokens: 800, requests: 70 }, + { id: "hash-b", label: "Search", spend: 70, tokens: 200, requests: 30 }, + ], +}); + +const modelColors = new Map([ + ["alpha", stackedUsageColor(3)], + ["beta", stackedUsageColor(4)], +]); +const valueAt = (series: Series, label: string, date = "2026-11-01"): number => { + const index = series.labels.indexOf(label); + return Number(series.data.find((row) => row.date === date)?.[series.keys[index]] ?? 0); +}; +const colorOf = (series: Series, label: string): string => series.colors[series.labels.indexOf(label)]; +const other = () => i18n.t("usage.overview.otherSeries"); + +describe("customerUsageSeries", () => { + it.each<[CustomerUsageDimension, string, number]>([ + ["models", "alpha", 60], + ["api_keys", "Support", 30], + ["teams", "Product", 25], + ])("should use %s values without adding the other breakdowns", (dimension, label, value) => { + const series = customerUsageSeries(snapshot(), dimension, "spend", modelColors); + + expect(valueAt(series, label)).toBe(value); + expect(bucketTotals(series).get("2026-11-01")).toBe(100); + expect(series.labels).not.toContain("deployed"); + }); + + it.each<[UsageMetric, number, number, number]>([ + ["spend", 25, 50, 25], + ["tokens", 100, 700, 200], + ["requests", 40, 20, 40], + ])("should use %s for both named teams and the unattributed remainder", (metric, first, second, remainder) => { + const series = customerUsageSeries(snapshot(), "teams", metric, modelColors); + + expect(valueAt(series, "Product")).toBe(first); + expect(valueAt(series, "Operations")).toBe(second); + expect(valueAt(series, other())).toBe(remainder); + expect(colorOf(series, other())).toBe(OTHER_COLOR); + }); + + it("should preserve all selected-scope dates and totals when companion records are missing", () => { + const series = customerUsageSeries(snapshot(), "teams", "spend", modelColors); + + expect(series.data.map((row) => row.date)).toEqual(["2026-11-01", "2026-11-02", "2026-11-03"]); + expect(valueAt(series, "Product", "2026-11-02")).toBe(10); + expect(valueAt(series, other(), "2026-11-02")).toBe(50); + expect(valueAt(series, "Product", "2026-11-03")).toBe(0); + expect(valueAt(series, other(), "2026-11-03")).toBe(50); + expect([...bucketTotals(series).values()]).toEqual([100, 60, 50]); + }); + + it("should include removed and unnamed descendants in Other rather than a misleading named bucket", () => { + const data = snapshot(); + const removed = day("2026-11-01", metrics(100, 1_000, 100), { + entities: { "team-a": entry(25, 100, 40), removed: entry(10, 150, 10), "": entry(5, 50, 5) }, + }); + const series = customerUsageSeries({ ...data, teamResults: [removed] }, "teams", "tokens", modelColors); + + expect(series.labels).toEqual(["Product", other()]); + expect(valueAt(series, "Product")).toBe(100); + expect(valueAt(series, other())).toBe(900); + }); + + it("should keep spend-ranked palette colors attached to API keys when the selected metric changes", () => { + const spend = customerUsageSeries(snapshot(), "api_keys", "spend", modelColors); + const tokens = customerUsageSeries(snapshot(), "api_keys", "tokens", modelColors); + + expect(spend.labels.slice(0, 2)).toEqual(["Search", "Support"]); + expect(tokens.labels.slice(0, 2)).toEqual(["Support", "Search"]); + expect(colorOf(spend, "Search")).toBe(stackedUsageColor(0)); + expect(colorOf(tokens, "Search")).toBe(stackedUsageColor(0)); + expect(colorOf(spend, "Support")).toBe(stackedUsageColor(1)); + expect(colorOf(tokens, "Support")).toBe(stackedUsageColor(1)); + }); + + it("should use key aliases and hide hashes even when a label is missing or repeats the hash", () => { + const data = snapshot(); + const series = customerUsageSeries( + { ...data, api_keys: [{ ...data.api_keys[0], label: " hash-a " }] }, + "api_keys", + "spend", + modelColors, + ); + + expect(series.labels).not.toContain("hash-a"); + expect(series.labels).not.toContain("hash-b"); + expect(series.labels.filter((label) => label === i18n.t("customerPortal.analytics.unnamedKey"))).toHaveLength(2); + expect(bucketTotals(series).get("2026-11-01")).toBe(100); + }); + + it("should preserve model colors and fall back to the models breakdown when groups are absent", () => { + const data = snapshot(); + const color = stackedUsageColor(5); + const series = customerUsageSeries( + { ...data, results: [day("2026-11-01", metrics(7, 70, 2), { models: { deployed: entry(7, 70, 2) } })] }, + "models", + "tokens", + new Map([["deployed", color]]), + ); + + expect(series.labels).toEqual(["deployed"]); + expect(valueAt(series, "deployed")).toBe(70); + expect(colorOf(series, "deployed")).toBe(color); + }); +}); diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageChartData.ts b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageChartData.ts new file mode 100644 index 00000000000..804d50bfbe8 --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageChartData.ts @@ -0,0 +1,71 @@ +import { + rankValue, + rollUpBreakdown, + seriesBy, + type Series, + type UsageMetric, +} from "@/app/(dashboard)/usage/_components/components/overview/overviewData"; +import { stackedUsageColor } from "@/components/shared/charts"; +import type { DailyData } from "@/components/UsagePage/types"; +import i18n from "@/i18n/i18n"; +import type { CustomerHierarchyUsage } from "./customerUsageApi"; +import type { CustomerUsageDimension } from "./CustomerUsageRankings"; + +const TOP_SERIES = 8; + +const entityDays = ( + primary: readonly DailyData[], + entities: readonly DailyData[], + labels: ReadonlyMap, +): DailyData[] => { + const byDate = new Map(entities.map((day) => [day.date, day])); + return primary.map((day) => ({ + ...day, + breakdown: { + ...day.breakdown, + model_groups: Object.fromEntries( + Object.entries(byDate.get(day.date)?.breakdown.entities ?? {}).filter(([id]) => id && labels.has(id)), + ), + }, + })); +}; + +export const customerUsageSeries = ( + data: CustomerHierarchyUsage, + dimension: CustomerUsageDimension, + metric: UsageMetric, + colorByModel: ReadonlyMap, +): Series => { + if (dimension === "models") { + const days = data.results.map((day) => ({ + ...day, + breakdown: { + ...day.breakdown, + model_groups: + Object.keys(day.breakdown.model_groups).length > 0 ? day.breakdown.model_groups : day.breakdown.models, + }, + })); + return seriesBy(days, "model_groups", metric, { top: TOP_SERIES, colorByKey: colorByModel }); + } + + const unnamedKey = i18n.t("customerPortal.analytics.unnamedKey"); + const rows = data[dimension]; + const labels = new Map( + rows + .filter((row) => row.id) + .map((row) => { + const label = row.label.trim(); + return [row.id, dimension === "api_keys" && (!label || label === row.id) ? unnamedKey : label]; + }), + ); + const days = dimension === "api_keys" ? data.results : entityDays(data.results, data.teamResults, labels); + const breakdown = dimension === "api_keys" ? "api_keys" : "model_groups"; + const ranking = rollUpBreakdown(days, breakdown).sort((a, b) => b.spend - a.spend || a.key.localeCompare(b.key)); + const colors = new Map(ranking.map((row, index) => [row.key, stackedUsageColor(index)])); + const series = seriesBy(days, breakdown, metric, { top: TOP_SERIES, colorByKey: colors }); + const namedCount = Math.min(TOP_SERIES, ranking.filter((row) => rankValue(row, metric) > 0).length); + return { + ...series, + labels: series.labels.map((id, index) => (index < namedCount ? labels.get(id) ?? unnamedKey : id)), + }; +}; diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageData.test.ts b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageData.test.ts new file mode 100644 index 00000000000..c8e6d638c79 --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageData.test.ts @@ -0,0 +1,104 @@ +import { describe, expect, it } from "vitest"; +import type { DailyData, MetricWithMetadata, SpendMetrics } from "@/components/UsagePage/types"; +import { rankCustomerModels, rankCustomerUsage, type CustomerUsageRow } from "./customerUsageData"; + +const metrics = (spend: number, tokens: number, requests: number): SpendMetrics => ({ + spend, + prompt_tokens: tokens, + completion_tokens: 0, + total_tokens: tokens, + api_requests: requests, + successful_requests: requests, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, +}); + +const model = (spend: number, tokens: number, requests: number): MetricWithMetadata => ({ + metrics: metrics(spend, tokens, requests), + metadata: {}, + api_key_breakdown: {}, +}); + +const day = ( + date: string, + modelGroups: Record, + models: Record = {}, +): DailyData => ({ + date, + metrics: metrics(0, 0, 0), + breakdown: { + models, + model_groups: modelGroups, + providers: {}, + mcp_servers: {}, + api_keys: {}, + entities: {}, + }, +}); + +const totals = { spend: 10, tokens: 1_000, requests: 20 }; + +describe("rankCustomerUsage", () => { + it("should sort by the selected metric and calculate shares against overall customer totals", () => { + const rows = [ + { id: "assistant", label: "Assistant", spend: 3, tokens: 100, requests: 8 }, + { id: "search", label: "Search", spend: 1, tokens: 500, requests: 2 }, + ]; + + expect(rankCustomerUsage(rows, "spend", totals).map((row) => [row.label, row.share])).toEqual([ + ["Assistant", 0.3], + ["Search", 0.1], + ]); + expect(rankCustomerUsage(rows, "tokens", totals).map((row) => [row.label, row.share])).toEqual([ + ["Search", 0.5], + ["Assistant", 0.1], + ]); + expect(rankCustomerUsage(rows, "requests", totals).map((row) => [row.label, row.share])).toEqual([ + ["Assistant", 0.4], + ["Search", 0.1], + ]); + }); + + it("should preserve zero-cost token usage while removing rows with no activity", () => { + const freeUsage = { id: "assistant", label: "Assistant", spend: 0, tokens: 200, requests: 3 }; + const inactive = { id: "inactive", label: "Inactive", spend: 0, tokens: 0, requests: 0 }; + + expect(rankCustomerUsage([inactive, freeUsage], "tokens", totals)).toEqual([{ ...freeUsage, share: 0.2 }]); + expect(rankCustomerUsage([freeUsage], "spend", totals)).toEqual([{ ...freeUsage, share: 0 }]); + }); + + it("should use stable label and id order for ties without modifying the source rows", () => { + const rows: readonly Omit[] = Object.freeze([ + { id: "b", label: "Search", spend: 1, tokens: 100, requests: 1 }, + { id: "c", label: "Assistant", spend: 1, tokens: 100, requests: 1 }, + { id: "a", label: "Assistant", spend: 1, tokens: 100, requests: 1 }, + ]); + expect(rankCustomerUsage(rows, "tokens", totals).map((row) => row.id)).toEqual(["a", "c", "b"]); + expect(rows.map((row) => row.id)).toEqual(["b", "c", "a"]); + }); + + it("should return finite zero shares when the overall metric is zero", () => { + const rows = [{ id: "assistant", label: "Assistant", spend: 0, tokens: 200, requests: 1 }]; + expect(rankCustomerUsage(rows, "spend", { spend: 0, tokens: 200, requests: 1 })[0].share).toBe(0); + expect(rankCustomerUsage([], "spend", totals)).toEqual([]); + }); +}); + +describe("rankCustomerModels", () => { + it("should sum models across days without adding model and model-group totals together", () => { + const days = [ + day("2026-11-01", { assistant: model(2, 100, 3) }, { assistant: model(2, 100, 3) }), + day("2026-11-02", { assistant: model(1, 200, 2), search: model(0, 400, 1) }), + day("2026-11-03", {}, { assistant: model(1, 100, 1) }), + ]; + expect(rankCustomerModels(days, "tokens", totals)).toEqual([ + { id: "assistant", label: "assistant", spend: 4, tokens: 400, requests: 6, share: 0.4 }, + { id: "search", label: "search", spend: 0, tokens: 400, requests: 1, share: 0.4 }, + ]); + }); + + it("should return no model rows for an empty date range", () => { + expect(rankCustomerModels([], "spend", totals)).toEqual([]); + }); +}); diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageData.ts b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageData.ts new file mode 100644 index 00000000000..afa0d3c7767 --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/customerUsageData.ts @@ -0,0 +1,262 @@ +import type { DailyData, SpendMetrics, BreakdownMetrics, MetricWithMetadata } from "@/components/UsagePage/types"; +import { + EMPTY_DAILY_ACTIVITY_METADATA, + toDailyData, + type DailyActivityAggregatedResponse, + type DailyActivityMetadata, +} from "@/components/UsagePage/dailyActivityApi"; +import { + rollUpBreakdown, + type UsageMetric, +} from "@/app/(dashboard)/usage/_components/components/overview/overviewData"; + +export interface CustomerUsageTotals { + spend: number; + tokens: number; + requests: number; +} + +export interface CustomerUsageRow extends CustomerUsageTotals { + id: string; + label: string; + share: number; +} + +export const rankCustomerUsage = ( + rows: readonly Omit[], + metric: UsageMetric, + totals: CustomerUsageTotals, +): CustomerUsageRow[] => + rows + .filter((row) => row.spend > 0 || row.tokens > 0 || row.requests > 0) + .map((row) => ({ ...row, share: totals[metric] > 0 ? row[metric] / totals[metric] : 0 })) + .sort((a, b) => b[metric] - a[metric] || a.label.localeCompare(b.label) || a.id.localeCompare(b.id)); + +export const rankCustomerModels = ( + days: readonly DailyData[], + metric: UsageMetric, + totals: CustomerUsageTotals, +): CustomerUsageRow[] => { + const withModelGroups = days.map((day) => ({ + ...day, + breakdown: { + ...day.breakdown, + model_groups: + Object.keys(day.breakdown.model_groups).length > 0 ? day.breakdown.model_groups : day.breakdown.models, + }, + })); + return rankCustomerUsage( + rollUpBreakdown(withModelGroups, "model_groups").map((model) => ({ + id: model.key, + label: model.key, + spend: model.spend, + tokens: model.tokens, + requests: model.requests, + })), + metric, + totals, + ); +}; + +const ZERO_METRICS: SpendMetrics = { + spend: 0, + prompt_tokens: 0, + completion_tokens: 0, + total_tokens: 0, + api_requests: 0, + successful_requests: 0, + failed_requests: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, +}; + +const sumMetrics = (values: readonly SpendMetrics[]): SpendMetrics => + values.reduce( + (total, value) => ({ + ...total, + ...Object.fromEntries( + (Object.keys(value) as (keyof SpendMetrics)[]).map((field) => [ + field, + (total[field] ?? 0) + (value[field] ?? 0), + ]), + ), + }), + ZERO_METRICS, + ); + +const mergeKeys = ( + left: BreakdownMetrics["api_keys"], + right: BreakdownMetrics["api_keys"], +): BreakdownMetrics["api_keys"] => { + const result = { ...left }; + for (const [id, entry] of Object.entries(right)) { + const previous = result[id]; + result[id] = { + ...entry, + metrics: sumMetrics([...(previous ? [previous.metrics] : []), entry.metrics]), + metadata: { + ...previous?.metadata, + ...entry.metadata, + key_alias: entry.metadata.key_alias ?? previous?.metadata.key_alias ?? null, + }, + }; + } + return result; +}; + +const mergeMetrics = ( + left: Record, + right: Record, +): Record => { + const result = { ...left }; + for (const [id, entry] of Object.entries(right)) { + const previous = result[id]; + result[id] = { + ...entry, + metrics: sumMetrics([...(previous ? [previous.metrics] : []), entry.metrics]), + metadata: { ...previous?.metadata, ...entry.metadata }, + api_key_breakdown: mergeKeys(previous?.api_key_breakdown ?? {}, entry.api_key_breakdown), + }; + } + return result; +}; + +const mergeBreakdowns = (left: BreakdownMetrics, right: BreakdownMetrics): BreakdownMetrics => ({ + models: mergeMetrics(left.models, right.models), + model_groups: mergeMetrics(left.model_groups, right.model_groups), + providers: mergeMetrics(left.providers, right.providers), + mcp_servers: mergeMetrics(left.mcp_servers, right.mcp_servers), + endpoints: mergeMetrics(left.endpoints ?? {}, right.endpoints ?? {}), + entities: mergeMetrics(left.entities, right.entities), + api_keys: mergeKeys(left.api_keys, right.api_keys), +}); + +export const mergeCustomerActivityPages = ( + pages: readonly DailyActivityAggregatedResponse[], +): { results: DailyData[]; metadata: DailyActivityMetadata } => { + const days = new Map(); + for (const day of pages.flatMap(toDailyData)) { + const previous = days.get(day.date); + days.set( + day.date, + previous + ? { + date: day.date, + metrics: sumMetrics([previous.metrics, day.metrics]), + breakdown: mergeBreakdowns(previous.breakdown, day.breakdown), + } + : day, + ); + } + const results = [...days.values()].sort((a, b) => a.date.localeCompare(b.date)); + const total = sumMetrics(results.map((day) => day.metrics)); + return { + results, + metadata: { + ...EMPTY_DAILY_ACTIVITY_METADATA, + total_spend: total.spend, + total_flat_cost: total.flat_cost ?? 0, + total_ptu_hours: total.ptu_hours ?? 0, + total_prompt_tokens: total.prompt_tokens, + total_completion_tokens: total.completion_tokens, + total_tokens: total.total_tokens, + total_api_requests: total.api_requests, + total_successful_requests: total.successful_requests, + total_failed_requests: total.failed_requests, + total_cache_read_input_tokens: total.cache_read_input_tokens, + total_cache_creation_input_tokens: total.cache_creation_input_tokens, + total_compression_saved_tokens: total.compression_saved_tokens ?? 0, + total_compression_savings_spend: total.compression_savings_spend ?? 0, + total_prompt_caching_savings_spend: total.prompt_caching_savings_spend ?? 0, + total_gateway_injected_caching_savings_spend: total.gateway_injected_caching_savings_spend ?? 0, + total_autorouter_savings_spend: total.autorouter_savings_spend ?? 0, + total_response_time_ms: total.total_response_time_ms ?? 0, + total_timed_requests: total.timed_requests ?? 0, + total_api_keys: new Set(results.flatMap((day) => Object.keys(day.breakdown.api_keys))).size, + }, + }; +}; + +const usageTotals = (value: SpendMetrics): CustomerUsageTotals => ({ + spend: value.spend, + tokens: value.total_tokens, + requests: value.api_requests, +}); +const addTotals = (left: CustomerUsageTotals, right: CustomerUsageTotals): CustomerUsageTotals => ({ + spend: left.spend + right.spend, + tokens: left.tokens + right.tokens, + requests: left.requests + right.requests, +}); + +export const customerKeyRankings = ( + days: readonly DailyData[], + unnamedKey: string, +): Omit[] => { + const keys = new Map>(); + for (const [id, entry] of days.flatMap((day) => Object.entries(day.breakdown.api_keys))) { + const previous = keys.get(id); + keys.set(id, { + id, + label: entry.metadata.key_alias?.trim() || previous?.label || unnamedKey, + ...addTotals(previous ?? { spend: 0, tokens: 0, requests: 0 }, usageTotals(entry.metrics)), + }); + } + return [...keys.values()]; +}; + +interface EntityRankingOptions { + unassignedLabel: string; + onlyEntityId?: string; +} + +export const customerEntityDailyData = ( + primary: readonly DailyData[], + days: readonly DailyData[], + labels: ReadonlyMap, + onlyEntityId?: string, +): DailyData[] => { + const byDate = new Map(days.map((day) => [day.date, day])); + return primary.map((day) => { + const entities = onlyEntityId + ? { + [onlyEntityId]: { + metrics: day.metrics, + metadata: { alias: labels.get(onlyEntityId) }, + api_key_breakdown: day.breakdown.api_keys, + }, + } + : Object.fromEntries( + Object.entries(byDate.get(day.date)?.breakdown.entities ?? {}) + .filter(([id]) => labels.has(id)) + .map(([id, entry]) => [id, { ...entry, metadata: { ...entry.metadata, alias: labels.get(id) } }]), + ); + return { ...day, breakdown: { ...day.breakdown, entities } }; + }); +}; + +export const customerEntityRankings = ( + days: readonly DailyData[], + labels: ReadonlyMap, + metadata: DailyActivityMetadata, + options: EntityRankingOptions, +): Omit[] => { + const entries = new Map(); + for (const [id, entry] of days.flatMap((day) => Object.entries(day.breakdown.entities))) { + const identity = options.onlyEntityId ?? id; + if (labels.has(identity)) + entries.set( + identity, + addTotals(entries.get(identity) ?? { spend: 0, tokens: 0, requests: 0 }, usageTotals(entry.metrics)), + ); + } + const rows = [...entries].map(([id, total]) => ({ id, label: labels.get(id)!, ...total })); + const assigned = rows.reduce(addTotals, { spend: 0, tokens: 0, requests: 0 }); + const unassigned = { + spend: Math.max(0, metadata.total_spend - assigned.spend), + tokens: Math.max(0, metadata.total_tokens - assigned.tokens), + requests: Math.max(0, metadata.total_api_requests - assigned.requests), + }; + return unassigned.spend > 0.000001 || unassigned.tokens > 0 || unassigned.requests > 0 + ? [...rows, { id: "", label: options.unassignedLabel, ...unassigned }] + : rows; +}; diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useCustomerHierarchyUsage.ts b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useCustomerHierarchyUsage.ts new file mode 100644 index 00000000000..05fa467d0e5 --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/_components/useCustomerHierarchyUsage.ts @@ -0,0 +1,122 @@ +import { useMemo } from "react"; +import { skipToken, useQuery } from "@tanstack/react-query"; +import { parseAsInteger, parseAsString, parseAsStringLiteral, useQueryStates } from "nuqs"; +import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; +import { + dateRangeFromParams, + dateRangePatch, + defaultUsageDateRange, + usageUrlParsers, +} from "@/app/(dashboard)/usage/_components/usageUrlState"; +import { + modelColorMap, + overviewTotals, + type Granularity, + type UsageMetric, +} from "@/app/(dashboard)/usage/_components/components/overview/overviewData"; +import { formatDate } from "@/components/networking"; +import type { DateRangePickerValue } from "@/components/shared/date_picker_types"; +import { EMPTY_DAILY_ACTIVITY_METADATA } from "@/components/UsagePage/dailyActivityApi"; +import type { DailyData } from "@/components/UsagePage/types"; +import { customerHierarchyUsageCall, type CustomerUsageScope } from "./customerUsageApi"; +import { rankCustomerModels, rankCustomerUsage } from "./customerUsageData"; +import { CUSTOMER_USAGE_DIMENSIONS, type CustomerUsageDimension } from "./CustomerUsageRankings"; + +const usageParsers = { + org: parseAsString, + team: parseAsString, + member: parseAsString, + metric: parseAsStringLiteral(["spend", "tokens", "requests"] as const).withDefault("spend"), + breakdown: parseAsStringLiteral(CUSTOMER_USAGE_DIMENSIONS).withDefault("models"), + bucket: parseAsStringLiteral(["day", "week"] as const).withDefault("day"), + search: parseAsString.withDefault(""), + usagePage: parseAsInteger.withDefault(1), + range: usageUrlParsers.range, + from: usageUrlParsers.from, + to: usageUrlParsers.to, +}; +const EMPTY_RESULTS: readonly DailyData[] = []; +const ORGANIZATION_PATCH = { team: null, member: null, search: null, usagePage: null }; +const TEAM_PATCH = { member: null, search: null, usagePage: null }; + +const usageRequest = ( + scope: CustomerUsageScope | null, + accessToken: string | null, + startDate: string | undefined, + endDate: string | undefined, +) => { + if (!scope || !accessToken) return null; + if (!startDate || !endDate) return null; + return { scope, accessToken, startDate, endDate }; +}; + +export default function useCustomerHierarchyUsage() { + const { accessToken } = useAuthorized(); + const [params, setParams] = useQueryStates(usageParsers); + const defaultRange = useMemo(() => defaultUsageDateRange("month_to_date"), []); + const dateValue = useMemo(() => dateRangeFromParams(params) ?? defaultRange.value, [params, defaultRange]); + const preset = params.range ?? (params.from === null && params.to === null ? defaultRange.preset : null); + const startDate = dateValue.from ? formatDate(dateValue.from) : undefined; + const endDate = dateValue.to ? formatDate(dateValue.to) : undefined; + const scope = params.org ? { org: params.org, team: params.team, member: params.member } : null; + const request = usageRequest(scope, accessToken, startDate, endDate); + const usage = useQuery({ + queryKey: ["customer-hierarchy", "usage", params.org, params.team, params.member, startDate, endDate], + queryFn: request ? ({ signal }) => customerHierarchyUsageCall({ ...request, signal }) : skipToken, + }); + const data = usage.data; + const results = data?.results ?? EMPTY_RESULTS; + const totals = useMemo(() => overviewTotals(data?.metadata ?? EMPTY_DAILY_ACTIVITY_METADATA, null), [data]); + const modelColors = useMemo(() => modelColorMap(results), [results]); + const rows = useMemo( + () => + params.breakdown === "models" + ? rankCustomerModels(results, params.metric, totals) + : rankCustomerUsage(data?.[params.breakdown] ?? [], params.metric, totals), + [data, results, params.breakdown, params.metric, totals], + ); + const drillDown = (dimension: "teams" | "members", id: string) => { + const patch = { + ...(dimension === "teams" ? { team: id, member: null } : { member: id }), + search: null, + usagePage: null, + breakdown: null, + }; + void setParams(patch, { history: "push" }); + }; + const setDateRange = (value: DateRangePickerValue, selectedPreset: string | null) => { + const { range, from, to } = dateRangePatch(value, selectedPreset); + const patch = { range, from, to, usagePage: null }; + void setParams(patch); + }; + const setDimension = (breakdown: CustomerUsageDimension) => + void setParams({ breakdown, search: null, usagePage: null }); + const selectOrganization = () => void setParams(ORGANIZATION_PATCH, { history: "push" }); + const selectTeam = () => void setParams(TEAM_PATCH, { history: "push" }); + const setMetric = (metric: UsageMetric) => void setParams({ metric, usagePage: null }); + const setBucket = (bucket: Granularity) => void setParams({ bucket }); + const setSearch = (search: string) => void setParams({ search, usagePage: null }); + const setPage = (usagePage: number) => void setParams({ usagePage }); + + return { + params, + scope, + dateValue, + preset, + usage, + data, + results, + totals, + modelColors, + rows, + drillDown, + setDateRange, + setDimension, + selectOrganization, + selectTeam, + setMetric, + setBucket, + setSearch, + setPage, + }; +} diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/usage/page.integration.test.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/usage/page.integration.test.tsx new file mode 100644 index 00000000000..852d85f3320 --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/usage/page.integration.test.tsx @@ -0,0 +1,379 @@ +import userEvent from "@testing-library/user-event"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { registerBaseUrlGetter } from "@/lib/http/runtime"; +import { EMPTY_DAILY_ACTIVITY_METADATA } from "@/components/UsagePage/dailyActivityApi"; +import type { SpendMetrics } from "@/components/UsagePage/types"; +import { + act, + chooseSelectOption, + fireEvent, + renderWithProviders, + screen, + testQueryClient, + waitFor, + within, +} from "../../../../../tests/test-utils"; +import CustomerHierarchyUsagePage from "./page"; + +const metrics = (spend: number, tokens: number, requests: number): SpendMetrics => ({ + spend, + api_requests: requests, + successful_requests: requests, + failed_requests: 0, + total_tokens: tokens, + prompt_tokens: tokens, + completion_tokens: 0, + cache_read_input_tokens: 0, + cache_creation_input_tokens: 0, +}); +const audit = { + created_at: "2026-11-01T00:00:00Z", + created_by: "admin", + updated_at: "2026-11-01T00:00:00Z", + updated_by: "admin", + spend: 40, +}; +const organization = { + ...audit, + customer_org_id: "org-1", + alias: "Customer organization", + budget: { max_budget: 100 }, + default_budgets: [], +}; +const teams = [ + { ...audit, customer_team_id: "team-a", parent_org_id: "org-1", alias: "Product", budget: null }, + { ...audit, customer_team_id: "team-b", parent_org_id: "org-1", alias: "Operations", budget: null }, +]; +const members = [ + { + ...audit, + customer_member_id: "member-a", + parent_team_id: "team-a", + alias: "Member One", + effective_budget: { max_budget: 10 }, + }, + { + ...audit, + customer_member_id: "member-b", + parent_team_id: "team-b", + alias: "Member Two", + effective_budget: { max_budget: 20 }, + }, +]; +const records = [ + { + organization: "org-1", + team: "team-a", + member: "member-a", + model: "model-a", + key: "hash-a", + alias: "Assistant key", + metrics: metrics(12, 100, 2), + date: "2026-11-01", + }, + { + organization: "org-1", + team: "team-b", + member: "member-b", + model: "model-b", + key: "hash-b", + alias: "Search key", + metrics: metrics(3, 600, 5), + date: "2026-11-02", + }, +]; +const fetchMock = vi.fn(); +const json = (body: unknown, status = 200) => + new Response(JSON.stringify(body), { status, headers: { "Content-Type": "application/json" } }); +const requestUrl = (input: Parameters[0]) => + new URL(input instanceof Request ? input.url : String(input), "http://localhost"); +const requestedUrls = () => fetchMock.mock.calls.map(([input]) => requestUrl(input)); +const initialUrl = "?org=org-1&from=2026-11-01&to=2026-11-03&source=portal"; +const tableRows = (name: string) => within(screen.getByRole("table", { name })).getAllByRole("row").slice(1); +const list = (data: T[]) => ({ data, meta: { total_count: data.length, page: 1, page_size: 100, total_pages: 1 } }); +const selectBreakdown = (user: ReturnType, name: string) => + chooseSelectOption(user, screen.getByRole("combobox", { name: "Break down by" }), name); +const activityDimension = (resource: string | undefined): "organization" | "team" | "member" => { + if (resource === "customer_organizations") return "organization"; + if (resource === "customer_teams") return "team"; + return "member"; +}; +const filterFor = (resource: string | undefined) => + ({ + organization: "customer_org_ids", + team: "customer_team_ids", + member: "customer_member_ids", + })[activityDimension(resource)]; +const activityUrl = (resource: string, id: string) => { + const filter = filterFor(resource); + return requestedUrls().findLast( + (url) => url.pathname === `/management/v1/${resource}/daily/activity` && url.searchParams.get(filter) === id, + ); +}; +const serve = (input: Parameters[0]): Response => { + const url = requestUrl(input); + if (url.pathname.endsWith("/daily/activity")) { + const resource = url.pathname.split("/").at(-3); + const dimension = activityDimension(resource); + const filter = filterFor(resource); + const ids = url.searchParams.get(filter)?.split(",") ?? []; + const selected = records.filter((record) => ids.includes(record[dimension])); + const results = selected.map((record) => { + const key = { metrics: record.metrics, metadata: { key_alias: record.alias, team_id: null } }; + const entry = { metrics: record.metrics, metadata: {}, api_key_breakdown: { [record.key]: key } }; + return { + date: record.date, + metrics: record.metrics, + breakdown: { + model_groups: { [record.model]: entry }, + api_keys: { [record.key]: key }, + entities: { [record[dimension]]: entry }, + }, + }; + }); + return json({ results, metadata: EMPTY_DAILY_ACTIVITY_METADATA }); + } + if (url.pathname === "/management/v1/customer_organizations/org-1") return json(organization); + if (url.pathname === "/management/v1/customer_teams") return json(list(teams)); + if (url.pathname === "/management/v1/customer_members") + return json(list(members.filter((member) => member.parent_team_id === url.searchParams.get("parent_team_id")))); + const team = teams.find((team) => url.pathname.endsWith(`/customer_teams/${team.customer_team_id}`)); + if (team) return json(team); + const member = members.find((member) => url.pathname.endsWith(`/customer_members/${member.customer_member_id}`)); + if (member) return json(member); + throw new Error(`Unexpected usage request ${url.pathname}`); +}; + +describe("CustomerHierarchyUsagePage", () => { + beforeEach(() => { + testQueryClient.clear(); + registerBaseUrlGetter(() => "http://localhost"); + fetchMock.mockReset(); + fetchMock.mockImplementation((input) => Promise.resolve(serve(input))); + vi.stubGlobal("fetch", fetchMock); + vi.stubGlobal("requestIdleCallback", (callback: IdleRequestCallback) => { + callback({ didTimeout: false, timeRemaining: () => 50 }); + return 1; + }); + }); + + afterEach(() => { + testQueryClient.clear(); + vi.unstubAllGlobals(); + }); + + it("should render organization totals and model rankings from the selected UTC date range", async () => { + renderWithProviders(, { searchParams: initialUrl }); + + expect(await screen.findByRole("heading", { name: "Customer organization" })).toBeInTheDocument(); + expect(screen.getByText("$15.00")).toBeInTheDocument(); + expect(screen.getByTitle("700 tokens")).toBeInTheDocument(); + expect(screen.getByRole("heading", { name: "Usage over time" })).toBeInTheDocument(); + expect(tableRows("Models")[0]).toHaveTextContent("model-a"); + expect(within(tableRows("Models")[0]).getByText("80%")).toBeInTheDocument(); + expect(tableRows("Models")[1]).toHaveTextContent("model-b"); + expect(activityUrl("customer_organizations", "org-1")?.pathname).toBe( + "/management/v1/customer_organizations/daily/activity", + ); + expect(activityUrl("customer_organizations", "org-1")?.searchParams.get("start_date")).toBe("2026-11-01"); + expect(activityUrl("customer_organizations", "org-1")?.searchParams.get("end_date")).toBe("2026-11-03"); + }); + + it("should apply the global metric to sorting and shares across models, keys and teams", async () => { + const user = userEvent.setup(); + renderWithProviders(, { searchParams: initialUrl }); + await screen.findByRole("table", { name: "Models" }); + + await user.click(screen.getByRole("radio", { name: "Tokens" })); + expect(tableRows("Models")[0]).toHaveTextContent("model-b"); + expect(within(tableRows("Models")[0]).getByText("85.7%")).toBeInTheDocument(); + + for (const [tab, leader] of [ + ["API keys", "Search key"], + ["Teams", "Operations"], + ]) { + await selectBreakdown(user, tab); + expect(tableRows(tab)[0]).toHaveTextContent(leader); + expect(within(tableRows(tab)[0]).getByText("85.7%")).toBeInTheDocument(); + expect(within(screen.getByRole("region", { name: "Usage over time" })).getByText(leader)).toBeInTheDocument(); + } + + await user.click(screen.getByRole("radio", { name: "Requests" })); + expect(tableRows("Teams")[0]).toHaveTextContent("Operations"); + expect(within(tableRows("Teams")[0]).getByText("71.4%")).toBeInTheDocument(); + expect(screen.getByRole("columnheader", { name: "Requests" })).toHaveAttribute("aria-sort", "descending"); + }); + + it("should drill into customer teams and members and retain dates and unrelated URL state", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { searchParams: initialUrl, onUrlUpdate }); + await screen.findByRole("heading", { name: "Customer organization" }); + + await selectBreakdown(user, "Teams"); + await user.click(within(screen.getByRole("table", { name: "Teams" })).getByRole("button", { name: "Product" })); + expect(await screen.findByRole("heading", { name: "Product" })).toBeInTheDocument(); + expect(activityUrl("customer_teams", "team-a")?.pathname).toBe("/management/v1/customer_teams/daily/activity"); + expect(activityUrl("customer_teams", "team-a")?.searchParams.get("customer_team_ids")).toBe("team-a"); + + await user.click( + within(screen.getByRole("table", { name: "Member spend" })).getByRole("button", { name: "Member One" }), + ); + expect(await screen.findByRole("heading", { name: "Member One" })).toBeInTheDocument(); + const memberUrl = activityUrl("customer_members", "member-a")!; + expect(memberUrl.pathname).toBe("/management/v1/customer_members/daily/activity"); + expect(memberUrl.searchParams.get("customer_member_ids")).toBe("member-a"); + expect(memberUrl.searchParams.has("customer_org_ids")).toBe(false); + expect(memberUrl.searchParams.has("customer_team_ids")).toBe(false); + expect(memberUrl.searchParams.get("start_date")).toBe("2026-11-01"); + expect(memberUrl.searchParams.get("end_date")).toBe("2026-11-03"); + expect(onUrlUpdate).toHaveBeenLastCalledWith( + expect.objectContaining({ + queryString: "?org=org-1&from=2026-11-01&to=2026-11-03&source=portal&team=team-a&member=member-a", + }), + ); + + await user.click(screen.getByRole("button", { name: "Customer organization" })); + expect(await screen.findByRole("heading", { name: "Customer organization" })).toBeInTheDocument(); + expect(onUrlUpdate).toHaveBeenLastCalledWith(expect.objectContaining({ queryString: initialUrl })); + }); + + it("should fall back to Models for a removed Members breakdown and offer only chart breakdowns", async () => { + const user = userEvent.setup(); + renderWithProviders(, { searchParams: `${initialUrl}&breakdown=members` }); + + expect(await screen.findByRole("table", { name: "Models" })).toBeInTheDocument(); + expect(screen.getByRole("heading", { name: "Top models" })).toBeInTheDocument(); + expect(screen.getByRole("table", { name: "Member spend" })).toBeInTheDocument(); + await user.click(screen.getByRole("combobox", { name: "Break down by" })); + expect((await screen.findAllByRole("option")).map((option) => option.textContent)).toEqual([ + "Models", + "API keys", + "Teams", + ]); + }); + + it("should retain a separate member spend table with range usage and current budget status across chart breakdowns", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { searchParams: initialUrl, onUrlUpdate }); + const table = await screen.findByRole("table", { name: "Member spend" }); + const firstMember = within(table).getAllByRole("row")[1]; + expect(firstMember).toHaveTextContent("Member One"); + expect(within(firstMember).getByText("$12.00")).toBeInTheDocument(); + expect(within(firstMember).getByText("$40.00 / $10.00")).toBeInTheDocument(); + expect(within(firstMember).getByText("Over budget")).toBeInTheDocument(); + + await selectBreakdown(user, "Teams"); + await user.click(screen.getByRole("radio", { name: "Tokens" })); + + expect(tableRows("Teams")[0]).toHaveTextContent("Operations"); + expect(within(table).getAllByRole("row")[1]).toHaveTextContent("Member One"); + const requestCount = fetchMock.mock.calls.length; + act(() => + fireEvent.change(screen.getByRole("textbox", { name: "Search members" }), { target: { value: "Member One" } }), + ); + expect(within(table).queryByText("Member Two")).not.toBeInTheDocument(); + expect(within(table).getByText("Member One")).toBeInTheDocument(); + await waitFor(() => + expect(new URLSearchParams(onUrlUpdate.mock.lastCall?.[0].queryString).get("memberSearch")).toBe("Member One"), + ); + const updatedUrl = new URLSearchParams(onUrlUpdate.mock.lastCall?.[0].queryString); + expect(updatedUrl.get("org")).toBe("org-1"); + expect(updatedUrl.get("breakdown")).toBe("teams"); + expect(updatedUrl.get("metric")).toBe("tokens"); + expect(updatedUrl.get("source")).toBe("portal"); + expect(fetchMock.mock.calls).toHaveLength(requestCount); + }); + + it("should clear the previous scope values while customer team usage loads", async () => { + const user = userEvent.setup(); + renderWithProviders(, { searchParams: initialUrl }); + await screen.findByRole("heading", { name: "Customer organization" }); + await selectBreakdown(user, "Teams"); + fetchMock.mockImplementation(() => new Promise(() => {})); + + await user.click(within(screen.getByRole("table", { name: "Teams" })).getByRole("button", { name: "Product" })); + + expect(screen.getByRole("status", { name: "Loading customer usage" })).toBeInTheDocument(); + expect(screen.queryByText("$15.00")).not.toBeInTheDocument(); + expect(screen.queryByRole("heading", { name: "Customer organization" })).not.toBeInTheDocument(); + expect(screen.queryByRole("table")).not.toBeInTheDocument(); + await waitFor(() => + expect(requestedUrls().some((url) => url.pathname === "/management/v1/customer_teams/team-a")).toBe(true), + ); + }); + + it("should request a changed date range without showing stale totals or losing the customer scope", async () => { + const user = userEvent.setup(); + const onUrlUpdate = vi.fn(); + renderWithProviders(, { searchParams: initialUrl, onUrlUpdate }); + await screen.findByRole("heading", { name: "Customer organization" }); + fetchMock.mockImplementation((input) => + requestUrl(input).pathname.endsWith("/daily/activity") && + requestUrl(input).searchParams.get("start_date") === "2026-10-01" + ? new Promise(() => {}) + : Promise.resolve(serve(input)), + ); + + await user.click(screen.getByRole("button", { name: / - / })); + act(() => { + fireEvent.change(screen.getByLabelText("Start date"), { target: { value: "2026-10-01" } }); + fireEvent.change(screen.getByLabelText("End date"), { target: { value: "2026-10-07" } }); + }); + await user.click(screen.getByRole("button", { name: "Apply" })); + + await waitFor(() => + expect(activityUrl("customer_organizations", "org-1")?.searchParams.get("start_date")).toBe("2026-10-01"), + ); + expect(activityUrl("customer_organizations", "org-1")?.searchParams.get("end_date")).toBe("2026-10-07"); + expect(activityUrl("customer_organizations", "org-1")?.pathname).toBe( + "/management/v1/customer_organizations/daily/activity", + ); + expect(screen.getByRole("status", { name: "Loading customer usage" })).toBeInTheDocument(); + expect(screen.queryByText("$15.00")).not.toBeInTheDocument(); + expect(onUrlUpdate).toHaveBeenLastCalledWith( + expect.objectContaining({ queryString: "?org=org-1&from=2026-10-01&to=2026-10-07&source=portal" }), + ); + }); + + it("should show a failed request as an error and load usage when retried", async () => { + const user = userEvent.setup(); + fetchMock.mockResolvedValueOnce(json({ error: "Usage temporarily unavailable" }, 503)); + renderWithProviders(, { searchParams: initialUrl }); + + expect(await screen.findByRole("alert")).toHaveTextContent("Could not load customer usage"); + expect(screen.getByRole("alert")).toHaveTextContent("Usage temporarily unavailable"); + expect(screen.queryByText("$0.00")).not.toBeInTheDocument(); + expect(screen.queryByText("No usage in this range")).not.toBeInTheDocument(); + + await user.click(screen.getByRole("button", { name: "Retry" })); + + expect(await screen.findByText("$15.00")).toBeInTheDocument(); + expect(screen.queryByRole("alert")).not.toBeInTheDocument(); + }); + + it("should show an empty date range after the scoped request succeeds without activity", async () => { + fetchMock.mockImplementation((input) => + Promise.resolve( + requestUrl(input).pathname.endsWith("/daily/activity") + ? json({ results: [], metadata: EMPTY_DAILY_ACTIVITY_METADATA }) + : serve(input), + ), + ); + renderWithProviders(, { searchParams: initialUrl }); + + expect(await screen.findByText("No usage in this range")).toBeInTheDocument(); + expect(screen.getByTitle("0 tokens")).toBeInTheDocument(); + expect(screen.queryByRole("status", { name: "Loading customer usage" })).not.toBeInTheDocument(); + expect(screen.getByText("No usage recorded for this breakdown.")).toBeInTheDocument(); + }); + + it("should request an organization selection without issuing an unscoped usage request", () => { + renderWithProviders(, { searchParams: "?source=portal" }); + + expect(screen.getByText("Open a customer organization to explore its usage.")).toBeInTheDocument(); + expect(screen.getByRole("link", { name: "Back to customers" })).toHaveAttribute("href", "/ui/customers"); + expect(fetchMock).not.toHaveBeenCalled(); + }); +}); diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/usage/page.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/usage/page.tsx new file mode 100644 index 00000000000..3fa99627bee --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/customers/usage/page.tsx @@ -0,0 +1,136 @@ +"use client"; + +import { useTranslation } from "react-i18next"; +import UsageStatStrip from "@/app/(dashboard)/usage/_components/components/overview/UsageStatStrip"; +import { ErrorState } from "@/components/shared/ErrorState"; +import { Button } from "@/components/ui/button"; +import { extractProxyErrorMessage } from "@/lib/http/client"; +import { uiHref } from "@/utils/uiHref"; +import CustomerMemberSpend from "../_components/CustomerMemberSpend"; +import CustomerUsageHeader from "../_components/CustomerUsageHeader"; +import CustomerUsageRankings from "../_components/CustomerUsageRankings"; +import CustomerUsageTrend from "../_components/CustomerUsageTrend"; +import useCustomerHierarchyUsage from "../_components/useCustomerHierarchyUsage"; + +const METRIC_LABELS = { + spend: "customerPortal.analytics.spend", + tokens: "customerPortal.analytics.tokens", + requests: "customerPortal.analytics.requests", +} as const; +const TOP_DIMENSION_LABELS = { + models: "customerPortal.analytics.topModels", + api_keys: "customerPortal.analytics.topKeys", + teams: "customerPortal.analytics.topTeams", +} as const; + +export default function CustomerHierarchyUsagePage() { + const { t } = useTranslation(); + const { + params, + scope, + dateValue, + preset, + usage, + data, + results, + totals, + modelColors, + rows, + drillDown, + setDateRange, + setDimension, + selectOrganization, + selectTeam, + setMetric, + setBucket, + setSearch, + setPage, + } = useCustomerHierarchyUsage(); + const metricLabel = t(METRIC_LABELS[params.metric]); + + if (!scope) + return ( +
+

{t("customerPortal.analytics.title")}

+

{t("customerPortal.analytics.selectOrganization")}

+ + {t("customerPortal.analytics.backToCustomers")} + +
+ ); + + return ( +
+ + {usage.isError ? ( + void usage.refetch()} + /> + ) : ( + <> + + + {!usage.isPending && ( + <> +
+

{t(TOP_DIMENSION_LABELS[params.breakdown])}

+

+ {t("customerPortal.analytics.sortedBy", { metric: metricLabel.toLocaleLowerCase() })} +

+
+ + + )} + {data && drillDown("members", id)} />} + {(params.team || params.member) && ( + + )} +

{t("customerPortal.analytics.freshness")}

+ + )} +
+ ); +} diff --git a/litellm/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/BreakdownChart.tsx b/litellm/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/BreakdownChart.tsx index 79c64e6079b..04b6ddbc2b0 100644 --- a/litellm/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/BreakdownChart.tsx +++ b/litellm/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/overview/BreakdownChart.tsx @@ -84,7 +84,7 @@ export function BreakdownControls({ } /** Provider logo on a soft tile with a series-colored dot, so mixed logo styles still read as a set. */ -function ModelMark({ model, color }: { model: string; color: string }) { +export function ModelMark({ model, color }: { model: string; color: string }) { const provider = providerForModel(model); return ( diff --git a/litellm/ui/litellm-dashboard/src/i18n/locales/en.json b/litellm/ui/litellm-dashboard/src/i18n/locales/en.json index bdade69fef6..a2369c9f077 100644 --- a/litellm/ui/litellm-dashboard/src/i18n/locales/en.json +++ b/litellm/ui/litellm-dashboard/src/i18n/locales/en.json @@ -20681,6 +20681,72 @@ "maxBudget": "Max budget (USD)", "maxBudgetPlaceholder": "Leave empty for no budget", "sharedBudgetWarning": "This customer is on budget {{budgetId}}. Changing the budget here changes it for every customer on that budget." + }, + "analytics": { + "title": "Customer usage", + "customers": "Customers", + "breadcrumb": "Customer hierarchy", + "selectOrganization": "Open a customer organization to explore its usage.", + "backToCustomers": "Back to customers", + "organizationUsage": "Organization usage", + "teamUsage": "Team usage", + "memberUsage": "Member usage", + "subtitle": "Understand where usage goes, from your organization to each member.", + "trend": "Usage over time", + "trendHint": "Daily activity, broken down by {{dimension}}.", + "metric": "Usage metric", + "spend": "Spend", + "tokens": "Tokens", + "requests": "Requests", + "bucket": "Time interval", + "daily": "Daily", + "weekly": "Weekly", + "loading": "Loading customer usage", + "empty": "No usage in this range", + "error": "Could not load customer usage", + "breakdown": "Usage breakdown", + "models": "Models", + "keys": "API keys", + "useCases": "Use cases", + "teams": "Teams", + "members": "Members", + "sortedBy": "Ranked by {{metric}}", + "search": "Search breakdown", + "rank": "Rank", + "share": "Share", + "rankingHint": "Share of the selected metric across this customer scope. Click a team or member to explore its usage.", + "useCaseHint": "Use cases recorded with requests. Unattributed activity appears as Unassigned.", + "noMatches": "No results match your search.", + "noBreakdown": "No usage recorded for this breakdown.", + "pagination": "{{from}}โ€“{{to}} of {{total}}", + "previous": "Previous", + "next": "Next", + "allOrganizationUsage": "View all organization usage", + "freshness": "Usage appears after accounting completes. Date ranges use UTC. Customer budgets are managed separately.", + "scopeMismatch": "This team or member does not belong to the selected organization.", + "invalidIdentifier": "Select a valid customer identifier.", + "incompletePages": "Could not load the complete usage range. Please choose fewer days.", + "unassignedTeam": "Unassigned or removed team", + "unassignedMember": "Unassigned or removed member", + "unnamedKey": "Unnamed key", + "currentSpend": "Budget spend", + "noBudget": "No budget set", + "breakDownBy": "Break down by", + "topModels": "Top models", + "topKeys": "Top API keys", + "topTeams": "Top teams", + "topMembers": "Top members", + "memberSpend": "Member spend", + "memberSpendHint": "Usage in the selected date range. Cap status uses each memberโ€™s recorded budget spend.", + "searchMembers": "Search members", + "member": "Member", + "currentBudget": "Spend / cap", + "status": "Status", + "noCap": "No cap", + "atBudget": "At budget", + "withinBudget": "Within budget", + "noMembers": "No members in this customer scope.", + "overBudget": "Over budget" } } } diff --git a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts index 97f169bf9f5..ff97a43c136 100644 --- a/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/litellm/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9608,6 +9608,26 @@ export interface paths { patch?: never; trace?: never; }; + "/management/v1/customer_members/daily/activity": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Customer Member Daily Activity + * @description Get paginated daily customer member activity with model, provider, key and member breakdowns. Filters accept comma-separated member IDs. Proxy admins only. + */ + get: operations["customer_member_daily_activity_management_v1_customer_members_daily_activity_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/management/v1/customer_members/{customer_member_id}": { parameters: { query?: never; @@ -9660,6 +9680,26 @@ export interface paths { patch?: never; trace?: never; }; + "/management/v1/customer_organizations/daily/activity": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Customer Organization Daily Activity + * @description Get paginated daily customer organization activity with model, provider, key and organization breakdowns. Filters accept comma-separated organization IDs. Proxy admins only. + */ + get: operations["customer_organization_daily_activity_management_v1_customer_organizations_daily_activity_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/management/v1/customer_organizations/{customer_org_id}": { parameters: { query?: never; @@ -9756,6 +9796,26 @@ export interface paths { patch?: never; trace?: never; }; + "/management/v1/customer_teams/daily/activity": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + /** + * Customer Team Daily Activity + * @description Get paginated daily customer team activity with model, provider, key and team breakdowns. Filters accept comma-separated team IDs. Proxy admins only. + */ + get: operations["customer_team_daily_activity_management_v1_customer_teams_daily_activity_get"]; + put?: never; + post?: never; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/management/v1/customer_teams/{customer_team_id}": { parameters: { query?: never; @@ -68072,6 +68132,44 @@ export interface operations { }; }; }; + customer_member_daily_activity_management_v1_customer_members_daily_activity_get: { + parameters: { + query?: { + customer_member_ids?: string | null; + start_date?: string | null; + end_date?: string | null; + model?: string | null; + api_key?: string | null; + page?: number; + page_size?: number; + exclude_customer_member_ids?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["SpendAnalyticsPaginatedResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_customer_member_management_v1_customer_members__customer_member_id__get: { parameters: { query?: never; @@ -68232,6 +68330,44 @@ export interface operations { }; }; }; + customer_organization_daily_activity_management_v1_customer_organizations_daily_activity_get: { + parameters: { + query?: { + customer_org_ids?: string | null; + start_date?: string | null; + end_date?: string | null; + model?: string | null; + api_key?: string | null; + page?: number; + page_size?: number; + exclude_customer_org_ids?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["SpendAnalyticsPaginatedResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_customer_organization_management_v1_customer_organizations__customer_org_id__get: { parameters: { query?: never; @@ -68494,6 +68630,44 @@ export interface operations { }; }; }; + customer_team_daily_activity_management_v1_customer_teams_daily_activity_get: { + parameters: { + query?: { + customer_team_ids?: string | null; + start_date?: string | null; + end_date?: string | null; + model?: string | null; + api_key?: string | null; + page?: number; + page_size?: number; + exclude_customer_team_ids?: string | null; + }; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["SpendAnalyticsPaginatedResponse"]; + }; + }; + /** @description Validation Error */ + 422: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": components["schemas"]["HTTPValidationError"]; + }; + }; + }; + }; get_customer_team_management_v1_customer_teams__customer_team_id__get: { parameters: { query?: never; diff --git a/tests/unit/accounting/test_application.py b/tests/unit/accounting/test_application.py new file mode 100644 index 00000000000..d60893a147d --- /dev/null +++ b/tests/unit/accounting/test_application.py @@ -0,0 +1,103 @@ +from collections.abc import Mapping +from dataclasses import dataclass, replace +from datetime import datetime, timezone +from typing import Final, cast + +import pytest +from pydantic import TypeAdapter + +from litellm.accounting.application import _apply_customer_usage +from litellm.accounting.delivery_projection import project_deliveries +from litellm.accounting.projection import Contribution +from litellm.proxy.common_utils.customer_usage import CUSTOMER_USAGE_METADATA_KEY, CustomerUsageAttribution +from litellm.repositories.accounting_delivery_repository import DeliveryProjection, LedgerTransaction + +_FIELDS: Final = TypeAdapter(Mapping[str, object]) +_TEXT: Final = TypeAdapter(str) +_UPDATES: Final = TypeAdapter(Mapping[str, Mapping[str, float]]) +_IDENTIFIER: Final = "\U0001f680" * 256 +_ATTRIBUTION: Final = CustomerUsageAttribution( + customer_org_id=_IDENTIFIER, customer_team_id=_IDENTIFIER, customer_member_id=_IDENTIFIER +) + + +@dataclass(frozen=True, slots=True) +class Bucket: + values: Mapping[str, object] + counters: Mapping[str, float] + + +class Buckets: + def __init__(self) -> None: + self.rows: Mapping[str, Bucket] = {} + + async def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: + row_id: Final = _TEXT.validate_python(where["id"]) + create: Final = _FIELDS.validate_python(data["create"]) + previous: Final = self.rows.get(row_id) + increments: Final = ( + {name: fields["increment"] for name, fields in _UPDATES.validate_python(data["update"]).items()} + if previous is not None + else {name: float(value) for name, value in create.items() if isinstance(value, (int, float))} + ) + self.rows = { + **self.rows, + row_id: Bucket( + values=previous.values if previous is not None else create, + counters={ + name: (previous.counters.get(name, 0) if previous is not None else 0) + value + for name, value in increments.items() + }, + ), + } + + +@dataclass(frozen=True, slots=True) +class Transaction: + litellm_dailycustomermemberspend: Buckets + + +def contribution(alias: str, provider: str, spend: float) -> Contribution: + delivery: Final = DeliveryProjection( + { + "request_id": "long-customer-id-request", + "startTime": datetime(2026, 1, 1, tzinfo=timezone.utc).isoformat(), + "api_key": "hashed-key", + "model": "shared-model", + "model_group": alias, + "custom_llm_provider": provider, + "spend": spend, + "prompt_tokens": 10, + "completion_tokens": 5, + "metadata": {CUSTOMER_USAGE_METADATA_KEY: _ATTRIBUTION.model_dump_json()}, + } + ) + return next(item for item in project_deliveries((delivery,)) if item.table == "LiteLLM_DailyCustomerMemberSpend") + + +@pytest.mark.asyncio +async def test_long_customer_ids_roll_up_without_truncation_and_keep_alias_provider_buckets_distinct() -> None: + buckets: Final = Buckets() + transaction: Final = cast(LedgerTransaction, Transaction(buckets)) # cast-ok: in-memory Prisma upsert boundary + first: Final = contribution("alias-a", "provider-a", 1) + second: Final = replace(contribution("alias-a", "provider-a", 2), key=tuple(reversed(first.key))) + + await _apply_customer_usage(transaction, (first,)) + await _apply_customer_usage(transaction, (second, contribution("alias-b", "provider-a", 4))) + await _apply_customer_usage(transaction, (contribution("alias-a", "provider-b", 8),)) + + assert len(buckets.rows) == 3 + assert all(len(row_id.encode("utf-8")) == 64 for row_id in buckets.rows) + assert all( + row.values["customer_org_id"] + == row.values["customer_team_id"] + == row.values["customer_member_id"] + == _IDENTIFIER + for row in buckets.rows.values() + ) + assert { + (row.values["model_group"], row.values["custom_llm_provider"]): row.counters["spend"] + for row in buckets.rows.values() + } == {("alias-a", "provider-a"): 3, ("alias-b", "provider-a"): 4, ("alias-a", "provider-b"): 8} + assert sum(row.counters["api_requests"] for row in buckets.rows.values()) == 4 + assert sum(row.counters["prompt_tokens"] + row.counters["completion_tokens"] for row in buckets.rows.values()) == 60 diff --git a/tests/unit/accounting/test_projection.py b/tests/unit/accounting/test_projection.py index 090825dca75..75b3f76de26 100644 --- a/tests/unit/accounting/test_projection.py +++ b/tests/unit/accounting/test_projection.py @@ -6,20 +6,120 @@ from typing import Final import pytest from pydantic import TypeAdapter +from litellm.accounting.delivery_projection import project_deliveries from litellm.accounting.projection import project from litellm.accounting.snapshot import build_accounting_payload from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.proxy._types import SpendLogsPayload +from litellm.proxy.common_utils.customer_usage import CUSTOMER_USAGE_METADATA_KEY, CustomerUsageAttribution from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.proxy_server import proxy_logging_obj from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload from litellm.proxy.utils import PrismaClient +from litellm.repositories.accounting_delivery_repository import DeliveryProjection from litellm.types.utils import ModelResponse, Usage _NOW: Final = datetime(2026, 11, 2, tzinfo=timezone.utc) _FIELDS: Final = TypeAdapter(dict[str, object]) +def test_customer_model_aliases_and_providers_keep_separate_daily_totals() -> None: + stamp: Final = CustomerUsageAttribution(customer_org_id="org-1", customer_member_id="member-1").model_dump_json() + base: Final = { + "startTime": _NOW, + "api_key": "hashed-key", + "model": "shared-model", + "metadata": {CUSTOMER_USAGE_METADATA_KEY: stamp}, + } + deliveries: Final = tuple( + DeliveryProjection({**base, "model_group": alias, "custom_llm_provider": provider, "spend": spend}) + for alias, provider, spend in ( + ("alias-a", "provider-a", 1), + ("alias-b", "provider-a", 2), + ("alias-a", "provider-b", 3), + ) + ) + totals: Final = { + (dict(item.key)["model_group"], dict(item.key)["custom_llm_provider"]): dict(item.increments)["spend"] + for item in project_deliveries(deliveries) + if item.table == "LiteLLM_DailyCustomerMemberSpend" + } + assert totals == {("alias-a", "provider-a"): 1, ("alias-b", "provider-a"): 2, ("alias-a", "provider-b"): 3} + + +@pytest.mark.parametrize("retry", [False, True]) +def test_customer_usage_survives_logging_and_accounting_with_correct_request_counts(retry: bool) -> None: + attribution: Final = CustomerUsageAttribution( + customer_org_id="org-1", customer_team_id="team-1", customer_member_id="member-1" + ) + payload: Final = get_logging_payload( + kwargs={ + "model": "usage-model", + "custom_llm_provider": "openai", + "call_type": "acompletion", + "litellm_call_id": "usage-request", + "litellm_params": { + "metadata": { + "user_api_key": "hashed-key", + "status": "success", + CUSTOMER_USAGE_METADATA_KEY: attribution.model_dump_json(), + **({BILLED_RETRY_ATTEMPT_METADATA_KEY: True} if retry else {}), + } + }, + }, + response_obj=ModelResponse( + id="usage-request", model="usage-model", usage=Usage(prompt_tokens=10, completion_tokens=5) + ), + start_time=_NOW, + end_time=_NOW, + ) + accounting: Final = build_accounting_payload( + {**payload, "spend": 0.75}, entities={}, model_access_groups=(), windows=(), router=lambda: None + ) + contributions: Final = project(accounting) + usage: Final = next(item for item in contributions if item.table == "LiteLLM_DailyCustomerMemberSpend") + assert dict(usage.key) == { + "date": _NOW.date().isoformat(), + "customer_org_id": attribution.customer_org_id, + "customer_team_id": attribution.customer_team_id, + "customer_member_id": attribution.customer_member_id, + "model": payload["model"], + "model_group": payload["model_group"], + "custom_llm_provider": payload["custom_llm_provider"], + "api_key": payload["api_key"], + "mcp_namespaced_tool_name": "", + "endpoint": "/chat/completions", + } + assert dict(usage.increments) == { + "spend": 0.75, + "prompt_tokens": 10, + "completion_tokens": 5, + "api_requests": 0 if retry else 1, + "successful_requests": 0 if retry else 1, + "failed_requests": 0, + "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, + "guardrail_spend": 0, + "total_response_time_ms": 0, + "timed_requests": 1, + } + assert { + item.table: dict(item.increments)["spend"] + for item in contributions + if item.table + in ("LiteLLM_CustomerOrganizationTable", "LiteLLM_CustomerTeamTable", "LiteLLM_CustomerMemberTable") + } == { + "LiteLLM_CustomerOrganizationTable": 0.75, + "LiteLLM_CustomerTeamTable": 0.75, + "LiteLLM_CustomerMemberTable": 0.75, + } + + @pytest.mark.parametrize("status", ["success", "failure"]) @pytest.mark.parametrize("origin", [None, "auto_router", "billed_retry"]) @pytest.mark.parametrize("spend", [0.0, 0.25]) diff --git a/tests/unit/accounting/test_worker.py b/tests/unit/accounting/test_worker.py index d3b5d71cff9..372ae90b158 100644 --- a/tests/unit/accounting/test_worker.py +++ b/tests/unit/accounting/test_worker.py @@ -1,5 +1,9 @@ import json +from collections.abc import AsyncIterator, Mapping +from contextlib import asynccontextmanager from dataclasses import dataclass +from datetime import datetime, timezone +from functools import partial from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock @@ -10,11 +14,18 @@ from fakeredis.aioredis import FakeRedis from pydantic import TypeAdapter from litellm import in_memory_llm_clients_cache +from litellm.accounting.application import apply_deliveries from litellm.accounting.worker import flush_accounting_deliveries from litellm.caching.redis_cache import RedisCache +from litellm.proxy.common_utils.customer_usage import CUSTOMER_USAGE_METADATA_KEY, CustomerUsageAttribution from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter from litellm.proxy.utils import PrismaClient, ProxyLogging +from litellm.repositories.accounting_delivery_repository import ( + AccountingDeliveryRepository, + DeliveryProjection, + LedgerTransaction, +) _FIELDS: Final = TypeAdapter(dict[str, object]) _INCREMENTS: Final = TypeAdapter(dict[str, dict[str, float]]) @@ -30,7 +41,7 @@ class _Ledger: spend: float = 0.0 attempts: int = 0 - async def query(self, query: str, *args: object) -> tuple[dict[str, object], ...]: + async def query(self, query: str, *args: object) -> tuple[Mapping[str, object], ...]: if query.startswith("SELECT delivery_id, attempts"): if not self.pending or self.attempts or args[1] is True: return () @@ -96,3 +107,119 @@ async def test_delivery_charges_the_creating_key_despite_an_orphaned_maintenance in_memory_llm_clients_cache.delete_cache(key=cache_key) await redis.aclose() sync_redis.close() + + +class _CustomerLedger: + def __init__(self, fail_acknowledgement: bool) -> None: + self.pending = True + self.attempts = 0 + self.fail_acknowledgement = fail_acknowledgement + self.daily: Mapping[str, int | float] = {} + self.entities: Mapping[str, float] = {} + + async def query(self, query: str, *args: object) -> tuple[Mapping[str, object], ...]: + if query.startswith("SELECT delivery_id, attempts"): + if not self.pending or (self.attempts > 0) is not args[1]: + return () + return ( + { + "delivery_id": "customer-delivery", + "attempts": self.attempts, + "accounting_payload": _CUSTOMER_PAYLOAD, + }, + ) + if "SET attempts = attempts + 1" in query: + self.attempts += 1 + return ({"delivery_id": "customer-delivery", "attempts": self.attempts},) + return () + + async def execute(self, query: str, *args: object) -> int: + if "SET state = 'applied'" not in query: + return 0 + if self.fail_acknowledgement: + self.fail_acknowledgement = False + raise OSError("database connection lost before delivery acknowledgement") + self.pending = False + return 1 + + def upsert_daily(self, *, where: object, data: object) -> None: + values: Final = _FIELDS.validate_python(_FIELDS.validate_python(data)["create"]) + counters: Final = {name: value for name, value in values.items() if isinstance(value, (int, float))} + self.daily = {name: self.daily.get(name, 0) + value for name, value in counters.items()} + + def update_entity(self, entity: str, *, where: object, data: object) -> int: + spend: Final = _INCREMENTS.validate_python(data)["spend"]["increment"] + self.entities = {**self.entities, entity: self.entities.get(entity, 0.0) + spend} + return 1 + + +_CUSTOMER_PAYLOAD: Final = json.dumps( + { + "request_id": "customer-completed-request", + "startTime": datetime(2026, 1, 1, tzinfo=timezone.utc).isoformat(), + "api_key": "hashed-key", + "spend": 0.75, + "prompt_tokens": 10, + "completion_tokens": 5, + "metadata": { + CUSTOMER_USAGE_METADATA_KEY: CustomerUsageAttribution( + customer_org_id="customer-org", customer_team_id="customer-team", customer_member_id="customer-member" + ).model_dump_json() + }, + } +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_acknowledgement", (False, True)) +async def test_customer_daily_and_lifetime_totals_commit_once_after_delivery_replay( + fail_acknowledgement: bool, +) -> None: + ledger: Final = _CustomerLedger(fail_acknowledgement) + transaction: Final = MagicMock() + transaction.query_raw = AsyncMock(side_effect=ledger.query) + transaction.execute_raw = AsyncMock(side_effect=ledger.execute) + transaction.litellm_dailycustomermemberspend.upsert = AsyncMock(side_effect=ledger.upsert_daily) + transaction.litellm_customerorganizationtable.update_many = AsyncMock( + side_effect=partial(ledger.update_entity, "organization") + ) + transaction.litellm_customerteamtable.update_many = AsyncMock(side_effect=partial(ledger.update_entity, "team")) + transaction.litellm_customermembertable.update_many = AsyncMock(side_effect=partial(ledger.update_entity, "member")) + transaction.batch_.return_value.__aenter__ = AsyncMock(return_value=MagicMock()) + transaction.batch_.return_value.__aexit__ = AsyncMock(return_value=False) + + @asynccontextmanager + async def database_transaction(*, timeout: object) -> AsyncIterator[LedgerTransaction]: + before: Final = (ledger.pending, ledger.daily, ledger.entities) + try: + yield transaction + except Exception: + ledger.pending, ledger.daily, ledger.entities = before + raise + + database: Final = MagicMock(spec=PrismaClient) + database.tx.side_effect = database_transaction + repository: Final = AccountingDeliveryRepository(database) + writer: Final = DBSpendUpdateWriter(accounting_clock=lambda: 0.0) + logger: Final = ProxyLogging(user_api_key_cache=UserApiKeyCache()) + + async def apply(tx: LedgerTransaction, deliveries: tuple[DeliveryProjection, ...]) -> None: + await apply_deliveries(tx, deliveries, prisma_client=database, writer=writer, proxy_logging_obj=logger) + + first: Final = await repository.apply_pending(apply) + if fail_acknowledgement: + assert first.failed == 1 + assert ledger.daily == {} + assert ledger.entities == {} + retried: Final = await repository.apply_pending(apply) + assert retried.applied == 1 + else: + assert first.applied == 1 + await repository.apply_pending(apply) + await repository.apply_pending(apply) + assert ledger.daily["spend"] == 0.75 + assert ledger.daily["prompt_tokens"] == 10 + assert ledger.daily["completion_tokens"] == 5 + assert ledger.daily["api_requests"] == 1 + assert ledger.entities == {"organization": 0.75, "team": 0.75, "member": 0.75} + assert ledger.pending is False diff --git a/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json b/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json index 6e5461c3b60..40674758371 100644 --- a/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json +++ b/tests/unit/integrations/gcs_pubsub/spend_logs_payload.json @@ -11,7 +11,7 @@ "user": "", "team_id": "", "organization_id": "", - "metadata": "{\"actor_agent_id\": null, \"target_agent_id\": null, \"billing_agent_id\": null, \"agent_execution_mode\": null, \"verified_human_user_id\": null, \"used_client_oauth_token\": null, \"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"litellm_roi_estimator\": false, \"billed_retry_attempt\": false, \"router_metadata\": null, \"autorouter_savings_estimate\": null, \"autorouter_baseline_observation\": null, \"azure_spillover\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_user_email\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"user_agent\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null, \"kubernetes_pod_routing\": null}", + "metadata": "{\"actor_agent_id\": null, \"target_agent_id\": null, \"billing_agent_id\": null, \"agent_execution_mode\": null, \"verified_human_user_id\": null, \"used_client_oauth_token\": null, \"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"litellm_roi_estimator\": false, \"billed_retry_attempt\": false, \"router_metadata\": null, \"autorouter_savings_estimate\": null, \"autorouter_baseline_observation\": null, \"azure_spillover\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_customer_usage\": null, \"user_api_key_user_id\": null, \"user_api_key_user_email\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"user_agent\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null, \"kubernetes_pod_routing\": null}", "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, diff --git a/tests/unit/proxy/common_utils/test_customer_usage.py b/tests/unit/proxy/common_utils/test_customer_usage.py new file mode 100644 index 00000000000..79b16cafa82 --- /dev/null +++ b/tests/unit/proxy/common_utils/test_customer_usage.py @@ -0,0 +1,225 @@ +from collections.abc import Mapping +from dataclasses import dataclass +from types import SimpleNamespace +from typing import Final, cast # noqa: TID251 # the fake exposes only the Prisma lookup methods exercised here + +import pytest + +from litellm.proxy.common_utils.customer_usage import ( + CustomerUsageAttribution, + CustomerUsageDatabase, + CustomerUsageError, + CustomerUsageFailure, + resolve_customer_usage, +) + + +@dataclass(frozen=True, slots=True) +class Table: + identifier: str + rows: Mapping[str, Mapping[str, object]] + + async def find_unique(self, where: Mapping[str, object]) -> Mapping[str, object] | None: + return self.rows.get(str(where[self.identifier])) + + +@dataclass(frozen=True, slots=True) +class Database: + litellm_customerorganizationtable: Table + litellm_customerteamtable: Table + litellm_customermembertable: Table + + +DATABASE: Final = cast( # cast-ok: the external Prisma fake implements the lookup subset used by attribution + CustomerUsageDatabase, + Database( + Table("customer_org_id", {"org-1": {}, "org-2": {}}), + Table("customer_team_id", {"team-1": {"parent_org_id": "org-1"}, "team-2": {"parent_org_id": "org-2"}}), + Table("customer_member_id", {"member-1": {"parent_team_id": "team-1"}}), + ), +) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata", "expected"), + ( + ( + {"customer_member_id": "member-1"}, + CustomerUsageAttribution(customer_org_id="org-1", customer_team_id="team-1", customer_member_id="member-1"), + ), + ({"customer_team_id": "team-1"}, CustomerUsageAttribution(customer_org_id="org-1", customer_team_id="team-1")), + ({"customer_org_id": "org-1"}, CustomerUsageAttribution(customer_org_id="org-1")), + ), +) +async def test_resolves_ancestors_without_requiring_redundant_request_ids( + metadata: Mapping[str, object], expected: CustomerUsageAttribution +) -> None: + actual: Final = await resolve_customer_usage( + metadata=metadata, key_metadata={}, administrator=True, database=DATABASE, ledger_enabled=True + ) + assert actual == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata", "key_metadata", "administrator", "expected"), + ( + ({"customer_member_id": "member-1", "customer_team_id": "team-2"}, {}, True, CustomerUsageError.MISMATCH), + ({"customer_member_id": "member-1"}, {"customer_org_id": "org-2"}, False, CustomerUsageError.FORBIDDEN), + ({"customer_team_id": "team-2", "customer_org_id": "org-1"}, {}, True, CustomerUsageError.MISMATCH), + ({"customer_org_id": "org-2"}, {"customer_org_id": "org-1"}, False, CustomerUsageError.FORBIDDEN), + ({"customer_org_id": "org-1"}, {}, False, CustomerUsageError.UNBOUND), + ({"customer_member_id": "missing"}, {}, True, CustomerUsageError.NOT_FOUND), + ({"customer_team_id": "missing"}, {}, True, CustomerUsageError.NOT_FOUND), + ({"customer_org_id": "missing"}, {}, True, CustomerUsageError.NOT_FOUND), + ({"customer_org_id": 123}, {}, True, CustomerUsageError.INVALID), + ), +) +async def test_rejects_cross_organization_attribution_and_invalid_identifiers( + metadata: Mapping[str, object], + key_metadata: Mapping[str, object], + administrator: bool, + expected: CustomerUsageError, +) -> None: + actual: Final = await resolve_customer_usage( + metadata=metadata, + key_metadata=key_metadata, + administrator=administrator, + database=DATABASE, + ledger_enabled=True, + ) + assert actual == CustomerUsageFailure(expected) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("metadata", "key_metadata"), + ( + ({"customer_member_id": "missing"}, {"customer_org_id": "org-1"}), + ({"customer_member_id": "member-1"}, {"customer_org_id": "org-2"}), + ( + {"customer_member_id": "member-1"}, + {"customer_org_id": "org-2", "customer_team_id": "team-2"}, + ), + ({"customer_team_id": "missing"}, {"customer_org_id": "org-1"}), + ({"customer_team_id": "team-2"}, {"customer_org_id": "org-1"}), + ({}, {"customer_org_id": "missing"}), + ), +) +async def test_non_admin_cannot_distinguish_missing_and_other_organization_identifiers( + metadata: Mapping[str, object], key_metadata: Mapping[str, object] +) -> None: + actual: Final = await resolve_customer_usage( + metadata=metadata, + key_metadata=key_metadata, + administrator=False, + database=DATABASE, + ledger_enabled=True, + ) + assert actual == CustomerUsageFailure(CustomerUsageError.FORBIDDEN) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("administrator", (False, True)) +async def test_invalid_identifier_types_keep_validation_errors(administrator: bool) -> None: + actual: Final = await resolve_customer_usage( + metadata={"customer_member_id": 123}, + key_metadata={"customer_org_id": "org-1"}, + administrator=administrator, + database=DATABASE, + ledger_enabled=True, + ) + assert actual == CustomerUsageFailure(CustomerUsageError.INVALID) + + +@pytest.mark.asyncio +async def test_key_binding_tracks_requests_with_no_repeated_customer_metadata() -> None: + actual: Final = await resolve_customer_usage( + metadata={}, + key_metadata={"customer_org_id": "org-1"}, + administrator=False, + database=DATABASE, + ledger_enabled=True, + ) + assert actual == CustomerUsageAttribution(customer_org_id="org-1") + + +@pytest.mark.asyncio +async def test_requires_accounting_before_accepting_a_tracked_request() -> None: + actual: Final = await resolve_customer_usage( + metadata={"customer_org_id": "org-1"}, + key_metadata={}, + administrator=True, + database=DATABASE, + ledger_enabled=False, + ) + assert actual == CustomerUsageFailure(CustomerUsageError.UNAVAILABLE) + + +@pytest.mark.asyncio +async def test_untracked_requests_do_not_require_a_database_or_ledger() -> None: + actual: Final = await resolve_customer_usage( + metadata={}, key_metadata={}, administrator=False, database=None, ledger_enabled=False + ) + assert actual is None + + +@pytest.mark.asyncio +async def test_untracked_request_enrichment_does_not_touch_customer_database_tables( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from fastapi import Request + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(db=object())) + request: Final = Request( + {"type": "http", "path": "/v1/chat/completions", "method": "POST", "headers": (), "query_string": b""} + ) + enriched: Final = await add_litellm_data_to_request( + data={"model": "usage-model", "metadata": {"application": "support"}}, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key"), + proxy_config=ProxyConfig(), + general_settings={}, + ) + assert enriched["metadata"]["application"] == "support" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("path", ["/v1/chat/completions", "/v1/responses"]) +async def test_request_enrichment_replaces_forged_stamps_and_preserves_valid_attribution( + monkeypatch: pytest.MonkeyPatch, path: str +) -> None: + from fastapi import Request + from pydantic import TypeAdapter + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.common_utils.customer_usage import CUSTOMER_USAGE_METADATA_KEY + from litellm.proxy.litellm_pre_call_utils import add_litellm_data_to_request, get_metadata_variable_name + from litellm.proxy.proxy_server import ProxyConfig + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", SimpleNamespace(db=DATABASE)) + request: Final = Request({"type": "http", "path": path, "method": "POST", "headers": (), "query_string": b""}) + metadata_name: Final = get_metadata_variable_name(request) + forged: Final = CustomerUsageAttribution(customer_org_id="org-2").model_dump_json() + metadata: Final = {"customer_member_id": "member-1", CUSTOMER_USAGE_METADATA_KEY: forged} + enriched: Final = await add_litellm_data_to_request( + data={ + "model": "usage-model", + "metadata": {CUSTOMER_USAGE_METADATA_KEY: forged}, + "litellm_metadata": {CUSTOMER_USAGE_METADATA_KEY: forged}, + metadata_name: metadata, + }, + request=request, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", metadata={"customer_org_id": "org-1"}), + proxy_config=ProxyConfig(), + general_settings={}, + ) + stamp: Final = TypeAdapter(str).validate_python(enriched[metadata_name][CUSTOMER_USAGE_METADATA_KEY]) + assert CustomerUsageAttribution.model_validate_json(stamp) == CustomerUsageAttribution( + customer_org_id="org-1", customer_team_id="team-1", customer_member_id="member-1" + ) diff --git a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py index 554a8bb3215..a54912c44ba 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -3488,7 +3488,7 @@ async def test_a_failed_request_settles_the_team_share_gauges_too(monkeypatch: p registry = _prometheus_sink(monkeypatch) await _admit_then_settle(handler, _team_key(), tokens=90) claim_request_stash_for_data({"litellm_call_id": "failed-call"}) - await _check(handler, _team_key(), saturation=0.0, data=_prompt(words=40), priority=None) + await _check(handler, _team_key(), saturation=0.5) await handler.async_log_failure_event( kwargs={ "litellm_params": {"metadata": {"model_group": "shared-model"}}, @@ -3504,10 +3504,10 @@ async def test_a_failed_request_settles_the_team_share_gauges_too(monkeypatch: p @pytest.mark.asyncio @pytest.mark.parametrize( - "auth, burns, saturation_threshold, reason", + "auth, burns, saturation_threshold, reason, enforced_share_tokens", [ - (_team_key(), (60, 70), 0.1, "model_per_team_saturated"), - (_capped_team_key(2), (150, 60), 0.5, "model_per_team_max_borrow_multiple"), + (_team_key(), (60, 70), 0.0, "model_per_team_saturated", 100), + (_capped_team_key(2), (150, 60), 0.5, "model_per_team_max_borrow_multiple", 200), ], ids=["saturated", "borrow_cap"], ) @@ -3517,6 +3517,7 @@ async def test_a_refused_request_is_counted_under_the_reason_its_429_header_carr burns: tuple[int, ...], saturation_threshold: float, reason: str, + enforced_share_tokens: int, ): handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=saturation_threshold) registry = _prometheus_sink(monkeypatch) @@ -3532,7 +3533,10 @@ async def test_a_refused_request_is_counted_under_the_reason_its_429_header_carr ) == 1 ) - assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **TEAM_LABELS) == sum(burns) - 100 + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **TEAM_LABELS) == enforced_share_tokens + assert ( + _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **TEAM_LABELS) == sum(burns) - enforced_share_tokens + ) @pytest.mark.asyncio @@ -3596,7 +3600,7 @@ async def test_an_alternates_refusal_is_not_counted_while_the_preferred_model_ad @pytest.mark.asyncio async def test_share_gauges_and_throttles_carry_the_teams_priority_so_the_series_join(monkeypatch: pytest.MonkeyPatch): - handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.1) + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) registry = _prometheus_sink(monkeypatch) team: Final = UserAPIKeyAuth( api_key="sk-interactive", @@ -3619,7 +3623,7 @@ async def test_share_gauges_and_throttles_carry_the_teams_priority_so_the_series async def test_a_fallback_served_request_publishes_the_teams_share_on_the_fallback_model( monkeypatch: pytest.MonkeyPatch, ): - handler = _borrowing_handler(monkeypatch, borrows=True) + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) registry = _prometheus_sink(monkeypatch) team: Final = _team_key(team_limits={"model_tpm_limit": {"other-model": 40, "shared-model": 100}}) assert await handler.async_pre_call_hook(team, DualCache(), {"model": "other-model"}, "completion") is None @@ -3642,7 +3646,7 @@ async def test_a_failing_metrics_sink_never_changes_admission(monkeypatch: pytes def increment_dynamic_limiter_throttled(self, *args: object, **kwargs: object) -> None: raise RuntimeError("metrics backend down") - handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.1) + handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.0) _prometheus_sink(monkeypatch, BrokenSink) await _admit_at_the_hook(handler, _team_key(), tokens=60) await _admit_at_the_hook(handler, _team_key(), tokens=70) diff --git a/tests/unit/proxy/management_endpoints/management_v1/test_customer_usage.py b/tests/unit/proxy/management_endpoints/management_v1/test_customer_usage.py new file mode 100644 index 00000000000..af16d11373b --- /dev/null +++ b/tests/unit/proxy/management_endpoints/management_v1/test_customer_usage.py @@ -0,0 +1,278 @@ +from collections.abc import Mapping, Sequence +from dataclasses import dataclass +from functools import reduce +from typing import Final, cast + +import pytest +from fastapi import FastAPI, Request +from fastapi.testclient import TestClient +from pydantic import BaseModel, ConfigDict + +from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth +from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.list_api.common import ManagementProblem, problem_response +from litellm.proxy.management_endpoints.management_v1.customer_usage import get_customer_usage_prisma, router +from litellm.proxy.utils import PrismaClient +from litellm.types.proxy.management_endpoints.common_daily_activity import SpendMetrics + + +class Record(BaseModel): + model_config = ConfigDict(frozen=True, extra="allow") + + +class DailyRecord(SpendMetrics): + model_config = ConfigDict(frozen=True) + + id: str + date: str + customer_org_id: str + customer_team_id: str + customer_member_id: str + api_key: str + model: str | None = None + model_group: str | None = None + custom_llm_provider: str | None = None + mcp_namespaced_tool_name: str | None = None + endpoint: str | None = None + + +def matches(values: Mapping[str, object], where: Mapping[str, object]) -> bool: + for key, condition in where.items(): + if key == "OR": + if not any(matches(values, clause) for clause in condition): + return False + elif key == "AND": + if not all(matches(values, clause) for clause in condition): + return False + elif isinstance(condition, Mapping): + if "in" in condition and values.get(key) not in condition["in"]: + return False + if "not" in condition and matches(values, {key: condition["not"]}): + return False + if "gte" in condition and str(values.get(key)) < str(condition["gte"]): + return False + if "lte" in condition and str(values.get(key)) > str(condition["lte"]): + return False + elif values.get(key) != condition: + return False + return True + + +@dataclass(frozen=True, slots=True) +class Table: + rows: tuple[Record | DailyRecord, ...] = () + + async def find_many( + self, + where: Mapping[str, object] | None = None, + take: int | None = None, + skip: int = 0, + order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None, + ) -> tuple[Record | DailyRecord, ...]: + selected: Final = tuple(row for row in self.rows if matches(row.model_dump(), where or {})) + orders: Final = (order,) if isinstance(order, Mapping) else tuple(order or ()) + ordered: Final = reduce( + lambda rows, item: tuple( + sorted( + rows, + key=lambda row: str(row.model_dump()[next(iter(item))]), + reverse=next(iter(item.values())) == "desc", + ) + ), + reversed(orders), + selected, + ) + return ordered[skip:] if take is None else ordered[skip : skip + take] + + async def count(self, where: Mapping[str, object] | None = None) -> int: + return len(await self.find_many(where=where)) + + +@dataclass(frozen=True, slots=True) +class Database: + litellm_dailycustomermemberspend: Table + litellm_customerorganizationtable: Table = Table( + ( + Record(customer_org_id="org-a", alias="Organization A"), + Record(customer_org_id="org-b", alias="Organization B"), + ) + ) + litellm_customerteamtable: Table = Table( + (Record(customer_team_id="team-a", alias="Product"), Record(customer_team_id="team-b", alias="Operations")) + ) + litellm_customermembertable: Table = Table( + ( + Record(customer_member_id="member-a", alias="Member A"), + Record(customer_member_id="member-b", alias="Member B"), + ) + ) + litellm_verificationtoken: Table = Table( + ( + Record(token="hash-a", key_alias="Support", team_id=None, user_id=None, metadata=None), + Record(token="hash-b", key_alias="Search", team_id=None, user_id=None, metadata=None), + ) + ) + litellm_deletedverificationtoken: Table = Table() + litellm_usertable: Table = Table() + + +@dataclass(frozen=True, slots=True) +class Client: + db: Database + + +def usage_row( + *, + id: str = "usage-a", + date: str = "2026-01-01", + customer_org_id: str = "org-a", + customer_team_id: str = "team-a", + customer_member_id: str = "member-a", + api_key: str = "hash-a", + model: str = "model-a", + model_group: str = "alias-a", + provider: str = "provider-a", + spend: float = 2.0, +) -> DailyRecord: + return DailyRecord( + id=id, + date=date, + customer_org_id=customer_org_id, + customer_team_id=customer_team_id, + customer_member_id=customer_member_id, + api_key=api_key, + model=model, + model_group=model_group, + custom_llm_provider=provider, + spend=spend, + prompt_tokens=10, + completion_tokens=5, + total_tokens=15, + api_requests=3, + successful_requests=2, + failed_requests=1, + ) + + +def client(rows: tuple[DailyRecord, ...] = (), role: LitellmUserRoles = LitellmUserRoles.PROXY_ADMIN) -> TestClient: + app: Final = FastAPI() + app.include_router(router) + + @app.exception_handler(ManagementProblem) + async def problem(request: Request, exc: ManagementProblem): + return problem_response(exc.problem) + + auth: Final = UserAPIKeyAuth(user_id="admin", user_role=role) + prisma: Final = cast(PrismaClient, Client(Database(Table(rows)))) # cast-ok: in-memory Prisma boundary + app.dependency_overrides[user_api_key_auth] = lambda: auth + app.dependency_overrides[get_customer_usage_prisma] = lambda: prisma + return TestClient(app) + + +DATES: Final = {"start_date": "2026-01-01", "end_date": "2026-01-03"} +SCOPES: Final = ( + ("customer_organizations", "customer_org_ids", "exclude_customer_org_ids", "org-a", "org-b", "Organization A"), + ("customer_teams", "customer_team_ids", "exclude_customer_team_ids", "team-a", "team-b", "Product"), + ("customer_members", "customer_member_ids", "exclude_customer_member_ids", "member-a", "member-b", "Member A"), +) + + +@pytest.mark.parametrize("resource,parameter,excluded,first,second,label", SCOPES) +def test_standard_daily_activity_filters_entities_and_preserves_aliases_and_dimensions( + resource: str, parameter: str, excluded: str, first: str, second: str, label: str +) -> None: + rows: Final = ( + usage_row(), + usage_row(id="usage-c", date="2026-01-03", model_group="alias-b", provider="provider-b", spend=4), + usage_row( + id="usage-b", customer_org_id="org-b", customer_team_id="team-b", customer_member_id="member-b", spend=10 + ), + usage_row(id="outside-range", date="2025-12-31", spend=50), + ) + response: Final = client(rows).get(f"/management/v1/{resource}/daily/activity", params={**DATES, parameter: first}) + assert response.status_code == 200, response.text + body: Final = response.json() + assert set(body) == {"results", "metadata"} + assert body["metadata"]["total_spend"] == 6 + assert body["metadata"]["total_tokens"] == 30 + assert body["metadata"]["total_api_requests"] == 6 + assert body["metadata"]["has_more"] is False + day: Final = body["results"][0] + assert day["breakdown"]["entities"][first]["metadata"]["alias"] == label + assert day["breakdown"]["models"]["model-a"]["metrics"]["spend"] == 4 + assert day["breakdown"]["model_groups"]["alias-b"]["metrics"]["spend"] == 4 + assert day["breakdown"]["providers"]["provider-b"]["metrics"]["spend"] == 4 + assert day["breakdown"]["api_keys"]["hash-a"]["metadata"]["key_alias"] == "Support" + + +@pytest.mark.parametrize("resource,parameter,excluded,first,second,label", SCOPES) +def test_inclusion_and_exclusion_lists_are_comma_separated( + resource: str, parameter: str, excluded: str, first: str, second: str, label: str +) -> None: + rows: Final = ( + usage_row(), + usage_row( + id="usage-b", customer_org_id="org-b", customer_team_id="team-b", customer_member_id="member-b", spend=10 + ), + ) + response: Final = client(rows).get( + f"/management/v1/{resource}/daily/activity", params={**DATES, parameter: f"{first},{second}", excluded: second} + ) + assert response.status_code == 200, response.text + assert response.json()["metadata"]["total_spend"] == 2 + + +@pytest.mark.parametrize("resource,parameter,excluded,first,second,label", SCOPES) +def test_standard_model_key_and_page_filters_use_the_shared_response_contract( + resource: str, parameter: str, excluded: str, first: str, second: str, label: str +) -> None: + rows: Final = ( + usage_row(id="page-a"), + usage_row(id="page-b", date="2026-01-02", spend=4), + usage_row(id="other-model", model="model-b", spend=20), + usage_row(id="other-key", api_key="hash-b", spend=30), + ) + response: Final = client(rows).get( + f"/management/v1/{resource}/daily/activity", + params={**DATES, parameter: first, "model": "model-a", "api_key": "hash-a", "page": 1, "page_size": 1}, + ) + assert response.status_code == 200, response.text + body: Final = response.json() + assert body["metadata"]["page"] == 1 + assert body["metadata"]["total_pages"] == 2 + assert body["metadata"]["has_more"] is True + assert body["metadata"]["total_spend"] == 4 + assert body["results"][0]["date"] == "2026-01-02" + next_page: Final = client(rows).get( + f"/management/v1/{resource}/daily/activity", + params={**DATES, parameter: first, "model": "model-a", "api_key": "hash-a", "page": 2, "page_size": 1}, + ) + assert next_page.status_code == 200, next_page.text + assert next_page.json()["metadata"]["total_spend"] == 2 + assert next_page.json()["metadata"]["has_more"] is False + + +@pytest.mark.parametrize("resource,parameter,excluded,first,second,label", SCOPES) +@pytest.mark.parametrize("role", (LitellmUserRoles.INTERNAL_USER, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY)) +def test_all_hierarchy_daily_activity_routes_are_admin_only( + resource: str, parameter: str, excluded: str, first: str, second: str, label: str, role: LitellmUserRoles +) -> None: + response: Final = client(role=role).get(f"/management/v1/{resource}/daily/activity", params=DATES) + assert response.status_code == 403, response.text + + +@pytest.mark.parametrize("params", ({}, {"start_date": "invalid", "end_date": "2026-01-03"}, {**DATES, "page": 0})) +def test_shared_helper_rejects_missing_dates_invalid_dates_and_nonpositive_pagination( + params: Mapping[str, object], +) -> None: + response: Final = client().get("/management/v1/customer_members/daily/activity", params=params) + assert response.status_code == 400, response.text + + +def test_unmatched_entities_return_the_standard_empty_response() -> None: + response: Final = client((usage_row(),)).get( + "/management/v1/customer_members/daily/activity", params={**DATES, "customer_member_ids": "missing"} + ) + assert response.status_code == 200, response.text + assert response.json()["results"] == [] + assert response.json()["metadata"]["total_spend"] == 0 diff --git a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py index ea36ffad044..e7bc5675e06 100644 --- a/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_internal_user_endpoints.py @@ -41,6 +41,7 @@ from litellm.proxy.management_endpoints.internal_user_endpoints import ( ) from litellm.proxy.proxy_server import app from litellm.types.proxy.auth.auth_checks import UserNotFoundError +from litellm.types.proxy.common_proxy_errors import CommonProxyErrors from litellm.types.proxy.management_endpoints.internal_user_endpoints import InsensitiveContains from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import ( CascadingJWTMappingTable, @@ -50,6 +51,30 @@ from tests.unit.proxy.management_endpoints.jwt_key_mapping_doubles import ( client = TestClient(app) +@pytest.mark.asyncio +@pytest.mark.parametrize("binding", ["customer_org_id", "customer_team_id", "customer_member_id"]) +async def test_new_user_cannot_mint_customer_bound_keys_without_proxy_admin(binding: str) -> None: + with pytest.raises(HTTPException) as caught: + await new_user( + data=NewUserRequest(metadata={binding: "another-customer"}), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.INTERNAL_USER), + ) + assert caught.value.status_code == 403 + assert caught.value.detail == "Only proxy admins can change a key's customer organization, team, or member binding." + + +@pytest.mark.asyncio +async def test_new_user_reports_missing_database_after_customer_binding_validation(mocker: MockerFixture) -> None: + mocker.patch("litellm.proxy.proxy_server.prisma_client", None) + with pytest.raises(ProxyException) as caught: + await new_user( + data=NewUserRequest(metadata={"customer_org_id": "org-1"}), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert str(caught.value.code) == "400" + assert caught.value.message == CommonProxyErrors.db_not_connected_error.value + + @pytest.mark.parametrize("user_id", ["", " \t\n", "user\x00id", 123, True, [], {}]) def test_new_user_rejects_invalid_custom_id(user_id: object) -> None: with pytest.raises(ValidationError) as exc: diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index c67cb974583..9e0e123436a 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -5,7 +5,7 @@ from collections.abc import Mapping from contextlib import ExitStack from datetime import datetime, timedelta, timezone from types import SimpleNamespace -from typing import Final +from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, call, patch from uuid import UUID, uuid4 @@ -24618,3 +24618,89 @@ def test_user_new_without_a_key_does_not_check_guardrail_names(spend_permission_ assert response.status_code == 200, response.text find_many.assert_not_awaited() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("field", ("customer_org_id", "customer_team_id", "customer_member_id")) +async def test_customer_binding_requires_admin_on_key_creation(field: str) -> None: + with pytest.raises(HTTPException) as rejected: + await _common_key_generation_helper( + data=GenerateKeyRequest(metadata={field: "customer-id"}), + user_api_key_dict=UserAPIKeyAuth(user_id="member", user_role=LitellmUserRoles.INTERNAL_USER), + litellm_changed_by=None, + team_table=None, + ) + assert rejected.value.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("request_type", ("key", "user")) +async def test_direct_key_mint_cannot_bypass_customer_binding_authorization(request_type: Literal["key", "user"]) -> None: + with pytest.raises(HTTPException) as rejected: + await generate_key_helper_fn(request_type=request_type, metadata={"customer_org_id": "another-customer"}) + assert rejected.value.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ("update", "regenerate")) +@pytest.mark.parametrize("field", ("customer_org_id", "customer_team_id", "customer_member_id")) +@pytest.mark.parametrize("requested", ("another-customer-id", None)) +async def test_customer_binding_cannot_be_reassigned_or_cleared_by_key_holder( + operation: Literal["update", "regenerate"], field: str, requested: str | None +) -> None: + existing: Final = LiteLLM_VerificationToken(token="hashed-key", metadata={field: "customer-id"}) + data: Final = ( + UpdateKeyRequest(key="hashed-key", metadata={field: requested}) + if operation == "update" + else RegenerateKeyRequest(metadata={field: requested}) + ) + with pytest.raises(HTTPException) as rejected: + await prepare_key_update_data( + data=data, + existing_key_row=existing, + user_api_key_dict=UserAPIKeyAuth(user_id="member", user_role=LitellmUserRoles.INTERNAL_USER), + ) + assert rejected.value.status_code == 403 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ("update", "regenerate")) +@pytest.mark.parametrize("metadata", ({"application": "support"}, {}, None)) +async def test_unrelated_key_metadata_changes_preserve_existing_customer_binding( + operation: Literal["update", "regenerate"], metadata: dict[str, object] | None +) -> None: + bindings: Final = {"customer_org_id": "org-id", "customer_team_id": "team-id", "customer_member_id": "member-id"} + existing: Final = LiteLLM_VerificationToken(token="hashed-key", metadata={**bindings, "old": "discard"}) + data: Final = ( + UpdateKeyRequest(key="hashed-key", metadata=metadata) + if operation == "update" + else RegenerateKeyRequest(metadata=metadata) + ) + actual: Final = await prepare_key_update_data( + data=data, + existing_key_row=existing, + user_api_key_dict=UserAPIKeyAuth(user_id="member", user_role=LitellmUserRoles.INTERNAL_USER), + ) + assert actual["metadata"] == {**(metadata or {}), **bindings} + assert existing.metadata == {**bindings, "old": "discard"} + assert data.metadata == metadata + + +@pytest.mark.asyncio +@pytest.mark.parametrize("operation", ("update", "regenerate")) +@pytest.mark.parametrize("metadata", ({"customer_org_id": "another-org"}, {}, None)) +async def test_proxy_admin_can_replace_or_clear_customer_key_binding( + operation: Literal["update", "regenerate"], metadata: dict[str, object] | None +) -> None: + existing: Final = LiteLLM_VerificationToken(token="hashed-key", metadata={"customer_org_id": "org-id"}) + data: Final = ( + UpdateKeyRequest(key="hashed-key", metadata=metadata) + if operation == "update" + else RegenerateKeyRequest(metadata=metadata) + ) + actual: Final = await prepare_key_update_data( + data=data, + existing_key_row=existing, + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + assert actual["metadata"] == metadata diff --git a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py index 337aa2939dd..f4138260e28 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/unit/proxy/spend_tracking/test_spend_management_endpoints.py @@ -4097,7 +4097,7 @@ class TestSpendLogsPayload: "model": "gpt-4o", "user": "", "team_id": "", - "metadata": '{"user_api_key_user_email": null, "actor_agent_id": null, "target_agent_id": null, "billing_agent_id": null, "agent_execution_mode": null, "verified_human_user_id": null, "applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "used_client_oauth_token": null, "litellm_roi_estimator": false, "billed_retry_attempt": false, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', + "metadata": '{"user_api_key_customer_usage": null, "user_api_key_user_email": null, "actor_agent_id": null, "target_agent_id": null, "billing_agent_id": null, "agent_execution_mode": null, "verified_human_user_id": null, "applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "autorouter_savings_estimate": null, "autorouter_baseline_observation": null, "azure_spillover": null, "used_client_oauth_token": null, "litellm_roi_estimator": false, "billed_retry_attempt": false, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30,