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 b1d7d9e7357..d757cfed0e4 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, utc_date_from_start_time @@ -221,10 +222,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 ab45935da0b..86fbf94cb90 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/cost_calculator.py b/litellm/cost_calculator.py index 91529ae9ff5..03feb8dc1e6 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -1144,11 +1144,29 @@ def _served_or_priced_requested_model( if served in (requested, requested.partition("/")[2]): return served pricing_provider: Final = _get_provider_for_cost_calc(requested, custom_llm_provider) + profile_of_served: Final = _is_bedrock_cross_region_profile_of(requested_model, served_model, pricing_provider) + if profile_of_served and _resolves_to_rates(requested, pricing_provider): + return requested if _resolves_to_rates(served, pricing_provider): return served return requested if _resolves_to_rates(requested, pricing_provider) else served +def _is_bedrock_cross_region_profile_of(profile: str, served_model: str, custom_llm_provider: str | None) -> bool: + if custom_llm_provider != "bedrock": + return False + from litellm.llms.bedrock.common_utils import ( + get_bedrock_cross_region_inference_regions, + strip_bedrock_routing_prefix, + ) + + region, _, profiled_model = strip_bedrock_routing_prefix(profile).partition(".") + return ( + region in get_bedrock_cross_region_inference_regions() + and strip_bedrock_routing_prefix(served_model) == profiled_model + ) + + def _strip_unregistered_leading_segments(model: str, region_name: str | None) -> str: """Resolve a provider-prefixed slash alias like "vertex_ai/vertex/claude-opus-5" to the registered cost key ("vertex_ai/claude-opus-5"), keeping the model unchanged when it already diff --git a/litellm/enterprise/pyproject.toml b/litellm/enterprise/pyproject.toml index 7a8056e8f15..0fdc674aea7 100644 --- a/litellm/enterprise/pyproject.toml +++ b/litellm/enterprise/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm-enterprise" -version = "0.1.83" +version = "0.1.84" description = "Package for LiteLLM Enterprise features" readme = "README.md" requires-python = ">=3.9" @@ -26,7 +26,7 @@ required-version = ">=0.10.9" module-root = "" [tool.commitizen] -version = "0.1.83" +version = "0.1.84" version_files = [ "pyproject.toml:^version", "../../pyproject.toml:litellm-enterprise==", diff --git a/litellm/integrations/otel/model/spans.py b/litellm/integrations/otel/model/spans.py index 5a3fbc93a33..b88a5d75b6e 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/llms/base_llm/embedding/transformation.py b/litellm/llms/base_llm/embedding/transformation.py index 46ac3ccf433..d25a0c3d524 100644 --- a/litellm/llms/base_llm/embedding/transformation.py +++ b/litellm/llms/base_llm/embedding/transformation.py @@ -1,7 +1,9 @@ from abc import ABC, abstractmethod -from typing import TYPE_CHECKING, Any +from collections.abc import Mapping +from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter from litellm.llms.base_llm.chat.transformation import BaseConfig from litellm.types.llms.openai import AllEmbeddingInputValues, AllMessageValues @@ -16,6 +18,24 @@ else: LiteLLMLoggingObj = Any +_BODY_FIELDS: Final = TypeAdapter(dict[str, object]) + + +def extra_body_fields(optional_params: Mapping[str, object] | None) -> Mapping[str, object]: + """The caller's ``extra_body`` fields, merged under the provider body after the transform so they add fields but never + replace the input or a mapped param.""" + extra_body: Final = (optional_params or {}).get("extra_body") + return _BODY_FIELDS.validate_python(extra_body) if isinstance(extra_body, Mapping) else {} + + +def split_extra_body( + optional_params: Mapping[str, object] | None, +) -> tuple[Mapping[str, object], Mapping[str, object]]: + """The params without ``extra_body``, and the ``extra_body`` fields to merge into the provider body.""" + params: Final = _BODY_FIELDS.validate_python(optional_params or {}) + return {k: v for k, v in params.items() if k != "extra_body"}, extra_body_fields(params) + + class BaseEmbeddingConfig(BaseConfig, ABC): @abstractmethod def transform_embedding_request( diff --git a/litellm/llms/bedrock/embed/embedding.py b/litellm/llms/bedrock/embed/embedding.py index e28e39ef2d4..0743bd6e7e7 100644 --- a/litellm/llms/bedrock/embed/embedding.py +++ b/litellm/llms/bedrock/embed/embedding.py @@ -6,6 +6,7 @@ import copy import json import urllib.parse from collections.abc import Callable, Mapping +from types import MappingProxyType from typing import TYPE_CHECKING, ClassVar, Final, get_args, overload import httpx @@ -15,6 +16,7 @@ import litellm from litellm.constants import BEDROCK_EMBEDDING_PROVIDERS_LITERAL from litellm.litellm_core_utils.hidden_params import set_hidden_param from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.base_llm.embedding.transformation import extra_body_fields from litellm.llms.cohere.embed.handler import embedding as cohere_embedding from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -103,7 +105,7 @@ class BedrockEmbedding(BaseAWSLLM): region_from_model: str | None = None, ) -> tuple[Credentials | None, str]: auth_params: Final = pop_aws_auth_params(optional_params) - explicit_region_name: Final = optional_params.pop("aws_region_name", None) + explicit_region_name: Final = _OPTIONAL_STR.validate_python(optional_params.pop("aws_region_name", None)) aws_region_name = region_from_model if explicit_region_name is None else explicit_region_name ### SET REGION NAME ### @@ -291,9 +293,10 @@ class BedrockEmbedding(BaseAWSLLM): provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL, api_key: str | None = None, is_async_invoke: bool | None = False, + extra_body: Mapping[str, object] = MappingProxyType({}), ): responses: Final[list[dict]] = [] - for data in batch_data: + for data in ({**extra_body, **item} for item in batch_data): headers = {"Content-Type": "application/json"} if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} @@ -360,9 +363,10 @@ class BedrockEmbedding(BaseAWSLLM): provider: BEDROCK_EMBEDDING_PROVIDERS_LITERAL, api_key: str | None = None, is_async_invoke: bool | None = False, + extra_body: Mapping[str, object] = MappingProxyType({}), ): responses: Final[list[dict]] = [] - for data in batch_data: + for data in ({**extra_body, **item} for item in batch_data): headers = {"Content-Type": "application/json"} if extra_headers is not None: headers = {"Content-Type": "application/json", **extra_headers} @@ -462,9 +466,12 @@ class BedrockEmbedding(BaseAWSLLM): ) inference_params = copy.deepcopy(optional_params) inference_params = { - k: v for k, v in inference_params.items() if k.lower() not in self.aws_authentication_params + k: v + for k, v in inference_params.items() + if k.lower() not in self.aws_authentication_params and k != "extra_body" } inference_params.pop("user", None) # make sure user is not passed in for bedrock call + extra_body: Final = extra_body_fields(optional_params) data: CohereEmbeddingRequest | None = None batch_data: list | None = None @@ -535,6 +542,8 @@ class BedrockEmbedding(BaseAWSLLM): ) batch_data.append(nova_request) + request_data: Final = None if data is None else {**extra_body, **data} + ### SET RUNTIME ENDPOINT ### endpoint_url, proxy_endpoint_url = self.get_runtime_endpoint( api_base=api_base, @@ -552,6 +561,7 @@ class BedrockEmbedding(BaseAWSLLM): client=(client if client is not None and isinstance(client, AsyncHTTPHandler) else None), timeout=timeout, batch_data=batch_data, + extra_body=extra_body, credentials=credentials, extra_headers=extra_headers, endpoint_url=endpoint_url, @@ -566,6 +576,7 @@ class BedrockEmbedding(BaseAWSLLM): client=(client if client is not None and isinstance(client, HTTPHandler) else None), timeout=timeout, batch_data=batch_data, + extra_body=extra_body, credentials=credentials, extra_headers=extra_headers, endpoint_url=endpoint_url, @@ -579,7 +590,7 @@ class BedrockEmbedding(BaseAWSLLM): if returned_response is None: raise Exception("Unable to map Bedrock request to provider") return returned_response - elif data is None: + elif request_data is None: raise Exception("Unable to map Bedrock request to provider") headers = {"Content-Type": "application/json"} @@ -591,7 +602,7 @@ class BedrockEmbedding(BaseAWSLLM): aws_region_name=aws_region_name, extra_headers=extra_headers, endpoint_url=endpoint_url, - data=json.dumps(data), + data=json.dumps(request_data), headers=headers, api_key=api_key, ) @@ -607,7 +618,7 @@ class BedrockEmbedding(BaseAWSLLM): logging_obj=logging_obj, optional_params=optional_params, encoding=encoding, - data=data, + data=request_data, complete_api_base=prepped.url, api_key=None, aembedding=aembedding, diff --git a/litellm/llms/nvidia_nim/embed.py b/litellm/llms/nvidia_nim/embed.py index d2ef446299e..7368831f7a8 100644 --- a/litellm/llms/nvidia_nim/embed.py +++ b/litellm/llms/nvidia_nim/embed.py @@ -66,7 +66,7 @@ class NvidiaNimEmbeddingConfig: self, non_default_params: dict, optional_params: dict, - kwargs: dict | None = None, + kwargs: Mapping[str, object] | None = None, ): if "extra_body" not in optional_params: optional_params["extra_body"] = {} @@ -80,5 +80,11 @@ class NvidiaNimEmbeddingConfig: if kwargs is not None: # pass kwargs in extra_body - optional_params["extra_body"].update(kwargs) + nested_extra_body: Final = kwargs.get("extra_body") + optional_params["extra_body"].update( + { + **{k: v for k, v in kwargs.items() if k != "extra_body"}, + **(nested_extra_body if isinstance(nested_extra_body, Mapping) else {}), + } + ) return optional_params diff --git a/litellm/llms/ollama/completion/handler.py b/litellm/llms/ollama/completion/handler.py index e1fc9a8d262..fd874246980 100644 --- a/litellm/llms/ollama/completion/handler.py +++ b/litellm/llms/ollama/completion/handler.py @@ -14,6 +14,7 @@ from typing_extensions import NotRequired, ReadOnly import litellm from litellm.litellm_core_utils.core_helpers import require_attributes from litellm.litellm_core_utils.url_utils import ensure_api_route +from litellm.llms.base_llm.embedding.transformation import split_extra_body from litellm.types.mcp import has_header from litellm.types.utils import EmbeddingResponse @@ -45,14 +46,15 @@ def _prepare_ollama_embedding_payload( data: Final[dict[str, object]] = {"model": model, "input": prompts} special_optional_params: Final = ["truncate", "options", "keep_alive", "dimensions"] - for k, v in optional_params.items(): + body_params, extra_body = split_extra_body(optional_params) + for k, v in body_params.items(): if k in special_optional_params: data[k] = v else: data.setdefault("options", {}) if isinstance(data["options"], dict): data["options"].update({k: v}) - return data + return {**extra_body, **data} def _process_ollama_embedding_response( diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index a852ae8b1a1..c20f0df6768 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -14,6 +14,7 @@ from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.optional_imports import ensure_optional_import +from litellm.llms.base_llm.embedding.transformation import extra_body_fields from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM, pop_aws_auth_params from litellm.llms.bedrock.web_identity_session_policy import SAGEMAKER_SESSION_POLICY, WebIdentitySessionPolicy from litellm.llms.custom_httpx.http_handler import ( @@ -594,8 +595,11 @@ class SagemakerLLM(BaseAWSLLM): inference_component_name: Final = embedding_params.get("model_id") if inference_component_name is not None and not isinstance(inference_component_name, str): raise SagemakerError(status_code=400, message="model_id must be the inference component name as a string") - body_params: Final = {k: v for k, v in embedding_params.items() if k != "model_id"} - request_data: Final = provider_config.transform_embedding_request(endpoint_name, input, body_params, {}) + body_params: Final = {k: v for k, v in embedding_params.items() if k not in ("model_id", "extra_body")} + request_data: Final = { + **extra_body_fields(embedding_params), + **provider_config.transform_embedding_request(endpoint_name, input, body_params, {}), + } data: Final = json.dumps(request_data).encode("utf-8") invoke_args: Final = _InvokeEndpointArgs( EndpointName=endpoint_name, diff --git a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py index 96a785c06c9..b27003369f7 100644 --- a/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py +++ b/litellm/llms/vertex_ai/gemini_embeddings/batch_embed_content_transformation.py @@ -10,6 +10,7 @@ from typing import Annotated, Final, Literal, cast from pydantic import ConfigDict, Field, TypeAdapter, ValidationError from litellm.exceptions import BadRequestError +from litellm.llms.base_llm.embedding.transformation import extra_body_fields from litellm.llms.vertex_ai.common_utils import GEMINI_FILES_API_URI_PREFIX, gemini_video_metadata_from_openai from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.vertex_ai import ( @@ -414,6 +415,7 @@ def transform_openai_input_gemini_content( gemini_model_name: Final = f"models/{model}" gemini_params: Final = _filter_embed_params(optional_params) + extra_body: Final = extra_body_fields(optional_params) input_list: Final = _top_level_elements(input) requests: Final[list[EmbedContentRequest]] = [] @@ -432,7 +434,11 @@ def transform_openai_input_gemini_content( content=ContentType(parts=parts), **gemini_params, ) - requests.append(request) + requests.append( + cast( # cast-ok: extra_body carries raw Gemini fields the TypedDict does not list + EmbedContentRequest, {**extra_body, **request} + ) + ) return VertexAIBatchEmbeddingsRequestBody(requests=requests) @@ -472,6 +478,7 @@ def transform_openai_input_gemini_embed_content( parts.append(_build_part_for_input(element, resolved_files=resolved_files, drop_params=drop_params)) request_body: Final[dict] = { + **extra_body_fields(optional_params), "content": ContentType(parts=parts), **gemini_params, } diff --git a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py index c282ae22ea6..7c276b43dc7 100644 --- a/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/multimodal_embeddings/transformation.py @@ -21,7 +21,7 @@ from litellm.types.utils import ( ) from litellm.utils import count_characters, is_base64_encoded -from ...base_llm.embedding.transformation import BaseEmbeddingConfig +from ...base_llm.embedding.transformation import BaseEmbeddingConfig, extra_body_fields from ..common_utils import VertexAIError @@ -189,8 +189,11 @@ class VertexAIMultimodalEmbeddingConfig(BaseEmbeddingConfig): request_data["instances"] = [vertex_request_instance] + extra_body: Final = extra_body_fields(optional_params) if "outputDimensionality" in optional_params: - request_data["parameters"] = {"dimension": optional_params["outputDimensionality"]} + request_data["parameters"] = {**extra_body, "dimension": optional_params["outputDimensionality"]} + elif extra_body: + request_data["parameters"] = dict(extra_body) return cast(dict, request_data) diff --git a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py index 96fccf007a5..5d99ce6d90f 100644 --- a/litellm/llms/vertex_ai/vertex_embeddings/transformation.py +++ b/litellm/llms/vertex_ai/vertex_embeddings/transformation.py @@ -1,10 +1,11 @@ import types from collections.abc import Mapping -from typing import Final, Literal +from typing import Final, Literal, cast # noqa: TID251 # extra_body carries raw fields the TypedDict lacks from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict +from litellm.llms.base_llm.embedding.transformation import split_extra_body from litellm.llms.vertex_ai.common_utils import pop_vertex_request_labels from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.utils import EmbeddingResponse, PromptTokensDetailsWrapper, Usage @@ -164,7 +165,12 @@ class VertexAITextEmbeddingConfig(LiteLLMBaseModel): vertex_text_embedding_input_list.append(embedding_input) vertex_request["instances"] = vertex_text_embedding_input_list - vertex_request["parameters"] = EmbeddingParameters(**optional_params) + body_params, extra_body = split_extra_body(optional_params) + vertex_request["parameters"] = ( + cast( # cast-ok: extra_body carries raw Vertex fields the TypedDict does not list + EmbeddingParameters, {**extra_body, **body_params} + ) + ) if labels: vertex_request["labels"] = labels @@ -205,7 +211,12 @@ class VertexAITextEmbeddingConfig(LiteLLMBaseModel): vertex_text_embedding_input_list.append(embedding_input) vertex_request["instances"] = vertex_text_embedding_input_list - vertex_request["parameters"] = TextEmbeddingFineTunedParameters(**optional_params) + body_params, extra_body = split_extra_body(optional_params) + vertex_request["parameters"] = ( + cast( # cast-ok: extra_body carries raw Vertex fields the TypedDict does not list + TextEmbeddingFineTunedParameters, {**extra_body, **body_params} + ) + ) # Remove 'shared_session' from parameters if present if vertex_request["parameters"] is not None and "shared_session" in vertex_request["parameters"]: del vertex_request["parameters"]["shared_session"] diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 4429fffe652..7c8a48f5aae 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -28650,7 +28650,7 @@ "mode": "chat", "output_cost_per_token": 2.8e-07, "source": "https://docs.fireworks.ai/serverless/pricing", - "successor_model": "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-0731", + "successor_model": "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -29124,7 +29124,7 @@ "mode": "chat", "output_cost_per_token": 2.8e-07, "source": "https://docs.fireworks.ai/serverless/pricing", - "successor_model": "fireworks_ai/deepseek-v4-flash-0731", + "successor_model": "fireworks_ai/deepseek-v4p1-flash", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -83697,7 +83697,7 @@ "mode": "chat", "output_cost_per_token": 2.1e-07, "source": "https://api.fireworks.ai/v1/serverless/models", - "successor_model": "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-0731", + "successor_model": "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, 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 9fd5b3238a2..95ce61ff7b8 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 23016445c50..d4171887c65 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/hooks/dynamic_rate_limiter_v3.py b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py index 73718c356f1..3bfc5ae5aa2 100644 --- a/litellm/proxy/hooks/dynamic_rate_limiter_v3.py +++ b/litellm/proxy/hooks/dynamic_rate_limiter_v3.py @@ -393,6 +393,20 @@ def _tightest_team_share(descriptors: Iterable[RateLimitDescriptor | None]) -> t return min(shares, key=lambda share: share[1], default=None) +def _tightest_of(*shares: tuple[RateLimitDescriptor, int] | None) -> tuple[RateLimitDescriptor, int] | None: + return min((share for share in shares if share is not None), key=lambda share: share[1], default=None) + + +def _capped_idle_share( + share: RateLimitDescriptor | None, cap: RateLimitDescriptor | None +) -> tuple[RateLimitDescriptor, int] | None: + """A capped team's usage while the model is idle sits on its cap counter, read against its share.""" + tightest: Final = _tightest_team_share((share,)) + if tightest is None or cap is None: + return None + return cap, tightest[1] + + def _usage_of(kwargs: Mapping[str, object], response_obj: object) -> Usage | None: """Usage of a modeled response, a Responses WebSocket turn's combined usage, or its raw event list.""" usage: Final = getattr(response_obj, "usage", None) @@ -431,7 +445,7 @@ class _BorrowingTeamShare: ceiling: RateLimitDescriptor | None def enforced(self, saturated: bool) -> RateLimitDescriptor | None: - return self.share if saturated else self.ceiling + return self.share if saturated else None def _team_limit_name(rate_limit_type: str) -> str: @@ -483,15 +497,25 @@ def _team_model_settlement( return () served: Final = stash.borrowing_team_model_fallback_limits.get(model_group) if model_group is not None else None admitted: Final = stash.borrowing_team_model_limit + cap: Final = stash.borrowing_team_model_cap reserved_tokens: Final = stash.borrowing_team_model_reserved_tokens + cap_reserved_tokens: Final = stash.borrowing_team_model_cap_reserved_tokens window: Final = stash.borrowing_team_model_reserved_window + cap_window: Final = stash.borrowing_team_model_cap_reserved_window if served is None: - return _window_guarded(_team_model_token_increment(admitted, total_tokens - reserved_tokens, ttl), window) + return ( + *_window_guarded(_team_model_token_increment(admitted, total_tokens - reserved_tokens, ttl), window), + *_window_guarded(_team_model_token_increment(cap, total_tokens - cap_reserved_tokens, ttl), cap_window), + ) return ( *_window_guarded(_team_model_token_increment(admitted, -reserved_tokens, ttl), window), + *_window_guarded(_team_model_token_increment(cap, -cap_reserved_tokens, ttl), cap_window), *_window_guarded( _team_model_token_increment(served.descriptor, total_tokens - served.reserved_tokens, ttl), served.window ), + *_window_guarded( + _team_model_token_increment(served.cap, total_tokens - served.cap_reserved_tokens, ttl), served.cap_window + ), ) @@ -861,12 +885,16 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return _BorrowingTeamShare( limits=limits, share=self._team_share_descriptor(value, limits), - ceiling=None if ceiling is None else self._team_share_descriptor(value, ceiling), + ceiling=None + if ceiling is None + else self._team_share_descriptor(value, ceiling, TEAM_MODEL_LIMIT_MAX_BORROW_MULTIPLE), ) - def _team_share_descriptor(self, value: str, limits: BorrowingTeamModelLimits) -> RateLimitDescriptor: + def _team_share_descriptor( + self, value: str, limits: BorrowingTeamModelLimits, key: str = TEAM_MODEL_LIMIT_AT_SATURATION + ) -> RateLimitDescriptor: return RateLimitDescriptor( - key=TEAM_MODEL_LIMIT_AT_SATURATION, + key=key, value=value, rate_limit={ "tokens_per_unit": limits.tpm, @@ -884,7 +912,7 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): user_api_key_dict: UserAPIKeyAuth, team_share: _BorrowingTeamShare, ) -> ProxyRateLimitError: - if saturation >= _get_priority_settings().saturation_threshold: + if status["descriptor_key"] != TEAM_MODEL_LIMIT_MAX_BORROW_MULTIPLE: return self._team_limit_error(model, status, saturation, priority, user_api_key_dict) return self._team_borrow_ceiling_error( model, status, saturation, priority, user_api_key_dict, team_share.limits @@ -925,8 +953,13 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): loggers: Final = _prometheus_loggers() if not loggers or stash is None or model_group is None or settlement.user_api_key_team_id is None: return - share: Final = _tightest_team_share( - (stash.borrowing_team_model_limit, stash.priority_class_share) if hold is None else (hold.descriptor,) + share: Final = ( + _tightest_of(_tightest_team_share((hold.descriptor,)), _capped_idle_share(hold.share, hold.cap)) + if hold is not None + else _tightest_of( + _tightest_team_share((stash.borrowing_team_model_limit, stash.priority_class_share)), + _capped_idle_share(stash.borrowing_team_model_share, stash.borrowing_team_model_cap), + ) ) if share is None: return @@ -1081,10 +1114,19 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): saturation: Final = await self._check_model_saturation(requested_model, model_group_info) saturated: Final = saturation >= _get_priority_settings().saturation_threshold enforced_team: Final = None if borrowing is None else borrowing.enforced(saturated) + team_cap: Final = None if borrowing is None else borrowing.ceiling if connection_stash is not None and borrowing is not None: - connection_stash.borrowing_team_model_limit = enforced_team - descriptors: Final = (share_descriptors if saturated else ()) + ( - () if enforced_team is None else (enforced_team,) + connection_stash.borrowing_team_model_limit = ( + connection_stash.borrowing_team_model_limit + if enforced_team is None and connection_stash.borrowing_team_model_reserved_tokens + else enforced_team + ) + connection_stash.borrowing_team_model_share = borrowing.share + connection_stash.borrowing_team_model_cap = team_cap + descriptors: Final = ( + (share_descriptors if saturated else ()) + + (() if enforced_team is None else (enforced_team,)) + + (() if team_cap is None else (team_cap,)) ) if not descriptors: return @@ -1229,28 +1271,34 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): if team_share is not None: descriptors_to_check.append(team_share.share) enforced_team_share: Final = None if team_share is None else team_share.enforced(should_enforce_priority) + team_cap: Final = None if team_share is None else team_share.ceiling + held_team_share: Final = enforced_team_share if enforced_team_share is not None else team_cap get_or_create_request_stash().borrowing_team_model_limit = enforced_team_share + get_or_create_request_stash().borrowing_team_model_share = None if team_share is None else team_share.share + get_or_create_request_stash().borrowing_team_model_cap = team_cap enforced_descriptors: Final[list[RateLimitDescriptor]] = [model_wide_descriptor] if priority_descriptors and should_enforce_priority: enforced_descriptors.extend(priority_descriptors) if enforced_team_share is not None: enforced_descriptors.append(enforced_team_share) + if team_cap is not None: + enforced_descriptors.append(team_cap) per_request_increment: Final[dict[Literal["requests", "tokens"], int]] = { "requests": 1, "tokens": 0, } team_reservation: Final = ( - await self._team_token_reservation(enforced_team_share, data, model, call_type, user_api_key_dict) - if enforced_team_share is not None + await self._team_token_reservation(team_share.share, data, model, call_type, user_api_key_dict) + if team_share is not None and held_team_share is not None else 0 ) atomic_response: Final = await self.v3_limiter.atomic_check_and_increment_by_n( descriptors=enforced_descriptors, increments=[ {"requests": 1, "tokens": team_reservation} - if descriptor is enforced_team_share + if descriptor is enforced_team_share or descriptor is team_cap else per_request_increment for descriptor in enforced_descriptors ], @@ -1310,7 +1358,10 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): llm_provider=llm_provider, ) - if descriptor_key == TEAM_MODEL_LIMIT_AT_SATURATION and team_share is not None: + if ( + descriptor_key in (TEAM_MODEL_LIMIT_AT_SATURATION, TEAM_MODEL_LIMIT_MAX_BORROW_MULTIPLE) + and team_share is not None + ): raise self._team_share_error(model, status, saturation, priority, user_api_key_dict, team_share) if descriptor_key == PRIORITY_CLASS_SHARE_AT_SATURATION and share is not None: raise self._class_share_error(model, status, saturation, priority, user_api_key_dict, share) @@ -1355,11 +1406,15 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): # enforced. The increment-then-check semantics of should_rate_limit # are fine because we ignore the OVER_LIMIT response. stash: Final = get_or_create_request_stash() - stash.borrowing_team_model_reserved_tokens = team_reservation + stash.borrowing_team_model_reserved_tokens = team_reservation if enforced_team_share is not None else 0 + stash.borrowing_team_model_cap_reserved_tokens = team_reservation if team_cap is not None else 0 stash.borrowing_team_model_reserved_window = ( _reserved_window(atomic_response, enforced_team_share) if team_reservation else None ) - self._cap_admitted_output(enforced_team_share, team_reservation, data, call_type) + stash.borrowing_team_model_cap_reserved_window = ( + _reserved_window(atomic_response, team_cap) if team_reservation else None + ) + self._cap_admitted_output(None if team_share is None else team_share.share, team_reservation, data, call_type) if priority_descriptors and share is not None and not should_enforce_priority: get_or_create_request_stash().rate_limit_response = atomic_response elif priority_descriptors and not should_enforce_priority: @@ -1434,33 +1489,48 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): return True saturation: Final = await self._check_model_saturation(model, model_group_info) enforced: Final = team_share.enforced(saturation >= _get_priority_settings().saturation_threshold) + cap: Final = team_share.ceiling + held: Final = enforced if enforced is not None else cap reservation: Final = ( await self._team_token_reservation( - enforced, + team_share.share, {**request_data, "model": model}, model, get_or_create_request_stash().call_type, user_api_key_dict, ) - if enforced is not None + if held is not None else 0 ) + team_descriptors: Final = tuple(d for d in (enforced, cap) if d is not None) if not commit: - return enforced is None or await self._share_has_room(enforced, reservation, user_api_key_dict) + for descriptor in team_descriptors: + if not await self._share_has_room(descriptor, reservation, user_api_key_dict): + return False + return True model_wide_descriptor: Final = self._create_model_tracking_descriptor( model=model, model_group_info=model_group_info, high_limit_multiplier=1 ) response: Final = await self.v3_limiter.atomic_check_and_increment_by_n( - descriptors=[model_wide_descriptor, *(() if enforced is None else (enforced,))], + descriptors=[model_wide_descriptor, *team_descriptors], increments=[ {"requests": 1, "tokens": 0}, - *(() if enforced is None else ({"requests": 1, "tokens": reservation},)), + *({"requests": 1, "tokens": reservation} for _ in team_descriptors), ], parent_otel_span=user_api_key_dict.parent_otel_span, ) if response["overall_code"] == "OVER_LIMIT": return False - self._hold_fallback_share(model, team_share.share, enforced, reservation, _reserved_window(response, enforced)) + self._hold_fallback_share( + model, + team_share.share, + enforced, + reservation if enforced is not None else 0, + _reserved_window(response, enforced), + cap, + reservation if cap is not None else 0, + _reserved_window(response, cap), + ) return True async def _share_has_room( @@ -1482,16 +1552,34 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): charged: RateLimitDescriptor | None, reserved_tokens: int, window: ReservationWindowIdentity | None, + cap: RateLimitDescriptor | None = None, + cap_reserved_tokens: int = 0, + cap_window: ReservationWindowIdentity | None = None, ) -> None: """Remember the hold under the requested name and the resolved group, so settlement finds it either way. ``charged`` is the share the target's usage settles to, or None for an idle target whose usage is forgiven. + ``cap`` is the team's borrow ceiling on the target, charged whether or not the target is saturated. Each + counter carries an earlier hold on the same target only while that hold sits in the counter's current window. """ stash: Final = get_or_create_request_stash() earlier: Final = stash.borrowing_team_model_fallback_limits.get(model) - carried: Final = earlier.reserved_tokens if earlier is not None and earlier.window == window else 0 + carried: Final = ( + earlier.reserved_tokens if earlier is not None and window is not None and earlier.window == window else 0 + ) + carried_cap: Final = ( + earlier.cap_reserved_tokens + if earlier is not None and cap_window is not None and earlier.cap_window == cap_window + else 0 + ) hold: Final = BorrowingTeamModelHold( - descriptor=charged, reserved_tokens=reserved_tokens + carried, window=window + descriptor=charged, + reserved_tokens=reserved_tokens + carried, + window=window, + cap=cap, + cap_reserved_tokens=cap_reserved_tokens + carried_cap, + cap_window=cap_window, + share=share, ) stash.borrowing_team_model_fallback_limits = MappingProxyType( {**stash.borrowing_team_model_fallback_limits, model: hold, share["value"].rpartition(":")[2]: hold} @@ -1814,7 +1902,9 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): ) if stash is not None: stash.borrowing_team_model_reserved_tokens = 0 + stash.borrowing_team_model_cap_reserved_tokens = 0 stash.borrowing_team_model_reserved_window = None + stash.borrowing_team_model_cap_reserved_window = None stash.borrowing_team_model_fallback_limits = MappingProxyType( {k: h for k, h in stash.borrowing_team_model_fallback_limits.items() if h is not served} ) @@ -1918,28 +2008,41 @@ class _PROXY_DynamicRateLimitHandlerV3(CustomLogger): _usage_of(callback_kwargs, None if isinstance(response_obj, BaseException) else response_obj), self.v3_limiter.get_rate_limit_type(), ) + own_usage: Final = used if failed_hold is None else 0 + own_delta: Final = own_usage - stash.borrowing_team_model_reserved_tokens + own_cap_delta: Final = own_usage - stash.borrowing_team_model_cap_reserved_tokens refund: Final = ( *_window_guarded( - _team_model_token_increment( - stash.borrowing_team_model_limit, - (used if failed_hold is None else 0) - stash.borrowing_team_model_reserved_tokens, - self.v3_limiter.window_size, - ), + _team_model_token_increment(stash.borrowing_team_model_limit, own_delta, self.v3_limiter.window_size), stash.borrowing_team_model_reserved_window, ), + *_window_guarded( + _team_model_token_increment(stash.borrowing_team_model_cap, own_cap_delta, self.v3_limiter.window_size), + stash.borrowing_team_model_cap_reserved_window, + ), *( () if failed_hold is None - else _window_guarded( - _team_model_token_increment( - failed_hold.descriptor, used - failed_hold.reserved_tokens, self.v3_limiter.window_size + else ( + *_window_guarded( + _team_model_token_increment( + failed_hold.descriptor, used - failed_hold.reserved_tokens, self.v3_limiter.window_size + ), + failed_hold.window, + ), + *_window_guarded( + _team_model_token_increment( + failed_hold.cap, used - failed_hold.cap_reserved_tokens, self.v3_limiter.window_size + ), + failed_hold.cap_window, ), - failed_hold.window, ) ), ) stash.borrowing_team_model_reserved_tokens = 0 + stash.borrowing_team_model_cap_reserved_tokens = 0 stash.borrowing_team_model_reserved_window = None + stash.borrowing_team_model_cap_reserved_window = None if failed_hold is not None: stash.borrowing_team_model_fallback_limits = MappingProxyType( {k: h for k, h in stash.borrowing_team_model_fallback_limits.items() if h is not failed_hold} diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 709b7d6f454..b8eeaeab217 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -581,6 +581,10 @@ class BorrowingTeamModelHold: descriptor: RateLimitDescriptor | None reserved_tokens: int = 0 window: ReservationWindowIdentity | None = None + cap: RateLimitDescriptor | None = None + cap_reserved_tokens: int = 0 + cap_window: ReservationWindowIdentity | None = None + share: RateLimitDescriptor | None = None @dataclass(slots=True) @@ -631,6 +635,10 @@ class RequestRateLimiterStash: borrowing_team_model_limit: RateLimitDescriptor | None = None borrowing_team_model_reserved_tokens: int = 0 borrowing_team_model_reserved_window: ReservationWindowIdentity | None = None + borrowing_team_model_share: RateLimitDescriptor | None = None + borrowing_team_model_cap: RateLimitDescriptor | None = None + borrowing_team_model_cap_reserved_tokens: int = 0 + borrowing_team_model_cap_reserved_window: ReservationWindowIdentity | None = None borrowing_team_model_fallback_limits: Mapping[str, BorrowingTeamModelHold] = field( default_factory=lambda: MappingProxyType({}) ) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 98025e0c4bf..5da902ee4ec 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, @@ -2716,7 +2727,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. @@ -2993,6 +3004,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 9fd5b3238a2..95ce61ff7b8 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 9fd5b3238a2..95ce61ff7b8 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/llms/bedrock.py b/litellm/types/llms/bedrock.py index e42a1a8f057..170f90c1adb 100644 --- a/litellm/types/llms/bedrock.py +++ b/litellm/types/llms/bedrock.py @@ -505,6 +505,8 @@ class CohereEmbeddingRequest(TypedDict, total=False): truncate: Literal["NONE", "START", "END"] embedding_types: Literal["float", "int8", "uint8", "binary", "ubinary"] output_dimension: int + max_tokens: ReadOnly[int] + priority: ReadOnly[int] class CohereEmbeddingRequestWithModel(CohereEmbeddingRequest): 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/components/key_info_utils.test.tsx b/litellm/ui/litellm-dashboard/src/components/key_info_utils.test.tsx index 490e188e21a..2f73eb4b3df 100644 --- a/litellm/ui/litellm-dashboard/src/components/key_info_utils.test.tsx +++ b/litellm/ui/litellm-dashboard/src/components/key_info_utils.test.tsx @@ -4,6 +4,7 @@ import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata, + stripFieldManagedKeysFromMetadata, } from "./key_info_utils"; describe("filterSensitiveMetadata", () => { @@ -57,3 +58,15 @@ describe("stripTagsFromMetadata", () => { expect(input).toEqual(originalCopy); }); }); + +describe("stripFieldManagedKeysFromMetadata", () => { + it.each(["ops@qa.test", ["ops@qa.test"], 80])( + "leaves a max_budget_alert_emails value that is not a threshold map in the Metadata box: %j", + (value) => { + expect(stripFieldManagedKeysFromMetadata({ cost_center: "ops", max_budget_alert_emails: value })).toEqual({ + cost_center: "ops", + max_budget_alert_emails: value, + }); + }, + ); +}); diff --git a/litellm/ui/litellm-dashboard/src/components/key_info_utils.tsx b/litellm/ui/litellm-dashboard/src/components/key_info_utils.tsx index e346878c24e..2f573a8e84e 100644 --- a/litellm/ui/litellm-dashboard/src/components/key_info_utils.tsx +++ b/litellm/ui/litellm-dashboard/src/components/key_info_utils.tsx @@ -2,6 +2,8 @@ * Utility functions for handling key information and metadata */ +import { budgetAlertEmailsMetadata, KEY_MAX_BUDGET_ALERT_EMAILS_KEY } from "./key_team_helpers/budgetAlertEmails"; + // Fields that should be filtered out from metadata display due to security concerns const SENSITIVE_METADATA_FIELDS = ["logging"] as const; @@ -61,3 +63,14 @@ export const stripTagsFromMetadata = (metadata: any) => { const { tags, ...rest } = metadata as Record; return rest; }; + +export const stripFieldManagedKeysFromMetadata = ( + metadata: Record | null | undefined, +): Record | null | undefined => { + const withoutTags: Record | null | undefined = stripTagsFromMetadata(metadata); + if (!withoutTags || typeof withoutTags !== "object") { + return withoutTags; + } + const { [KEY_MAX_BUDGET_ALERT_EMAILS_KEY]: _maxBudgetAlertEmails, ...rest } = withoutTags; + return { ...rest, ...budgetAlertEmailsMetadata({}, KEY_MAX_BUDGET_ALERT_EMAILS_KEY, withoutTags) }; +}; diff --git a/litellm/ui/litellm-dashboard/src/components/key_team_helpers/BudgetAlertEmailsEditor.tsx b/litellm/ui/litellm-dashboard/src/components/key_team_helpers/BudgetAlertEmailsEditor.tsx new file mode 100644 index 00000000000..87d40a55b4b --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/components/key_team_helpers/BudgetAlertEmailsEditor.tsx @@ -0,0 +1,89 @@ +import React from "react"; +import { type ArrayPath, type Control, type FieldPath, type FieldValues, useFieldArray } from "react-hook-form"; +import { CircleMinus, Plus } from "lucide-react"; + +import { Button } from "@/components/ui/button"; +import { Field, FieldLabel } from "@/components/ui/field"; +import { Input } from "@/components/ui/input"; +import { FormField } from "@/components/shared/form/FormField"; +import NumericalInput from "../shared/numerical_input"; +import type { BudgetAlertRow } from "./budgetAlertEmails"; + +export interface BudgetAlertEmailsEditorText { + readonly thresholdPlaceholder: string; + readonly recipientsPlaceholder: string; + readonly removeThreshold: string; + readonly addThreshold: string; +} + +type RowsPath = string extends keyof TFieldValues ? string : ArrayPath; + +interface BudgetAlertEmailsEditorProps { + readonly control: Control; + readonly name: RowsPath; + readonly label: React.ReactNode; + readonly text: BudgetAlertEmailsEditorText; +} + +const EMPTY_ROW: BudgetAlertRow = { threshold: null, emails: "" }; + +export const BudgetAlertEmailsEditor = ({ + control, + name, + label, + text, +}: BudgetAlertEmailsEditorProps) => { + const { fields, append, remove } = useFieldArray({ control, name: name as ArrayPath }); + const rowField = (index: number, column: keyof BudgetAlertRow) => + `${name}.${index}.${column}` as FieldPath; + + return ( + + {label} + {fields.map((row, index) => ( +
+ + {({ ref, value, onChange, ...field }) => ( + ) => + onChange(event.target.value === "" ? null : Number(event.target.value)) + } + placeholder={text.thresholdPlaceholder} + min={1} + max={100} + step={1} + /> + )} + + + {({ ref, value, ...field }) => ( + + )} + + +
+ ))} + + + ); +}; diff --git a/litellm/ui/litellm-dashboard/src/components/key_team_helpers/budgetAlertEmails.test.ts b/litellm/ui/litellm-dashboard/src/components/key_team_helpers/budgetAlertEmails.test.ts new file mode 100644 index 00000000000..23e372fff80 --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/components/key_team_helpers/budgetAlertEmails.test.ts @@ -0,0 +1,228 @@ +import { describe, expect, it } from "vitest"; +import { + budgetAlertEmailsFromRows, + budgetAlertEmailsMetadata, + budgetAlertEmailsRowsCanShow, + budgetAlertRowsFromMetadata, + isValidThreshold, + KEY_MAX_BUDGET_ALERT_EMAILS_KEY, + TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY, + teamMemberBudgetAlertSummary, +} from "./budgetAlertEmails"; + +describe("budgetAlertRowsFromMetadata", () => { + it("turns the stored threshold map into rows sorted by threshold", () => { + const metadata = { + team_member_max_budget_alert_emails: { "100": ["finance@example.com", "cto@example.com"], "50": [] }, + }; + expect(budgetAlertRowsFromMetadata(metadata, TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY)).toEqual([ + { threshold: 50, emails: "" }, + { threshold: 100, emails: "finance@example.com, cto@example.com" }, + ]); + }); + + it("drops non-numeric thresholds and non-email recipients instead of crashing", () => { + const metadata = { + team_member_max_budget_alert_emails: { fifty: [], "90": { a: 1 }, "100": ["a@b.c"] }, + }; + expect(budgetAlertRowsFromMetadata(metadata, TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY)).toEqual([ + { threshold: 100, emails: "a@b.c" }, + ]); + }); + + it("loads a null recipient list as an owner-only row and keeps the strings of a mixed list", () => { + const metadata = { max_budget_alert_emails: { "75": null, "80": ["b@example.com", null, 3] } }; + expect(budgetAlertRowsFromMetadata(metadata, KEY_MAX_BUDGET_ALERT_EMAILS_KEY)).toEqual([ + { threshold: 75, emails: "" }, + { threshold: 80, emails: "b@example.com" }, + ]); + }); + + it("lists API-stored string recipients in the team member summary", () => { + expect(teamMemberBudgetAlertSummary({ team_member_max_budget_alert_emails: { "80": "a@x.com, b@x.com" } })).toEqual( + ["80%: member, a@x.com, b@x.com"], + ); + }); + + it("leaves a zero-padded threshold to the Metadata box so it never duplicates a row", () => { + const metadata = { max_budget_alert_emails: { "080": ["padded@example.com"], "80": ["plain@example.com"] } }; + expect(budgetAlertRowsFromMetadata(metadata, KEY_MAX_BUDGET_ALERT_EMAILS_KEY)).toEqual([ + { threshold: 80, emails: "plain@example.com" }, + ]); + }); + + it("still loads a zero-padded team member threshold as a row, since the team form has no Metadata box fallback", () => { + const metadata = { team_member_max_budget_alert_emails: { "080": ["finance@example.com"] } }; + expect(budgetAlertRowsFromMetadata(metadata, TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY)).toEqual([ + { threshold: 80, emails: "finance@example.com" }, + ]); + }); + + it("keeps API-stored comma-separated string recipients, which the backend also accepts", () => { + const metadata = { + max_budget_alert_emails: { "80": "finance@example.com, cto@example.com", "90": ["ops@example.com"] }, + }; + const rows = budgetAlertRowsFromMetadata(metadata, KEY_MAX_BUDGET_ALERT_EMAILS_KEY); + expect(rows).toEqual([ + { threshold: 80, emails: "finance@example.com, cto@example.com" }, + { threshold: 90, emails: "ops@example.com" }, + ]); + expect(budgetAlertEmailsFromRows(rows)).toEqual({ + "80": ["finance@example.com", "cto@example.com"], + "90": ["ops@example.com"], + }); + }); + + it("drops API-stored thresholds outside 1 to 100 so they never block the form", () => { + const metadata = { + team_member_max_budget_alert_emails: { "0": ["a@b.c"], "50": [], "101": ["a@b.c"] }, + }; + expect(budgetAlertRowsFromMetadata(metadata, TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY)).toEqual([ + { threshold: 50, emails: "" }, + ]); + }); + + it.each([undefined, null, "50", { team_member_max_budget_alert_emails: "50" }, { soft_budget_alerting_emails: [] }])( + "returns no rows for unrelated or malformed metadata %j", + (metadata) => { + expect(budgetAlertRowsFromMetadata(metadata, TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY)).toEqual([]); + }, + ); + + it("reads only the metadata key it is asked for", () => { + const metadata = { + team_member_max_budget_alert_emails: { "50": ["team@example.com"] }, + max_budget_alert_emails: { "80": ["shared-mailbox@example.com"] }, + }; + expect(budgetAlertRowsFromMetadata(metadata, KEY_MAX_BUDGET_ALERT_EMAILS_KEY)).toEqual([ + { threshold: 80, emails: "shared-mailbox@example.com" }, + ]); + }); +}); + +describe("budgetAlertEmailsMetadata", () => { + it("writes the threshold map under the given metadata key", () => { + expect( + budgetAlertEmailsMetadata( + budgetAlertEmailsFromRows([{ threshold: 80, emails: "dl@example.com" }]), + KEY_MAX_BUDGET_ALERT_EMAILS_KEY, + {}, + ), + ).toEqual({ max_budget_alert_emails: { "80": ["dl@example.com"] } }); + }); + + it.each([{ rows: [] }, { rows: [{ threshold: null, emails: "dl@example.com" }] }])( + "leaves the key out when no row has a valid threshold $rows", + ({ rows }) => { + expect(budgetAlertEmailsMetadata(budgetAlertEmailsFromRows(rows), KEY_MAX_BUDGET_ALERT_EMAILS_KEY, {})).toEqual( + {}, + ); + }, + ); + + it("keeps the entries the rows cannot show and lets the rows own every threshold they can", () => { + const boxMetadata = { + max_budget_alert_emails: { + "50": ["old@example.com"], + "80": ["removed@example.com"], + "90": null, + "150": ["over@example.com"], + "0": ["zero@example.com"], + "12.5": ["fraction@example.com"], + soon: "label@example.com", + }, + }; + + expect( + budgetAlertEmailsMetadata( + budgetAlertEmailsFromRows([{ threshold: 50, emails: "new@example.com" }]), + KEY_MAX_BUDGET_ALERT_EMAILS_KEY, + boxMetadata, + ), + ).toEqual({ + max_budget_alert_emails: { + "50": ["new@example.com"], + "150": ["over@example.com"], + "0": ["zero@example.com"], + "12.5": ["fraction@example.com"], + soon: "label@example.com", + }, + }); + }); +}); + +describe("budgetAlertEmailsRowsCanShow", () => { + it("rebuilds the stored map together with the entries the rows cannot show, values untouched", () => { + const stored = { + max_budget_alert_emails: { + "50": "a@example.com, b@example.com", + "75": null, + "80": ["c@example.com", null], + "150": ["over@example.com"], + }, + }; + + expect( + budgetAlertEmailsMetadata( + budgetAlertEmailsRowsCanShow(stored, KEY_MAX_BUDGET_ALERT_EMAILS_KEY), + KEY_MAX_BUDGET_ALERT_EMAILS_KEY, + stored, + ), + ).toEqual(stored); + expect(budgetAlertEmailsRowsCanShow(stored, KEY_MAX_BUDGET_ALERT_EMAILS_KEY)).not.toHaveProperty("150"); + }); +}); + +describe("budgetAlertEmailsFromRows", () => { + it("builds the threshold map, splitting, trimming and deduplicating recipients", () => { + expect( + budgetAlertEmailsFromRows([ + { threshold: 50, emails: "" }, + { threshold: 100, emails: " finance@example.com,cto@example.com , finance@example.com, " }, + ]), + ).toEqual({ "50": [], "100": ["finance@example.com", "cto@example.com"] }); + }); + + it("skips rows without a valid threshold", () => { + expect( + budgetAlertEmailsFromRows([ + { threshold: null, emails: "finance@example.com" }, + { threshold: 0, emails: "" }, + { threshold: 101, emails: "" }, + { threshold: 12.5, emails: "" }, + { threshold: 80, emails: "" }, + ]), + ).toEqual({ "80": [] }); + }); + + it("round-trips the stored config", () => { + const stored = { team_member_max_budget_alert_emails: { "50": [], "100": ["finance@example.com"] } }; + expect( + budgetAlertEmailsFromRows(budgetAlertRowsFromMetadata(stored, TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY)), + ).toEqual(stored.team_member_max_budget_alert_emails); + }); +}); + +describe("isValidThreshold", () => { + it.each([ + [1, true], + [50, true], + [100, true], + [0, false], + [101, false], + [33.3, false], + [null, false], + ])("treats %s as valid=%s", (threshold, valid) => { + expect(isValidThreshold(threshold)).toBe(valid); + }); +}); + +describe("teamMemberBudgetAlertSummary", () => { + it("states that the member is always notified and lists extra recipients", () => { + expect( + teamMemberBudgetAlertSummary({ + team_member_max_budget_alert_emails: { "100": ["finance@example.com"], "50": [] }, + }), + ).toEqual(["50%: member", "100%: member, finance@example.com"]); + }); +}); diff --git a/litellm/ui/litellm-dashboard/src/components/key_team_helpers/budgetAlertEmails.ts b/litellm/ui/litellm-dashboard/src/components/key_team_helpers/budgetAlertEmails.ts new file mode 100644 index 00000000000..b8f7ba4020e --- /dev/null +++ b/litellm/ui/litellm-dashboard/src/components/key_team_helpers/budgetAlertEmails.ts @@ -0,0 +1,143 @@ +import { z } from "zod"; + +export const TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY = "team_member_max_budget_alert_emails" as const; +export const KEY_MAX_BUDGET_ALERT_EMAILS_KEY = "max_budget_alert_emails" as const; + +export type BudgetAlertEmailsMetadataKey = + | typeof TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY + | typeof KEY_MAX_BUDGET_ALERT_EMAILS_KEY; + +export interface BudgetAlertRow { + readonly threshold: number | null; + readonly emails: string; +} + +export type BudgetAlertEmails = Readonly>; + +export type BudgetAlertRowIssue = "thresholdRange" | "duplicateThreshold"; + +const splitEmails = (emails: string): readonly string[] => + Array.from( + new Set( + emails + .split(",") + .map((email) => email.trim()) + .filter((email) => email.length > 0), + ), + ); + +const storedRecipients = (value: unknown): readonly string[] | null => { + if (value === null) return []; + if (typeof value === "string") return splitEmails(value); + if (Array.isArray(value)) return value.filter((email): email is string => typeof email === "string"); + return null; +}; + +const THRESHOLD_MIN = 1; +const THRESHOLD_MAX = 100; + +export const isValidThreshold = (threshold: number | null): threshold is number => { + const isWholeNumber = threshold !== null && Number.isInteger(threshold); + return isWholeNumber && threshold >= THRESHOLD_MIN && threshold <= THRESHOLD_MAX; +}; + +const isBudgetAlertMap = (value: unknown): value is Readonly> => + typeof value === "object" && value !== null && !Array.isArray(value); + +const storedBudgetAlertValue = (metadata: unknown, metadataKey: BudgetAlertEmailsMetadataKey): unknown => + typeof metadata === "object" && metadata !== null ? (metadata as Record)[metadataKey] : undefined; + +const storedBudgetAlertConfig = ( + metadata: unknown, + metadataKey: BudgetAlertEmailsMetadataKey, +): Readonly> => { + const config = storedBudgetAlertValue(metadata, metadataKey); + return isBudgetAlertMap(config) ? config : {}; +}; + +const ROW_THRESHOLD_KEY_PATTERN: Readonly> = { + [KEY_MAX_BUDGET_ALERT_EMAILS_KEY]: /^[1-9]\d*$/, + [TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY]: /^\d+$/, +}; + +const budgetAlertRowFromEntry = ( + [key, emails]: readonly [string, unknown], + metadataKey: BudgetAlertEmailsMetadataKey, +): BudgetAlertRow | null => { + const threshold = Number(key); + const recipients = storedRecipients(emails); + return ROW_THRESHOLD_KEY_PATTERN[metadataKey].test(key) && isValidThreshold(threshold) && recipients !== null + ? { threshold, emails: recipients.join(", ") } + : null; +}; + +export const budgetAlertRowsFromMetadata = ( + metadata: unknown, + metadataKey: BudgetAlertEmailsMetadataKey, +): readonly BudgetAlertRow[] => + Object.entries(storedBudgetAlertConfig(metadata, metadataKey)) + .flatMap((entry) => { + const row = budgetAlertRowFromEntry(entry, metadataKey); + return row === null ? [] : [row]; + }) + .sort((a, b) => (a.threshold ?? 0) - (b.threshold ?? 0)); + +const storedBudgetAlertEntries = ( + metadata: unknown, + metadataKey: BudgetAlertEmailsMetadataKey, + rowsCanShow: boolean, +): Readonly> => + Object.fromEntries( + Object.entries(storedBudgetAlertConfig(metadata, metadataKey)).filter( + (entry) => (budgetAlertRowFromEntry(entry, metadataKey) !== null) === rowsCanShow, + ), + ); + +export const budgetAlertEmailsRowsCanShow = ( + metadata: unknown, + metadataKey: BudgetAlertEmailsMetadataKey, +): Readonly> => storedBudgetAlertEntries(metadata, metadataKey, true); + +export const budgetAlertEmailsRowsCannotShow = ( + metadata: unknown, + metadataKey: BudgetAlertEmailsMetadataKey, +): Readonly> => storedBudgetAlertEntries(metadata, metadataKey, false); + +export const budgetAlertEmailsFromRows = (rows: readonly BudgetAlertRow[]): BudgetAlertEmails => + Object.fromEntries( + rows + .filter((row) => isValidThreshold(row.threshold)) + .map((row) => [String(row.threshold), splitEmails(row.emails)]), + ); + +export const budgetAlertEmailsMetadata = ( + rowEntries: Readonly>, + metadataKey: BudgetAlertEmailsMetadataKey, + metadataWithHiddenEntries: unknown, +): Readonly>> => { + const boxValue = storedBudgetAlertValue(metadataWithHiddenEntries, metadataKey); + if (Object.keys(rowEntries).length === 0 && boxValue !== undefined && !isBudgetAlertMap(boxValue)) { + return { [metadataKey]: boxValue }; + } + const emails = { + ...budgetAlertEmailsRowsCannotShow(metadataWithHiddenEntries, metadataKey), + ...rowEntries, + }; + return Object.keys(emails).length > 0 ? { [metadataKey]: emails } : {}; +}; + +export const budgetAlertRowsSchema = (message: (issue: BudgetAlertRowIssue) => string) => + z.array(z.object({ threshold: z.number().nullable(), emails: z.string() })).superRefine((rows, ctx) => { + rows.forEach((row, index) => { + if (!isValidThreshold(row.threshold)) { + ctx.addIssue({ code: "custom", message: message("thresholdRange"), path: [index, "threshold"] }); + } else if (rows.filter((other) => other.threshold === row.threshold).length > 1) { + ctx.addIssue({ code: "custom", message: message("duplicateThreshold"), path: [index, "threshold"] }); + } + }); + }); + +export const teamMemberBudgetAlertSummary = (metadata: unknown): readonly string[] => + budgetAlertRowsFromMetadata(metadata, TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY).map((row) => + row.emails.length > 0 ? `${row.threshold}%: member, ${row.emails}` : `${row.threshold}%: member`, + ); diff --git a/litellm/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts b/litellm/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts index c544a89930d..6104da1cf65 100644 --- a/litellm/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts +++ b/litellm/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts @@ -29,6 +29,78 @@ it.each([false, true])("preserves the prompt storage restriction %s in key metad expect(JSON.parse(String(result.payload.metadata))).toEqual({ note: "keep", store_prompts_in_spend_logs: store }); }); +it("merges the budget alert rows into the metadata JSON box content", () => { + const result = build( + { + metadata: JSON.stringify({ cost_center: "ops" }), + max_budget_alert_emails: [ + { threshold: 80, emails: "shared-mailbox@example.com" }, + { threshold: null, emails: "ignored@example.com" }, + ], + }, + { disabledCallbacks: ["langfuse"] }, + ); + expect(result.kind).toBe("ok"); + if (result.kind !== "ok") throw new Error(result.kind); + expect(result.payload).not.toHaveProperty("max_budget_alert_emails"); + expect(JSON.parse(String(result.payload.metadata))).toEqual({ + cost_center: "ops", + litellm_disabled_callbacks: ["langfuse"], + max_budget_alert_emails: { "80": ["shared-mailbox@example.com"] }, + }); +}); + +it("keeps typed recipient entries the rows cannot show and lets the rows own every threshold they can", () => { + const result = build({ + metadata: JSON.stringify({ + cost_center: "ops", + max_budget_alert_emails: { "500": "typed@example.com", "80": ["stale@example.com"] }, + }), + max_budget_alert_emails: [{ threshold: 50, emails: "row@example.com" }], + }); + expect(result.kind).toBe("ok"); + if (result.kind !== "ok") throw new Error(result.kind); + expect(JSON.parse(String(result.payload.metadata))).toEqual({ + cost_center: "ops", + max_budget_alert_emails: { "500": "typed@example.com", "50": ["row@example.com"] }, + }); +}); + +it("stores a typed max_budget_alert_emails string as typed when no row is set", () => { + const result = build({ metadata: JSON.stringify({ cost_center: "ops", max_budget_alert_emails: "ops@qa.test" }) }); + expect(result.kind).toBe("ok"); + if (result.kind !== "ok") throw new Error(result.kind); + expect(JSON.parse(String(result.payload.metadata))).toEqual({ + cost_center: "ops", + max_budget_alert_emails: "ops@qa.test", + }); +}); + +it("replaces a typed max_budget_alert_emails string with the threshold rows", () => { + const result = build({ + metadata: JSON.stringify({ cost_center: "ops", max_budget_alert_emails: "ops@qa.test" }), + max_budget_alert_emails: [{ threshold: 80, emails: "row@example.com" }], + }); + expect(result.kind).toBe("ok"); + if (result.kind !== "ok") throw new Error(result.kind); + expect(JSON.parse(String(result.payload.metadata))).toEqual({ + cost_center: "ops", + max_budget_alert_emails: { "80": ["row@example.com"] }, + }); +}); + +it("keeps the budget alert rows when the metadata JSON box holds null", () => { + const result = build({ + metadata: "null", + max_budget_alert_emails: [{ threshold: 80, emails: "shared-mailbox@example.com" }], + }); + expect(result.kind).toBe("ok"); + if (result.kind !== "ok") throw new Error(result.kind); + expect(JSON.parse(String(result.payload.metadata))).toEqual({ + max_budget_alert_emails: { "80": ["shared-mailbox@example.com"] }, + }); +}); + const payloadOf = (result: KeyPayloadResult): Record => { expect(result.kind).toBe("ok"); if (result.kind !== "ok") throw new Error("unreachable"); diff --git a/litellm/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts b/litellm/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts index 311e9357825..4cd2e35a2fa 100644 --- a/litellm/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts +++ b/litellm/ui/litellm-dashboard/src/components/organisms/createKeyPayload.ts @@ -1,6 +1,12 @@ import { mapDisplayToInternalNames } from "../callback_info_helpers"; import { NEVER_RESETS_BUDGET_DURATION } from "../common_components/budget_duration_dropdown"; import type { RouterSettingsAccordionValue } from "../common_components/RouterSettingsAccordion"; +import { + type BudgetAlertRow, + budgetAlertEmailsFromRows, + budgetAlertEmailsMetadata, + KEY_MAX_BUDGET_ALERT_EMAILS_KEY, +} from "../key_team_helpers/budgetAlertEmails"; import type { BudgetWindowEntry } from "../key_team_helpers/BudgetWindowsEditor"; import type { ModelMaxBudget } from "../key_team_helpers/ModelMaxBudgetEditor"; import { tagRowsToLimits, type TagRateLimitEntry } from "../key_team_helpers/TagRateLimitEditor"; @@ -103,7 +109,18 @@ const buildMetadataJson = (values: Record, input: KeyCreateInpu input.disabledCallbacks.length > 0 ? { ...(logged as object), litellm_disabled_callbacks: mapDisplayToInternalNames(input.disabledCallbacks) } : logged; - return JSON.stringify(disabled); + const metadataObject = disabled ?? {}; + const alertEmails = budgetAlertEmailsMetadata( + budgetAlertEmailsFromRows((values[KEY_MAX_BUDGET_ALERT_EMAILS_KEY] as readonly BudgetAlertRow[] | undefined) ?? []), + KEY_MAX_BUDGET_ALERT_EMAILS_KEY, + metadataObject, + ); + if (typeof metadataObject !== "object" || Array.isArray(metadataObject)) return JSON.stringify(metadataObject); + const { [KEY_MAX_BUDGET_ALERT_EMAILS_KEY]: _typedAlertEmails, ...withoutTypedAlertEmails } = metadataObject as Record< + string, + unknown + >; + return JSON.stringify({ ...withoutTypedAlertEmails, ...alertEmails }); }; interface PermissionSources { @@ -153,6 +170,7 @@ const consumedSourceKeys = ( new Set([ "mcp_tool_permissions", "allowed_skills", + KEY_MAX_BUDGET_ALERT_EMAILS_KEY, ...(values.disable_global_guardrails ? [] : ["disable_global_guardrails"]), ...(vectorStores ? ["allowed_vector_store_ids"] : []), ...(mcp ? ["allowed_mcp_servers_and_groups"] : []), diff --git a/litellm/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/litellm/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index bbc849167cc..89cdb51d234 100644 --- a/litellm/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/litellm/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -1,6 +1,6 @@ import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; -import { act, renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils"; +import { act, fireEvent, renderWithProviders, screen, testQueryClient, waitFor } from "../../../tests/test-utils"; import type { Team } from "../key_team_helpers/key_list"; import { keyCreateCall, keyCreateServiceAccountCall, modelAvailableCall, userFilterUICall } from "../networking"; import { toast } from "@/lib/toast"; @@ -593,6 +593,68 @@ describe("CreateKey", () => { expect((await createdPayload()).metadata).toBe('{"team":"research"}'); }); + const addBudgetAlertRow = async (threshold: string, emails: string) => { + await userEvent.click(await screen.findByRole("button", { name: "Add Budget Alert Threshold" })); + const thresholds = screen.getAllByPlaceholderText("% of budget"); + fireEvent.change(thresholds[thresholds.length - 1], { target: { value: threshold } }); + const recipients = screen.getAllByPlaceholderText("Additional recipients, e.g. finance@example.com"); + fireEvent.change(recipients[recipients.length - 1], { target: { value: emails } }); + }; + + it("merges the budget alert rows into the metadata JSON box content", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + fireEvent.change(await screen.findByLabelText("Metadata"), { target: { value: '{"team":"research"}' } }); + await addBudgetAlertRow("80", "shared-mailbox@example.com, finance-dl@example.com"); + await submit(); + + const payload = await createdPayload(); + expect(JSON.parse(String(payload.metadata))).toEqual({ + team: "research", + max_budget_alert_emails: { "80": ["shared-mailbox@example.com", "finance-dl@example.com"] }, + }); + expect(payload).not.toHaveProperty("max_budget_alert_emails"); + }); + + it("moves a recipient map typed into the metadata JSON box into the rows on blur", async () => { + vi.mocked(toast.info).mockClear(); + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await addBudgetAlertRow("80", "row@example.com"); + const metadataBox = await screen.findByLabelText("Metadata"); + fireEvent.change(metadataBox, { + target: { value: JSON.stringify({ team: "research", max_budget_alert_emails: { "80": "typed@example.com" } }) }, + }); + fireEvent.blur(metadataBox); + + expect(toast.info).toHaveBeenCalledWith( + "Moved alert thresholds 80% from metadata to the Budget Alert Thresholds rows", + ); + expect(screen.getByPlaceholderText("Additional recipients, e.g. finance@example.com")).toHaveValue( + "typed@example.com", + ); + await submit(); + + expect(JSON.parse(String((await createdPayload()).metadata))).toEqual({ + team: "research", + max_budget_alert_emails: { "80": ["typed@example.com"] }, + }); + }); + + it("blocks the create on a duplicate budget alert threshold", async () => { + await openModal(); + await nameTheKey(); + await openSection(/Optional Settings/i); + await addBudgetAlertRow("80", "a@example.com"); + await addBudgetAlertRow("80", "b@example.com"); + await submit(); + + expect(await screen.findAllByText("Duplicate threshold")).toHaveLength(2); + expect(vi.mocked(keyCreateCall)).not.toHaveBeenCalled(); + }); + it("carries the typed key alias and the chosen team onto the wire", async () => { state.teams = [{ team_id: "team-1", team_alias: "Team One", models: [] }]; await openModal({ teams: state.teams as unknown as Team[] }); diff --git a/litellm/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/litellm/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 2e82e1ca4f5..cc2b75d10a0 100644 --- a/litellm/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/litellm/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -42,6 +42,7 @@ import { MountedFormField, MountedFormProvider, projectMountedValues, + useMountedName, useMountRegistry, type MountedFormValues, } from "../common_components/MountedFormField"; @@ -56,6 +57,13 @@ import TeamDropdown from "../common_components/team_dropdown"; import OrganizationDropdown from "../common_components/OrganizationDropdown"; import ProjectDropdown from "../common_components/ProjectDropdown"; import { CreateUserButton } from "../CreateUserButton"; +import { moveMetadataBudgetAlertsToRows } from "../templates/KeyEditViewControls"; +import { BudgetAlertEmailsEditor, type BudgetAlertEmailsEditorText } from "../key_team_helpers/BudgetAlertEmailsEditor"; +import { + type BudgetAlertRow, + budgetAlertRowsSchema, + KEY_MAX_BUDGET_ALERT_EMAILS_KEY, +} from "../key_team_helpers/budgetAlertEmails"; import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor"; import { END_USER_BUDGET_HINT, EndUserBudgetSelect } from "../key_team_helpers/EndUserBudgetSelect"; import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor"; @@ -132,6 +140,15 @@ type McpSelectorValue = { servers: string[]; accessGroups: string[]; toolsets?: type AgentSelectorValue = { agents: string[]; accessGroups: string[] }; +const MountedBudgetAlertEmailsEditor = (props: { + readonly control: Control; + readonly label: React.ReactNode; + readonly text: BudgetAlertEmailsEditorText; +}) => { + useMountedName(KEY_MAX_BUDGET_ALERT_EMAILS_KEY); + return ; +}; + const isBlank = (value: unknown): boolean => value === undefined || value === null || value === ""; const requiredRule = (required: boolean, message: string) => ({ @@ -532,7 +549,21 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp }; const handleSubmit = (event: React.FormEvent) => - void form.handleSubmit(() => handleCreate(projectMountedValues(registry, form.getValues)))(event); + void form.handleSubmit(() => { + const values = projectMountedValues(registry, form.getValues); + const alertRows = budgetAlertRowsSchema((issue) => + issue === "thresholdRange" + ? t("keyDetail.maxBudgetAlertEmails.thresholdRange") + : t("keyDetail.maxBudgetAlertEmails.duplicateThreshold"), + ).safeParse(values[KEY_MAX_BUDGET_ALERT_EMAILS_KEY] ?? []); + if (!alertRows.success) { + alertRows.error.issues.forEach((issue) => + form.setError([KEY_MAX_BUDGET_ALERT_EMAILS_KEY, ...issue.path].join("."), { message: issue.message }), + ); + return; + } + return handleCreate(values); + })(event); // Fetch available models when team or auth changes. // Note: Model prefill from URL params is handled by the useEffect below, which @@ -1093,6 +1124,25 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp /> )} +
+ + {t("keyDetail.maxBudgetAlertEmails.budgetAlertThresholds")}{" "} + + + + + } + text={{ + thresholdPlaceholder: t("keyDetail.maxBudgetAlertEmails.thresholdPlaceholder"), + recipientsPlaceholder: t("keyDetail.maxBudgetAlertEmails.recipientsPlaceholder"), + removeThreshold: t("keyDetail.maxBudgetAlertEmails.removeThreshold"), + addThreshold: t("keyDetail.maxBudgetAlertEmails.addThreshold"), + }} + /> +
= ({ team, teams, data, addKey, autoOp value={(control.value as string | undefined) ?? ""} rows={4} placeholder={t("createKey.metadata.placeholder")} + onBlur={() => { + control.onBlur(); + moveMetadataBudgetAlertsToRows( + form.getValues("metadata") as string | undefined, + form.getValues(KEY_MAX_BUDGET_ALERT_EMAILS_KEY) as BudgetAlertRow[] | undefined, + (metadata, rows) => { + form.setValue("metadata", metadata, { shouldDirty: true }); + form.setValue(KEY_MAX_BUDGET_ALERT_EMAILS_KEY, rows, { shouldDirty: true }); + }, + t, + ); + }} /> )} diff --git a/litellm/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/litellm/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 24e6dd75abf..ad4c88dd75f 100644 --- a/litellm/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/litellm/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -130,12 +130,13 @@ import { import TeamMembersComponent from "./TeamMemberTab"; import { useModelAccessGroupNames } from "@/app/(dashboard)/hooks/models/useModels"; import { - isValidThreshold, + budgetAlertEmailsFromRows, + budgetAlertRowsFromMetadata, + budgetAlertRowsSchema, TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY, - teamMemberBudgetAlertEmailsFromRows, - teamMemberBudgetAlertRowsFromMetadata, teamMemberBudgetAlertSummary, -} from "./teamMemberBudgetAlertEmails"; +} from "../key_team_helpers/budgetAlertEmails"; +import { BudgetAlertEmailsEditor } from "../key_team_helpers/BudgetAlertEmailsEditor"; import { TeamVirtualKeysTable } from "./TeamVirtualKeysTable"; import ResetMemberBudgetsDialog from "./ResetMemberBudgetsDialog"; import { customBudgetMemberUserIds, shouldPromptMemberBudgetReset } from "./memberBudgetReset"; @@ -387,26 +388,11 @@ const buildTeamUpdateFieldsSchema = (t: TFunction) => team_member_key_duration: z.string().optional(), team_member_tpm_limit: numericInputSchema, team_member_rpm_limit: numericInputSchema, - team_member_max_budget_alert_emails: z - .array(z.object({ threshold: z.number().nullable(), emails: z.string() })) - .superRefine((rows, ctx) => { - rows.forEach((row, index) => { - if (!isValidThreshold(row.threshold)) { - ctx.addIssue({ - code: "custom", - message: t("teams.teamInfo.validation.thresholdRange"), - path: [index, "threshold"], - }); - } else if (rows.filter((other) => other.threshold === row.threshold).length > 1) { - ctx.addIssue({ - code: "custom", - message: t("teams.teamInfo.validation.duplicateThreshold"), - path: [index, "threshold"], - }); - } - }); - }) - .optional(), + team_member_max_budget_alert_emails: budgetAlertRowsSchema((issue) => + issue === "thresholdRange" + ? t("teams.teamInfo.validation.thresholdRange") + : t("teams.teamInfo.validation.duplicateThreshold"), + ).optional(), budget_duration: z.string().nullish(), tpm_limit: numericInputSchema, rpm_limit: numericInputSchema, @@ -551,7 +537,9 @@ const toTeamFormValues = (info: TeamInfoRecord, effectiveGuardrails: string[]): team_member_key_duration: info.metadata?.team_member_key_duration, team_member_tpm_limit: info.team_member_budget_table?.tpm_limit, team_member_rpm_limit: info.team_member_budget_table?.rpm_limit, - team_member_max_budget_alert_emails: [...teamMemberBudgetAlertRowsFromMetadata(info.metadata)], + team_member_max_budget_alert_emails: [ + ...budgetAlertRowsFromMetadata(info.metadata, TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY), + ], budget_duration: info.budget_duration, tpm_limit: info.tpm_limit, rpm_limit: info.rpm_limit, @@ -639,11 +627,6 @@ const TeamInfoView: React.FC = ({ append: appendModelLimit, remove: removeModelLimit, } = useFieldArray({ control: form.control, name: "modelLimits" }); - const { - fields: memberBudgetAlertRows, - append: appendMemberBudgetAlertRow, - remove: removeMemberBudgetAlertRow, - } = useFieldArray({ control: form.control, name: "team_member_max_budget_alert_emails" }); const [teamMemberSettingsOpen, setTeamMemberSettingsOpen] = useState(false); const [searchToolSettingsOpen, setSearchToolSettingsOpen] = useState(false); const [isEditMemberModalVisible, setIsEditMemberModalVisible] = useState(false); @@ -1081,7 +1064,7 @@ const TeamInfoView: React.FC = ({ const memberBudgetAlertEmails = values.team_member_max_budget_alert_emails !== undefined - ? teamMemberBudgetAlertEmailsFromRows(values.team_member_max_budget_alert_emails) + ? budgetAlertEmailsFromRows(values.team_member_max_budget_alert_emails) : info.metadata?.[TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY]; const memberBudgetAlertEmailsMetadata = memberBudgetAlertEmails !== undefined && Object.keys(memberBudgetAlertEmails).length > 0 @@ -1877,71 +1860,22 @@ const TeamInfoView: React.FC = ({ /> )} - - - {labelWithHint( - t("teams.teamInfo.editForm.teamMemberSettings.budgetAlertThresholds"), - t("teams.teamInfo.editForm.teamMemberSettings.budgetAlertThresholdsHint"), - )} - - {memberBudgetAlertRows.map((row, index) => ( -
- - {({ ref, value, onChange, ...field }) => ( - ) => - onChange(event.target.value === "" ? null : Number(event.target.value)) - } - placeholder={t("teams.teamInfo.editForm.teamMemberSettings.thresholdPlaceholder")} - min={1} - max={100} - step={1} - /> - )} - - - {({ ref, value, ...field }) => ( - - )} - - -
- ))} - - + diff --git a/litellm/ui/litellm-dashboard/src/components/team/teamMemberBudgetAlertEmails.test.ts b/litellm/ui/litellm-dashboard/src/components/team/teamMemberBudgetAlertEmails.test.ts deleted file mode 100644 index 3faa051047b..00000000000 --- a/litellm/ui/litellm-dashboard/src/components/team/teamMemberBudgetAlertEmails.test.ts +++ /dev/null @@ -1,94 +0,0 @@ -import { describe, expect, it } from "vitest"; -import { - isValidThreshold, - teamMemberBudgetAlertEmailsFromRows, - teamMemberBudgetAlertRowsFromMetadata, - teamMemberBudgetAlertSummary, -} from "./teamMemberBudgetAlertEmails"; - -describe("teamMemberBudgetAlertRowsFromMetadata", () => { - it("turns the stored threshold map into rows sorted by threshold", () => { - const metadata = { - team_member_max_budget_alert_emails: { "100": ["finance@example.com", "cto@example.com"], "50": [] }, - }; - expect(teamMemberBudgetAlertRowsFromMetadata(metadata)).toEqual([ - { threshold: 50, emails: "" }, - { threshold: 100, emails: "finance@example.com, cto@example.com" }, - ]); - }); - - it("drops non-numeric thresholds and non-list recipients instead of crashing", () => { - const metadata = { - team_member_max_budget_alert_emails: { fifty: [], "75": "finance@example.com", "90": [1], "100": ["a@b.c"] }, - }; - expect(teamMemberBudgetAlertRowsFromMetadata(metadata)).toEqual([{ threshold: 100, emails: "a@b.c" }]); - }); - - it("drops API-stored thresholds outside 1 to 100 so they never block the form", () => { - const metadata = { - team_member_max_budget_alert_emails: { "0": ["a@b.c"], "50": [], "101": ["a@b.c"] }, - }; - expect(teamMemberBudgetAlertRowsFromMetadata(metadata)).toEqual([{ threshold: 50, emails: "" }]); - }); - - it.each([undefined, null, "50", { team_member_max_budget_alert_emails: "50" }, { soft_budget_alerting_emails: [] }])( - "returns no rows for unrelated or malformed metadata %j", - (metadata) => { - expect(teamMemberBudgetAlertRowsFromMetadata(metadata)).toEqual([]); - }, - ); -}); - -describe("teamMemberBudgetAlertEmailsFromRows", () => { - it("builds the threshold map, splitting, trimming and deduplicating recipients", () => { - expect( - teamMemberBudgetAlertEmailsFromRows([ - { threshold: 50, emails: "" }, - { threshold: 100, emails: " finance@example.com,cto@example.com , finance@example.com, " }, - ]), - ).toEqual({ "50": [], "100": ["finance@example.com", "cto@example.com"] }); - }); - - it("skips rows without a valid threshold", () => { - expect( - teamMemberBudgetAlertEmailsFromRows([ - { threshold: null, emails: "finance@example.com" }, - { threshold: 0, emails: "" }, - { threshold: 101, emails: "" }, - { threshold: 12.5, emails: "" }, - { threshold: 80, emails: "" }, - ]), - ).toEqual({ "80": [] }); - }); - - it("round-trips the stored config", () => { - const stored = { team_member_max_budget_alert_emails: { "50": [], "100": ["finance@example.com"] } }; - expect(teamMemberBudgetAlertEmailsFromRows(teamMemberBudgetAlertRowsFromMetadata(stored))).toEqual( - stored.team_member_max_budget_alert_emails, - ); - }); -}); - -describe("isValidThreshold", () => { - it.each([ - [1, true], - [50, true], - [100, true], - [0, false], - [101, false], - [33.3, false], - [null, false], - ])("treats %s as valid=%s", (threshold, valid) => { - expect(isValidThreshold(threshold)).toBe(valid); - }); -}); - -describe("teamMemberBudgetAlertSummary", () => { - it("states that the member is always notified and lists extra recipients", () => { - expect( - teamMemberBudgetAlertSummary({ - team_member_max_budget_alert_emails: { "100": ["finance@example.com"], "50": [] }, - }), - ).toEqual(["50%: member", "100%: member, finance@example.com"]); - }); -}); diff --git a/litellm/ui/litellm-dashboard/src/components/team/teamMemberBudgetAlertEmails.ts b/litellm/ui/litellm-dashboard/src/components/team/teamMemberBudgetAlertEmails.ts deleted file mode 100644 index 36d5ddbac02..00000000000 --- a/litellm/ui/litellm-dashboard/src/components/team/teamMemberBudgetAlertEmails.ts +++ /dev/null @@ -1,57 +0,0 @@ -export const TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY = "team_member_max_budget_alert_emails" as const; - -export interface TeamMemberBudgetAlertRow { - readonly threshold: number | null; - readonly emails: string; -} - -export type TeamMemberBudgetAlertEmails = Readonly>; - -const isEmailList = (value: unknown): value is readonly string[] => - Array.isArray(value) && value.every((email) => typeof email === "string"); - -const splitEmails = (emails: string): readonly string[] => - Array.from( - new Set( - emails - .split(",") - .map((email) => email.trim()) - .filter((email) => email.length > 0), - ), - ); - -const THRESHOLD_MIN = 1; -const THRESHOLD_MAX = 100; - -export const isValidThreshold = (threshold: number | null): threshold is number => { - const isWholeNumber = threshold !== null && Number.isInteger(threshold); - return isWholeNumber && threshold >= THRESHOLD_MIN && threshold <= THRESHOLD_MAX; -}; - -export const teamMemberBudgetAlertRowsFromMetadata = (metadata: unknown): readonly TeamMemberBudgetAlertRow[] => { - if (typeof metadata !== "object" || metadata === null) return []; - const config: unknown = (metadata as Record)[TEAM_MEMBER_MAX_BUDGET_ALERT_EMAILS_KEY]; - if (typeof config !== "object" || config === null || Array.isArray(config)) return []; - return Object.entries(config as Record) - .flatMap(([key, emails]) => { - const threshold = Number(key); - return /^\d+$/.test(key) && isValidThreshold(threshold) && isEmailList(emails) - ? [{ threshold, emails: emails.join(", ") }] - : []; - }) - .sort((a, b) => (a.threshold ?? 0) - (b.threshold ?? 0)); -}; - -export const teamMemberBudgetAlertEmailsFromRows = ( - rows: readonly TeamMemberBudgetAlertRow[], -): TeamMemberBudgetAlertEmails => - Object.fromEntries( - rows - .filter((row) => isValidThreshold(row.threshold)) - .map((row) => [String(row.threshold), splitEmails(row.emails)]), - ); - -export const teamMemberBudgetAlertSummary = (metadata: unknown): readonly string[] => - teamMemberBudgetAlertRowsFromMetadata(metadata).map((row) => - row.emails.length > 0 ? `${row.threshold}%: member, ${row.emails}` : `${row.threshold}%: member`, - ); diff --git a/litellm/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx b/litellm/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx index 3b96134f746..d44877a39e8 100644 --- a/litellm/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx +++ b/litellm/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx @@ -14,7 +14,8 @@ import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"; import { StorePromptsMetadataToggle } from "../common_components/StorePromptsToggle"; import NumericalInput from "../shared/numerical_input"; import SkillSelector from "../skills/SkillSelector"; -import { moveTagsOutOfMetadataJson } from "./keyEditFieldNormalizers"; +import type { BudgetAlertRow } from "../key_team_helpers/budgetAlertEmails"; +import { moveBudgetAlertEmailsOutOfMetadataJson, moveTagsOutOfMetadataJson } from "./keyEditFieldNormalizers"; import { AgentsAndGroups, KeyEditFormValues } from "./keyEditFormValues"; export const labelWithHint = (label: React.ReactNode, hint: string): React.ReactNode => ( @@ -186,8 +187,42 @@ export const moveMetadataTagsToTagsField = (form: KeyEditForm, t: TFunction): vo } }; -export const KeyMetadataField = ({ form }: { form: KeyEditForm }) => { +export const moveMetadataBudgetAlertsToRows = ( + metadata: string | undefined, + rows: readonly BudgetAlertRow[] | undefined, + write: (metadata: string, rows: BudgetAlertRow[]) => void, + t: TFunction, +): void => { + const moved = moveBudgetAlertEmailsOutOfMetadataJson(metadata, rows); + if (moved === null) return; + write(moved.metadata, [...moved.rows]); + const thresholds = moved.movedThresholds.map((threshold) => `${threshold}%`).join(", "); + toast.info(t("keyDetail.maxBudgetAlertEmails.movedThresholds", { thresholds })); +}; + +export const moveMetadataBudgetAlertsToKeyRows = (form: KeyEditForm, t: TFunction): void => + moveMetadataBudgetAlertsToRows( + form.getValues("metadata"), + form.getValues("max_budget_alert_emails"), + (metadata, rows) => { + form.setValue("metadata", metadata, { shouldDirty: true }); + form.setValue("max_budget_alert_emails", rows, { shouldDirty: true }); + }, + t, + ); + +export const KeyMetadataField = ({ + form, + budgetAlertRowsShown, +}: { + form: KeyEditForm; + budgetAlertRowsShown: boolean; +}) => { const { t } = useTranslation(); + const moveTypedFieldsOutOfMetadata = () => { + moveMetadataTagsToTagsField(form, t); + if (budgetAlertRowsShown) moveMetadataBudgetAlertsToKeyRows(form, t); + }; return ( { rows={10} onBlur={() => { field.onBlur(); - moveMetadataTagsToTagsField(form, t); + moveTypedFieldsOutOfMetadata(); }} /> diff --git a/litellm/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx b/litellm/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx index 2c57445104f..4ddb5001a06 100644 --- a/litellm/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx +++ b/litellm/ui/litellm-dashboard/src/components/templates/KeyInfoView.handleKeyUpdate.test.tsx @@ -196,7 +196,8 @@ vi.mock("./key_edit_view", async () => { React.createElement( "button", { - onClick: () => props.onSubmit((globalThis as any).__TEST_FORM_VALUES ?? {}), + onClick: () => + props.onSubmit((globalThis as any).__TEST_FORM_VALUES ?? {}, (globalThis as any).__TEST_DIRTY_FIELDS ?? []), }, "Mock Submit", ), @@ -231,7 +232,7 @@ const baseKeyData = { next_rotation_at: null as any, }; -const renderView = (premiumUser: boolean) => { +const renderView = (premiumUser: boolean, keyData: Record = baseKeyData) => { // Configure the mock for this test mockUseAuthorized.mockReturnValue({ accessToken: "access_abc", @@ -245,19 +246,14 @@ const renderView = (premiumUser: boolean) => { }); return render( - {}} - keyData={baseKeyData as any} - onKeyDataUpdate={() => {}} - teams={[]} - />, + {}} keyData={keyData as any} onKeyDataUpdate={() => {}} teams={[]} />, ); }; beforeEach(() => { vi.clearAllMocks(); (globalThis as any).__TEST_FORM_VALUES = undefined; + (globalThis as any).__TEST_DIRTY_FIELDS = undefined; }); // ---- Tests ---- @@ -715,3 +711,171 @@ describe("KeyInfoView handleKeyUpdate organization", () => { expect(shownOrganization()).toHaveTextContent("org-1"); }); }); + +describe("KeyInfoView handleKeyUpdate max_budget_alert_emails", () => { + const submit = async ( + formValues: Record, + dirtyFields: readonly string[] = ["max_budget_alert_emails"], + keyData: Record = baseKeyData, + ) => { + renderView(true, keyData); + fireEvent.click(screen.getByText("Settings")); + fireEvent.click(screen.getByText("Edit Settings")); + Object.assign(globalThis, { + __TEST_FORM_VALUES: { token: "tok_123", ...formValues }, + __TEST_DIRTY_FIELDS: dirtyFields, + }); + fireEvent.click(screen.getByText("Mock Submit")); + await waitFor(() => expect(keyUpdateCallMock).toHaveBeenCalled()); + return keyUpdateCallMock.mock.calls[0][1]; + }; + + it("writes the alert rows into metadata next to the JSON box content and the dedicated metadata fields", async () => { + const formValues = { + metadata: JSON.stringify({ cost_center: "ops", nested: { owner: "platform" } }), + tags: ["prod"], + guardrails: ["gr-1"], + logging_settings: [{ callback_name: "langfuse" }], + disabled_callbacks: ["langsmith"], + max_budget_alert_emails: [ + { threshold: 80, emails: "shared-mailbox@example.com, finance-dl@example.com" }, + { threshold: 50, emails: "" }, + ], + }; + const expectedMetadata = { + cost_center: "ops", + nested: { owner: "platform" }, + tags: ["prod"], + guardrails: ["gr-1"], + logging: [{ callback_name: "langfuse" }], + litellm_disabled_callbacks: ["langsmith"], + max_budget_alert_emails: { "80": ["shared-mailbox@example.com", "finance-dl@example.com"], "50": [] }, + }; + + const sentPayload = await submit(formValues); + + expect(sentPayload.metadata).toEqual(expectedMetadata); + expect(sentPayload).not.toHaveProperty("max_budget_alert_emails"); + }); + + it("clears the stored recipients when every row is removed, even if the JSON box still names them", async () => { + const sentPayload = await submit({ + metadata: JSON.stringify({ cost_center: "ops", max_budget_alert_emails: { "80": ["old@example.com"] } }), + max_budget_alert_emails: [], + }); + + expect(sentPayload.metadata).toEqual({ cost_center: "ops" }); + }); + + it("writes the alert rows when the metadata JSON box is left empty", async () => { + const sentPayload = await submit({ + metadata: "", + max_budget_alert_emails: [{ threshold: 90, emails: "dl@example.com" }], + }); + + expect(sentPayload.metadata).toEqual({ max_budget_alert_emails: { "90": ["dl@example.com"] } }); + }); + + it("sends the stored map verbatim on a save that does not edit the rows", async () => { + const storedAlertEmails = { + "50": ["a@example.com"], + "75": null, + "80": ["b@example.com", null], + "150": ["c@example.com"], + }; + const sentPayload = await submit( + { + key_alias: "renamed", + metadata: JSON.stringify({ cost_center: "ops", max_budget_alert_emails: { "150": ["c@example.com"] } }), + max_budget_alert_emails: [{ threshold: 50, emails: "a@example.com" }], + }, + ["key_alias"], + { ...baseKeyData, metadata: { cost_center: "ops", max_budget_alert_emails: storedAlertEmails } }, + ); + + expect(sentPayload.metadata).toEqual({ cost_center: "ops", max_budget_alert_emails: storedAlertEmails }); + }); + const KEY_WITH_HIDDEN_ALERT = { + ...baseKeyData, + metadata: { + cost_center: "ops", + max_budget_alert_emails: { + "50": ["old@example.com"], + "80": ["removed@example.com"], + "150": ["c@example.com"], + }, + }, + }; + + it("keeps the entry the rows cannot show that the JSON box still holds when the rows were edited", async () => { + const sentPayload = await submit( + { + metadata: JSON.stringify({ cost_center: "ops", max_budget_alert_emails: { "150": ["c@example.com"] } }), + max_budget_alert_emails: [{ threshold: 50, emails: "new@example.com" }], + }, + ["max_budget_alert_emails"], + KEY_WITH_HIDDEN_ALERT, + ); + + expect(sentPayload.metadata).toEqual({ + cost_center: "ops", + max_budget_alert_emails: { "50": ["new@example.com"], "150": ["c@example.com"] }, + }); + }); + + it("keeps only the box's hidden entry when every row is removed", async () => { + const sentPayload = await submit( + { + metadata: JSON.stringify({ cost_center: "ops", max_budget_alert_emails: { "150": ["c@example.com"] } }), + max_budget_alert_emails: [], + }, + ["max_budget_alert_emails"], + KEY_WITH_HIDDEN_ALERT, + ); + + expect(sentPayload.metadata).toEqual({ + cost_center: "ops", + max_budget_alert_emails: { "150": ["c@example.com"] }, + }); + }); + + it("clears the map when every row is removed and the hidden entry is deleted from the JSON box", async () => { + const sentPayload = await submit( + { metadata: JSON.stringify({ cost_center: "ops" }), max_budget_alert_emails: [] }, + ["max_budget_alert_emails"], + KEY_WITH_HIDDEN_ALERT, + ); + + expect(sentPayload.metadata).toEqual({ cost_center: "ops" }); + }); + + it("saves an entry the rows cannot show typed into the JSON box on an untouched-rows edit", async () => { + const sentPayload = await submit( + { + metadata: JSON.stringify({ cost_center: "ops", max_budget_alert_emails: { "0": ["zero@example.com"] } }), + max_budget_alert_emails: [], + }, + ["metadata"], + { ...baseKeyData, metadata: { cost_center: "ops" } }, + ); + + expect(sentPayload.metadata).toEqual({ + cost_center: "ops", + max_budget_alert_emails: { "0": ["zero@example.com"] }, + }); + }); + + it("saves a stored max_budget_alert_emails string back unchanged on an unrelated edit", async () => { + const sentPayload = await submit( + { + key_alias: "renamed", + metadata: JSON.stringify({ cost_center: "ops", max_budget_alert_emails: "ops@qa.test" }), + max_budget_alert_emails: [], + }, + ["key_alias"], + { ...baseKeyData, metadata: { cost_center: "ops", max_budget_alert_emails: "ops@qa.test" } }, + ); + + expect(sentPayload.metadata).toEqual({ cost_center: "ops", max_budget_alert_emails: "ops@qa.test" }); + }); +}); diff --git a/litellm/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.test.ts b/litellm/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.test.ts index 9223cc6dbc8..cbae3e67540 100644 --- a/litellm/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.test.ts +++ b/litellm/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.test.ts @@ -1,6 +1,6 @@ import { describe, expect, it } from "vitest"; -import { moveTagsOutOfMetadataJson } from "./keyEditFieldNormalizers"; +import { moveBudgetAlertEmailsOutOfMetadataJson, moveTagsOutOfMetadataJson } from "./keyEditFieldNormalizers"; describe("moveTagsOutOfMetadataJson", () => { it("moves a tags array out of the JSON and appends it to the current tags", () => { @@ -48,3 +48,45 @@ describe("moveTagsOutOfMetadataJson", () => { expect(moveTagsOutOfMetadataJson("null", ["a"])).toBeNull(); }); }); + +describe("moveBudgetAlertEmailsOutOfMetadataJson", () => { + it("moves the typed thresholds into the rows, letting a typed entry replace the row with its threshold", () => { + const typed = JSON.stringify({ + env: "prod", + max_budget_alert_emails: { "80": "typed@example.com, dl@example.com", "90": null }, + }); + const rows = [ + { threshold: 50, emails: "kept@example.com" }, + { threshold: 80, emails: "replaced@example.com" }, + ]; + + expect(moveBudgetAlertEmailsOutOfMetadataJson(typed, rows)).toEqual({ + metadata: '{\n "env": "prod"\n}', + rows: [ + { threshold: 50, emails: "kept@example.com" }, + { threshold: 80, emails: "typed@example.com, dl@example.com" }, + { threshold: 90, emails: "" }, + ], + movedThresholds: [80, 90], + }); + }); + + it("leaves the typed entries the rows cannot show in the JSON", () => { + const typed = JSON.stringify({ max_budget_alert_emails: { "80": ["a@example.com"], "150": ["b@example.com"] } }); + + expect(moveBudgetAlertEmailsOutOfMetadataJson(typed, [])).toEqual({ + metadata: JSON.stringify({ max_budget_alert_emails: { "150": ["b@example.com"] } }, null, 2), + rows: [{ threshold: 80, emails: "a@example.com" }], + movedThresholds: [80], + }); + }); + + it("returns null when nothing typed can become a row", () => { + expect( + moveBudgetAlertEmailsOutOfMetadataJson('{"max_budget_alert_emails": {"150": ["b@example.com"]}}', []), + ).toBeNull(); + expect(moveBudgetAlertEmailsOutOfMetadataJson('{"env": "prod"}', [])).toBeNull(); + expect(moveBudgetAlertEmailsOutOfMetadataJson('{"max_budget_alert_emails": {', [])).toBeNull(); + expect(moveBudgetAlertEmailsOutOfMetadataJson(undefined, [])).toBeNull(); + }); +}); diff --git a/litellm/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts b/litellm/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts index 529bffb9a84..8dec643b747 100644 --- a/litellm/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts +++ b/litellm/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts @@ -1,5 +1,11 @@ import type { TFunction } from "i18next"; import i18n from "@/i18n/i18n"; +import { + type BudgetAlertRow, + budgetAlertEmailsRowsCannotShow, + budgetAlertRowsFromMetadata, + KEY_MAX_BUDGET_ALERT_EMAILS_KEY, +} from "../key_team_helpers/budgetAlertEmails"; const WORD_FORM_BUDGET_DURATIONS: Record = { hourly: "1h", @@ -75,6 +81,27 @@ export const moveTagsOutOfMetadataJson = ( return { metadata: JSON.stringify(rest, null, 2), tags: [...existing, ...movedTags], movedTags }; }; +export type MovedMetadataBudgetAlerts = { + readonly metadata: string; + readonly rows: readonly BudgetAlertRow[]; + readonly movedThresholds: readonly (number | null)[]; +}; + +export const moveBudgetAlertEmailsOutOfMetadataJson = ( + metadataJson: string | undefined, + currentRows: readonly BudgetAlertRow[] | undefined, +): MovedMetadataBudgetAlerts | null => { + const parsed = metadataJson === undefined ? null : parseJsonObject(metadataJson); + const typedRows = budgetAlertRowsFromMetadata(parsed, KEY_MAX_BUDGET_ALERT_EMAILS_KEY); + if (parsed === null || typedRows.length === 0) return null; + const movedThresholds = typedRows.map((row) => row.threshold); + const unmoved = budgetAlertEmailsRowsCannotShow(parsed, KEY_MAX_BUDGET_ALERT_EMAILS_KEY); + const { [KEY_MAX_BUDGET_ALERT_EMAILS_KEY]: _moved, ...rest } = parsed; + const metadata = Object.keys(unmoved).length > 0 ? { ...rest, [KEY_MAX_BUDGET_ALERT_EMAILS_KEY]: unmoved } : rest; + const keptRows = (currentRows ?? []).filter((row) => !movedThresholds.includes(row.threshold)); + return { metadata: JSON.stringify(metadata, null, 2), rows: [...keptRows, ...typedRows], movedThresholds }; +}; + export const currentValuePlaceholder = ( premiumUser: boolean, current: unknown, diff --git a/litellm/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts b/litellm/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts index 9ad3efe84ce..800a5bdf079 100644 --- a/litellm/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts +++ b/litellm/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts @@ -4,7 +4,13 @@ import { UseFormReturn, UseFormSetValue } from "react-hook-form"; import { KeyResponse } from "../key_team_helpers/key_list"; import { ToolPermissionsChangeOrigin } from "../mcp_server_management/MCPToolPermissions"; -import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata } from "../key_info_utils"; +import { extractLoggingSettings, formatMetadataForDisplay, stripFieldManagedKeysFromMetadata } from "../key_info_utils"; +import { + type BudgetAlertRow, + budgetAlertRowsFromMetadata, + budgetAlertRowsSchema, + KEY_MAX_BUDGET_ALERT_EMAILS_KEY, +} from "../key_team_helpers/budgetAlertEmails"; import { mapInternalToDisplayNames } from "../callback_info_helpers"; import { estimateChecks, estimateFields } from "./estimatedOutputTokens"; import { canonicalBudgetDuration, keyOrganizationId } from "./keyEditFieldNormalizers"; @@ -26,6 +32,7 @@ export interface KeyEditFormValues { allowed_routes?: string; max_budget?: number | string | null; soft_budget?: number | string | null; + max_budget_alert_emails?: BudgetAlertRow[]; budget_duration?: string | null; tpm_limit?: number | string | null; tpm_limit_type?: string | null; @@ -76,6 +83,7 @@ export const toKeyEditFormValues = (keyData: KeyResponse): KeyEditFormValues => max_budget: keyData.max_budget, soft_budget: (keyData.litellm_budget_table as { soft_budget?: number | null } | null | undefined)?.soft_budget ?? null, + max_budget_alert_emails: [...budgetAlertRowsFromMetadata(keyData.metadata, KEY_MAX_BUDGET_ALERT_EMAILS_KEY)], budget_duration: canonicalBudgetDuration(keyData.budget_duration), tpm_limit: keyData.tpm_limit, tpm_limit_type: (keyData as { tpm_limit_type?: string | null }).tpm_limit_type ?? null, @@ -114,7 +122,7 @@ export const toKeyEditFormValues = (keyData: KeyResponse): KeyEditFormValues => team_id: keyData.team_id, project_id: keyData.project_id, logging_settings: extractLoggingSettings(keyData.metadata), - metadata: formatMetadataForDisplay(stripTagsFromMetadata(keyData.metadata)), + metadata: formatMetadataForDisplay(stripFieldManagedKeysFromMetadata(keyData.metadata)), duration: (keyData as { duration?: string }).duration ?? "", token: keyData.token || keyData.token_id, disabled_callbacks: Array.isArray(readMetadata(keyData, "litellm_disabled_callbacks")) @@ -130,6 +138,11 @@ export const keyEditFormSchema = z.object({ allowed_routes: z.custom().optional(), max_budget: z.custom().optional(), soft_budget: z.custom().optional(), + max_budget_alert_emails: budgetAlertRowsSchema((issue) => + issue === "thresholdRange" + ? i18n.t("keyDetail.maxBudgetAlertEmails.thresholdRange") + : i18n.t("keyDetail.maxBudgetAlertEmails.duplicateThreshold"), + ).optional(), budget_duration: z.custom().optional(), tpm_limit: z.custom().optional(), tpm_limit_type: z.custom().optional(), @@ -187,6 +200,7 @@ export const toSubmittedValues = ( allowed_routes: values.allowed_routes, max_budget: values.max_budget, soft_budget: values.soft_budget, + max_budget_alert_emails: values.max_budget_alert_emails, budget_duration: values.budget_duration, tpm_limit: values.tpm_limit, tpm_limit_type: values.tpm_limit_type, diff --git a/litellm/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx b/litellm/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx index 88d063358ca..ae21dfe7100 100644 --- a/litellm/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx +++ b/litellm/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx @@ -15,6 +15,7 @@ import { modelAvailableCall, vectorStoreListCall, } from "../networking"; +import { toast } from "@/lib/toast"; import { KeyEditView } from "./key_edit_view"; import { teamAdminMemberKeyPayload } from "./teamAdminMemberKeyPayload"; @@ -2577,6 +2578,7 @@ describe("KeyEditView", () => { models: [], max_budget: 0, soft_budget: null, + max_budget_alert_emails: [], budget_duration: "30d", tpm_limit: 10, tpm_limit_type: null, @@ -3308,4 +3310,185 @@ describe("KeyEditView", () => { }); }); }); + + describe("budget alert recipients", () => { + const KEY_WITH_ALERTS: KeyResponse = { + ...MOCK_KEY_DATA, + max_budget: 100, + metadata: { + ...MOCK_KEY_DATA.metadata, + cost_center: "ops", + max_budget_alert_emails: { "80": ["shared-mailbox@example.com", "finance-dl@example.com"], "50": [] }, + }, + }; + + const renderAlerts = ( + onSubmit: (values: Record) => Promise, + editingMemberKeyAsTeamAdmin = false, + keyData: KeyResponse = KEY_WITH_ALERTS, + ) => + renderWithProviders( + {}} + onSubmit={onSubmit} + accessToken={"test-token"} + userID={"test-user"} + userRole={"Admin"} + premiumUser={true} + editingMemberKeyAsTeamAdmin={editingMemberKeyAsTeamAdmin} + />, + ); + + const save = async (onSubmit: ReturnType) => { + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + await waitFor(() => expect(onSubmit).toHaveBeenCalledTimes(1)); + return onSubmit.mock.calls[0][0]; + }; + + it("shows the stored thresholds as rows and keeps them out of the metadata JSON box", async () => { + renderAlerts(vi.fn()); + + const thresholds = await screen.findAllByPlaceholderText("% of budget"); + expect(thresholds.map((input) => (input as HTMLInputElement).value)).toEqual(["50", "80"]); + expect(screen.getAllByPlaceholderText("Additional recipients, e.g. finance@example.com")[1]).toHaveValue( + "shared-mailbox@example.com, finance-dl@example.com", + ); + expect(JSON.parse((screen.getByLabelText("Metadata") as HTMLTextAreaElement).value)).toEqual({ + cost_center: "ops", + }); + }); + + it("shows only the stored entries the rows cannot show in the metadata JSON box", async () => { + renderAlerts(vi.fn(), false, { + ...KEY_WITH_ALERTS, + metadata: { + cost_center: "ops", + max_budget_alert_emails: { "80": ["row@example.com"], "150": ["over@example.com"], "0": [] }, + }, + }); + + const thresholds = await screen.findAllByPlaceholderText("% of budget"); + expect(thresholds.map((input) => (input as HTMLInputElement).value)).toEqual(["80"]); + expect(JSON.parse((screen.getByLabelText("Metadata") as HTMLTextAreaElement).value)).toEqual({ + cost_center: "ops", + max_budget_alert_emails: { "150": ["over@example.com"], "0": [] }, + }); + }); + + it("submits an untouched key's rows unchanged", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderAlerts(onSubmit); + await screen.findAllByPlaceholderText("% of budget"); + + expect((await save(onSubmit)).max_budget_alert_emails).toEqual([ + { threshold: 50, emails: "" }, + { threshold: 80, emails: "shared-mailbox@example.com, finance-dl@example.com" }, + ]); + expect(onSubmit.mock.calls[0][1]).not.toContain("max_budget_alert_emails"); + }); + + it("submits an added row and drops a removed one", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderAlerts(onSubmit); + await screen.findAllByPlaceholderText("% of budget"); + + await userEvent.click(screen.getAllByRole("button", { name: "Remove budget alert threshold" })[0]); + await userEvent.click(screen.getByRole("button", { name: "Add Budget Alert Threshold" })); + const thresholds = screen.getAllByPlaceholderText("% of budget"); + fireEvent.change(thresholds[thresholds.length - 1], { target: { value: "100" } }); + const recipients = screen.getAllByPlaceholderText("Additional recipients, e.g. finance@example.com"); + fireEvent.change(recipients[recipients.length - 1], { target: { value: "ap-team@example.com" } }); + + expect((await save(onSubmit)).max_budget_alert_emails).toEqual([ + { threshold: 80, emails: "shared-mailbox@example.com, finance-dl@example.com" }, + { threshold: 100, emails: "ap-team@example.com" }, + ]); + expect(onSubmit.mock.calls[0][1]).toContain("max_budget_alert_emails"); + }); + + it("reports the rows as edited when every row is removed, so the save clears them", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderAlerts(onSubmit); + await screen.findAllByPlaceholderText("% of budget"); + + for (const button of screen.getAllByRole("button", { name: "Remove budget alert threshold" }).reverse()) { + await userEvent.click(button); + } + + expect((await save(onSubmit)).max_budget_alert_emails).toEqual([]); + expect(onSubmit.mock.calls[0][1]).toContain("max_budget_alert_emails"); + }); + + it("moves a recipient map typed into the metadata JSON box into the rows on blur", async () => { + const info = vi.spyOn(toast, "info"); + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderAlerts(onSubmit); + await screen.findAllByPlaceholderText("% of budget"); + + const metadataBox = screen.getByLabelText("Metadata"); + fireEvent.change(metadataBox, { + target: { + value: JSON.stringify({ + cost_center: "ops", + max_budget_alert_emails: { "80": ["typed@example.com"], "95": [] }, + }), + }, + }); + fireEvent.blur(metadataBox); + + expect(JSON.parse((metadataBox as HTMLTextAreaElement).value)).toEqual({ cost_center: "ops" }); + expect(info).toHaveBeenCalledWith( + "Moved alert thresholds 80%, 95% from metadata to the Budget Alert Thresholds rows", + ); + expect((await save(onSubmit)).max_budget_alert_emails).toEqual([ + { threshold: 50, emails: "" }, + { threshold: 80, emails: "typed@example.com" }, + { threshold: 95, emails: "" }, + ]); + expect(onSubmit.mock.calls[0][1]).toContain("max_budget_alert_emails"); + }); + + it("moves a recipient map typed into the metadata JSON box into the rows even without a blur", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderAlerts(onSubmit); + await screen.findAllByPlaceholderText("% of budget"); + + fireEvent.change(screen.getByLabelText("Metadata"), { + target: { + value: JSON.stringify({ cost_center: "ops", max_budget_alert_emails: { "80": ["typed@example.com"] } }), + }, + }); + fireEvent.submit(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => expect(onSubmit).toHaveBeenCalled()); + expect(onSubmit.mock.calls[0][0].max_budget_alert_emails).toEqual([ + { threshold: 50, emails: "" }, + { threshold: 80, emails: "typed@example.com" }, + ]); + expect(onSubmit.mock.calls[0][1]).toContain("max_budget_alert_emails"); + }); + + it("hides the rows from a team admin editing a member's key", async () => { + renderAlerts(vi.fn(), true); + + expect(await screen.findByLabelText("Metadata")).toBeInTheDocument(); + expect(screen.queryByPlaceholderText("% of budget")).not.toBeInTheDocument(); + expect(screen.queryByRole("button", { name: "Add Budget Alert Threshold" })).not.toBeInTheDocument(); + }); + + it("blocks the save on a duplicate threshold", async () => { + const onSubmit = vi.fn().mockResolvedValue(undefined); + renderAlerts(onSubmit); + await screen.findAllByPlaceholderText("% of budget"); + + await userEvent.click(screen.getByRole("button", { name: "Add Budget Alert Threshold" })); + const thresholds = screen.getAllByPlaceholderText("% of budget"); + fireEvent.change(thresholds[thresholds.length - 1], { target: { value: "80" } }); + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + expect(await screen.findAllByText("Duplicate threshold")).toHaveLength(2); + expect(onSubmit).not.toHaveBeenCalled(); + }); + }); }); diff --git a/litellm/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/litellm/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index bafeedd5be7..1bb6676470c 100644 --- a/litellm/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/litellm/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -38,6 +38,7 @@ import { KeyRateLimitFields, KeyTypeSelect, labelWithHint, + moveMetadataBudgetAlertsToKeyRows, moveMetadataTagsToTagsField, } from "./KeyEditViewControls"; import { @@ -49,6 +50,7 @@ import { clearMcpGrants, writeMcpToolPermissions, } from "./keyEditFormValues"; +import { BudgetAlertEmailsEditor } from "../key_team_helpers/BudgetAlertEmailsEditor"; import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor"; import { END_USER_BUDGET_HINT, EndUserBudgetSelect } from "../key_team_helpers/EndUserBudgetSelect"; import { @@ -90,6 +92,7 @@ interface KeyEditViewProps { userID: string | null; userRole: string | null; premiumUser?: boolean; + editingMemberKeyAsTeamAdmin?: boolean; } const sessionKey = (keyData: KeyResponse) => keyData.token || keyData.token_id; @@ -105,6 +108,7 @@ function KeyEditSession({ userID, userRole, premiumUser = false, + editingMemberKeyAsTeamAdmin = false, }: KeyEditViewProps) { const { t } = useTranslation(); const canEditGuardrails = premiumUser || (userRole != null && rolesWithWriteAccess.includes(userRole)); @@ -323,7 +327,7 @@ function KeyEditSession({ ...values, ...(detachProject && enableProjectsUI && canDetachProject ? { project_id: null } : {}), }), - [...Object.keys(dirtyFields), ...(budgetLimitsUnchanged ? [] : ["budget_limits"])], + [...Object.keys(form.formState.dirtyFields), ...(budgetLimitsUnchanged ? [] : ["budget_limits"])], ); } finally { setIsKeySaving(false); @@ -373,16 +377,17 @@ function KeyEditSession({ ); }; + const submitForm = (event: React.FormEvent) => { + moveMetadataTagsToTagsField(form, t); + if (!editingMemberKeyAsTeamAdmin) moveMetadataBudgetAlertsToKeyRows(form, t); + return form.handleSubmit((values) => handleSubmit(toSubmittedValues(values, { canViewPolicies, canViewPrompts })))( + event, + ); + }; + return ( -
{ - moveMetadataTagsToTagsField(form, t); - return form.handleSubmit((values) => - handleSubmit(toSubmittedValues(values, { canViewPolicies, canViewPrompts })), - )(event); - }} - > + {(field) => } @@ -467,6 +472,23 @@ function KeyEditSession({ placeholder={t("keyDetail.keyEditView.fields.softBudgetPlaceholder")} /> + {!editingMemberKeyAsTeamAdmin && ( + + )} + - +
diff --git a/litellm/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/litellm/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index 8a30c6ff7a7..78607a70058 100644 --- a/litellm/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/litellm/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -17,6 +17,7 @@ import { useModelAccessGroupNames } from "@/app/(dashboard)/hooks/models/useMode const editViewMocks = vi.hoisted(() => ({ onSubmit: undefined as ((v: Record) => Promise) | undefined, + editingMemberKeyAsTeamAdmin: undefined as boolean | undefined, })); vi.mock("next/navigation", () => ({ useRouter: () => ({ push: vi.fn() }), usePathname: () => "/ui/alerts" })); @@ -26,8 +27,15 @@ vi.mock("@/app/(dashboard)/hooks/organizations/useOrganizations", () => ({ })); vi.mock("./key_edit_view", () => ({ - KeyEditView: ({ onSubmit }: { onSubmit: (v: Record) => Promise }) => { + KeyEditView: ({ + onSubmit, + editingMemberKeyAsTeamAdmin, + }: { + onSubmit: (v: Record) => Promise; + editingMemberKeyAsTeamAdmin?: boolean; + }) => { editViewMocks.onSubmit = onSubmit; + editViewMocks.editingMemberKeyAsTeamAdmin = editingMemberKeyAsTeamAdmin; return
; }, })); @@ -892,7 +900,9 @@ describe("KeyInfoView", () => { await renderAndOpenSettingsTab({ ...MOCK_KEY_DATA, user_id: "someone-else-id" }); - expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + await userEvent.click(screen.getByRole("button", { name: /edit settings/i })); + expect(await screen.findByTestId("key-edit-view-stub")).toBeInTheDocument(); + expect(editViewMocks.editingMemberKeyAsTeamAdmin).toBe(false); }); it("should show the Edit Settings button when the user is the key owner", async () => { @@ -961,7 +971,9 @@ describe("KeyInfoView", () => { await renderAndOpenSettingsTab({ ...MOCK_KEY_DATA, team_id: teamId, user_id: "other-user-id" }); - expect(screen.getByRole("button", { name: /edit settings/i })).toBeInTheDocument(); + await userEvent.click(screen.getByRole("button", { name: /edit settings/i })); + expect(await screen.findByTestId("key-edit-view-stub")).toBeInTheDocument(); + expect(editViewMocks.editingMemberKeyAsTeamAdmin).toBe(true); }); }); diff --git a/litellm/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/litellm/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 317a3c309ab..8e5f9262101 100644 --- a/litellm/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/litellm/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -39,6 +39,12 @@ import DeleteResourceModal from "../common_components/DeleteResourceModal"; import RouterSettingsSummary from "../common_components/RouterSettingsSummary"; import { hasRouterSettings } from "../common_components/routerSettingsPayload"; import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata } from "../key_info_utils"; +import { + budgetAlertEmailsFromRows, + budgetAlertEmailsMetadata, + budgetAlertEmailsRowsCanShow, + KEY_MAX_BUDGET_ALERT_EMAILS_KEY, +} from "../key_team_helpers/budgetAlertEmails"; import { KeyResponse } from "../key_team_helpers/key_list"; import LoggingSettingsView from "../logging_settings_view"; import { toast } from "@/lib/toast"; @@ -219,6 +225,15 @@ export default function KeyInfoView({ ); } + const memberKeyEditContext = { + userRole: userRole || "", + userId: userID || "", + keyUserId: currentKeyData.user_id, + keyTeamId: currentKeyData.team_id, + teamMembers: teamsData?.find((team) => team.team_id === currentKeyData.team_id)?.members_with_roles, + }; + const editingMemberKeyAsTeamAdmin = isTeamAdminEditingMemberKey(memberKeyEditContext); + const handleKeyUpdate = async (formValues: Record, dirtyFields: readonly string[] = []) => { try { if (!accessToken) return; @@ -334,6 +349,12 @@ export default function KeyInfoView({ formValues.tpd_limit = mapEmptyStringToNull(formValues.tpd_limit); formValues.max_parallel_requests = mapEmptyStringToNull(formValues.max_parallel_requests); + const budgetAlertRowEntries = dirtyFields.includes(KEY_MAX_BUDGET_ALERT_EMAILS_KEY) + ? budgetAlertEmailsFromRows(formValues.max_budget_alert_emails ?? []) + : budgetAlertEmailsRowsCanShow(currentKeyData.metadata, KEY_MAX_BUDGET_ALERT_EMAILS_KEY); + const budgetAlertEmailsWithBoxEntries = (boxMetadata: unknown) => + budgetAlertEmailsMetadata(budgetAlertRowEntries, KEY_MAX_BUDGET_ALERT_EMAILS_KEY, boxMetadata); + // Convert metadata back to an object if it exists and is a string if (formValues.metadata && typeof formValues.metadata === "string") { try { @@ -342,9 +363,12 @@ export default function KeyInfoView({ if ("tags" in parsedMetadata) { delete parsedMetadata["tags"]; } + const budgetAlertEmails = budgetAlertEmailsWithBoxEntries(parsedMetadata); + delete parsedMetadata[KEY_MAX_BUDGET_ALERT_EMAILS_KEY]; formValues.metadata = { ...parsedMetadata, ...(Array.isArray(formValues.tags) && formValues.tags.length > 0 ? { tags: formValues.tags } : {}), + ...budgetAlertEmails, ...(formValues.guardrails?.length > 0 ? { guardrails: formValues.guardrails } : {}), ...(Array.isArray(formValues.logging_settings) && formValues.logging_settings.length > 0 ? { logging: formValues.logging_settings } @@ -362,10 +386,11 @@ export default function KeyInfoView({ } } else { const baseMetadata = formValues.metadata || {}; - const { tags: _omitTags, ...rest } = baseMetadata; + const { tags: _omitTags, [KEY_MAX_BUDGET_ALERT_EMAILS_KEY]: _boxAlertEmails, ...rest } = baseMetadata; formValues.metadata = { ...rest, ...(Array.isArray(formValues.tags) && formValues.tags.length > 0 ? { tags: formValues.tags } : {}), + ...budgetAlertEmailsWithBoxEntries(baseMetadata), ...(formValues.guardrails?.length > 0 ? { guardrails: formValues.guardrails } : {}), ...(Array.isArray(formValues.logging_settings) && formValues.logging_settings.length > 0 ? { logging: formValues.logging_settings } @@ -383,6 +408,7 @@ export default function KeyInfoView({ delete formValues.tags; } delete formValues.logging_settings; + delete formValues.max_budget_alert_emails; // Normalize any legacy word-form budget_duration to the canonical API format if (formValues.budget_duration) { @@ -395,14 +421,6 @@ export default function KeyInfoView({ formValues.budget_duration = wordToCanonical[formValues.budget_duration] ?? formValues.budget_duration; } - const memberKeyEditContext = { - userRole: userRole || "", - userId: userID || "", - keyUserId: currentKeyData.user_id, - keyTeamId: currentKeyData.team_id, - teamMembers: teamsData?.find((team) => team.team_id === currentKeyData.team_id)?.members_with_roles, - }; - const editingMemberKeyAsTeamAdmin = isTeamAdminEditingMemberKey(memberKeyEditContext); if (editingMemberKeyAsTeamAdmin) { const trimmed = teamAdminMemberKeyPayload(formValues, dirtyFields); if (trimmed.kind === "blocked") { @@ -1010,6 +1028,7 @@ export default function KeyInfoView({ userID={userID} userRole={userRole} premiumUser={premiumUser} + editingMemberKeyAsTeamAdmin={editingMemberKeyAsTeamAdmin} /> ) : (
diff --git a/litellm/ui/litellm-dashboard/src/i18n/locales/en.json b/litellm/ui/litellm-dashboard/src/i18n/locales/en.json index 6e09e3eb785..a2369c9f077 100644 --- a/litellm/ui/litellm-dashboard/src/i18n/locales/en.json +++ b/litellm/ui/litellm-dashboard/src/i18n/locales/en.json @@ -2391,6 +2391,17 @@ "description": "Tags are managed by the Tags field above. A tags array typed here is moved to that field." } }, + "maxBudgetAlertEmails": { + "budgetAlertThresholds": "Budget Alert Thresholds", + "budgetAlertThresholdsHint": "Thresholds are a percentage of the key's max budget and replace this key's default 80% alert. Global alert thresholds configured on the proxy still apply. Thresholds outside 1 to 100 can be edited in the Metadata box. The key owner's email is always included. Add comma-separated addresses, such as a shared mailbox or distribution list, to notify others as well. Requires email alerting to be configured on the proxy.", + "thresholdPlaceholder": "% of budget", + "recipientsPlaceholder": "Additional recipients, e.g. finance@example.com", + "removeThreshold": "Remove budget alert threshold", + "addThreshold": "Add Budget Alert Threshold", + "thresholdRange": "Enter a whole number from 1 to 100", + "duplicateThreshold": "Duplicate threshold", + "movedThresholds": "Moved alert thresholds {{thresholds}} from metadata to the Budget Alert Thresholds rows" + }, "estimatedOutputTokens": { "invalidPerModel": "Enter a JSON object of positive integers, e.g. {\"gpt-4\": 4096}", "invalidPositive": "Enter a positive integer", @@ -20670,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 2d9f7cb42ad..72a2c20824f 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; @@ -68150,6 +68210,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; @@ -68310,6 +68408,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; @@ -68572,6 +68708,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/litellm/utils.py b/litellm/utils.py index 67b7b4f3d1e..118f89766a9 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5024,9 +5024,7 @@ def get_optional_params_embeddings( optional_params["dimensions"] = non_default_params.pop("dimensions") if len(non_default_params.keys()) > 0: if litellm.drop_params is True or drop_params is True: # drop the unsupported non-default values - keys = list(non_default_params.keys()) - for k in keys: - non_default_params.pop(k, None) + non_default_params.clear() else: raise UnsupportedParamsError( status_code=500, @@ -5039,9 +5037,7 @@ def get_optional_params_embeddings( ): if len(non_default_params.keys()) > 0: if litellm.drop_params is True or drop_params is True: # drop the unsupported non-default values - keys = list(non_default_params.keys()) - for k in keys: - non_default_params.pop(k, None) + non_default_params.clear() else: raise UnsupportedParamsError( status_code=500, @@ -5060,7 +5056,7 @@ def get_optional_params_embeddings( additional_drop_params=kwargs.get("additional_drop_params", None), ) - if "extra_body" in final_params and len(final_params["extra_body"]) == 0: + if "extra_body" in final_params and not final_params["extra_body"]: final_params.pop("extra_body", None) return final_params diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 4429fffe652..7c8a48f5aae 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -28650,7 +28650,7 @@ "mode": "chat", "output_cost_per_token": 2.8e-07, "source": "https://docs.fireworks.ai/serverless/pricing", - "successor_model": "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-0731", + "successor_model": "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -29124,7 +29124,7 @@ "mode": "chat", "output_cost_per_token": 2.8e-07, "source": "https://docs.fireworks.ai/serverless/pricing", - "successor_model": "fireworks_ai/deepseek-v4-flash-0731", + "successor_model": "fireworks_ai/deepseek-v4p1-flash", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -83697,7 +83697,7 @@ "mode": "chat", "output_cost_per_token": 2.1e-07, "source": "https://api.fireworks.ai/v1/serverless/models", - "successor_model": "fireworks_ai/accounts/fireworks/models/deepseek-v4-flash-0731", + "successor_model": "fireworks_ai/accounts/fireworks/models/deepseek-v4p1-flash", "supports_function_calling": true, "supports_prompt_caching": true, "supports_reasoning": true, diff --git a/pyproject.toml b/pyproject.toml index ccfbd5d7d9b..3b7e3e154b9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.110.0" +version = "1.111.0" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.16" @@ -88,7 +88,7 @@ proxy = [ "httpx2>=2.10.0,<3", "pydantic>=2.12.0,<3", "litellm-proxy-extras==0.4.115", - "litellm-enterprise==0.1.83", + "litellm-enterprise==0.1.84", "RestrictedPython>=8.5,<9.0", "rich>=13.9.4,<14.0", "InquirerPy>=0.3.4,<1.0", @@ -441,7 +441,7 @@ litellm-enterprise = { workspace = true } members = ["litellm/enterprise", "litellm/proxy-extras", "litellm/packaging/litellm-proxy"] [tool.commitizen] -version = "1.110.0" +version = "1.111.0" version_files = [ "pyproject.toml:^version", ] diff --git a/tests/e2e/ui/tests/proxy-admin/keyBudgetAlertRecipients.spec.ts b/tests/e2e/ui/tests/proxy-admin/keyBudgetAlertRecipients.spec.ts new file mode 100644 index 00000000000..af9c0942331 --- /dev/null +++ b/tests/e2e/ui/tests/proxy-admin/keyBudgetAlertRecipients.spec.ts @@ -0,0 +1,164 @@ +import { test, expect, type Page as PlaywrightPage } from "@playwright/test"; +import { + ADMIN_STORAGE_PATH, + E2E_TEAM_CRUD_ALIAS, + E2E_TEAM_CRUD_ID, +} from "../../constants"; +import { Page } from "../../fixtures/pages"; +import { + dismissFeedbackPopup, + navigateToPage, + openKeyDetail, +} from "../../helpers/navigation"; +import { readBack } from "../../helpers/roundTrip"; +import { + createVirtualKey, + deleteVirtualKey, + uniqueSuffix, +} from "../../helpers/traffic"; + +const SHARED_MAILBOX = "shared-mailbox@example.com"; +const DISTRIBUTION_LIST = "finance-dl@example.com"; +const RECIPIENTS_PLACEHOLDER = + "Additional recipients, e.g. finance@example.com"; + +async function keyMetadata( + page: PlaywrightPage, + token: string, +): Promise> { + const body = await readBack<{ info: { metadata: Record } }>( + page, + `/key/info?key=${encodeURIComponent(token)}`, + ); + return body.info.metadata; +} + +async function tokenForAlias( + page: PlaywrightPage, + alias: string, +): Promise { + const body = await readBack<{ keys: { key_alias: string; token: string }[] }>( + page, + `/key/list?key_alias=${encodeURIComponent(alias)}&return_full_object=true&size=100`, + ); + return body.keys.find((row) => row.key_alias === alias)?.token; +} + +async function addAlertRow( + page: PlaywrightPage, + threshold: string, + emails: string, +): Promise { + await page + .getByRole("button", { name: "Add Budget Alert Threshold" }) + .click(); + await page.getByPlaceholder("% of budget").last().fill(threshold); + await page.getByPlaceholder(RECIPIENTS_PLACEHOLDER).last().fill(emails); +} + +test.describe("Proxy Admin - Key budget alert recipients", () => { + test.use({ storageState: ADMIN_STORAGE_PATH }); + + test("Create Key stores the alert recipients next to the metadata JSON", async ({ + page, + }) => { + const alias = `e2e-alert-create-${uniqueSuffix()}`; + await navigateToPage(page, Page.ApiKeys); + await dismissFeedbackPopup(page); + await page.getByRole("button", { name: /Create New Key/i }).click(); + await page.getByLabel(/Key Name/).fill(alias); + await page.getByTestId("team-dropdown").getByRole("combobox").click(); + await page.keyboard.type(E2E_TEAM_CRUD_ALIAS); + await page + .getByRole("option", { name: E2E_TEAM_CRUD_ALIAS }) + .first() + .click(); + await page.getByRole("combobox", { name: "Select models" }).click(); + await page + .getByRole("option", { name: "All Team Models", exact: true }) + .click(); + await page.keyboard.press("Escape"); + + await page.getByRole("button", { name: /Optional Settings/i }).click(); + await page.getByLabel(/Max Budget \(USD\)/).fill("10"); + await addAlertRow(page, "80", `${SHARED_MAILBOX}, ${DISTRIBUTION_LIST}`); + await page.getByLabel("Metadata").fill('{"cost_center": "ops"}'); + await page.getByRole("button", { name: "Create Key", exact: true }).click(); + await expect(page.getByText("Save your Key")).toBeVisible({ + timeout: 10_000, + }); + + const token = await tokenForAlias(page, alias); + expect(token, `key ${alias} readable from /key/list`).toBeTruthy(); + try { + expect(await keyMetadata(page, token!)).toMatchObject({ + cost_center: "ops", + max_budget_alert_emails: { "80": [SHARED_MAILBOX, DISTRIBUTION_LIST] }, + }); + } finally { + await deleteVirtualKey(page.request, token!); + } + }); + + test("Edit Key round-trips the rows, keeps the rest of the metadata, and clears them when every row is removed", async ({ + page, + }) => { + const alias = `e2e-alert-edit-${uniqueSuffix()}`; + const created = await createVirtualKey(page.request, { + key_alias: alias, + team_id: E2E_TEAM_CRUD_ID, + max_budget: 10, + metadata: { + cost_center: "ops", + max_budget_alert_emails: { "80": [SHARED_MAILBOX] }, + }, + }); + try { + await navigateToPage(page, Page.ApiKeys); + await dismissFeedbackPopup(page); + await openKeyDetail(page, alias); + await page.getByRole("tab", { name: "Settings" }).click(); + await page.getByRole("button", { name: "Edit Settings" }).click(); + + await expect(page.getByPlaceholder("% of budget")).toHaveValue("80"); + await expect(page.getByPlaceholder(RECIPIENTS_PLACEHOLDER)).toHaveValue( + SHARED_MAILBOX, + ); + expect( + JSON.parse(await page.getByLabel("Metadata").inputValue()), + ).toEqual({ cost_center: "ops" }); + + await page + .getByRole("button", { name: "Remove budget alert threshold" }) + .click(); + await addAlertRow(page, "50", DISTRIBUTION_LIST); + await page.getByRole("button", { name: "Save Changes" }).click(); + await expect + .poll(() => keyMetadata(page, created.token), { timeout: 20_000 }) + .toMatchObject({ + cost_center: "ops", + max_budget_alert_emails: { "50": [DISTRIBUTION_LIST] }, + }); + + await page.getByRole("button", { name: "Edit Settings" }).click(); + await expect(page.getByPlaceholder("% of budget")).toHaveValue("50"); + await page + .getByRole("button", { name: "Remove budget alert threshold" }) + .click(); + await page.getByRole("button", { name: "Save Changes" }).click(); + await expect + .poll( + async () => + "max_budget_alert_emails" in + (await keyMetadata(page, created.token)), + { timeout: 20_000 }, + ) + .toBe(false); + expect(await keyMetadata(page, created.token)).toMatchObject({ + cost_center: "ops", + }); + } finally { + await deleteVirtualKey(page.request, created.token); + } + }); +}); diff --git a/tests/integration/database/image/test_image_admin_mcp.py b/tests/integration/database/image/test_image_admin_mcp.py index aab22f354e4..a7fd33f1417 100644 --- a/tests/integration/database/image/test_image_admin_mcp.py +++ b/tests/integration/database/image/test_image_admin_mcp.py @@ -186,7 +186,7 @@ def test_image_admin_mcp_personal_admin_manages_team(offline_postgres: tuple[str "--env", "HOME=/opt/prisma", "--env", "XDG_CACHE_HOME=/opt/prisma/.cache", "--env", "PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries", "--entrypoint", "prisma", SCHEMA_IMAGE, - "db", "push", "--schema", "/app/schema.prisma", "--skip-generate", "--accept-data-loss", + "db", "push", "--schema", "/app/litellm/schema.prisma", "--skip-generate", "--accept-data-loss", ], capture_output=True, text=True, timeout=180, check=False, ) diff --git a/tests/integration/spend/test_key_budget_alerts.py b/tests/integration/spend/test_key_budget_alerts.py new file mode 100644 index 00000000000..a904d2ba03a --- /dev/null +++ b/tests/integration/spend/test_key_budget_alerts.py @@ -0,0 +1,74 @@ +import uuid +from pathlib import Path +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.mail import smtp_sink +from integration._support.process import owned_proxy + +pytestmark: Final = pytest.mark.shares_worker_state + +KEY_BUDGET: Final = 0.10 +CALL_COST: Final = 20 * 0.001 + 20 * 0.002 + + +def _key_spend(key: str) -> float: + rows: Final = read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE key_alias = %s', (key,)) + return float(str(rows[0]["spend"])) if rows else 0.0 + + +def test_key_threshold_emails_owner_and_shared_mailbox_from_admin_ui_metadata(gateway: Gateway, tmp_path: Path) -> None: + owner_email: Final = f"owner-{uuid.uuid4().hex}@integration.test" + shared_mailbox: Final = f"shared-mailbox-{uuid.uuid4().hex}@integration.test" + alias: Final = f"alert-key-{uuid.uuid4().hex}" + configuration: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + configuration["general_settings"]["alerting"] = ["email"] + configuration["general_settings"]["user_api_key_cache_ttl"] = 1 + path: Final = tmp_path / "email-alerting.yaml" + path.write_text(yaml.safe_dump(configuration)) + with smtp_sink() as mailbox: + overrides: Final = { + "SMTP_HOST": mailbox.host, + "SMTP_PORT": str(mailbox.port), + "SMTP_TLS": "False", + "SMTP_SENDER_EMAIL": "alerts@integration.test", + } + with owned_proxy(gateway, tmp_path, overrides, config=path) as candidate, candidate.scenario() as scenario: + model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002) + user_id: Final = scenario.user(user_email=owner_email) + key: Final = scenario.key( + key_alias=alias, + user_id=user_id, + models=[model], + max_budget=KEY_BUDGET, + metadata={"cost_center": "ops", "max_budget_alert_emails": {"50": [shared_mailbox]}}, + ) + + first: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "first call"}]}, + key=key, + ) + assert first.status_code == 200, first.text + assert float(first.headers["x-litellm-response-cost"]) == pytest.approx(CALL_COST) + eventually(lambda: _key_spend(alias), lambda spend: spend == pytest.approx(CALL_COST), seconds=70) + assert mailbox.deliveries() == (), "no threshold is reached before the first call is recorded" + + second: Final = candidate.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": "second call"}]}, + key=key, + ) + assert second.status_code == 200, second.text + halfway: Final = eventually(mailbox.deliveries, lambda found: len(found) >= 2, seconds=30) + assert {recipient for delivery in halfway for recipient in delivery.recipients} == { + owner_email, + shared_mailbox, + }, halfway + assert all("50%" in delivery.subject for delivery in halfway), halfway + assert len(halfway) == 2, halfway 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 edf0fc4110d..62e4938f489 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/integrations/langfuse/langfuse_expected_request_body/completion_with_bedrock_call.json b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_bedrock_call.json index 6f359380245..96548edc60f 100644 --- a/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_bedrock_call.json +++ b/tests/unit/integrations/langfuse/langfuse_expected_request_body/completion_with_bedrock_call.json @@ -2,9 +2,6 @@ "name": "litellm-acompletion", "parent_span_id": null, "attributes": { - "langfuse.observation.cost_details": { - "total": 6e-05 - }, "langfuse.observation.input": { "messages": [ { diff --git a/tests/unit/integrations/langfuse/test_langfuse.py b/tests/unit/integrations/langfuse/test_langfuse.py index 7a374c1affe..e7fcad5c748 100644 --- a/tests/unit/integrations/langfuse/test_langfuse.py +++ b/tests/unit/integrations/langfuse/test_langfuse.py @@ -97,7 +97,9 @@ def _comparable(span: Mapping[str, object]) -> dict[str, object]: } -def _expected_generation(route: respx.Route, fixture_name: str, trace_id: str) -> dict[str, object]: +def _expected_generation( + route: respx.Route, fixture_name: str, trace_id: str, cost_total: float | None = None +) -> dict[str, object]: otel_trace_id: Final = resolve_trace_id(trace_id) generations: Final = tuple( span @@ -108,7 +110,14 @@ def _expected_generation(route: respx.Route, fixture_name: str, trace_id: str) - ) assert len(generations) == 1 expected_path: Final = Path(__file__).with_name("langfuse_expected_request_body") / fixture_name - expected: Final = TypeAdapter(dict[str, object]).validate_python(json.loads(expected_path.read_text())) + fixture: Final = TypeAdapter(dict[str, object]).validate_python(json.loads(expected_path.read_text())) + fixture_attributes: Final = TypeAdapter(dict[str, object]).validate_python(fixture["attributes"]) + priced_attributes: Final = ( + fixture_attributes + if cost_total is None + else {**fixture_attributes, "langfuse.observation.cost_details": {"total": cost_total}} + ) + expected: Final = {**fixture, "attributes": priced_attributes} actual: Final = _comparable(generations[0]) expected_comparable: Final = _comparable(expected) assert actual == expected_comparable, json.dumps( @@ -355,6 +364,12 @@ async def test_langfuse_logging_completion_with_bedrock_llm_response( langfuse_export: tuple[respx.Route, str], ) -> None: route, trace_id = langfuse_export + profile: Final = litellm.get_model_info( + model="us.anthropic.claude-haiku-4-5-20251001-v1:0", custom_llm_provider="bedrock" + ) + input_rate: Final = profile["input_cost_per_token"] + output_rate: Final = profile["output_cost_per_token"] + assert isinstance(input_rate, float) and isinstance(output_rate, float) response: Final = litellm.ModelResponse( choices=[], usage=litellm.Usage(prompt_tokens=10, completion_tokens=10, total_tokens=20), @@ -373,7 +388,9 @@ async def test_langfuse_logging_completion_with_bedrock_llm_response( ) await _flush_langfuse() assert route.call_count == 1 - _expected_generation(route, "completion_with_bedrock_call.json", trace_id) + _expected_generation( + route, "completion_with_bedrock_call.json", trace_id, cost_total=10 * input_rate + 10 * output_rate + ) @pytest.mark.asyncio diff --git a/tests/unit/llms/base_llm/embedding/test_provider_embedding_translation.py b/tests/unit/llms/base_llm/embedding/test_provider_embedding_translation.py index 24157263daa..316fdad5493 100644 --- a/tests/unit/llms/base_llm/embedding/test_provider_embedding_translation.py +++ b/tests/unit/llms/base_llm/embedding/test_provider_embedding_translation.py @@ -106,3 +106,4 @@ def test_image_embedding(case: _Case, respx_mock: MockRouter) -> None: else: assert body["inputImage"] == png_b64 assert response.data[0]["embedding"] == [0.11, 0.22, 0.33] + diff --git a/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py index 904ca3bb8ca..63526e917a6 100644 --- a/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py +++ b/tests/unit/llms/bedrock/embed/test_bedrock_embedding.py @@ -190,6 +190,34 @@ def test_bedrock_cohere_embedding_returns_the_scripted_vector(respx_mock: respx. assert response.data == [{"object": "embedding", "index": 0, "embedding": [0.3, 0.2, 0.1]}] +def test_bedrock_cohere_embedding_sends_max_tokens_and_priority(respx_mock: respx.MockRouter) -> None: + route = respx_mock.post("https://bedrock-runtime.us-east-1.amazonaws.com/model/cohere.embed-v4:0/invoke").mock( + return_value=httpx.Response( + 200, + json={"embeddings": {"float": [[0.3, 0.2, 0.1]]}, "response_type": "embeddings_by_type"}, + ) + ) + + litellm.embedding( + model="bedrock/cohere.embed-v4:0", + input=["hello"], + api_key="test-bearer-token", + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint="https://bedrock-runtime.us-east-1.amazonaws.com", + max_tokens=16, + priority=5, + truncate="END", + ) + + assert json.loads(route.calls.last.request.read()) == { + "input_type": "search_document", + "texts": ["hello"], + "max_tokens": 16, + "priority": 5, + "truncate": "END", + } + + GOVCLOUD_WEST_RUNTIME: Final = "https://bedrock-runtime.us-gov-west-1.amazonaws.com" @@ -2130,3 +2158,122 @@ def test_bedrock_embedding_region_bug_reproduction(): img_base_64 = "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAGQAAABkBAMAAACCzIhnAAAAG1BMVEURAAD///+ln5/h39/Dv79qX18uHx+If39MPz9oMSdmAAAACXBIWXMAAA7EAAAOxAGVKw4bAAABB0lEQVRYhe2SzWrEIBCAh2A0jxEs4j6GLDS9hqWmV5Flt0cJS+lRwv742DXpEjY1kOZW6HwHFZnPmVEBEARBEARB/jd0KYA/bcUYbPrRLh6amXHJ/K+ypMoyUaGthILzw0l+xI0jsO7ZcmCcm4ILd+QuVYgpHOmDmz6jBeJImdcUCmeBqQpuqRIbVmQsLCrAalrGpfoEqEogqbLTWuXCPCo+Ki1XGqgQ+jVVuhB8bOaHkvmYuzm/b0KYLWwoK58oFqi6XfxQ4Uz7d6WeKpna6ytUs5e8betMcqAv5YPC5EZB2Lm9FIn0/VP6R58+/GEY1X1egVoZ/3bt/EqF6malgSAIgiDIH+QL41409QMY0LMAAAAASUVORK5CYII=" + + +@pytest.mark.parametrize( + ("model_id", "embed_response", "extra_body", "expected_body"), + [ + ( + "amazon.titan-embed-text-v2:0", + titan_embedding_response, + {"normalize": False, "dimensions": 256}, + {"inputText": test_input, "normalize": False, "dimensions": 256}, + ), + ( + "cohere.embed-english-v3", + cohere_embedding_response, + {"truncate": "END", "embedding_types": ["float"]}, + {"texts": [test_input], "truncate": "END", "embedding_types": ["float"]}, + ), + ], +) +def test_bedrock_embedding_sends_extra_body_fields_in_the_request_body( + respx_mock: respx.MockRouter, + model_id: str, + embed_response: dict[str, object], + extra_body: dict[str, object], + expected_body: dict[str, object], +) -> None: + route = respx_mock.post(f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{model_id}/invoke").mock( + return_value=httpx.Response(200, json=embed_response) + ) + + litellm.embedding( + model=f"bedrock/{model_id}", + input=[test_input], + api_key="test-bearer-token", + aws_region_name="us-east-1", + extra_body=extra_body, + ) + + body = json.loads(route.calls.last.request.read()) + assert "extra_body" not in body + assert expected_body.items() <= body.items() + + +def test_bedrock_nova_embedding_sends_extra_body_fields_in_the_request_body(respx_mock: respx.MockRouter) -> None: + model_id: Final = "amazon.nova-2-multimodal-embeddings-v1:0" + route = respx_mock.post(f"https://bedrock-runtime.us-east-1.amazonaws.com/model/{model_id}/invoke").mock( + return_value=httpx.Response(200, json={"embeddings": [{"embeddingType": "TEXT", "embedding": [0.1, 0.2]}]}) + ) + + litellm.embedding( + model=f"bedrock/{model_id}", + input=[test_input], + api_key="test-bearer-token", + aws_region_name="us-east-1", + extra_body={"segmentationConfig": {"durationSeconds": 10}}, + ) + + body = json.loads(route.calls.last.request.read()) + assert "extra_body" not in body + assert body["segmentationConfig"] == {"durationSeconds": 10} + + +@pytest.mark.parametrize( + ("extra_body", "field", "expected"), + [ + ({"dimensions": 512}, "dimensions", 256), + ({"inputText": "text the guardrails never saw"}, "inputText", test_input), + ], + ids=["mapped_param", "input"], +) +def test_bedrock_embedding_extra_body_cannot_replace_the_input_or_a_mapped_param( + respx_mock: respx.MockRouter, extra_body: dict[str, object], field: str, expected: object +) -> None: + route = respx_mock.post( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-embed-text-v2:0/invoke" + ).mock(return_value=httpx.Response(200, json=titan_embedding_response)) + + litellm.embedding( + model="bedrock/amazon.titan-embed-text-v2:0", + input=[test_input], + api_key="test-bearer-token", + aws_region_name="us-east-1", + dimensions=256, + extra_body=extra_body, + ) + + assert json.loads(route.calls.last.request.read())[field] == expected + + +def test_bedrock_embedding_signs_the_body_it_sends(respx_mock: respx.MockRouter) -> None: + from botocore.auth import SigV4Auth + from botocore.awsrequest import AWSRequest + from botocore.credentials import Credentials + + route = respx_mock.post( + "https://bedrock-runtime.us-east-1.amazonaws.com/model/amazon.titan-embed-text-v2:0/invoke" + ).mock(return_value=httpx.Response(200, json=titan_embedding_response)) + + litellm.embedding( + model="bedrock/amazon.titan-embed-text-v2:0", + input=[test_input], + aws_access_key_id="AKIAIOSFODNN7EXAMPLE", + aws_secret_access_key="wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY", + aws_region_name="us-east-1", + dimensions=256, + extra_body={"dimensions": 512}, + ) + + sent: Final = route.calls.last.request + replay: Final = AWSRequest( + method="POST", url=str(sent.url), data=sent.content, headers={"Content-Type": "application/json"} + ) + replay.context["timestamp"] = sent.headers["x-amz-date"] + SigV4Auth( + Credentials("AKIAIOSFODNN7EXAMPLE", "wJalrXUtnFEMI/K7MDENG/bPxRfiCYEXAMPLEKEY"), "bedrock", "us-east-1" + ).add_auth(replay) + + assert json.loads(sent.content)["dimensions"] == 256 + assert replay.headers["Authorization"] == sent.headers["authorization"] diff --git a/tests/unit/llms/nvidia_nim/test_nvidia_nim.py b/tests/unit/llms/nvidia_nim/test_nvidia_nim.py index 3633ac97265..04bd26d53ef 100644 --- a/tests/unit/llms/nvidia_nim/test_nvidia_nim.py +++ b/tests/unit/llms/nvidia_nim/test_nvidia_nim.py @@ -237,3 +237,28 @@ def test_completion_nvidia_nim(): assert request_body["model"] == "databricks/dbrx-instruct" assert request_body["frequency_penalty"] == 0.1 assert request_body["presence_penalty"] == 0.5 + + +def test_embedding_nvidia_nim_sends_extra_body_fields_without_an_extra_body_key(): + from openai import OpenAI + + transport: Final = CapturingTransport( + CreateEmbeddingResponse( + object="list", + data=(Embedding(object="embedding", index=0, embedding=(0.1, 0.2, 0.3)),), + model="nvidia/nv-embedqa-e5-v5", + usage=EmbeddingUsage(prompt_tokens=6, total_tokens=6), + ) + ) + client: Final = OpenAI(api_key="fake-api-key", http_client=httpx.Client(transport=transport)) + + litellm.embedding( + model="nvidia_nim/nvidia/nv-embedqa-e5-v5", + input="What is the meaning of life?", + extra_body={"input_type": "query", "truncate": "END"}, + client=client, + ) + + request_body: Final = transport.request_bodies[0] + assert "extra_body" not in request_body + assert (request_body["input_type"], request_body["truncate"]) == ("query", "END") diff --git a/tests/unit/llms/ollama/test_ollama_embedding.py b/tests/unit/llms/ollama/test_ollama_embedding.py index b32bb9c4415..abbe449bf4a 100644 --- a/tests/unit/llms/ollama/test_ollama_embedding.py +++ b/tests/unit/llms/ollama/test_ollama_embedding.py @@ -348,3 +348,21 @@ async def test_sdk_embedding_rejects_non_string_batch(embedding_route: respx.Rou litellm.embedding(model="ollama/nomic-embed-text", input=[123], api_base="http://ollama.test") assert not embedding_route.called + + +def test_ollama_embeddings_sends_extra_body_fields_at_the_documented_places(embedding_route: respx.Route) -> None: + response: Final = litellm.embedding( + model="ollama/nomic-embed-text", + input=["hello", "world"], + api_base="http://ollama.test", + extra_body={"truncate": False, "keep_alive": "5m", "options": {"num_ctx": 2048}}, + ) + + assert json.loads(embedding_route.calls.last.request.content) == { + "model": "nomic-embed-text", + "input": ["hello", "world"], + "truncate": False, + "keep_alive": "5m", + "options": {"num_ctx": 2048}, + } + assert response.data[0]["embedding"] == [0.1, 0.2] diff --git a/tests/unit/llms/sagemaker/test_sagemaker_embedding_voyage.py b/tests/unit/llms/sagemaker/test_sagemaker_embedding_voyage.py index 3951b17db92..3d8c5ccfa21 100644 --- a/tests/unit/llms/sagemaker/test_sagemaker_embedding_voyage.py +++ b/tests/unit/llms/sagemaker/test_sagemaker_embedding_voyage.py @@ -7,6 +7,7 @@ transformation, and model type detection. """ import json +from typing import Final from unittest.mock import MagicMock, patch import httpx @@ -105,6 +106,21 @@ class TestSagemakerCohereEmbeddingConfig: ) assert result["texts"] == ["hello"] + def test_transform_request_keeps_max_tokens_and_priority(self): + result: Final = self.config.transform_embedding_request( + model="cohere-embed-v4", + input=["hello"], + optional_params={"max_tokens": 16, "priority": 5, "truncate": "END"}, + headers={}, + ) + assert result == { + "texts": ["hello"], + "input_type": "search_document", + "max_tokens": 16, + "priority": 5, + "truncate": "END", + } + def test_map_openai_params_dimensions_to_output_dimension(self): params = self.config.map_openai_params( non_default_params={"dimensions": 512, "encoding_format": "float"}, @@ -626,3 +642,40 @@ class TestTEIEmbeddingResponse: if __name__ == "__main__": pytest.main([__file__]) + + +def test_voyage_endpoint_sends_extra_body_fields_in_the_invoke_body(): + client = MagicMock() + client.invoke_endpoint.return_value = { + "Body": MagicMock( + read=MagicMock( + return_value=json.dumps( + { + "data": [{"object": "embedding", "embedding": [0.1, 0.2], "index": 0}], + "usage": {"total_tokens": 3}, + } + ).encode() + ) + ) + } + session = MagicMock() + session.client.return_value = client + + with patch("boto3.Session", return_value=session): + response = embedding( + model="sagemaker/voyage-3-5-endpoint", + input=["hello"], + aws_region_name="us-east-1", + aws_access_key_id="test-access-key", + aws_secret_access_key="test-secret-key", + extra_body={"input_type": "query", "truncation": False, "output_dimension": 512}, + ) + + assert json.loads(client.invoke_endpoint.call_args.kwargs["Body"]) == { + "input": ["hello"], + "model": "voyage-3-5-endpoint", + "input_type": "query", + "truncation": False, + "output_dimension": 512, + } + assert response.data[0]["embedding"] == [0.1, 0.2] diff --git a/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py b/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py index d7d921ab35c..06306981daf 100644 --- a/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py +++ b/tests/unit/llms/vertex_ai/gemini_embeddings/test_batch_embed_content_transformation.py @@ -652,3 +652,32 @@ def test_gemini_embedding(respx_mock: respx.MockRouter) -> None: expected_prompt_tokens: Final = litellm.token_counter(model="gemini/gemini-embedding-001", text="Hello, world!") assert response.usage.prompt_tokens == expected_prompt_tokens assert response.usage.total_tokens == expected_prompt_tokens + + +def test_gemini_embedding_sends_extra_body_fields_on_each_request(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.post( + "https://generativelanguage.googleapis.com/v1beta/models/gemini-embedding-001:batchEmbedContents" + ).mock(return_value=httpx.Response(200, json={"embeddings": [{"values": [0.1, 0.2]}]})) + + litellm.embedding( + model="gemini/gemini-embedding-001", + input="Hello, world!", + api_key="gemini-test-key", + extra_body={ + "title": "Greeting", + "taskType": "RETRIEVAL_DOCUMENT", + "outputDimensionality": 256, + "documentOcr": True, + }, + ) + + assert json.loads(route.calls.last.request.content)["requests"] == [ + { + "model": "models/gemini-embedding-001", + "content": {"parts": [{"text": "Hello, world!"}]}, + "title": "Greeting", + "taskType": "RETRIEVAL_DOCUMENT", + "outputDimensionality": 256, + "documentOcr": True, + } + ] diff --git a/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py b/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py index edb6e889814..123d4128679 100644 --- a/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py +++ b/tests/unit/llms/vertex_ai/multimodal_embeddings/test_vertex_ai_multimodal_embedding_transformation.py @@ -1,8 +1,13 @@ import json from unittest.mock import AsyncMock, MagicMock, patch +import google.auth +import google.auth.credentials import httpx import pytest +import respx + +import litellm from litellm.llms.vertex_ai.multimodal_embeddings.transformation import ( @@ -174,3 +179,30 @@ class TestVertexMultimodalEmbedding: headers={}, ) assert "parameters" not in request + + +class _StaticGoogleCredentials(google.auth.credentials.Credentials): + def refresh(self, request: object) -> None: + self.token = "test-token" + + +def test_extra_body_fields_reach_the_multimodal_predict_body( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(google.auth, "default", lambda scopes: (_StaticGoogleCredentials(), "test-project")) + route = respx_mock.post(host="us-central1-aiplatform.googleapis.com").mock( + return_value=httpx.Response(200, json={"predictions": [{"textEmbedding": [0.1, 0.2]}]}) + ) + + litellm.embedding( + model="vertex_ai/multimodalembedding@001", + input="Hello world", + vertex_project="test-project", + vertex_location="us-central1", + extra_body={"dimension": 128}, + ) + + assert json.loads(route.calls.last.request.content) == { + "instances": [{"text": "Hello world"}], + "parameters": {"dimension": 128}, + } diff --git a/tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py b/tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py index 1058d4fb82d..0e8d49b01e5 100644 --- a/tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py +++ b/tests/unit/llms/vertex_ai/vertex_embeddings/test_embedding_handler.py @@ -305,3 +305,20 @@ async def test_vertex_embedding_usage_carries_the_billable_character_count_googl 8, character_count, ) + + +def test_vertex_embedding_call_sends_extra_body_fields_as_predict_parameters(predict_endpoint: respx.Route) -> None: + predict_endpoint.mock(return_value=_predict_response({"predictions": [_gecko_prediction([0.1, 2], 3)]})) + + litellm.embedding( + model=f"vertex_ai/{GECKO}", + input=["hello"], + vertex_project=PUBLIC_CALL_PROJECT, + vertex_location="us-central1", + extra_body={"autoTruncate": False}, + ) + + body: Final = json.loads(predict_endpoint.calls.last.request.content) + assert body["instances"] == [{"content": "hello"}] + assert "extra_body" not in body["parameters"] + assert body["parameters"]["autoTruncate"] is False 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 3932fae0737..93026c3ea17 100644 --- a/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py +++ b/tests/unit/proxy/hooks/test_dynamic_rate_limiter_v3.py @@ -3172,6 +3172,62 @@ def _held_to_borrow_cap(exc: ProxyRateLimitError, rate_limit_type: str, limit_se ), exc.detail +async def _cap_tokens_remaining(handler: _PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth) -> int: + share: Final = handler._borrowing_team_share("shared-model", auth) + assert share is not None and share.ceiling is not None + response: Final = await handler.v3_limiter.should_rate_limit( + [share.ceiling], read_only=True, skip_parallel_check=True + ) + return next(s["limit_remaining"] for s in response["statuses"] if s["rate_limit_type"] == "tokens") + + +@pytest.mark.asyncio +async def test_a_capped_borrowers_idle_usage_counts_toward_its_cap_but_not_toward_its_share( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _check(handler, _capped_team_key(2), saturation=0.0) + await _settle_tokens(handler, 150) + assert await _cap_tokens_remaining(handler, _capped_team_key(2)) == 50 + assert await _team_tokens_remaining(handler, _capped_team_key(2)) == 100 + await _check(handler, _capped_team_key(2), saturation=0.5) + await _settle_tokens(handler, 30) + assert await _cap_tokens_remaining(handler, _capped_team_key(2)) == 20 + assert await _team_tokens_remaining(handler, _capped_team_key(2)) == 70 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("saturation", [0.0, 0.5], ids=["idle", "saturated"]) +async def test_a_capped_borrowers_reservation_is_sized_by_its_share_not_its_cap( + monkeypatch: pytest.MonkeyPatch, saturation: float +): + """A prompt without max_tokens reserves a floor that scales with the limit it is sized against. Sizing it + against the cap would refuse a request its own share would let through, and would admit different + requests depending on whether the model happened to be saturated.""" + handler: Final = _borrowing_handler(monkeypatch, borrows=True) + await _admit_at_the_hook(handler, _capped_team_key(2), tokens=160) + short_prompt_without_max_tokens: Final[dict[str, object]] = {"messages": [{"role": "user", "content": "hi"}]} + await _check(handler, _capped_team_key(2), saturation=saturation, data=short_prompt_without_max_tokens) + assert short_prompt_without_max_tokens["max_tokens"] == handler.v3_limiter.no_max_tokens_output_floor(100) + + +@pytest.mark.asyncio +async def test_a_capped_borrower_is_held_to_its_cap_while_idle_after_a_saturated_stretch( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True) + await _check(handler, _capped_team_key(2), saturation=0.0) + await _settle_tokens(handler, 150) + await _check(handler, _capped_team_key(2), saturation=0.5) + await _settle_tokens(handler, 60) + with pytest.raises(ProxyRateLimitError) as exc: + await _check(handler, _capped_team_key(2), saturation=0.49) + _held_to_borrow_cap(exc.value, "tokens", "model_tpm_limit") + with pytest.raises(ProxyRateLimitError) as saturated: + await _check(handler, _capped_team_key(2), saturation=0.5) + _held_to_borrow_cap(saturated.value, "tokens", "model_tpm_limit") + + @pytest.mark.asyncio async def test_team_borrows_past_its_share_while_it_stays_under_its_max_borrow_multiple( monkeypatch: pytest.MonkeyPatch, @@ -3280,6 +3336,44 @@ async def test_a_websocket_turn_is_refused_at_the_teams_max_borrow_multiple_whil _held_to_borrow_cap(exc.value, "tokens", "model_tpm_limit") +@pytest.mark.asyncio +async def test_a_websocket_admitted_while_saturated_gives_its_share_hold_back_when_the_model_goes_idle( + monkeypatch: pytest.MonkeyPatch, +): + handler: Final = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.5) + await _check(handler, _capped_team_key(2), saturation=0.5, data=_prompt(20)) + assert await _team_tokens_remaining(handler, _capped_team_key(2)) < 100 + await handler.enforce_token_limits_for_frame(user_api_key_dict=_capped_team_key(2), requested_model="shared-model") + await _settle_tokens( + handler, + 90, + response=[], + combined_usage_object=litellm.Usage(prompt_tokens=60, completion_tokens=30, total_tokens=90), + ) + assert await _team_tokens_remaining(handler, _capped_team_key(2)) == 10 + assert await _cap_tokens_remaining(handler, _capped_team_key(2)) == 110 + + +@pytest.mark.asyncio +async def test_a_capped_websocket_admitted_while_idle_settles_its_first_saturated_turn_on_a_fresh_share( + monkeypatch: pytest.MonkeyPatch, +): + """The idle admission held tokens on the cap only, so the share the turn settles to once the model saturates + starts from zero rather than being refunded a hold it never carried.""" + handler: Final = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.1) + await _check(handler, _capped_team_key(2), saturation=0.0, data=_prompt(20)) + await asyncio.create_task(_saturate_model_from_another_request(handler, 200)) + await handler.enforce_token_limits_for_frame(user_api_key_dict=_capped_team_key(2), requested_model="shared-model") + await _settle_tokens( + handler, + 90, + response=[], + combined_usage_object=litellm.Usage(prompt_tokens=60, completion_tokens=30, total_tokens=90), + ) + assert await _team_tokens_remaining(handler, _capped_team_key(2)) == 10 + assert await _cap_tokens_remaining(handler, _capped_team_key(2)) == 110 + + @pytest.mark.asyncio async def test_a_websocket_turn_under_the_teams_max_borrow_multiple_still_borrows(monkeypatch: pytest.MonkeyPatch): handler = _borrowing_handler(monkeypatch, borrows=True, saturation_threshold=0.5) @@ -3476,6 +3570,22 @@ async def _admit_at_the_hook(handler: _PROXY_DynamicRateLimitHandlerV3, auth: Us await _settle_tokens(handler, tokens) +async def _saturate_model_from_another_request(handler: _PROXY_DynamicRateLimitHandlerV3, tokens: int) -> None: + """``_saturate_model`` on a request stash of its own, as a second request in the proxy would have, so the + stash of the request under test keeps its holds.""" + from litellm.proxy.hooks import parallel_request_limiter_v3 + + parallel_request_limiter_v3._request_stash.set(None) # pyright: ignore[reportPrivateUsage] # fresh request context + await _saturate_model(handler, tokens) + + +async def _saturate_model(handler: _PROXY_DynamicRateLimitHandlerV3, tokens: int) -> None: + """Idle usage that pushes the model past its saturation threshold without touching any team's share, from a + priority class of its own so the default pool the measured team sits in stays untouched.""" + warm_up: Final = UserAPIKeyAuth(api_key="sk-warm-up", team_id="team-warm-up", team_metadata={"priority": "batch"}) + await _admit_at_the_hook(handler, warm_up, tokens=tokens) + + async def _refused_at_the_hook( handler: _PROXY_DynamicRateLimitHandlerV3, auth: UserAPIKeyAuth, model: str = "shared-model" ) -> ProxyRateLimitError: @@ -3520,7 +3630,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"}}, @@ -3536,10 +3646,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", 100), ], ids=["saturated", "borrow_cap"], ) @@ -3549,6 +3659,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) @@ -3564,7 +3675,11 @@ 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 @@ -3628,13 +3743,14 @@ 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", team_id="team-a", team_metadata={"model_tpm_limit": {"shared-model": 100}, "priority": "interactive"}, ) + await _saturate_model(handler, 200) await _admit_at_the_hook(handler, team, tokens=60) await _admit_at_the_hook(handler, team, tokens=70) labels: Final = {"model_group": "shared-model", "team": "team-a", "priority": "interactive"} @@ -3651,9 +3767,10 @@ 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}}) + await _saturate_model(handler, 500) assert await handler.async_pre_call_hook(team, DualCache(), {"model": "other-model"}, "completion") is None assert await handler.async_fallback_target_allowed(team, "shared-model", {}) is True await _settle_tokens(handler, 150, model_group="shared-model") @@ -3663,6 +3780,35 @@ async def test_a_fallback_served_request_publishes_the_teams_share_on_the_fallba assert _sample(registry, "litellm_dynamic_limiter_share_tokens", model_group="other-model") is None +@pytest.mark.asyncio +async def test_a_capped_idle_teams_borrowing_is_read_against_its_share_not_the_wider_class_share( + monkeypatch: pytest.MonkeyPatch, +): + handler = _borrowing_handler(monkeypatch, borrows=True, splits=True) + registry = _prometheus_sink(monkeypatch) + await _admit_at_the_hook(handler, _capped_team_key(2), tokens=150) + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **TEAM_LABELS) == 100 + assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **TEAM_LABELS) == 50 + + +@pytest.mark.asyncio +async def test_an_idle_fallback_served_for_a_capped_team_still_publishes_its_borrowing(monkeypatch: pytest.MonkeyPatch): + handler = _borrowing_handler(monkeypatch, borrows=True) + registry = _prometheus_sink(monkeypatch) + team: Final = _team_key( + team_limits={ + "model_tpm_limit": {"other-model": 40, "shared-model": 100}, + "team_model_limit_max_borrow_multiple": 2, + } + ) + assert await handler.async_pre_call_hook(team, DualCache(), {"model": "other-model"}, "completion") is None + assert await handler.async_fallback_target_allowed(team, "shared-model", {}) is True + await _settle_tokens(handler, 150, model_group="shared-model") + shared: Final = {"model_group": "shared-model", "team": "team-a", "priority": "default"} + assert _sample(registry, "litellm_dynamic_limiter_share_tokens", **shared) == 100 + assert _sample(registry, "litellm_dynamic_limiter_borrowed_tokens", **shared) == 50 + + @pytest.mark.asyncio async def test_a_failing_metrics_sink_never_changes_admission(monkeypatch: pytest.MonkeyPatch): from litellm.integrations.prometheus import PrometheusLogger @@ -3674,8 +3820,9 @@ 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 _saturate_model(handler, 200) await _admit_at_the_hook(handler, _team_key(), tokens=60) await _admit_at_the_hook(handler, _team_key(), tokens=70) assert await _team_tokens_remaining(handler) == -30 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, diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index 71efa1352e4..23748790869 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -2,7 +2,7 @@ import json import asyncio import datetime import time -from collections.abc import Iterator, Mapping +from collections.abc import Callable, Iterator, Mapping from pathlib import Path from types import MappingProxyType, SimpleNamespace from typing import Final, Literal, cast @@ -526,6 +526,204 @@ def test_completion_cost_bills_priced_response_model_over_priced_requested_model assert cost == pytest.approx(_token_rates_cost(litellm.model_cost["openai/sol-requested-test-2099-01-01"], usage)) +_BEDROCK_PROFILE_GLOBAL_RATES: Final = MappingProxyType( + { + "input_cost_per_token": 1e-6, + "output_cost_per_token": 5e-6, + "cache_read_input_token_cost": 1e-7, + "cache_creation_input_token_cost": 1.25e-6, + } +) +_BEDROCK_PROFILE_MODEL: Final = "anthropic.bedrock-profile-test-v1:0" +_BEDROCK_PROFILE_RATE_MULTIPLIERS: Final = MappingProxyType( + { + _BEDROCK_PROFILE_MODEL: 1.0, + f"eu.{_BEDROCK_PROFILE_MODEL}": 1.1, + f"us.{_BEDROCK_PROFILE_MODEL}": 1.2, + f"jp.{_BEDROCK_PROFILE_MODEL}": 1.3, + f"global.{_BEDROCK_PROFILE_MODEL}": 0.9, + "anthropic.bedrock-profile-other-test-v1:0": 3.0, + "anthropic.bedrock-profile-unpriced-test-v1:0": 1.0, + "bedrock-profile-router-test": 5.0, + } +) +_BEDROCK_PROFILE_CACHE_USAGE: Final = MappingProxyType( + {"cache_creation_input_tokens": 2000, "cache_read_input_tokens": 24000} +) + + +@pytest.fixture +def _bedrock_profile_rows(_local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch) -> None: + for key, multiplier in _BEDROCK_PROFILE_RATE_MULTIPLIERS.items(): + monkeypatch.setitem( + litellm.model_cost, + key, + { + "litellm_provider": "bedrock_converse", + "mode": "chat", + **{field: rate * multiplier for field, rate in _BEDROCK_PROFILE_GLOBAL_RATES.items()}, + }, + ) + monkeypatch.setitem( + litellm.model_cost, + "eu.anthropic.bedrock-profile-unpriced-test-v1:0", + {"litellm_provider": "bedrock_converse", "mode": "chat"}, + ) + + +def _bedrock_messages_body(model: str) -> dict[str, object]: + return { + "id": "msg_bdrk_profile", + "type": "message", + "role": "assistant", + "model": model, + "content": [{"type": "text", "text": "hi"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 18, "output_tokens": 4, **_BEDROCK_PROFILE_CACHE_USAGE}, + } + + +def _bedrock_stream_response(model: str) -> ModelResponse: + cached: Final = _BEDROCK_PROFILE_CACHE_USAGE["cache_read_input_tokens"] + written: Final = _BEDROCK_PROFILE_CACHE_USAGE["cache_creation_input_tokens"] + return ModelResponse( + model=model, + choices=[], + usage=Usage( + prompt_tokens=18 + cached + written, + completion_tokens=4, + total_tokens=22 + cached + written, + cache_creation_input_tokens=written, + cache_read_input_tokens=cached, + prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=cached), + ), + ) + + +def _bedrock_private_route_response(model: str) -> ModelResponse: + response: Final = _bedrock_stream_response("bedrock-profile-alias") + response._hidden_params = {"custom_llm_provider": "bedrock", "provider_response_model": model} + return response + + +def _bedrock_global_cost(served_shape: Callable[[str], object]) -> float: + return completion_cost( + completion_response=served_shape(_BEDROCK_PROFILE_MODEL), + model=f"bedrock/{_BEDROCK_PROFILE_MODEL}", + custom_llm_provider="bedrock", + ) + + +_BEDROCK_SERVED_SHAPES: Final = ( + pytest.param(_bedrock_messages_body, id="non-streaming"), + pytest.param(_bedrock_stream_response, id="streaming"), + pytest.param(_bedrock_private_route_response, id="private-provider-response-model"), +) + + +@pytest.mark.parametrize("served_shape", _BEDROCK_SERVED_SHAPES) +@pytest.mark.parametrize("route", ["bedrock/", "bedrock/invoke/"]) +@pytest.mark.parametrize("region", ["eu", "us", "jp"]) +def test_completion_cost_bills_bedrock_cross_region_profile_when_bedrock_serves_the_region_stripped_id( + _bedrock_profile_rows: None, region: str, route: str, served_shape: Callable[[str], object] +) -> None: + profile: Final = f"{region}.{_BEDROCK_PROFILE_MODEL}" + + cost: Final = completion_cost( + completion_response=served_shape(_BEDROCK_PROFILE_MODEL), + model=f"{route}{profile}", + custom_llm_provider="bedrock", + ) + + assert cost == pytest.approx(_bedrock_global_cost(served_shape) * _BEDROCK_PROFILE_RATE_MULTIPLIERS[profile]) + + +@pytest.mark.parametrize("served_shape", _BEDROCK_SERVED_SHAPES) +@pytest.mark.parametrize( + ("deployment", "served_model"), + [ + pytest.param(f"bedrock/eu.{_BEDROCK_PROFILE_MODEL}", f"us.{_BEDROCK_PROFILE_MODEL}", id="other-region"), + pytest.param(f"bedrock/eu.{_BEDROCK_PROFILE_MODEL}", f"global.{_BEDROCK_PROFILE_MODEL}", id="global-profile"), + pytest.param( + f"bedrock/jp.{_BEDROCK_PROFILE_MODEL}", "anthropic.bedrock-profile-other-test-v1:0", id="other-model" + ), + pytest.param( + "bedrock/eu.anthropic.bedrock-profile-unpriced-test-v1:0", + "anthropic.bedrock-profile-unpriced-test-v1:0", + id="unpriced-profile", + ), + ], +) +def test_completion_cost_keeps_the_served_bedrock_row_unless_it_is_the_region_stripped_profile( + _bedrock_profile_rows: None, deployment: str, served_model: str, served_shape: Callable[[str], object] +) -> None: + cost: Final = completion_cost( + completion_response=served_shape(served_model), + model=deployment, + custom_llm_provider="bedrock", + ) + + assert cost == pytest.approx(_bedrock_global_cost(served_shape) * _BEDROCK_PROFILE_RATE_MULTIPLIERS[served_model]) + + +@pytest.mark.parametrize( + ("base_model", "custom_pricing", "router_model_id", "priced_row"), + [ + pytest.param( + "bedrock/anthropic.bedrock-profile-other-test-v1:0", + None, + None, + "anthropic.bedrock-profile-other-test-v1:0", + id="base-model", + ), + pytest.param(None, True, "bedrock-profile-router-test", "bedrock-profile-router-test", id="custom-pricing"), + ], +) +def test_completion_cost_explicit_pricing_precedes_the_bedrock_cross_region_profile( + _bedrock_profile_rows: None, + base_model: str | None, + custom_pricing: bool | None, + router_model_id: str | None, + priced_row: str, +) -> None: + cost: Final = completion_cost( + completion_response=_bedrock_messages_body(_BEDROCK_PROFILE_MODEL), + model=f"bedrock/eu.{_BEDROCK_PROFILE_MODEL}", + custom_llm_provider="bedrock", + base_model=base_model, + custom_pricing=custom_pricing, + router_model_id=router_model_id, + ) + + assert cost == pytest.approx( + _bedrock_global_cost(_bedrock_messages_body) * _BEDROCK_PROFILE_RATE_MULTIPLIERS[priced_row] + ) + + +def test_completion_cost_keeps_the_served_row_for_a_region_dotted_name_outside_bedrock( + _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "openai/eu.profile-dotted-test", + {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, + ) + monkeypatch.setitem( + litellm.model_cost, + "openai/profile-dotted-test", + {"litellm_provider": "openai", "mode": "chat", "input_cost_per_token": 5e-6, "output_cost_per_token": 11e-6}, + ) + usage: Final = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) + + cost: Final = completion_cost( + completion_response=ModelResponse(model="profile-dotted-test", choices=[], usage=usage), + model="openai/eu.profile-dotted-test", + custom_llm_provider="openai", + ) + + assert cost == pytest.approx(_token_rates_cost(litellm.model_cost["openai/profile-dotted-test"], usage)) + + def test_completion_cost_keeps_response_pricing_when_requested_alias_has_no_rates( _local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch ) -> None: diff --git a/uv.lock b/uv.lock index 6dc41480bf5..9c23d77ea3b 100644 --- a/uv.lock +++ b/uv.lock @@ -5032,7 +5032,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.110.0" +version = "1.111.0" source = { editable = "." } dependencies = [ { name = "aiohttp" }, @@ -5571,7 +5571,7 @@ provides-extras = ["test"] [[package]] name = "litellm-enterprise" -version = "0.1.83" +version = "0.1.84" source = { editable = "litellm/enterprise" } [[package]]