diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0e0f1558bb5..c8ca9dcf57e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1,7 +1,7 @@ import enum import json import os -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime from typing import TYPE_CHECKING, Any, Final, Literal @@ -11,6 +11,7 @@ from pydantic import ( ConfigDict, Field, Json, + PositiveInt, field_validator, model_validator, ) @@ -1102,6 +1103,8 @@ class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase): class KeyRequestBase(GenerateRequestBase): key: str | None = None + default_estimated_output_tokens: PositiveInt | None = None + default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None budget_id: str | None = None tags: list[str] | None = None disable_global_guardrails: bool | None = None @@ -1819,6 +1822,8 @@ class NewTeamRequest(TeamBase): ) model_tpm_limit: dict[str, int] | None = None + default_estimated_output_tokens: PositiveInt | None = None + default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None mcp_rpm_limit: dict[str, int] | None = None team_member_budget: float | None = None # allow user to set a budget for all team members team_member_rpm_limit: int | None = None # allow user to set RPM limit for all team members @@ -1883,6 +1888,8 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): prompts: list[str] | None = None model_rpm_limit: dict[str, int] | None = None model_tpm_limit: dict[str, int] | None = None + default_estimated_output_tokens: PositiveInt | None = None + default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None mcp_rpm_limit: dict[str, int] | None = None allowed_vector_store_indexes: list[AllowedVectorStoreIndexItem] | None = None enforced_batch_output_expires_after: dict | None = None @@ -4103,6 +4110,8 @@ class PassThroughEndpointLoggingTypedDict(TypedDict): LiteLLM_ManagementEndpoint_MetadataFields: Final = [ "model_rpm_limit", "model_tpm_limit", + "default_estimated_output_tokens", + "default_estimated_output_tokens_per_model", "mcp_rpm_limit", "tag_rpm_limit", "rpm_limit_type", diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index b4e634b5eb1..0bfe4b685e1 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -1,12 +1,13 @@ import os import re import sys -from collections.abc import Iterator, Mapping +from collections.abc import Collection, Iterator, Mapping from functools import lru_cache from logging import Logger -from typing import Any, Final +from typing import Any, Final, Protocol from fastapi import HTTPException, Request, status +from pydantic import PositiveInt, TypeAdapter, ValidationError import litellm from litellm import Router, provider_list @@ -999,6 +1000,167 @@ def get_key_model_tpm_limit( return None +ESTIMATED_OUTPUT_TOKENS_FIELD: Final = "default_estimated_output_tokens" +ESTIMATED_OUTPUT_TOKENS_PER_MODEL_FIELD: Final = "default_estimated_output_tokens_per_model" +ESTIMATED_OUTPUT_TOKENS_METADATA_FIELDS: Final = frozenset( + {ESTIMATED_OUTPUT_TOKENS_FIELD, ESTIMATED_OUTPUT_TOKENS_PER_MODEL_FIELD} +) + +_ESTIMATED_OUTPUT_TOKENS_ADAPTER: Final = TypeAdapter(PositiveInt) +_ESTIMATED_OUTPUT_TOKENS_PER_MODEL_ADAPTER: Final = TypeAdapter(Mapping[str, PositiveInt]) + + +def _validated_output_token_estimate(raw: object) -> int | None: + """Coerce one declared estimate to a positive int, or ignore it.""" + if raw is None: + return None + try: + return _ESTIMATED_OUTPUT_TOKENS_ADAPTER.validate_python(raw) + except ValidationError as validation_error: + verbose_proxy_logger.warning( + "Ignoring malformed %s in metadata: %s", + ESTIMATED_OUTPUT_TOKENS_FIELD, + validation_error, + ) + return None + + +def _validated_output_token_estimates_per_model(raw: object) -> Mapping[str, int] | None: + """Coerce a declared per-model estimate map, or ignore it.""" + if raw is None: + return None + try: + return _ESTIMATED_OUTPUT_TOKENS_PER_MODEL_ADAPTER.validate_python(raw) + except ValidationError as validation_error: + verbose_proxy_logger.warning( + "Ignoring malformed %s in metadata: %s", + ESTIMATED_OUTPUT_TOKENS_PER_MODEL_FIELD, + validation_error, + ) + return None + + +def _estimated_output_tokens_from_metadata( + metadata: Mapping[str, Any] | None, + model_name: str | None, +) -> int | None: + """Resolve the per-model, then global, estimate out of one metadata blob. + + The two fields are validated independently so a malformed per-model map + cannot discard a valid global estimate, or the other way round. + """ + if not metadata or ESTIMATED_OUTPUT_TOKENS_METADATA_FIELDS.isdisjoint(metadata): + return None + + if model_name is not None: + per_model: Final = _validated_output_token_estimates_per_model( + metadata.get(ESTIMATED_OUTPUT_TOKENS_PER_MODEL_FIELD) + ) + per_model_estimate: Final = per_model.get(model_name) if per_model is not None else None + if per_model_estimate is not None: + return per_model_estimate + + return _validated_output_token_estimate(metadata.get(ESTIMATED_OUTPUT_TOKENS_FIELD)) + + +def get_estimated_output_tokens( + user_api_key_dict: UserAPIKeyAuth, + model_name: str | None = None, +) -> int | None: + """Resolve the operator-declared output-token estimate for TPM reservation. + + Priority order (returns first found): + 1. Key metadata ``default_estimated_output_tokens_per_model[model_name]`` + 2. Key metadata ``default_estimated_output_tokens`` + 3. Team metadata ``default_estimated_output_tokens_per_model[model_name]`` + 4. Team metadata ``default_estimated_output_tokens`` + + Returns ``None`` when nothing is configured, which leaves the static + heuristic floor in place. + """ + key_estimate: Final = _estimated_output_tokens_from_metadata(user_api_key_dict.metadata, model_name) + if key_estimate is not None: + return key_estimate + return _estimated_output_tokens_from_metadata(user_api_key_dict.team_metadata, model_name) + + +class OutputTokenEstimateRequest(Protocol): + """The shape of any management request that can carry an output-token estimate. + + Read-only members: the gate inspects a request, it never writes one back. + """ + + @property + def metadata(self) -> Mapping[str, object] | None: ... + + @property + def default_estimated_output_tokens(self) -> int | None: ... + + @property + def default_estimated_output_tokens_per_model(self) -> Mapping[str, int] | None: ... + + @property + def model_fields_set(self) -> Collection[str]: ... + + +def _requested_output_token_estimates( + data: OutputTokenEstimateRequest, + existing_metadata: Mapping[str, object], +) -> tuple[object, object]: + """The output-token estimates this request would leave stored on the entity. + + Mirrors how the management endpoints merge metadata: a supplied ``metadata`` + replaces the stored blob wholesale, an omitted one preserves it, and the + dedicated top-level fields overlay whatever survives. Both sources are read + because the same declaration reaches the same stored field either way. + """ + base: Final[Mapping[str, object]] = ( + (data.metadata or {}) if "metadata" in data.model_fields_set else existing_metadata + ) + return ( + data.default_estimated_output_tokens + if data.default_estimated_output_tokens is not None + else base.get(ESTIMATED_OUTPUT_TOKENS_FIELD), + data.default_estimated_output_tokens_per_model + if data.default_estimated_output_tokens_per_model is not None + else base.get(ESTIMATED_OUTPUT_TOKENS_PER_MODEL_FIELD), + ) + + +def enforce_output_token_estimates_are_admin_only( + data: OutputTokenEstimateRequest, + existing_metadata: Mapping[str, object] | None, + user_api_key_dict: UserAPIKeyAuth, + entity: Literal["key", "team"], +) -> None: + """Only a proxy admin may change what a key or team declares its models emit. + + That declaration is what the TPM limiter reserves for a request omitting + ``max_tokens``, so lowering or clearing it under-reserves against every + window the request is charged against, including the team and organization + ones the writer may not own. A key's metadata is writable by its holder and + a team's by its team admin, so neither is a trustworthy source for a value + that weakens a limit set above them. Gated on the resulting value rather + than on presence, so a form resending the stored declaration stays a no-op. + """ + if user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN.value: + return + stored: Final[Mapping[str, object]] = existing_metadata or {} + if _requested_output_token_estimates(data, stored) == ( + stored.get(ESTIMATED_OUTPUT_TOKENS_FIELD), + stored.get(ESTIMATED_OUTPUT_TOKENS_PER_MODEL_FIELD), + ): + return + raise HTTPException( + status_code=403, + detail={ + "error": f"Only proxy admins can set {ESTIMATED_OUTPUT_TOKENS_FIELD} or " + f"{ESTIMATED_OUTPUT_TOKENS_PER_MODEL_FIELD} on a {entity}. They decide how many output tokens " + "the rate limiter reserves for a request that omits max_tokens." + }, + ) + + def get_model_rate_limit_from_metadata( user_api_key_dict: UserAPIKeyAuth, metadata_accessor_key: Literal["team_metadata", "organization_metadata", "project_metadata"], diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index 53c5112d1d7..3fd2adda480 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -31,6 +31,8 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.auth_utils import ( + ESTIMATED_OUTPUT_TOKENS_FIELD, + get_estimated_output_tokens, get_key_tag_rpm_limit, get_model_rate_limit_from_metadata, ) @@ -562,6 +564,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): data: dict, model: str | None = None, min_configured_tpm_limit: int | None = None, + configured_output_tokens: int | None = None, ) -> int: """ Estimate total tokens this request will consume so we can reserve them @@ -575,6 +578,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): provided, the no-``max_tokens`` output-budget floor is capped at a fraction of that limit so small TPM caps remain usable. Omit to preserve the unconstrained floor. + + ``configured_output_tokens`` is the operator-declared estimate resolved + from key or team metadata. When provided it replaces the heuristic + floor entirely, so the reservation reflects what this tenant's model + actually emits rather than one constant shared by every tenant. """ messages = data.get("messages") prompt: Final = data.get("prompt") @@ -604,7 +612,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): case (_, embeddings_input) if embeddings_input: # Embeddings have no output tokens max_tokens_estimate = 0 - case _ if total_chars == 0: + case _ if total_chars == 0 and configured_output_tokens is None: # Fully contentless request (no messages, prompt, or input). # Don't apply the conservative output-budget floor here — it # would over-reserve and could push small TPM limits into a @@ -619,7 +627,11 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): # so a small per-tenant TPM cap can't be tripped by the floor # alone. output_floor: Final = self._no_max_tokens_output_floor(min_configured_tpm_limit) - max_tokens_estimate = max(estimated_input_tokens, output_floor) + max_tokens_estimate = ( + configured_output_tokens + if configured_output_tokens is not None + else max(estimated_input_tokens, output_floor) + ) total_estimated: Final = estimated_input_tokens + max_tokens_estimate @@ -2586,8 +2598,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): data.get("max_tokens") is not None or data.get("max_completion_tokens") is not None ) is_embedding: Final = data.get("input") is not None + configured_output_tokens: Final = get_estimated_output_tokens( + user_api_key_dict=user_api_key_dict, + model_name=requested_model, + ) if capped_floor < baseline_floor and not has_explicit_max_tokens and not is_embedding: - data["max_tokens"] = capped_floor + data["max_tokens"] = max(capped_floor, configured_output_tokens or 0) # Floor at 1 token so contentless requests (/responses, # tool-call continuations, empty messages) still flow @@ -2601,10 +2617,23 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): data=data, model=requested_model, min_configured_tpm_limit=min_configured_tpm_limit, + configured_output_tokens=configured_output_tokens, ), 1, ) + if configured_output_tokens is not None and estimated_tokens > min_configured_tpm_limit: + verbose_proxy_logger.debug( + "Reserving %s tokens for model %s (declared %s=%s plus the input estimate) exceeds the " + "smallest TPM limit this request is charged against (%s), so it cannot be admitted even " + "against an empty window. Lower the declared estimate or raise the TPM limit.", + estimated_tokens, + requested_model, + ESTIMATED_OUTPUT_TOKENS_FIELD, + configured_output_tokens, + min_configured_tpm_limit, + ) + tpm_response: Final = await self.reserve_tpm_tokens( descriptors=descriptors, estimated_tokens=estimated_tokens, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2de1d177b33..3a97558fcbe 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -55,7 +55,10 @@ from litellm.proxy.auth.auth_checks import ( get_project_object, get_team_object, ) -from litellm.proxy.auth.auth_utils import abbreviate_api_key +from litellm.proxy.auth.auth_utils import ( + abbreviate_api_key, + enforce_output_token_estimates_are_admin_only, +) from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import ( decrypt_callback_vars, @@ -847,6 +850,13 @@ async def _common_key_generation_helper( detail={"error": "Only proxy admins can enable throttle_on_budget_exceeded on a key."}, ) + enforce_output_token_estimates_are_admin_only( + data=data, + existing_metadata=None, + user_api_key_dict=user_api_key_dict, + entity="key", + ) + if data.metadata is not None and data.metadata.get("service_account_id") is not None and data.team_id is None: await validate_team_id_used_in_service_account_request( team_id=data.team_id, @@ -1584,6 +1594,8 @@ async def generate_key_fn( - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit. + - default_estimated_output_tokens: Optional[int] - Proxy admin only. Expected output tokens reserved for TPM limiting when a request omits max_tokens. Positive integer. Falls back to the team setting, then to the built-in estimate. + - default_estimated_output_tokens_per_model: Optional[dict] - Proxy admin only. Per-model override of the above. Example - {"gpt-4": 4096, "gpt-3.5-turbo": 1024}. Takes precedence over the key-wide value. - mcp_rpm_limit: Optional[dict] - key-specific per-MCP-server rpm limit, keyed by MCP server name (alias if set, else the configured name). Example - {"github": 100, "slack": 200}. IF null or {} then no MCP-specific rpm limit. - tag_rpm_limit: Optional[dict] - key-specific per-request-tag rpm limit, keyed by request tag. Example - {"cell-1": 1000, "cell-2": 500}. Each tag gets an independent counter; requests whose tag is absent fall back to the key-level rpm limit. - tpm_limit_type: Optional[str] - Type of tpm limit. Options: "best_effort_throughput" (no error if we're overallocating tpm), "guaranteed_throughput" (raise an error if we're overallocating tpm), "dynamic" (dynamically exceed limit when no 429 errors). Defaults to "best_effort_throughput". @@ -1793,6 +1805,8 @@ async def generate_service_account_key_fn( - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. - model_rpm_limit: Optional[dict] - key-specific model rpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific rpm limit. - model_tpm_limit: Optional[dict] - key-specific model tpm limit. Example - {"text-davinci-002": 1000, "gpt-3.5-turbo": 1000}. IF null or {} then no model specific tpm limit. + - default_estimated_output_tokens: Optional[int] - Proxy admin only. Expected output tokens reserved for TPM limiting when a request omits max_tokens. Positive integer. Falls back to the team setting, then to the built-in estimate. + - default_estimated_output_tokens_per_model: Optional[dict] - Proxy admin only. Per-model override of the above. Example - {"gpt-4": 4096, "gpt-3.5-turbo": 1024}. Takes precedence over the key-wide value. - mcp_rpm_limit: Optional[dict] - key-specific per-MCP-server rpm limit, keyed by MCP server name (alias if set, else the configured name). Example - {"github": 100, "slack": 200}. IF null or {} then no MCP-specific rpm limit. - tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" @@ -2473,6 +2487,13 @@ async def _validate_update_key_data( detail={"error": "Only proxy admins can enable throttle_on_budget_exceeded on a key."}, ) + enforce_output_token_estimates_are_admin_only( + data=data, + existing_metadata=_existing_metadata if isinstance(_existing_metadata, dict) else None, + user_api_key_dict=user_api_key_dict, + entity="key", + ) + # Personal-key bypass: the caller both created the key AND still owns it # (user_id == caller). Checking only created_by would let a demoted admin # who originally created a key for another user continue editing it without @@ -2655,6 +2676,8 @@ async def update_key_fn( - mcp_rpm_limit: Optional[dict] - Per-MCP-server RPM limits, keyed by MCP server name {"github": 100, "slack": 200} - tag_rpm_limit: Optional[dict] - Per-request-tag RPM limits, keyed by request tag {"cell-1": 1000, "cell-2": 500}. Each tag gets an independent counter; absent tags fall back to the key-level rpm limit. - model_tpm_limit: Optional[dict] - Model-specific TPM limits {"gpt-4": 100000, "claude-v1": 200000} + - default_estimated_output_tokens: Optional[int] - Proxy admin only. Expected output tokens reserved for TPM limiting when a request omits max_tokens. Positive integer. + - default_estimated_output_tokens_per_model: Optional[dict] - Proxy admin only. Per-model override of the above {"gpt-4": 4096, "gpt-3.5-turbo": 1024} - tpm_limit_type: Optional[str] - TPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - rpm_limit_type: Optional[str] - RPM rate limit type - "best_effort_throughput", "guaranteed_throughput", or "dynamic" - allowed_cache_controls: Optional[list] - List of allowed cache control values @@ -4629,6 +4652,15 @@ async def _execute_virtual_key_regeneration( prisma_client=prisma_client, ) + if data is not None: + _existing_key_metadata: Final = getattr(key_in_db, "metadata", None) + enforce_output_token_estimates_are_admin_only( + data=data, + existing_metadata=_existing_key_metadata if isinstance(_existing_key_metadata, dict) else None, + user_api_key_dict=user_api_key_dict, + entity="key", + ) + new_token: Final = await get_new_token(data=data) new_token_hash: Final = hash_token(new_token) new_token_key_name: Final = abbreviate_api_key(api_key=new_token) diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index 97f494c51de..f5d0d63e311 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -82,6 +82,7 @@ from litellm.proxy.auth.auth_checks import ( get_team_object, get_user_object, ) +from litellm.proxy.auth.auth_utils import enforce_output_token_estimates_are_admin_only from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.callback_utils import encrypt_callback_vars from litellm.proxy.common_utils.json_merge_patch import apply_json_merge_patch @@ -1153,6 +1154,8 @@ async def new_team( - metadata: Optional[dict] - Metadata for team, store information for team. Example metadata = {"extra_info": "some info"} - model_rpm_limit: Optional[Dict[str, int]] - The RPM (Requests Per Minute) limit for this team - applied across all keys for this team. - model_tpm_limit: Optional[Dict[str, int]] - The TPM (Tokens Per Minute) limit for this team - applied across all keys for this team. + - default_estimated_output_tokens: Optional[int] - Expected output tokens reserved for TPM limiting when a request omits max_tokens, for keys on this team that do not set their own. Positive integer. + - default_estimated_output_tokens_per_model: Optional[Dict[str, int]] - Per-model override of the above. Example: {"gpt-4": 4096, "gpt-3.5-turbo": 1024} - mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team. - tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for this team - all keys with this team_id will have at max this TPM limit - rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit @@ -1266,6 +1269,13 @@ async def new_team( }, ) + enforce_output_token_estimates_are_admin_only( + data=data, + existing_metadata=None, + user_api_key_dict=user_api_key_dict, + entity="team", + ) + # Check if license is over limit total_teams: Final = await _team_db(prisma_client).count() if total_teams and _license_check.is_team_count_over_limit(team_count=total_teams): @@ -1863,6 +1873,8 @@ async def update_team( - allowed_passthrough_routes: Optional[List[str]] - List of allowed pass through routes for the team. - model_rpm_limit: Optional[Dict[str, int]] - The RPM (Requests Per Minute) limit per model for this team. Example: {"gpt-4": 100, "gpt-3.5-turbo": 200} - model_tpm_limit: Optional[Dict[str, int]] - The TPM (Tokens Per Minute) limit per model for this team. Example: {"gpt-4": 10000, "gpt-3.5-turbo": 20000} + - default_estimated_output_tokens: Optional[int] - Expected output tokens reserved for TPM limiting when a request omits max_tokens, for keys on this team that do not set their own. Positive integer. + - default_estimated_output_tokens_per_model: Optional[Dict[str, int]] - Per-model override of the above. Example: {"gpt-4": 4096, "gpt-3.5-turbo": 1024} - mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team. Example - update team TPM Limit - allowed_vector_store_indexes: Optional[List[dict]] - List of allowed vector store indexes for the key. Example - [{"index_name": "my-index", "index_permissions": ["write", "read"]}]. If specified, the key will only be able to use these specific vector store indexes. Create index, using `/v1/indexes` endpoint. @@ -1949,6 +1961,14 @@ async def update_team( user_api_key_dict=user_api_key_dict, ) + _existing_team_metadata: Final[object] = getattr(existing_team_row, "metadata", None) + enforce_output_token_estimates_are_admin_only( + data=data, + existing_metadata=_existing_team_metadata if isinstance(_existing_team_metadata, dict) else None, + user_api_key_dict=user_api_key_dict, + entity="team", + ) + _check_passthrough_routes_caller_permission(data, user_api_key_dict, entity="team") if data.soft_budget is not None: diff --git a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py index 1c29c287c3a..076151fcd3b 100644 --- a/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py +++ b/tests/test_litellm/proxy/hooks/test_parallel_request_limiter_v3.py @@ -3,6 +3,7 @@ Unit Tests for the max parallel request limiter v3 for the proxy """ import asyncio +import logging import os import sys import time @@ -5100,3 +5101,456 @@ async def test_reserve_tpm_tokens_never_evaluates_the_requests_dimension(): f"reservation pass, got: {response}" ) assert [s["rate_limit_type"] for s in response["statuses"]] == ["tokens"] + + +STATIC_OUTPUT_FLOOR = 1024 +ONE_TOKEN_PROMPT = [{"role": "user", "content": "hello"}] +ONE_TOKEN_PROMPT_INPUT_ESTIMATE = 1 + + +async def _reserved_tokens_for( + handler, + local_cache, + user_api_key_dict, + data, + call_type="completion", +): + """Drive the pre-call hook and read back what landed on the :tokens counter.""" + tokens_key = handler.create_rate_limit_keys( + key="api_key", value=user_api_key_dict.api_key, rate_limit_type="tokens" + ) + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data=data, + call_type=call_type, + ) + return int(await local_cache.async_get_cache(key=tokens_key) or 0) + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_metadata, team_metadata, expected_output_estimate, tier", + [ + ( + { + "default_estimated_output_tokens_per_model": {"gpt-4o-mini": 3001}, + "default_estimated_output_tokens": 2002, + }, + { + "default_estimated_output_tokens_per_model": {"gpt-4o-mini": 1503}, + "default_estimated_output_tokens": 777, + }, + 3001, + "key per-model wins over every other tier", + ), + ( + {"default_estimated_output_tokens": 2002}, + { + "default_estimated_output_tokens_per_model": {"gpt-4o-mini": 1503}, + "default_estimated_output_tokens": 777, + }, + 2002, + "key global wins over team config", + ), + ( + {"default_estimated_output_tokens_per_model": {"some-other-model": 9999}}, + { + "default_estimated_output_tokens_per_model": {"gpt-4o-mini": 1503}, + "default_estimated_output_tokens": 777, + }, + 1503, + "team per-model wins when the key has no applicable entry", + ), + ( + {}, + {"default_estimated_output_tokens": 777}, + 777, + "team global is the last configured tier", + ), + ({}, {}, STATIC_OUTPUT_FLOOR, "unconfigured falls back to the static floor"), + ( + {"unrelated": "value"}, + {"unrelated": "value"}, + STATIC_OUTPUT_FLOOR, + "unrelated metadata changes nothing", + ), + ( + {"default_estimated_output_tokens": "not-a-number"}, + {}, + STATIC_OUTPUT_FLOOR, + "malformed config falls back to the static floor instead of erroring", + ), + ( + {"default_estimated_output_tokens": 0}, + {}, + STATIC_OUTPUT_FLOOR, + "a non-positive estimate is rejected, not reserved", + ), + ], +) +async def test_estimated_output_tokens_resolution_precedence( + monkeypatch, key_metadata, team_metadata, expected_output_estimate, tier +): + """The no-max_tokens output reservation resolves per key / team / model. + + Every configured value here is distinct from the static 1024 floor and + from the input estimate, so the reserved amount identifies which tier the + resolver picked. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token(f"sk-estimate-{expected_output_estimate}-{tier}"), + tpm_limit=1_000_000, + metadata=key_metadata, + team_metadata=team_metadata, + ) + + reserved = await _reserved_tokens_for( + handler, + local_cache, + user_api_key_dict, + {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT}, + ) + + assert reserved == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + expected_output_estimate, tier + + +@pytest.mark.asyncio +async def test_request_max_tokens_outranks_configured_estimate(monkeypatch): + """An explicit request-level max_tokens stays the top of the precedence order.""" + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-estimate-explicit-max-tokens"), + tpm_limit=1_000_000, + metadata={"default_estimated_output_tokens": 2002}, + ) + + reserved = await _reserved_tokens_for( + handler, + local_cache, + user_api_key_dict, + {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT, "max_tokens": 42}, + ) + + assert reserved == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + 42 + + +@pytest.mark.asyncio +async def test_configured_estimate_does_not_apply_to_embeddings(monkeypatch): + """Embeddings generate no output, so a declared output estimate must not be reserved.""" + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token("sk-estimate-embeddings"), + tpm_limit=1_000_000, + metadata={"default_estimated_output_tokens": 2002}, + ) + + reserved = await _reserved_tokens_for( + handler, + local_cache, + user_api_key_dict, + {"model": "text-embedding-3-small", "input": "hello"}, + call_type="embeddings", + ) + + assert reserved == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + + +@pytest.mark.asyncio +async def test_configured_estimate_applies_to_contentless_requests(monkeypatch): + """A declared estimate describes generation, so it holds even with no prompt body. + + Without config such a request reserves the 1-token floor only; the + declaration is what makes concurrent tool-call continuations countable. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + configured = UserAPIKeyAuth( + api_key=hash_token("sk-estimate-contentless-configured"), + tpm_limit=1_000_000, + metadata={"default_estimated_output_tokens": 2002}, + ) + unconfigured = UserAPIKeyAuth( + api_key=hash_token("sk-estimate-contentless-plain"), + tpm_limit=1_000_000, + ) + + assert ( + await _reserved_tokens_for( + handler, local_cache, configured, {"model": "gpt-4o-mini", "messages": []} + ) + == 2002 + ) + assert ( + await _reserved_tokens_for( + handler, local_cache, unconfigured, {"model": "gpt-4o-mini", "messages": []} + ) + == 1 + ) + + +@pytest.mark.asyncio +async def test_declared_estimate_never_tightens_the_small_tpm_clamp(monkeypatch): + """The small-TPM clamp can only be loosened by a declaration, never tightened. + + That clamp is the one place the proxy rewrites the caller's generation + budget, and it only fires below a 4096 TPM limit. A declaration above it + raises it, so the tenant is not truncated below what they said their + model emits; a declaration below it changes nothing, because an estimate + describes the typical response and must not become a hard cap that + truncates the tail. The reservation tracks whatever the clamp settles on, + so a small tenant can never generate more than was reserved. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + raised_data: Dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} + raised_reserved = await _reserved_tokens_for( + handler, + local_cache, + UserAPIKeyAuth( + api_key=hash_token("sk-estimate-hard-cap-raised"), + tpm_limit=2000, + metadata={"default_estimated_output_tokens": 900}, + ), + raised_data, + ) + assert raised_data["max_tokens"] == 900 + assert raised_reserved == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + 900 + + lowered_data: Dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} + lowered_reserved = await _reserved_tokens_for( + handler, + local_cache, + UserAPIKeyAuth( + api_key=hash_token("sk-estimate-hard-cap-lowered"), + tpm_limit=2000, + metadata={"default_estimated_output_tokens": 120}, + ), + lowered_data, + ) + assert lowered_data["max_tokens"] == 500 + assert lowered_reserved == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + 500 + + unconfigured_data: Dict[str, Any] = {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT} + unconfigured_reserved = await _reserved_tokens_for( + handler, + local_cache, + UserAPIKeyAuth( + api_key=hash_token("sk-estimate-hard-cap-plain"), + tpm_limit=2000, + ), + unconfigured_data, + ) + assert unconfigured_data["max_tokens"] == 500 + assert unconfigured_reserved == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + 500 + + +@pytest.mark.asyncio +async def test_one_malformed_estimate_field_does_not_discard_the_other(monkeypatch): + """Each declared field is validated on its own. + + A per-model map with a bad entry must not take a valid global estimate + down with it, and a bad global must not hide a valid per-model entry. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + broken_map = await _reserved_tokens_for( + handler, + local_cache, + UserAPIKeyAuth( + api_key=hash_token("sk-estimate-broken-map"), + tpm_limit=1_000_000, + metadata={ + "default_estimated_output_tokens_per_model": {"gpt-4o-mini": "huge"}, + "default_estimated_output_tokens": 2002, + }, + ), + {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT}, + ) + assert broken_map == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + 2002 + + broken_global = await _reserved_tokens_for( + handler, + local_cache, + UserAPIKeyAuth( + api_key=hash_token("sk-estimate-broken-global"), + tpm_limit=1_000_000, + metadata={ + "default_estimated_output_tokens_per_model": {"gpt-4o-mini": 3001}, + "default_estimated_output_tokens": -5, + }, + ), + {"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT}, + ) + assert broken_global == ONE_TOKEN_PROMPT_INPUT_ESTIMATE + 3001 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("declared", [100_000, 5000]) +async def test_declared_estimate_over_the_tpm_budget_is_honored_and_explained(monkeypatch, caplog, declared): + """A declaration bigger than the budget must not be silently shrunk. + + Capping it against the TPM limit would re-admit exactly the traffic this + feature exists to hold back, so the request is refused instead and the + reservation is explained rather than leaving an unexplained 429 loop. + + ``declared == tpm_limit`` is the boundary case: the declaration alone + equals the limit, so only adding the input estimate tips the reservation + over. Comparing the declaration against the limit rather than the + reservation would refuse this request while saying nothing. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token(f"sk-estimate-over-budget-{declared}"), + tpm_limit=5000, + metadata={"default_estimated_output_tokens": declared}, + ) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT}, + call_type="completion", + ) + + assert exc_info.value.status_code == 429 + explained = [ + record.getMessage() + for record in caplog.records + if "cannot be admitted even against an empty window" in record.getMessage() + ] + assert len(explained) == 1, f"expected exactly one explanation, got {explained}" + assert str(declared) in explained[0] + assert str(ONE_TOKEN_PROMPT_INPUT_ESTIMATE + declared) in explained[0] + assert "5000" in explained[0] + + +@pytest.mark.asyncio +async def test_a_key_that_declared_nothing_is_never_blamed_for_a_declaration(monkeypatch, caplog): + """A request can outgrow its budget on prompt size alone, with no declaration. + + The heuristic path reserves input plus the injected clamp, so a long + prompt against a small limit is refused without anyone having declared + anything. Blaming the declared field there would point an operator at a + setting they never set, to fix a 429 whose real cause is prompt size. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + with pytest.raises(HTTPException) as exc_info: + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth( + api_key=hash_token("sk-undeclared-long-prompt"), + tpm_limit=1000, + ), + cache=local_cache, + data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "x" * 3600}]}, + call_type="completion", + ) + + assert exc_info.value.status_code == 429 + assert not [ + record for record in caplog.records if "cannot be admitted even against an empty window" in record.getMessage() + ] + + +@pytest.mark.asyncio +async def test_declared_estimate_inside_the_tpm_budget_is_not_explained(monkeypatch, caplog): + """The explanation is for requests that cannot fit, not for every request. + + Without this, a correctly configured key would emit one line per call. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + + with caplog.at_level(logging.DEBUG, logger="LiteLLM Proxy"): + await handler.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth( + api_key=hash_token("sk-estimate-within-budget"), + tpm_limit=5000, + metadata={"default_estimated_output_tokens": 1000}, + ), + cache=local_cache, + data={"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT}, + call_type="completion", + ) + + assert not [ + record for record in caplog.records if "cannot be admitted even against an empty window" in record.getMessage() + ] + + +@pytest.mark.asyncio +async def test_configured_estimate_blocks_the_overrun_the_static_floor_admits(monkeypatch): + """Concurrent unbounded requests must stop at the declared budget. + + A key with tpm_limit=8000 whose model really emits ~3000 output tokens + admits 7 concurrent requests under the 1024 floor (7 * 1025 <= 8000), so + once they all report actual usage the window carries ~21000 tokens + against an 8000 limit. Declaring the real output size admits only the two + requests the budget actually covers. + """ + monkeypatch.delenv("LITELLM_TPM_TOKEN_RESERVATION_ENABLED", raising=False) + + async def admitted(metadata): + local_cache = DualCache() + handler = _PROXY_MaxParallelRequestsHandler( + internal_usage_cache=InternalUsageCache(local_cache) + ) + user_api_key_dict = UserAPIKeyAuth( + api_key=hash_token(f"sk-overrun-{metadata}"), + tpm_limit=8000, + metadata=metadata, + ) + accepted = 0 + for _ in range(10): + try: + await handler.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=local_cache, + data={"model": "gpt-4o-mini", "messages": ONE_TOKEN_PROMPT}, + call_type="completion", + ) + except HTTPException: + break + accepted += 1 + return accepted + + assert await admitted({}) == 7 + assert await admitted({"default_estimated_output_tokens": 3000}) == 2 diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index e8709f3af34..3ce25f7a934 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -15442,3 +15442,312 @@ async def test_migrate_encryption_endpoint_rejects_proxy_admin_viewer(): assert exc_info.value.status_code == 403 mock_migrate.assert_not_awaited() + + +_ESTIMATE = "default_estimated_output_tokens" +_ESTIMATE_PER_MODEL = "default_estimated_output_tokens_per_model" + + +@pytest.mark.parametrize( + "label, request_body, existing_metadata, allowed", + [ + ("nothing declared", {}, None, True), + ("declared top-level on a key with none stored", {_ESTIMATE: 1}, None, False), + ("declared inside metadata on a key with none stored", {"metadata": {_ESTIMATE: 1}}, None, False), + ( + "per-model map declared inside metadata", + {"metadata": {_ESTIMATE_PER_MODEL: {"gpt-4": 1}}}, + None, + False, + ), + ("unrelated edit, metadata omitted", {"models": ["gpt-4"]}, {_ESTIMATE: 2000}, True), + ("stored value resent unchanged", {_ESTIMATE: 2000}, {_ESTIMATE: 2000}, True), + ("stored value lowered", {_ESTIMATE: 1}, {_ESTIMATE: 2000}, False), + ("stored value raised", {_ESTIMATE: 9000}, {_ESTIMATE: 2000}, False), + ( + "stored value cleared by sending a metadata blob without it", + {"metadata": {"other": "keep"}}, + {_ESTIMATE: 2000, "other": "keep"}, + False, + ), + ( + "stored value resent inside the metadata blob", + {"metadata": {_ESTIMATE: 2000, "other": "keep"}}, + {_ESTIMATE: 2000, "other": "keep"}, + True, + ), + ( + "per-model map resent unchanged", + {_ESTIMATE_PER_MODEL: {"gpt-4": 4096}}, + {_ESTIMATE_PER_MODEL: {"gpt-4": 4096}}, + True, + ), + ( + "one model in the per-model map lowered", + {_ESTIMATE_PER_MODEL: {"gpt-4": 1}}, + {_ESTIMATE_PER_MODEL: {"gpt-4": 4096}}, + False, + ), + ], +) +def test_output_token_estimate_admin_gate_matrix(label, request_body, existing_metadata, allowed): + """A non-admin may only leave a key's stored output-token estimate exactly as it is. + + The estimate decides what the TPM limiter reserves for a request that omits + max_tokens, so lowering, raising or clearing it moves a reservation charged + against team and organization windows the key holder does not own. Key + metadata is writable by the key holder, and the declaration can be written + either as a dedicated top-level field or nested in the metadata blob, so + both routes are gated. Resending the stored value is what the edit form + produces on every save and has to stay allowed. + """ + from litellm.proxy.auth.auth_utils import ( + enforce_output_token_estimates_are_admin_only, + ) + + def _call(caller): + enforce_output_token_estimates_are_admin_only( + data=UpdateKeyRequest(key="sk-1", **request_body), + existing_metadata=existing_metadata, + user_api_key_dict=caller, + entity="key", + ) + + non_admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-non-admin", + user_id="alice", + ) + if allowed: + _call(non_admin) + else: + with pytest.raises(HTTPException) as exc: + _call(non_admin) + assert exc.value.status_code == 403 + assert "Only proxy admins can set" in str(exc.value.detail) + + _call( + UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin", + ) + ) + + +@pytest.mark.asyncio +async def test_generate_key_output_token_estimate_rejected_for_non_admin(): + """The /key/update gate does not cover generate, so without its own check a + non-admin could self-mint a key that reserves one output token per + unbounded request and overrun the TPM window it is charged against.""" + with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()): + with pytest.raises(HTTPException) as exc: + await _common_key_generation_helper( + data=GenerateKeyRequest(default_estimated_output_tokens=1, tpm_limit=100000), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + litellm_changed_by=None, + team_table=None, + ) + assert int(getattr(exc.value, "status_code", 0)) == 403 + assert "Only proxy admins can set" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_generate_key_output_token_estimate_in_metadata_rejected_for_non_admin(): + """Writing the declaration into the raw metadata blob lands in the same + stored field, so gating only the dedicated top-level field leaves the + bypass wide open.""" + with patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()): + with pytest.raises(HTTPException) as exc: + await _common_key_generation_helper( + data=GenerateKeyRequest(metadata={"default_estimated_output_tokens": 1}), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + litellm_changed_by=None, + team_table=None, + ) + assert int(getattr(exc.value, "status_code", 0)) == 403 + + +@pytest.mark.asyncio +async def test_generate_key_output_token_estimate_allowed_for_admin(): + """A proxy admin declaring the estimate must reach key creation.""" + with ( + patch("litellm.proxy.proxy_server.prisma_client", AsyncMock()), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.premium_user", False), + patch( + "litellm.proxy.management_endpoints.key_management_endpoints.generate_key_helper_fn" + ) as mock_generate_key, + ): + mock_generate_key.return_value = { + "key": "sk-test-key", + "expires": None, + "user_id": "admin", + "team_id": None, + } + await _common_key_generation_helper( + data=GenerateKeyRequest(default_estimated_output_tokens=200), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-1"), + litellm_changed_by=None, + team_table=None, + ) + assert mock_generate_key.called + + +def _estimate_key_row(token: str, metadata: dict): + existing_key = MagicMock() + existing_key.token = token + existing_key.user_id = "internal_user" + existing_key.created_by = "internal_user" + existing_key.team_id = None + existing_key.project_id = None + existing_key.max_budget = 10.0 + existing_key.key_alias = None + existing_key.models = [] + existing_key.metadata = metadata + existing_key.model_dump.return_value = { + "token": token, + "user_id": "internal_user", + "team_id": None, + "max_budget": 10.0, + } + return existing_key + + +def _wire_update_key_fn(monkeypatch, existing_key): + mock_prisma_client = AsyncMock() + updated_key = MagicMock() + updated_key.token = existing_key.token + updated_key.key_alias = "my-alias" + + mock_prisma_client.get_data = AsyncMock(return_value=existing_key) + mock_prisma_client.update_data = AsyncMock(return_value=updated_key) + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing_key) + + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", AsyncMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", MagicMock()) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.store_audit_logs", False) + monkeypatch.setattr("litellm.proxy.proxy_server.hash_token", lambda token: existing_key.token) + + async def _noop(**kwargs): + pass + + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._delete_cache_key_object", + _noop, + ) + monkeypatch.setattr( + "litellm.proxy.management_endpoints.key_management_endpoints._enforce_unique_key_alias", + _noop, + ) + + +@pytest.mark.asyncio +async def test_update_key_output_token_estimate_lowered_rejected_for_non_admin(monkeypatch): + """End-to-end wiring: a key's owner reaches /key/update without any admin + check because metadata is a non-budget field, so the gate has to fire + inside the update path itself rather than only in a helper.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + token = "a1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + _wire_update_key_fn(monkeypatch, _estimate_key_row(token, {_ESTIMATE: 4000})) + + mock_request = MagicMock() + mock_request.query_params = {} + + with pytest.raises(ProxyException) as exc: + await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key=token, default_estimated_output_tokens=1), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + ) + + assert str(exc.value.code) == "403" + assert "Only proxy admins can set" in str(exc.value.message) + + +@pytest.mark.asyncio +async def test_update_key_output_token_estimate_unchanged_allows_non_admin_edit(monkeypatch): + """The edit form resends every field it renders, so gating on presence + would 403 a key owner renaming their own key.""" + from litellm.proxy.management_endpoints.key_management_endpoints import ( + update_key_fn, + ) + + token = "b1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + _wire_update_key_fn(monkeypatch, _estimate_key_row(token, {_ESTIMATE: 4000})) + + mock_request = MagicMock() + mock_request.query_params = {} + + result = await update_key_fn( + request=mock_request, + data=UpdateKeyRequest(key=token, key_alias="my-alias", default_estimated_output_tokens=4000), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + ) + + assert result is not None + + +@pytest.mark.asyncio +async def test_regenerate_key_output_token_estimate_lowered_rejected_for_non_admin(): + """/key/regenerate is a third write path into the same stored metadata. + + can_modify_verification_token lets a key's own holder regenerate it, and + the request body runs through prepare_key_update_data exactly as an update + does, so gating only generate and update leaves the declaration writable. + """ + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + token = "c1b2c3d4e5f6789012345678901234567890123456789012345678901234abcd" + key_in_db = LiteLLM_VerificationToken( + token=token, + user_id="internal_user", + metadata={_ESTIMATE: 4000}, + ) + + with pytest.raises(HTTPException) as exc: + await _execute_virtual_key_regeneration( + prisma_client=AsyncMock(), + key_in_db=key_in_db, + hashed_api_key=token, + key="sk-original", + data=RegenerateKeyRequest(key="sk-original", default_estimated_output_tokens=1), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-internal", + user_id="internal_user", + ), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + + assert exc.value.status_code == 403 + assert "Only proxy admins can set" in str(exc.value.detail) diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 1e47010b57c..f17a6dbd380 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -11008,3 +11008,207 @@ def test_validate_member_user_id_provisioning_caps_the_ids_it_echoes_back(): assert f"u{_MAX_REPORTED_UNKNOWN_USER_IDS}" not in detail assert f"and {500 - _MAX_REPORTED_UNKNOWN_USER_IDS} more" in detail assert len(detail) < 1000 + + +_TEAM_ESTIMATE = "default_estimated_output_tokens" +_TEAM_ESTIMATE_PER_MODEL = "default_estimated_output_tokens_per_model" + + +@pytest.mark.parametrize( + "label, request_body, existing_metadata, allowed", + [ + ("nothing declared", {}, None, True), + ("declared top-level with none stored", {_TEAM_ESTIMATE: 1}, None, False), + ("declared inside metadata with none stored", {"metadata": {_TEAM_ESTIMATE: 1}}, None, False), + ( + "per-model map declared inside metadata", + {"metadata": {_TEAM_ESTIMATE_PER_MODEL: {"gpt-4": 1}}}, + None, + False, + ), + ("unrelated edit, metadata omitted", {"tpm_limit": 99}, {_TEAM_ESTIMATE: 2000}, True), + ("stored value resent unchanged", {_TEAM_ESTIMATE: 2000}, {_TEAM_ESTIMATE: 2000}, True), + ("stored value lowered", {_TEAM_ESTIMATE: 1}, {_TEAM_ESTIMATE: 2000}, False), + ("stored value raised", {_TEAM_ESTIMATE: 9000}, {_TEAM_ESTIMATE: 2000}, False), + ( + "stored value cleared by sending a metadata blob without it", + {"metadata": {"other": "keep"}}, + {_TEAM_ESTIMATE: 2000, "other": "keep"}, + False, + ), + ( + "stored value resent inside the metadata blob", + {"metadata": {_TEAM_ESTIMATE: 2000, "other": "keep"}}, + {_TEAM_ESTIMATE: 2000, "other": "keep"}, + True, + ), + ( + "per-model map resent unchanged", + {_TEAM_ESTIMATE_PER_MODEL: {"gpt-4": 4096}}, + {_TEAM_ESTIMATE_PER_MODEL: {"gpt-4": 4096}}, + True, + ), + ( + "one model in the per-model map lowered", + {_TEAM_ESTIMATE_PER_MODEL: {"gpt-4": 1}}, + {_TEAM_ESTIMATE_PER_MODEL: {"gpt-4": 4096}}, + False, + ), + ], +) +def test_team_output_token_estimate_admin_gate_matrix(label, request_body, existing_metadata, allowed): + """A team admin may only leave a team's stored output-token estimate exactly as it is. + + A team admin can write team metadata, and every key on the team inherits the + team declaration, so without this a team admin could shrink the reservation + for the whole team and under-reserve against an organization TPM window the + organization set above them. Same value-transition rule as the key gate, + including the raw-metadata route and clearing by omission. + """ + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.auth.auth_utils import ( + enforce_output_token_estimates_are_admin_only, + ) + + def _call(caller): + enforce_output_token_estimates_are_admin_only( + data=UpdateTeamRequest(team_id="t", **request_body), + existing_metadata=existing_metadata, + user_api_key_dict=caller, + entity="team", + ) + + team_admin = UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-team-admin", + user_id="team-admin", + ) + if allowed: + _call(team_admin) + else: + with pytest.raises(HTTPException) as exc: + _call(team_admin) + assert exc.value.status_code == 403 + assert "on a team" in str(exc.value.detail) + + _call( + UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin", + ) + ) + + +def _wire_update_team(stack, existing_metadata): + """Mock just enough of update_team to reach (or pass) the estimate gate.""" + from unittest.mock import AsyncMock, MagicMock, patch + + mock_prisma_client = stack.enter_context(patch("litellm.proxy.proxy_server.prisma_client")) + stack.enter_context(patch("litellm.proxy.proxy_server.llm_router")) + stack.enter_context(patch("litellm.proxy.proxy_server.user_api_key_cache")) + stack.enter_context(patch("litellm.proxy.proxy_server.proxy_logging_obj")) + stack.enter_context(patch("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")) + stack.enter_context(patch("litellm.proxy.management_endpoints.team_endpoints._cache_team_object")) + + existing_team = MagicMock() + existing_team.metadata = existing_metadata + existing_team.model_dump.return_value = { + "team_id": "test_team_id", + "team_alias": "test_team", + "metadata": existing_metadata, + "members_with_roles": [{"user_id": "team-admin", "role": "admin"}], + } + mock_prisma_client.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + + updated_team = MagicMock() + updated_team.team_id = "test_team_id" + updated_team.model_dump.return_value = {"team_id": "test_team_id"} + mock_prisma_client.db.litellm_teamtable.update = AsyncMock(return_value=updated_team) + mock_prisma_client.jsonify_team_object = MagicMock(side_effect=lambda db_data: db_data) + return mock_prisma_client + + +@pytest.mark.asyncio +async def test_update_team_output_token_estimate_lowered_rejected_for_team_admin(): + """End-to-end wiring: _verify_team_access admits a team admin, so the gate + has to fire inside update_team itself.""" + import contextlib + from unittest.mock import Mock + + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import update_team + + with contextlib.ExitStack() as stack: + _wire_update_team(stack, {_TEAM_ESTIMATE: 4000}) + with pytest.raises(ProxyException) as exc: + await update_team( + data=UpdateTeamRequest(team_id="test_team_id", default_estimated_output_tokens=1), + http_request=Mock(spec=Request), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-team-admin", + user_id="team-admin", + ), + ) + + assert str(exc.value.code) == "403" + assert "on a team" in str(exc.value.message) + + +@pytest.mark.asyncio +async def test_update_team_output_token_estimate_unchanged_allows_team_admin_edit(): + """The team settings form resends every field it renders, so gating on + presence would break a team admin editing an unrelated setting.""" + import contextlib + from unittest.mock import Mock + + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import update_team + + with contextlib.ExitStack() as stack: + prisma = _wire_update_team(stack, {_TEAM_ESTIMATE: 4000}) + await update_team( + data=UpdateTeamRequest( + team_id="test_team_id", + team_alias="renamed", + default_estimated_output_tokens=4000, + ), + http_request=Mock(spec=Request), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-team-admin", + user_id="team-admin", + ), + ) + + assert prisma.db.litellm_teamtable.update.called + + +@pytest.mark.asyncio +async def test_new_team_output_token_estimate_rejected_for_non_admin(): + """/team/new is the other write path into the same stored declaration.""" + from unittest.mock import Mock + + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + with pytest.raises(ProxyException) as exc: + await new_team( + data=NewTeamRequest(team_alias="t", default_estimated_output_tokens=1), + http_request=Mock(spec=Request), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.INTERNAL_USER, + api_key="sk-alice", + user_id="alice", + ), + ) + + assert str(exc.value.code) == "403" + assert "on a team" in str(exc.value.message) diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index 50e10285148..16248fba6a3 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -1,11 +1,26 @@ import { useTeamMetadataSchema } from "@/app/(dashboard)/hooks/teams/useTeamMetadataSchema"; import * as networking from "@/components/networking"; -import { screen, waitFor, within } from "@testing-library/react"; +import { fireEvent, screen, waitFor, within } from "@testing-library/react"; import userEvent from "@testing-library/user-event"; import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { renderWithProviders, testQueryClient } from "../../../tests/test-utils"; import TeamInfoView from "./TeamInfo"; +const authState = vi.hoisted(() => ({ userRole: "Admin" })); + +vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ + default: () => ({ + token: "123", + accessToken: "123", + userId: "user-1", + userEmail: "user@example.com", + userRole: authState.userRole, + premiumUser: false, + disabledPersonalKeyCreation: null, + showSSOBanner: false, + }), +})); + vi.mock("@/components/networking", () => ({ teamInfoCall: vi.fn(), teamMemberDeleteCall: vi.fn(), @@ -238,6 +253,7 @@ describe("TeamInfoView", () => { afterEach(() => { vi.clearAllMocks(); + authState.userRole = "Admin"; }); describe("display and rendering", () => { @@ -964,6 +980,106 @@ describe("TeamInfoView", () => { expect(updateArg.model_tpm_limit).toEqual({ "gpt-4": 100 }); }); + it("prefills the estimated output token controls, hides them from the pair editor, and saves edits", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + metadata: { + department: "research", + default_estimated_output_tokens: 512, + default_estimated_output_tokens_per_model: { "gpt-4": 4096 }, + }, + models: ["gpt-4"], + }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + await openSettingsEditor(user); + + expect(screen.getByLabelText("Estimated Output Tokens")).toHaveValue(512); + expect(screen.getByLabelText("Estimated Output Tokens Per Model")).toHaveValue('{"gpt-4":4096}'); + const keyValues = screen.queryAllByPlaceholderText("Key").map((input) => (input as HTMLInputElement).value); + expect(keyValues).toEqual(["department"]); + + fireEvent.change(screen.getByLabelText("Estimated Output Tokens"), { target: { value: "999" } }); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + const updateArg = vi.mocked(networking.teamUpdateCall).mock.calls[0][1]; + expect(updateArg.metadata.default_estimated_output_tokens).toBe(999); + expect(updateArg.metadata.default_estimated_output_tokens_per_model).toEqual({ "gpt-4": 4096 }); + }); + + it("omits the estimated output token settings when both controls are blank", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["gpt-4"] })); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + await openSettingsEditor(user); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + const updateArg = vi.mocked(networking.teamUpdateCall).mock.calls[0][1]; + expect(updateArg.metadata).not.toHaveProperty("default_estimated_output_tokens"); + expect(updateArg.metadata).not.toHaveProperty("default_estimated_output_tokens_per_model"); + }); + + it.each(["Internal User", "Admin Viewer", "org_admin"])( + "leaves both estimate controls read-only for %s and still resubmits the stored values", + async (userRole) => { + authState.userRole = userRole; + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + metadata: { + default_estimated_output_tokens: 512, + default_estimated_output_tokens_per_model: { "gpt-4": 4096 }, + }, + models: ["gpt-4"], + }), + ); + vi.mocked(networking.teamUpdateCall).mockResolvedValue({ data: {}, team_id: "123" } as any); + + renderWithProviders(); + await openSettingsEditor(user); + + expect(screen.getByLabelText("Estimated Output Tokens")).toBeDisabled(); + expect(screen.getByLabelText("Estimated Output Tokens Per Model")).toBeDisabled(); + + await user.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(networking.teamUpdateCall).toHaveBeenCalled(); + }); + + const updateArg = vi.mocked(networking.teamUpdateCall).mock.calls[0][1]; + expect(updateArg.metadata.default_estimated_output_tokens).toBe(512); + expect(updateArg.metadata.default_estimated_output_tokens_per_model).toEqual({ "gpt-4": 4096 }); + }, + ); + + it.each(["Admin", "proxy_admin"])("leaves both estimate controls editable for %s", async (userRole) => { + authState.userRole = userRole; + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ models: ["gpt-4"] })); + + renderWithProviders(); + await openSettingsEditor(user); + + expect(screen.getByLabelText("Estimated Output Tokens")).toBeEnabled(); + expect(screen.getByLabelText("Estimated Output Tokens Per Model")).toBeEnabled(); + }); + it("should keep declared keys as ordinary prefilled rows and submit the edited value", async () => { const user = userEvent.setup({ delay: null }); vi.mocked(useTeamMetadataSchema).mockReturnValue({ @@ -1057,6 +1173,34 @@ describe("TeamInfoView", () => { expect(screen.getByText("gpt-3.5-turbo")).toBeInTheDocument(); }); + it("should render the estimated output token settings in the overview and read-only settings views", async () => { + const user = userEvent.setup({ delay: null }); + vi.mocked(networking.teamInfoCall).mockResolvedValue( + createMockTeamData({ + metadata: { + default_estimated_output_tokens: 512, + default_estimated_output_tokens_per_model: { "gpt-4": 4096 }, + }, + }), + ); + + renderWithProviders(); + + await waitFor(() => { + const teamNameElements = screen.queryAllByText("Test Team"); + expect(teamNameElements.length).toBeGreaterThan(0); + }); + + await user.click(screen.getByRole("tab", { name: "Settings" })); + + await waitFor(() => { + expect(screen.getByText("Team Settings")).toBeInTheDocument(); + }); + + expect(screen.getAllByText("Estimated Output Tokens: 512")).toHaveLength(2); + expect(screen.getAllByText('Estimated Output Tokens Per Model: {"gpt-4":4096}')).toHaveLength(2); + }); + it("should show an empty state when the team has no model aliases", async () => { const user = userEvent.setup({ delay: null }); vi.mocked(networking.teamInfoCall).mockResolvedValue(createMockTeamData({ litellm_model_table: null })); diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index 6c6fdcaedd6..715bfdd83dc 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -58,6 +58,7 @@ import MCPServerSelector from "../mcp_server_management/MCPServerSelector"; import MCPToolPermissions from "../mcp_server_management/MCPToolPermissions"; import { ModelSelect } from "../ModelSelect/ModelSelect"; import NotificationsManager from "../molecules/notifications_manager"; +import { estimateRules, estimateTooltips } from "../templates/estimatedOutputTokens"; import ObjectPermissionsView from "../object_permissions_view"; import NumericalInput from "../shared/numerical_input"; import VectorStoreSelector from "../vector_store_management/VectorStoreSelector"; @@ -82,6 +83,8 @@ const UI_MANAGED_METADATA_KEYS: ReadonlySet = new Set([ "soft_budget_alerting_emails", "model_tpm_limit", "model_rpm_limit", + "default_estimated_output_tokens", + "default_estimated_output_tokens_per_model", "allowed_passthrough_routes", "guardrails", "opted_out_global_guardrails", @@ -207,6 +210,8 @@ const TeamInfoView: React.FC = ({ const routerSettingsRef = React.useRef(null); const [organization, setOrganization] = useState(null); const { userRole, userId } = useAuthorized(); + const canEditTeamEstimates = isProxyAdminRole(userRole); + const teamEstimateTooltip = estimateTooltips(canEditTeamEstimates, "team"); const { data: userOrganizations = [] } = useOrganizations(); const { data: teamMetadataSchemaFields = [], isLoading: isTeamMetadataSchemaLoading } = useTeamMetadataSchema(); const queryClient = useQueryClient(); @@ -472,6 +477,21 @@ const TeamInfoView: React.FC = ({ return v; }; + const estimatedOutputTokens = sanitizeNumeric(values.default_estimated_output_tokens); + + let estimatedOutputTokensPerModel: Record | undefined; + if (typeof values.default_estimated_output_tokens_per_model === "string") { + const trimmedEstimates = values.default_estimated_output_tokens_per_model.trim(); + if (trimmedEstimates.length > 0) { + try { + estimatedOutputTokensPerModel = JSON.parse(trimmedEstimates); + } catch (e) { + NotificationsManager.fromBackend("Invalid JSON in estimated output tokens per model"); + return; + } + } + } + const modelTpmLimit: Record = {}; const modelRpmLimit: Record = {}; for (const entry of (values.modelLimits ?? []) as { model?: string; tpm?: number; rpm?: number }[]) { @@ -512,6 +532,10 @@ const TeamInfoView: React.FC = ({ opted_out_global_guardrails: optedOutGlobalGuardrails, ...(values.logging_settings?.length > 0 ? { logging: values.logging_settings } : {}), disable_global_guardrails: killSwitchOnAtSave, + ...(estimatedOutputTokens !== null ? { default_estimated_output_tokens: Number(estimatedOutputTokens) } : {}), + ...(estimatedOutputTokensPerModel !== undefined + ? { default_estimated_output_tokens_per_model: estimatedOutputTokensPerModel } + : {}), soft_budget_alerting_emails: typeof values.soft_budget_alerting_emails === "string" ? values.soft_budget_alerting_emails @@ -772,6 +796,13 @@ const TeamInfoView: React.FC = ({ ); })()} + Estimated Output Tokens: {info.metadata?.default_estimated_output_tokens ?? "Default"} + + Estimated Output Tokens Per Model:{" "} + {info.metadata?.default_estimated_output_tokens_per_model + ? JSON.stringify(info.metadata.default_estimated_output_tokens_per_model) + : "Default"} + @@ -953,6 +984,11 @@ const TeamInfoView: React.FC = ({ soft_budget_alerting_emails: Array.isArray(info.metadata?.soft_budget_alerting_emails) ? info.metadata.soft_budget_alerting_emails.join(", ") : "", + default_estimated_output_tokens: info.metadata?.default_estimated_output_tokens, + default_estimated_output_tokens_per_model: info.metadata + ?.default_estimated_output_tokens_per_model + ? JSON.stringify(info.metadata.default_estimated_output_tokens_per_model) + : "", metadata: metadataObjectToPairs(info.metadata, UI_MANAGED_METADATA_KEYS), logging_settings: info.metadata?.logging || [], secret_manager_settings: info.metadata?.secret_manager_settings @@ -1211,6 +1247,24 @@ const TeamInfoView: React.FC = ({ + + + + + + + + = ({ ); })()} +
Estimated Output Tokens: {info.metadata?.default_estimated_output_tokens ?? "Default"}
+
+ Estimated Output Tokens Per Model:{" "} + {info.metadata?.default_estimated_output_tokens_per_model + ? JSON.stringify(info.metadata.default_estimated_output_tokens_per_model) + : "Default"} +
Team Budget diff --git a/ui/litellm-dashboard/src/components/templates/estimatedOutputTokens.test.ts b/ui/litellm-dashboard/src/components/templates/estimatedOutputTokens.test.ts new file mode 100644 index 00000000000..f857ec3efef --- /dev/null +++ b/ui/litellm-dashboard/src/components/templates/estimatedOutputTokens.test.ts @@ -0,0 +1,120 @@ +import { describe, expect, it } from "vitest"; + +import { estimateFields, estimateRules, withNormalizedEstimates } from "./estimatedOutputTokens"; + +const expectRejects = async (value: unknown) => + expect(estimateRules.perModel.validator(null, value)).rejects.toThrow(/JSON object of positive integers/); + +describe("estimateFields", () => { + it("renders a stored per-model map as editable JSON text", () => { + expect( + estimateFields({ + default_estimated_output_tokens: 2048, + default_estimated_output_tokens_per_model: { "gpt-4": 4096 }, + }), + ).toEqual({ + default_estimated_output_tokens: 2048, + default_estimated_output_tokens_per_model: '{"gpt-4":4096}', + }); + }); + + it("leaves the controls blank when metadata carries neither setting", () => { + expect(estimateFields({ unrelated: true })).toEqual({ + default_estimated_output_tokens: undefined, + default_estimated_output_tokens_per_model: "", + }); + }); + + it("tolerates absent metadata", () => { + expect(estimateFields(null).default_estimated_output_tokens_per_model).toBe(""); + expect(estimateFields(undefined).default_estimated_output_tokens_per_model).toBe(""); + }); +}); + +describe("estimateRules.perModel", () => { + it("accepts a blank control", async () => { + await expect(estimateRules.perModel.validator(null, "")).resolves.toBeUndefined(); + await expect(estimateRules.perModel.validator(null, " ")).resolves.toBeUndefined(); + await expect(estimateRules.perModel.validator(null, undefined)).resolves.toBeUndefined(); + }); + + it("accepts a per-model object", async () => { + await expect(estimateRules.perModel.validator(null, '{"gpt-4": 4096}')).resolves.toBeUndefined(); + }); + + it("rejects text that is not JSON", async () => { + await expectRejects("gpt-4: 4096"); + }); + + it("rejects JSON that is not an object, which the API would refuse", async () => { + await expectRejects("4096"); + await expectRejects('"gpt-4"'); + await expectRejects("[4096]"); + await expectRejects("null"); + }); + + it("rejects a per-model map whose values the runtime would ignore", async () => { + await expectRejects('{"gpt-4": -5}'); + await expectRejects('{"gpt-4": 0}'); + await expectRejects('{"gpt-4": 4.5}'); + await expectRejects('{"gpt-4": "4096"}'); + await expectRejects("{}"); + }); +}); + +describe("withNormalizedEstimates", () => { + it("coerces the numeric control and parses the per-model control without mutating the input", () => { + const values = { + default_estimated_output_tokens: "2048", + default_estimated_output_tokens_per_model: '{"gpt-4": 4096}', + other: "untouched", + }; + const before = { ...values }; + + expect(withNormalizedEstimates(values)).toEqual({ + default_estimated_output_tokens: 2048, + default_estimated_output_tokens_per_model: { "gpt-4": 4096 }, + other: "untouched", + }); + expect(values).toEqual(before); + }); + + it("drops blank controls so a save never sends an empty value", () => { + expect( + withNormalizedEstimates({ + default_estimated_output_tokens: "", + default_estimated_output_tokens_per_model: " ", + }), + ).toEqual({}); + }); + + it("drops each control independently", () => { + expect( + withNormalizedEstimates({ + default_estimated_output_tokens: 900, + default_estimated_output_tokens_per_model: "", + }), + ).toEqual({ default_estimated_output_tokens: 900 }); + }); + + it("drops a per-model map the API would reject rather than sending it", () => { + expect( + withNormalizedEstimates({ + default_estimated_output_tokens_per_model: '{"gpt-4": -5}', + }), + ).toEqual({}); + }); +}); + +describe("estimateRules.positive", () => { + it("accepts a blank control and a positive integer", async () => { + await expect(estimateRules.positive.validator(null, "")).resolves.toBeUndefined(); + await expect(estimateRules.positive.validator(null, 2048)).resolves.toBeUndefined(); + }); + + it("rejects values the runtime would ignore", async () => { + await expect(estimateRules.positive.validator(null, 0)).rejects.toThrow(/positive integer/); + await expect(estimateRules.positive.validator(null, -5)).rejects.toThrow(/positive integer/); + await expect(estimateRules.positive.validator(null, 12.5)).rejects.toThrow(/positive integer/); + }); +}); diff --git a/ui/litellm-dashboard/src/components/templates/estimatedOutputTokens.ts b/ui/litellm-dashboard/src/components/templates/estimatedOutputTokens.ts new file mode 100644 index 00000000000..842d7d2bb68 --- /dev/null +++ b/ui/litellm-dashboard/src/components/templates/estimatedOutputTokens.ts @@ -0,0 +1,77 @@ +type Metadata = Record | null | undefined; + +type FormValues = Record; + +const ESTIMATE_FIELD = "default_estimated_output_tokens"; +const PER_MODEL_FIELD = "default_estimated_output_tokens_per_model"; + +const INVALID_PER_MODEL_MESSAGE = 'Enter a JSON object of positive integers, e.g. {"gpt-4": 4096}'; + +const perModelEstimateToText = (value: unknown): string => + value != null && typeof value === "object" ? JSON.stringify(value) : ""; + +const isPositiveInteger = (value: unknown): boolean => + typeof value === "number" && Number.isInteger(value) && value > 0; + +const parsePerModelEstimates = (value: string): Record | null => { + let parsed: unknown; + try { + parsed = JSON.parse(value); + } catch { + return null; + } + if (parsed == null || typeof parsed !== "object" || Array.isArray(parsed)) return null; + const entries = Object.entries(parsed as Record); + if (entries.length === 0 || !entries.every(([, v]) => isPositiveInteger(v))) return null; + return Object.fromEntries(entries) as Record; +}; + +export const estimateFields = (metadata: Metadata) => ({ + [ESTIMATE_FIELD]: metadata?.[ESTIMATE_FIELD], + [PER_MODEL_FIELD]: perModelEstimateToText(metadata?.[PER_MODEL_FIELD]), +}); + +const ADMIN_ONLY_TOOLTIP = + "Only a proxy admin can change this. It sets how many output tokens the rate limiter reserves for a request " + + "that omits max_tokens, which is charged against the team and organization TPM windows."; + +export const estimateTooltips = (canEdit: boolean, entity: "key" | "team" = "key") => ({ + estimate: canEdit + ? `Expected output tokens reserved for TPM limiting when a request omits max_tokens. Overrides the built-in estimate for this ${entity}.` + : ADMIN_ONLY_TOOLTIP, + perModel: canEdit + ? `Per-model expected output tokens reserved for TPM limiting when a request omits max_tokens. Takes precedence over the ${entity}-wide estimate.` + : ADMIN_ONLY_TOOLTIP, +}); + +export const estimateRules = { + perModel: { + validator: (_: unknown, value: unknown) => { + if (typeof value !== "string" || value.trim() === "") return Promise.resolve(); + return parsePerModelEstimates(value) === null + ? Promise.reject(new Error(INVALID_PER_MODEL_MESSAGE)) + : Promise.resolve(); + }, + }, + positive: { + validator: (_: unknown, value: unknown) => { + if (value === "" || value === null || value === undefined) return Promise.resolve(); + return isPositiveInteger(Number(value)) + ? Promise.resolve() + : Promise.reject(new Error("Enter a positive integer")); + }, + }, +}; + +export const withNormalizedEstimates = (values: T): FormValues => { + const { [ESTIMATE_FIELD]: estimate, [PER_MODEL_FIELD]: perModel, ...rest } = values; + + const normalizedEstimate = estimate === "" || estimate === null || estimate === undefined ? null : Number(estimate); + const normalizedPerModel = typeof perModel === "string" ? parsePerModelEstimates(perModel) : null; + + return { + ...rest, + ...(normalizedEstimate === null ? {} : { [ESTIMATE_FIELD]: normalizedEstimate }), + ...(normalizedPerModel === null ? {} : { [PER_MODEL_FIELD]: normalizedPerModel }), + }; +}; diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts b/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts new file mode 100644 index 00000000000..ee8d9ebf685 --- /dev/null +++ b/ui/litellm-dashboard/src/components/templates/keyEditFieldNormalizers.ts @@ -0,0 +1,19 @@ +const WORD_FORM_BUDGET_DURATIONS: Record = { + hourly: "1h", + daily: "24h", + weekly: "7d", + monthly: "30d", +}; + +// Normalize any legacy word-form budget duration to the canonical value the dropdown uses +export const canonicalBudgetDuration = (duration: string | null | undefined): string | null => + duration ? WORD_FORM_BUDGET_DURATIONS[duration] ?? duration : null; + +// Determine the key_type display value from allowed_routes +export const keyTypeFromRoutes = (allowedRoutes: string[] | null | undefined): string => { + if (!allowedRoutes || allowedRoutes.length === 0) return "default"; + if (allowedRoutes.includes("llm_api_routes")) return "llm_api"; + if (allowedRoutes.includes("management_routes")) return "management"; + if (allowedRoutes.includes("info_routes")) return "read_only"; + return "default"; +}; diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index 9030a027928..3c75982612e 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -1348,4 +1348,135 @@ describe("KeyEditView", () => { }); }); }); + + describe("estimated output tokens", () => { + const renderEditView = ( + keyData: KeyResponse, + onSubmit: (values: any) => Promise, + userRole: string = "Admin", + ) => + renderWithProviders( + {}} + onSubmit={onSubmit} + accessToken={"test-token"} + userID={"test-user"} + userRole={userRole} + premiumUser={false} + />, + ); + + it("loads the estimates from key metadata and resubmits them unchanged", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderEditView( + { + ...MOCK_KEY_DATA, + metadata: { + ...MOCK_KEY_DATA.metadata, + default_estimated_output_tokens: 512, + default_estimated_output_tokens_per_model: { "gpt-4": 4096 }, + }, + }, + onSubmitMock, + ); + + await waitFor(() => { + expect(screen.getByLabelText("Estimated Output Tokens")).toHaveValue(512); + }); + expect(screen.getByLabelText("Estimated Output Tokens Per Model")).toHaveValue('{"gpt-4":4096}'); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + const callArgs = onSubmitMock.mock.calls[0][0]; + expect(callArgs.default_estimated_output_tokens).toBe(512); + expect(callArgs.default_estimated_output_tokens_per_model).toEqual({ "gpt-4": 4096 }); + }); + + it("submits edited estimates as a number and a parsed object", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderEditView(MOCK_KEY_DATA, onSubmitMock); + + await waitFor(() => { + expect(screen.getByLabelText("Estimated Output Tokens")).toBeInTheDocument(); + }); + + fireEvent.change(screen.getByLabelText("Estimated Output Tokens"), { target: { value: "2048" } }); + fireEvent.change(screen.getByLabelText("Estimated Output Tokens Per Model"), { + target: { value: '{"gpt-5": 8192}' }, + }); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + const callArgs = onSubmitMock.mock.calls[0][0]; + expect(callArgs.default_estimated_output_tokens).toBe(2048); + expect(callArgs.default_estimated_output_tokens_per_model).toEqual({ "gpt-5": 8192 }); + }); + + it("omits both estimates from the payload when the controls are blank", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderEditView(MOCK_KEY_DATA, onSubmitMock); + + await waitFor(() => { + expect(screen.getByLabelText("Estimated Output Tokens Per Model")).toHaveValue(""); + }); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + const callArgs = onSubmitMock.mock.calls[0][0]; + expect(callArgs).not.toHaveProperty("default_estimated_output_tokens"); + expect(callArgs).not.toHaveProperty("default_estimated_output_tokens_per_model"); + }); + + it.each(["Internal User", "Admin Viewer", "org_admin"])( + "leaves both controls read-only for %s and still resubmits the stored values", + async (userRole) => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + renderEditView( + { + ...MOCK_KEY_DATA, + metadata: { + ...MOCK_KEY_DATA.metadata, + default_estimated_output_tokens: 512, + default_estimated_output_tokens_per_model: { "gpt-4": 4096 }, + }, + }, + onSubmitMock, + userRole, + ); + + await waitFor(() => { + expect(screen.getByLabelText("Estimated Output Tokens")).toBeDisabled(); + }); + expect(screen.getByLabelText("Estimated Output Tokens Per Model")).toBeDisabled(); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalled(); + }); + const callArgs = onSubmitMock.mock.calls[0][0]; + expect(callArgs.default_estimated_output_tokens).toBe(512); + expect(callArgs.default_estimated_output_tokens_per_model).toEqual({ "gpt-4": 4096 }); + }, + ); + + it.each(["Admin", "proxy_admin"])("leaves both controls editable for %s", async (userRole) => { + renderEditView(MOCK_KEY_DATA, vi.fn().mockResolvedValue(undefined), userRole); + + await waitFor(() => { + expect(screen.getByLabelText("Estimated Output Tokens")).toBeEnabled(); + }); + expect(screen.getByLabelText("Estimated Output Tokens Per Model")).toBeEnabled(); + }); + }); }); diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index 2a02cf3edd0..b5538f0ca59 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -7,7 +7,7 @@ import { InfoCircleOutlined } from "@ant-design/icons"; import { TextInput, Button as TremorButton } from "@tremor/react"; import { Form, Input, Select, Switch, Tooltip } from "antd"; import { useEffect, useState } from "react"; -import { rolesWithWriteAccess } from "../../utils/roles"; +import { isProxyAdminRole, rolesWithWriteAccess } from "../../utils/roles"; import AgentSelector from "../agent_management/AgentSelector"; import AccessGroupSelector from "../common_components/AccessGroupSelector"; import BudgetDurationDropdown from "../common_components/budget_duration_dropdown"; @@ -17,6 +17,8 @@ import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSel import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"; import OrganizationDropdown from "../common_components/OrganizationDropdown"; import { extractLoggingSettings, formatMetadataForDisplay, stripTagsFromMetadata } from "../key_info_utils"; +import { estimateFields, estimateRules, estimateTooltips, withNormalizedEstimates } from "./estimatedOutputTokens"; +import { canonicalBudgetDuration, keyTypeFromRoutes } from "./keyEditFieldNormalizers"; import { BudgetFallbacksEditor } from "../key_team_helpers/BudgetFallbacksEditor"; import { BudgetWindowEntry, BudgetWindowsEditor } from "../key_team_helpers/BudgetWindowsEditor"; import { @@ -49,29 +51,6 @@ interface KeyEditViewProps { premiumUser?: boolean; } -// Add this helper function - -// Helper function to determine key_type display value from allowed_routes -const getKeyTypeFromRoutes = (allowedRoutes: string[] | null | undefined): string => { - if (!allowedRoutes || allowedRoutes.length === 0) { - return "default"; - } - - if (allowedRoutes.includes("llm_api_routes")) { - return "llm_api"; - } - - if (allowedRoutes.includes("management_routes")) { - return "management"; - } - - if (allowedRoutes.includes("info_routes")) { - return "read_only"; - } - - return "default"; -}; - export function KeyEditView({ keyData, onCancel, @@ -83,6 +62,8 @@ export function KeyEditView({ premiumUser = false, }: KeyEditViewProps) { const canEditGuardrails = premiumUser || (userRole != null && rolesWithWriteAccess.includes(userRole)); + const canEditEstimates = userRole != null && isProxyAdminRole(userRole); + const estimateTooltip = estimateTooltips(canEditEstimates); const [form] = Form.useForm(); const [promptsList, setPromptsList] = useState([]); const [tagsList, setTagsList] = useState>({}); @@ -157,27 +138,16 @@ export function KeyEditView({ form.setFieldValue("disabled_callbacks", disabledCallbacks); }, [form, disabledCallbacks]); - // Normalize any legacy word-form budget duration to the canonical value the dropdown uses - const getBudgetDuration = (duration: string | null) => { - if (!duration) return null; - const wordToCanonical: Record = { - hourly: "1h", - daily: "24h", - weekly: "7d", - monthly: "30d", - }; - return wordToCanonical[duration] ?? duration; - }; - // Set initial form values const initialValues = { ...keyData, token: keyData.token || keyData.token_id, - budget_duration: getBudgetDuration(keyData.budget_duration), + budget_duration: canonicalBudgetDuration(keyData.budget_duration), metadata: formatMetadataForDisplay(stripTagsFromMetadata(keyData.metadata)), guardrails: keyData.metadata?.guardrails, disable_global_guardrails: keyData.metadata?.disable_global_guardrails || false, throttle_on_budget_exceeded: keyData.metadata?.throttle_on_budget_exceeded || false, + ...estimateFields(keyData.metadata), prompts: keyData.metadata?.prompts, tags: keyData.metadata?.tags, vector_stores: keyData.object_permission?.vector_stores || [], @@ -208,7 +178,7 @@ export function KeyEditView({ form.setFieldsValue({ ...keyData, token: keyData.token || keyData.token_id, - budget_duration: getBudgetDuration(keyData.budget_duration), + budget_duration: canonicalBudgetDuration(keyData.budget_duration), metadata: formatMetadataForDisplay(stripTagsFromMetadata(keyData.metadata)), guardrails: keyData.metadata?.guardrails, disable_global_guardrails: keyData.metadata?.disable_global_guardrails || false, @@ -222,6 +192,7 @@ export function KeyEditView({ }, mcp_tool_permissions: keyData.object_permission?.mcp_tool_permissions || {}, throttle_on_budget_exceeded: keyData.metadata?.throttle_on_budget_exceeded || false, + ...estimateFields(keyData.metadata), logging_settings: extractLoggingSettings(keyData.metadata), disabled_callbacks: Array.isArray(keyData.metadata?.litellm_disabled_callbacks) ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks) @@ -339,7 +310,7 @@ export function KeyEditView({ values.budget_fallbacks = {}; } - await onSubmit(values); + await onSubmit(withNormalizedEstimates(values)); } finally { setIsKeySaving(false); } @@ -418,7 +389,7 @@ export function KeyEditView({ > {({ getFieldValue, setFieldValue }) => { const allowedRoutesValue = getFieldValue("allowed_routes") || ""; - // Convert string to array for getKeyTypeFromRoutes + // Convert string to array for keyTypeFromRoutes const allowedRoutes = typeof allowedRoutesValue === "string" && allowedRoutesValue.trim() !== "" ? allowedRoutesValue @@ -426,7 +397,7 @@ export function KeyEditView({ .map((r: string) => r.trim()) .filter((r: string) => r.length > 0) : []; - const keyTypeValue = getKeyTypeFromRoutes(allowedRoutes); + const keyTypeValue = keyTypeFromRoutes(allowedRoutes); return (