diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260913000000_add_tpd_limit/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260913000000_add_tpd_limit/migration.sql new file mode 100644 index 00000000000..cdf8f4975c1 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260913000000_add_tpd_limit/migration.sql @@ -0,0 +1,14 @@ +-- AlterTable +ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT; + +-- AlterTable +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT; + +-- AlterTable +ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT; + +-- AlterTable +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT; + +-- AlterTable +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index dd7967aafe3..8072df5aa5b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -17,6 +17,7 @@ model LiteLLM_BudgetTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? model_max_budget Json? budget_duration String? budget_reset_at DateTime? @@ -133,6 +134,7 @@ model LiteLLM_TeamTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? blocked Boolean @default(false) @@ -203,6 +205,7 @@ model LiteLLM_DeletedTeamTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? blocked Boolean @default(false) @@ -438,6 +441,7 @@ model LiteLLM_VerificationToken { blocked Boolean? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? max_budget Float? budget_duration String? budget_reset_at DateTime? @@ -534,6 +538,7 @@ model LiteLLM_DeletedVerificationToken { blocked Boolean? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? max_budget Float? budget_duration String? budget_reset_at DateTime? diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 0e8b8136c19..39600328074 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -21,6 +21,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator from litellm.a2a_protocol.utils import A2ARequestUtils from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -507,7 +508,7 @@ async def asend_message( prompt_tokens, completion_tokens, _, - ) = A2ARequestUtils.calculate_usage_from_request_response( + ) = await asyncify(A2ARequestUtils.calculate_usage_from_request_response)( request=request, response_dict=response_dict, ) diff --git a/litellm/a2a_protocol/streaming_iterator.py b/litellm/a2a_protocol/streaming_iterator.py index 67db8e905e3..d936caeb75e 100644 --- a/litellm/a2a_protocol/streaming_iterator.py +++ b/litellm/a2a_protocol/streaming_iterator.py @@ -11,6 +11,7 @@ import litellm from litellm._logging import verbose_logger from litellm.a2a_protocol.cost_calculator import A2ACostCalculator from litellm.a2a_protocol.utils import A2ARequestUtils +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj if TYPE_CHECKING: @@ -99,11 +100,11 @@ class A2AStreamingIterator: # Calculate tokens from collected text input_message: Final = A2ARequestUtils.get_input_message_from_request(self.request) input_text: Final = A2ARequestUtils.extract_text_from_message(input_message) - prompt_tokens: Final = A2ARequestUtils.count_tokens(input_text) + prompt_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(input_text) # Use the last (most complete) text from chunks output_text: Final = self.collected_text_parts[-1] if self.collected_text_parts else "" - completion_tokens: Final = A2ARequestUtils.count_tokens(output_text) + completion_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(output_text) total_tokens: Final = prompt_tokens + completion_tokens diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index c5876e993d3..058cc8a1579 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -21,6 +21,7 @@ from litellm.constants import ( QDRANT_VECTOR_SIZE, SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS, ) +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -255,7 +256,7 @@ class QdrantSemanticCache(BaseCache): llm_router = None router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) - embedding_input: Final = self._embedding_input(prompt, router) + embedding_input: Final = await asyncify(self._embedding_input)(prompt, router) embedding_call: Final = ( router.aembedding( model=self.embedding_model, diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index 9a70bfc1418..d4c815e15b7 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -522,7 +523,7 @@ class RedisSemanticCache(BaseCache): llm_router = None router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) - embedding_input: Final = self._embedding_input(prompt, router) + embedding_input: Final = await asyncify(self._embedding_input)(prompt, router) embedding_call: Final = ( router.aembedding( model=self.embedding_model, diff --git a/litellm/constants.py b/litellm/constants.py index 09442d6151e..ba5ec73d435 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1566,6 +1566,8 @@ BASE_MCP_ROUTE: Final = "/mcp" BATCH_STATUS_POLL_INTERVAL_SECONDS: Final = int(os.getenv("BATCH_STATUS_POLL_INTERVAL_SECONDS", 3600)) # 1 hour BATCH_STATUS_POLL_MAX_ATTEMPTS: Final = int(os.getenv("BATCH_STATUS_POLL_MAX_ATTEMPTS", 24)) # for 24 hours +BATCH_TPD_WINDOW_SECONDS: Final = 86400 +BATCH_TPD_DESCRIPTOR_SUFFIX: Final = "_tpd" HEALTH_CHECK_TIMEOUT_SECONDS: Final = int(os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60)) # 60 seconds _background_health_check_max_tokens_env: Final = os.getenv("BACKGROUND_HEALTH_CHECK_MAX_TOKENS") diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index 1be7a01ba3a..5352ce6b6a0 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -15,6 +15,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.compression import compress from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.types.integrations.compression_interception import ( CompressionInterceptionConfig, CompressionSavingsMetadata, @@ -153,7 +154,7 @@ class CompressionInterceptionLogger(CustomLogger): self._prune_expired_cache() - compressed: Final = compress( + compressed: Final = await asyncify(compress)( messages=messages, model=model, call_type=CallTypes.anthropic_messages, diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4ddb9ce5b8e..abac624d5ec 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1991,6 +1991,12 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost + def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None: + """Bill a fully streamed response on the failure log when a post-call hook rejects it.""" + usage: Final = getattr(assembled, "usage", None) + if isinstance(usage, Usage): + self.record_partial_usage_for_failure(usage, self._response_cost_calculator(result=assembled) or 0.0) + async def dispatch_failure_handlers( self, exception: Exception, diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6a5a8832cc6..766d60ad180 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -19,6 +19,7 @@ from typing_extensions import NotRequired, TypedDict import litellm from litellm import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.model_response_utils import ( is_model_response_stream_empty, ) @@ -2247,7 +2248,7 @@ class CustomStreamWrapper: if self.sent_last_chunk is True: # log the final chunk with accurate streaming values try: - complete_streaming_response = litellm.stream_chunk_builder( + complete_streaming_response = await asyncify(litellm.stream_chunk_builder)( chunks=self.chunks, messages=self.messages, logging_obj=self.logging_obj, diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py index 902808647c0..ad33e5e0592 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py @@ -5,6 +5,7 @@ from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Final, TypeAlias from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.types.llms.anthropic import AppliedEdit from .constants import CLEAR_TOOL_USES_EDIT_TYPE, COMPACT_EDIT_TYPE @@ -82,9 +83,9 @@ async def apply_context_management( """Run edits in order; return a single ``PolyfillResult``. The dispatcher is async so async editors (``compact_20260112``) can - ``await`` the configured summarization model. Sync editors are called - inline — ``inspect.iscoroutinefunction`` decides how each editor is - invoked. + ``await`` the configured summarization model. Sync editors run in a + worker thread so their token counts stay off the event loop; + ``inspect.iscoroutinefunction`` decides how each editor is invoked. """ edits: Final = _normalize_spec(context_management_spec) if not edits: @@ -121,7 +122,7 @@ async def apply_context_management( user_api_key_auth=user_api_key_auth, ) if editor_is_async - else editor( + else await asyncify(editor)( model=model, messages=current_messages, tools=tools, diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index 050ab67c86c..fb6a1c40253 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -20,6 +20,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.types.llms.anthropic import ( AppliedEdit, CompactionBlock, @@ -1157,7 +1158,7 @@ async def apply_compact_20260112( # Phase B: threshold check. try: - current_tokens = _count_effective_tokens( + current_tokens = await asyncify(_count_effective_tokens)( model=model, effective_messages=effective_messages, # ``augmented_system`` already carries the prior compaction summary diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 7d01aee5d98..98c5c6d6d4e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -678,7 +678,7 @@ class BaseAnthropicMessagesStreamingIterator: """ from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler - PassThroughStreamingHandler.schedule_stream_failure_logging( + await PassThroughStreamingHandler.schedule_stream_failure_logging( litellm_logging_obj=self.litellm_logging_obj, endpoint_type=EndpointType.ANTHROPIC, request_body=self.request_body, diff --git a/litellm/main.py b/litellm/main.py index 8b64c29d8c4..9fbc5881b4f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -67,7 +67,7 @@ from litellm.constants import ( ) from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.asyncify import asyncify, run_async_function from litellm.litellm_core_utils.audio_utils.utils import ( calculate_request_duration, get_audio_file_for_health_check, @@ -9127,7 +9127,7 @@ async def acount_tokens( fallback_messages = messages or [] if system and fallback_messages: fallback_messages = [{"role": "system", "content": system}] + fallback_messages - local_count: Final = litellm.token_counter( + local_count: Final = await asyncify(litellm.token_counter)( model=model, messages=fallback_messages, tools=tools, diff --git a/litellm/models/budget.py b/litellm/models/budget.py index 335800a49a8..125ce739d6a 100644 --- a/litellm/models/budget.py +++ b/litellm/models/budget.py @@ -26,6 +26,7 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): max_parallel_requests: int | None = None tpm_limit: int | None = None rpm_limit: int | None = None + tpd_limit: int | None = None model_max_budget: dict | None = None budget_duration: str | None = None allowed_models: list[str] | None = None # per-member model scope; empty = inherit team models diff --git a/litellm/models/team.py b/litellm/models/team.py index da526515e6e..8edf10703b1 100644 --- a/litellm/models/team.py +++ b/litellm/models/team.py @@ -71,6 +71,7 @@ class TeamBase(LiteLLMPydanticObjectBase): metadata: dict | None = None tpm_limit: int | None = None rpm_limit: int | None = None + tpd_limit: int | None = None max_budget: float | None = None soft_budget: float | None = None budget_duration: str | None = None diff --git a/litellm/models/verification_token.py b/litellm/models/verification_token.py index fec3caec457..06ff877a41a 100644 --- a/litellm/models/verification_token.py +++ b/litellm/models/verification_token.py @@ -31,6 +31,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): metadata: dict = {} tpm_limit: int | None = None rpm_limit: int | None = None + tpd_limit: int | None = None budget_duration: str | None = None budget_reset_at: datetime | None = None allowed_cache_controls: list | None = [] diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 993f705a918..afc2150934e 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1206,6 +1206,7 @@ class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase): class KeyRequestBase(GenerateRequestBase): key: str | None = None + tpd_limit: int | 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 @@ -1891,6 +1892,9 @@ class BudgetNewRequest(LiteLLMPydanticObjectBase): ) tpm_limit: int | None = Field(default=None, description="Max tokens per minute, allowed for this budget id.") rpm_limit: int | None = Field(default=None, description="Max requests per minute, allowed for this budget id.") + tpd_limit: int | None = Field( + default=None, description="Max tokens per day, charged by batch submissions, allowed for this budget id." + ) budget_duration: str | None = Field( default=None, description="Max duration budget should be set for (e.g. '1hr', '1d', '28d')", @@ -2067,6 +2071,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): metadata: dict | None = None tpm_limit: int | None = None rpm_limit: int | None = None + tpd_limit: int | None = None max_budget: float | None = None soft_budget: float | None = None models: list | None = None @@ -3022,6 +3027,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): team_alias: str | None = None team_tpm_limit: int | None = None team_rpm_limit: int | None = None + team_tpd_limit: int | None = None team_max_budget: float | None = None team_soft_budget: float | None = None team_models: list = [] @@ -3041,6 +3047,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): end_user_id: str | None = None end_user_tpm_limit: int | None = None end_user_rpm_limit: int | None = None + end_user_tpd_limit: int | None = None end_user_max_budget: float | None = None end_user_model_max_budget: dict | None = None diff --git a/litellm/proxy/auth/team_grants.py b/litellm/proxy/auth/team_grants.py index 1196011dcdd..0421659c331 100644 --- a/litellm/proxy/auth/team_grants.py +++ b/litellm/proxy/auth/team_grants.py @@ -56,6 +56,7 @@ class TeamGrants(TypedDict, total=False): team_alias: ReadOnly[str | None] team_tpm_limit: ReadOnly[int | None] team_rpm_limit: ReadOnly[int | None] + team_tpd_limit: ReadOnly[int | None] team_max_budget: ReadOnly[float | None] team_soft_budget: ReadOnly[float | None] team_spend: ReadOnly[float | None] @@ -97,6 +98,7 @@ def team_grants( team_alias=team_object.team_alias, team_tpm_limit=team_object.tpm_limit, team_rpm_limit=team_object.rpm_limit, + team_tpd_limit=team_object.tpd_limit, team_max_budget=team_object.max_budget, team_soft_budget=team_object.soft_budget, team_spend=team_object.spend, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 0618875be55..7d62baf39a8 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -537,6 +537,9 @@ def _apply_budget_limits_to_end_user_params( if budget_info.rpm_limit is not None: end_user_params["end_user_rpm_limit"] = budget_info.rpm_limit + if budget_info.tpd_limit is not None: + end_user_params["end_user_tpd_limit"] = budget_info.tpd_limit + if budget_info.max_budget is not None: end_user_params["end_user_max_budget"] = budget_info.max_budget @@ -621,6 +624,8 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use valid_token.end_user_tpm_limit = end_user_params["end_user_tpm_limit"] if end_user_params.get("end_user_rpm_limit") is not None: valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"] + if end_user_params.get("end_user_tpd_limit") is not None: + valid_token.end_user_tpd_limit = end_user_params["end_user_tpd_limit"] if end_user_params.get("allowed_model_region") is not None: valid_token.allowed_model_region = end_user_params["allowed_model_region"] if end_user_params.get("end_user_model_max_budget") is not None: @@ -2026,6 +2031,7 @@ async def _user_api_key_auth_builder( valid_token.end_user_id = end_user_params.get("end_user_id") valid_token.end_user_tpm_limit = end_user_params.get("end_user_tpm_limit") valid_token.end_user_rpm_limit = end_user_params.get("end_user_rpm_limit") + valid_token.end_user_tpd_limit = end_user_params.get("end_user_tpd_limit") valid_token.allowed_model_region = end_user_params.get("allowed_model_region") if valid_token is not None: @@ -2302,6 +2308,7 @@ async def _user_api_key_auth_builder( spend=valid_token.team_spend, tpm_limit=valid_token.team_tpm_limit, rpm_limit=valid_token.team_rpm_limit, + tpd_limit=valid_token.team_tpd_limit, blocked=valid_token.team_blocked, models=token_team_models, metadata=valid_token.team_metadata, @@ -2455,6 +2462,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached spend=valid_token.team_spend, tpm_limit=valid_token.team_tpm_limit, rpm_limit=valid_token.team_rpm_limit, + tpd_limit=valid_token.team_tpd_limit, blocked=valid_token.team_blocked, models=token_team_models, metadata=valid_token.team_metadata, diff --git a/litellm/proxy/db/create_views.py b/litellm/proxy/db/create_views.py index 10daeee4e7b..d3f3de730ab 100644 --- a/litellm/proxy/db/create_views.py +++ b/litellm/proxy/db/create_views.py @@ -80,6 +80,7 @@ async def create_missing_views(db: SupportsRawQueries) -> None: t.max_budget AS team_max_budget, t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, + t.tpd_limit AS team_tpd_limit, p.project_alias AS project_alias FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index dcd34a1d9cb..ab6e10ca76b 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -18,12 +18,13 @@ Quick summary: """ import json -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence +from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, TypeAlias from fastapi import HTTPException -from pydantic import BaseModel, Field, TypeAdapter +from pydantic import BaseModel, Field, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -33,6 +34,7 @@ from litellm.batches.batch_utils import ( _extract_file_access_credentials, _iter_batch_input_lines, ) +from litellm.constants import BATCH_TPD_DESCRIPTOR_SUFFIX, BATCH_TPD_WINDOW_SECONDS from litellm.exceptions import RateLimitErrorCategory from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ( @@ -55,6 +57,7 @@ from litellm.proxy.hooks.batch_enqueued_tokens import ( from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY, + ReservationAwareIncrementOperation, get_or_create_request_stash, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit @@ -92,6 +95,7 @@ else: _BATCH_BODY_ADAPTER: Final = TypeAdapter(dict[str, object]) +_WINDOW_START_ADAPTER: Final[TypeAdapter[int | float | str | None]] = TypeAdapter(int | float | str | None) IncrementAmounts: TypeAlias = dict[Literal["requests", "tokens"], int] @@ -128,6 +132,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): self, internal_usage_cache: InternalUsageCache, parallel_request_limiter: ParallelRequestLimiter, + time_provider: Callable[[], datetime] | None = None, ): """ Initialize the batch rate limiter. @@ -138,9 +143,11 @@ class _PROXY_BatchRateLimiter(CustomLogger): Args: internal_usage_cache: Cache for storing rate limit data (auto-injected) parallel_request_limiter: Existing rate limiter to integrate with (needs custom injection) + time_provider: Clock used for rate limit reset times (defaults to ``datetime.now``) """ self.internal_usage_cache = internal_usage_cache self.parallel_request_limiter = parallel_request_limiter + self._time_provider: Final = time_provider or datetime.now self._warned_unsupported_model_skip = False def _get_file_bound_batch_model(self, data: dict) -> str | None: @@ -236,14 +243,48 @@ class _PROXY_BatchRateLimiter(CustomLogger): file-bound/top-level routing model this function resolves. Charging project quotas here would let a caller bind the file to a model without a quota while rows execute against a quota-limited model. + + Scopes with a ``tpd_limit`` (key, team, end user) are charged against a + daily token descriptor instead of their per-minute RPM/TPM descriptor, + because a batch's rows are scheduled by the provider and never share a + minute with the submission. The daily descriptor uses its own key so + its 24h window never collides with the online limiter's counters. """ - return self.parallel_request_limiter._create_rate_limit_descriptors( + descriptors: Final = self.parallel_request_limiter._create_rate_limit_descriptors( user_api_key_dict=user_api_key_dict, data=data, rpm_limit_type=None, tpm_limit_type=None, model_has_failures=False, ) + tpd_limits: Final[Mapping[str, tuple[str, int]]] = MappingProxyType( + { + key: (value, limit) + for key, value, limit in ( + ("api_key", user_api_key_dict.api_key, user_api_key_dict.tpd_limit), + ("team", user_api_key_dict.team_id, user_api_key_dict.team_tpd_limit), + ("end_user", user_api_key_dict.end_user_id, user_api_key_dict.end_user_tpd_limit), + ) + if value and limit is not None + } + ) + if not tpd_limits: + return descriptors + return [ + *(d for d in descriptors if d["key"] not in tpd_limits), + *( + RateLimitDescriptor( + key=f"{key}{BATCH_TPD_DESCRIPTOR_SUFFIX}", + value=value, + rate_limit={ + "requests_per_unit": None, + "tokens_per_unit": limit, + "window_size": BATCH_TPD_WINDOW_SECONDS, + }, + ) + for key, (value, limit) in tpd_limits.items() + ), + ] @staticmethod def _project_has_any_io_token_limits(user_api_key_dict: UserAPIKeyAuth) -> bool: @@ -583,9 +624,14 @@ class _PROXY_BatchRateLimiter(CustomLogger): batch_usage: BatchFileUsage, limit_type: str, requested_model: str | None = None, + window_start: int | None = None, ) -> NoReturn: - """Raise :class:`ProxyRateLimitError` (a 429) for batch rate limit exceeded.""" - from datetime import datetime + """Raise :class:`ProxyRateLimitError` (a 429) for batch rate limit exceeded. + + ``window_start`` is the active counter window's start (unix seconds) when + known, so the reset time reflects that window's actual end rather than a + full window from now. + """ # Find the descriptor for this status. Matching on (key, value) is # required, not key alone: a batch can carry several project ITPM/OTPM @@ -609,9 +655,12 @@ class _PROXY_BatchRateLimiter(CustomLogger): descriptors[descriptor_index] if descriptors else {"key": "", "value": "", "rate_limit": None} ) - now: Final = datetime.now().timestamp() - window_size: Final = self.parallel_request_limiter.window_size - reset_time: Final = now + window_size + now: Final = self._time_provider().timestamp() + window_size: Final = (descriptor.get("rate_limit") or {}).get( + "window_size" + ) or self.parallel_request_limiter.window_size + reset_time: Final = now + window_size if window_start is None else window_start + window_size + retry_after: Final = max(0, int(reset_time - now)) reset_time_formatted: Final = datetime.fromtimestamp(reset_time).strftime("%Y-%m-%d %H:%M:%S UTC") remaining_display: Final = max(0, status["limit_remaining"]) @@ -643,10 +692,13 @@ class _PROXY_BatchRateLimiter(CustomLogger): if descriptor.get("key") == PROJECT_ITPM_DESCRIPTOR_KEY else batch_usage.total_tokens ) + token_limit_label: Final = ( + "TPD" if descriptor.get("key", "").endswith(BATCH_TPD_DESCRIPTOR_SUFFIX) else "TPM" + ) detail = ( f"Batch rate limit exceeded for {descriptor.get('key', 'unknown')}: {descriptor.get('value', 'unknown')}. " f"Batch contains {batch_token_count} tokens but only {remaining_display} tokens remaining " - f"out of {current_limit} TPM limit. " + f"out of {current_limit} {token_limit_label} limit. " f"Limit resets at: {reset_time_formatted}" ) @@ -654,7 +706,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): raise ProxyRateLimitError( detail=detail, headers={ - "retry-after": str(window_size), + "retry-after": str(retry_after), "rate_limit_type": limit_type, "reset_at": reset_time_formatted, }, @@ -712,6 +764,8 @@ class _PROXY_BatchRateLimiter(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) + stash: Final = get_or_create_request_stash() + stash.batch_tpd_refund_ops = () if rate_limit_response["overall_code"] == "OVER_LIMIT": requested_model: Final = data.get("model") if data else None for status in rate_limit_response["statuses"]: @@ -722,8 +776,70 @@ class _PROXY_BatchRateLimiter(CustomLogger): batch_usage, status["rate_limit_type"], requested_model=requested_model, + window_start=await self._read_tpd_window_start( + status=status, parent_otel_span=user_api_key_dict.parent_otel_span + ), ) + stash.batch_tpd_refund_ops = self._build_tpd_refund_ops( + descriptors=descriptors, + tokens=batch_usage.total_tokens, + reservation_windows=rate_limit_response.get("reservation_windows", frozenset()), + ) + + async def _read_tpd_window_start(self, status: "RateLimitStatus", parent_otel_span: "Span | None") -> int | None: + descriptor_key: Final = status.get("descriptor_key") or "" + if not descriptor_key.endswith(BATCH_TPD_DESCRIPTOR_SUFFIX): + return None + try: + window_start: Final = _WINDOW_START_ADAPTER.validate_python( + await self.parallel_request_limiter.internal_usage_cache.async_get_cache( + key=f"{{{descriptor_key}:{status.get('descriptor_value') or ''}}}:window", + litellm_parent_otel_span=parent_otel_span, + ), + strict=True, + ) + return None if window_start is None else int(float(window_start)) + except (ValidationError, ValueError): + return None + + def _build_tpd_refund_ops( + self, + descriptors: Sequence["RateLimitDescriptor"], + tokens: int, + reservation_windows: frozenset[tuple[str, str, Literal["redis", "local"]]], + ) -> tuple[ReservationAwareIncrementOperation, ...]: + """Refund operations for the daily token counters this batch charged. + + The v3 limiter's failure hook applies them when the submission fails + after the counters were incremented. Each operation carries the window + identity the charge landed in, so the refund is skipped once that + window has rolled over. + """ + if tokens <= 0 or not reservation_windows: + return () + tpd_descriptors_by_counter: Final[Mapping[str, RateLimitDescriptor]] = MappingProxyType( + { + self.parallel_request_limiter.create_rate_limit_keys( + descriptor["key"], descriptor["value"], "tokens" + ): descriptor + for descriptor in descriptors + if descriptor["key"].endswith(BATCH_TPD_DESCRIPTOR_SUFFIX) + } + ) + return tuple( + ReservationAwareIncrementOperation( + key=counter_key, + increment_value=-tokens, + ttl=BATCH_TPD_WINDOW_SECONDS, + window_key=f"{{{descriptor['key']}:{descriptor['value']}}}:window", + expected_window_start=window_start, + reservation_backend=backend, + ) + for counter_key, window_start, backend in sorted(reservation_windows) + if (descriptor := tpd_descriptors_by_counter.get(counter_key)) is not None + ) + async def count_input_file_usage( self, file_id: str, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index a34dc99e472..8ca4124521a 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -396,6 +396,8 @@ CacheCounterValue: TypeAlias = int | float | str | bytes CacheCounterValues: TypeAlias = Sequence[CacheCounterValue | None] +ReservationWindowIdentity: TypeAlias = tuple[str, str, Literal["redis", "local"]] + ParallelGaugeCacheValue: TypeAlias = dict[str, object] | int | float | str | bytes @@ -542,6 +544,7 @@ class RequestRateLimiterStash: default_factory=frozenset ) batch_enqueued_reservation: BatchEnqueuedTokenReservation | None = None + batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = () reservation_released: bool = False @@ -683,6 +686,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self._batch_rate_limiter = _PROXY_BatchRateLimiter( internal_usage_cache=self.internal_usage_cache, parallel_request_limiter=self, + time_provider=self._time_provider, ) except Exception as e: verbose_proxy_logger.debug("Could not load batch rate limiter: %s", e) @@ -1823,6 +1827,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) applied: Final[list[list[AtomicCounterMeta]]] = [] statuses: Final[list[RateLimitStatus]] = [] + reservation_windows: Final[set[ReservationWindowIdentity]] = set() # mutable-ok: filled by the group loop raw: list[CacheCounterValue] for _idx, (keys, args, meta) in enumerate(descriptor_groups): @@ -1860,11 +1865,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return response applied.append(meta) statuses.extend(response["statuses"]) + reservation_windows.update(response.get("reservation_windows", frozenset())) return RateLimitResponse( overall_code="OK", statuses=statuses, - reservation_windows=frozenset(), + reservation_windows=frozenset(reservation_windows), ) async def _refund_applied_descriptor_groups( @@ -4824,6 +4830,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) stash.batch_enqueued_reservation = None + if stash.batch_tpd_refund_ops: + await self.async_increment_reservation_aware_tokens( + pipeline_operations=stash.batch_tpd_refund_ops, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + stash.batch_tpd_refund_ops = () + if stash.reservation_released: return reserved_tokens: Final = stash.reserved_tokens diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 81a607aaa43..e16ea4a812e 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -52,6 +52,7 @@ async def new_budget( - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. - tpm_limit: Optional[int] - The tokens per minute limit for the budget. - rpm_limit: Optional[int] - The requests per minute limit for the budget. + - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit. - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} - budget_reset_at: Optional[datetime] - Datetime when the initial budget is reset. Default is now. """ @@ -135,6 +136,7 @@ async def update_budget( - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. - tpm_limit: Optional[int] - The tokens per minute limit for the budget. - rpm_limit: Optional[int] - The requests per minute limit for the budget. + - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit. - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} - budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset. """ @@ -272,6 +274,7 @@ async def budget_settings( "max_parallel_requests": {"type": "Integer"}, "tpm_limit": {"type": "Integer"}, "rpm_limit": {"type": "Integer"}, + "tpd_limit": {"type": "Integer"}, "budget_duration": {"type": "String"}, "max_budget": {"type": "Float"}, "soft_budget": {"type": "Float"}, diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index d2d87331d55..b35bc01b4d0 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -335,6 +335,7 @@ async def new_end_user( - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - tpm_limit: Optional[int] - [Not Implemented Yet] Specify tpm limit for a given customer (Tokens per minute) - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute) + - tpd_limit: Optional[int] - Specify tpd limit for a given customer (Tokens per day). Batch submissions are charged against it instead of tpm_limit/rpm_limit - model_max_budget: Optional[dict] - [Not Implemented Yet] Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}} - max_parallel_requests: Optional[int] - [Not Implemented Yet] Specify max parallel requests for a given customer. - soft_budget: Optional[float] - [Not Implemented Yet] Get alerts when customer crosses given budget, doesn't block requests. diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 62019f462a2..ee8ae66ea11 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1078,7 +1078,9 @@ async def validate_team_id_used_in_service_account_request( return True -_BUDGET_NUMERIC_KEYS = frozenset(["max_budget", "soft_budget", "max_parallel_requests", "tpm_limit", "rpm_limit"]) +_BUDGET_NUMERIC_KEYS = frozenset( + ["max_budget", "soft_budget", "max_parallel_requests", "tpm_limit", "rpm_limit", "tpd_limit"] +) def _enforce_upperbound_key_params( @@ -1957,6 +1959,7 @@ async def generate_key_fn( - blocked: Optional[bool] - Whether the key is blocked. - rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute) - tpm_limit: Optional[int] - Specify tpm limit for a given key (Tokens per minute) + - tpd_limit: Optional[int] - Specify tpd limit for a given key (Tokens per day). Charged by batch submissions instead of tpm_limit/rpm_limit. - soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached. - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. @@ -2163,6 +2166,7 @@ async def generate_service_account_key_fn( - blocked: Optional[bool] - Whether the key is blocked. - rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute) - tpm_limit: Optional[int] - Specify tpm limit for a given key (Tokens per minute) + - tpd_limit: Optional[int] - Specify tpd limit for a given key (Tokens per day). Charged by batch submissions instead of tpm_limit/rpm_limit. - soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached. - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests) @@ -3192,6 +3196,7 @@ async def update_key_fn( - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"} - tpm_limit: Optional[int] - Tokens per minute limit - rpm_limit: Optional[int] - Requests per minute limit + - tpd_limit: Optional[int] - Tokens per day limit, charged by batch submissions instead of tpm_limit/rpm_limit - model_rpm_limit: Optional[dict] - Model-specific RPM limits {"gpt-4": 100, "claude-v1": 200} - 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. @@ -4355,6 +4360,7 @@ async def generate_key_helper_fn( metadata: dict | None = {}, tpm_limit: int | None = None, rpm_limit: int | None = None, + tpd_limit: int | None = None, query_type: Literal["insert_data", "update_data"] = "insert_data", update_key_values: dict | None = None, key_alias: str | None = None, @@ -4503,6 +4509,7 @@ async def generate_key_helper_fn( "metadata": metadata_json, "tpm_limit": tpm_limit, "rpm_limit": rpm_limit, + "tpd_limit": tpd_limit, "budget_duration": key_budget_duration, "budget_reset_at": key_reset_at, "allowed_cache_controls": allowed_cache_controls, diff --git a/litellm/proxy/management_endpoints/management_v1/budgets.py b/litellm/proxy/management_endpoints/management_v1/budgets.py index cc2fefc426f..ea13e4547bd 100644 --- a/litellm/proxy/management_endpoints/management_v1/budgets.py +++ b/litellm/proxy/management_endpoints/management_v1/budgets.py @@ -58,6 +58,7 @@ class BudgetListItem(BaseModel): soft_budget: float | None = None tpm_limit: int | None = None rpm_limit: int | None = None + tpd_limit: int | None = None budget_duration: str | None = None budget_reset_at: datetime | None = None created_at: datetime @@ -123,7 +124,7 @@ BUDGET_FILTERS: Final[Mapping[str, FilterSpec]] = MappingProxyType( BUDGETS_LIST_SPEC: Final[ListSpec[BudgetListItem, BudgetListItem]] = ListSpec( resource="budgets", - sortable=frozenset(("budget_id", "max_budget", "tpm_limit", "rpm_limit", "created_at")), + sortable=frozenset(("budget_id", "max_budget", "tpm_limit", "rpm_limit", "tpd_limit", "created_at")), searchable=frozenset(("budget_id",)), filters=BUDGET_FILTERS, default_sort=(SortKey(field="created_at", descending=True),), @@ -154,7 +155,7 @@ async def list_budgets( way to page, sort or filter it. `sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`, - `rpm_limit` or `created_at`, each optionally prefixed with `-` for descending, + `rpm_limit`, `tpd_limit` or `created_at`, each optionally prefixed with `-` for descending, and defaults to `-created_at`. `budget_id` is appended to every sort as the tiebreaker. `q` is a case-insensitive substring match on `budget_id`. `page_size` defaults to 50 and is capped at 100. Filters are diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 96e946424bd..c6a76a920f6 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -362,6 +362,7 @@ async def new_organization( - max_budget: *Optional[float]* - Max budget for org - tpm_limit: *Optional[int]* - Max tpm limit for org - rpm_limit: *Optional[int]* - Max rpm limit for org + - tpd_limit: *Optional[int]* - Max tokens per day stored on the org budget. Batch submissions enforce tpd_limit at the key, team and end user scopes only. - model_rpm_limit: *Optional[Dict[str, int]]* - The RPM (Requests Per Minute) limit per model for this organization. - model_tpm_limit: *Optional[Dict[str, int]]* - The TPM (Tokens Per Minute) limit per model for this organization. - max_parallel_requests: *Optional[int]* - [Not Implemented Yet] Max parallel requests for org diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index bc4b3610b1c..e719d6d761a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1217,6 +1217,7 @@ async def new_team( - 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 + - tpd_limit: Optional[int] - The TPD (Tokens Per Day) limit for this team. Batch submissions are charged against it instead of tpm_limit/rpm_limit - rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of RPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating RPM, or "best_effort_throughput" for best effort enforcement. - tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of TPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating TPM, or "best_effort_throughput" for best effort enforcement. - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget @@ -1969,6 +1970,7 @@ async def update_team( - metadata: Optional[dict] - Metadata for team, store information for team. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } - 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 + - tpd_limit: Optional[int] - The TPD (Tokens Per Day) limit for this team. Batch submissions are charged against it instead of tpm_limit/rpm_limit - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set (either in the request or existing), soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set. - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets) diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index b310fc661c4..fe9e104789b 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -8,6 +8,7 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.proxy._types import PassThroughEndpointLoggingResultValues @@ -60,7 +61,7 @@ class PassThroughStreamingHandler: litellm_logging_obj._update_completion_start_time(completion_start_time=datetime.now()) @staticmethod - def schedule_stream_failure_logging( + async def schedule_stream_failure_logging( litellm_logging_obj: LiteLLMLoggingObj, endpoint_type: EndpointType, request_body: dict[str, object], @@ -68,7 +69,7 @@ class PassThroughStreamingHandler: exception: Exception, stream_context: PassThroughStreamContext | None = None, ) -> None: - PassThroughStreamingHandler._record_partial_usage_for_failure( + await asyncify(PassThroughStreamingHandler._record_partial_usage_for_failure)( litellm_logging_obj=litellm_logging_obj, endpoint_type=endpoint_type, request_body=request_body, @@ -222,7 +223,7 @@ class PassThroughStreamingHandler: verbose_proxy_logger.error("Error in chunk_processor: %s", e) if response.status_code < 400: logging_scheduled = True - PassThroughStreamingHandler.schedule_stream_failure_logging( + await PassThroughStreamingHandler.schedule_stream_failure_logging( litellm_logging_obj=litellm_logging_obj, endpoint_type=endpoint_type, request_body=resolved_request_body, @@ -292,7 +293,7 @@ class PassThroughStreamingHandler: ( standard_logging_response_object, kwargs, - ) = PassThroughStreamingHandler._build_passthrough_logging_result( + ) = await asyncify(PassThroughStreamingHandler._build_passthrough_logging_result)( litellm_logging_obj=litellm_logging_obj, passthrough_success_handler_obj=passthrough_success_handler_obj, url_route=url_route, @@ -334,8 +335,8 @@ class PassThroughStreamingHandler: Synchronous, CPU-bound reconstruction of the standard logging payload from collected raw SSE bytes. Extracted from _route_streaming_logging_to_handler so the per-endpoint dispatch can - be unit-tested in isolation. Still invoked synchronously on the event - loop; an off-loop dispatch is a future change, not part of this PR. + be unit-tested in isolation. The async callers run it in a worker + thread so the token counts inside stay off the event loop. """ all_chunks: Final = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes) standard_logging_response_object: PassThroughEndpointLoggingResultValues | None = None diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index dd7967aafe3..8072df5aa5b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -17,6 +17,7 @@ model LiteLLM_BudgetTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? model_max_budget Json? budget_duration String? budget_reset_at DateTime? @@ -133,6 +134,7 @@ model LiteLLM_TeamTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? blocked Boolean @default(false) @@ -203,6 +205,7 @@ model LiteLLM_DeletedTeamTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? blocked Boolean @default(false) @@ -438,6 +441,7 @@ model LiteLLM_VerificationToken { blocked Boolean? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? max_budget Float? budget_duration String? budget_reset_at DateTime? @@ -534,6 +538,7 @@ model LiteLLM_DeletedVerificationToken { blocked Boolean? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? max_budget Float? budget_duration String? budget_reset_at DateTime? diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e611984bf4c..479bd0a55af 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -85,7 +85,11 @@ from litellm._logging import _redact_string, verbose_proxy_logger from litellm._service_logger import ServiceLogging, ServiceTypes from litellm.caching.caching import DualCache, RedisCache from litellm.caching.dual_cache import LimitedSizeOrderedDict -from litellm.exceptions import RejectedRequestError, SensitiveDataRouteException +from litellm.exceptions import ( + GuardrailRaisedException, + RejectedRequestError, + SensitiveDataRouteException, +) from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, @@ -901,6 +905,9 @@ def _call_type_for_route(route: str | None) -> str | None: return call_types[0].value if len(operations) == 1 else None +_PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException) + + def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]: """Failure-path callbacks run after ``litellm_logging_obj`` is popped from request_data (it is not serialisable), so the caller merges these fields @@ -2991,6 +2998,7 @@ class ProxyLogging: - Authentication Errors from user_api_key_auth - HTTP HTTPException (rate limit errors) - ProxyException (guardrail blocks, budget / rate-limit errors) + - GuardrailRaisedException (guardrail blocks / guardrail failures) """ ######################################################### @@ -3005,9 +3013,7 @@ class ProxyLogging: if not (RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route)): return False - return isinstance(original_exception, (HTTPException, ProxyException)) or ( - error_type == ProxyErrorTypes.auth_error - ) + return isinstance(original_exception, _PROXY_ONLY_LLM_API_ERRORS) or (error_type == ProxyErrorTypes.auth_error) async def _handle_logging_proxy_only_error( self, @@ -3563,8 +3569,9 @@ class ProxyLogging: yield chunk except (GeneratorExit, asyncio.CancelledError): raise - except Exception: - ProxyLogging._fire_deferred_stream_logging(request_data) + except Exception as e: + if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, e): + ProxyLogging._fire_deferred_stream_logging(request_data) raise ProxyLogging._fire_deferred_stream_logging(request_data) return @@ -3638,8 +3645,9 @@ class ProxyLogging: yield chunk except (GeneratorExit, asyncio.CancelledError): raise - except Exception: - ProxyLogging._fire_deferred_stream_logging(request_data) + except Exception as e: + if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, e): + ProxyLogging._fire_deferred_stream_logging(request_data) raise # Fire deferred logging AFTER all guardrail end-of-stream blocks @@ -3735,6 +3743,23 @@ class ProxyLogging: logging_obj._deferred_stream_complete_args = None asyncio.create_task(_deferred_cb(*_args)) + @staticmethod + def _discard_deferred_stream_logging_for_failure(request_data: Mapping[str, object], error: Exception) -> bool: + """Drop the parked success dispatch for an assembled chat stream that ends in an error + ``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead. + Returns False when the parked dispatch should still be flushed by the caller.""" + logging_obj: Final = request_data.get("litellm_logging_obj") + if not isinstance(logging_obj, Logging): + return False + _args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None) + assembled: Final = _args[0] if _args else None + if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS) or not isinstance(assembled, ModelResponse): + return False + logging_obj._on_deferred_stream_complete = None + logging_obj._deferred_stream_complete_args = None + logging_obj.record_assembled_response_for_failure(assembled) + return True + async def _arelease_max_parallel_requests_on_disconnect( self, user_api_key_dict: UserAPIKeyAuth, @@ -4295,7 +4320,8 @@ class PrismaClient: t.spend AS team_spend, t.max_budget AS team_max_budget, t.tpm_limit AS team_tpm_limit, - t.rpm_limit AS team_rpm_limit + t.rpm_limit AS team_rpm_limit, + t.tpd_limit AS team_tpd_limit FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id; """, @@ -4734,6 +4760,7 @@ class PrismaClient: t.soft_budget AS team_soft_budget, t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, + t.tpd_limit AS team_tpd_limit, t.models AS team_models, t.metadata AS team_metadata, t.blocked AS team_blocked, @@ -4751,6 +4778,7 @@ class PrismaClient: b.max_budget AS litellm_budget_table_max_budget, b.tpm_limit AS litellm_budget_table_tpm_limit, b.rpm_limit AS litellm_budget_table_rpm_limit, + b.tpd_limit AS litellm_budget_table_tpd_limit, b.model_max_budget as litellm_budget_table_model_max_budget, b.soft_budget as litellm_budget_table_soft_budget, o.metadata as organization_metadata, diff --git a/schema.prisma b/schema.prisma index dd7967aafe3..8072df5aa5b 100644 --- a/schema.prisma +++ b/schema.prisma @@ -17,6 +17,7 @@ model LiteLLM_BudgetTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? model_max_budget Json? budget_duration String? budget_reset_at DateTime? @@ -133,6 +134,7 @@ model LiteLLM_TeamTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? blocked Boolean @default(false) @@ -203,6 +205,7 @@ model LiteLLM_DeletedTeamTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? blocked Boolean @default(false) @@ -438,6 +441,7 @@ model LiteLLM_VerificationToken { blocked Boolean? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? max_budget Float? budget_duration String? budget_reset_at DateTime? @@ -534,6 +538,7 @@ model LiteLLM_DeletedVerificationToken { blocked Boolean? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? max_budget Float? budget_duration String? budget_reset_at DateTime? diff --git a/tests/pass_through_unit_tests/conftest.py b/tests/pass_through_unit_tests/conftest.py index e6e98f790e8..df8196f1785 100644 --- a/tests/pass_through_unit_tests/conftest.py +++ b/tests/pass_through_unit_tests/conftest.py @@ -1,6 +1,8 @@ +import asyncio import pytest +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, @@ -45,6 +47,21 @@ def _vcr_outcome_gate(request, vcr): record_vcr_outcome(request, vcr) +@pytest.fixture(autouse=True) +async def _drain_logging_worker(): + """ + The logging queue is bound to the running loop, so anything left queued when a test's loop + goes away is carried onto the next loop and fires against that test's callbacks. + """ + GLOBAL_LOGGING_WORKER.start() + try: + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + except asyncio.TimeoutError: + pass + await GLOBAL_LOGGING_WORKER.stop() + yield + + def pytest_configure(config): _verbose_state.remember_pluginmanager(config) reset_vcr_diag_dir() diff --git a/tests/test_litellm/a2a_protocol/test_a2a_streaming_iterator.py b/tests/test_litellm/a2a_protocol/test_a2a_streaming_iterator.py index d86cbb94a91..2603d135dce 100644 --- a/tests/test_litellm/a2a_protocol/test_a2a_streaming_iterator.py +++ b/tests/test_litellm/a2a_protocol/test_a2a_streaming_iterator.py @@ -100,3 +100,57 @@ async def test_custom_logger_only_never_submits_sync_success_handler(monkeypatch assert recorder.async_hook_fired is True assert recording_executor.submitted_for(logging_obj) == [] + + +class _AgentChunk: + def __init__(self, text: str): + self._text = text + + def model_dump(self, mode: str, exclude_none: bool) -> dict: + return {"result": {"kind": "message", "role": "agent", "parts": [{"kind": "text", "text": self._text}]}} + + +@pytest.mark.asyncio +async def test_stream_completion_counts_tokens_off_the_event_loop(monkeypatch): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("gpt-5.6-luna") + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + logging_obj = LitellmLogging( + model="a2a/test-agent", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="a2a_send_message_streaming", + start_time=time.time(), + litellm_call_id="lit-7190-test", + function_id="lit-7190-test", + ) + + async def _stream(): + yield _AgentChunk(text * 100) + + iterator = A2AStreamingIterator( + stream=_stream(), + request=SimpleNamespace( + params=SimpleNamespace(message={"role": "user", "parts": [{"kind": "text", "text": text * 100}]}) + ), + logging_obj=logging_obj, + agent_name="test-agent", + ) + + async def drain() -> int: + return len([chunk async for chunk in iterator]) + + yielded, took, lags = await timed_with_loop_lags(drain) + + assert yielded == 1 + usage = logging_obj.model_call_details["usage"] + assert usage.prompt_tokens > 100_000 + assert usage.completion_tokens > 100_000 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/a2a_protocol/test_main.py b/tests/test_litellm/a2a_protocol/test_main.py index 8850a2eca6c..318b40138ed 100644 --- a/tests/test_litellm/a2a_protocol/test_main.py +++ b/tests/test_litellm/a2a_protocol/test_main.py @@ -1,5 +1,7 @@ """Tests for litellm/a2a_protocol/main.py non-streaming send behavior.""" +import asyncio + import httpx import pytest @@ -13,7 +15,8 @@ from a2a.compat.v0_3.types import ( ) import litellm -from litellm.a2a_protocol.main import _send_message, _stream_messages, create_a2a_client +from litellm.integrations.custom_logger import CustomLogger +from litellm.a2a_protocol.main import _send_message, _stream_messages, asend_message, create_a2a_client from litellm.caching.llm_caching_handler import LLMClientCache from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT from litellm.llms.custom_httpx.http_handler import ( @@ -413,3 +416,51 @@ async def test_the_pooled_a2a_client_arrives_with_cookie_persistence_disabled(is assert dict(handler.client.cookies) == {}, "the pooled A2A client kept an upstream's cookie" await handler.close() + + +class _UsageRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.logged = asyncio.Event() + self.payload = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.payload = kwargs["standard_logging_object"] + self.logged.set() + + +@pytest.mark.asyncio +async def test_asend_message_counts_usage_off_the_event_loop(monkeypatch): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("gpt-5.6-luna") + recorder = _UsageRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "success_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + + reply = _conv.pb2_v10.StreamResponse() + reply.message.message_id = "reply-1" + reply.message.role = _conv.pb2_v10.Role.ROLE_AGENT + reply.message.parts.add().text = text * 100 + request = SendMessageRequest( + id="r1", + params=MessageSendParams( + message={"messageId": "m1", "role": "user", "parts": [{"kind": "text", "text": text * 100}]} + ), + ) + + response, took, lags = await timed_with_loop_lags( + lambda: asend_message(a2a_client=_FakeClient(reply), request=request) + ) + + assert response.id == "r1" + await asyncio.wait_for(recorder.logged.wait(), timeout=10) + assert recorder.payload["prompt_tokens"] > 100_000 + assert recorder.payload["completion_tokens"] > 100_000 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/caching/test_qdrant_semantic_cache.py b/tests/test_litellm/caching/test_qdrant_semantic_cache.py index e07578dd7e5..a0a9b71787c 100644 --- a/tests/test_litellm/caching/test_qdrant_semantic_cache.py +++ b/tests/test_litellm/caching/test_qdrant_semantic_cache.py @@ -1026,3 +1026,36 @@ def test_qdrant_semantic_cache_defaults_embedding_timeout(): cache = QdrantSemanticCache.__new__(QdrantSemanticCache) assert cache.embedding_timeout == SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS assert SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS < 60 + + +@pytest.mark.asyncio +async def test_qdrant_async_embedding_truncates_off_the_event_loop(monkeypatch): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + + warm_tokenizer("sem-embed") + cache = QdrantSemanticCache.__new__(QdrantSemanticCache) + cache.embedding_model = "sem-embed" + cache.embedding_max_input_tokens = 5 + cache.embedding_timeout = 5 + + router = MagicMock() + router.get_configured_token_limits.return_value = (8191, None) + router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]}) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + _router_proxy_module(router, "sem-embed"), + ) + + response, took, lags = await timed_with_loop_lags(lambda: cache._get_async_embedding(text * 100)) + + assert response["data"][0]["embedding"] == [0.1, 0.2] + assert _token_count("sem-embed", router.aembedding.call_args.kwargs["input"]) == 5 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index 9884e9d9bc0..de253b4f10b 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -1387,3 +1387,32 @@ def test_redis_semantic_cache_defaults_embedding_timeout(): cache = RedisSemanticCache.__new__(RedisSemanticCache) assert cache.embedding_timeout == SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS assert SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS < 60 + + +@pytest.mark.asyncio +async def test_redis_async_embedding_truncates_off_the_event_loop(monkeypatch): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + warm_tokenizer("sem-embed") + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "sem-embed" + cache.embedding_max_input_tokens = 5 + cache.embedding_timeout = 5 + + router = MagicMock() + router.get_configured_token_limits.return_value = (8191, None) + router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]}) + _proxy_with_router(monkeypatch, router, "sem-embed") + + embedding, took, lags = await timed_with_loop_lags(lambda: cache._get_async_embedding(text * 100)) + + assert embedding == [0.1, 0.2] + assert _token_count("sem-embed", router.aembedding.call_args.kwargs["input"]) == 5 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py b/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py index bc4dccc7d70..e66cd654f93 100644 --- a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py +++ b/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py @@ -523,3 +523,28 @@ async def test_pre_call_hook_no_compression_records_no_savings(monkeypatch): await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.anthropic_messages) assert "compression_savings" not in litellm_metadata + + +@pytest.mark.asyncio +async def test_pre_call_hook_counts_tokens_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + model = "anthropic/claude-fable-5" + warm_tokenizer(model) + logger = CompressionInterceptionLogger(compression_trigger=10_000_000) + messages = [{"role": "user", "content": text * 100}] + kwargs = {"model": model, "messages": messages} + + result, took, lags = await timed_with_loop_lags( + lambda: logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.anthropic_messages) + ) + + assert result is not None + assert result["messages"] is messages + assert "tools" not in result + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 784468c839b..47efbe7f19a 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4927,3 +4927,50 @@ class TestStableStreamingResponseId: ) wrapper.response_id = "chatcmpl-from-provider" assert wrapper.model_response_creator().id == "chatcmpl-from-provider" + + +@pytest.mark.asyncio +async def test_async_stream_without_usage_counts_tokens_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + model = "gpt-5.6-luna" + warm_tokenizer(model) + messages = [{"role": "user", "content": text * 100}] + content_chunks = [_make_chunk(text) for _ in range(100)] + stop_chunk = ModelResponseStream( + id="test", + created=1741037890, + model=model, + choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")], + ) + logging_obj = Logging( + model=model, + messages=messages, + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="12345", + function_id="1245", + ) + wrapper = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=content_chunks + [stop_chunk]), + model=model, + custom_llm_provider="openai", + logging_obj=logging_obj, + stream_options={"include_usage": True}, + ) + + async def consume() -> list[ModelResponseStream]: + return [chunk async for chunk in wrapper] + + chunks, took, lags = await timed_with_loop_lags(consume) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == text * 100 + assert chunks[-1].usage.prompt_tokens > 100_000 + assert chunks[-1].usage.completion_tokens > 100_000 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py index 28c82fdf528..7660a8649b5 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -2605,3 +2605,34 @@ def test_build_summary_messages_keeps_midturn_system_correction_in_place(): assert summary_messages[0]["content"] == "caller system prompt" assert summary_messages[2]["content"] == "use the corrected result" assert summary_messages[-1]["content"] == "summarize the conversation" + + +async def test_threshold_check_counts_tokens_off_the_event_loop(monkeypatch): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + from litellm.llms.anthropic.experimental_pass_through.context_management.constants import ( + COMPACT_SUMMARY_MODEL_SETTING_KEY, + ) + from litellm.proxy.proxy_server import general_settings + + monkeypatch.setitem(general_settings, COMPACT_SUMMARY_MODEL_SETTING_KEY, "claude-haiku-4-5") + warm_tokenizer(MODEL) + messages = [{"role": "user", "content": text * 100}, *_simple_messages()] + result, took, lags = await timed_with_loop_lags( + lambda: apply_compact_20260112( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec={"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 10_000_000}}, + ) + ) + + assert result.messages == messages + assert result.compaction_block is None + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py index 50c72cfe8d0..a21c22cf5fa 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py @@ -129,3 +129,35 @@ async def test_malformed_edit_entries_are_skipped(): ) assert result.applied_edits == [] assert result.messages == messages + + +async def test_sync_editor_counts_tokens_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer(MODEL) + messages = [{"role": "user", "content": text * 100}, *_history_with_two_tool_pairs()] + + result, took, lags = await timed_with_loop_lags( + lambda: apply_context_management( + model=MODEL, + messages=messages, + tools=None, + system=None, + context_management_spec={ + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 10_000_000}, + } + ] + }, + ) + ) + + assert result.messages == messages + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py index cf1f665ad21..5263cf2774c 100644 --- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -277,3 +277,18 @@ def test_update_valid_token_db_values_override_custom_auth_when_set(): # DB values should win assert result.end_user_tpm_limit == 500 assert result.end_user_model_max_budget == db_budget + + +def test_end_user_budget_tpd_limit_reaches_the_token(): + from litellm.proxy.auth.user_api_key_auth import _apply_budget_limits_to_end_user_params + + end_user_params = {"end_user_id": "user_1"} + _apply_budget_limits_to_end_user_params( + end_user_params=end_user_params, + budget_info=LiteLLM_BudgetTable(rpm_limit=5, tpd_limit=750000), + end_user_id="user_1", + ) + result = update_valid_token_with_end_user_params(UserAPIKeyAuth(token="test_token"), end_user_params) + + assert result.end_user_rpm_limit == 5 + assert result.end_user_tpd_limit == 750000 diff --git a/tests/test_litellm/proxy/auth/test_team_grants.py b/tests/test_litellm/proxy/auth/test_team_grants.py index 447fc1c93a1..7b6717f804f 100644 --- a/tests/test_litellm/proxy/auth/test_team_grants.py +++ b/tests/test_litellm/proxy/auth/test_team_grants.py @@ -27,6 +27,7 @@ def _full_team(model_aliases=ALIASES) -> LiteLLM_TeamTable: team_alias="grants-team", tpm_limit=1000, rpm_limit=10, + tpd_limit=200000, max_budget=50.0, soft_budget=25.0, spend=12.5, @@ -67,6 +68,7 @@ def test_team_grants_cover_every_team_field_the_key_path_gets(): assert token.team_alias == "grants-team" assert token.team_tpm_limit == 1000 assert token.team_rpm_limit == 10 + assert token.team_tpd_limit == 200000 assert token.team_max_budget == 50.0 assert token.team_soft_budget == 25.0 assert token.team_spend == 12.5 diff --git a/tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py b/tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py new file mode 100644 index 00000000000..919e9c79828 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py @@ -0,0 +1,259 @@ +""" +Tests for `tpd_limit` (tokens per day) enforcement on batch submissions. + +A batch's rows are scheduled by the provider, so a caller cannot keep a large +batch under a per-minute RPM/TPM budget. Scopes that configure `tpd_limit` +are charged against a 24h token window instead of their minute counters. +""" + +from datetime import datetime + +import pytest +from fastapi import HTTPException + +from litellm import DualCache +from litellm.constants import BATCH_TPD_WINDOW_SECONDS +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.batch_rate_limiter import BatchFileUsage +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, +) +from litellm.proxy.utils import InternalUsageCache, hash_token + + +class _Clock: + def __init__(self, start: datetime): + self.now = start + + def __call__(self) -> datetime: + return self.now + + +def _make_limiters(clock: _Clock | None = None): + internal_usage_cache = InternalUsageCache(dual_cache=DualCache()) + rate_limiter = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=internal_usage_cache, time_provider=clock) + batch_limiter = rate_limiter._get_batch_rate_limiter() + assert batch_limiter is not None + return internal_usage_cache, rate_limiter, batch_limiter + + +async def _counter(internal_usage_cache, rate_limiter, descriptor_key, value, rate_limit_type): + cache_key = rate_limiter.create_rate_limit_keys(descriptor_key, value, rate_limit_type) + raw = await internal_usage_cache.async_get_cache(key=cache_key, litellm_parent_otel_span=None, local_only=True) + return int(raw or 0) + + +@pytest.mark.asyncio +async def test_batch_over_rpm_and_tpm_but_under_tpd_is_accepted(): + internal_usage_cache, rate_limiter, batch_limiter = _make_limiters() + api_key = hash_token("tpd-key") + user_api_key_dict = UserAPIKeyAuth(api_key=api_key, rpm_limit=1, tpm_limit=10, tpd_limit=1000) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=500, request_count=50), + ) + + assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 500 + assert await _counter(internal_usage_cache, rate_limiter, "api_key", api_key, "requests") == 0 + assert await _counter(internal_usage_cache, rate_limiter, "api_key", api_key, "tokens") == 0 + + +@pytest.mark.asyncio +async def test_cumulative_batch_tokens_over_tpd_returns_429_with_remaining_daily_window(): + window_start = datetime(2026, 9, 13, 8, 0, 0) + clock = _Clock(window_start) + _internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters(clock) + user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("tpd-key-2"), rpm_limit=1, tpd_limit=1000) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=600, request_count=6), + ) + clock.now = datetime(2026, 9, 13, 11, 0, 0) + with pytest.raises(HTTPException) as exc: + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=600, request_count=6), + ) + + assert exc.value.status_code == 429 + assert "api_key_tpd" in str(exc.value.detail) + assert "600 tokens but only 400 tokens remaining out of 1000 TPD limit" in str(exc.value.detail) + assert exc.value.headers["retry-after"] == str(BATCH_TPD_WINDOW_SECONDS - 3 * 3600) + assert exc.value.headers["reset_at"] == "2026-09-14 08:00:00 UTC" + + +@pytest.mark.asyncio +async def test_failed_batch_submission_refunds_tpd_tokens(): + internal_usage_cache, rate_limiter, batch_limiter = _make_limiters() + api_key = hash_token("tpd-refund-key") + user_api_key_dict = UserAPIKeyAuth(api_key=api_key, rpm_limit=1, tpd_limit=1000) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=600, request_count=6), + ) + await rate_limiter.async_post_call_failure_hook( + request_data={}, + original_exception=RuntimeError("provider rejected the file"), + user_api_key_dict=user_api_key_dict, + ) + + assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 0 + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=1000, request_count=10), + ) + assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 1000 + + +@pytest.mark.asyncio +async def test_tpd_refund_applies_once_and_only_to_daily_counters(): + internal_usage_cache, rate_limiter, batch_limiter = _make_limiters() + team_key = UserAPIKeyAuth( + api_key=hash_token("tpd-refund-team-key"), + rpm_limit=100, + tpm_limit=10_000, + team_id="team-r", + team_tpd_limit=5000, + ) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=team_key, + data={}, + batch_usage=BatchFileUsage(total_tokens=800, request_count=8), + ) + await rate_limiter.async_post_call_failure_hook( + request_data={}, original_exception=RuntimeError("boom"), user_api_key_dict=team_key + ) + await rate_limiter.async_post_call_failure_hook( + request_data={}, original_exception=RuntimeError("boom"), user_api_key_dict=team_key + ) + + assert await _counter(internal_usage_cache, rate_limiter, "team_tpd", "team-r", "tokens") == 0 + assert await _counter(internal_usage_cache, rate_limiter, "api_key", team_key.api_key, "tokens") == 800 + assert await _counter(internal_usage_cache, rate_limiter, "api_key", team_key.api_key, "requests") == 8 + + +@pytest.mark.asyncio +async def test_rejected_batch_leaves_nothing_to_refund(): + internal_usage_cache, rate_limiter, batch_limiter = _make_limiters() + api_key = hash_token("tpd-rejected-key") + user_api_key_dict = UserAPIKeyAuth(api_key=api_key, tpd_limit=100) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, data={}, batch_usage=BatchFileUsage(total_tokens=90, request_count=9) + ) + with pytest.raises(HTTPException): + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, data={}, batch_usage=BatchFileUsage(total_tokens=20, request_count=2) + ) + await rate_limiter.async_post_call_failure_hook( + request_data={}, original_exception=RuntimeError("429 bubbled up"), user_api_key_dict=user_api_key_dict + ) + + assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 90 + + +@pytest.mark.asyncio +async def test_batch_without_tpd_still_enforces_minute_rpm(): + _internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters() + user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("rpm-only-key"), rpm_limit=1, tpm_limit=1000) + + with pytest.raises(HTTPException) as exc: + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=50, request_count=5), + ) + + assert exc.value.status_code == 429 + assert "RPM limit" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_team_tpd_replaces_team_minute_limits_but_key_minute_limits_still_apply(): + internal_usage_cache, rate_limiter, batch_limiter = _make_limiters() + team_key = UserAPIKeyAuth( + api_key=hash_token("team-key"), + team_id="team-1", + team_rpm_limit=1, + team_tpm_limit=10, + team_tpd_limit=5000, + ) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=team_key, + data={}, + batch_usage=BatchFileUsage(total_tokens=800, request_count=8), + ) + assert await _counter(internal_usage_cache, rate_limiter, "team_tpd", "team-1", "tokens") == 800 + assert await _counter(internal_usage_cache, rate_limiter, "team", "team-1", "requests") == 0 + + key_rpm_in_team_with_tpd = UserAPIKeyAuth( + api_key=hash_token("team-key-2"), + rpm_limit=1, + team_id="team-1", + team_tpd_limit=5000, + ) + with pytest.raises(HTTPException) as exc: + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=key_rpm_in_team_with_tpd, + data={}, + batch_usage=BatchFileUsage(total_tokens=10, request_count=2), + ) + assert exc.value.status_code == 429 + assert "api_key:" in str(exc.value.detail) + assert "RPM limit" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_end_user_tpd_is_enforced_per_end_user(): + _internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters() + first_customer = UserAPIKeyAuth( + api_key=hash_token("shared-key"), end_user_id="customer-a", end_user_rpm_limit=1, end_user_tpd_limit=100 + ) + second_customer = UserAPIKeyAuth( + api_key=hash_token("shared-key"), end_user_id="customer-b", end_user_rpm_limit=1, end_user_tpd_limit=100 + ) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=first_customer, data={}, batch_usage=BatchFileUsage(total_tokens=90, request_count=9) + ) + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=second_customer, data={}, batch_usage=BatchFileUsage(total_tokens=90, request_count=9) + ) + with pytest.raises(HTTPException) as exc: + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=first_customer, data={}, batch_usage=BatchFileUsage(total_tokens=20, request_count=2) + ) + assert exc.value.status_code == 429 + assert "end_user_tpd: customer-a" in str(exc.value.detail) + + +def test_tpd_only_key_is_not_skipped_as_having_no_limits(): + _internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters() + descriptors = batch_limiter._create_batch_rate_limit_descriptors( + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("tpd-only"), tpd_limit=100), + data={}, + ) + assert batch_limiter._has_applicable_batch_rate_limits(descriptors) is True + + +def test_online_descriptors_ignore_tpd_limit(): + _internal_usage_cache, rate_limiter, _batch_limiter = _make_limiters() + api_key = hash_token("online-key") + descriptors = rate_limiter._create_rate_limit_descriptors( + user_api_key_dict=UserAPIKeyAuth(api_key=api_key, rpm_limit=5, tpd_limit=100, team_id="t", team_tpd_limit=9), + data={"model": "gpt-4o"}, + rpm_limit_type=None, + tpm_limit_type=None, + model_has_failures=False, + ) + assert [(d["key"], d["rate_limit"]["window_size"]) for d in descriptors] == [("api_key", rate_limiter.window_size)] diff --git a/tests/test_litellm/proxy/management_endpoints/management_v1/test_budgets.py b/tests/test_litellm/proxy/management_endpoints/management_v1/test_budgets.py index add2126ac7b..2b438a9d370 100644 --- a/tests/test_litellm/proxy/management_endpoints/management_v1/test_budgets.py +++ b/tests/test_litellm/proxy/management_endpoints/management_v1/test_budgets.py @@ -52,7 +52,7 @@ app.include_router(router) client = TestClient(app) BUDGETS_PATH = f"{MANAGEMENT_V1_PREFIX}/budgets" -SORTABLE = ["budget_id", "created_at", "max_budget", "rpm_limit", "tpm_limit"] +SORTABLE = ["budget_id", "created_at", "max_budget", "rpm_limit", "tpd_limit", "tpm_limit"] def _row(budget_id: str, **overrides: Any) -> dict[str, Any]: @@ -62,6 +62,7 @@ def _row(budget_id: str, **overrides: Any) -> dict[str, Any]: "soft_budget": None, "tpm_limit": None, "rpm_limit": None, + "tpd_limit": None, "budget_duration": "30d", "budget_reset_at": None, "created_at": "2026-07-20T12:00:00+00:00", @@ -123,7 +124,7 @@ def test_returns_flat_rows_in_the_control_plane_envelope(query_raw, as_proxy_adm def test_serves_the_columns_the_budgets_page_renders(query_raw, as_proxy_admin): - _serve(query_raw, [_row("b-1", soft_budget=5.0, budget_reset_at="2026-08-01T00:00:00+00:00")]) + _serve(query_raw, [_row("b-1", soft_budget=5.0, tpd_limit=250000, budget_reset_at="2026-08-01T00:00:00+00:00")]) row = _get().json()["data"][0] @@ -133,12 +134,14 @@ def test_serves_the_columns_the_budgets_page_renders(query_raw, as_proxy_admin): "soft_budget", "tpm_limit", "rpm_limit", + "tpd_limit", "budget_duration", "budget_reset_at", "created_at", "updated_at", } assert row["soft_budget"] == 5.0 + assert row["tpd_limit"] == 250000 assert row["budget_reset_at"].startswith("2026-08-01T00:00:00") diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py index 4b6815d7552..2f3be61d00f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py @@ -136,6 +136,21 @@ async def test_update_budget_success(client_and_mocks, monkeypatch): assert body["updated_by"] == "test_user" +@pytest.mark.asyncio +async def test_new_and_update_budget_persist_tpd_limit(client_and_mocks): + client, _, mock_table = client_and_mocks + + resp = client.post("/budget/new", json={"budget_id": "budget_tpd", "tpd_limit": 250000}) + assert resp.status_code == 200, resp.text + assert resp.json()["tpd_limit"] == 250000 + assert mock_table.create.await_args.kwargs["data"]["tpd_limit"] == 250000 + + resp = client.post("/budget/update", json={"budget_id": "budget_tpd", "tpd_limit": 500000}) + assert resp.status_code == 200, resp.text + assert resp.json()["tpd_limit"] == 500000 + assert mock_table.update.await_args.kwargs["data"]["tpd_limit"] == 500000 + + @pytest.mark.asyncio async def test_update_budget_missing_id(client_and_mocks, monkeypatch): client, mock_prisma, mock_table = client_and_mocks diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index bd59a82cbd2..9ce3a6fb4c2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -743,6 +743,7 @@ _EXPECTED_CUSTOMER = { "max_parallel_requests": None, "tpm_limit": None, "rpm_limit": None, + "tpd_limit": None, "model_max_budget": None, "budget_duration": "30d", "allowed_models": [], 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 c590a24203c..63055872aa1 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 @@ -474,6 +474,28 @@ async def test_key_expiration_exact_duration_hours(monkeypatch): ), f"Expected expiration to be approximately 12 hours from creation, got {hours_diff} hours" +@pytest.mark.asyncio +async def test_generate_key_persists_tpd_limit(monkeypatch): + mock_prisma_client = AsyncMock() + mock_prisma_client.insert_data = AsyncMock( + return_value=MagicMock(token="hashed_token_123", litellm_budget_table=None) + ) + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + data_json = GenerateKeyRequest(tpd_limit=250000, rpm_limit=5).model_dump(exclude_none=True) + response = await generate_key_helper_fn(request_type="key", **data_json, table_name="key") + + assert response["tpd_limit"] == 250000 + key_insert = mock_prisma_client.insert_data.await_args_list[-1].kwargs + assert key_insert["table_name"] == "key" + assert key_insert["data"]["tpd_limit"] == 250000 + assert key_insert["data"]["rpm_limit"] == 5 + + @pytest.mark.asyncio async def test_key_generation_with_object_permission(monkeypatch): """Ensure /key/generate correctly handles `object_permission` input by @@ -1823,6 +1845,18 @@ async def test_update_key_enable_prompt_caching_folds_into_metadata(flag_value): assert "enable_prompt_caching" not in {k for k in updated if k != "metadata"} +@pytest.mark.asyncio +@pytest.mark.parametrize("tpd_limit", [250000, None]) +async def test_update_key_writes_tpd_limit_as_a_column(tpd_limit): + data = UpdateKeyRequest(key="sk-1", tpd_limit=tpd_limit) + existing_key = LiteLLM_VerificationToken(token="hashed", tpd_limit=1) + + updated = await prepare_key_update_data(data=data, existing_key_row=existing_key) + + assert updated["tpd_limit"] == tpd_limit + assert "rpm_limit" not in updated + + @pytest.mark.asyncio async def test_update_preserves_service_account_id_when_metadata_replaced(): """ 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 1f17bc29f3b..ebbedc6541e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -675,6 +675,42 @@ async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth): assert "object_permission" not in team_data +@pytest.mark.asyncio +async def test_new_team_persists_tpd_limit(mock_db_client, mock_admin_auth): + mock_db_client.jsonify_team_object = lambda db_data: db_data + mock_db_client.get_data = AsyncMock(return_value=None) + mock_db_client.update_data = AsyncMock(return_value=MagicMock()) + mock_db_client.db = MagicMock() + mock_db_client.db.litellm_modeltable = MagicMock() + mock_db_client.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) + + team_create_result = MagicMock(team_id="team-tpd") + team_create_result.model_dump.return_value = {"team_id": "team-tpd", "tpd_limit": 250000} + mock_team_create = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_teamtable = MagicMock() + mock_db_client.db.litellm_teamtable.create = mock_team_create + _wire_team_create_tx(mock_db_client) + mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_usertable = MagicMock() + mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + await new_team( + data=NewTeamRequest(team_alias="tpd-team", rpm_limit=5, tpd_limit=250000), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + team_data = mock_team_create.call_args.kwargs["data"] + assert team_data["tpd_limit"] == 250000 + assert team_data["rpm_limit"] == 5 + + @pytest.mark.asyncio async def test_new_team_with_mcp_tool_permissions(mock_db_client, mock_admin_auth): """ @@ -7596,6 +7632,48 @@ async def test_update_team_rpm_limit_not_gated_by_user_limit( assert result is not None +@pytest.mark.asyncio +async def test_update_team_persists_tpd_limit(disable_audit_logging_for_mocked_team): + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import update_team + + with ( + patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point + "litellm.proxy.proxy_server.prisma_client" + ) as mock_prisma, + patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, + patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), + patch( # test-quality-ok: stubs the audit write so the test observes only the team column written + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ), + ): + existing_team = MagicMock(team_id="team-tpd", organization_id=None, model_id=None, tpd_limit=None) + existing_team.model_dump.return_value = {"team_id": "team-tpd", "organization_id": None} + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma.jsonify_team_object = lambda db_data: db_data + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + updated_team = MagicMock(team_id="team-tpd", organization_id=None, litellm_model_table=None) + updated_team.model_dump.return_value = {"team_id": "team-tpd", "organization_id": None, "tpd_limit": 250000} + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated_team) + + await update_team( + data=UpdateTeamRequest(team_id="team-tpd", tpd_limit=250000), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + ) + + written = mock_prisma.db.litellm_teamtable.update.call_args.kwargs["data"] + assert written["tpd_limit"] == 250000 + assert "rpm_limit" not in written + + @pytest.mark.asyncio async def test_new_team_org_scoped_tpm_exceeds_org_limit(): """ diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py index dd9fbd9161f..e91b7ef970c 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py @@ -7,6 +7,7 @@ import pytest import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, ) @@ -128,3 +129,104 @@ def test_vertex_generate_content_payload_prices_gemini_urls_at_gemini_rates(): assert result["kwargs"]["response_cost"] == pytest.approx(GEMINI_COST) assert logging_obj.model_call_details["custom_llm_provider"] == "gemini" + + +def _interrupted_anthropic_stream(model: str, output_text: str) -> list[bytes]: + def sse(event: str, data: dict) -> bytes: + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + message_start = { + "type": "message_start", + "message": { + "id": "msg_interrupted", + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 29, "output_tokens": 2}, + }, + } + block_start = {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} + delta = {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": output_text}} + return [ + sse("message_start", message_start), + sse("content_block_start", block_start), + sse("content_block_delta", delta), + ] + + +@pytest.mark.asyncio +async def test_interrupted_anthropic_stream_recovers_output_tokens_off_the_event_loop(): + from unittest.mock import AsyncMock + + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + model = "claude-fable-5" + warm_tokenizer(model) + logging_obj = _logging_obj() + logging_obj.model_call_details = {"model": model, "stream": True} + logging_obj.litellm_params = {} + logging_obj.get_router_model_id.return_value = None + logging_obj.dispatch_success_handlers = AsyncMock() + + _, took, lags = await timed_with_loop_lags( + lambda: PassThroughStreamingHandler._route_streaming_logging_to_handler( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=PassThroughEndpointLogging(), + url_route="/anthropic/v1/messages", + request_body={"model": model, "stream": True}, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + raw_bytes=_interrupted_anthropic_stream(model, text * 100), + end_time=datetime.now(), + model=model, + ) + ) + + logging_obj.dispatch_success_handlers.assert_awaited_once() + logged_usage = logging_obj.dispatch_success_handlers.await_args.kwargs["result"].usage + assert logged_usage.completion_tokens > 100_000 + assert_loop_stayed_free(took, lags) + + +@pytest.mark.asyncio +async def test_failed_anthropic_stream_records_partial_usage_off_the_event_loop(): + from unittest.mock import AsyncMock + + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + model = "claude-fable-5" + warm_tokenizer(model) + logging_obj = _logging_obj() + logging_obj.model_call_details = {"model": model, "stream": True} + logging_obj.litellm_params = {} + logging_obj.get_router_model_id.return_value = None + logging_obj.dispatch_failure_handlers = AsyncMock() + + _, took, lags = await timed_with_loop_lags( + lambda: PassThroughStreamingHandler.schedule_stream_failure_logging( + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + request_body={"model": model, "stream": True}, + raw_bytes=_interrupted_anthropic_stream(model, text * 100), + exception=RuntimeError("upstream closed the stream"), + ) + ) + await GLOBAL_LOGGING_WORKER.flush() + + logging_obj.dispatch_failure_handlers.assert_awaited_once() + partial_usage = logging_obj.record_partial_usage_for_failure.call_args.kwargs["usage"] + assert partial_usage.completion_tokens > 100_000 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 8b2b6b4c6ca..d7a6124dd97 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -4,13 +4,14 @@ and ``_handle_logging_proxy_only_error``.""" from __future__ import annotations import asyncio -from typing import Any +from datetime import datetime from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException import litellm +from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import AlertType, ProxyErrorTypes from litellm.proxy.utils import ProxyLogging @@ -47,12 +48,17 @@ def test_is_proxy_only_llm_api_truth_table(proxy_logging): error_type=ProxyErrorTypes.auth_error, route="/chat/completions", ), + "guardrail_raised_on_llm_route": proxy_logging._is_proxy_only_llm_api_error( + original_exception=GuardrailRaisedException(guardrail_name="g", message="blocked"), + route="/chat/completions", + ), } assert snapshot == { "no_route": False, "non_llm_route": False, "http_on_llm_route": True, "auth_short_circuit": True, + "guardrail_raised_on_llm_route": True, } @@ -318,3 +324,50 @@ async def test_post_call_failure_hook_keeps_the_route_for_multi_operation_routes route=route, ) assert request_data["call_type"] == route + + +@pytest.mark.asyncio +async def test_post_call_failure_hook_guardrail_block_fires_failure_callback( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """A ``GuardrailRaisedException`` on an LLM route must reach the logging + object's ``async_failure_handler`` so custom loggers see a ``failure`` + status - without this, guardrail blocks produce only + ``post_call_failure_hook`` and no failure logging event.""" + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + recorded: list[object] = [] + + class _StatusRecorder(CustomLogger): + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + standard_logging_object = kwargs.get("standard_logging_object") + recorded.append(standard_logging_object.get("status") if isinstance(standard_logging_object, dict) else None) + + monkeypatch.setattr(litellm, "_async_failure_callback", [_StatusRecorder()]) + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="test_guardrail_block_failure_cb", + function_id="test_guardrail_block_failure_cb", + ) + request_data = { + "litellm_logging_obj": logging_obj, + "litellm_call_id": "test_guardrail_block_failure_cb", + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "metadata": {}, + } + proxy_logging.alert_types = [] + await proxy_logging.post_call_failure_hook( + request_data=request_data, + original_exception=GuardrailRaisedException(guardrail_name="g", message="blocked"), + user_api_key_dict=make_user_api_key_auth(request_route="/chat/completions"), + ) + await asyncio.sleep(0) + await asyncio.sleep(0) + assert recorded == ["failure"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py index ec5b994f147..ebc831b4102 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -10,6 +10,7 @@ Covers ``_wrap_streaming_iterator_with_enrichment``, from __future__ import annotations import asyncio +from collections.abc import AsyncGenerator, AsyncIterator from datetime import datetime from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock @@ -18,12 +19,15 @@ import pytest from fastapi import HTTPException import litellm +from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( BaseAnthropicMessagesStreamingIterator, ) +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging +from litellm.types.utils import Usage @pytest.fixture(autouse=True) @@ -479,6 +483,135 @@ async def test_native_messages_stream_logging_fires_when_guardrail_blocks_after_ assert logging_obj._deferred_stream_complete_args is None +def _armed_chat_stream( + test_name: str, request_data: dict[str, object], events: list[str] +) -> tuple[LiteLLMLoggingObj, AsyncIterator[dict[str, object]]]: + """A /chat/completions stream whose CSW shape parks ``(assembled ModelResponse, cache_hit)`` + at upstream exhaustion, with the deferred dispatch recording into ``events``.""" + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id=test_name, + function_id=test_name, + ) + logging_obj.optional_params = {} + logging_obj.litellm_params = {} + logging_obj.standard_built_in_tools_params = None + + async def _dispatch_deferred_logging(*args: object) -> None: + events.append("success_dispatched") + + logging_obj._on_deferred_stream_complete = _dispatch_deferred_logging + request_data["litellm_logging_obj"] = logging_obj + + assembled = litellm.ModelResponse( + model="gpt-4o-mini", + choices=[{"index": 0, "message": {"role": "assistant", "content": "BANANA"}}], + usage=Usage(prompt_tokens=3, completion_tokens=5, total_tokens=8), + ) + + async def _upstream() -> AsyncIterator[dict[str, object]]: + yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "BAN"}}]} + yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "ANA"}}]} + logging_obj._deferred_stream_complete_args = (assembled, False) + + return logging_obj, _upstream() + + +def _raising_at_end_of_stream(error: Exception) -> CustomLogger: + class _EndOfStreamRaiser(CustomLogger): + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: + async for chunk in response: + yield chunk + raise error + + return _EndOfStreamRaiser() + + +@pytest.mark.asyncio +async def test_chat_stream_guardrail_block_after_stream_end_logs_failure_not_success( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """ + A guardrail that raises ``GuardrailRaisedException`` at end of a + /chat/completions stream must NOT dispatch the parked success logging: + the request is logged via the failure path instead, with the consumed + usage carried over so the failure row bills correctly. + """ + events: list[str] = [] + request_data: dict[str, object] = {"metadata": {}} + logging_obj, upstream = _armed_chat_stream("test_chat_stream_guardrail_block", request_data, events) + monkeypatch.setattr( + litellm, + "callbacks", + [_raising_at_end_of_stream(GuardrailRaisedException(guardrail_name="g", message="blocked"))], + ) + + with pytest.raises(GuardrailRaisedException): + async for _ in proxy_logging.async_post_call_streaming_iterator_hook( + response=upstream, + user_api_key_dict=make_user_api_key_auth(), + request_data=request_data, + ): + pass + await asyncio.sleep(0) + await asyncio.sleep(0) + + snapshot = { + "events": events, + "callback_cleared": logging_obj._on_deferred_stream_complete is None, + "args_cleared": logging_obj._deferred_stream_complete_args is None, + "combined_usage_total_tokens": logging_obj.model_call_details["combined_usage_object"].total_tokens, + "response_cost_positive": logging_obj.model_call_details["response_cost"] > 0, + } + assert snapshot == { + "events": [], + "callback_cleared": True, + "args_cleared": True, + "combined_usage_total_tokens": 8, + "response_cost_positive": True, + } + + +@pytest.mark.asyncio +async def test_chat_stream_generic_callback_error_after_stream_end_still_flushes_success_logging( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """ + ``post_call_failure_hook`` only routes proxy-level errors (HTTPException, + ProxyException, GuardrailRaisedException) through failure logging. A + callback that dies with any other exception after the stream completed + must keep flushing the parked success dispatch, or the request ends with + no terminal log at all. + """ + events: list[str] = [] + request_data: dict[str, object] = {"metadata": {}} + logging_obj, upstream = _armed_chat_stream("test_chat_stream_generic_callback_error", request_data, events) + monkeypatch.setattr(litellm, "callbacks", [_raising_at_end_of_stream(RuntimeError("callback crashed"))]) + + with pytest.raises(RuntimeError): + async for _ in proxy_logging.async_post_call_streaming_iterator_hook( + response=upstream, + user_api_key_dict=make_user_api_key_auth(), + request_data=request_data, + ): + pass + await asyncio.sleep(0) + await asyncio.sleep(0) + + snapshot = { + "events": events, + "args_cleared": logging_obj._deferred_stream_complete_args is None, + "failure_usage_recorded": "combined_usage_object" in logging_obj.model_call_details, + } + assert snapshot == {"events": ["success_dispatched"], "args_cleared": True, "failure_usage_recorded": False} + + # --------------------------------------------------------------------------- # _fire_deferred_stream_logging # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/test_count_tokens_public_api.py b/tests/test_litellm/test_count_tokens_public_api.py index 86c33c3e8f7..2918d0aa522 100644 --- a/tests/test_litellm/test_count_tokens_public_api.py +++ b/tests/test_litellm/test_count_tokens_public_api.py @@ -155,3 +155,23 @@ def test_acount_tokens_no_api_key_falls_back(monkeypatch): # Should fall back to local tokenizer since no API key assert result.total_tokens > 0 assert result.tokenizer_type == "local_tokenizer" + + +async def test_acount_tokens_local_fallback_counts_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + model = "together_ai/meta-llama/Llama-3-8b-chat-hf" + warm_tokenizer(model) + + result, took, lags = await timed_with_loop_lags( + lambda: litellm.acount_tokens(model=model, messages=[{"role": "user", "content": text * 100}]) + ) + + assert result.tokenizer_type == "local_tokenizer" + assert result.total_tokens > 100_000 + assert_loop_stayed_free(took, lags) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx index c425e766f2d..ae898645de4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx @@ -129,7 +129,7 @@ describe("BudgetTable", () => { const user = userEvent.setup(); renderWithProviders(); await showColumn(user, "created_at"); - for (const field of ["budget_id", "max_budget", "tpm_limit", "rpm_limit", "created_at"]) { + for (const field of ["budget_id", "max_budget", "tpm_limit", "rpm_limit", "tpd_limit", "created_at"]) { expect(screen.getByTestId(`sort-header-${field}`)).toBeInTheDocument(); } }); @@ -152,9 +152,10 @@ describe("BudgetTable", () => { }); it("should show n/a for missing rate limits and Unlimited for a missing max budget", () => { - const list = makeList({ rows: [makeBudget({ max_budget: null, tpm_limit: null, rpm_limit: null })] }); + const noLimits = { max_budget: null, tpm_limit: null, rpm_limit: null, tpd_limit: null }; + const list = makeList({ rows: [makeBudget(noLimits)] }); renderWithProviders(); - expect(screen.getAllByText("n/a")).toHaveLength(2); + expect(screen.getAllByText("n/a")).toHaveLength(3); expect(screen.getByText("Unlimited")).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx index e5cd9043492..fb894322208 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx @@ -126,6 +126,14 @@ export const getBudgetTableColumns = ({ size: 100, cell: ({ row }) => , }, + { + id: "tpd_limit", + accessorKey: "tpd_limit", + meta: { title: "TPD (batch)", numeric: true }, + header: ({ column }) => , + size: 110, + cell: ({ row }) => , + }, { id: "budget_duration", accessorKey: "budget_duration", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx index 492a6b5c630..5068cbed453 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx @@ -17,6 +17,7 @@ const budgetShape = { budget_id: z.string().min(1, "Please input a human-friendly name for the budget"), tpm_limit: z.number().nullish(), rpm_limit: z.number().nullish(), + tpd_limit: z.number().nullish(), max_budget: z.number().nullish(), budget_duration: z.string().nullish(), }; @@ -112,6 +113,23 @@ const BudgetModal: React.FC = ({ isModalVisible, setIsModalVis /> )} + + {({ ref, value, onChange, ...field }) => ( + onChange(event.target.value === "" ? null : event.target.valueAsNumber)} + /> + )} + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx index 25344c52847..7455c252e26 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx @@ -133,6 +133,7 @@ const BudgetPanel: React.FC = ({ accessToken }) => { { label: "Max Budget", value: selectedBudget?.max_budget }, { label: "TPM", value: selectedBudget?.tpm_limit }, { label: "RPM", value: selectedBudget?.rpm_limit }, + { label: "TPD (batch)", value: selectedBudget?.tpd_limit }, ]} onCancel={handleDeleteCancel} onOk={handleDeleteConfirm} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx index 71fce2de836..1931a88f096 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx @@ -15,13 +15,14 @@ import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/u type EditBudgetFormValues = Pick< budgetItem, - "budget_id" | "tpm_limit" | "rpm_limit" | "max_budget" | "budget_duration" + "budget_id" | "tpm_limit" | "rpm_limit" | "tpd_limit" | "max_budget" | "budget_duration" >; const toFormValues = (budget: budgetItem): EditBudgetFormValues => ({ budget_id: budget.budget_id, tpm_limit: budget.tpm_limit, rpm_limit: budget.rpm_limit, + tpd_limit: budget.tpd_limit, max_budget: budget.max_budget, budget_duration: budget.budget_duration, }); @@ -118,6 +119,23 @@ const EditBudgetModal: React.FC = ({ isModalVisible, setIs /> )} + + {({ ref, value, onChange, ...field }) => ( + onChange(event.target.value === "" ? null : event.target.valueAsNumber)} + /> + )} + diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index 70c30596c75..851c9e6d487 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -1187,6 +1187,7 @@ describe("Teams - which fields reach the create payload depends on the open sect "organization_id", "rpm_limit", "team_alias", + "tpd_limit", "tpm_limit", ]); expect(payload.team_alias).toBe("Closed Sections Team"); @@ -1314,6 +1315,7 @@ describe("Teams - the exact bytes the create call sends", () => { budget_duration: undefined, tpm_limit: undefined, rpm_limit: undefined, + tpd_limit: undefined, metadata: undefined, }); expect(wireBody(payload)).toStrictEqual({ @@ -1341,6 +1343,7 @@ describe("Teams - the exact bytes the create call sends", () => { budget_duration: undefined, tpm_limit: undefined, rpm_limit: undefined, + tpd_limit: undefined, metadata: undefined, team_id: undefined, team_member_budget: undefined, @@ -1513,6 +1516,7 @@ describe("Teams - the exact bytes the create call sends", () => { budget_duration: undefined, tpm_limit: undefined, rpm_limit: undefined, + tpd_limit: undefined, metadata: undefined, team_id: undefined, team_member_budget: undefined, diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index dc531ea5dad..4f3367d8b98 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -77,6 +77,7 @@ const teamCreateFieldsSchema = z.object({ budget_duration: z.string().nullish(), tpm_limit: numericInputSchema, rpm_limit: numericInputSchema, + tpd_limit: numericInputSchema, metadata: metadataPairsSchema.optional(), team_id: z.string().optional(), team_member_budget: z.number().optional(), @@ -113,6 +114,7 @@ const EMPTY_TEAM_CREATE_VALUES: TeamCreateFormValues = { budget_duration: undefined, tpm_limit: undefined, rpm_limit: undefined, + tpd_limit: undefined, metadata: [], team_id: undefined, team_member_budget: undefined, @@ -821,6 +823,18 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser )} + + {({ ref, value, ...field }) => ( + + )} + Metadata | null; budget_reset_at?: string | null; @@ -50,6 +51,7 @@ export interface KeyResponse { metadata: Record; tpm_limit: number; rpm_limit: number; + tpd_limit?: number | null; duration: string; budget_duration: string; budget_reset_at: string; diff --git a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts index fef94cc3c2b..7c6d5def8da 100644 --- a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts +++ b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts @@ -45,6 +45,7 @@ const DROPPED_AT_SERIALISATION = [ "rpm_limit", "tags", "throttle_on_budget_exceeded", + "tpd_limit", "tpm_limit", ]; @@ -64,6 +65,7 @@ const OPTIONAL_SETTINGS_VALUES = { tpm_limit_type: "key", rpm_limit: undefined, rpm_limit_type: "key", + tpd_limit: undefined, throttle_on_budget_exceeded: undefined, enable_prompt_caching: undefined, guardrails: undefined, @@ -456,6 +458,18 @@ describe("budget duration", () => { }); }); +describe("tpd_limit", () => { + it("forwards the daily batch token budget alongside the minute limits", () => { + expect(payloadOf(build({ key_alias: "my-key", tpd_limit: 250000, rpm_limit: 5 }))).toStrictEqual( + aliasOnly({ tpd_limit: 250000, rpm_limit: 5 }), + ); + }); + + it("keeps a zero tpd_limit rather than treating it as unset", () => { + expect(payloadOf(build({ key_alias: "my-key", tpd_limit: 0 }))).toStrictEqual(aliasOnly({ tpd_limit: 0 })); + }); +}); + describe("purity", () => { it("leaves the submitted form values untouched", () => { const values = { @@ -499,9 +513,9 @@ describe("serialised wire shape", () => { expect(payloadOf(build({ ...CLOSED_SECTIONS_VALUES, team_id: "team-1" })).team_id).toBe("team-1"); }); - it("adds fifteen keys to the object and only the two limit types to the wire when Optional Settings opens", () => { + it("adds sixteen keys to the object and only the two limit types to the wire when Optional Settings opens", () => { const payload = payloadOf(build(OPTIONAL_SETTINGS_VALUES)); - expect(Object.keys(payload)).toHaveLength(23); + expect(Object.keys(payload)).toHaveLength(24); expect(wireKeys(payload)).toStrictEqual([ "team_id", "key_alias", diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index 0d5d9f5ec8d..3e3c29e330d 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -143,6 +143,7 @@ const OPTIONAL_OPEN_PAYLOAD = { tpm_limit_type: null, rpm_limit: undefined, rpm_limit_type: null, + tpd_limit: undefined, throttle_on_budget_exceeded: undefined, enable_prompt_caching: undefined, guardrails: undefined, @@ -395,6 +396,7 @@ describe("CreateKey", () => { it.each([ ["Tokens per minute Limit (TPM)", "tpm_limit"], ["Requests per minute Limit (RPM)", "rpm_limit"], + ["Tokens per day Limit (TPD)", "tpd_limit"], ])("routes a typed %s into the %s payload key", async (label, key) => { await openModal(); await nameTheKey(); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index b5789101f77..b8ea8de7f59 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1150,6 +1150,32 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp /> )} + + Tokens per day Limit (TPD){" "} + + + + + } + name="tpd_limit" + help={`TPD cannot exceed team TPD limit: ${team?.tpd_limit !== null && team?.tpd_limit !== undefined ? team?.tpd_limit : "unlimited"}`} + rules={ceilingRule( + team?.tpd_limit, + (limit) => `TPD limit cannot exceed team TPD limit: ${limit}`, + )} + > + {(control) => ( + + )} + @@ -1760,6 +1786,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp "budget_duration", "tpm_limit", "rpm_limit", + "tpd_limit", ...(disableCustomApiKeys ? ["key"] : []), ]} /> diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index a25ca28651a..eb912ffa3cc 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -1968,6 +1968,7 @@ describe("TeamInfoView - the exact bytes the update call sends", () => { models: ["gpt-4"], tpm_limit: 1000, rpm_limit: 1000, + tpd_limit: null, model_tpm_limit: {}, model_rpm_limit: {}, max_budget: 100, diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index ffc83d0165e..ce705008678 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -264,6 +264,7 @@ export interface TeamData { metadata: Record; tpm_limit: number | null; rpm_limit: number | null; + tpd_limit?: number | null; max_budget: number | null; soft_budget?: number | null; budget_duration: string | null; @@ -330,6 +331,7 @@ const teamUpdateFieldsSchema = z.object({ budget_duration: z.string().nullish(), tpm_limit: numericInputSchema, rpm_limit: numericInputSchema, + tpd_limit: numericInputSchema, modelLimits: z .array( z.object({ @@ -411,6 +413,7 @@ const EMPTY_TEAM_UPDATE_VALUES: TeamUpdateFormValues = { budget_duration: undefined, tpm_limit: undefined, rpm_limit: undefined, + tpd_limit: undefined, modelLimits: [], default_estimated_output_tokens: undefined, default_estimated_output_tokens_per_model: "", @@ -460,6 +463,7 @@ const toTeamFormValues = (info: TeamInfoRecord, effectiveGuardrails: string[]): budget_duration: info.budget_duration, tpm_limit: info.tpm_limit, rpm_limit: info.rpm_limit, + tpd_limit: info.tpd_limit, modelLimits: Array.from( new Set([ ...Object.keys(info.metadata?.model_tpm_limit ?? {}), @@ -918,6 +922,7 @@ const TeamInfoView: React.FC = ({ models: normalizeTeamModelSelection(values.models), tpm_limit: sanitizeNumeric(values.tpm_limit), rpm_limit: sanitizeNumeric(values.rpm_limit), + tpd_limit: sanitizeNumeric(values.tpd_limit), model_tpm_limit: modelTpmLimit, model_rpm_limit: modelRpmLimit, max_budget: values.max_budget, @@ -1168,6 +1173,7 @@ const TeamInfoView: React.FC = ({

TPM: {info.tpm_limit ?? "Unlimited"}

RPM: {info.rpm_limit ?? "Unlimited"}

+

TPD (batch): {info.tpd_limit ?? "Unlimited"}

{info.max_parallel_requests &&

Max Parallel Requests: {info.max_parallel_requests}

} {(() => { const modelTpm = (info.metadata?.model_tpm_limit ?? {}) as Record; @@ -1538,6 +1544,17 @@ const TeamInfoView: React.FC = ({ {({ ref, value, ...field }) => } + + {({ ref, value, ...field }) => } + + Metadata = ({

Rate Limits

TPM: {info.tpm_limit ?? "Unlimited"}
RPM: {info.rpm_limit ?? "Unlimited"}
+
TPD (batch): {info.tpd_limit ?? "Unlimited"}
{(() => { const modelTpm = (info.metadata?.model_tpm_limit ?? {}) as Record; const modelRpm = (info.metadata?.model_rpm_limit ?? {}) as Record; diff --git a/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx b/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx index 39542882798..99220af7c6f 100644 --- a/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx +++ b/ui/litellm-dashboard/src/components/templates/KeyEditViewControls.tsx @@ -7,6 +7,7 @@ import { CircleHelp } from "lucide-react"; import { FormField } from "@/components/shared/form/FormField"; import { toast } from "@/lib/toast"; import AgentSelector from "../agent_management/AgentSelector"; +import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"; import NumericalInput from "../shared/numerical_input"; import SkillSelector from "../skills/SkillSelector"; import { moveTagsOutOfMetadataJson } from "./keyEditFieldNormalizers"; @@ -61,6 +62,51 @@ export const KeyTypeSelect = ({ const SKILLS_HINT = "Enabled skills are visible to every key. Grant disabled (private) Claude Code plugins to this key here."; +const TPD_HINT = + "Daily token budget for batch submissions (/v1/batches). When set, batch input files are charged against this 24h window instead of the key's TPM/RPM limits. Online requests keep using TPM/RPM."; + +export const KeyRateLimitFields = ({ control }: { control: Control }) => ( + <> + + {({ ref: _ref, ...field }) => } + + + + {({ value, onChange, id }) => ( + + )} + + + + {({ ref: _ref, ...field }) => } + + + + {({ value, onChange, id }) => ( + + )} + + + + {({ ref: _ref, ...field }) => } + + +); + export const KeyAgentAndSkillFields = ({ control, accessToken, diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.test.ts b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.test.ts index 948ed659cd5..f12088bea19 100644 --- a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.test.ts +++ b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.test.ts @@ -1,8 +1,30 @@ import { describe, expect, it } from "vitest"; -import { keyEditFormSchema } from "./keyEditFormValues"; +import type { KeyResponse } from "../key_team_helpers/key_list"; +import { keyEditFormSchema, toKeyEditFormValues, toSubmittedValues } from "./keyEditFormValues"; const parse = (values: Record) => keyEditFormSchema.safeParse(values); +describe("tpd_limit round trip", () => { + const keyData = { token: "tok", models: [], rpm_limit: 5, tpd_limit: 250000 } as unknown as KeyResponse; + + it("hydrates the stored daily batch budget into the edit form", () => { + expect(toKeyEditFormValues(keyData)).toMatchObject({ rpm_limit: 5, tpd_limit: 250000 }); + }); + + it("submits tpd_limit next to the minute limits", () => { + const submitted = toSubmittedValues(toKeyEditFormValues(keyData), { canViewPolicies: true, canViewPrompts: true }); + expect(submitted).toMatchObject({ rpm_limit: 5, tpd_limit: 250000 }); + }); + + it("submits null when the operator cleared tpd_limit", () => { + const submitted = toSubmittedValues( + { ...toKeyEditFormValues(keyData), tpd_limit: null }, + { canViewPolicies: true, canViewPrompts: true }, + ); + expect(submitted.tpd_limit).toBeNull(); + }); +}); + describe("keyEditFormSchema", () => { it("accepts an empty form", () => { expect(parse({}).success).toBe(true); diff --git a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts index 233b58b48ab..7436380d6ee 100644 --- a/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts +++ b/ui/litellm-dashboard/src/components/templates/keyEditFormValues.ts @@ -28,6 +28,7 @@ export interface KeyEditFormValues { tpm_limit_type?: string | null; rpm_limit?: number | string | null; rpm_limit_type?: string | null; + tpd_limit?: number | string | null; throttle_on_budget_exceeded?: boolean; enable_prompt_caching?: boolean; max_parallel_requests?: number | string | null; @@ -77,6 +78,7 @@ export const toKeyEditFormValues = (keyData: KeyResponse): KeyEditFormValues => tpm_limit_type: (keyData as { tpm_limit_type?: string | null }).tpm_limit_type ?? null, rpm_limit: keyData.rpm_limit, rpm_limit_type: (keyData as { rpm_limit_type?: string | null }).rpm_limit_type ?? null, + tpd_limit: keyData.tpd_limit, throttle_on_budget_exceeded: Boolean(readMetadata(keyData, "throttle_on_budget_exceeded")), enable_prompt_caching: Boolean(readMetadata(keyData, "enable_prompt_caching")), max_parallel_requests: keyData.max_parallel_requests, @@ -130,6 +132,7 @@ export const keyEditFormSchema = z.object({ tpm_limit_type: z.custom(), rpm_limit: z.custom(), rpm_limit_type: z.custom(), + tpd_limit: z.custom(), throttle_on_budget_exceeded: z.custom(), enable_prompt_caching: z.custom(), max_parallel_requests: z.custom(), @@ -184,6 +187,7 @@ export const toSubmittedValues = ( tpm_limit_type: values.tpm_limit_type, rpm_limit: values.rpm_limit, rpm_limit_type: values.rpm_limit_type, + tpd_limit: values.tpd_limit, throttle_on_budget_exceeded: values.throttle_on_budget_exceeded, enable_prompt_caching: values.enable_prompt_caching, max_parallel_requests: values.max_parallel_requests, diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx index d3a7bf10c2f..9efcff04832 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.integration.test.tsx @@ -188,6 +188,7 @@ describe("KeyEditView", () => { }, tpm_limit: 10, rpm_limit: 10, + tpd_limit: 250000, duration: "30d", budget_duration: "30d", budget_reset_at: "never", @@ -1986,6 +1987,7 @@ describe("KeyEditView", () => { tpm_limit_type: null, rpm_limit: 10, rpm_limit_type: null, + tpd_limit: 250000, throttle_on_budget_exceeded: false, enable_prompt_caching: false, max_parallel_requests: 10, 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 022e85e51d7..6a94cbcf2a0 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -20,7 +20,6 @@ import BudgetDurationDropdown from "../common_components/budget_duration_dropdow import { mapInternalToDisplayNames } from "../callback_info_helpers"; import KeyLifecycleSettings from "../common_components/KeyLifecycleSettings"; import PassThroughRoutesSelector from "../common_components/PassThroughRoutesSelector"; -import RateLimitTypeFormItem from "../common_components/RateLimitTypeFormItem"; import OrganizationDropdown from "../common_components/OrganizationDropdown"; import RouterSettingsAccordion, { RouterSettingsAccordionRef } from "../common_components/RouterSettingsAccordion"; import { routerSettingsEditorValue, routerSettingsUpdate } from "../common_components/routerSettingsPayload"; @@ -35,6 +34,7 @@ import { KeyAgentAndSkillFields, KeyBudgetNumberField, KeyMetadataField, + KeyRateLimitFields, KeyTypeSelect, labelWithHint, moveMetadataTagsToTagsField, @@ -484,39 +484,7 @@ export function KeyEditView({ /> - - {({ ref: _ref, ...field }) => } - - - - {({ value, onChange, id }) => ( - - )} - - - - {({ ref: _ref, ...field }) => } - - - - {({ value, onChange, id }) => ( - - )} - + RPM: {currentKeyData.rpm_limit !== null ? currentKeyData.rpm_limit : "Unlimited"}

+

TPD (batch): {currentKeyData.tpd_limit ?? "Unlimited"}

{Boolean(currentKeyData.metadata?.throttle_on_budget_exceeded) && (

Throttle on budget exceeded: Yes

)} @@ -1064,6 +1066,7 @@ export default function KeyInfoView({

RPM: {currentKeyData.rpm_limit !== null ? currentKeyData.rpm_limit : "Unlimited"}

+

TPD (batch): {currentKeyData.tpd_limit ?? "Unlimited"}

Max Parallel Requests:{" "} {currentKeyData.max_parallel_requests !== null diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index ba7909b3ab1..4c16a8613eb 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -1843,6 +1843,7 @@ export interface paths { * - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. * - tpm_limit: Optional[int] - The tokens per minute limit for the budget. * - rpm_limit: Optional[int] - The requests per minute limit for the budget. + * - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit. * - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} * - budget_reset_at: Optional[datetime] - Datetime when the initial budget is reset. Default is now. */ @@ -1899,6 +1900,7 @@ export interface paths { * - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. * - tpm_limit: Optional[int] - The tokens per minute limit for the budget. * - rpm_limit: Optional[int] - The requests per minute limit for the budget. + * - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit. * - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} * - budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset. */ @@ -3980,6 +3982,7 @@ export interface paths { * - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). * - tpm_limit: Optional[int] - [Not Implemented Yet] Specify tpm limit for a given customer (Tokens per minute) * - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute) + * - tpd_limit: Optional[int] - Specify tpd limit for a given customer (Tokens per day). Batch submissions are charged against it instead of tpm_limit/rpm_limit * - model_max_budget: Optional[dict] - [Not Implemented Yet] Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}} * - max_parallel_requests: Optional[int] - [Not Implemented Yet] Specify max parallel requests for a given customer. * - soft_budget: Optional[float] - [Not Implemented Yet] Get alerts when customer crosses given budget, doesn't block requests. @@ -4514,6 +4517,7 @@ export interface paths { * - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). * - tpm_limit: Optional[int] - [Not Implemented Yet] Specify tpm limit for a given customer (Tokens per minute) * - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute) + * - tpd_limit: Optional[int] - Specify tpd limit for a given customer (Tokens per day). Batch submissions are charged against it instead of tpm_limit/rpm_limit * - model_max_budget: Optional[dict] - [Not Implemented Yet] Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}} * - max_parallel_requests: Optional[int] - [Not Implemented Yet] Specify max parallel requests for a given customer. * - soft_budget: Optional[float] - [Not Implemented Yet] Get alerts when customer crosses given budget, doesn't block requests. @@ -7736,6 +7740,7 @@ export interface paths { * - blocked: Optional[bool] - Whether the key is blocked. * - rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute) * - tpm_limit: Optional[int] - Specify tpm limit for a given key (Tokens per minute) + * - tpd_limit: Optional[int] - Specify tpd limit for a given key (Tokens per day). Charged by batch submissions instead of tpm_limit/rpm_limit. * - soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached. * - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). * - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. @@ -8049,6 +8054,7 @@ export interface paths { * - blocked: Optional[bool] - Whether the key is blocked. * - rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute) * - tpm_limit: Optional[int] - Specify tpm limit for a given key (Tokens per minute) + * - tpd_limit: Optional[int] - Specify tpd limit for a given key (Tokens per day). Charged by batch submissions instead of tpm_limit/rpm_limit. * - soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached. * - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). * - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests) @@ -8175,6 +8181,7 @@ export interface paths { * - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"} * - tpm_limit: Optional[int] - Tokens per minute limit * - rpm_limit: Optional[int] - Requests per minute limit + * - tpd_limit: Optional[int] - Tokens per day limit, charged by batch submissions instead of tpm_limit/rpm_limit * - model_rpm_limit: Optional[dict] - Model-specific RPM limits {"gpt-4": 100, "claude-v1": 200} * - 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. @@ -8424,7 +8431,7 @@ export interface paths { * way to page, sort or filter it. * * `sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`, - * `rpm_limit` or `created_at`, each optionally prefixed with `-` for descending, + * `rpm_limit`, `tpd_limit` or `created_at`, each optionally prefixed with `-` for descending, * and defaults to `-created_at`. `budget_id` is appended to every sort as the * tiebreaker. `q` is a case-insensitive substring match on `budget_id`. * `page_size` defaults to 50 and is capped at 100. Filters are @@ -10619,6 +10626,7 @@ export interface paths { * - max_budget: *Optional[float]* - Max budget for org * - tpm_limit: *Optional[int]* - Max tpm limit for org * - rpm_limit: *Optional[int]* - Max rpm limit for org + * - tpd_limit: *Optional[int]* - Max tokens per day stored on the org budget. Batch submissions enforce tpd_limit at the key, team and end user scopes only. * - model_rpm_limit: *Optional[Dict[str, int]]* - The RPM (Requests Per Minute) limit per model for this organization. * - model_tpm_limit: *Optional[Dict[str, int]]* - The TPM (Tokens Per Minute) limit per model for this organization. * - max_parallel_requests: *Optional[int]* - [Not Implemented Yet] Max parallel requests for org @@ -15605,6 +15613,7 @@ export interface paths { * - 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 + * - tpd_limit: Optional[int] - The TPD (Tokens Per Day) limit for this team. Batch submissions are charged against it instead of tpm_limit/rpm_limit * - rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of RPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating RPM, or "best_effort_throughput" for best effort enforcement. * - tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of TPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating TPM, or "best_effort_throughput" for best effort enforcement. * - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget @@ -15833,6 +15842,7 @@ export interface paths { * - metadata: Optional[dict] - Metadata for team, store information for team. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - 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 + * - tpd_limit: Optional[int] - The TPD (Tokens Per Day) limit for this team. Batch submissions are charged against it instead of tpm_limit/rpm_limit * - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget * - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set (either in the request or existing), soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set. * - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets) @@ -24582,6 +24592,8 @@ export interface components { rpm_limit?: number | null; /** Soft Budget */ soft_budget?: number | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** @@ -24634,6 +24646,11 @@ export interface components { * @description Requests will NOT fail if this is exceeded. Will fire alerting though. */ soft_budget?: number | null; + /** + * Tpd Limit + * @description Max tokens per day, charged by batch submissions, allowed for this budget id. + */ + tpd_limit?: number | null; /** * Tpm Limit * @description Max tokens per minute, allowed for this budget id. @@ -28335,6 +28352,8 @@ export interface components { team_id?: string | null; /** Throttle On Budget Exceeded */ throttle_on_budget_exceeded?: boolean | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Tpm Limit Type */ @@ -28495,6 +28514,8 @@ export interface components { token?: string | null; /** Token Id */ token_id?: string | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Tpm Limit Type */ @@ -29226,6 +29247,8 @@ export interface components { rpm_limit?: number | null; /** Soft Budget */ soft_budget?: number | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; }; @@ -29259,6 +29282,8 @@ export interface components { rpm_limit?: number | null; /** Soft Budget */ soft_budget?: number | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; }; @@ -29367,6 +29392,8 @@ export interface components { team_id: string; /** Team Member Permissions */ team_member_permissions?: string[] | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Updated At */ @@ -29534,6 +29561,8 @@ export interface components { team_id?: string | null; /** Token */ token?: string | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Updated At */ @@ -30727,6 +30756,8 @@ export interface components { team_id: string; /** Team Member Permissions */ team_member_permissions?: string[] | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Updated At */ @@ -31103,6 +31134,8 @@ export interface components { team_id?: string | null; /** Token */ token?: string | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Updated At */ @@ -32752,6 +32785,11 @@ export interface components { soft_budget?: number | null; /** Spend */ spend?: number | null; + /** + * Tpd Limit + * @description Max tokens per day, charged by batch submissions, allowed for this budget id. + */ + tpd_limit?: number | null; /** * Tpm Limit * @description Max tokens per minute, allowed for this budget id. @@ -32967,6 +33005,8 @@ export interface components { rpm_limit?: number | null; /** Soft Budget */ soft_budget?: number | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; }; @@ -33088,6 +33128,8 @@ export interface components { tags?: string[] | null; /** Team Id */ team_id: string; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; }; @@ -33272,6 +33314,8 @@ export interface components { team_member_rpm_limit?: number | null; /** Team Member Tpm Limit */ team_member_tpm_limit?: number | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Tpm Limit Type */ @@ -33565,6 +33609,8 @@ export interface components { token?: string | null; /** Token Id */ token_id?: string | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Tpm Limit Type */ @@ -34019,6 +34065,8 @@ export interface components { team_member_rpm_limit?: number | null; /** Team Member Tpm Limit */ team_member_tpm_limit?: number | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; }; @@ -35428,6 +35476,8 @@ export interface components { team_id?: string | null; /** Throttle On Budget Exceeded */ throttle_on_budget_exceeded?: boolean | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Tpm Limit Type */ @@ -37421,6 +37471,8 @@ export interface components { team_id: string; /** Team Member Permissions */ team_member_permissions?: string[] | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Updated At */ @@ -37561,6 +37613,8 @@ export interface components { team_id: string; /** Team Member Permissions */ team_member_permissions?: string[] | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Updated At */ @@ -38667,6 +38721,8 @@ export interface components { temp_budget_increase?: number | null; /** Throttle On Budget Exceeded */ throttle_on_budget_exceeded?: boolean | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Tpm Limit Type */ @@ -38931,6 +38987,8 @@ export interface components { tags?: string[] | null; /** Team Id */ team_id?: string | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; }; @@ -39130,6 +39188,8 @@ export interface components { team_member_rpm_limit?: number | null; /** Team Member Tpm Limit */ team_member_tpm_limit?: number | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; }; @@ -39640,6 +39700,8 @@ export interface components { end_user_object_permission?: components["schemas"]["LiteLLM_ObjectPermissionTable"] | null; /** End User Rpm Limit */ end_user_rpm_limit?: number | null; + /** End User Tpd Limit */ + end_user_tpd_limit?: number | null; /** End User Tpm Limit */ end_user_tpm_limit?: number | null; /** Expires */ @@ -39808,10 +39870,14 @@ export interface components { team_soft_budget?: number | null; /** Team Spend */ team_spend?: number | null; + /** Team Tpd Limit */ + team_tpd_limit?: number | null; /** Team Tpm Limit */ team_tpm_limit?: number | null; /** Token */ token?: string | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Tpm Limit Per Model */