mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-15 23:31:29 +00:00
Merge remote-tracking branch 'origin/main' into litellm_fix_xai_chat_web_search_live_search
This commit is contained in:
commit
113e43b87a
74 changed files with 1602 additions and 90 deletions
|
|
@ -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;
|
||||
|
|
@ -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?
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 = []
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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"},
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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?
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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?
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
259
tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py
Normal file
259
tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py
Normal 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)]
|
||||
|
|
@ -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")
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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": [],
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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():
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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"]
|
||||
|
|
|
|||
|
|
@ -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
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
});
|
||||
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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">
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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">
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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;
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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();
|
||||
|
|
|
|||
|
|
@ -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"] : []),
|
||||
]}
|
||||
/>
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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>;
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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);
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
68
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
68
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -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 */
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue