Merge remote-tracking branch 'origin/main' into litellm_fix_xai_chat_web_search_live_search

This commit is contained in:
yassin 2026-09-15 20:46:51 +00:00
commit 113e43b87a
74 changed files with 1602 additions and 90 deletions

View file

@ -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;

View file

@ -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?

View file

@ -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,
)

View file

@ -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

View file

@ -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,

View file

@ -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,

View file

@ -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")

View file

@ -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,

View file

@ -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,

View file

@ -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,

View file

@ -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,

View file

@ -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

View file

@ -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,

View file

@ -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,

View file

@ -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

View file

@ -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

View file

@ -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 = []

View file

@ -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

View file

@ -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,

View file

@ -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,

View file

@ -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

View file

@ -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,

View file

@ -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

View file

@ -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"},

View file

@ -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.

View file

@ -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,

View file

@ -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

View file

@ -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

View file

@ -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)

View file

@ -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

View file

@ -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?

View file

@ -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,

View file

@ -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?

View file

@ -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()

View file

@ -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)

View file

@ -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)

View file

@ -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)

View file

@ -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)

View file

@ -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)

View file

@ -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)

View file

@ -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)

View file

@ -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)

View file

@ -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

View file

@ -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

View file

@ -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)]

View file

@ -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")

View file

@ -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

View file

@ -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": [],

View file

@ -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():
"""

View file

@ -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():
"""

View file

@ -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)

View file

@ -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"]

View file

@ -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
# ---------------------------------------------------------------------------

View file

@ -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)

View file

@ -129,7 +129,7 @@ describe("BudgetTable", () => {
const user = userEvent.setup();
renderWithProviders(<BudgetTable {...defaultProps} list={makeList()} />);
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(<BudgetTable {...defaultProps} list={list} />);
expect(screen.getAllByText("n/a")).toHaveLength(2);
expect(screen.getAllByText("n/a")).toHaveLength(3);
expect(screen.getByText("Unlimited")).toBeInTheDocument();
});

View file

@ -126,6 +126,14 @@ export const getBudgetTableColumns = ({
size: 100,
cell: ({ row }) => <RateLimitCell value={row.original.rpm_limit} />,
},
{
id: "tpd_limit",
accessorKey: "tpd_limit",
meta: { title: "TPD (batch)", numeric: true },
header: ({ column }) => <DataTableSortHeader column={column} title="TPD (batch)" />,
size: 110,
cell: ({ row }) => <RateLimitCell value={row.original.tpd_limit} />,
},
{
id: "budget_duration",
accessorKey: "budget_duration",

View file

@ -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<BudgetModalProps> = ({ isModalVisible, setIsModalVis
/>
)}
</FormField>
<FormField
control={form.control}
name="tpd_limit"
label="Max Tokens per day (batch)"
description="Daily token budget for batch submissions. When set, batches are charged against this instead of TPM/RPM."
>
{({ ref, value, onChange, ...field }) => (
<Input
{...field}
ref={ref}
type="number"
step={1}
value={value ?? ""}
onChange={(event) => onChange(event.target.value === "" ? null : event.target.valueAsNumber)}
/>
)}
</FormField>
<Collapsible open={optionalSettingsOpen} onOpenChange={setOptionalSettingsOpen} className="mt-20 mb-8">
<CollapsibleTrigger className="group flex w-full items-center justify-between py-2 text-left">

View file

@ -133,6 +133,7 @@ const BudgetPanel: React.FC<BudgetSettingsPageProps> = ({ 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}

View file

@ -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<EditBudgetModalProps> = ({ isModalVisible, setIs
/>
)}
</FormField>
<FormField
control={form.control}
name="tpd_limit"
label="Max Tokens per day (batch)"
description="Daily token budget for batch submissions. When set, batches are charged against this instead of TPM/RPM."
>
{({ ref, value, onChange, ...field }) => (
<Input
{...field}
ref={ref}
type="number"
step={1}
value={value ?? ""}
onChange={(event) => onChange(event.target.value === "" ? null : event.target.valueAsNumber)}
/>
)}
</FormField>
<Collapsible open={optionalSettingsOpen} onOpenChange={setOptionalSettingsOpen} className="mt-20 mb-8">
<CollapsibleTrigger className="group flex w-full items-center justify-between py-2 text-left">

View file

@ -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,

View file

@ -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<TeamProps> = ({ accessToken, userID, userRole, premiumUser
<NumericalInput {...field} ref={ref} value={value ?? ""} step={1} width={400} />
)}
</FormField>
<FormField
control={form.control}
name="tpd_limit"
label={labelWithHint(
"Tokens per day Limit (TPD)",
"Daily token budget for batch submissions (/v1/batches). When set, batch input files are charged against this 24h window instead of the team's TPM/RPM limits. Online requests keep using TPM/RPM.",
)}
>
{({ ref, value, ...field }) => (
<NumericalInput {...field} ref={ref} value={value ?? ""} step={1} width={400} />
)}
</FormField>
<Field>
<FieldLabel>Metadata</FieldLabel>
<MetadataKeyValueFields

View file

@ -12,6 +12,7 @@ export interface Team {
budget_duration: string | null;
tpm_limit: number | null;
rpm_limit: number | null;
tpd_limit?: number | null;
organization_id: string;
metadata?: Record<string, unknown> | null;
budget_reset_at?: string | null;
@ -50,6 +51,7 @@ export interface KeyResponse {
metadata: Record<string, unknown>;
tpm_limit: number;
rpm_limit: number;
tpd_limit?: number | null;
duration: string;
budget_duration: string;
budget_reset_at: string;

View file

@ -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",

View file

@ -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();

View file

@ -1150,6 +1150,32 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
/>
)}
</MountedFormField>
<MountedFormField
className="mt-4"
label={
<span>
Tokens per day Limit (TPD){" "}
<SimpleTooltip content="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.">
<Info className="ml-1 inline size-3.5 align-text-bottom" />
</SimpleTooltip>
</span>
}
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) => (
<NumericalInput
{...control}
value={control.value as number | string | undefined}
step={1}
width={400}
/>
)}
</MountedFormField>
<Field className="mt-4">
<FieldLabel>
<span>
@ -1760,6 +1786,7 @@ const CreateKey: React.FC<CreateKeyProps> = ({ team, teams, data, addKey, autoOp
"budget_duration",
"tpm_limit",
"rpm_limit",
"tpd_limit",
...(disableCustomApiKeys ? ["key"] : []),
]}
/>

View file

@ -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,

View file

@ -264,6 +264,7 @@ export interface TeamData {
metadata: Record<string, any>;
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<TeamInfoProps> = ({
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<TeamInfoProps> = ({
<div className="mt-2">
<p>TPM: {info.tpm_limit ?? "Unlimited"}</p>
<p>RPM: {info.rpm_limit ?? "Unlimited"}</p>
<p>TPD (batch): {info.tpd_limit ?? "Unlimited"}</p>
{info.max_parallel_requests && <p>Max Parallel Requests: {info.max_parallel_requests}</p>}
{(() => {
const modelTpm = (info.metadata?.model_tpm_limit ?? {}) as Record<string, number>;
@ -1538,6 +1544,17 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
{({ ref, value, ...field }) => <NumericalInput {...field} ref={ref} value={value ?? ""} step={1} />}
</FormField>
<FormField
control={form.control}
name="tpd_limit"
label={labelWithHint(
"Tokens per day Limit (TPD)",
"Daily token budget for batch submissions (/v1/batches). When set, batch input files are charged against this 24h window instead of the team's TPM/RPM limits. Online requests keep using TPM/RPM.",
)}
>
{({ ref, value, ...field }) => <NumericalInput {...field} ref={ref} value={value ?? ""} step={1} />}
</FormField>
<Field>
<FieldLabel>Metadata</FieldLabel>
<MetadataKeyValueFields
@ -1997,6 +2014,7 @@ const TeamInfoView: React.FC<TeamInfoProps> = ({
<p className="font-medium">Rate Limits</p>
<div>TPM: {info.tpm_limit ?? "Unlimited"}</div>
<div>RPM: {info.rpm_limit ?? "Unlimited"}</div>
<div>TPD (batch): {info.tpd_limit ?? "Unlimited"}</div>
{(() => {
const modelTpm = (info.metadata?.model_tpm_limit ?? {}) as Record<string, number>;
const modelRpm = (info.metadata?.model_rpm_limit ?? {}) as Record<string, number>;

View file

@ -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<KeyEditFormValues> }) => (
<>
<FormField control={control} name="tpm_limit" label="TPM Limit">
{({ ref: _ref, ...field }) => <NumericalInput {...field} value={field.value ?? ""} min={0} />}
</FormField>
<FormField control={control} name="tpm_limit_type">
{({ value, onChange, id }) => (
<RateLimitTypeFormItem
id={id}
type="tpm"
name="tpm_limit_type"
showDetailedDescriptions={false}
value={value as string | null}
onChange={onChange}
/>
)}
</FormField>
<FormField control={control} name="rpm_limit" label="RPM Limit">
{({ ref: _ref, ...field }) => <NumericalInput {...field} value={field.value ?? ""} min={0} />}
</FormField>
<FormField control={control} name="rpm_limit_type">
{({ value, onChange, id }) => (
<RateLimitTypeFormItem
id={id}
type="rpm"
name="rpm_limit_type"
showDetailedDescriptions={false}
value={value as string | null}
onChange={onChange}
/>
)}
</FormField>
<FormField control={control} name="tpd_limit" label={labelWithHint("TPD Limit (batch)", TPD_HINT)}>
{({ ref: _ref, ...field }) => <NumericalInput {...field} value={field.value ?? ""} min={0} />}
</FormField>
</>
);
export const KeyAgentAndSkillFields = ({
control,
accessToken,

View file

@ -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<string, unknown>) => 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);

View file

@ -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<string | null | undefined>(),
rpm_limit: z.custom<number | string | null | undefined>(),
rpm_limit_type: z.custom<string | null | undefined>(),
tpd_limit: z.custom<number | string | null | undefined>(),
throttle_on_budget_exceeded: z.custom<boolean | undefined>(),
enable_prompt_caching: z.custom<boolean | undefined>(),
max_parallel_requests: z.custom<number | string | null | undefined>(),
@ -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,

View file

@ -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,

View file

@ -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({
/>
</Field>
<FormField control={form.control} name="tpm_limit" label="TPM Limit">
{({ ref: _ref, ...field }) => <NumericalInput {...field} value={field.value ?? ""} min={0} />}
</FormField>
<FormField control={form.control} name="tpm_limit_type">
{({ value, onChange, id }) => (
<RateLimitTypeFormItem
id={id}
type="tpm"
name="tpm_limit_type"
showDetailedDescriptions={false}
value={value as string | null}
onChange={onChange}
/>
)}
</FormField>
<FormField control={form.control} name="rpm_limit" label="RPM Limit">
{({ ref: _ref, ...field }) => <NumericalInput {...field} value={field.value ?? ""} min={0} />}
</FormField>
<FormField control={form.control} name="rpm_limit_type">
{({ value, onChange, id }) => (
<RateLimitTypeFormItem
id={id}
type="rpm"
name="rpm_limit_type"
showDetailedDescriptions={false}
value={value as string | null}
onChange={onChange}
/>
)}
</FormField>
<KeyRateLimitFields control={form.control} />
<FormField
control={form.control}

View file

@ -288,6 +288,7 @@ export default function KeyInfoView({
formValues.max_budget = mapEmptyStringToNull(formValues.max_budget);
formValues.tpm_limit = mapEmptyStringToNull(formValues.tpm_limit);
formValues.rpm_limit = mapEmptyStringToNull(formValues.rpm_limit);
formValues.tpd_limit = mapEmptyStringToNull(formValues.tpd_limit);
formValues.max_parallel_requests = mapEmptyStringToNull(formValues.max_parallel_requests);
// Convert metadata back to an object if it exists and is a string
@ -688,6 +689,7 @@ export default function KeyInfoView({
<p className="text-sm">
RPM: {currentKeyData.rpm_limit !== null ? currentKeyData.rpm_limit : "Unlimited"}
</p>
<p className="text-sm">TPD (batch): {currentKeyData.tpd_limit ?? "Unlimited"}</p>
{Boolean(currentKeyData.metadata?.throttle_on_budget_exceeded) && (
<p className="text-sm">Throttle on budget exceeded: Yes</p>
)}
@ -1064,6 +1066,7 @@ export default function KeyInfoView({
<p className="text-sm">
RPM: {currentKeyData.rpm_limit !== null ? currentKeyData.rpm_limit : "Unlimited"}
</p>
<p className="text-sm">TPD (batch): {currentKeyData.tpd_limit ?? "Unlimited"}</p>
<p className="text-sm">
Max Parallel Requests:{" "}
{currentKeyData.max_parallel_requests !== null

View file

@ -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 */