diff --git a/CLAUDE.md b/CLAUDE.md index 41678432989..b9753ab864b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -25,6 +25,8 @@ Same thing for bug fixes. The tests should make it so that this specific bug can Never test structure of code only function of it +A test must only fail when litellm code changes. Never pin facts we don't own (a vendor's price, a third party's field, an upstream default, today's date) as literals or as "X must be absent"; assert the invariant our code guarantees instead, e.g. two rows agree, a value is within range, a field is derived from another. If an outside fact is truly load-bearing, cite its source and date next to the assertion so a reader can tell stale from broken + `tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_.py` if you're the first test there). One focused regression test beats many shallow ones End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md` diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260913000000_add_tpd_limit/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260913000000_add_tpd_limit/migration.sql new file mode 100644 index 00000000000..cdf8f4975c1 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260913000000_add_tpd_limit/migration.sql @@ -0,0 +1,14 @@ +-- AlterTable +ALTER TABLE "LiteLLM_BudgetTable" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT; + +-- AlterTable +ALTER TABLE "LiteLLM_TeamTable" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT; + +-- AlterTable +ALTER TABLE "LiteLLM_DeletedTeamTable" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT; + +-- AlterTable +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT; + +-- AlterTable +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "tpd_limit" BIGINT; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index dd7967aafe3..8072df5aa5b 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -17,6 +17,7 @@ model LiteLLM_BudgetTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? model_max_budget Json? budget_duration String? budget_reset_at DateTime? @@ -133,6 +134,7 @@ model LiteLLM_TeamTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? blocked Boolean @default(false) @@ -203,6 +205,7 @@ model LiteLLM_DeletedTeamTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? blocked Boolean @default(false) @@ -438,6 +441,7 @@ model LiteLLM_VerificationToken { blocked Boolean? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? max_budget Float? budget_duration String? budget_reset_at DateTime? @@ -534,6 +538,7 @@ model LiteLLM_DeletedVerificationToken { blocked Boolean? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? max_budget Float? budget_duration String? budget_reset_at DateTime? diff --git a/litellm/a2a_protocol/main.py b/litellm/a2a_protocol/main.py index 0e8b8136c19..39600328074 100644 --- a/litellm/a2a_protocol/main.py +++ b/litellm/a2a_protocol/main.py @@ -21,6 +21,7 @@ from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.a2a_protocol.streaming_iterator import A2AStreamingIterator from litellm.a2a_protocol.utils import A2ARequestUtils from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, @@ -507,7 +508,7 @@ async def asend_message( prompt_tokens, completion_tokens, _, - ) = A2ARequestUtils.calculate_usage_from_request_response( + ) = await asyncify(A2ARequestUtils.calculate_usage_from_request_response)( request=request, response_dict=response_dict, ) diff --git a/litellm/a2a_protocol/streaming_iterator.py b/litellm/a2a_protocol/streaming_iterator.py index 67db8e905e3..d936caeb75e 100644 --- a/litellm/a2a_protocol/streaming_iterator.py +++ b/litellm/a2a_protocol/streaming_iterator.py @@ -11,6 +11,7 @@ import litellm from litellm._logging import verbose_logger from litellm.a2a_protocol.cost_calculator import A2ACostCalculator from litellm.a2a_protocol.utils import A2ARequestUtils +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj if TYPE_CHECKING: @@ -99,11 +100,11 @@ class A2AStreamingIterator: # Calculate tokens from collected text input_message: Final = A2ARequestUtils.get_input_message_from_request(self.request) input_text: Final = A2ARequestUtils.extract_text_from_message(input_message) - prompt_tokens: Final = A2ARequestUtils.count_tokens(input_text) + prompt_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(input_text) # Use the last (most complete) text from chunks output_text: Final = self.collected_text_parts[-1] if self.collected_text_parts else "" - completion_tokens: Final = A2ARequestUtils.count_tokens(output_text) + completion_tokens: Final = await asyncify(A2ARequestUtils.count_tokens)(output_text) total_tokens: Final = prompt_tokens + completion_tokens diff --git a/litellm/caching/qdrant_semantic_cache.py b/litellm/caching/qdrant_semantic_cache.py index c5876e993d3..058cc8a1579 100644 --- a/litellm/caching/qdrant_semantic_cache.py +++ b/litellm/caching/qdrant_semantic_cache.py @@ -21,6 +21,7 @@ from litellm.constants import ( QDRANT_VECTOR_SIZE, SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS, ) +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -255,7 +256,7 @@ class QdrantSemanticCache(BaseCache): llm_router = None router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) - embedding_input: Final = self._embedding_input(prompt, router) + embedding_input: Final = await asyncify(self._embedding_input)(prompt, router) embedding_call: Final = ( router.aembedding( model=self.embedding_model, diff --git a/litellm/caching/redis_semantic_cache.py b/litellm/caching/redis_semantic_cache.py index 9a70bfc1418..d4c815e15b7 100644 --- a/litellm/caching/redis_semantic_cache.py +++ b/litellm/caching/redis_semantic_cache.py @@ -19,6 +19,7 @@ from typing import TYPE_CHECKING, Any, Final, cast import litellm from litellm._logging import print_verbose, verbose_logger from litellm.constants import SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -522,7 +523,7 @@ class RedisSemanticCache(BaseCache): llm_router = None router: Final = resolve_embedding_router(self.embedding_model, llm_router, llm_model_list) - embedding_input: Final = self._embedding_input(prompt, router) + embedding_input: Final = await asyncify(self._embedding_input)(prompt, router) embedding_call: Final = ( router.aembedding( model=self.embedding_model, diff --git a/litellm/constants.py b/litellm/constants.py index 09442d6151e..ba5ec73d435 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1566,6 +1566,8 @@ BASE_MCP_ROUTE: Final = "/mcp" BATCH_STATUS_POLL_INTERVAL_SECONDS: Final = int(os.getenv("BATCH_STATUS_POLL_INTERVAL_SECONDS", 3600)) # 1 hour BATCH_STATUS_POLL_MAX_ATTEMPTS: Final = int(os.getenv("BATCH_STATUS_POLL_MAX_ATTEMPTS", 24)) # for 24 hours +BATCH_TPD_WINDOW_SECONDS: Final = 86400 +BATCH_TPD_DESCRIPTOR_SUFFIX: Final = "_tpd" HEALTH_CHECK_TIMEOUT_SECONDS: Final = int(os.getenv("HEALTH_CHECK_TIMEOUT_SECONDS", 60)) # 60 seconds _background_health_check_max_tokens_env: Final = os.getenv("BACKGROUND_HEALTH_CHECK_MAX_TOKENS") diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index 3dc6d81256b..852713595d5 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -814,6 +814,7 @@ def _select_model_name_for_cost_calc( if ( entry.get("input_cost_per_token") is not None or entry.get("input_cost_per_second") is not None + or entry.get("input_cost_per_query") is not None or entry.get("tiered_pricing") is not None ): return_model = router_model_id diff --git a/litellm/integrations/compression_interception/handler.py b/litellm/integrations/compression_interception/handler.py index 1be7a01ba3a..5352ce6b6a0 100644 --- a/litellm/integrations/compression_interception/handler.py +++ b/litellm/integrations/compression_interception/handler.py @@ -15,6 +15,7 @@ from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_logger from litellm.compression import compress from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.types.integrations.compression_interception import ( CompressionInterceptionConfig, CompressionSavingsMetadata, @@ -153,7 +154,7 @@ class CompressionInterceptionLogger(CustomLogger): self._prune_expired_cache() - compressed: Final = compress( + compressed: Final = await asyncify(compress)( messages=messages, model=model, call_type=CallTypes.anthropic_messages, diff --git a/litellm/integrations/langsmith.py b/litellm/integrations/langsmith.py index 9607eccef52..32664ed75d2 100644 --- a/litellm/integrations/langsmith.py +++ b/litellm/integrations/langsmith.py @@ -39,6 +39,8 @@ def is_serializable(value): class LangsmithLogger(CustomBatchLogger): + preserve_events_added_during_flush = True + def __init__( self, langsmith_api_key: str | None = None, diff --git a/litellm/integrations/prometheus.py b/litellm/integrations/prometheus.py index 69a38e83835..09be00f2b7b 100644 --- a/litellm/integrations/prometheus.py +++ b/litellm/integrations/prometheus.py @@ -2610,12 +2610,6 @@ class PrometheusLogger(CustomLogger): StandardLoggingPayloadSetup, ) - if self._should_skip_metrics_for_invalid_key( - user_api_key_dict=user_api_key_dict, - exception=original_exception, - ): - return - status_code: Final = self._extract_status_code(exception=original_exception) try: @@ -2633,7 +2627,7 @@ class PrometheusLogger(CustomLogger): end_user=user_api_key_dict.end_user_id, user=user_api_key_dict.user_id, user_email=user_api_key_dict.user_email, - hashed_api_key=user_api_key_dict.api_key, + hashed_api_key=None if status_code == 401 else user_api_key_dict.api_key, api_key_alias=user_api_key_dict.key_alias, team=user_api_key_dict.team_id, team_alias=user_api_key_dict.team_alias, diff --git a/litellm/litellm_core_utils/get_litellm_params.py b/litellm/litellm_core_utils/get_litellm_params.py index edd2e88f95c..49fc9abc525 100644 --- a/litellm/litellm_core_utils/get_litellm_params.py +++ b/litellm/litellm_core_utils/get_litellm_params.py @@ -41,6 +41,7 @@ OPTIONAL_KWARGS_KEYS: Final = ( "azure_password", "azure_scope", "timeout", + "client_side_timeout", "gcs_bucket_name", "bucket_name", "vertex_credentials", diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index 02681d8b499..3a1dbd24e86 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -238,6 +238,8 @@ def get_llm_provider( if dynamic_api_key is not None and not isinstance(dynamic_api_key, str): raise Exception(f"dynamic_api_key needs to be a string. Got type={type(dynamic_api_key).__name__}") return model, custom_llm_provider, dynamic_api_key, api_base + if "/" in model and is_registered_custom_provider(provider_prefix): + return model.split("/", 1)[1], provider_prefix, dynamic_api_key, api_base # check if api base is a known openai compatible endpoint if api_base: for endpoint in litellm.openai_compatible_endpoints: @@ -536,6 +538,10 @@ def get_llm_provider( ) +def is_registered_custom_provider(custom_llm_provider: str | None) -> bool: + return any(item["provider"] == custom_llm_provider for item in litellm.custom_provider_map) + + def _dashscope_family_chat_config(custom_llm_provider: str) -> "litellm.DashScopeChatConfig": if custom_llm_provider == "qwencloud": return litellm.QwenCloudChatConfig() diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4ddb9ce5b8e..abac624d5ec 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -1991,6 +1991,12 @@ class Logging(LiteLLMLoggingBaseClass): self.model_call_details["combined_usage_object"] = usage self.model_call_details["response_cost"] = response_cost + def record_assembled_response_for_failure(self, assembled: ModelResponse) -> None: + """Bill a fully streamed response on the failure log when a post-call hook rejects it.""" + usage: Final = getattr(assembled, "usage", None) + if isinstance(usage, Usage): + self.record_partial_usage_for_failure(usage, self._response_cost_calculator(result=assembled) or 0.0) + async def dispatch_failure_handlers( self, exception: Exception, diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 6a5a8832cc6..766d60ad180 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -19,6 +19,7 @@ from typing_extensions import NotRequired, TypedDict import litellm from litellm import verbose_logger from litellm._uuid import uuid +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.model_response_utils import ( is_model_response_stream_empty, ) @@ -2247,7 +2248,7 @@ class CustomStreamWrapper: if self.sent_last_chunk is True: # log the final chunk with accurate streaming values try: - complete_streaming_response = litellm.stream_chunk_builder( + complete_streaming_response = await asyncify(litellm.stream_chunk_builder)( chunks=self.chunks, messages=self.messages, logging_obj=self.logging_obj, diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py index 902808647c0..ad33e5e0592 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/dispatcher.py @@ -5,6 +5,7 @@ from collections.abc import Awaitable, Callable from typing import TYPE_CHECKING, Final, TypeAlias from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.types.llms.anthropic import AppliedEdit from .constants import CLEAR_TOOL_USES_EDIT_TYPE, COMPACT_EDIT_TYPE @@ -82,9 +83,9 @@ async def apply_context_management( """Run edits in order; return a single ``PolyfillResult``. The dispatcher is async so async editors (``compact_20260112``) can - ``await`` the configured summarization model. Sync editors are called - inline — ``inspect.iscoroutinefunction`` decides how each editor is - invoked. + ``await`` the configured summarization model. Sync editors run in a + worker thread so their token counts stay off the event loop; + ``inspect.iscoroutinefunction`` decides how each editor is invoked. """ edits: Final = _normalize_spec(context_management_spec) if not edits: @@ -121,7 +122,7 @@ async def apply_context_management( user_api_key_auth=user_api_key_auth, ) if editor_is_async - else editor( + else await asyncify(editor)( model=model, messages=current_messages, tools=tools, diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index 050ab67c86c..fb6a1c40253 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -20,6 +20,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict, Unpack import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.types.llms.anthropic import ( AppliedEdit, CompactionBlock, @@ -1157,7 +1158,7 @@ async def apply_compact_20260112( # Phase B: threshold check. try: - current_tokens = _count_effective_tokens( + current_tokens = await asyncify(_count_effective_tokens)( model=model, effective_messages=effective_messages, # ``augmented_system`` already carries the prior compaction summary diff --git a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py index 7d01aee5d98..98c5c6d6d4e 100644 --- a/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py +++ b/litellm/llms/anthropic/experimental_pass_through/messages/streaming_iterator.py @@ -678,7 +678,7 @@ class BaseAnthropicMessagesStreamingIterator: """ from litellm.proxy.pass_through_endpoints.streaming_handler import PassThroughStreamingHandler - PassThroughStreamingHandler.schedule_stream_failure_logging( + await PassThroughStreamingHandler.schedule_stream_failure_logging( litellm_logging_obj=self.litellm_logging_obj, endpoint_type=EndpointType.ANTHROPIC, request_body=self.request_body, diff --git a/litellm/llms/fireworks_ai/rerank/transformation.py b/litellm/llms/fireworks_ai/rerank/transformation.py index e142622aa1b..509dbd5ff24 100644 --- a/litellm/llms/fireworks_ai/rerank/transformation.py +++ b/litellm/llms/fireworks_ai/rerank/transformation.py @@ -250,8 +250,7 @@ class FireworksAIRerankConfig(FireworksAIMixin, BaseRerankConfig): rerank_results.append(rerank_result) - # Use model name as id if no id is provided - response_id: Final = raw_response_json.get("id") or raw_response_json.get("model") or str(uuid.uuid4()) + response_id: Final = raw_response_json.get("id") or str(uuid.uuid4()) return RerankResponse( id=response_id, diff --git a/litellm/llms/vertex_ai/rerank/transformation.py b/litellm/llms/vertex_ai/rerank/transformation.py index 2ec4f2da79b..b0c6add69fd 100644 --- a/litellm/llms/vertex_ai/rerank/transformation.py +++ b/litellm/llms/vertex_ai/rerank/transformation.py @@ -4,6 +4,8 @@ Translates from Cohere's `/v1/rerank` input format to Vertex AI Discovery Engine Why separate file? Make it easy to see how transformation works """ +import math +import uuid from collections.abc import Mapping from typing import Any, Final @@ -32,6 +34,8 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): Reference: https://cloud.google.com/generative-ai-app-builder/docs/ranking#rank_or_rerank_a_set_of_records_according_to_a_query """ + MAX_RECORDS_PER_SEARCH_UNIT = 100 + def __init__(self) -> None: super().__init__() @@ -208,10 +212,11 @@ class VertexAIRerankConfig(BaseRerankConfig, VertexBase): RerankResponseResult(index=result["index"], relevance_score=result["relevance_score"]) ) - # Create meta object - meta: Final = RerankResponseMeta(billed_units=RerankBilledUnits(search_units=len(records))) + input_record_count: Final = len(request_data.get("records", ())) + search_units: Final = math.ceil(input_record_count / self.MAX_RECORDS_PER_SEARCH_UNIT) + meta: Final = RerankResponseMeta(billed_units=RerankBilledUnits(search_units=search_units)) - return RerankResponse(id=f"vertex_ai_rerank_{model}", results=rerank_results, meta=meta) + return RerankResponse(id=f"vertex_ai_rerank_{uuid.uuid4()}", results=rerank_results, meta=meta) def get_supported_cohere_rerank_params(self, model: str) -> list: return [ diff --git a/litellm/llms/voyage/rerank/transformation.py b/litellm/llms/voyage/rerank/transformation.py index fea8452d934..0f57ac11028 100644 --- a/litellm/llms/voyage/rerank/transformation.py +++ b/litellm/llms/voyage/rerank/transformation.py @@ -9,6 +9,7 @@ from typing import Any, Final import httpx +from litellm._uuid import uuid from litellm.llms.base_llm.chat.transformation import LiteLLMLoggingObj from litellm.llms.base_llm.rerank.transformation import BaseRerankConfig from litellm.secret_managers.main import get_secret_str @@ -127,7 +128,7 @@ class VoyageRerankConfig(BaseRerankConfig): rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) return RerankResponse( - id=_json_response.get("id", f"voyage-rerank-{model}"), + id=_json_response.get("id") or str(uuid.uuid4()), results=transformed_results, meta=rerank_meta, ) diff --git a/litellm/llms/watsonx/rerank/transformation.py b/litellm/llms/watsonx/rerank/transformation.py index 293880b188d..bd6b23ff2be 100644 --- a/litellm/llms/watsonx/rerank/transformation.py +++ b/litellm/llms/watsonx/rerank/transformation.py @@ -191,7 +191,7 @@ class IBMWatsonXRerankConfig(IBMWatsonXMixin, BaseRerankConfig): transformed_results.append(transformed_result) - response_id: Final = raw_response_json.get("id") or raw_response_json.get("model_id") or str(uuid.uuid4()) + response_id: Final = raw_response_json.get("id") or str(uuid.uuid4()) # Extract usage information _tokens: Final = RerankTokens( diff --git a/litellm/llms/xai/chat/transformation.py b/litellm/llms/xai/chat/transformation.py index 747ee0b6c49..91bf697487d 100644 --- a/litellm/llms/xai/chat/transformation.py +++ b/litellm/llms/xai/chat/transformation.py @@ -219,13 +219,19 @@ class XAIChatConfig(OpenAIGPTConfig): litellm_params: dict, headers: dict, ) -> dict: - """ - Handle https://github.com/BerriAI/litellm/issues/9720 + """Handle https://github.com/BerriAI/litellm/issues/9720""" + if "web_search_options" in optional_params: + verbose_logger.warning( + "XAI no longer supports web search on /chat/completions (Live Search is deprecated). " + "Dropping 'web_search_options'. Use the Responses API for XAI web search." + ) - Filter out 'name' from messages - """ - messages = strip_name_from_messages(messages) - return super().transform_request(model, messages, optional_params, litellm_params, headers) + chat_params: Final = { # mutable-ok: base transform_request takes a plain dict of optional params + key: value for key, value in optional_params.items() if key != "web_search_options" + } + return super().transform_request( + model, strip_name_from_messages(messages), chat_params, litellm_params, headers + ) @staticmethod def _fix_choice_finish_reason_for_tool_calls(choice: Choices) -> None: diff --git a/litellm/llms/xai/responses/transformation.py b/litellm/llms/xai/responses/transformation.py index 646d6798783..1f977a66186 100644 --- a/litellm/llms/xai/responses/transformation.py +++ b/litellm/llms/xai/responses/transformation.py @@ -3,6 +3,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter import litellm from litellm._logging import verbose_logger @@ -32,6 +33,8 @@ if TYPE_CHECKING: else: LiteLLMLoggingObj = Any +_STR_MAPPING_ADAPTER: Final = TypeAdapter(Mapping[str, object]) + def _usage_restated_from_xai_ticks(usage: ResponseAPIUsage | None) -> ResponseAPIUsage | None: reported_cost: Final = xai_reported_cost_in_usd(getattr(usage, "cost_in_usd_ticks", None)) @@ -81,30 +84,25 @@ class XAIResponsesAPIConfig(OpenAIResponsesAPIConfig): - enable_image_understanding XAI does NOT support search_context_size (OpenAI-specific). + + Domains may come nested under 'filters' (the OpenAI/XAI documented shape) or flat on the tool. """ xai_tool: Final[dict[str, object]] = {"type": "web_search"} - # Remove search_context_size if present (not supported by XAI) if "search_context_size" in tool: verbose_logger.info( "XAI does not support 'search_context_size' parameter. Removing it from web_search tool." ) - # Handle filters (XAI-specific structure) - filters: Final = {} - if "allowed_domains" in tool: - allowed_domains: Final = tool["allowed_domains"] - filters["allowed_domains"] = allowed_domains + nested_filters: Final = tool.get("filters") + domains: Final = ( + _STR_MAPPING_ADAPTER.validate_python(nested_filters) if isinstance(nested_filters, Mapping) else tool + ) + filters: Final = {key: domains[key] for key in ("allowed_domains", "excluded_domains") if key in domains} - if "excluded_domains" in tool: - excluded_domains: Final = tool["excluded_domains"] - filters["excluded_domains"] = excluded_domains - - # Add filters if any were specified if filters: xai_tool["filters"] = filters - # Handle enable_image_understanding (top-level in XAI format) if "enable_image_understanding" in tool: xai_tool["enable_image_understanding"] = tool["enable_image_understanding"] diff --git a/litellm/main.py b/litellm/main.py index f6f4ec1bf63..9fbc5881b4f 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -67,7 +67,7 @@ from litellm.constants import ( ) from litellm.exceptions import LiteLLMUnknownProvider from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.asyncify import asyncify, run_async_function from litellm.litellm_core_utils.audio_utils.utils import ( calculate_request_duration, get_audio_file_for_health_check, @@ -1072,10 +1072,6 @@ def responses_api_bridge_check( mode = "responses" model_info["mode"] = mode - if web_search_options is not None and custom_llm_provider == "xai": - model_info["mode"] = "responses" - model = model.replace("responses/", "") - except Exception as e: verbose_logger.debug("Error getting model info: %s", e) @@ -1084,6 +1080,10 @@ def responses_api_bridge_check( mode = "responses" model_info["mode"] = mode + if web_search_options is not None and custom_llm_provider == "xai": + model_info["mode"] = "responses" + model = model.replace("responses/", "") + # OpenAI/Azure GPT-5 chat-completions that need Responses-only fields (e.g. # ``reasoningSummary`` in ``extra_body``) must be bridged; Chat Completions rejects # those keys. @@ -9127,7 +9127,7 @@ async def acount_tokens( fallback_messages = messages or [] if system and fallback_messages: fallback_messages = [{"role": "system", "content": system}] + fallback_messages - local_count: Final = litellm.token_counter( + local_count: Final = await asyncify(litellm.token_counter)( model=model, messages=fallback_messages, tools=tools, diff --git a/litellm/models/budget.py b/litellm/models/budget.py index 335800a49a8..125ce739d6a 100644 --- a/litellm/models/budget.py +++ b/litellm/models/budget.py @@ -26,6 +26,7 @@ class LiteLLM_BudgetTable(LiteLLMPydanticObjectBase): max_parallel_requests: int | None = None tpm_limit: int | None = None rpm_limit: int | None = None + tpd_limit: int | None = None model_max_budget: dict | None = None budget_duration: str | None = None allowed_models: list[str] | None = None # per-member model scope; empty = inherit team models diff --git a/litellm/models/team.py b/litellm/models/team.py index da526515e6e..8edf10703b1 100644 --- a/litellm/models/team.py +++ b/litellm/models/team.py @@ -71,6 +71,7 @@ class TeamBase(LiteLLMPydanticObjectBase): metadata: dict | None = None tpm_limit: int | None = None rpm_limit: int | None = None + tpd_limit: int | None = None max_budget: float | None = None soft_budget: float | None = None budget_duration: str | None = None diff --git a/litellm/models/verification_token.py b/litellm/models/verification_token.py index fec3caec457..06ff877a41a 100644 --- a/litellm/models/verification_token.py +++ b/litellm/models/verification_token.py @@ -31,6 +31,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): metadata: dict = {} tpm_limit: int | None = None rpm_limit: int | None = None + tpd_limit: int | None = None budget_duration: str | None = None budget_reset_at: datetime | None = None allowed_cache_controls: list | None = [] diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 993f705a918..d4eda1c9540 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1206,6 +1206,7 @@ class AllowedVectorStoreIndexItem(LiteLLMPydanticObjectBase): class KeyRequestBase(GenerateRequestBase): key: str | None = None + tpd_limit: int | None = None default_estimated_output_tokens: PositiveInt | None = None default_estimated_output_tokens_per_model: Mapping[str, PositiveInt] | None = None budget_id: str | None = None @@ -1891,6 +1892,9 @@ class BudgetNewRequest(LiteLLMPydanticObjectBase): ) tpm_limit: int | None = Field(default=None, description="Max tokens per minute, allowed for this budget id.") rpm_limit: int | None = Field(default=None, description="Max requests per minute, allowed for this budget id.") + tpd_limit: int | None = Field( + default=None, description="Max tokens per day, charged by batch submissions, allowed for this budget id." + ) budget_duration: str | None = Field( default=None, description="Max duration budget should be set for (e.g. '1hr', '1d', '28d')", @@ -2067,6 +2071,7 @@ class UpdateTeamRequest(LiteLLMPydanticObjectBase): metadata: dict | None = None tpm_limit: int | None = None rpm_limit: int | None = None + tpd_limit: int | None = None max_budget: float | None = None soft_budget: float | None = None models: list | None = None @@ -3022,6 +3027,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): team_alias: str | None = None team_tpm_limit: int | None = None team_rpm_limit: int | None = None + team_tpd_limit: int | None = None team_max_budget: float | None = None team_soft_budget: float | None = None team_models: list = [] @@ -3041,6 +3047,7 @@ class LiteLLM_VerificationTokenView(LiteLLM_VerificationToken): end_user_id: str | None = None end_user_tpm_limit: int | None = None end_user_rpm_limit: int | None = None + end_user_tpd_limit: int | None = None end_user_max_budget: float | None = None end_user_model_max_budget: dict | None = None @@ -3839,6 +3846,7 @@ class SpendLogsMetadata(TypedDict): user_api_key_team_alias: str | None spend_logs_metadata: dict | None # special param to log k,v pairs to spendlogs for a call requester_ip_address: str | None + user_agent: ReadOnly[str | None] litellm_call_id: str | None applied_guardrails: list[str] | None mcp_tool_call_metadata: StandardLoggingMCPToolCall | None diff --git a/litellm/proxy/auth/auth_exception_handler.py b/litellm/proxy/auth/auth_exception_handler.py index b36c8a038fc..661b6a83c38 100644 --- a/litellm/proxy/auth/auth_exception_handler.py +++ b/litellm/proxy/auth/auth_exception_handler.py @@ -23,6 +23,7 @@ from litellm.proxy.auth.auth_utils import ( _get_request_ip_address, is_invalid_virtual_key_error, mark_invalid_virtual_key_error, + normalize_request_route, ) from litellm.proxy.db.exception_handler import PrismaDBExceptionHandler from litellm.types.services import ServiceTypes @@ -74,17 +75,28 @@ def _as_proxy_exception(e: Exception) -> ProxyException: ) -def _with_requester_ip_address(request_data: dict[str, object], requester_ip: str | None) -> dict[str, object]: +def _get_user_agent(request: Request) -> str | None: + if "headers" not in request.scope: + return None + return request.headers.get("user-agent") + + +def _with_client_context( + request_data: dict[str, object], requester_ip: str | None, user_agent: str | None +) -> dict[str, object]: """Auth gate rejections are raised before `add_litellm_data_to_request` records the - caller IP, so their failure logs would otherwise carry no IP nor key/user identity.""" - if not requester_ip: - return request_data + caller IP and User-Agent, so their failure logs would otherwise carry neither.""" key: Final = "litellm_metadata" if "litellm_metadata" in request_data else "metadata" metadata: Final = request_data.get(key) base: Final[Mapping[str, object]] = metadata if isinstance(metadata, Mapping) else EMPTY_MAPPING - if base.get("requester_ip_address"): + stamped: Final = { + name: value + for name, value in (("requester_ip_address", requester_ip), ("user_agent", user_agent)) + if value and not base.get(name) + } + if not stamped: return request_data - return {**request_data, key: {**base, "requester_ip_address": requester_ip}} # mutable-ok: logging needs dicts + return {**request_data, key: {**base, **stamped}} # mutable-ok: logging needs dicts class UserAPIKeyAuthExceptionHandler: @@ -148,6 +160,7 @@ class UserAPIKeyAuthExceptionHandler: request=request, use_x_forwarded_for=general_settings.get("use_x_forwarded_for") is True, ) + user_agent: Final = _get_user_agent(request) # Log authentication failures before identity seeding and callbacks, so the log # survives a raising callback pipeline. Classify and route malformed virtual-key @@ -172,7 +185,7 @@ class UserAPIKeyAuthExceptionHandler: # so the handler is side-effect-free for the caller's identity object. user_api_key_dict = resolved_identity.model_copy() if resolved_identity is not None else UserAPIKeyAuth() user_api_key_dict.parent_otel_span = parent_otel_span - user_api_key_dict.request_route = route + user_api_key_dict.request_route = normalize_request_route(route) user_api_key_dict.api_key = user_api_key_dict.api_key or UserAPIKeyAuth(api_key=api_key).api_key # Stamp identity onto the request's server span now, before the request @@ -200,7 +213,7 @@ class UserAPIKeyAuthExceptionHandler: # Allow callbacks to transform the error response transformed_exception: Final = await proxy_logging_obj.post_call_failure_hook( - request_data=_with_requester_ip_address(request_data, requester_ip), + request_data=_with_client_context(request_data, requester_ip, user_agent), original_exception=e, user_api_key_dict=user_api_key_dict, error_type=ProxyErrorTypes.auth_error, diff --git a/litellm/proxy/auth/team_grants.py b/litellm/proxy/auth/team_grants.py index 1196011dcdd..0421659c331 100644 --- a/litellm/proxy/auth/team_grants.py +++ b/litellm/proxy/auth/team_grants.py @@ -56,6 +56,7 @@ class TeamGrants(TypedDict, total=False): team_alias: ReadOnly[str | None] team_tpm_limit: ReadOnly[int | None] team_rpm_limit: ReadOnly[int | None] + team_tpd_limit: ReadOnly[int | None] team_max_budget: ReadOnly[float | None] team_soft_budget: ReadOnly[float | None] team_spend: ReadOnly[float | None] @@ -97,6 +98,7 @@ def team_grants( team_alias=team_object.team_alias, team_tpm_limit=team_object.tpm_limit, team_rpm_limit=team_object.rpm_limit, + team_tpd_limit=team_object.tpd_limit, team_max_budget=team_object.max_budget, team_soft_budget=team_object.soft_budget, team_spend=team_object.spend, diff --git a/litellm/proxy/auth/user_api_key_auth.py b/litellm/proxy/auth/user_api_key_auth.py index 0618875be55..7d62baf39a8 100644 --- a/litellm/proxy/auth/user_api_key_auth.py +++ b/litellm/proxy/auth/user_api_key_auth.py @@ -537,6 +537,9 @@ def _apply_budget_limits_to_end_user_params( if budget_info.rpm_limit is not None: end_user_params["end_user_rpm_limit"] = budget_info.rpm_limit + if budget_info.tpd_limit is not None: + end_user_params["end_user_tpd_limit"] = budget_info.tpd_limit + if budget_info.max_budget is not None: end_user_params["end_user_max_budget"] = budget_info.max_budget @@ -621,6 +624,8 @@ def update_valid_token_with_end_user_params(valid_token: UserAPIKeyAuth, end_use valid_token.end_user_tpm_limit = end_user_params["end_user_tpm_limit"] if end_user_params.get("end_user_rpm_limit") is not None: valid_token.end_user_rpm_limit = end_user_params["end_user_rpm_limit"] + if end_user_params.get("end_user_tpd_limit") is not None: + valid_token.end_user_tpd_limit = end_user_params["end_user_tpd_limit"] if end_user_params.get("allowed_model_region") is not None: valid_token.allowed_model_region = end_user_params["allowed_model_region"] if end_user_params.get("end_user_model_max_budget") is not None: @@ -2026,6 +2031,7 @@ async def _user_api_key_auth_builder( valid_token.end_user_id = end_user_params.get("end_user_id") valid_token.end_user_tpm_limit = end_user_params.get("end_user_tpm_limit") valid_token.end_user_rpm_limit = end_user_params.get("end_user_rpm_limit") + valid_token.end_user_tpd_limit = end_user_params.get("end_user_tpd_limit") valid_token.allowed_model_region = end_user_params.get("allowed_model_region") if valid_token is not None: @@ -2302,6 +2308,7 @@ async def _user_api_key_auth_builder( spend=valid_token.team_spend, tpm_limit=valid_token.team_tpm_limit, rpm_limit=valid_token.team_rpm_limit, + tpd_limit=valid_token.team_tpd_limit, blocked=valid_token.team_blocked, models=token_team_models, metadata=valid_token.team_metadata, @@ -2455,6 +2462,7 @@ def _team_obj_from_token(valid_token: UserAPIKeyAuth) -> LiteLLM_TeamTableCached spend=valid_token.team_spend, tpm_limit=valid_token.team_tpm_limit, rpm_limit=valid_token.team_rpm_limit, + tpd_limit=valid_token.team_tpd_limit, blocked=valid_token.team_blocked, models=token_team_models, metadata=valid_token.team_metadata, diff --git a/litellm/proxy/common_utils/http_parsing_utils.py b/litellm/proxy/common_utils/http_parsing_utils.py index 845589aee7a..9c2767c7771 100644 --- a/litellm/proxy/common_utils/http_parsing_utils.py +++ b/litellm/proxy/common_utils/http_parsing_utils.py @@ -189,8 +189,9 @@ async def _read_request_body(request: Request | None) -> dict: try: parsed_body = json.loads(body_str) - except json.JSONDecodeError: - # If both orjson and json.loads fail, throw a proper error + json.dumps(parsed_body, ensure_ascii=False).encode("utf-8") + except (json.JSONDecodeError, UnicodeEncodeError): + # json.loads accepts lone surrogate escapes that no provider can encode verbose_proxy_logger.error("Invalid JSON payload received: %s", e) raise ProxyException( message=f"Invalid JSON payload: {e}", diff --git a/litellm/proxy/common_utils/reset_budget_job.py b/litellm/proxy/common_utils/reset_budget_job.py index 35e74418628..acb51e73daf 100644 --- a/litellm/proxy/common_utils/reset_budget_job.py +++ b/litellm/proxy/common_utils/reset_budget_job.py @@ -7,7 +7,7 @@ from dataclasses import dataclass, field from datetime import datetime, timedelta, timezone from enum import Enum from types import MappingProxyType -from typing import Final, Literal, Protocol, TypeVar +from typing import Final, Generic, Literal, Protocol, TypeVar from typing_extensions import assert_never @@ -68,6 +68,13 @@ from litellm.types.services import ServiceTypes _RowT = TypeVar("_RowT") + +@dataclass(frozen=True, slots=True) +class _RowReset(Generic[_RowT]): + row: _RowT + spend_decrement: float + + _LINKED_KEYS_WHERE: Final[Mapping[str, object]] = MappingProxyType({"budget_duration": None, "spend": {"gt": 0}}) _SPENT_ROWS_WHERE: Final[Mapping[str, object]] = MappingProxyType({"spend": {"gt": 0}}) @@ -530,10 +537,9 @@ class ResetBudgetJob: ) @staticmethod - async def _invalidate_spend_counter(counter_key: str, new_spend: float = 0.0) -> None: - """Overwrite a spend counter with the post-reset value (0, or the carried - overage when budget rollover is enabled) so a DB-row reset takes effect - immediately. + async def _invalidate_spend_counter(counter_key: str) -> None: + """Drop a spend counter so the next read reseeds from the committed DB + row, the only value that includes increments that raced the reset. Call AFTER the DB write commits. Clearing Redis before the DB commit opens a window where get_current_spend reads 0 from Redis @@ -542,10 +548,10 @@ class ResetBudgetJob: try: from litellm.proxy.proxy_server import spend_counter_cache - spend_counter_cache.in_memory_cache.set_cache(key=counter_key, value=new_spend, ttl=60) + spend_counter_cache.in_memory_cache.delete_cache(key=counter_key) if spend_counter_cache.redis_cache is not None: try: - await spend_counter_cache.redis_cache.async_set_cache(key=counter_key, value=new_spend, ttl=60) + await spend_counter_cache.redis_cache.async_delete_cache(key=counter_key) except Exception as redis_err: verbose_proxy_logger.warning( "Failed to reset spend counter %s in Redis: %s. " @@ -730,8 +736,8 @@ class ResetBudgetJob: uow.budgets.queue_window_advance(budget_id=budget_id, budget_reset_at=budget_reset_at) async def _invalidate_budget_cascade_caches(self, cascade: _BudgetCascade) -> None: - for counter_key, new_spend in cascade.counter_resets: - await self._invalidate_spend_counter(counter_key, new_spend=new_spend) + for counter_key, _ in cascade.counter_resets: + await self._invalidate_spend_counter(counter_key) for cache_key in cascade.cache_keys: await self._invalidate_user_api_key_cache_entry(cache_key) @@ -842,7 +848,7 @@ class ResetBudgetJob: ) return [LiteLLM_EndUserTable.model_validate(row.model_dump()) for row in rows] - async def _write_key_reset_updates(self, updated_keys: list[LiteLLM_VerificationToken]) -> None: + async def _write_key_reset_updates(self, updated_keys: Sequence[_RowReset[LiteLLM_VerificationToken]]) -> None: """ Write per-row {spend, budget_reset_at} updates for keys. @@ -858,18 +864,18 @@ class ResetBudgetJob: reason="reset_budget_write_keys_failure", ) - async def _write_key_reset_updates_once(self, updated_keys: list[LiteLLM_VerificationToken]) -> None: + async def _write_key_reset_updates_once(self, updated_keys: Sequence[_RowReset[LiteLLM_VerificationToken]]) -> None: async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow: for k in updated_keys: - if k.token is None: + if k.row.token is None: continue uow.keys.queue_spend_reset( - token=k.token, - budget_reset_at=k.budget_reset_at, - spend_decrement=k.max_budget if (k.spend or 0.0) > 0.0 else None, + token=k.row.token, + budget_reset_at=k.row.budget_reset_at, + spend_decrement=k.spend_decrement, ) - async def _write_user_reset_updates(self, updated_users: list[LiteLLM_UserTable]) -> None: + async def _write_user_reset_updates(self, updated_users: Sequence[_RowReset[LiteLLM_UserTable]]) -> None: """ Write per-row {spend, budget_reset_at} updates for users. @@ -882,16 +888,16 @@ class ResetBudgetJob: reason="reset_budget_write_users_failure", ) - async def _write_user_reset_updates_once(self, updated_users: list[LiteLLM_UserTable]) -> None: + async def _write_user_reset_updates_once(self, updated_users: Sequence[_RowReset[LiteLLM_UserTable]]) -> None: async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow: for u in updated_users: uow.users.queue_spend_reset( - user_id=u.user_id, - budget_reset_at=u.budget_reset_at, - spend_decrement=u.max_budget if (u.spend or 0.0) > 0.0 else None, + user_id=u.row.user_id, + budget_reset_at=u.row.budget_reset_at, + spend_decrement=u.spend_decrement, ) - async def _write_team_reset_updates(self, updated_teams: list[LiteLLM_TeamTable]) -> None: + async def _write_team_reset_updates(self, updated_teams: Sequence[_RowReset[LiteLLM_TeamTable]]) -> None: """ Write per-row {spend, budget_reset_at} updates for teams. @@ -904,13 +910,13 @@ class ResetBudgetJob: reason="reset_budget_write_teams_failure", ) - async def _write_team_reset_updates_once(self, updated_teams: list[LiteLLM_TeamTable]) -> None: + async def _write_team_reset_updates_once(self, updated_teams: Sequence[_RowReset[LiteLLM_TeamTable]]) -> None: async with spend_reset_unit_of_work(self.prisma_client.db.batch_) as uow: for t in updated_teams: uow.teams.queue_spend_reset( - team_id=t.team_id, - budget_reset_at=t.budget_reset_at, - spend_decrement=t.max_budget if (t.spend or 0.0) > 0.0 else None, + team_id=t.row.team_id, + budget_reset_at=t.row.budget_reset_at, + spend_decrement=t.spend_decrement, ) def _emit_phase_failure( @@ -962,18 +968,24 @@ class ResetBudgetJob: reason="reset_budget_read_keys_failure", ) verbose_proxy_logger.debug("Keys to reset %s", _LazyJson(keys_to_reset)) - updated_keys: Final[list[LiteLLM_VerificationToken]] = [] + updated_keys: Final[list[_RowReset[LiteLLM_VerificationToken]]] = [] failed_keys: Final = [] if keys_to_reset is not None and len(keys_to_reset) > 0: for key in keys_to_reset: try: + pre_reset_spend = float(key.spend or 0.0) updated_key = await ResetBudgetJob._reset_budget_for_key( key=key, current_time=now, reset_settings=self.reset_settings, ) if updated_key is not None: - updated_keys.append(updated_key) + updated_keys.append( + _RowReset( + row=updated_key, + spend_decrement=pre_reset_spend - float(updated_key.spend or 0.0), + ) + ) else: failed_keys.append({"key": key, "error": "Returned None without exception"}) except Exception as e: @@ -985,15 +997,15 @@ class ResetBudgetJob: if updated_keys: await self._write_key_reset_updates(updated_keys=updated_keys) for k in updated_keys: - token = getattr(k, "token", None) + token = getattr(k.row, "token", None) if token: - await self._invalidate_spend_counter(f"spend:key:{token}", new_spend=k.spend or 0.0) + await self._invalidate_spend_counter(f"spend:key:{token}") end_time = time.time() outcome: Final = _ChunkOutcome( fetched=len(keys_to_reset) if keys_to_reset else 0, advanced=_count_advanced( - (k.budget_reset_at for k in updated_keys), + (k.row.budget_reset_at for k in updated_keys), cutoff=datetime.now(timezone.utc), ), ) @@ -1063,18 +1075,24 @@ class ResetBudgetJob: ), reason="reset_budget_read_users_failure", ) - updated_users: Final[list[LiteLLM_UserTable]] = [] + updated_users: Final[list[_RowReset[LiteLLM_UserTable]]] = [] failed_users: Final = [] if users_to_reset is not None and len(users_to_reset) > 0: for user in users_to_reset: try: + pre_reset_spend = float(user.spend or 0.0) updated_user = await ResetBudgetJob._reset_budget_for_user( user=user, current_time=now, reset_settings=self.reset_settings, ) if updated_user is not None: - updated_users.append(updated_user) + updated_users.append( + _RowReset( + row=updated_user, + spend_decrement=pre_reset_spend - float(updated_user.spend or 0.0), + ) + ) else: failed_users.append( { @@ -1090,9 +1108,9 @@ class ResetBudgetJob: if updated_users: await self._write_user_reset_updates(updated_users=updated_users) for u in updated_users: - user_id = getattr(u, "user_id", None) + user_id = getattr(u.row, "user_id", None) if user_id: - await self._invalidate_spend_counter(f"spend:user:{user_id}", new_spend=u.spend or 0.0) + await self._invalidate_spend_counter(f"spend:user:{user_id}") if user_id == LITELLM_PROXY_BUDGET_NAME: await self._invalidate_global_proxy_spend_cache() @@ -1100,7 +1118,7 @@ class ResetBudgetJob: outcome: Final = _ChunkOutcome( fetched=len(users_to_reset) if users_to_reset else 0, advanced=_count_advanced( - (u.budget_reset_at for u in updated_users), + (u.row.budget_reset_at for u in updated_users), cutoff=datetime.now(timezone.utc), ), ) @@ -1172,18 +1190,24 @@ class ResetBudgetJob: ), reason="reset_budget_read_teams_failure", ) - updated_teams: Final[list[LiteLLM_TeamTable]] = [] + updated_teams: Final[list[_RowReset[LiteLLM_TeamTable]]] = [] failed_teams: Final = [] if teams_to_reset is not None and len(teams_to_reset) > 0: for team in teams_to_reset: try: + pre_reset_spend = float(team.spend or 0.0) updated_team = await ResetBudgetJob._reset_budget_for_team( team=team, current_time=now, reset_settings=self.reset_settings, ) if updated_team is not None: - updated_teams.append(updated_team) + updated_teams.append( + _RowReset( + row=updated_team, + spend_decrement=pre_reset_spend - float(updated_team.spend or 0.0), + ) + ) else: failed_teams.append( { @@ -1199,15 +1223,15 @@ class ResetBudgetJob: if updated_teams: await self._write_team_reset_updates(updated_teams=updated_teams) for t in updated_teams: - team_id = getattr(t, "team_id", None) + team_id = getattr(t.row, "team_id", None) if team_id: - await self._invalidate_spend_counter(f"spend:team:{team_id}", new_spend=t.spend or 0.0) + await self._invalidate_spend_counter(f"spend:team:{team_id}") end_time = time.time() outcome: Final = _ChunkOutcome( fetched=len(teams_to_reset) if teams_to_reset else 0, advanced=_count_advanced( - (t.budget_reset_at for t in updated_teams), + (t.row.budget_reset_at for t in updated_teams), cutoff=datetime.now(timezone.utc), ), ) diff --git a/litellm/proxy/db/create_views.py b/litellm/proxy/db/create_views.py index 10daeee4e7b..d3f3de730ab 100644 --- a/litellm/proxy/db/create_views.py +++ b/litellm/proxy/db/create_views.py @@ -80,6 +80,7 @@ async def create_missing_views(db: SupportsRawQueries) -> None: t.max_budget AS team_max_budget, t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, + t.tpd_limit AS team_tpd_limit, p.project_alias AS project_alias FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id diff --git a/litellm/proxy/db/routing_prisma_wrapper.py b/litellm/proxy/db/routing_prisma_wrapper.py index be515392a17..0eb378b2fe9 100644 --- a/litellm/proxy/db/routing_prisma_wrapper.py +++ b/litellm/proxy/db/routing_prisma_wrapper.py @@ -81,6 +81,11 @@ class WriterPinnedClient: self.db: Final = db.writer if isinstance(db, RoutingPrismaWrapper) and not db.writer_unavailable else db +def writer_wrapper(db: "PrismaWrapper | RoutingPrismaWrapper") -> PrismaWrapper: + """Unlike `WriterPinnedClient`, ignores `writer_unavailable`: a raw SQL write has no replica fallback.""" + return db.writer if isinstance(db, RoutingPrismaWrapper) else db + + class RoutingPrismaWrapper: """ Routes Prisma operations between a writer and a reader Prisma client. diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index dcd34a1d9cb..ab6e10ca76b 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -18,12 +18,13 @@ Quick summary: """ import json -from collections.abc import Iterable, Mapping, Sequence +from collections.abc import Callable, Iterable, Mapping, Sequence +from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, NoReturn, TypeAlias from fastapi import HTTPException -from pydantic import BaseModel, Field, TypeAdapter +from pydantic import BaseModel, Field, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger @@ -33,6 +34,7 @@ from litellm.batches.batch_utils import ( _extract_file_access_credentials, _iter_batch_input_lines, ) +from litellm.constants import BATCH_TPD_DESCRIPTOR_SUFFIX, BATCH_TPD_WINDOW_SECONDS from litellm.exceptions import RateLimitErrorCategory from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import ( @@ -55,6 +57,7 @@ from litellm.proxy.hooks.batch_enqueued_tokens import ( from litellm.proxy.hooks.parallel_request_limiter_v3 import ( PROJECT_ITPM_DESCRIPTOR_KEY, PROJECT_OTPM_DESCRIPTOR_KEY, + ReservationAwareIncrementOperation, get_or_create_request_stash, ) from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit @@ -92,6 +95,7 @@ else: _BATCH_BODY_ADAPTER: Final = TypeAdapter(dict[str, object]) +_WINDOW_START_ADAPTER: Final[TypeAdapter[int | float | str | None]] = TypeAdapter(int | float | str | None) IncrementAmounts: TypeAlias = dict[Literal["requests", "tokens"], int] @@ -128,6 +132,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): self, internal_usage_cache: InternalUsageCache, parallel_request_limiter: ParallelRequestLimiter, + time_provider: Callable[[], datetime] | None = None, ): """ Initialize the batch rate limiter. @@ -138,9 +143,11 @@ class _PROXY_BatchRateLimiter(CustomLogger): Args: internal_usage_cache: Cache for storing rate limit data (auto-injected) parallel_request_limiter: Existing rate limiter to integrate with (needs custom injection) + time_provider: Clock used for rate limit reset times (defaults to ``datetime.now``) """ self.internal_usage_cache = internal_usage_cache self.parallel_request_limiter = parallel_request_limiter + self._time_provider: Final = time_provider or datetime.now self._warned_unsupported_model_skip = False def _get_file_bound_batch_model(self, data: dict) -> str | None: @@ -236,14 +243,48 @@ class _PROXY_BatchRateLimiter(CustomLogger): file-bound/top-level routing model this function resolves. Charging project quotas here would let a caller bind the file to a model without a quota while rows execute against a quota-limited model. + + Scopes with a ``tpd_limit`` (key, team, end user) are charged against a + daily token descriptor instead of their per-minute RPM/TPM descriptor, + because a batch's rows are scheduled by the provider and never share a + minute with the submission. The daily descriptor uses its own key so + its 24h window never collides with the online limiter's counters. """ - return self.parallel_request_limiter._create_rate_limit_descriptors( + descriptors: Final = self.parallel_request_limiter._create_rate_limit_descriptors( user_api_key_dict=user_api_key_dict, data=data, rpm_limit_type=None, tpm_limit_type=None, model_has_failures=False, ) + tpd_limits: Final[Mapping[str, tuple[str, int]]] = MappingProxyType( + { + key: (value, limit) + for key, value, limit in ( + ("api_key", user_api_key_dict.api_key, user_api_key_dict.tpd_limit), + ("team", user_api_key_dict.team_id, user_api_key_dict.team_tpd_limit), + ("end_user", user_api_key_dict.end_user_id, user_api_key_dict.end_user_tpd_limit), + ) + if value and limit is not None + } + ) + if not tpd_limits: + return descriptors + return [ + *(d for d in descriptors if d["key"] not in tpd_limits), + *( + RateLimitDescriptor( + key=f"{key}{BATCH_TPD_DESCRIPTOR_SUFFIX}", + value=value, + rate_limit={ + "requests_per_unit": None, + "tokens_per_unit": limit, + "window_size": BATCH_TPD_WINDOW_SECONDS, + }, + ) + for key, (value, limit) in tpd_limits.items() + ), + ] @staticmethod def _project_has_any_io_token_limits(user_api_key_dict: UserAPIKeyAuth) -> bool: @@ -583,9 +624,14 @@ class _PROXY_BatchRateLimiter(CustomLogger): batch_usage: BatchFileUsage, limit_type: str, requested_model: str | None = None, + window_start: int | None = None, ) -> NoReturn: - """Raise :class:`ProxyRateLimitError` (a 429) for batch rate limit exceeded.""" - from datetime import datetime + """Raise :class:`ProxyRateLimitError` (a 429) for batch rate limit exceeded. + + ``window_start`` is the active counter window's start (unix seconds) when + known, so the reset time reflects that window's actual end rather than a + full window from now. + """ # Find the descriptor for this status. Matching on (key, value) is # required, not key alone: a batch can carry several project ITPM/OTPM @@ -609,9 +655,12 @@ class _PROXY_BatchRateLimiter(CustomLogger): descriptors[descriptor_index] if descriptors else {"key": "", "value": "", "rate_limit": None} ) - now: Final = datetime.now().timestamp() - window_size: Final = self.parallel_request_limiter.window_size - reset_time: Final = now + window_size + now: Final = self._time_provider().timestamp() + window_size: Final = (descriptor.get("rate_limit") or {}).get( + "window_size" + ) or self.parallel_request_limiter.window_size + reset_time: Final = now + window_size if window_start is None else window_start + window_size + retry_after: Final = max(0, int(reset_time - now)) reset_time_formatted: Final = datetime.fromtimestamp(reset_time).strftime("%Y-%m-%d %H:%M:%S UTC") remaining_display: Final = max(0, status["limit_remaining"]) @@ -643,10 +692,13 @@ class _PROXY_BatchRateLimiter(CustomLogger): if descriptor.get("key") == PROJECT_ITPM_DESCRIPTOR_KEY else batch_usage.total_tokens ) + token_limit_label: Final = ( + "TPD" if descriptor.get("key", "").endswith(BATCH_TPD_DESCRIPTOR_SUFFIX) else "TPM" + ) detail = ( f"Batch rate limit exceeded for {descriptor.get('key', 'unknown')}: {descriptor.get('value', 'unknown')}. " f"Batch contains {batch_token_count} tokens but only {remaining_display} tokens remaining " - f"out of {current_limit} TPM limit. " + f"out of {current_limit} {token_limit_label} limit. " f"Limit resets at: {reset_time_formatted}" ) @@ -654,7 +706,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): raise ProxyRateLimitError( detail=detail, headers={ - "retry-after": str(window_size), + "retry-after": str(retry_after), "rate_limit_type": limit_type, "reset_at": reset_time_formatted, }, @@ -712,6 +764,8 @@ class _PROXY_BatchRateLimiter(CustomLogger): parent_otel_span=user_api_key_dict.parent_otel_span, ) + stash: Final = get_or_create_request_stash() + stash.batch_tpd_refund_ops = () if rate_limit_response["overall_code"] == "OVER_LIMIT": requested_model: Final = data.get("model") if data else None for status in rate_limit_response["statuses"]: @@ -722,8 +776,70 @@ class _PROXY_BatchRateLimiter(CustomLogger): batch_usage, status["rate_limit_type"], requested_model=requested_model, + window_start=await self._read_tpd_window_start( + status=status, parent_otel_span=user_api_key_dict.parent_otel_span + ), ) + stash.batch_tpd_refund_ops = self._build_tpd_refund_ops( + descriptors=descriptors, + tokens=batch_usage.total_tokens, + reservation_windows=rate_limit_response.get("reservation_windows", frozenset()), + ) + + async def _read_tpd_window_start(self, status: "RateLimitStatus", parent_otel_span: "Span | None") -> int | None: + descriptor_key: Final = status.get("descriptor_key") or "" + if not descriptor_key.endswith(BATCH_TPD_DESCRIPTOR_SUFFIX): + return None + try: + window_start: Final = _WINDOW_START_ADAPTER.validate_python( + await self.parallel_request_limiter.internal_usage_cache.async_get_cache( + key=f"{{{descriptor_key}:{status.get('descriptor_value') or ''}}}:window", + litellm_parent_otel_span=parent_otel_span, + ), + strict=True, + ) + return None if window_start is None else int(float(window_start)) + except (ValidationError, ValueError): + return None + + def _build_tpd_refund_ops( + self, + descriptors: Sequence["RateLimitDescriptor"], + tokens: int, + reservation_windows: frozenset[tuple[str, str, Literal["redis", "local"]]], + ) -> tuple[ReservationAwareIncrementOperation, ...]: + """Refund operations for the daily token counters this batch charged. + + The v3 limiter's failure hook applies them when the submission fails + after the counters were incremented. Each operation carries the window + identity the charge landed in, so the refund is skipped once that + window has rolled over. + """ + if tokens <= 0 or not reservation_windows: + return () + tpd_descriptors_by_counter: Final[Mapping[str, RateLimitDescriptor]] = MappingProxyType( + { + self.parallel_request_limiter.create_rate_limit_keys( + descriptor["key"], descriptor["value"], "tokens" + ): descriptor + for descriptor in descriptors + if descriptor["key"].endswith(BATCH_TPD_DESCRIPTOR_SUFFIX) + } + ) + return tuple( + ReservationAwareIncrementOperation( + key=counter_key, + increment_value=-tokens, + ttl=BATCH_TPD_WINDOW_SECONDS, + window_key=f"{{{descriptor['key']}:{descriptor['value']}}}:window", + expected_window_start=window_start, + reservation_backend=backend, + ) + for counter_key, window_start, backend in sorted(reservation_windows) + if (descriptor := tpd_descriptors_by_counter.get(counter_key)) is not None + ) + async def count_input_file_usage( self, file_id: str, diff --git a/litellm/proxy/hooks/parallel_request_limiter_v3.py b/litellm/proxy/hooks/parallel_request_limiter_v3.py index a34dc99e472..8ca4124521a 100644 --- a/litellm/proxy/hooks/parallel_request_limiter_v3.py +++ b/litellm/proxy/hooks/parallel_request_limiter_v3.py @@ -396,6 +396,8 @@ CacheCounterValue: TypeAlias = int | float | str | bytes CacheCounterValues: TypeAlias = Sequence[CacheCounterValue | None] +ReservationWindowIdentity: TypeAlias = tuple[str, str, Literal["redis", "local"]] + ParallelGaugeCacheValue: TypeAlias = dict[str, object] | int | float | str | bytes @@ -542,6 +544,7 @@ class RequestRateLimiterStash: default_factory=frozenset ) batch_enqueued_reservation: BatchEnqueuedTokenReservation | None = None + batch_tpd_refund_ops: tuple[ReservationAwareIncrementOperation, ...] = () reservation_released: bool = False @@ -683,6 +686,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): self._batch_rate_limiter = _PROXY_BatchRateLimiter( internal_usage_cache=self.internal_usage_cache, parallel_request_limiter=self, + time_provider=self._time_provider, ) except Exception as e: verbose_proxy_logger.debug("Could not load batch rate limiter: %s", e) @@ -1823,6 +1827,7 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) applied: Final[list[list[AtomicCounterMeta]]] = [] statuses: Final[list[RateLimitStatus]] = [] + reservation_windows: Final[set[ReservationWindowIdentity]] = set() # mutable-ok: filled by the group loop raw: list[CacheCounterValue] for _idx, (keys, args, meta) in enumerate(descriptor_groups): @@ -1860,11 +1865,12 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): return response applied.append(meta) statuses.extend(response["statuses"]) + reservation_windows.update(response.get("reservation_windows", frozenset())) return RateLimitResponse( overall_code="OK", statuses=statuses, - reservation_windows=frozenset(), + reservation_windows=frozenset(reservation_windows), ) async def _refund_applied_descriptor_groups( @@ -4824,6 +4830,13 @@ class _PROXY_MaxParallelRequestsHandler_v3(CustomLogger): ) stash.batch_enqueued_reservation = None + if stash.batch_tpd_refund_ops: + await self.async_increment_reservation_aware_tokens( + pipeline_operations=stash.batch_tpd_refund_ops, + parent_otel_span=user_api_key_dict.parent_otel_span, + ) + stash.batch_tpd_refund_ops = () + if stash.reservation_released: return reserved_tokens: Final = stash.reserved_tokens diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 3a61f773001..1ae106be390 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -652,6 +652,10 @@ async def _update_database_and_spend_counters( request_tags: list[str] | None = None, model_access_groups: Sequence[str] | None = None, ) -> bool: + if budget_reservation is not None: + await _reconcile_budget_reservation_before_db_update( + budget_reservation=budget_reservation, response_cost=response_cost + ) try: charged: Final = await proxy_logging_obj.db_spend_update_writer.update_database( token=user_api_key, @@ -709,6 +713,30 @@ async def _update_database_and_spend_counters( return True +async def _reconcile_budget_reservation_before_db_update( + budget_reservation: dict, # mutable-ok: reconcile_budget_reservation stamps applied_adjustment on the caller's shared reservation dict + response_cost: float, +) -> None: + from litellm.proxy.spend_tracking.budget_reservation import reconcile_budget_reservation + + try: + await reconcile_budget_reservation( + budget_reservation=budget_reservation, actual_cost=response_cost, finalize=False + ) + except Exception: # noqa: BLE001 # a failed reconcile must not block the spend write; the counters are dropped instead + verbose_proxy_logger.warning( + "Failed to reconcile budget reservation before persisting spend; invalidating reserved counters" + ) + try: + await _invalidate_budget_reservation_counters(budget_reservation=budget_reservation) + except Exception: # noqa: BLE001 # nothing left to try; the finalized stamp below keeps it from being reprocessed + verbose_proxy_logger.exception( + "Failed to invalidate budget reservation counters after pre-persist reconcile failed" + ) + finally: + budget_reservation["finalized"] = True # rebind-ok: the counter update reads the stamp off the shared dict + + async def _release_budget_reservation(budget_reservation: dict | None) -> None: if budget_reservation is None: return diff --git a/litellm/proxy/management_endpoints/budget_management_endpoints.py b/litellm/proxy/management_endpoints/budget_management_endpoints.py index 81a607aaa43..e16ea4a812e 100644 --- a/litellm/proxy/management_endpoints/budget_management_endpoints.py +++ b/litellm/proxy/management_endpoints/budget_management_endpoints.py @@ -52,6 +52,7 @@ async def new_budget( - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. - tpm_limit: Optional[int] - The tokens per minute limit for the budget. - rpm_limit: Optional[int] - The requests per minute limit for the budget. + - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit. - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} - budget_reset_at: Optional[datetime] - Datetime when the initial budget is reset. Default is now. """ @@ -135,6 +136,7 @@ async def update_budget( - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. - tpm_limit: Optional[int] - The tokens per minute limit for the budget. - rpm_limit: Optional[int] - The requests per minute limit for the budget. + - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit. - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} - budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset. """ @@ -272,6 +274,7 @@ async def budget_settings( "max_parallel_requests": {"type": "Integer"}, "tpm_limit": {"type": "Integer"}, "rpm_limit": {"type": "Integer"}, + "tpd_limit": {"type": "Integer"}, "budget_duration": {"type": "String"}, "max_budget": {"type": "Float"}, "soft_budget": {"type": "Float"}, diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index d2d87331d55..b35bc01b4d0 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -335,6 +335,7 @@ async def new_end_user( - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). - tpm_limit: Optional[int] - [Not Implemented Yet] Specify tpm limit for a given customer (Tokens per minute) - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute) + - tpd_limit: Optional[int] - Specify tpd limit for a given customer (Tokens per day). Batch submissions are charged against it instead of tpm_limit/rpm_limit - model_max_budget: Optional[dict] - [Not Implemented Yet] Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}} - max_parallel_requests: Optional[int] - [Not Implemented Yet] Specify max parallel requests for a given customer. - soft_budget: Optional[float] - [Not Implemented Yet] Get alerts when customer crosses given budget, doesn't block requests. diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 62019f462a2..ee8ae66ea11 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1078,7 +1078,9 @@ async def validate_team_id_used_in_service_account_request( return True -_BUDGET_NUMERIC_KEYS = frozenset(["max_budget", "soft_budget", "max_parallel_requests", "tpm_limit", "rpm_limit"]) +_BUDGET_NUMERIC_KEYS = frozenset( + ["max_budget", "soft_budget", "max_parallel_requests", "tpm_limit", "rpm_limit", "tpd_limit"] +) def _enforce_upperbound_key_params( @@ -1957,6 +1959,7 @@ async def generate_key_fn( - blocked: Optional[bool] - Whether the key is blocked. - rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute) - tpm_limit: Optional[int] - Specify tpm limit for a given key (Tokens per minute) + - tpd_limit: Optional[int] - Specify tpd limit for a given key (Tokens per day). Charged by batch submissions instead of tpm_limit/rpm_limit. - soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached. - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. @@ -2163,6 +2166,7 @@ async def generate_service_account_key_fn( - blocked: Optional[bool] - Whether the key is blocked. - rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute) - tpm_limit: Optional[int] - Specify tpm limit for a given key (Tokens per minute) + - tpd_limit: Optional[int] - Specify tpd limit for a given key (Tokens per day). Charged by batch submissions instead of tpm_limit/rpm_limit. - soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached. - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests) @@ -3192,6 +3196,7 @@ async def update_key_fn( - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"} - tpm_limit: Optional[int] - Tokens per minute limit - rpm_limit: Optional[int] - Requests per minute limit + - tpd_limit: Optional[int] - Tokens per day limit, charged by batch submissions instead of tpm_limit/rpm_limit - model_rpm_limit: Optional[dict] - Model-specific RPM limits {"gpt-4": 100, "claude-v1": 200} - mcp_rpm_limit: Optional[dict] - Per-MCP-server RPM limits, keyed by MCP server name {"github": 100, "slack": 200} - tag_rpm_limit: Optional[dict] - Per-request-tag RPM limits, keyed by request tag {"cell-1": 1000, "cell-2": 500}. Each tag gets an independent counter; absent tags fall back to the key-level rpm limit. @@ -4355,6 +4360,7 @@ async def generate_key_helper_fn( metadata: dict | None = {}, tpm_limit: int | None = None, rpm_limit: int | None = None, + tpd_limit: int | None = None, query_type: Literal["insert_data", "update_data"] = "insert_data", update_key_values: dict | None = None, key_alias: str | None = None, @@ -4503,6 +4509,7 @@ async def generate_key_helper_fn( "metadata": metadata_json, "tpm_limit": tpm_limit, "rpm_limit": rpm_limit, + "tpd_limit": tpd_limit, "budget_duration": key_budget_duration, "budget_reset_at": key_reset_at, "allowed_cache_controls": allowed_cache_controls, diff --git a/litellm/proxy/management_endpoints/management_v1/budgets.py b/litellm/proxy/management_endpoints/management_v1/budgets.py index cc2fefc426f..ea13e4547bd 100644 --- a/litellm/proxy/management_endpoints/management_v1/budgets.py +++ b/litellm/proxy/management_endpoints/management_v1/budgets.py @@ -58,6 +58,7 @@ class BudgetListItem(BaseModel): soft_budget: float | None = None tpm_limit: int | None = None rpm_limit: int | None = None + tpd_limit: int | None = None budget_duration: str | None = None budget_reset_at: datetime | None = None created_at: datetime @@ -123,7 +124,7 @@ BUDGET_FILTERS: Final[Mapping[str, FilterSpec]] = MappingProxyType( BUDGETS_LIST_SPEC: Final[ListSpec[BudgetListItem, BudgetListItem]] = ListSpec( resource="budgets", - sortable=frozenset(("budget_id", "max_budget", "tpm_limit", "rpm_limit", "created_at")), + sortable=frozenset(("budget_id", "max_budget", "tpm_limit", "rpm_limit", "tpd_limit", "created_at")), searchable=frozenset(("budget_id",)), filters=BUDGET_FILTERS, default_sort=(SortKey(field="created_at", descending=True),), @@ -154,7 +155,7 @@ async def list_budgets( way to page, sort or filter it. `sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`, - `rpm_limit` or `created_at`, each optionally prefixed with `-` for descending, + `rpm_limit`, `tpd_limit` or `created_at`, each optionally prefixed with `-` for descending, and defaults to `-created_at`. `budget_id` is appended to every sort as the tiebreaker. `q` is a case-insensitive substring match on `budget_id`. `page_size` defaults to 50 and is capped at 100. Filters are diff --git a/litellm/proxy/management_endpoints/organization_endpoints.py b/litellm/proxy/management_endpoints/organization_endpoints.py index 96e946424bd..c6a76a920f6 100644 --- a/litellm/proxy/management_endpoints/organization_endpoints.py +++ b/litellm/proxy/management_endpoints/organization_endpoints.py @@ -362,6 +362,7 @@ async def new_organization( - max_budget: *Optional[float]* - Max budget for org - tpm_limit: *Optional[int]* - Max tpm limit for org - rpm_limit: *Optional[int]* - Max rpm limit for org + - tpd_limit: *Optional[int]* - Max tokens per day stored on the org budget. Batch submissions enforce tpd_limit at the key, team and end user scopes only. - model_rpm_limit: *Optional[Dict[str, int]]* - The RPM (Requests Per Minute) limit per model for this organization. - model_tpm_limit: *Optional[Dict[str, int]]* - The TPM (Tokens Per Minute) limit per model for this organization. - max_parallel_requests: *Optional[int]* - [Not Implemented Yet] Max parallel requests for org diff --git a/litellm/proxy/management_endpoints/team_endpoints.py b/litellm/proxy/management_endpoints/team_endpoints.py index bc4b3610b1c..e719d6d761a 100644 --- a/litellm/proxy/management_endpoints/team_endpoints.py +++ b/litellm/proxy/management_endpoints/team_endpoints.py @@ -1217,6 +1217,7 @@ async def new_team( - mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team. - tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for this team - all keys with this team_id will have at max this TPM limit - rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit + - tpd_limit: Optional[int] - The TPD (Tokens Per Day) limit for this team. Batch submissions are charged against it instead of tpm_limit/rpm_limit - rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of RPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating RPM, or "best_effort_throughput" for best effort enforcement. - tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of TPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating TPM, or "best_effort_throughput" for best effort enforcement. - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget @@ -1969,6 +1970,7 @@ async def update_team( - metadata: Optional[dict] - Metadata for team, store information for team. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } - tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for this team - all keys with this team_id will have at max this TPM limit - rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit + - tpd_limit: Optional[int] - The TPD (Tokens Per Day) limit for this team. Batch submissions are charged against it instead of tpm_limit/rpm_limit - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set (either in the request or existing), soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set. - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets) diff --git a/litellm/proxy/management_helpers/access_group_key_sync.py b/litellm/proxy/management_helpers/access_group_key_sync.py index c9f93fae0d9..b9a28a2ebb3 100644 --- a/litellm/proxy/management_helpers/access_group_key_sync.py +++ b/litellm/proxy/management_helpers/access_group_key_sync.py @@ -38,7 +38,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.auth_checks import ( _delete_cache_access_object, # pyright: ignore[reportPrivateUsage] # the access-group endpoints reach for this same cache primitive ) -from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient +from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper from litellm.repositories.table_repositories import AccessGroupRepository @@ -75,7 +75,7 @@ _REPOINT_KEY_SQL: Final = ( def _raw_executor(prisma_client: object) -> _RawExecutor: """Narrow the untyped Prisma client down to the raw-query call this module makes, pinned to the writer.""" db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client - return WriterPinnedClient(db).db # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin + return writer_wrapper(db) # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin async def _invalidate_access_group_cache(access_group_id: str) -> None: diff --git a/litellm/proxy/management_helpers/access_group_model_sync.py b/litellm/proxy/management_helpers/access_group_model_sync.py index b9d81f2981f..7a8dcc2939c 100644 --- a/litellm/proxy/management_helpers/access_group_model_sync.py +++ b/litellm/proxy/management_helpers/access_group_model_sync.py @@ -10,7 +10,7 @@ from typing import Final, Protocol from pydantic import BaseModel -from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient +from litellm.proxy.db.routing_prisma_wrapper import writer_wrapper from litellm.proxy.management_helpers.access_group_team_sync import invalidate_access_group_caches from litellm.repositories.table_repositories import AccessGroupRepository from litellm.router import Router @@ -56,7 +56,7 @@ _REMOVE_MODEL_NAME_SQL: Final = ( def _raw_executor(prisma_client: object) -> _RawExecutor: db: Final = AccessGroupRepository(prisma_client).prisma_client.db # pyright: ignore[reportAny] # untyped Prisma client - return WriterPinnedClient(db).db # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin + return writer_wrapper(db) # pyright: ignore[reportAny, reportReturnType] # untyped Prisma client behind the pin def _config_sourced_sibling(llm_router: Router, deployment_id: str, model_id: str) -> bool: diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index b310fc661c4..fe9e104789b 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -8,6 +8,7 @@ import httpx import litellm from litellm._logging import verbose_proxy_logger +from litellm.litellm_core_utils.asyncify import asyncify from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.proxy._types import PassThroughEndpointLoggingResultValues @@ -60,7 +61,7 @@ class PassThroughStreamingHandler: litellm_logging_obj._update_completion_start_time(completion_start_time=datetime.now()) @staticmethod - def schedule_stream_failure_logging( + async def schedule_stream_failure_logging( litellm_logging_obj: LiteLLMLoggingObj, endpoint_type: EndpointType, request_body: dict[str, object], @@ -68,7 +69,7 @@ class PassThroughStreamingHandler: exception: Exception, stream_context: PassThroughStreamContext | None = None, ) -> None: - PassThroughStreamingHandler._record_partial_usage_for_failure( + await asyncify(PassThroughStreamingHandler._record_partial_usage_for_failure)( litellm_logging_obj=litellm_logging_obj, endpoint_type=endpoint_type, request_body=request_body, @@ -222,7 +223,7 @@ class PassThroughStreamingHandler: verbose_proxy_logger.error("Error in chunk_processor: %s", e) if response.status_code < 400: logging_scheduled = True - PassThroughStreamingHandler.schedule_stream_failure_logging( + await PassThroughStreamingHandler.schedule_stream_failure_logging( litellm_logging_obj=litellm_logging_obj, endpoint_type=endpoint_type, request_body=resolved_request_body, @@ -292,7 +293,7 @@ class PassThroughStreamingHandler: ( standard_logging_response_object, kwargs, - ) = PassThroughStreamingHandler._build_passthrough_logging_result( + ) = await asyncify(PassThroughStreamingHandler._build_passthrough_logging_result)( litellm_logging_obj=litellm_logging_obj, passthrough_success_handler_obj=passthrough_success_handler_obj, url_route=url_route, @@ -334,8 +335,8 @@ class PassThroughStreamingHandler: Synchronous, CPU-bound reconstruction of the standard logging payload from collected raw SSE bytes. Extracted from _route_streaming_logging_to_handler so the per-endpoint dispatch can - be unit-tested in isolation. Still invoked synchronously on the event - loop; an off-loop dispatch is a future change, not part of this PR. + be unit-tested in isolation. The async callers run it in a worker + thread so the token counts inside stay off the event loop. """ all_chunks: Final = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes) standard_logging_response_object: PassThroughEndpointLoggingResultValues | None = None diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 34b742a082b..f63e088ebf7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3061,7 +3061,7 @@ async def _reconcile_budget_reservation_for_counter_update( budget_reservation: dict | None, response_cost: float | None, ) -> set[str]: - if budget_reservation is None: + if budget_reservation is None or budget_reservation.get("finalized") is True: return set() from litellm.proxy.spend_tracking.budget_reservation import ( diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index dd7967aafe3..8072df5aa5b 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -17,6 +17,7 @@ model LiteLLM_BudgetTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? model_max_budget Json? budget_duration String? budget_reset_at DateTime? @@ -133,6 +134,7 @@ model LiteLLM_TeamTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? blocked Boolean @default(false) @@ -203,6 +205,7 @@ model LiteLLM_DeletedTeamTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? blocked Boolean @default(false) @@ -438,6 +441,7 @@ model LiteLLM_VerificationToken { blocked Boolean? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? max_budget Float? budget_duration String? budget_reset_at DateTime? @@ -534,6 +538,7 @@ model LiteLLM_DeletedVerificationToken { blocked Boolean? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? max_budget Float? budget_duration String? budget_reset_at DateTime? diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 6074a50a69b..373f2d0fe36 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -959,8 +959,9 @@ async def _set_reserved_entries_actual_cost( async def _reseed_reserved_entry(item: _EntryAdjustment, actual_cost: float) -> None: """Post-call reconcile / release of a counter that was flushed, expired or reseeded between reservation and - reconcile: the optimistic delta no longer applies, so reseed from the DB floor (which cannot include this - request's cost yet) and add the settled cost, since increment_spend_counters skips reserved keys.""" + reconcile: the optimistic delta no longer applies, so reseed from the DB floor and add the settled cost, since + increment_spend_counters skips reserved keys. The reconcile runs before this request's spend is enqueued to the + DB, so the reseeded floor excludes it.""" from litellm.proxy.proxy_server import _increment_spend_counter_cache, reseed_spend_counter_from_db reseeded: Final = await reseed_spend_counter_from_db(counter_key=item.counter_key) diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 4bcdf6aad22..56438fe45bd 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -157,6 +157,7 @@ def _get_spend_logs_metadata( user_api_key_team_alias=None, spend_logs_metadata=None, requester_ip_address=None, + user_agent=None, additional_usage_values=None, applied_guardrails=None, status="success", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index e611984bf4c..479bd0a55af 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -85,7 +85,11 @@ from litellm._logging import _redact_string, verbose_proxy_logger from litellm._service_logger import ServiceLogging, ServiceTypes from litellm.caching.caching import DualCache, RedisCache from litellm.caching.dual_cache import LimitedSizeOrderedDict -from litellm.exceptions import RejectedRequestError, SensitiveDataRouteException +from litellm.exceptions import ( + GuardrailRaisedException, + RejectedRequestError, + SensitiveDataRouteException, +) from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, @@ -901,6 +905,9 @@ def _call_type_for_route(route: str | None) -> str | None: return call_types[0].value if len(operations) == 1 else None +_PROXY_ONLY_LLM_API_ERRORS: Final = (HTTPException, ProxyException, GuardrailRaisedException) + + def _failure_fields_to_lift(request_data: Mapping[str, object]) -> Mapping[str, object]: """Failure-path callbacks run after ``litellm_logging_obj`` is popped from request_data (it is not serialisable), so the caller merges these fields @@ -2991,6 +2998,7 @@ class ProxyLogging: - Authentication Errors from user_api_key_auth - HTTP HTTPException (rate limit errors) - ProxyException (guardrail blocks, budget / rate-limit errors) + - GuardrailRaisedException (guardrail blocks / guardrail failures) """ ######################################################### @@ -3005,9 +3013,7 @@ class ProxyLogging: if not (RouteChecks.is_llm_api_route(route) or RouteChecks.is_info_route(route)): return False - return isinstance(original_exception, (HTTPException, ProxyException)) or ( - error_type == ProxyErrorTypes.auth_error - ) + return isinstance(original_exception, _PROXY_ONLY_LLM_API_ERRORS) or (error_type == ProxyErrorTypes.auth_error) async def _handle_logging_proxy_only_error( self, @@ -3563,8 +3569,9 @@ class ProxyLogging: yield chunk except (GeneratorExit, asyncio.CancelledError): raise - except Exception: - ProxyLogging._fire_deferred_stream_logging(request_data) + except Exception as e: + if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, e): + ProxyLogging._fire_deferred_stream_logging(request_data) raise ProxyLogging._fire_deferred_stream_logging(request_data) return @@ -3638,8 +3645,9 @@ class ProxyLogging: yield chunk except (GeneratorExit, asyncio.CancelledError): raise - except Exception: - ProxyLogging._fire_deferred_stream_logging(request_data) + except Exception as e: + if not ProxyLogging._discard_deferred_stream_logging_for_failure(request_data, e): + ProxyLogging._fire_deferred_stream_logging(request_data) raise # Fire deferred logging AFTER all guardrail end-of-stream blocks @@ -3735,6 +3743,23 @@ class ProxyLogging: logging_obj._deferred_stream_complete_args = None asyncio.create_task(_deferred_cb(*_args)) + @staticmethod + def _discard_deferred_stream_logging_for_failure(request_data: Mapping[str, object], error: Exception) -> bool: + """Drop the parked success dispatch for an assembled chat stream that ends in an error + ``post_call_failure_hook`` logs as a failure, billing its usage on the failure row instead. + Returns False when the parked dispatch should still be flushed by the caller.""" + logging_obj: Final = request_data.get("litellm_logging_obj") + if not isinstance(logging_obj, Logging): + return False + _args: Final[tuple[object, ...] | None] = getattr(logging_obj, "_deferred_stream_complete_args", None) + assembled: Final = _args[0] if _args else None + if not isinstance(error, _PROXY_ONLY_LLM_API_ERRORS) or not isinstance(assembled, ModelResponse): + return False + logging_obj._on_deferred_stream_complete = None + logging_obj._deferred_stream_complete_args = None + logging_obj.record_assembled_response_for_failure(assembled) + return True + async def _arelease_max_parallel_requests_on_disconnect( self, user_api_key_dict: UserAPIKeyAuth, @@ -4295,7 +4320,8 @@ class PrismaClient: t.spend AS team_spend, t.max_budget AS team_max_budget, t.tpm_limit AS team_tpm_limit, - t.rpm_limit AS team_rpm_limit + t.rpm_limit AS team_rpm_limit, + t.tpd_limit AS team_tpd_limit FROM "LiteLLM_VerificationToken" v LEFT JOIN "LiteLLM_TeamTable" t ON v.team_id = t.team_id; """, @@ -4734,6 +4760,7 @@ class PrismaClient: t.soft_budget AS team_soft_budget, t.tpm_limit AS team_tpm_limit, t.rpm_limit AS team_rpm_limit, + t.tpd_limit AS team_tpd_limit, t.models AS team_models, t.metadata AS team_metadata, t.blocked AS team_blocked, @@ -4751,6 +4778,7 @@ class PrismaClient: b.max_budget AS litellm_budget_table_max_budget, b.tpm_limit AS litellm_budget_table_tpm_limit, b.rpm_limit AS litellm_budget_table_rpm_limit, + b.tpd_limit AS litellm_budget_table_tpd_limit, b.model_max_budget as litellm_budget_table_model_max_budget, b.soft_budget as litellm_budget_table_soft_budget, o.metadata as organization_metadata, diff --git a/litellm/repositories/unit_of_work.py b/litellm/repositories/unit_of_work.py index a497d0580db..0cdce307f9b 100644 --- a/litellm/repositories/unit_of_work.py +++ b/litellm/repositories/unit_of_work.py @@ -24,12 +24,8 @@ from typing import Final from litellm.repositories.prisma_protocols import BatchTable, PrismaBatch -def _spend_reset_data(budget_reset_at: datetime | None, spend_decrement: float | None) -> Mapping[str, object]: - spend: Final[object] = ( - {"decrement": spend_decrement} # mutable-ok: prisma update payload must be a dict - if spend_decrement is not None - else 0 - ) +def _spend_reset_data(budget_reset_at: datetime | None, spend_decrement: float) -> Mapping[str, object]: + spend: Final[object] = {"decrement": spend_decrement} # mutable-ok: prisma update payload must be a dict return {"spend": spend, "budget_reset_at": budget_reset_at} # mutable-ok: prisma update payload must be a dict @@ -37,9 +33,7 @@ def _spend_reset_data(budget_reset_at: datetime | None, spend_decrement: float | class KeySpendResetWrites: table: BatchTable - def queue_spend_reset( - self, token: str, budget_reset_at: datetime | None, spend_decrement: float | None = None - ) -> None: + def queue_spend_reset(self, token: str, budget_reset_at: datetime | None, spend_decrement: float) -> None: self.table.update( where={"token": token}, # mutable-ok: prisma where filter must be a dict data=_spend_reset_data(budget_reset_at, spend_decrement), @@ -50,9 +44,7 @@ class KeySpendResetWrites: class UserSpendResetWrites: table: BatchTable - def queue_spend_reset( - self, user_id: str, budget_reset_at: datetime | None, spend_decrement: float | None = None - ) -> None: + def queue_spend_reset(self, user_id: str, budget_reset_at: datetime | None, spend_decrement: float) -> None: self.table.update( where={"user_id": user_id}, # mutable-ok: prisma where filter must be a dict data=_spend_reset_data(budget_reset_at, spend_decrement), @@ -63,9 +55,7 @@ class UserSpendResetWrites: class TeamSpendResetWrites: table: BatchTable - def queue_spend_reset( - self, team_id: str, budget_reset_at: datetime | None, spend_decrement: float | None = None - ) -> None: + def queue_spend_reset(self, team_id: str, budget_reset_at: datetime | None, spend_decrement: float) -> None: self.table.update( where={"team_id": team_id}, # mutable-ok: prisma where filter must be a dict data=_spend_reset_data(budget_reset_at, spend_decrement), diff --git a/litellm/router.py b/litellm/router.py index ef89d611075..5a89fdb72f1 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -79,7 +79,10 @@ from litellm.litellm_core_utils.core_helpers import ( from litellm.litellm_core_utils.coroutine_checker import coroutine_checker from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dd_tracing import tracer -from litellm.litellm_core_utils.get_llm_provider_logic import declared_authenticating_provider +from litellm.litellm_core_utils.get_llm_provider_logic import ( + declared_authenticating_provider, + is_registered_custom_provider, +) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.ptu_pricing import ( PTU_COST_ATTRIBUTION_ENV_VAR, @@ -167,6 +170,7 @@ from litellm.router_utils.cooldown_handlers import ( _get_cooldown_deployments, _set_cooldown_deployments, is_advisor_orchestration_failure, + is_caller_timeout_408, ) from litellm.router_utils.fallback_event_handlers import ( AttemptedFallbackTargets, @@ -3773,7 +3777,16 @@ class Router: self, deployment: dict, kwargs: dict, function_name: str | None = None ) -> Deployment: """ - Handle clientside credential + Build a per-request Deployment carrying the caller-supplied api_key/api_base, + with its own stable id for cooldown, logging, and cost-map identity. + + This deployment is deliberately never registered with the router (no + upsert_deployment/add_deployment call): doing so used to add it to + self.model_list under the shared model_name, which made a request-scoped, + caller-supplied provider credential a permanent, load-balanced deployment + that every other caller of that model group could be routed onto. Its + pricing is still registered directly, so a custom price configured on the + underlying deployment still applies to this call. """ model_info: Final = deployment.get("model_info", {}).copy() litellm_params: Final = deployment["litellm_params"].copy() @@ -3792,7 +3805,7 @@ class Router: litellm_params=LiteLLM_Params(**dynamic_litellm_params), model_info=model_info, ) - self.upsert_deployment(deployment=deployment_pydantic_obj) # add new deployment to router + Router._register_deployment_pricing(deployment=deployment_pydantic_obj) return deployment_pydantic_obj @staticmethod @@ -8297,6 +8310,13 @@ class Router: litellm_params: Final = kwargs.get("litellm_params", {}) _model_info: Final = litellm_params.get("model_info", {}) + if is_caller_timeout_408(kwargs, exception_status): + verbose_router_logger.debug( + "Router: Exiting 'deployment_callback_on_failure' without cooldown. " + "A timeout the caller set caused this 408, not the deployment's health." + ) + return False + exception_headers: Final = litellm.litellm_core_utils.exception_mapping_utils._get_response_headers( original_exception=exception ) @@ -9546,8 +9566,10 @@ class Router: ) # done reading model["litellm_params"] # Check if provider is supported: either in enum or JSON-configured - if custom_llm_provider not in litellm.provider_list and not JSONProviderRegistry.exists( - custom_llm_provider + if ( + custom_llm_provider not in litellm.provider_list + and not JSONProviderRegistry.exists(custom_llm_provider) + and not is_registered_custom_provider(custom_llm_provider) ): raise Exception(f"Unsupported provider - {custom_llm_provider}") @@ -9693,40 +9715,7 @@ class Router: # initialize client self._add_deployment(deployment=deployment) - _model_info_dict: Final[dict] = deployment.model_info.model_dump(exclude_none=True) - for field in CustomPricingLiteLLMParams.model_fields: - field_value = deployment.litellm_params.get(field) - if field_value is not None: - _model_info_dict[field] = field_value - - Router._inherit_builtin_base_rates_for_off_peak( - model_info=_model_info_dict, - backend_model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.custom_llm_provider, - ) - if _model_info_dict.get("input_cost_per_token") is not None: - Router._inherit_builtin_cache_pricing( - model_info=_model_info_dict, - backend_model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.custom_llm_provider, - ) - Router._inherit_builtin_tiered_output_rate( - model_info=_model_info_dict, - backend_model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.custom_llm_provider, - ) - - # Register custom pricing in litellm.model_cost. - # Mirrors _create_deployment() logic to ensure dynamically-added deployments - # (e.g., loaded from DB) also have their custom pricing registered. - # Without this, _is_model_cost_zero() cannot detect explicitly-configured - # zero-cost models, causing budget checks to block free models. - Router._register_deployment_in_model_cost( - model_id=deployment.model_info.id, - model_info=_model_info_dict, - model=deployment.litellm_params.model, - custom_llm_provider=deployment.litellm_params.custom_llm_provider, - ) + Router._register_deployment_pricing(deployment=deployment) # add to model names self._add_model_to_list_and_index_map(model=_deployment, model_id=deployment.model_info.id) @@ -9988,6 +9977,21 @@ class Router: ) return model_info + @staticmethod + def _register_deployment_pricing(deployment: Deployment) -> None: + """Register a deployment's custom/inherited pricing in ``litellm.model_cost``. + + Takes only a ``Deployment``, so it registers pricing for a deployment that + is never added to ``self.model_list`` (a per-request client-side-credential + deployment) just as readily as one that is. + """ + Router._register_deployment_in_model_cost( + model_id=deployment.model_info.id, + model_info=Router._deployment_model_cost_payload(deployment), + model=deployment.litellm_params.model, + custom_llm_provider=deployment.litellm_params.custom_llm_provider, + ) + @staticmethod def _register_deployment_in_model_cost( *, diff --git a/litellm/router_strategy/complexity_router/capability_classifier.py b/litellm/router_strategy/complexity_router/capability_classifier.py index 66ed9c36ed8..21046ff3421 100644 --- a/litellm/router_strategy/complexity_router/capability_classifier.py +++ b/litellm/router_strategy/complexity_router/capability_classifier.py @@ -202,10 +202,15 @@ def capability_classifier_system_prompt(mode: Literal["json_schema", "json_objec ) -def parse_capability_classifier_verdict(content: str) -> CapabilityClassifierVerdict: - """Parse raw JSON or the fenced JSON shape tolerated by Switchyard.""" +def unwrap_classifier_json(content: str) -> str: + """Remove the optional Markdown fence without repairing or weakening verdict JSON.""" text: Final = content.strip() if not text.startswith("```"): - return CapabilityClassifierVerdict.model_validate_json(text) + return text unfenced: Final = text.removeprefix("```").removeprefix("json").lstrip("\n\r") - return CapabilityClassifierVerdict.model_validate_json(unfenced.removesuffix("```").strip()) + return unfenced.removesuffix("```").strip() + + +def parse_capability_classifier_verdict(content: str) -> CapabilityClassifierVerdict: + """Parse raw JSON or the fenced JSON shape tolerated by Switchyard.""" + return CapabilityClassifierVerdict.model_validate_json(unwrap_classifier_json(content)) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 590ca03c863..6abad578517 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -81,6 +81,7 @@ from .capability_classifier import ( capability_classifier_response_format, capability_classifier_system_prompt, parse_capability_classifier_verdict, + unwrap_classifier_json, ) from .classification_rubrics import BUSINESS_TIER_CRITERIA, calibration_examples_section from .config import ( @@ -103,7 +104,7 @@ from .config import ( CustomDimension, TierDefinition, ) -from .llm_v2 import LLMV2TaskContext, LLMV2Verdict, llm_v2_response_format +from .llm_v2 import LLM_V2_PROMPT_VERSION, LLMV2Decision, LLMV2TaskContext, LLMV2Verdict, llm_v2_response_format from .stall_detector import detect_stalled_task if TYPE_CHECKING: @@ -1017,16 +1018,41 @@ class ClassificationOutcome(NamedTuple): ] classifier_cost: float | None = None capability_forecast: CapabilityClassifierForecast | None = None + llm_v2_forecast: LLMV2Decision | None = None def _with_signal(outcome: ClassificationOutcome, signal: str | None) -> ClassificationOutcome: return outcome if signal is None else outcome._replace(signals=(*outcome.signals, signal)) -def _with_capability_forecast( +def _with_llm_v2_forecast( + decision: StandardLoggingRoutingDecision, forecast: LLMV2Decision +) -> StandardLoggingRoutingDecision: + """Preserve full numeric precision for both solver forecasts and the applied policy.""" + enriched: Final[StandardLoggingRoutingDecision] = { + **decision, + "classifier_efficient_p_solve": forecast.verdict.forecasts.efficient.p_solve, + "classifier_capable_p_solve": forecast.verdict.forecasts.capable.p_solve, + **({"classifier_max_quality_gap": forecast.max_quality_gap} if forecast.selective_decision is None else {}), + "classifier_prompt_version": LLM_V2_PROMPT_VERSION, + } + if forecast.calibration_version is None: + return enriched + calibrated: Final[StandardLoggingRoutingDecision] = { + **enriched, + "classifier_calibrated_efficient_p_solve": forecast.efficient, + "classifier_calibrated_capable_p_solve": forecast.capable, + "classifier_calibration_version": forecast.calibration_version, + } + return calibrated + + +def _with_classifier_forecast( decision: StandardLoggingRoutingDecision, outcome: ClassificationOutcome ) -> StandardLoggingRoutingDecision: - """Attach the validated capability verdict and applied threshold to its decision record.""" + """Attach validated forecasts and their applied policy to the routing decision.""" + if outcome.llm_v2_forecast is not None: + return _with_llm_v2_forecast(decision, outcome.llm_v2_forecast) forecast: Final = outcome.capability_forecast if forecast is None: return decision @@ -1320,12 +1346,12 @@ class ComplexityRouter(CustomLogger): capability_config: Final = self.config.capability_classifier_config self._classifier_response_format: Mapping[str, object] | None = ( ( - llm_v2_response_format(self.config.llm_v2_config.response_format) - if self.config.llm_v2_config is not None - else capability_classifier_response_format( + capability_classifier_response_format( capability_config.response_format if capability_config is not None else "json_schema" ) if self.config.classifier_type == "capability" + else llm_v2_response_format(self.config.llm_v2_config.response_format) + if self.config.llm_v2_config is not None else type_to_response_format_param(_tier_classification_model(self.config.classifier_wire_labels())) ) if llm_classifier_configured @@ -1353,15 +1379,15 @@ class ComplexityRouter(CustomLogger): llm_config: Final = self.config.classifier_llm_config if llm_config is None: raise ValueError("classifier_llm_config is not set") - v2: Final = self.config.llm_v2_config - if v2 is not None: - pools: Final = self._tier_pools() - return v2.system_prompt(pools[v2.efficient_tier][0], pools[v2.capable_tier][0]) if self.config.classifier_type == "capability": capability: Final = self.config.capability_classifier_config return capability_classifier_system_prompt( capability.response_format if capability is not None else "json_schema" ) + v2: Final = self.config.llm_v2_config + if v2 is not None: + pools: Final = self._tier_pools() + return v2.system_prompt(pools[v2.efficient_tier][0], pools[v2.capable_tier][0]) definitions: Final = self.config.tier_definitions if definitions is not None: return custom_tier_classification_prompt( @@ -2001,7 +2027,9 @@ class ComplexityRouter(CustomLogger): except Exception as e: # noqa: BLE001 -- external LLM call can fail in many distinct ways (timeout, provider error, validation, parse error); any failure must fall back to the configured fallback path if breaker is not None and permit is not None: breaker.record_failure(permit, is_timeout=_is_classifier_timeout(e)) - return self._classifier_failure_outcome(f"LLM classifier failed ({e})", prompt, system_prompt, scored) + return self._classifier_failure_outcome( + f"LLM classifier failed ({type(e).__name__})", prompt, system_prompt, scored + ) def _classifier_failure_outcome( self, @@ -2140,6 +2168,20 @@ class ComplexityRouter(CustomLogger): tier=tier, score=None, signals=("classifier-failed:default-model",), cause="default_model_fallback" ) + def _classifier_caller_constraints( + self, system_prompt: str | None, request_kwargs: Mapping[str, object] | None + ) -> str | None: + """Exclude Claude Code's environment and skill catalogs from task forecasts.""" + return ( + None + if any( + is_claude_code_user_agent(user_agent) + for metadata in (self._iter_metadata_dicts(request_kwargs) if request_kwargs is not None else ()) + if isinstance(user_agent := metadata.get("user_agent"), str) + ) + else system_prompt + ) + async def _classify_with_llm( self, prompt: str, @@ -2189,15 +2231,7 @@ class ComplexityRouter(CustomLogger): ) encrypted_task: Final = _encrypted_classifier_task(request_kwargs, marker_pairs) - caller_system_prompt: Final = ( - None - if any( - is_claude_code_user_agent(user_agent) - for metadata in (self._iter_metadata_dicts(request_kwargs) if request_kwargs is not None else ()) - if isinstance(user_agent := metadata.get("user_agent"), str) - ) - else system_prompt - ) + caller_system_prompt: Final = self._classifier_caller_constraints(system_prompt, request_kwargs) user_payload: Final = self._build_classifier_user_payload( prompt="The delegated task in the following agent_message." if encrypted_task is not None else prompt, system_prompt=caller_system_prompt, @@ -2229,57 +2263,6 @@ class ComplexityRouter(CustomLogger): raise ValueError(f"LLM classifier returned an unrecognized tier: {raw_tier!r}") return tier, classifier_cost - async def _classify_with_llm_v2( - self, - prompt: str, - system_prompt: str | None, - request_kwargs: Mapping[str, object] | None, - messages: Sequence[Mapping[str, object]] | None, - ) -> ClassificationOutcome: - v2: Final = self.config.llm_v2_config - if v2 is None or self._classifier_system_prompt is None: - raise ValueError("llm_v2_config is not set") - request: Final[Mapping[str, object]] = request_kwargs or MappingProxyType({}) - markers: Final = self._reminder_markers_for_request(request) - asks: Final = tuple(reversed(tuple(_iter_human_asks_newest_first(messages or (), markers)))) - encrypted: Final = _encrypted_classifier_task(request_kwargs, markers) - task_context: Final[LLMV2TaskContext] = { - "caller_constraints": system_prompt, - "task_and_follow_ups": asks or (prompt,), - } - task: Final = json.dumps(task_context) - image_parts: Final = self._classifier_image_parts(messages) - text_part: Final[ChatCompletionTextObject] = {"type": "text", "text": task} - user_content: Final[str | Sequence[ChatCompletionTextObject | ChatCompletionImageObject]] = ( - [text_part, *image_parts] if image_parts else task # mutable-ok: provider adapters require content arrays - ) - system_message: Final[ChatCompletionSystemMessage] = { - "role": "system", - "content": self._classifier_system_prompt, - } - user_message: Final[ChatCompletionUserMessage] = {"role": "user", "content": user_content} - messages_for_call: Final[list[AllMessageValues]] = [ # mutable-ok: Router requires an SDK message list - system_message, - user_message, - ] - content, classifier_cost = await self._call_classifier_model( - messages_for_call, request_kwargs, encrypted_task=encrypted, max_output_tokens=v2.max_output_tokens - ) - try: - verdict: Final = LLMV2Verdict.model_validate_json(content) - except ValidationError: - return self._classifier_failure_outcome("Invalid LLM V2 forecast", prompt, system_prompt)._replace( - classifier_cost=classifier_cost - ) - decision: Final = v2.classify(verdict) - return ClassificationOutcome( - tier=ComplexityTier(v2.efficient_tier if decision.use_efficient else v2.capable_tier), - score=None, - signals=decision.signals, - cause="llm_v2_classifier", - classifier_cost=classifier_cost, - ) - async def _classify_with_capability_llm( self, prompt: str, @@ -2347,6 +2330,62 @@ class ComplexityRouter(CustomLogger): ) return ComplexityTier(selected_tier), classifier_cost, forecast + async def _classify_with_llm_v2( + self, + prompt: str, + system_prompt: str | None, + request_kwargs: Mapping[str, object] | None, + messages: Sequence[Mapping[str, object]] | None, + ) -> ClassificationOutcome: + v2: Final = self.config.llm_v2_config + if v2 is None or self._classifier_system_prompt is None: + raise ValueError("llm_v2_config is not set") + request: Final[Mapping[str, object]] = request_kwargs or MappingProxyType({}) + markers: Final = self._reminder_markers_for_request(request) + encrypted: Final = _encrypted_classifier_task(request_kwargs, markers) + asks: Final = ( + ("The delegated task in the following agent_message.",) + if encrypted is not None + else tuple(reversed(tuple(_iter_human_asks_newest_first(messages or (), markers)))) + ) + task_context: Final[LLMV2TaskContext] = { + "caller_constraints": self._classifier_caller_constraints(system_prompt, request_kwargs), + "task_and_follow_ups": asks or (prompt,), + } + task: Final = json.dumps(task_context) + image_parts: Final = self._classifier_image_parts(messages) + text_part: Final[ChatCompletionTextObject] = {"type": "text", "text": task} + user_content: Final[str | Sequence[ChatCompletionTextObject | ChatCompletionImageObject]] = ( + [text_part, *image_parts] if image_parts else task # mutable-ok: provider adapters require content arrays + ) + system_message: Final[ChatCompletionSystemMessage] = { + "role": "system", + "content": self._classifier_system_prompt, + } + user_message: Final[ChatCompletionUserMessage] = {"role": "user", "content": user_content} + messages_for_call: Final[list[AllMessageValues]] = [ # mutable-ok: Router requires an SDK message list + system_message, + user_message, + ] + content, classifier_cost = await self._call_classifier_model( + messages_for_call, request_kwargs, encrypted_task=encrypted, max_output_tokens=v2.max_output_tokens + ) + try: + verdict: Final = LLMV2Verdict.model_validate_json(unwrap_classifier_json(content)) + except ValidationError: + return self._classifier_failure_outcome("Invalid LLM V2 forecast", prompt, system_prompt)._replace( + classifier_cost=classifier_cost + ) + decision: Final = v2.classify(verdict) + return ClassificationOutcome( + tier=ComplexityTier(v2.efficient_tier if decision.use_efficient else v2.capable_tier), + score=None, + signals=decision.signals, + cause="llm_v2_classifier", + classifier_cost=classifier_cost, + llm_v2_forecast=decision, + ) + async def _call_classifier_model( self, messages_for_call: list[AllMessageValues], # mutable-ok: SDKs require a list @@ -4470,5 +4509,5 @@ class ComplexityRouter(CustomLogger): model=routed_model, messages=messages if has_original_messages else None, litellm_params=tier_litellm_params, - routing_decision=_with_capability_forecast(routing_decision, outcome), + routing_decision=_with_classifier_forecast(routing_decision, outcome), ) diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 8e3ab01332c..370589d7da4 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -970,9 +970,8 @@ class ComplexityRouterConfig(BaseModel): default="heuristic", description=( "Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, " - "an LLM tier-selection call, a Switchyard-compatible capability forecast, a joint task-demand and " - "capability forecast, a custom classifier plugin, 'heuristic_first', which scores locally and only " - "pays for the LLM classifier when the " + "an LLM tier-selection call, a Switchyard-compatible capability forecast, a joint Fuse V2 forecast, " + "a custom classifier plugin, 'heuristic_first', which scores locally and only pays for the LLM classifier when the " "local scorer does not confidently land a cheap tier, or 'hybrid', which trusts the local scorer " "everywhere except when its score lands near a tier boundary" ), @@ -1594,6 +1593,8 @@ class ComplexityRouterConfig(BaseModel): return self if v2 is None: raise ValueError("llm_v2_config is required when classifier_type is llm_v2") + if self.classifier_fallback != "heuristic": + raise ValueError("llm_v2 always fails closed to capable_tier; classifier_fallback cannot override it") llm: Final = self.classifier_llm_config if self.adaptive or self.tier_definitions is not None or self.enable_non_reasoning_tier: raise ValueError("llm_v2 requires two built-in tiers and adaptive=false") diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index 027f0a9ca05..6e6d4c253e9 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -9,6 +9,7 @@ Router cooldown handlers import asyncio import math from collections.abc import Mapping +from datetime import datetime from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final @@ -637,3 +638,23 @@ def cast_exception_status_to_int(exception_status: str | int) -> int: ) exception_status = 500 return exception_status + + +def is_caller_timeout_408( + model_call_details: Mapping[str, object], exception_status: str | int, ended: datetime | None = None +) -> bool: + """A 408 that arrives before the caller-set timeout could have fired came from the provider. + + ``ended`` overrides ``model_call_details["end_time"]`` for callers that run before the + failure logger has stamped the current API call's end time.""" + if cast_exception_status_to_int(exception_status) != 408: + return False + litellm_params: Final = model_call_details.get("litellm_params") + if not isinstance(litellm_params, Mapping) or not litellm_params.get("client_side_timeout"): + return False + timeout: Final = litellm_params.get("timeout") + started: Final = model_call_details.get("api_call_start_time") or model_call_details.get("start_time") + finished: Final = ended if ended is not None else model_call_details.get("end_time") + if not isinstance(timeout, (int, float)) or not isinstance(started, datetime) or not isinstance(finished, datetime): + return False + return (finished - started).total_seconds() >= timeout diff --git a/litellm/router_utils/fallback_event_handlers.py b/litellm/router_utils/fallback_event_handlers.py index 7fda5d96fb0..94164d0ea0c 100644 --- a/litellm/router_utils/fallback_event_handlers.py +++ b/litellm/router_utils/fallback_event_handlers.py @@ -2,7 +2,9 @@ import hashlib import json from collections.abc import Mapping, Sequence from dataclasses import dataclass +from datetime import datetime from enum import Enum +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import litellm @@ -20,6 +22,7 @@ from litellm.router_utils.cooldown_handlers import ( _set_cooldown_deployments, # pyright: ignore[reportPrivateUsage] - shared helper, used across router_utils cast_exception_status_to_int, is_advisor_orchestration_failure, + is_caller_timeout_408, ) from litellm.router_utils.router_callbacks.track_deployment_metrics import ( increment_deployment_failures_for_current_minute, @@ -36,12 +39,14 @@ else: # Status codes a generic API call's caller-supplied resource id can trigger on its own # (e.g. a nonexistent file/batch/thread id), independent of the selected deployment's health. _REQUEST_SCOPED_STATUS_CODES: Final = frozenset((404,)) +_NO_MODEL_CALL_DETAILS: Final[Mapping[str, object]] = MappingProxyType({}) def _trigger_cooldown_for_failed_deployment( litellm_router: LitellmRouter, kwargs: Mapping[str, object], exception: Exception, + model_call_details: Mapping[str, object] = _NO_MODEL_CALL_DETAILS, ) -> None: """ Trigger cooldown for a failed fallback deployment. @@ -80,7 +85,11 @@ def _trigger_cooldown_for_failed_deployment( # timeout, which litellm.Timeout reports as status 408 regardless of the deployment's # actual health. Left unguarded, a caller could force a 408 on every deployment in # the fallback chain from a single request with a near-zero timeout. - if kwargs.get("client_side_timeout") and cast_exception_status_to_int(exception_status) == 408: + if is_caller_timeout_408( + model_call_details, + exception_status, + ended=datetime.now(), # noqa: DTZ005 # naive to match the logging pipeline's api_call_start_time + ): verbose_router_logger.debug( "Not triggering cooldown for fallback deployment: a caller-supplied " "x-litellm-timeout caused this 408, not deployment health." @@ -579,6 +588,7 @@ async def run_async_fallback( litellm_router=litellm_router, kwargs=kwargs, exception=e, + model_call_details=logging_obj.model_call_details, ) raise error_from_fallbacks diff --git a/litellm/types/utils.py b/litellm/types/utils.py index ac62ce42bb5..fdf533fb4e9 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -2990,6 +2990,12 @@ class StandardLoggingRoutingDecision(TypedDict, total=False): classifier_p_solve: float # writable-ok: added only when a capability verdict is available classifier_calibrated_p_solve: ReadOnly[float] classifier_calibration_version: ReadOnly[str] + classifier_efficient_p_solve: ReadOnly[float] + classifier_capable_p_solve: ReadOnly[float] + classifier_calibrated_efficient_p_solve: ReadOnly[float] + classifier_calibrated_capable_p_solve: ReadOnly[float] + classifier_max_quality_gap: ReadOnly[float] + classifier_prompt_version: ReadOnly[str] classifier_threshold: float # writable-ok: added only when a capability verdict is available escalated: bool context_escalated: bool # writable-ok: Pydantic warns on ReadOnly TypedDict fields @@ -3026,6 +3032,12 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset( "classifier_p_solve", "classifier_calibrated_p_solve", "classifier_calibration_version", + "classifier_efficient_p_solve", + "classifier_capable_p_solve", + "classifier_calibrated_efficient_p_solve", + "classifier_calibrated_capable_p_solve", + "classifier_max_quality_gap", + "classifier_prompt_version", "classifier_threshold", "escalated", "context_escalated", diff --git a/schema.prisma b/schema.prisma index dd7967aafe3..8072df5aa5b 100644 --- a/schema.prisma +++ b/schema.prisma @@ -17,6 +17,7 @@ model LiteLLM_BudgetTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? model_max_budget Json? budget_duration String? budget_reset_at DateTime? @@ -133,6 +134,7 @@ model LiteLLM_TeamTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? blocked Boolean @default(false) @@ -203,6 +205,7 @@ model LiteLLM_DeletedTeamTable { max_parallel_requests Int? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? budget_duration String? budget_reset_at DateTime? blocked Boolean @default(false) @@ -438,6 +441,7 @@ model LiteLLM_VerificationToken { blocked Boolean? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? max_budget Float? budget_duration String? budget_reset_at DateTime? @@ -534,6 +538,7 @@ model LiteLLM_DeletedVerificationToken { blocked Boolean? tpm_limit BigInt? rpm_limit BigInt? + tpd_limit BigInt? max_budget Float? budget_duration String? budget_reset_at DateTime? diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 06a39d8da20..ee5b42fe0b7 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -16,6 +16,7 @@ longer signal it. ### Added +- **team_member_add**: `tpm_limit`, `rpm_limit`, `budget_duration`, and `allowed_models` attributes on `litellm_team_member_add`, applied to every member of the resource; `budget_duration` and `allowed_models` ride on `/team/member_add`, while the limits are sent through `/team/member_update`, which is where the proxy accepts them - **team**: Optional `team_id` argument on `litellm_team`, so teams can be created with a stable, human-readable ID instead of a provider-generated UUID; changing it forces replacement - **jwt_key_mapping**: New `litellm_jwt_key_mapping` resource for the proxy's JWT to virtual key mappings, so JWT clients identified by a claim (`client_id`, `azp`, `sub`) map to virtual keys and inherit their models, budgets and rate limits. Supports `description` and `is_active`, rotating the mapped key in place, and forces replacement when the claim name or value changes - **team**: `soft_budget`, `tags`, and `soft_budget_alerting_emails` attributes on `litellm_team`, matching what `/team/new` and `/team/update` already accept; `soft_budget_alerting_emails` is sent under `metadata`, where the proxy reads it diff --git a/terraform/provider/docs/resources/team_member_add.md b/terraform/provider/docs/resources/team_member_add.md index f5398e49d9c..bad241cddec 100644 --- a/terraform/provider/docs/resources/team_member_add.md +++ b/terraform/provider/docs/resources/team_member_add.md @@ -27,6 +27,10 @@ resource "litellm_team_member_add" "example" { } max_budget_in_team = 100.0 + budget_duration = "30d" + tpm_limit = 100000 + rpm_limit = 100 + allowed_models = ["gpt-4"] } ``` @@ -152,6 +156,12 @@ resource "litellm_team_member_add" "budget_example" { * `user_email` - (Optional) The email of the user to add to the team. * `role` - (Required) The role of the user in the team. Must be one of: "admin" or "user". * `max_budget_in_team` - (Optional) The maximum budget allocated for the team members. +* `budget_duration` - (Optional) Duration after which each member's budget resets, for example "1h", "24h", "7d", "30d". If not set, the budget never resets. +* `tpm_limit` - (Optional) Tokens per minute limit applied to each team member. Sent via `/team/member_update` after members are added, since `/team/member_add` does not accept it. +* `rpm_limit` - (Optional) Requests per minute limit applied to each team member. Sent via `/team/member_update` after members are added, since `/team/member_add` does not accept it. +* `allowed_models` - (Optional) List of models each team member can access. If not set, members inherit the team's `default_team_member_models` or all team models. + +Removing `budget_duration`, `tpm_limit`, `rpm_limit`, or `allowed_models` from the configuration clears that setting on every member through `/team/member_update`. ## Import diff --git a/terraform/provider/litellm/resource_team_member_add.go b/terraform/provider/litellm/resource_team_member_add.go index da5c7a6ebd7..ca3541408ba 100644 --- a/terraform/provider/litellm/resource_team_member_add.go +++ b/terraform/provider/litellm/resource_team_member_add.go @@ -49,10 +49,105 @@ func resourceLiteLLMTeamMemberAdd() *schema.Resource { Type: schema.TypeFloat, Optional: true, }, + "tpm_limit": { + Type: schema.TypeInt, + Optional: true, + }, + "rpm_limit": { + Type: schema.TypeInt, + Optional: true, + }, + "budget_duration": { + Type: schema.TypeString, + Optional: true, + }, + "allowed_models": { + Type: schema.TypeList, + Optional: true, + Elem: &schema.Schema{Type: schema.TypeString}, + }, }, } } +func expandAllowedModels(raw []interface{}) []string { + models := make([]string, 0, len(raw)) + for _, m := range raw { + models = append(models, m.(string)) + } + return models +} + +func applyAddOnlySettings(d *schema.ResourceData, payload map[string]interface{}) { + if v, ok := d.GetOk("budget_duration"); ok { + payload["budget_duration"] = v.(string) + } + if v, ok := d.GetOk("allowed_models"); ok { + payload["allowed_models"] = expandAllowedModels(v.([]interface{})) + } +} + +func applyLimits(d *schema.ResourceData, payload map[string]interface{}) { + for _, key := range []string{"tpm_limit", "rpm_limit"} { + if v, ok := d.GetOk(key); ok { + payload[key] = v.(int) + } + } +} + +func applyUpdateSettings(d *schema.ResourceData, payload map[string]interface{}) { + applyAddOnlySettings(d, payload) + applyLimits(d, payload) + for _, key := range []string{"tpm_limit", "rpm_limit", "budget_duration"} { + if _, ok := d.GetOk(key); !ok && d.HasChange(key) { + payload[key] = nil + } + } + if _, ok := d.GetOk("allowed_models"); !ok && d.HasChange("allowed_models") { + payload["allowed_models"] = []string{} + } +} + +func memberIdentity(member map[string]interface{}, payload map[string]interface{}) { + if userID, ok := member["user_id"].(string); ok && userID != "" { + payload["user_id"] = userID + } + if userEmail, ok := member["user_email"].(string); ok && userEmail != "" { + payload["user_email"] = userEmail + } +} + +// tpm/rpm limits are only accepted by /team/member_update, not /team/member_add +func setMemberLimits(client *Client, d *schema.ResourceData, teamID string, members []map[string]interface{}) error { + limits := map[string]interface{}{} + applyLimits(d, limits) + if len(limits) == 0 { + return nil + } + for _, member := range members { + updateData := map[string]interface{}{ + "team_id": teamID, + } + for k, v := range limits { + updateData[k] = v + } + memberIdentity(member, updateData) + + log.Printf("[DEBUG] Set team member limits request payload: %+v", updateData) + + resp, err := MakeRequest(client, "POST", "/team/member_update", updateData) + if err != nil { + return fmt.Errorf("error setting team member limits: %v", err) + } + defer resp.Body.Close() + + if err := handleResponse(resp, "setting team member limits"); err != nil { + return err + } + } + return nil +} + func resourceLiteLLMTeamMemberAddCreate(d *schema.ResourceData, m interface{}) error { client := m.(*Client) @@ -81,6 +176,7 @@ func resourceLiteLLMTeamMemberAddCreate(d *schema.ResourceData, m interface{}) e "team_id": teamID, "max_budget_in_team": maxBudget, } + applyAddOnlySettings(d, memberData) log.Printf("[DEBUG] Create team members request payload: %+v", memberData) @@ -94,9 +190,12 @@ func resourceLiteLLMTeamMemberAddCreate(d *schema.ResourceData, m interface{}) e return err } - // Set ID as team_id since this resource manages all members for a team d.SetId(teamID) + if err := setMemberLimits(client, d, teamID, membersList); err != nil { + return err + } + return resourceLiteLLMTeamMemberAddRead(d, m) } @@ -140,11 +239,13 @@ func resourceLiteLLMTeamMemberAddUpdate(d *schema.ResourceData, m interface{}) e // Track which members have been updated to avoid duplicates updatedMembers := make(map[string]bool) - // Check if max_budget_in_team has changed - if d.HasChange("max_budget_in_team") { - log.Printf("[DEBUG] max_budget_in_team changed, updating all existing members with new budget: %f", maxBudget) + // Check if any team-wide member setting has changed + settingsChanged := d.HasChange("max_budget_in_team") || d.HasChange("tpm_limit") || d.HasChange("rpm_limit") || + d.HasChange("budget_duration") || d.HasChange("allowed_models") + if settingsChanged { + log.Printf("[DEBUG] Member settings changed, updating all existing members") - // Update ALL existing members with the new budget + // Update ALL existing members with the new settings for key, newMember := range newMemberMap { if _, exists := oldMemberMap[key]; exists { updateData := map[string]interface{}{ @@ -152,22 +253,18 @@ func resourceLiteLLMTeamMemberAddUpdate(d *schema.ResourceData, m interface{}) e "role": newMember["role"].(string), "max_budget_in_team": maxBudget, } - if userID, ok := newMember["user_id"].(string); ok && userID != "" { - updateData["user_id"] = userID - } - if userEmail, ok := newMember["user_email"].(string); ok && userEmail != "" { - updateData["user_email"] = userEmail - } + applyUpdateSettings(d, updateData) + memberIdentity(newMember, updateData) - log.Printf("[DEBUG] Update team member budget request payload: %+v", updateData) + log.Printf("[DEBUG] Update team member settings request payload: %+v", updateData) resp, err := MakeRequest(client, "POST", "/team/member_update", updateData) if err != nil { - return fmt.Errorf("error updating team member budget: %v", err) + return fmt.Errorf("error updating team member settings: %v", err) } defer resp.Body.Close() - if err := handleResponse(resp, "updating team member budget"); err != nil { + if err := handleResponse(resp, "updating team member settings"); err != nil { return err } @@ -220,12 +317,8 @@ func resourceLiteLLMTeamMemberAddUpdate(d *schema.ResourceData, m interface{}) e "role": newMember["role"].(string), "max_budget_in_team": maxBudget, } - if userID, ok := newMember["user_id"].(string); ok && userID != "" { - updateData["user_id"] = userID - } - if userEmail, ok := newMember["user_email"].(string); ok && userEmail != "" { - updateData["user_email"] = userEmail - } + applyUpdateSettings(d, updateData) + memberIdentity(newMember, updateData) log.Printf("[DEBUG] Update team member request payload: %+v", updateData) @@ -265,6 +358,7 @@ func resourceLiteLLMTeamMemberAddUpdate(d *schema.ResourceData, m interface{}) e "team_id": teamID, "max_budget_in_team": maxBudget, } + applyAddOnlySettings(d, memberData) log.Printf("[DEBUG] Adding new team members request payload: %+v", memberData) @@ -277,6 +371,10 @@ func resourceLiteLLMTeamMemberAddUpdate(d *schema.ResourceData, m interface{}) e if err := handleResponse(resp, "adding team members"); err != nil { return err } + + if err := setMemberLimits(client, d, teamID, membersToAdd); err != nil { + return err + } } return resourceLiteLLMTeamMemberAddRead(d, m) diff --git a/terraform/provider/litellm/resource_team_member_add_test.go b/terraform/provider/litellm/resource_team_member_add_test.go new file mode 100644 index 00000000000..a2ddb0016bc --- /dev/null +++ b/terraform/provider/litellm/resource_team_member_add_test.go @@ -0,0 +1,274 @@ +package litellm + +import ( + "encoding/json" + "io" + "net/http" + "net/http/httptest" + "reflect" + "testing" + + "github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema" + "github.com/hashicorp/terraform-plugin-sdk/v2/terraform" +) + +func TestTeamMemberAddCreateSendsMemberSettings(t *testing.T) { + var addPayload map[string]interface{} + var updatePayloads []map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + var payload map[string]interface{} + json.Unmarshal(body, &payload) + switch r.URL.Path { + case "/team/member_add": + addPayload = payload + case "/team/member_update": + updatePayloads = append(updatePayloads, payload) + default: + t.Errorf("unexpected request path: %s", r.URL.Path) + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMTeamMemberAdd().Schema, map[string]interface{}{ + "team_id": "team-1", + "member": []interface{}{ + map[string]interface{}{ + "user_id": "user-1", + "role": "user", + }, + }, + "max_budget_in_team": 25.0, + "tpm_limit": 1000, + "rpm_limit": 10, + "budget_duration": "30d", + "allowed_models": []interface{}{"claude-opus-4-6-v1"}, + }) + + if err := resourceLiteLLMTeamMemberAddCreate(d, client); err != nil { + t.Fatalf("create failed: %v", err) + } + + if addPayload["budget_duration"] != "30d" { + t.Fatalf("member_add payload sent budget_duration %v, want 30d", addPayload["budget_duration"]) + } + wantModels := []interface{}{"claude-opus-4-6-v1"} + if !reflect.DeepEqual(addPayload["allowed_models"], wantModels) { + t.Fatalf("member_add payload sent allowed_models %v, want %v", addPayload["allowed_models"], wantModels) + } + if _, ok := addPayload["tpm_limit"]; ok { + t.Fatalf("member_add payload must not carry tpm_limit, got %v", addPayload["tpm_limit"]) + } + + if len(updatePayloads) != 1 { + t.Fatalf("expected 1 member_update call for limits, got %d", len(updatePayloads)) + } + update := updatePayloads[0] + if update["tpm_limit"] != float64(1000) { + t.Fatalf("member_update payload sent tpm_limit %v, want 1000", update["tpm_limit"]) + } + if update["rpm_limit"] != float64(10) { + t.Fatalf("member_update payload sent rpm_limit %v, want 10", update["rpm_limit"]) + } + if update["user_id"] != "user-1" { + t.Fatalf("member_update payload sent user_id %v, want user-1", update["user_id"]) + } +} + +func TestTeamMemberAddCreateOmitsUnsetSettings(t *testing.T) { + var addPayload map[string]interface{} + updateCalls := 0 + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + body, _ := io.ReadAll(r.Body) + switch r.URL.Path { + case "/team/member_add": + json.Unmarshal(body, &addPayload) + case "/team/member_update": + updateCalls++ + } + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMTeamMemberAdd().Schema, map[string]interface{}{ + "team_id": "team-1", + "member": []interface{}{ + map[string]interface{}{ + "user_id": "user-1", + "role": "user", + }, + }, + }) + + if err := resourceLiteLLMTeamMemberAddCreate(d, client); err != nil { + t.Fatalf("create failed: %v", err) + } + + for _, field := range []string{"tpm_limit", "rpm_limit", "budget_duration", "allowed_models"} { + if _, ok := addPayload[field]; ok { + t.Fatalf("member_add payload must not carry unset %s, got %v", field, addPayload[field]) + } + } + if updateCalls != 0 { + t.Fatalf("expected no member_update calls without limits, got %d", updateCalls) + } +} + +func TestTeamMemberAddCreateSetsIDBeforeLimitsFail(t *testing.T) { + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Content-Type", "application/json") + if r.URL.Path == "/team/member_update" { + w.WriteHeader(http.StatusInternalServerError) + w.Write([]byte(`{"error":"boom"}`)) + return + } + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + })) + defer srv.Close() + + client := NewClient(srv.URL, "test-key", true) + d := schema.TestResourceDataRaw(t, resourceLiteLLMTeamMemberAdd().Schema, map[string]interface{}{ + "team_id": "team-1", + "member": []interface{}{ + map[string]interface{}{ + "user_id": "user-1", + "role": "user", + }, + }, + "tpm_limit": 1000, + }) + + if err := resourceLiteLLMTeamMemberAddCreate(d, client); err == nil { + t.Fatal("create should fail when member_update fails") + } + if d.Id() != "team-1" { + t.Fatalf("resource ID = %q after failed limits call, want team-1 so Terraform can taint and recreate it", d.Id()) + } +} + +// newTeamMemberUpdateResourceData builds a ResourceData with one member in state +// and a real old -> new diff on the scalar settings, so d.HasChange and d.GetOk +// behave as they do during a real Update call +func newTeamMemberUpdateResourceData(t *testing.T, old, new map[string]string) *schema.ResourceData { + t.Helper() + attrs := map[string]string{ + "team_id": "team-1", + "member.#": "1", + "member.1.user_id": "user-1", + "member.1.user_email": "", + "member.1.role": "user", + "allowed_models.#": "0", + "max_budget_in_team": "25", + } + for k, v := range old { + attrs[k] = v + } + diffAttrs := map[string]*terraform.ResourceAttrDiff{} + for k, v := range new { + diffAttrs[k] = &terraform.ResourceAttrDiff{Old: attrs[k], New: v} + } + for k := range old { + if _, ok := new[k]; !ok { + diffAttrs[k] = &terraform.ResourceAttrDiff{Old: attrs[k], New: "", NewRemoved: true} + } + } + state := &terraform.InstanceState{ID: "team-1", Attributes: attrs} + d, err := schema.InternalMap(resourceLiteLLMTeamMemberAdd().Schema).Data(state, &terraform.InstanceDiff{Attributes: diffAttrs}) + if err != nil { + t.Fatalf("building ResourceData returned error: %v", err) + } + return d +} + +func runTeamMemberUpdate(t *testing.T, d *schema.ResourceData) []map[string]interface{} { + t.Helper() + var updatePayloads []map[string]interface{} + srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != "/team/member_update" { + t.Errorf("unexpected request path: %s", r.URL.Path) + } + body, _ := io.ReadAll(r.Body) + var payload map[string]interface{} + json.Unmarshal(body, &payload) + updatePayloads = append(updatePayloads, payload) + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(http.StatusOK) + w.Write([]byte(`{}`)) + })) + defer srv.Close() + + if err := resourceLiteLLMTeamMemberAddUpdate(d, NewClient(srv.URL, "test-key", true)); err != nil { + t.Fatalf("update failed: %v", err) + } + if len(updatePayloads) != 1 { + t.Fatalf("expected 1 member_update call, got %d", len(updatePayloads)) + } + return updatePayloads +} + +func TestTeamMemberAddUpdateSendsChangedSettings(t *testing.T) { + d := newTeamMemberUpdateResourceData(t, + map[string]string{"tpm_limit": "1000", "rpm_limit": "10", "budget_duration": "30d"}, + map[string]string{"tpm_limit": "500", "rpm_limit": "5", "budget_duration": "7d", "allowed_models.#": "1", "allowed_models.0": "gpt-5.2"}, + ) + + update := runTeamMemberUpdate(t, d)[0] + if update["tpm_limit"] != float64(500) || update["rpm_limit"] != float64(5) { + t.Fatalf("member_update payload limits = %v/%v, want 500/5", update["tpm_limit"], update["rpm_limit"]) + } + if update["budget_duration"] != "7d" { + t.Fatalf("member_update payload budget_duration = %v, want 7d", update["budget_duration"]) + } + if !reflect.DeepEqual(update["allowed_models"], []interface{}{"gpt-5.2"}) { + t.Fatalf("member_update payload allowed_models = %v, want [gpt-5.2]", update["allowed_models"]) + } + if update["user_id"] != "user-1" { + t.Fatalf("member_update payload user_id = %v, want user-1", update["user_id"]) + } +} + +func TestTeamMemberAddUpdateClearsRemovedSettings(t *testing.T) { + d := newTeamMemberUpdateResourceData(t, + map[string]string{"tpm_limit": "1000", "rpm_limit": "10", "budget_duration": "30d", "allowed_models.#": "1", "allowed_models.0": "gpt-5.2"}, + map[string]string{"allowed_models.#": "0"}, + ) + + update := runTeamMemberUpdate(t, d)[0] + for _, field := range []string{"tpm_limit", "rpm_limit", "budget_duration"} { + v, present := update[field] + if !present { + t.Fatalf("member_update payload omitted removed %s, so the proxy would keep the old value", field) + } + if v != nil { + t.Fatalf("member_update payload %s = %v, want explicit null", field, v) + } + } + if !reflect.DeepEqual(update["allowed_models"], []interface{}{}) { + t.Fatalf("member_update payload allowed_models = %v, want empty list", update["allowed_models"]) + } +} + +func TestTeamMemberAddUpdateLeavesUnchangedSettingsAlone(t *testing.T) { + d := newTeamMemberUpdateResourceData(t, + map[string]string{"budget_duration": "30d"}, + map[string]string{"budget_duration": "7d"}, + ) + + update := runTeamMemberUpdate(t, d)[0] + for _, field := range []string{"tpm_limit", "rpm_limit"} { + if v, present := update[field]; present { + t.Fatalf("member_update payload must not touch never-set %s, got %v", field, v) + } + } + if _, present := update["allowed_models"]; present { + t.Fatalf("member_update payload must not touch unchanged allowed_models, got %v", update["allowed_models"]) + } +} diff --git a/tests/local_testing/test_router_utils.py b/tests/local_testing/test_router_utils.py index 45fe42f4cd3..1b3e361bb1f 100644 --- a/tests/local_testing/test_router_utils.py +++ b/tests/local_testing/test_router_utils.py @@ -3,6 +3,7 @@ import sys, os, time import traceback, asyncio +import httpx import pytest import litellm @@ -402,6 +403,10 @@ def test_router_redis_cache(): def test_router_handle_clientside_credential(): + """A caller-supplied credential must stay scoped to the current call: it must + never be registered as a router deployment, or a later caller with no override + of their own can be load-balanced onto it and reach the provider with someone + else's credential (see LIT-7811).""" deployment = { "model_name": "gemini/*", "litellm_params": {"model": "gemini/*"}, @@ -421,7 +426,67 @@ def test_router_handle_clientside_credential(): ) assert new_deployment.litellm_params.api_key == "123" - assert len(router.get_model_list()) == 2 + assert len(router.get_model_list()) == 1 + assert router.get_deployment(model_id=new_deployment.model_info.id) is None + + +async def test_router_clientside_credential_not_reused_by_other_callers( + respx_mock, monkeypatch: pytest.MonkeyPatch +): + """End-to-end regression test for LIT-7811. + + One caller's request-scoped api_key must never leak into a later, unrelated + caller's request. Before the fix, the router registered the caller-supplied + credential as a second, permanent deployment for the shared model group, so + plain follow-up calls with no override of their own could be load-balanced + onto it and reach the provider with the first caller's key. + """ + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + route = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-1", + "object": "chat.completion", + "created": 0, + "model": "gpt-4o", + "choices": [{"index": 0, "message": {"role": "assistant", "content": "hi"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + ) + router = Router( + model_list=[ + { + "model_name": "shared-model", + "litellm_params": {"model": "openai/gpt-4o", "api_key": "configured-key"}, + "model_info": {"id": "configured-deployment"}, + } + ] + ) + + await router.acompletion( + model="shared-model", + messages=[{"role": "user", "content": "hi"}], + api_key="alternate-tenant-key", + ) + assert route.calls[-1].request.headers["authorization"] == "Bearer alternate-tenant-key" + + # The forwarded credential must never become a routable deployment for the + # model group other callers share. + assert [d["model_info"]["id"] for d in router.get_model_list(model_name="shared-model")] == [ + "configured-deployment" + ] + + for _ in range(20): + await router.acompletion( + model="shared-model", + messages=[{"role": "user", "content": "hi"}], + ) + + used_auth_headers = {call.request.headers["authorization"] for call in route.calls[1:]} + assert used_auth_headers == {"Bearer configured-key"} def test_router_get_async_openai_model_client(): diff --git a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json index 28912a27501..54d4ea85181 100644 --- a/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json +++ b/tests/logging_callback_tests/gcs_pub_sub_body/spend_logs_payload.json @@ -11,7 +11,7 @@ "user": "", "team_id": "", "organization_id": "", - "metadata": "{\"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"router_metadata\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", + "metadata": "{\"applied_guardrails\": [], \"attempted_fallbacks\": null, \"original_model_group\": null, \"batch_models\": null, \"batch_successful_requests\": null, \"batch_failed_requests\": null, \"mcp_tool_call_metadata\": null, \"vector_store_request_metadata\": null, \"routing_decision\": null, \"internal_call_origin\": null, \"router_metadata\": null, \"guardrail_information\": null, \"compression_savings\": null, \"litellm_gateway_injected_cache\": null, \"usage_object\": {\"completion_tokens\": 20, \"prompt_tokens\": 10, \"total_tokens\": 30, \"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"model_map_information\": {\"model_map_key\": \"gpt-4o\", \"model_map_value\": {\"key\": \"gpt-4o\", \"max_tokens\": 16384, \"max_input_tokens\": 128000, \"max_output_tokens\": 16384, \"input_cost_per_token\": 2.5e-06, \"cache_creation_input_token_cost\": null, \"cache_read_input_token_cost\": 1.25e-06, \"input_cost_per_character\": null, \"input_cost_per_token_above_128k_tokens\": null, \"input_cost_per_token_above_200k_tokens\": null, \"input_cost_per_query\": null, \"input_cost_per_second\": null, \"input_cost_per_audio_token\": null, \"input_cost_per_token_batches\": 1.25e-06, \"output_cost_per_token_batches\": 5e-06, \"output_cost_per_token\": 1e-05, \"output_cost_per_audio_token\": null, \"output_cost_per_character\": null, \"output_cost_per_token_above_128k_tokens\": null, \"output_cost_per_character_above_128k_tokens\": null, \"output_cost_per_token_above_200k_tokens\": null, \"output_cost_per_second\": null, \"output_cost_per_image\": null, \"output_vector_size\": null, \"litellm_provider\": \"openai\", \"mode\": \"chat\", \"supports_system_messages\": true, \"supports_response_schema\": true, \"supports_vision\": true, \"supports_function_calling\": true, \"supports_tool_choice\": true, \"supports_assistant_prefill\": false, \"supports_prompt_caching\": true, \"supports_audio_input\": false, \"supports_audio_output\": false, \"supports_pdf_input\": false, \"supports_embedding_image_input\": false, \"supports_native_streaming\": null, \"supports_web_search\": true, \"supports_reasoning\": false, \"search_context_cost_per_query\": {\"search_context_size_low\": 0.03, \"search_context_size_medium\": 0.035, \"search_context_size_high\": 0.05}, \"tpm\": null, \"rpm\": null, \"supported_openai_params\": [\"frequency_penalty\", \"logit_bias\", \"logprobs\", \"top_logprobs\", \"max_tokens\", \"max_completion_tokens\", \"modalities\", \"prediction\", \"n\", \"presence_penalty\", \"seed\", \"stop\", \"stream\", \"stream_options\", \"temperature\", \"top_p\", \"tools\", \"tool_choice\", \"function_call\", \"functions\", \"max_retries\", \"extra_headers\", \"parallel_tool_calls\", \"audio\", \"response_format\", \"user\"]}}, \"additional_usage_values\": {\"completion_tokens_details\": null, \"prompt_tokens_details\": null}, \"user_api_key\": null, \"user_api_key_alias\": null, \"user_api_key_team_id\": null, \"user_api_key_project_id\": null, \"user_api_key_project_alias\": null, \"user_api_key_org_id\": null, \"user_api_key_user_id\": null, \"user_api_key_team_alias\": null, \"spend_logs_metadata\": null, \"requester_ip_address\": null, \"user_agent\": null, \"status\": null, \"proxy_server_request\": null, \"error_information\": null, \"attempted_retries\": null, \"max_retries\": null}", "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, diff --git a/tests/pass_through_unit_tests/conftest.py b/tests/pass_through_unit_tests/conftest.py index e6e98f790e8..df8196f1785 100644 --- a/tests/pass_through_unit_tests/conftest.py +++ b/tests/pass_through_unit_tests/conftest.py @@ -1,6 +1,8 @@ +import asyncio import pytest +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from tests._vcr_conftest_common import ( # noqa: E402,F401 VerboseReporterState, @@ -45,6 +47,21 @@ def _vcr_outcome_gate(request, vcr): record_vcr_outcome(request, vcr) +@pytest.fixture(autouse=True) +async def _drain_logging_worker(): + """ + The logging queue is bound to the running loop, so anything left queued when a test's loop + goes away is carried onto the next loop and fires against that test's callbacks. + """ + GLOBAL_LOGGING_WORKER.start() + try: + await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) + except asyncio.TimeoutError: + pass + await GLOBAL_LOGGING_WORKER.stop() + yield + + def pytest_configure(config): _verbose_state.remember_pluginmanager(config) reset_vcr_diag_dir() diff --git a/tests/router_unit_tests/test_router_helper_utils.py b/tests/router_unit_tests/test_router_helper_utils.py index b18bf9351c8..14d86743557 100644 --- a/tests/router_unit_tests/test_router_helper_utils.py +++ b/tests/router_unit_tests/test_router_helper_utils.py @@ -2099,8 +2099,12 @@ def test_handle_clientside_credential_metadata_loading( assert result_deployment.model_info.id != "original-id-123" assert result_deployment.model_info.original_model_id == "original-id-123" - # Verify the deployment was added to the router - assert len(router.model_list) == len(model_list) + 1 + # The caller-supplied credential must stay scoped to this call: it must never be + # registered as a router deployment, or a later caller with no override of their + # own could be load-balanced onto it and reach the provider with this credential + # (see LIT-7811). + assert len(router.model_list) == len(model_list) + assert router.get_deployment(model_id=result_deployment.model_info.id) is None # Test that the function correctly uses the right metadata key # For acompletion, it should use "metadata" @@ -2260,14 +2264,63 @@ def test_handle_clientside_credential_with_responses_function(model_list): assert result_deployment.model_info.id != "original-id-responses" assert result_deployment.model_info.original_model_id == "original-id-responses" - # Verify the deployment was added to the router - assert len(router.model_list) == len(model_list) + 1 + # The caller-supplied credential must stay scoped to this call: it must never be + # registered as a router deployment (see LIT-7811). + assert len(router.model_list) == len(model_list) + assert router.get_deployment(model_id=result_deployment.model_info.id) is None print( "✓ Success with _ageneric_api_call_with_fallbacks function name and litellm_metadata" ) +def test_handle_clientside_credential_still_registers_custom_pricing(model_list): + """A clientside-credential call must still price against the deployment's own + custom rate, even though the call's ephemeral deployment is never added to the + router (see LIT-7811): losing that registration would silently fall back to + public catalog pricing for every clientside-credential call on a deployment + with a custom rate configured.""" + router = Router(model_list=model_list) + deployment = { + "model_name": "gpt-4.1", + "litellm_params": { + "model": "gpt-4.1", + "api_key": "test_key", + "input_cost_per_token": 0.0001234, + "output_cost_per_token": 0.0005678, + }, + "model_info": {"id": "original-id-pricing"}, + } + kwargs = {"api_key": "client_side_key", "metadata": {"model_group": "gpt-4.1"}} + + result_deployment = router._handle_clientside_credential( + deployment=deployment, kwargs=kwargs, function_name="acompletion" + ) + + registered = litellm.model_cost.get(result_deployment.model_info.id) + assert registered is not None + assert registered["input_cost_per_token"] == 0.0001234 + assert registered["output_cost_per_token"] == 0.0005678 + + +def test_register_deployment_pricing_direct_call(): + """Direct-call unit test for the pricing-registration helper `_handle_clientside_credential` + relies on, so it prices a deployment that is deliberately never added to `self.model_list`.""" + deployment = Deployment( + model_name="gpt-4.1", + litellm_params=LiteLLM_Params( + model="gpt-4.1", + api_key="test_key", + input_cost_per_token=0.0009999, + ), + model_info=ModelInfo(id="direct-call-pricing-id"), + ) + + Router._register_deployment_pricing(deployment=deployment) + + assert litellm.model_cost["direct-call-pricing-id"]["input_cost_per_token"] == 0.0009999 + + def test_get_metadata_variable_name_from_kwargs(model_list): """ Test _get_metadata_variable_name_from_kwargs method returns correct metadata variable name based on kwargs content. diff --git a/tests/test_litellm/a2a_protocol/test_a2a_streaming_iterator.py b/tests/test_litellm/a2a_protocol/test_a2a_streaming_iterator.py index d86cbb94a91..2603d135dce 100644 --- a/tests/test_litellm/a2a_protocol/test_a2a_streaming_iterator.py +++ b/tests/test_litellm/a2a_protocol/test_a2a_streaming_iterator.py @@ -100,3 +100,57 @@ async def test_custom_logger_only_never_submits_sync_success_handler(monkeypatch assert recorder.async_hook_fired is True assert recording_executor.submitted_for(logging_obj) == [] + + +class _AgentChunk: + def __init__(self, text: str): + self._text = text + + def model_dump(self, mode: str, exclude_none: bool) -> dict: + return {"result": {"kind": "message", "role": "agent", "parts": [{"kind": "text", "text": self._text}]}} + + +@pytest.mark.asyncio +async def test_stream_completion_counts_tokens_off_the_event_loop(monkeypatch): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("gpt-5.6-luna") + monkeypatch.setattr(litellm, "success_callback", []) + monkeypatch.setattr(litellm, "_async_success_callback", []) + logging_obj = LitellmLogging( + model="a2a/test-agent", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="a2a_send_message_streaming", + start_time=time.time(), + litellm_call_id="lit-7190-test", + function_id="lit-7190-test", + ) + + async def _stream(): + yield _AgentChunk(text * 100) + + iterator = A2AStreamingIterator( + stream=_stream(), + request=SimpleNamespace( + params=SimpleNamespace(message={"role": "user", "parts": [{"kind": "text", "text": text * 100}]}) + ), + logging_obj=logging_obj, + agent_name="test-agent", + ) + + async def drain() -> int: + return len([chunk async for chunk in iterator]) + + yielded, took, lags = await timed_with_loop_lags(drain) + + assert yielded == 1 + usage = logging_obj.model_call_details["usage"] + assert usage.prompt_tokens > 100_000 + assert usage.completion_tokens > 100_000 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/a2a_protocol/test_main.py b/tests/test_litellm/a2a_protocol/test_main.py index 8850a2eca6c..318b40138ed 100644 --- a/tests/test_litellm/a2a_protocol/test_main.py +++ b/tests/test_litellm/a2a_protocol/test_main.py @@ -1,5 +1,7 @@ """Tests for litellm/a2a_protocol/main.py non-streaming send behavior.""" +import asyncio + import httpx import pytest @@ -13,7 +15,8 @@ from a2a.compat.v0_3.types import ( ) import litellm -from litellm.a2a_protocol.main import _send_message, _stream_messages, create_a2a_client +from litellm.integrations.custom_logger import CustomLogger +from litellm.a2a_protocol.main import _send_message, _stream_messages, asend_message, create_a2a_client from litellm.caching.llm_caching_handler import LLMClientCache from litellm.constants import DEFAULT_A2A_AGENT_TIMEOUT from litellm.llms.custom_httpx.http_handler import ( @@ -413,3 +416,51 @@ async def test_the_pooled_a2a_client_arrives_with_cookie_persistence_disabled(is assert dict(handler.client.cookies) == {}, "the pooled A2A client kept an upstream's cookie" await handler.close() + + +class _UsageRecorder(CustomLogger): + def __init__(self): + super().__init__() + self.logged = asyncio.Event() + self.payload = None + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self.payload = kwargs["standard_logging_object"] + self.logged.set() + + +@pytest.mark.asyncio +async def test_asend_message_counts_usage_off_the_event_loop(monkeypatch): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer("gpt-5.6-luna") + recorder = _UsageRecorder() + monkeypatch.setattr(litellm, "callbacks", [recorder]) + monkeypatch.setattr(litellm, "success_callback", [recorder]) + monkeypatch.setattr(litellm, "_async_success_callback", [recorder]) + + reply = _conv.pb2_v10.StreamResponse() + reply.message.message_id = "reply-1" + reply.message.role = _conv.pb2_v10.Role.ROLE_AGENT + reply.message.parts.add().text = text * 100 + request = SendMessageRequest( + id="r1", + params=MessageSendParams( + message={"messageId": "m1", "role": "user", "parts": [{"kind": "text", "text": text * 100}]} + ), + ) + + response, took, lags = await timed_with_loop_lags( + lambda: asend_message(a2a_client=_FakeClient(reply), request=request) + ) + + assert response.id == "r1" + await asyncio.wait_for(recorder.logged.wait(), timeout=10) + assert recorder.payload["prompt_tokens"] > 100_000 + assert recorder.payload["completion_tokens"] > 100_000 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/caching/test_qdrant_semantic_cache.py b/tests/test_litellm/caching/test_qdrant_semantic_cache.py index e07578dd7e5..a0a9b71787c 100644 --- a/tests/test_litellm/caching/test_qdrant_semantic_cache.py +++ b/tests/test_litellm/caching/test_qdrant_semantic_cache.py @@ -1026,3 +1026,36 @@ def test_qdrant_semantic_cache_defaults_embedding_timeout(): cache = QdrantSemanticCache.__new__(QdrantSemanticCache) assert cache.embedding_timeout == SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS assert SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS < 60 + + +@pytest.mark.asyncio +async def test_qdrant_async_embedding_truncates_off_the_event_loop(monkeypatch): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + from litellm.caching.qdrant_semantic_cache import QdrantSemanticCache + + warm_tokenizer("sem-embed") + cache = QdrantSemanticCache.__new__(QdrantSemanticCache) + cache.embedding_model = "sem-embed" + cache.embedding_max_input_tokens = 5 + cache.embedding_timeout = 5 + + router = MagicMock() + router.get_configured_token_limits.return_value = (8191, None) + router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]}) + monkeypatch.setitem( + sys.modules, + "litellm.proxy.proxy_server", + _router_proxy_module(router, "sem-embed"), + ) + + response, took, lags = await timed_with_loop_lags(lambda: cache._get_async_embedding(text * 100)) + + assert response["data"][0]["embedding"] == [0.1, 0.2] + assert _token_count("sem-embed", router.aembedding.call_args.kwargs["input"]) == 5 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/caching/test_redis_semantic_cache.py b/tests/test_litellm/caching/test_redis_semantic_cache.py index 9884e9d9bc0..de253b4f10b 100644 --- a/tests/test_litellm/caching/test_redis_semantic_cache.py +++ b/tests/test_litellm/caching/test_redis_semantic_cache.py @@ -1387,3 +1387,32 @@ def test_redis_semantic_cache_defaults_embedding_timeout(): cache = RedisSemanticCache.__new__(RedisSemanticCache) assert cache.embedding_timeout == SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS assert SEMANTIC_CACHE_EMBEDDING_TIMEOUT_SECONDS < 60 + + +@pytest.mark.asyncio +async def test_redis_async_embedding_truncates_off_the_event_loop(monkeypatch): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + from litellm.caching.redis_semantic_cache import RedisSemanticCache + + warm_tokenizer("sem-embed") + cache = RedisSemanticCache.__new__(RedisSemanticCache) + cache.embedding_model = "sem-embed" + cache.embedding_max_input_tokens = 5 + cache.embedding_timeout = 5 + + router = MagicMock() + router.get_configured_token_limits.return_value = (8191, None) + router.aembedding = AsyncMock(return_value={"data": [{"embedding": [0.1, 0.2]}]}) + _proxy_with_router(monkeypatch, router, "sem-embed") + + embedding, took, lags = await timed_with_loop_lags(lambda: cache._get_async_embedding(text * 100)) + + assert embedding == [0.1, 0.2] + assert _token_count("sem-embed", router.aembedding.call_args.kwargs["input"]) == 5 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py b/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py index bc4dccc7d70..e66cd654f93 100644 --- a/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py +++ b/tests/test_litellm/integrations/compression_interception/test_compression_interception_handler.py @@ -523,3 +523,28 @@ async def test_pre_call_hook_no_compression_records_no_savings(monkeypatch): await logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.anthropic_messages) assert "compression_savings" not in litellm_metadata + + +@pytest.mark.asyncio +async def test_pre_call_hook_counts_tokens_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + model = "anthropic/claude-fable-5" + warm_tokenizer(model) + logger = CompressionInterceptionLogger(compression_trigger=10_000_000) + messages = [{"role": "user", "content": text * 100}] + kwargs = {"model": model, "messages": messages} + + result, took, lags = await timed_with_loop_lags( + lambda: logger.async_pre_call_deployment_hook(kwargs=kwargs, call_type=CallTypes.anthropic_messages) + ) + + assert result is not None + assert result["messages"] is messages + assert "tools" not in result + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/integrations/test_langsmith_init.py b/tests/test_litellm/integrations/test_langsmith_init.py index 0bc9e279fbf..f56d2310e73 100644 --- a/tests/test_litellm/integrations/test_langsmith_init.py +++ b/tests/test_litellm/integrations/test_langsmith_init.py @@ -1,5 +1,6 @@ import asyncio import os +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -7,6 +8,7 @@ import pytest import litellm from litellm.integrations.langsmith import LangsmithLogger +from litellm.types.integrations.langsmith import LangsmithQueueObject @pytest.fixture @@ -531,3 +533,44 @@ class TestLangsmithRootRunIdConsistency: assert data["trace_id"] == "trace-1" assert data["dotted_order"] == dotted + + +@pytest.mark.asyncio +async def test_events_appended_during_flush_are_not_dropped(): + logger = LangsmithLogger(langsmith_api_key="test-key", langsmith_project="test-project") + try: + sent_batches: Final[list[list[dict[str, str]]]] = [] + late_event: Final = LangsmithQueueObject( + credentials=logger.default_credentials, data={"id": "late"} + ) + + async def fake_post( + url: str, json: dict[str, list[dict[str, str]]], headers: dict[str, str] + ) -> MagicMock: + if not sent_batches: + logger.log_queue.append(late_event) + sent_batches.append(json["post"]) + response = MagicMock() + response.status_code = 200 + response.raise_for_status = MagicMock() + return response + + logger.async_httpx_client = MagicMock(post=AsyncMock(side_effect=fake_post)) + logger.log_queue = [ + LangsmithQueueObject(credentials=logger.default_credentials, data={"id": "a"}), + LangsmithQueueObject(credentials=logger.default_credentials, data={"id": "b"}), + ] + + await logger.flush_queue() + + assert [e["id"] for e in sent_batches[0]] == ["a", "b"] + assert logger.log_queue == [late_event] + + await logger.flush_queue() + + assert [e["id"] for e in sent_batches[1]] == ["late"] + assert logger.log_queue == [] + finally: + if logger._flush_task is not None: + logger._flush_task.cancel() + await asyncio.gather(logger._flush_task, return_exceptions=True) diff --git a/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py b/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py index 278a4ef1df6..9e8348e860a 100644 --- a/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py +++ b/tests/test_litellm/integrations/test_prometheus_invalid_key_filtering.py @@ -1,18 +1,18 @@ """ Unit tests for Prometheus invalid API key request filtering. -Tests functionality that prevents invalid API key requests (401 status codes) -from being recorded in Prometheus metrics. +Tests the 401 detection helpers, that LLM-level metrics skip invalid API key +requests, and that the proxy-level failed request counter still records them. """ from unittest.mock import Mock, patch import pytest +from fastapi import HTTPException from prometheus_client import REGISTRY - from litellm.integrations.prometheus import PrometheusLogger -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import ProxyErrorTypes, ProxyException, UserAPIKeyAuth @pytest.fixture(scope="function") @@ -129,28 +129,29 @@ class TestSkipMetricsValidation: class TestAsyncHooks: - """Test async hook methods skip metrics for invalid API keys.""" - - @pytest.fixture - def mock_user_api_key(self): - """Create a mock UserAPIKeyAuth object.""" - user_key = Mock(spec=UserAPIKeyAuth) - user_key.api_key = "test-key" - user_key.end_user_id = None - user_key.user_id = None - user_key.user_email = None - user_key.key_alias = None - user_key.team_id = None - user_key.team_alias = None - user_key.request_route = "/test" - return user_key + """Test how async hook methods treat invalid API key requests.""" @pytest.mark.asyncio - async def test_post_call_failure_hook_skips_401( - self, prometheus_logger, mock_user_api_key + @pytest.mark.parametrize( + "exception", + [ + HTTPException( + status_code=401, + detail="LiteLLM Virtual Key expected. Received=nota****tall, expected to start with 'sk-'.", + ), + ProxyException( + message="Authentication Error, Invalid proxy server token passed.", + type=ProxyErrorTypes.token_not_found_in_db, + param="key", + code=401, + ), + ], + ) + async def test_post_call_failure_hook_counts_401_without_key_hash( + self, prometheus_logger, exception ): - exception = ExceptionWithCode("401") - exception.__class__.__name__ = "ProxyException" + unauthenticated = UserAPIKeyAuth(request_route="/v1/chat/completions") + unauthenticated.api_key = "notakeyatall" with ( patch.object( @@ -160,15 +161,50 @@ class TestAsyncHooks: prometheus_logger, "litellm_proxy_total_requests_metric" ) as mock_total, ): - await prometheus_logger.async_post_call_failure_hook( request_data={"model": "test-model"}, original_exception=exception, - user_api_key_dict=mock_user_api_key, + user_api_key_dict=unauthenticated, ) - mock_failed.labels.assert_not_called() - mock_total.labels.assert_not_called() + failed_labels = mock_failed.labels.call_args.kwargs + assert failed_labels["exception_status"] == "401" + assert failed_labels["hashed_api_key"] is None + assert failed_labels["route"] == "/v1/chat/completions" + mock_failed.labels.return_value.inc.assert_called_once() + assert mock_total.labels.call_args.kwargs["status_code"] == "401" + mock_total.labels.return_value.inc.assert_called_once() + + @pytest.mark.asyncio + async def test_post_call_failure_hook_keeps_resolved_identity_labels_for_401( + self, prometheus_logger + ): + expired_key = UserAPIKeyAuth( + api_key="sk-expired", + key_alias="expired-alias", + team_id="team-1", + ) + exception = ProxyException( + message="Authentication Error - Expired Key.", + type=ProxyErrorTypes.expired_key, + param="key", + code=401, + ) + + with patch.object( + prometheus_logger, "litellm_proxy_failed_requests_metric" + ) as mock_failed: + await prometheus_logger.async_post_call_failure_hook( + request_data={"model": "test-model"}, + original_exception=exception, + user_api_key_dict=expired_key, + ) + + failed_labels = mock_failed.labels.call_args.kwargs + assert failed_labels["exception_status"] == "401" + assert failed_labels["hashed_api_key"] is None + assert failed_labels["api_key_alias"] == "expired-alias" + assert failed_labels["team"] == "team-1" @pytest.mark.asyncio async def test_log_failure_event_skips_401(self, prometheus_logger): diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py index 30b158e3b2c..a315b7003ad 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_llm_cost_calc_utils.py @@ -2321,36 +2321,6 @@ def test_gpt55_dated_variants_match_base_reasoning_effort_capabilities(_local_mo ) -@pytest.mark.parametrize( - "model,expected_mode,expected_input,expected_output,expected_cache_read", - [ - ("azure/gpt-5.5", "chat", 5e-6, 3e-5, 5e-7), - ("azure/gpt-5.5-2026-04-23", "chat", 5e-6, 3e-5, 5e-7), - ("azure/gpt-5.5-pro", "responses", 3e-5, 1.8e-4, 3e-6), - ("azure/gpt-5.5-pro-2026-04-23", "responses", 3e-5, 1.8e-4, 3e-6), - ], -) -def test_azure_gpt55_entries_present_with_correct_pricing(_local_model_cost_map, - model, expected_mode, expected_input, expected_output, expected_cache_read -): - """Day-0 Azure entries for GPT-5.5 mirror the OpenAI pricing structure. - - Pricing parity with openai/gpt-5.5* (verified against OpenAI's pricing page - on 2026-04-24): $5/$30 input/output per 1M for chat, $30/$180 for pro. - Cache discount is 10% of input. - """ - - m = litellm.model_cost[model] - assert m["litellm_provider"] == "azure" - assert m["mode"] == expected_mode - assert m["input_cost_per_token"] == expected_input - assert m["output_cost_per_token"] == expected_output - assert m["cache_read_input_token_cost"] == expected_cache_read - # Long-context window inherited from gpt-5.4 / openai gpt-5.5. - assert m["max_input_tokens"] == 1050000 - assert m["max_output_tokens"] == 128000 - - @pytest.mark.parametrize( "model,expected_none,expected_minimal,expected_xhigh", [ @@ -3414,8 +3384,6 @@ def test_query_count_is_free_without_a_per_query_price(_local_model_cost_map): # --------------------------------------------------------------------------- - - @pytest.mark.parametrize("model", ["gpt-5.4", "gpt-realtime-2.1", "gpt-realtime-2.1-mini"]) @pytest.mark.parametrize("data_residency", ["eu", "us"]) def test_data_residency_applies_uplift(data_residency, model, _local_model_cost_map): @@ -4556,20 +4524,6 @@ GEMINI_DAY0_LAUNCH_PRICING = [ ] -@pytest.mark.parametrize("model,input_cost,output_cost,cache_read_cost", GEMINI_DAY0_LAUNCH_PRICING) -def test_gemini_36_flash_and_35_flash_lite_launch_pricing(_local_model_cost_map, model, input_cost, output_cost, cache_read_cost): - - model_cost_map = litellm.model_cost[model] - assert model_cost_map["input_cost_per_token"] == input_cost - assert model_cost_map["output_cost_per_token"] == output_cost - assert model_cost_map["output_cost_per_reasoning_token"] == output_cost - assert model_cost_map["cache_read_input_token_cost"] == cache_read_cost - assert model_cost_map["mode"] == "chat" - assert model_cost_map["supports_reasoning"] is True - assert model_cost_map["supports_function_calling"] is True - assert model_cost_map["max_input_tokens"] == 1048576 - - def test_generic_cost_per_token_gemini_36_flash(_local_model_cost_map): usage = Usage( @@ -4598,44 +4552,6 @@ GEMINI_36_FLASH_SERVICE_TIER_PRICING = [ ] -@pytest.mark.parametrize( - "service_tier,input_rate,output_rate,cache_read_rate", GEMINI_36_FLASH_SERVICE_TIER_PRICING -) -@pytest.mark.parametrize( - "model", ["gemini-3.6-flash", "gemini/gemini-3.6-flash", "vertex_ai/gemini-3.6-flash"] -) -def test_gemini_36_flash_service_tier_introductory_pricing( - model, service_tier, input_rate, output_rate, cache_read_rate, _local_model_cost_map -): - """Regression: every 3.6 Flash tier is on Google's introductory rates through 2026-12-31, - so flex and priority requests must not be billed at the post-introductory rates.""" - usage = Usage( - prompt_tokens=1_000, - completion_tokens=500, - total_tokens=1_500, - prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=200, text_tokens=800), - ) - - prompt_cost, completion_cost = generic_cost_per_token( - model=model.split("/")[-1], - usage=usage, - custom_llm_provider=model.split("/")[0] if "/" in model else "gemini", - service_tier=service_tier, - ) - - assert prompt_cost == pytest.approx(800 * input_rate + 200 * cache_read_rate, rel=1e-9) - assert completion_cost == pytest.approx(500 * output_rate, rel=1e-9) - - -@pytest.mark.parametrize( - "model", ["gemini-3.6-flash", "gemini/gemini-3.6-flash", "vertex_ai/gemini-3.6-flash"] -) -def test_gemini_36_flash_batch_introductory_pricing(model, _local_model_cost_map): - model_cost_map = litellm.model_cost[model] - assert model_cost_map["input_cost_per_token_batches"] == 3.75e-07 - assert model_cost_map["output_cost_per_token_batches"] == 1.875e-06 - - def test_generic_cost_per_token_gemini_35_flash_lite(_local_model_cost_map): usage = Usage( @@ -4667,43 +4583,6 @@ GEMINI_35_FLASH_LITE_TIER_RATES_BY_SURFACE = [ ] -@pytest.mark.parametrize( - "custom_llm_provider,service_tier,input_rate,output_rate,cache_read_rate", - GEMINI_35_FLASH_LITE_TIER_RATES_BY_SURFACE, -) -def test_gemini_35_flash_lite_service_tier_pricing( - custom_llm_provider, service_tier, input_rate, output_rate, cache_read_rate, _local_model_cost_map -): - """Regression: Vertex publishes flash-lite flex context caching at $0.015/M while the - Gemini API publishes $0.02/M, so vertex_ai flex cache reads must bill 1.5e-08/token - instead of the 2e-08 the map used to carry, without disturbing the Gemini API rate.""" - usage = Usage( - prompt_tokens=1_000, - completion_tokens=500, - total_tokens=1_500, - prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=200, text_tokens=800), - ) - - prompt_cost, completion_cost = generic_cost_per_token( - model="gemini-3.5-flash-lite", - usage=usage, - custom_llm_provider=custom_llm_provider, - service_tier=service_tier, - ) - - assert prompt_cost == pytest.approx(800 * input_rate + 200 * cache_read_rate, rel=1e-9) - assert completion_cost == pytest.approx(500 * output_rate, rel=1e-9) - - -def test_gemini_35_flash_lite_flex_cache_read_map_entries(_local_model_cost_map): - """Each map entry carries its own surface's published flex cache-read rate: the bare - and vertex_ai keys are the Vertex surface at $0.015/M, the gemini key is the Gemini - API surface at $0.02/M.""" - assert litellm.model_cost["gemini-3.5-flash-lite"]["cache_read_input_token_cost_flex"] == 1.5e-08 - assert litellm.model_cost["vertex_ai/gemini-3.5-flash-lite"]["cache_read_input_token_cost_flex"] == 1.5e-08 - assert litellm.model_cost["gemini/gemini-3.5-flash-lite"]["cache_read_input_token_cost_flex"] == 2e-08 - - @pytest.mark.parametrize( "service_tier,input_rate,cache_read_rate,cache_write_rate,output_rate", [ @@ -4932,19 +4811,6 @@ GEMINI_37_FLASH_LAUNCH_PRICING = [ ] -@pytest.mark.parametrize("model,input_cost,output_cost,cache_read_cost", GEMINI_37_FLASH_LAUNCH_PRICING) -def test_gemini_37_flash_launch_pricing(model, input_cost, output_cost, cache_read_cost, _local_model_cost_map): - model_cost_map = litellm.model_cost[model] - assert model_cost_map["input_cost_per_token"] == input_cost - assert model_cost_map["output_cost_per_token"] == output_cost - assert model_cost_map["output_cost_per_reasoning_token"] == output_cost - assert model_cost_map["cache_read_input_token_cost"] == cache_read_cost - assert model_cost_map["mode"] == "chat" - assert model_cost_map["supports_reasoning"] is True - assert model_cost_map["supports_function_calling"] is True - assert model_cost_map["max_input_tokens"] == 1048576 - - def test_generic_cost_per_token_gemini_37_flash(_local_model_cost_map): usage = Usage( prompt_tokens=1000, @@ -4972,19 +4838,6 @@ GEMINI_38_FLASH_LAUNCH_PRICING = [ ] -@pytest.mark.parametrize("model,input_cost,output_cost,cache_read_cost", GEMINI_38_FLASH_LAUNCH_PRICING) -def test_gemini_38_flash_launch_pricing(model, input_cost, output_cost, cache_read_cost, _local_model_cost_map): - model_cost_map = litellm.model_cost[model] - assert model_cost_map["input_cost_per_token"] == input_cost - assert model_cost_map["output_cost_per_token"] == output_cost - assert model_cost_map["output_cost_per_reasoning_token"] == output_cost - assert model_cost_map["cache_read_input_token_cost"] == cache_read_cost - assert model_cost_map["mode"] == "chat" - assert model_cost_map["supports_reasoning"] is True - assert model_cost_map["supports_function_calling"] is True - assert model_cost_map["max_input_tokens"] == 1048576 - - GEMINI_38_FLASH_FIELDS_SHARED_WITH_37_FLASH = ( "input_cost_per_token", "output_cost_per_token", @@ -5045,20 +4898,6 @@ def test_generic_cost_per_token_gemini_38_flash(_local_model_cost_map): assert completion_cost == pytest.approx(0.001875) -def test_grok_46_launch_pricing(_local_model_cost_map): - model_cost_map = litellm.model_cost["xai/grok-4.6"] - assert model_cost_map["input_cost_per_token"] == 2e-06 - assert model_cost_map["output_cost_per_token"] == 6e-06 - assert model_cost_map["cache_read_input_token_cost"] == 5e-07 - assert model_cost_map["input_cost_per_token_above_200k_tokens"] == 4e-06 - assert model_cost_map["output_cost_per_token_above_200k_tokens"] == 1.2e-05 - assert model_cost_map["cache_read_input_token_cost_above_200k_tokens"] == 1e-06 - assert model_cost_map["mode"] == "chat" - assert model_cost_map["supports_reasoning"] is True - assert model_cost_map["supports_function_calling"] is True - assert model_cost_map["max_input_tokens"] == 500000 - - def test_generic_cost_per_token_grok_46(_local_model_cost_map): usage = Usage( prompt_tokens=1_000, diff --git a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index bbb7b5f9c35..37b985897da 100644 --- a/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/test_litellm/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -1,6 +1,4 @@ -import json from collections.abc import Mapping, Sequence -from pathlib import Path import pytest @@ -892,29 +890,6 @@ def test_gpt_4o_mini_snapshot_bills_web_search_like_its_alias( assert snapshot_cost == alias_cost == 0.025 -def test_gpt_4o_mini_web_search_price_matches_in_both_cost_maps(): - repo_root = Path(__file__).parents[4] - cost_maps = tuple( - json.loads((repo_root / path).read_text(encoding="utf-8")) - for path in ( - "model_prices_and_context_window.json", - "litellm/model_prices_and_context_window_backup.json", - ) - ) - canonical, backup = cost_maps - expected_search_price = { - "search_context_size_low": 0.025, - "search_context_size_medium": 0.025, - "search_context_size_high": 0.025, - } - for model_name in ("gpt-4o-mini", "gpt-4o-mini-2024-07-18"): - canonical_entry = canonical[model_name] - backup_entry = backup[model_name] - assert canonical_entry["search_context_cost_per_query"] == expected_search_price - assert backup_entry["search_context_cost_per_query"] == expected_search_price - assert canonical_entry == backup_entry - - # Note: File search integration test removed due to complex annotation detection logic # The unit tests in test_azure_assistant_cost_tracking.py provide comprehensive coverage diff --git a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py index b2cc3ebe4c6..057fa228562 100644 --- a/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py +++ b/tests/test_litellm/litellm_core_utils/test_fallback_generalizations.py @@ -627,17 +627,6 @@ def test_shipped_adaptive_rule_requires_claude_prefix(shipped_cost_map): litellm.get_model_info(model) -def test_shipped_exact_entry_beats_rules(shipped_cost_map): - model = "us.anthropic.claude-sonnet-4-6" - assert model in litellm.model_cost - info = litellm.get_model_info(model, custom_llm_provider="bedrock") - assert info["litellm_provider"] == "bedrock_converse" - assert info["input_cost_per_token"] == 3.3e-06 - assert info["max_input_tokens"] == 1000000 - assert info["supports_adaptive_thinking"] is True - assert info.get("supports_mid_conversation_system") is None - - def test_shipped_rules_lose_to_exact_entries_across_cost_ladder_variants(shipped_cost_map): """A route-mangled variant of an exactly-mapped model must never resolve from rules. The cost calculator tries model-name variants in order; a rule-derived diff --git a/tests/test_litellm/litellm_core_utils/test_get_llm_provider_logic.py b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_logic.py new file mode 100644 index 00000000000..1ecef9ffff7 --- /dev/null +++ b/tests/test_litellm/litellm_core_utils/test_get_llm_provider_logic.py @@ -0,0 +1,55 @@ +from typing import Final + +import pytest + +import litellm +from litellm import CustomLLM +from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_llm_provider, + is_registered_custom_provider, +) + +CUSTOM_PROVIDER: Final = "test-onprem-llm" + + +@pytest.fixture +def registered_custom_provider(monkeypatch: pytest.MonkeyPatch) -> str: + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": CUSTOM_PROVIDER, "custom_handler": CustomLLM()}]) + monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list)) + monkeypatch.setattr(litellm, "_custom_providers", list(litellm._custom_providers)) + return CUSTOM_PROVIDER + + +def test_get_llm_provider_resolves_custom_provider_map_prefix_before_first_completion( + registered_custom_provider: str, +) -> None: + assert registered_custom_provider not in litellm.provider_list + + model, provider, dynamic_api_key, api_base = get_llm_provider(model=f"{registered_custom_provider}/my-model") + + assert (model, provider, dynamic_api_key, api_base) == ("my-model", registered_custom_provider, None, None) + + +def test_get_llm_provider_strips_prefix_when_custom_provider_passed_explicitly( + registered_custom_provider: str, +) -> None: + model, provider, _, api_base = get_llm_provider( + model="my-model", + custom_llm_provider=registered_custom_provider, + api_base="http://onprem.internal:8080", + ) + + assert (model, provider, api_base) == ("my-model", registered_custom_provider, "http://onprem.internal:8080") + + +def test_get_llm_provider_still_rejects_unregistered_prefix(registered_custom_provider: str) -> None: + with pytest.raises(litellm.BadRequestError, match="LLM Provider NOT provided"): + get_llm_provider(model="not-registered-llm/my-model") + + +@pytest.mark.parametrize( + ("candidate", "expected"), + [(CUSTOM_PROVIDER, True), ("not-registered-llm", False), (None, False), ("", False)], +) +def test_is_registered_custom_provider(registered_custom_provider: str, candidate: str | None, expected: bool) -> None: + assert is_registered_custom_provider(candidate) is expected diff --git a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py index c509c8399c9..53fee36b3a8 100644 --- a/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py +++ b/tests/test_litellm/litellm_core_utils/test_get_model_cost_map.py @@ -225,36 +225,6 @@ def test_shipped_backup_marks_claude_4_6_plus_adaptive_not_4_0(): assert "supports_adaptive_thinking" not in backup[non_adaptive], non_adaptive -@pytest.mark.parametrize( - "cost_map", - [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], - ids=["root", "bundled_backup"], -) -def test_azure_ai_claude_1m_context_entries(cost_map: dict): - """Microsoft Foundry serves a 1M-token context window for Opus 4.6+ and Sonnet - 4.6+, so the ``azure_ai`` entries must not advertise the 200k cap that made - context-aware clients compact prompts early (LIT-4406). Both the root map (used - by default network loading) and the bundled fallback are checked so the two can - never drift apart.""" - for model in [ - "azure_ai/claude-opus-4-6", - "azure_ai/claude-opus-4-7", - "azure_ai/claude-opus-4-8", - "azure_ai/claude-opus-5", - "azure_ai/claude-sonnet-5", - "azure_ai/claude-sonnet-4-6", - ]: - assert cost_map[model]["max_input_tokens"] == 1000000, model - - for model in [ - "azure_ai/claude-opus-4-1", - "azure_ai/claude-opus-4-5", - "azure_ai/claude-sonnet-4-5", - "azure_ai/claude-haiku-4-5", - ]: - assert cost_map[model]["max_input_tokens"] == 200000, model - - # OpenRouter headline rates from GET https://openrouter.ai/api/v1/models. # These were the catalog values that disagreed with that API (and, for the # two spotlight models, the public model pages that their source fields cite). @@ -278,34 +248,6 @@ _OPENROUTER_STALE_COSTS = { } -@pytest.mark.parametrize( - "cost_map", - [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], - ids=["root", "bundled_backup"], -) -def test_openrouter_catalog_costs_match_live_headline_rates(cost_map: dict): - """openrouter/* spend tracking reads these catalog fields. The values must - stay aligned with OpenRouter's published headline rate, not the stale - figures that over/under-counted by up to 30x. Both maps are checked so - the root file and bundled backup cannot drift apart.""" - control = cost_map["openrouter/anthropic/claude-opus-5"] - assert control["input_cost_per_token"] == 5e-06 - assert control["output_cost_per_token"] == 2.5e-05 - assert control["cache_read_input_token_cost"] == 5e-07 - - for model, (inp, out, cache) in _OPENROUTER_LIVE_COSTS.items(): - entry = cost_map[model] - assert entry["input_cost_per_token"] == inp, model - assert entry["output_cost_per_token"] == out, model - if cache is not None: - assert entry["cache_read_input_token_cost"] == cache, model - - for model, (stale_in, stale_out) in _OPENROUTER_STALE_COSTS.items(): - entry = cost_map[model] - assert entry["input_cost_per_token"] != stale_in, model - assert entry["output_cost_per_token"] != stale_out, model - - def test_get_model_cost_map_stamps_loaded_at(): """The load time feeds each pod's reload-due decision; a load that does not stamp it would make manual reload requests race the proxy's startup""" diff --git a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py index 784468c839b..47efbe7f19a 100644 --- a/tests/test_litellm/litellm_core_utils/test_streaming_handler.py +++ b/tests/test_litellm/litellm_core_utils/test_streaming_handler.py @@ -4927,3 +4927,50 @@ class TestStableStreamingResponseId: ) wrapper.response_id = "chatcmpl-from-provider" assert wrapper.model_response_creator().id == "chatcmpl-from-provider" + + +@pytest.mark.asyncio +async def test_async_stream_without_usage_counts_tokens_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + model = "gpt-5.6-luna" + warm_tokenizer(model) + messages = [{"role": "user", "content": text * 100}] + content_chunks = [_make_chunk(text) for _ in range(100)] + stop_chunk = ModelResponseStream( + id="test", + created=1741037890, + model=model, + choices=[StreamingChoices(index=0, delta=Delta(content=""), finish_reason="stop")], + ) + logging_obj = Logging( + model=model, + messages=messages, + stream=True, + call_type="acompletion", + start_time=time.time(), + litellm_call_id="12345", + function_id="1245", + ) + wrapper = CustomStreamWrapper( + completion_stream=ModelResponseListIterator(model_responses=content_chunks + [stop_chunk]), + model=model, + custom_llm_provider="openai", + logging_obj=logging_obj, + stream_options={"include_usage": True}, + ) + + async def consume() -> list[ModelResponseStream]: + return [chunk async for chunk in wrapper] + + chunks, took, lags = await timed_with_loop_lags(consume) + + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks) == text * 100 + assert chunks[-1].usage.prompt_tokens > 100_000 + assert chunks[-1].usage.completion_tokens > 100_000 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py index 28c82fdf528..7660a8649b5 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_compact.py @@ -2605,3 +2605,34 @@ def test_build_summary_messages_keeps_midturn_system_correction_in_place(): assert summary_messages[0]["content"] == "caller system prompt" assert summary_messages[2]["content"] == "use the corrected result" assert summary_messages[-1]["content"] == "summarize the conversation" + + +async def test_threshold_check_counts_tokens_off_the_event_loop(monkeypatch): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + from litellm.llms.anthropic.experimental_pass_through.context_management.constants import ( + COMPACT_SUMMARY_MODEL_SETTING_KEY, + ) + from litellm.proxy.proxy_server import general_settings + + monkeypatch.setitem(general_settings, COMPACT_SUMMARY_MODEL_SETTING_KEY, "claude-haiku-4-5") + warm_tokenizer(MODEL) + messages = [{"role": "user", "content": text * 100}, *_simple_messages()] + result, took, lags = await timed_with_loop_lags( + lambda: apply_compact_20260112( + model=MODEL, + messages=messages, + tools=None, + system=None, + edit_spec={"type": "compact_20260112", "trigger": {"type": "input_tokens", "value": 10_000_000}}, + ) + ) + + assert result.messages == messages + assert result.compaction_block is None + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py index 50c72cfe8d0..a21c22cf5fa 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/context_management/test_dispatcher.py @@ -129,3 +129,35 @@ async def test_malformed_edit_entries_are_skipped(): ) assert result.applied_edits == [] assert result.messages == messages + + +async def test_sync_editor_counts_tokens_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + warm_tokenizer(MODEL) + messages = [{"role": "user", "content": text * 100}, *_history_with_two_tool_pairs()] + + result, took, lags = await timed_with_loop_lags( + lambda: apply_context_management( + model=MODEL, + messages=messages, + tools=None, + system=None, + context_management_spec={ + "edits": [ + { + "type": "clear_tool_uses_20250919", + "trigger": {"type": "input_tokens", "value": 10_000_000}, + } + ] + }, + ) + ) + + assert result.messages == messages + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py b/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py index 84d5cd2a7d4..9b20192c3f2 100644 --- a/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py +++ b/tests/test_litellm/llms/azure_ai/test_azure_ai_foundry_catalog_model_metadata.py @@ -12,7 +12,6 @@ REPO_ROOT: Final = Path(__file__).parents[4] MAIN_COST_MAP: Final = REPO_ROOT / "model_prices_and_context_window.json" BACKUP_COST_MAP: Final = REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json" COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) -AZURE_PRICING_PREFIX: Final = "https://azure.microsoft.com/en-us/pricing/details/" A_MILLION: Final = 1_000_000 AN_HOUR_IN_SECONDS: Final = 3600 @@ -76,7 +75,9 @@ def test_azure_ai_catalog_name_prices_the_same_in_any_casing(catalog_name: str) @pytest.mark.usefixtures("local_model_cost_map") @pytest.mark.parametrize("catalog_name", GROK_4_20_NAMES) def test_azure_ai_grok_4_20_bills_cached_prompt_tokens_at_the_input_price(catalog_name: str) -> None: - uncached_prompt_cost, _ = cost_per_token(model=f"azure_ai/{catalog_name}", prompt_tokens=A_MILLION, completion_tokens=0) + uncached_prompt_cost, _ = cost_per_token( + model=f"azure_ai/{catalog_name}", prompt_tokens=A_MILLION, completion_tokens=0 + ) cached_prompt_cost, _ = cost_per_token( model=f"azure_ai/{catalog_name}", prompt_tokens=A_MILLION, @@ -100,7 +101,6 @@ def test_azure_ai_catalog_entry_source_and_backup_match(catalog_name: str) -> No main_entry = _cost_map_entry(MAIN_COST_MAP, catalog_name) backup_entry = _cost_map_entry(BACKUP_COST_MAP, catalog_name) - assert str(main_entry["source"]).startswith(AZURE_PRICING_PREFIX) assert backup_entry == main_entry diff --git a/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_cost.py b/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_cost.py index dfa3c7a056e..1f878930207 100644 --- a/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_cost.py +++ b/tests/test_litellm/llms/cohere/ocr/test_cohere_parse_cost.py @@ -1,4 +1,3 @@ -import json from pathlib import Path import pytest @@ -24,19 +23,6 @@ def _ocr_response(model: str, pages_processed: int) -> OCRResponse: ) -@pytest.mark.parametrize("cost_map_path", COST_MAPS, ids=lambda path: path.name) -@pytest.mark.parametrize("model, provider", MODELS) -def test_pricing_entry(cost_map_path: Path, model: str, provider: str) -> None: - with open(cost_map_path) as f: - info = json.load(f).get(model) - - assert info is not None, f"{model} missing from {cost_map_path.name}" - assert info["litellm_provider"] == provider - assert info["mode"] == "ocr" - assert info["supported_endpoints"] == ["/v1/ocr"] - assert info["ocr_cost_per_page"] == COST_PER_PAGE - - @pytest.mark.parametrize("model, provider", MODELS) def test_model_info_resolves_ocr_mode_and_price(local_model_cost_map, model: str, provider: str) -> None: info = litellm.get_model_info(model=model, custom_llm_provider=provider) diff --git a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py index 0b251be5408..904a625ef86 100644 --- a/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py +++ b/tests/test_litellm/llms/databricks/test_databricks_cost_calculator.py @@ -163,23 +163,6 @@ def test_legacy_endpoint_names_still_resolve(local_model_cost_map: None) -> None assert completion_cost == pytest.approx(100 * info["output_cost_per_token"]) -@pytest.mark.parametrize("model", NEW_MODELS) -def test_new_models_price_at_published_dbu_rates(local_model_cost_map: None, model: str) -> None: - info: Final = _model_info(model) - - for field, dbu_per_million in zip(PRICE_FIELDS, PUBLISHED_DBU_PER_MILLION[model]): - assert info[field] == _dollars_per_token(dbu_per_million), field - - -@pytest.mark.parametrize("model", sorted(set(PUBLISHED_DBU_PER_MILLION) - set(ENTRIES_STORING_PROMOTIONAL_RATE))) -def test_cache_rates_derive_from_published_cache_dbu(local_model_cost_map: None, model: str) -> None: - info: Final = _model_info(model) - cache_dbu_per_million: Final = PUBLISHED_DBU_PER_MILLION[model][2:] - - for field, dbu_per_million in zip(CACHE_FIELDS, cache_dbu_per_million): - assert info[field] == _dollars_per_token(dbu_per_million), field - - @pytest.mark.parametrize("model", NEW_MODELS) def test_new_models_carry_cache_pricing(local_model_cost_map: None, model: str) -> None: info: Final = _model_info(model) @@ -255,38 +238,3 @@ def test_sonnet_5_ships_standard_rates_not_introductory(local_model_cost_map: No for field in PRICE_FIELDS: assert sonnet_5[field] == pytest.approx(sonnet_4_6[field]), field - - -@pytest.mark.parametrize("model", ENTRIES_STORING_PROMOTIONAL_RATE) -def test_entries_storing_the_promotional_rate_price_below_the_published_table( - local_model_cost_map: None, - model: str, -) -> None: - info: Final = _model_info(model) - input_dbu, output_dbu, _, _ = PUBLISHED_DBU_PER_MILLION[model] - expiry_hint: Final = f"the gemini promotion expires {PROMOTION_EXPIRES}, after which the list rate applies" - - assert info["input_cost_per_token"] == pytest.approx( - _dollars_per_token(input_dbu) * PROMOTIONAL_DISCOUNT, rel=2e-4 - ), expiry_hint - assert info["output_cost_per_token"] == pytest.approx( - _dollars_per_token(output_dbu) * PROMOTIONAL_DISCOUNT, rel=2e-4 - ), expiry_hint - assert info["cache_creation_input_token_cost"] == pytest.approx(info["input_cost_per_token"]) - assert info["cache_read_input_token_cost"] == pytest.approx(0.1 * info["input_cost_per_token"]) - - -@pytest.mark.parametrize("model", ENTRIES_STORING_LIST_RATE_DESPITE_PROMOTION) -def test_entries_storing_the_list_rate_bill_above_the_promotional_price( - local_model_cost_map: None, - model: str, -) -> None: - info: Final = _model_info(model) - input_dbu, _, _, _ = PUBLISHED_DBU_PER_MILLION[model] - list_rate: Final = _dollars_per_token(input_dbu) - - assert info["input_cost_per_token"] == pytest.approx(list_rate, rel=2e-4), ( - f"{model} moved off the list rate; if it now stores the discount that runs to " - f"{PROMOTION_EXPIRES}, move it into ENTRIES_STORING_PROMOTIONAL_RATE" - ) - assert info["cache_creation_input_token_cost"] == pytest.approx(info["input_cost_per_token"]) diff --git a/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py b/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py index 521ea4f8263..a03b7708238 100644 --- a/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py +++ b/tests/test_litellm/llms/fireworks_ai/rerank/test_fireworks_ai_rerank_transformation.py @@ -3,6 +3,7 @@ Tests for Fireworks AI rerank transformation functionality. """ import json +import uuid from unittest.mock import MagicMock import httpx @@ -181,8 +182,7 @@ class TestFireworksAIRerankTransform: ) # Verify response structure - # Fireworks AI doesn't return "id", so it uses "model" as the id - assert result.id == "accounts/fireworks/models/qwen3-reranker-8b" + assert uuid.UUID(result.id).version == 4 assert len(result.results) == 2 assert result.results[0]["index"] == 0 assert result.results[0]["relevance_score"] == 0.95 @@ -229,16 +229,14 @@ class TestFireworksAIRerankTransform: logging_obj=mock_logging, ) - # Fireworks AI doesn't return "id", so it uses "model" as the id - assert result.id == "accounts/fireworks/models/qwen3-reranker-8b" + assert uuid.UUID(result.id).version == 4 assert len(result.results) == 2 assert result.results[0]["index"] == 0 assert result.results[0]["relevance_score"] == 0.95 # Document should not be present assert "document" not in result.results[0] - def test_transform_rerank_response_missing_id(self): - """Test response transformation when id is missing (should use model name or generate UUID).""" + def test_transform_rerank_response_missing_id_stamps_a_fresh_id_per_call(self): response_data = { "object": "list", "model": "accounts/fireworks/models/qwen3-reranker-8b", @@ -248,23 +246,22 @@ class TestFireworksAIRerankTransform: "usage": {"total_tokens": 10}, } - mock_response = MagicMock(spec=httpx.Response) - mock_response.json.return_value = response_data - mock_response.status_code = 200 - mock_response.headers = {} + def transform() -> str: + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = response_data + mock_response.status_code = 200 + mock_response.headers = {} + return self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=RerankResponse(), + logging_obj=MagicMock(), + ).id - mock_logging = MagicMock() - model_response = RerankResponse() + first, second = transform(), transform() - result = self.config.transform_rerank_response( - model=self.model, - raw_response=mock_response, - model_response=model_response, - logging_obj=mock_logging, - ) - - # Should use model name when id is missing - assert result.id == "accounts/fireworks/models/qwen3-reranker-8b" + assert first != second + assert "accounts/fireworks/models/qwen3-reranker-8b" not in (first, second) def test_transform_rerank_response_missing_results(self): """Test that missing results raises ValueError.""" diff --git a/tests/test_litellm/llms/gemini/test_cost_calculator.py b/tests/test_litellm/llms/gemini/test_cost_calculator.py index 6d547b0dc55..2d56757c601 100644 --- a/tests/test_litellm/llms/gemini/test_cost_calculator.py +++ b/tests/test_litellm/llms/gemini/test_cost_calculator.py @@ -452,33 +452,6 @@ def test_map_traffic_type_to_service_tier( ) -@pytest.mark.parametrize( - "model,custom_llm_provider,expected_cache_read_cost", - [ - ("gemini/gemini-flash-latest", "gemini", 3e-08), - ("gemini/gemini-flash-lite-latest", "gemini", 1e-08), - ("gemini/gemini-2.5-flash-preview-09-2025", "gemini", 3e-08), - ("gemini/gemini-2.5-flash-lite-preview-06-17", "gemini", 1e-08), - ("vertex_ai/gemini-2.5-flash-preview-09-2025", "vertex_ai", 3e-08), - ("vertex_ai/gemini-2.5-flash-lite-preview-06-17", "vertex_ai", 1e-08), - ], -) -def test_flash_alias_cache_read_is_ten_percent_of_input( - monkeypatch, model, custom_llm_provider, expected_cache_read_cost -): - monkeypatch.setenv("LITELLM_LOCAL_MODEL_COST_MAP", "True") - monkeypatch.setattr(litellm, "model_cost", litellm.get_model_cost_map(url="")) - - model_info = litellm.get_model_info( - model=model, custom_llm_provider=custom_llm_provider - ) - - assert model_info["cache_read_input_token_cost"] == expected_cache_read_cost - assert model_info["cache_read_input_token_cost"] == pytest.approx( - 0.10 * model_info["input_cost_per_token"] - ) - - @pytest.mark.parametrize( "prefixed,bare", [ diff --git a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py index 40e54f71eeb..c894f92148d 100644 --- a/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py +++ b/tests/test_litellm/llms/mistral/ocr/test_mistral_ocr_cost.py @@ -5,7 +5,6 @@ for mistral-ocr-4-0 and mistral-ocr-latest, which now both resolve to OCR 4 at $4 / 1000 pages. """ -import json from pathlib import Path import pytest @@ -45,12 +44,6 @@ def _annotated_ocr_response(model: str, pages_processed: int | None, annotation_ ) -@pytest.mark.parametrize("model", ["mistral-ocr-4-0", "mistral-ocr-latest"]) -def test_model_info_ocr4_price(model: str) -> None: - info = litellm.get_model_info(model=f"mistral/{model}", custom_llm_provider="mistral") - assert info["ocr_cost_per_page"] == OCR4_COST_PER_PAGE - - @pytest.mark.parametrize("model", ["mistral-ocr-4-0", "mistral-ocr-latest"]) @pytest.mark.parametrize("pages_processed", [1, 3, 10]) def test_ocr4_cost_scales_with_pages(model: str, pages_processed: int) -> None: @@ -63,20 +56,6 @@ def test_ocr4_cost_scales_with_pages(model: str, pages_processed: int) -> None: assert cost == pytest.approx(OCR4_COST_PER_PAGE * pages_processed) - -@pytest.mark.parametrize("cost_map_path", [MAIN_COST_MAP, BACKUP_COST_MAP]) -def test_ocr3_pricing_entry(cost_map_path: Path) -> None: - with open(cost_map_path) as f: - info = json.load(f).get(OCR3_MODEL) - - assert info is not None, f"{OCR3_MODEL} missing from {cost_map_path.name}" - assert info["litellm_provider"] == "mistral" - assert info["mode"] == "ocr" - assert info["supported_endpoints"] == ["/v1/ocr"] - assert info["ocr_cost_per_page"] == OCR3_COST_PER_PAGE - assert info["annotation_cost_per_page"] == OCR3_ANNOTATION_COST_PER_PAGE - - def test_ocr3_model_info_price(local_model_cost_map) -> None: info = litellm.get_model_info(model=OCR3_MODEL, custom_llm_provider="mistral") assert info["ocr_cost_per_page"] == OCR3_COST_PER_PAGE diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py index 7fea5ac0965..3ec734611ef 100644 --- a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py +++ b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_integration.py @@ -104,10 +104,12 @@ class TestVertexAIRerankIntegration: raw_response=mock_response, model_response=model_response, logging_obj=mock_logging, + request_data=request_data, ) # Verify response structure - assert result.id == f"vertex_ai_rerank_{self.model}" + assert result.id.startswith("vertex_ai_rerank_") + assert result.id != f"vertex_ai_rerank_{self.model}" assert len(result.results) == 2 # Results should be sorted by relevance score (descending) @@ -116,8 +118,8 @@ class TestVertexAIRerankIntegration: assert result.results[1]["index"] == 0 # Second highest score assert result.results[1]["relevance_score"] == 0.92 - # Verify metadata - assert result.meta["billed_units"]["search_units"] == 2 + # Verify metadata: 4 input records bill as 1 search unit (ceil(4/100)) + assert result.meta["billed_units"]["search_units"] == 1 def test_return_documents_false_flow(self): """Test rerank flow when return_documents=False (ID-only response).""" diff --git a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py index c2ea6f6fab9..630b2e1eb34 100644 --- a/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py +++ b/tests/test_litellm/llms/vertex_ai/rerank/test_vertex_ai_rerank_transformation.py @@ -287,10 +287,11 @@ class TestVertexAIRerankTransform: raw_response=mock_response, model_response=model_response, logging_obj=mock_logging, + request_data={"records": [{"id": "0"}, {"id": "1"}]}, ) # Verify response structure - assert result.id == f"vertex_ai_rerank_{self.model}" + assert result.id.startswith("vertex_ai_rerank_") assert len(result.results) == 2 assert result.results[0]["index"] == 1 # Converted back to 0-based index assert result.results[0]["relevance_score"] == 0.98 @@ -298,7 +299,7 @@ class TestVertexAIRerankTransform: assert result.results[1]["relevance_score"] == 0.64 # Verify metadata - assert result.meta["billed_units"]["search_units"] == 2 + assert result.meta["billed_units"]["search_units"] == 1 def test_transform_rerank_response_with_ignore_record_details(self): """Test response transformation when ignoreRecordDetailsInResponse=true.""" @@ -326,6 +327,96 @@ class TestVertexAIRerankTransform: assert result.results[1]["index"] == 0 assert result.results[1]["relevance_score"] == 1.0 + def _build_response(self, num_records): + response_data = { + "records": [ + {"id": str(i), "score": 1.0 - i / 1000, "title": "t", "content": "c"} + for i in range(num_records) + ] + } + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = response_data + mock_response.text = json.dumps(response_data) + return mock_response + + def test_search_units_from_input_records_not_truncated_response(self): + """ + Regression for LIT-4995 part 1: search_units must be derived from the + billable input records (ceil(input / 100)), not from the response, which + Google truncates to topN. + """ + documents = [f"doc {i}" for i in range(5)] + request_data = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params={"query": "q", "documents": documents, "top_n": 2}, + headers={}, + ) + # Google truncates the response to top_n=2 records + mock_response = self._build_response(num_records=2) + + result = self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=RerankResponse(), + logging_obj=MagicMock(), + request_data=request_data, + ) + + assert result.meta["billed_units"]["search_units"] == 1 + + def test_search_units_rounds_up_per_hundred_input_records(self): + """ + Regression for LIT-4995 part 1: one query bills up to 100 input records, + so 150 input records is 2 search units regardless of the response size. + """ + documents = [f"doc {i}" for i in range(150)] + request_data = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params={"query": "q", "documents": documents, "top_n": 3}, + headers={}, + ) + mock_response = self._build_response(num_records=3) + + result = self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=RerankResponse(), + logging_obj=MagicMock(), + request_data=request_data, + ) + + assert result.meta["billed_units"]["search_units"] == 2 + + def test_response_id_is_unique_per_request(self): + """ + Regression for LIT-4995 part 2: response IDs must be unique per request, + not a constant derived only from the model name. + """ + request_data = self.config.transform_rerank_request( + model=self.model, + optional_rerank_params={"query": "q", "documents": ["a", "b"]}, + headers={}, + ) + mock_response = self._build_response(num_records=2) + + first = self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=RerankResponse(), + logging_obj=MagicMock(), + request_data=request_data, + ) + second = self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=RerankResponse(), + logging_obj=MagicMock(), + request_data=request_data, + ) + + assert first.id != second.id + assert first.id != f"vertex_ai_rerank_{self.model}" + def test_transform_rerank_response_json_error(self): """Test response transformation with JSON parsing error.""" mock_response = MagicMock(spec=httpx.Response) diff --git a/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py b/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py index f466b7e19b5..5eb4bf31845 100644 --- a/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py +++ b/tests/test_litellm/llms/voyage/rerank/test_voyage_rerank_transformation.py @@ -3,6 +3,7 @@ Tests for Voyage AI rerank transformation functionality. """ import json +import uuid from unittest.mock import MagicMock, patch import httpx @@ -258,6 +259,33 @@ class TestVoyageRerankTransform: assert "Failed to parse response" in str(exc_info.value) + def test_transform_rerank_response_without_id_stamps_a_fresh_id_per_call(self): + response_data = { + "object": "list", + "data": [{"relevance_score": 0.5, "index": 0}], + "model": "rerank-2.5", + "usage": {"total_tokens": 10}, + } + + def transform() -> str: + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = response_data + mock_response.status_code = 200 + mock_response.text = json.dumps(response_data) + mock_response.headers = {} + return self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=RerankResponse(), + logging_obj=MagicMock(), + ).id + + first, second = transform(), transform() + + assert uuid.UUID(first).version == 4 + assert first != second + assert f"voyage-rerank-{self.model}" not in (first, second) + def test_get_supported_cohere_rerank_params(self): """Test getting supported parameters for Voyage AI rerank.""" supported_params = self.config.get_supported_cohere_rerank_params(self.model) diff --git a/tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py b/tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py index c8f2c4dd87c..ccbd318959f 100644 --- a/tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py +++ b/tests/test_litellm/llms/watsonx/rerank/test_watsonx_rerank.py @@ -120,9 +120,7 @@ class TestIBMWatsonXRerankTransform: logging_obj=mock_logging, ) - # Verify response structure - # IBM watsonx.ai doesn't return "id", so it uses "model" as the id - assert result.id == "watsonx/cross-encoder/ms-marco-minilm-l-12-v2" + assert uuid.UUID(result.id).version == 4 assert len(result.results) == 2 assert result.results[0]["index"] == 0 assert result.results[0]["relevance_score"] == 6.53515625 @@ -172,9 +170,7 @@ class TestIBMWatsonXRerankTransform: logging_obj=mock_logging, ) - # Verify response structure - # IBM watsonx.ai doesn't return "id", so it uses "model" as the id - assert result.id == "watsonx/cross-encoder/ms-marco-minilm-l-12-v2" + assert uuid.UUID(result.id).version == 4 assert len(result.results) == 2 assert result.results[0]["index"] == 0 @@ -231,6 +227,30 @@ class TestIBMWatsonXRerankTransform: logging_obj=mock_logging, ) + def test_transform_rerank_response_without_id_stamps_a_fresh_id_per_call(self): + response_data = { + "model_id": self.model, + "results": [{"index": 0, "score": 1.5}], + "input_token_count": 12, + } + + def transform() -> str: + mock_response = MagicMock(spec=httpx.Response) + mock_response.json.return_value = response_data + mock_response.status_code = 200 + mock_response.headers = {} + return self.config.transform_rerank_response( + model=self.model, + raw_response=mock_response, + model_response=RerankResponse(), + logging_obj=MagicMock(), + ).id + + first, second = transform(), transform() + + assert first != second + assert self.model not in (first, second) + def test_get_supported_cohere_rerank_params(self): """Test getting supported parameters for IBM watsonx.ai rerank.""" supported_params = self.config.get_supported_cohere_rerank_params(self.model) diff --git a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py index 4cff5c76b9e..8f933f7e5c2 100644 --- a/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py +++ b/tests/test_litellm/llms/xai/responses/test_xai_responses_transformation.py @@ -119,6 +119,67 @@ class TestXAIResponsesAPITransformation: assert tool["filters"]["allowed_domains"] == ["wikipedia.org", "x.ai"] assert tool["enable_image_understanding"] is True + def test_web_search_nested_filters_preserved(self): + """The documented nested 'filters' shape must reach xAI instead of being dropped""" + config = XAIResponsesAPIConfig() + + params = ResponsesAPIOptionalRequestParams( + tools=[ + { + "type": "web_search", + "filters": {"allowed_domains": ["grokipedia.com"], "excluded_domains": ["example.com"]}, + } + ] + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="grok-4-1-fast", + drop_params=False, + ) + + tool = result["tools"][0] + assert tool["filters"]["allowed_domains"] == ["grokipedia.com"] + assert tool["filters"]["excluded_domains"] == ["example.com"] + + def test_web_search_nested_filters_win_over_flat(self): + """Nested filters take precedence when both shapes are sent""" + config = XAIResponsesAPIConfig() + + params = ResponsesAPIOptionalRequestParams( + tools=[ + { + "type": "web_search", + "allowed_domains": ["flat.com"], + "filters": {"allowed_domains": ["nested.com"]}, + } + ] + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="grok-4-1-fast", + drop_params=False, + ) + + assert result["tools"][0]["filters"] == {"allowed_domains": ["nested.com"]} + + def test_web_search_empty_nested_filters_win_over_flat(self): + """An explicit empty 'filters' object means unrestricted search, even when stale flat fields are present""" + config = XAIResponsesAPIConfig() + + params = ResponsesAPIOptionalRequestParams( + tools=[{"type": "web_search", "allowed_domains": ["flat.com"], "filters": {}}] + ) + + result = config.map_openai_params( + response_api_optional_params=params, + model="grok-4-1-fast", + drop_params=False, + ) + + assert result["tools"][0] == {"type": "web_search"} + def test_web_search_search_context_size_removed(self): """Test that search_context_size is removed from web_search tools""" config = XAIResponsesAPIConfig() diff --git a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py index c67dca11a56..524ca6a02d7 100644 --- a/tests/test_litellm/llms/xai/test_xai_chat_transformation.py +++ b/tests/test_litellm/llms/xai/test_xai_chat_transformation.py @@ -124,6 +124,24 @@ class TestXAIParallelToolCalls: assert result["messages"][0]["role"] == "user" +class TestXAIChatWebSearchOptions: + """XAI answers /chat/completions requests carrying web_search_options with a 410 (Live Search retired)""" + + def test_transform_request_drops_web_search_options(self): + config = XAIChatConfig() + + result = config.transform_request( + model="xai/grok-4.6", + messages=[{"role": "user", "content": "newest litellm version?"}], + optional_params={"web_search_options": {"search_context_size": "medium"}, "temperature": 0.5}, + litellm_params={}, + headers={}, + ) + + assert "web_search_options" not in result + assert result["temperature"] == 0.5 + + class TestXAIUsageNormalization: def test_preserves_reasoning_tokens_in_total_usage(self): usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=200) diff --git a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py index e591c1ae682..4c8231d357e 100644 --- a/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py +++ b/tests/test_litellm/llms/xai/test_xai_redirected_slug_pricing.py @@ -54,9 +54,6 @@ CODE_SLUGS = ( "xai/grok-code-fast-1", "xai/grok-code-fast-1-0825", ) -RETIREMENT_DATE = "2026-05-15" -GROK_3_MINI_RETIREMENT_DATE = "2026-02-28" - BASE_COST_FIELDS = ("input_cost_per_token", "output_cost_per_token", "cache_read_input_token_cost") TIER_COST_FIELDS = ( "input_cost_per_token_above_200k_tokens", @@ -65,10 +62,6 @@ TIER_COST_FIELDS = ( ) -def expected_retirement_date(slug: str) -> str: - return GROK_3_MINI_RETIREMENT_DATE if slug in GROK_3_MINI_SLUGS else RETIREMENT_DATE - - @pytest.fixture(scope="module", params=[p.name for p in MAP_PATHS]) def cost_map(request: pytest.FixtureRequest) -> dict: path = next(p for p in MAP_PATHS if p.name == request.param) @@ -92,15 +85,9 @@ def test_code_slug_bills_at_grok_build_rate(cost_map: dict, slug: str): assert entry[field] == target[field], field -@pytest.mark.parametrize("slug", (*REDIRECTED_SLUGS, *CODE_SLUGS)) -def test_redirected_slug_keeps_its_retirement_date(cost_map: dict, slug: str): - assert cost_map[slug]["deprecation_date"] == expected_retirement_date(slug) - - def test_a_live_xai_model_is_untouched(cost_map: dict): """Guard against the repricing leaking onto models xAI still serves directly.""" assert cost_map["xai/grok-4.6"]["input_cost_per_token"] != cost_map[REDIRECT_TARGET]["input_cost_per_token"] - assert "deprecation_date" not in cost_map["xai/grok-4.6"] @pytest.mark.parametrize("slug", REDIRECTED_SLUGS) diff --git a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py index 21e0b83791f..6e9770bced8 100644 --- a/tests/test_litellm/proxy/auth/test_auth_exception_handler.py +++ b/tests/test_litellm/proxy/auth/test_auth_exception_handler.py @@ -487,6 +487,34 @@ async def test_route_passed_to_post_call_failure_hook(): assert call_args["user_api_key_dict"].request_route == test_route +@pytest.mark.asyncio +async def test_dynamic_route_normalized_on_auth_failure(): + handler = UserAPIKeyAuthExceptionHandler() + + with ( + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + ) as mock_post_call_failure_hook, + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.general_settings", {} + ), + pytest.raises(ProxyException), + ): + await handler._handle_authentication_error( + HTTPException(status_code=401, detail="Authentication Error, Invalid proxy server token passed"), + MagicMock(), + {}, + "/v1/responses/resp_attacker_controlled_id", + None, + "sk-doesnotexist", + ) + + hook_kwargs = mock_post_call_failure_hook.call_args.kwargs + assert hook_kwargs["route"] == "/v1/responses/resp_attacker_controlled_id" + assert hook_kwargs["user_api_key_dict"].request_route == "/v1/responses/{response_id}" + + @pytest.mark.asyncio async def test_resolved_identity_exported_on_auth_failure(): """Regression: when auth fails AFTER the key/team/user identity is resolved @@ -795,6 +823,89 @@ async def test_auth_failure_ip_stamp_does_not_mutate_callers_request_data(): assert request_data == {"model": "gpt-4o"} +@pytest.mark.asyncio +@pytest.mark.parametrize( + "request_data, metadata_key, route", + [ + pytest.param({"model": "gpt-4o"}, "metadata", "/v1/chat/completions", id="chat_metadata"), + pytest.param({"litellm_metadata": {}}, "litellm_metadata", "/v1/responses", id="responses_litellm_metadata"), + ], +) +async def test_auth_failure_logs_user_agent(request_data: dict[str, object], metadata_key: str, route: str) -> None: + """Auth gate rejections never reach `add_litellm_data_to_request`, which is what + stamps `user_agent`, so the failure spend log and prometheus `user_agent` label + had nothing to identify an abusive client by.""" + with ( + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.auth.auth_exception_handler.seed_request_identity" + ), + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ) as mock_hook, + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + with pytest.raises(ProxyException): + await UserAPIKeyAuthExceptionHandler._handle_authentication_error( + ProxyException( + message="Invalid API key", + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ), + _http_request(headers={"user-agent": "abusive-client/9.9"}), + request_data, + route, + None, + "sk-bad-key", + ) + + logged_metadata = mock_hook.call_args[1]["request_data"][metadata_key] + assert logged_metadata["user_agent"] == "abusive-client/9.9" + assert logged_metadata["requester_ip_address"] == "10.1.2.3" + + +@pytest.mark.asyncio +async def test_auth_failure_without_headers_scope_still_raises_original_error() -> None: + """A request scope with no `headers` entry must surface the auth error itself, not a + `KeyError` from reading the User-Agent.""" + with ( + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.auth.auth_exception_handler.seed_request_identity" + ), + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.proxy_logging_obj.post_call_failure_hook", + new_callable=AsyncMock, + return_value=None, + ) as mock_hook, + patch( # test-quality-ok: handler reads proxy_server globals at call time + "litellm.proxy.proxy_server.general_settings", + {"allow_requests_on_db_unavailable": False}, + ), + ): + with pytest.raises(ProxyException) as exc_info: + await UserAPIKeyAuthExceptionHandler._handle_authentication_error( + ProxyException( + message="Invalid API key", + type=ProxyErrorTypes.auth_error, + param=None, + code=status.HTTP_401_UNAUTHORIZED, + ), + Request(scope={"type": "http"}), + {"model": "gpt-4o"}, + "/v1/chat/completions", + None, + "sk-bad-key", + ) + + assert str(exc_info.value.code) == str(status.HTTP_401_UNAUTHORIZED) + assert "user_agent" not in mock_hook.call_args[1]["request_data"].get("metadata", {}) + + def _marked_malformed_key_error() -> HTTPException: """Build the malformed-key 401 as its raise site does: marker stamped on it.""" error = HTTPException(status_code=401, detail="LiteLLM Virtual Key expected. Received=test") diff --git a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py index cf1f665ad21..5263cf2774c 100644 --- a/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py +++ b/tests/test_litellm/proxy/auth/test_custom_auth_end_user_budget.py @@ -277,3 +277,18 @@ def test_update_valid_token_db_values_override_custom_auth_when_set(): # DB values should win assert result.end_user_tpm_limit == 500 assert result.end_user_model_max_budget == db_budget + + +def test_end_user_budget_tpd_limit_reaches_the_token(): + from litellm.proxy.auth.user_api_key_auth import _apply_budget_limits_to_end_user_params + + end_user_params = {"end_user_id": "user_1"} + _apply_budget_limits_to_end_user_params( + end_user_params=end_user_params, + budget_info=LiteLLM_BudgetTable(rpm_limit=5, tpd_limit=750000), + end_user_id="user_1", + ) + result = update_valid_token_with_end_user_params(UserAPIKeyAuth(token="test_token"), end_user_params) + + assert result.end_user_rpm_limit == 5 + assert result.end_user_tpd_limit == 750000 diff --git a/tests/test_litellm/proxy/auth/test_team_grants.py b/tests/test_litellm/proxy/auth/test_team_grants.py index 447fc1c93a1..7b6717f804f 100644 --- a/tests/test_litellm/proxy/auth/test_team_grants.py +++ b/tests/test_litellm/proxy/auth/test_team_grants.py @@ -27,6 +27,7 @@ def _full_team(model_aliases=ALIASES) -> LiteLLM_TeamTable: team_alias="grants-team", tpm_limit=1000, rpm_limit=10, + tpd_limit=200000, max_budget=50.0, soft_budget=25.0, spend=12.5, @@ -67,6 +68,7 @@ def test_team_grants_cover_every_team_field_the_key_path_gets(): assert token.team_alias == "grants-team" assert token.team_tpm_limit == 1000 assert token.team_rpm_limit == 10 + assert token.team_tpd_limit == 200000 assert token.team_max_budget == 50.0 assert token.team_soft_budget == 25.0 assert token.team_spend == 12.5 diff --git a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py index bc4e756eb65..72cd7a218d3 100644 --- a/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py +++ b/tests/test_litellm/proxy/common_utils/test_http_parsing_utils.py @@ -512,8 +512,8 @@ async def test_surrogate_repair_skipped_above_size_limit(monkeypatch): the repair must be skipped and the existing 400 raised immediately, while bodies at or below the limit still get repaired. - `\\ud83d` is a lone high-surrogate escape: orjson rejects it, the json fallback - accepts it, so a body containing it is only salvaged when the repair path runs. + `NaN` is rejected by orjson and accepted by the json fallback, so a body containing + it is only salvaged when the repair path runs. """ import litellm.proxy.common_utils.http_parsing_utils as http_parsing_utils @@ -522,14 +522,14 @@ async def test_surrogate_repair_skipped_above_size_limit(monkeypatch): http_parsing_utils, "MAX_REQUEST_BODY_SIZE_TO_REPAIR_MB", 100 / (1024 * 1024) ) - small_body = b'{"model":"gpt-4o","x":"\\ud83d"}' + small_body = b'{"model":"gpt-4o","x":NaN}' assert len(small_body) <= 100 repaired = await _read_request_body(_make_json_request(small_body)) assert repaired["model"] == "gpt-4o" padding = "a" * 200 large_body = ( - b'{"model":"gpt-4o","pad":"' + padding.encode() + b'","x":"\\ud83d"}' + b'{"model":"gpt-4o","pad":"' + padding.encode() + b'","x":NaN}' ) assert len(large_body) > 100 with pytest.raises(ProxyException) as exc_info: @@ -546,6 +546,33 @@ async def test_surrogate_repair_skipped_above_size_limit(monkeypatch): assert repaired_large["model"] == "gpt-4o" +@pytest.mark.asyncio +@pytest.mark.parametrize( + "content", + [ + pytest.param(b"say ok \\ud83d", id="lone-high-surrogate"), + pytest.param(b"say ok \\ude00", id="lone-low-surrogate"), + pytest.param(b"\\ud83d\\ud83d\\ude00", id="lone-high-before-valid-pair"), + ], +) +async def test_lone_surrogate_escape_is_rejected_with_400(content: bytes): + """ + orjson rejects a lone surrogate escape, and the json fallback accepts it, so the + parsed body used to carry a code point no provider request can UTF-8 encode. That + surfaced as a 500 from the provider handler instead of a 400 for the bad input. + """ + body = b'{"model":"gpt-4o","messages":[{"role":"user","content":"' + content + b'"}]}' + with pytest.raises(ProxyException) as exc_info: + await _read_request_body(_make_json_request(body)) + assert exc_info.value.code == "400" + assert exc_info.value.type == "invalid_request_error" + assert "Invalid JSON payload" in exc_info.value.message + + paired = body.replace(content, b"say ok \\ud83d\\ude00") + parsed = await _read_request_body(_make_json_request(paired)) + assert parsed["messages"][0]["content"] == "say ok \U0001F600" + + @pytest.mark.asyncio async def test_get_form_data(): """ diff --git a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py index 560953f0b51..943a6c905c0 100644 --- a/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py +++ b/tests/test_litellm/proxy/common_utils/test_reset_budget_job.py @@ -19,7 +19,7 @@ from litellm.constants import ( RESET_BUDGET_JOB_LOCK_TTL_SECONDS, RESET_BUDGET_JOB_NAME, ) -from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob +from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob, _RowReset from litellm.proxy.common_utils.timezone_utils import BudgetResetSettings @@ -243,14 +243,18 @@ def test_write_key_reset_updates_skips_none_token_and_still_writes_the_rest(rese LiteLLM_VerificationToken(token="tok-ok", budget_reset_at=reset_at), ] - asyncio.run(reset_budget_job._write_key_reset_updates(updated_keys=keys)) + asyncio.run( + reset_budget_job._write_key_reset_updates( + updated_keys=[_RowReset(row=k, spend_decrement=(k.spend or 0.0)) for k in keys] + ) + ) assert _batch_writes(mock_prisma_client, "key") == [ { "table": "key", "op": "update", "where": {"token": "tok-ok"}, - "data": {"spend": 0, "budget_reset_at": reset_at}, + "data": {"spend": {"decrement": 0.0}, "budget_reset_at": reset_at}, } ] @@ -282,7 +286,7 @@ def test_reset_budget_for_key(reset_budget_job, mock_prisma_client): assert len(key_writes) == 1 write = key_writes[0] assert write["where"] == {"token": "tok-key-1"} - assert write["data"]["spend"] == 0 + assert write["data"]["spend"] == {"decrement": 100.0} assert write["data"]["budget_reset_at"] > now assert set(write["data"].keys()) == {"spend", "budget_reset_at"} @@ -345,7 +349,7 @@ def test_reset_budget_for_user(reset_budget_job, mock_prisma_client): assert len(user_writes) == 1 write = user_writes[0] assert write["where"] == {"user_id": "uid-1"} - assert write["data"]["spend"] == 0 + assert write["data"]["spend"] == {"decrement": 200.0} assert write["data"]["budget_reset_at"] > now assert set(write["data"].keys()) == {"spend", "budget_reset_at"} @@ -374,7 +378,7 @@ def test_reset_budget_for_team(reset_budget_job, mock_prisma_client): assert len(team_writes) == 1 write = team_writes[0] assert write["where"] == {"team_id": "tid-1"} - assert write["data"]["spend"] == 0 + assert write["data"]["spend"] == {"decrement": 500.0} assert write["data"]["budget_reset_at"] > now assert set(write["data"].keys()) == {"spend", "budget_reset_at"} @@ -488,15 +492,15 @@ def test_reset_budget_all(reset_budget_job, mock_prisma_client): # key/user/team rows are written via batch_()..update — verify each # one fired exactly once with the narrow {spend, budget_reset_at} payload. - for table_name, where in [ - ("key", {"token": "tok-all-1"}), - ("user", {"user_id": "uid-all-1"}), - ("team", {"team_id": "tid-all-1"}), + for table_name, where, decrement in [ + ("key", {"token": "tok-all-1"}, 100.0), + ("user", {"user_id": "uid-all-1"}, 200.0), + ("team", {"team_id": "tid-all-1"}, 500.0), ]: writes = _batch_writes(mock_prisma_client, table_name, op="update") assert len(writes) == 1, f"expected 1 {table_name} write, got {len(writes)}" assert writes[0]["where"] == where - assert writes[0]["data"]["spend"] == 0 + assert writes[0]["data"]["spend"] == {"decrement": decrement} assert set(writes[0]["data"].keys()) == {"spend", "budget_reset_at"} # The budget tier's cascade rides the same batch machinery. @@ -1226,6 +1230,7 @@ def _make_counter_invalidation_job(monkeypatch): spend_counter_cache.in_memory_cache.set_cache = MagicMock() spend_counter_cache.redis_cache = MagicMock() spend_counter_cache.redis_cache.async_set_cache = AsyncMock() + spend_counter_cache.redis_cache.async_delete_cache = AsyncMock() user_api_key_cache = MagicMock() user_api_key_cache.async_delete_cache = AsyncMock() @@ -1260,7 +1265,8 @@ def test_reset_budget_for_keys_invalidates_redis_counter(reset_budget_job, mock_ asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) - counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-abc", value=0.0, ttl=60) + counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:key:sk-abc") + counter_cache.redis_cache.async_delete_cache.assert_any_await(key="spend:key:sk-abc") def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock_prisma_client, monkeypatch): @@ -1284,7 +1290,8 @@ def test_reset_budget_for_users_invalidates_redis_counter(reset_budget_job, mock asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) - counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:user:alice", value=0.0, ttl=60) + counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:user:alice") + counter_cache.redis_cache.async_delete_cache.assert_any_await(key="spend:user:alice") def test_reset_budget_for_proxy_budget_row_invalidates_global_spend_cache( @@ -1368,7 +1375,8 @@ def test_reset_budget_for_teams_invalidates_redis_counter(reset_budget_job, mock asyncio.run(reset_budget_job.reset_budget_for_litellm_teams()) - counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:team:team-x", value=0.0, ttl=60) + counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:team:team-x") + counter_cache.redis_cache.async_delete_cache.assert_any_await(key="spend:team:team-x") def test_reset_does_not_zero_counter_when_db_write_fails(monkeypatch): @@ -1428,7 +1436,7 @@ def test_reset_does_not_zero_counter_when_db_write_fails(monkeypatch): # assert_not_called() instead of iterating call_args_list, because the # latter is vacuously true when the list is empty (would pass even if # the bypass were re-introduced via a different code path). - counter_cache.in_memory_cache.set_cache.assert_not_called() + counter_cache.in_memory_cache.delete_cache.assert_not_called() def test_reset_budget_for_keys_writes_only_spend_and_reset_at(reset_budget_job, mock_prisma_client): @@ -1526,8 +1534,8 @@ def test_budget_table_reset_invalidates_counters_and_management_cache( asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - counter_cache.in_memory_cache.set_cache.assert_any_call(key=counter_key, value=0.0, ttl=60) - counter_cache.redis_cache.async_set_cache.assert_any_await(key=counter_key, value=0.0, ttl=60) + counter_cache.in_memory_cache.delete_cache.assert_any_call(key=counter_key) + counter_cache.redis_cache.async_delete_cache.assert_any_await(key=counter_key) deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} assert cache_keys <= deleted @@ -1565,8 +1573,8 @@ def test_budget_table_reset_invalidates_enduser_counter_and_cache(reset_budget_j asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:end_user:customer-42", value=0.0, ttl=60) - counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:end_user:customer-42", value=0.0, ttl=60) + counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:end_user:customer-42") + counter_cache.redis_cache.async_delete_cache.assert_any_await(key="spend:end_user:customer-42") deleted: Final = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} assert "end_user_id:customer-42" in deleted @@ -1627,7 +1635,7 @@ def test_access_groups_are_untouched_when_no_budget_is_due(reset_budget_job, moc assert mock_prisma_client.db.litellm_modelaccessgroupbudgettable.find_many_calls == [] assert _batch_writes(mock_prisma_client, "model_access_group") == [] - counter_cache.in_memory_cache.set_cache.assert_not_called() + counter_cache.in_memory_cache.delete_cache.assert_not_called() counter_cache.user_api_key_cache.async_delete_cache.assert_not_awaited() @@ -1646,7 +1654,7 @@ def test_budget_table_reset_invalidates_every_access_group_not_just_the_first( deleted = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} assert deleted == {"model_access_group:group-a", "model_access_group:group-b", "model_access_group:group-c"} for name in ("group-a", "group-b", "group-c"): - counter_cache.in_memory_cache.set_cache.assert_any_call(key=f"spend:model_access_group:{name}", value=0.0, ttl=60) + counter_cache.in_memory_cache.delete_cache.assert_any_call(key=f"spend:model_access_group:{name}") def test_budget_cascade_carries_access_group_overage_when_rollover_enabled( @@ -1678,7 +1686,7 @@ def test_budget_cascade_carries_access_group_overage_when_rollover_enabled( } in writes assert _replay_spend_writes(writes, 15.0) == 5.0 assert _replay_spend_writes(writes, 8.0) == 0 - counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:model_access_group:gpt-4-group", value=5.0, ttl=60) + counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:model_access_group:gpt-4-group") # --------------------------------------------------------------------------- @@ -1769,7 +1777,7 @@ def test_budget_reset_at_is_not_advanced_when_the_cascade_fails(db_factory, monk assert prisma_client.db.batch_calls == [], "a failed cascade must not persist any write" assert prisma_client.db.batchers[0].committed is False assert prisma_client.updated_data["budget"] == [], "budget_reset_at must not be advanced outside the transaction" - counter_cache.in_memory_cache.set_cache.assert_not_called() + counter_cache.in_memory_cache.delete_cache.assert_not_called() counter_cache.user_api_key_cache.async_delete_cache.assert_not_awaited() @@ -1806,7 +1814,7 @@ def test_caches_are_invalidated_only_after_the_transaction_commits(monkeypatch): cap while the DB still holds the over-budget spend.""" events = [] counter_cache = _make_counter_invalidation_job(monkeypatch) - counter_cache.in_memory_cache.set_cache.side_effect = lambda **kwargs: events.append("counter") + counter_cache.in_memory_cache.delete_cache.side_effect = lambda **kwargs: events.append("counter") job, _ = _job_with_expired_budget(OrderRecordingDB(events)) @@ -2839,13 +2847,7 @@ _SPEND_ACCRUED_AFTER_COMMIT = 7.5 class AmbiguousCommitClient(MockPrismaClient): - """A client whose batch commit lands in the database and only then fails in - transit, so the caller cannot tell whether it committed. - - The queued spend-zero is applied to `key_spend`, and fresh usage accrues in - the window between that landed commit and any replay, so a replay is - observable as erased spend rather than merely as an extra commit. - """ + """A client whose batch commit lands in the database and only then fails in transit.""" def __init__(self, *, error: Exception, spend_accrued_after_commit: float): super().__init__() @@ -2864,7 +2866,12 @@ class AmbiguousCommitClient(MockPrismaClient): outer.commit_attempts += 1 result = await batch_commit() for call in batcher.calls: - if call["table"] == "key" and call["data"].get("spend") == 0: + if call["table"] != "key": + continue + spend_field = call["data"].get("spend") + if isinstance(spend_field, dict): + outer.key_spend -= spend_field["decrement"] + elif spend_field == 0: outer.key_spend = 0.0 if outer.commit_attempts > 1: return result @@ -2886,22 +2893,19 @@ class AmbiguousCommitClient(MockPrismaClient): [ (httpx.ReadError("response lost in transit"), 1, _SPEND_ACCRUED_AFTER_COMMIT, []), (httpx.ReadTimeout("response lost in transit"), 1, _SPEND_ACCRUED_AFTER_COMMIT, []), - (httpx.ConnectError("never left the client"), 2, 0.0, ["reset_budget_write_keys_failure"]), + ( + httpx.ConnectError("never left the client"), + 2, + _SPEND_ACCRUED_AFTER_COMMIT - _DUE_ROW_SPEND, + ["reset_budget_write_keys_failure"], + ), ], ids=["read_error", "read_timeout", "connect_error_erasure_control"], ) def test_ambiguous_commit_replay_does_not_erase_newly_accrued_spend( error, expected_commits, expected_spend, expected_reconnects ): - """A reset zeroes spend unconditionally, so replaying a commit that already - landed erases every dollar spent since it landed (LIT-5372 review finding). - - The `connect_error` case is the control: it is the one error class allowed - to replay, and driving it through this same land-then-fail harness proves - the spend assertion can actually observe an erasure. In production a - ConnectError means the statements never reached the database, so its replay - has nothing to erase. - """ + """Replaying a commit that already landed erases spend accrued since it landed.""" client = AmbiguousCommitClient(error=error, spend_accrued_after_commit=_SPEND_ACCRUED_AFTER_COMMIT) client.data["key"] = [_due_row("key", "tok-1")] job = ResetBudgetJob(proxy_logging_obj=MockProxyLogging(), prisma_client=client) @@ -2999,7 +3003,7 @@ def test_direct_reset_carries_overage_when_rollover_enabled( assert writes[0]["data"]["spend"] == {"decrement": 100.0} assert writes[0]["data"]["budget_reset_at"] > now counter_prefix = {"key": "spend:key", "user": "spend:user", "team": "spend:team"}[table] - counter_cache.in_memory_cache.set_cache.assert_any_call(key=f"{counter_prefix}:{id_value}", value=50.0, ttl=60) + counter_cache.in_memory_cache.delete_cache.assert_any_call(key=f"{counter_prefix}:{id_value}") def test_direct_reset_zeroes_under_budget_row_even_with_rollover( @@ -3017,8 +3021,8 @@ def test_direct_reset_zeroes_under_budget_row_even_with_rollover( asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) - assert _batch_writes(mock_prisma_client, "key")[0]["data"]["spend"] == 0 - counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:tok-under", value=0.0, ttl=60) + assert _batch_writes(mock_prisma_client, "key")[0]["data"]["spend"] == {"decrement": 40.0} + counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:key:tok-under") def test_direct_reset_zeroes_row_without_max_budget_even_with_rollover( @@ -3037,7 +3041,7 @@ def test_direct_reset_zeroes_row_without_max_budget_even_with_rollover( asyncio.run(reset_budget_job.reset_budget_for_litellm_keys()) - assert _batch_writes(mock_prisma_client, "key")[0]["data"]["spend"] == 0 + assert _batch_writes(mock_prisma_client, "key")[0]["data"]["spend"] == {"decrement": 150.0} def test_budget_cascade_carries_overage_per_tier_when_rollover_enabled( @@ -3071,7 +3075,7 @@ def test_budget_cascade_carries_overage_per_tier_when_rollover_enabled( "where": {"budget_id": "budget-roll", "spend": {"gt": 0, "lte": 10.0}}, "data": {"spend": 0}, } in membership_writes - counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:team_member:member-1:team-1", value=5.0, ttl=60) + counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:team_member:member-1:team-1") def test_budget_cascade_carries_enduser_overage_when_rollover_enabled( @@ -3131,8 +3135,8 @@ def test_budget_cascade_carries_default_tier_enduser_counter_when_rollover_enabl asyncio.run(reset_budget_job.reset_budget_for_litellm_budget_table()) - counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:end_user:enduser-implicit", value=5.0, ttl=60) - counter_cache.redis_cache.async_set_cache.assert_any_await(key="spend:end_user:enduser-implicit", value=5.0, ttl=60) + counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:end_user:enduser-implicit") + counter_cache.redis_cache.async_delete_cache.assert_any_await(key="spend:end_user:enduser-implicit") deleted: Final = {call.kwargs.get("key") for call in counter_cache.user_api_key_cache.async_delete_cache.await_args_list} assert "end_user_id:enduser-implicit" in deleted @@ -3243,3 +3247,138 @@ def test_window_reset_zeroes_counter_when_rollover_disabled(monkeypatch): spend_counter_cache.in_memory_cache.set_cache.assert_any_call(key="spend:key:sk-off:window:1d", value=0.0) spend_counter_cache.async_get_cache.assert_not_awaited() + + +def _apply_spend_payload(db_spend: float, spend_field: dict[str, float]) -> float: + return db_spend - spend_field["decrement"] + + +_RACE_TABLES = [ + ( + lambda job: job.reset_budget_for_litellm_keys(), + "key", + "token", + "tok-race", + lambda now: type( + "Key", + (), + {"spend": 5.0, "budget_duration": "1d", "budget_reset_at": now, "token": "tok-race"}, + ), + ), + ( + lambda job: job.reset_budget_for_litellm_users(), + "user", + "user_id", + "user-race", + lambda now: type( + "User", + (), + {"spend": 5.0, "budget_duration": "7d", "budget_reset_at": now, "user_id": "user-race"}, + ), + ), + ( + lambda job: job.reset_budget_for_litellm_teams(), + "team", + "team_id", + "team-race", + lambda now: type( + "Team", + (), + {"spend": 5.0, "budget_duration": "1mo", "budget_reset_at": now, "team_id": "team-race"}, + ), + ), +] + + +@pytest.mark.parametrize("run_phase, table, id_field, id_value, row_factory", _RACE_TABLES) +def test_reset_decrement_preserves_spend_landed_after_read( + reset_budget_job, mock_prisma_client, run_phase, table, id_field, id_value, row_factory +): + """LIT-7814: spend flushed between the read and the commit survives the reset.""" + now = datetime.now(timezone.utc) + mock_prisma_client.data[table] = [row_factory(now)] + + asyncio.run(run_phase(reset_budget_job)) + + writes = _batch_writes(mock_prisma_client, table) + assert len(writes) == 1 + assert writes[0]["where"] == {id_field: id_value} + assert writes[0]["data"]["spend"] == {"decrement": 5.0} + assert writes[0]["data"]["budget_reset_at"] > now + assert _apply_spend_payload(db_spend=5.4, spend_field=writes[0]["data"]["spend"]) == pytest.approx(0.4) + + +@pytest.mark.parametrize("run_phase, table, id_field, id_value, row_factory", _RACE_TABLES) +def test_reset_decrement_subsumes_rollover_cap( + rollover_enabled, reset_budget_job, mock_prisma_client, run_phase, table, id_field, id_value, row_factory +): + """Rollover on, spend over the cap decrements by the cap itself.""" + now = datetime.now(timezone.utc) + row = row_factory(now) + row.max_budget = 3.0 + mock_prisma_client.data[table] = [row] + + asyncio.run(run_phase(reset_budget_job)) + + writes = _batch_writes(mock_prisma_client, table) + assert len(writes) == 1 + assert writes[0]["data"]["spend"] == {"decrement": 3.0} + assert _apply_spend_payload(db_spend=5.4, spend_field=writes[0]["data"]["spend"]) == pytest.approx(2.4) + + +@pytest.mark.parametrize("run_phase, table, id_field, id_value, row_factory", _RACE_TABLES) +def test_reset_decrement_under_cap_with_rollover( + rollover_enabled, reset_budget_job, mock_prisma_client, run_phase, table, id_field, id_value, row_factory +): + """Rollover on, spend under the cap decrements by the read-time spend.""" + now = datetime.now(timezone.utc) + row = row_factory(now) + row.spend = 2.0 + row.max_budget = 3.0 + mock_prisma_client.data[table] = [row] + + asyncio.run(run_phase(reset_budget_job)) + + writes = _batch_writes(mock_prisma_client, table) + assert len(writes) == 1 + assert writes[0]["data"]["spend"] == {"decrement": 2.0} + assert _apply_spend_payload(db_spend=2.4, spend_field=writes[0]["data"]["spend"]) == pytest.approx(0.4) + + +@pytest.mark.parametrize("run_phase, table, id_field, id_value, row_factory", _RACE_TABLES) +def test_reset_zero_spend_row_writes_noop_decrement( + reset_budget_job, mock_prisma_client, run_phase, table, id_field, id_value, row_factory +): + """A spend=0 row gets a no-op decrement, never an absolute spend=0.""" + now = datetime.now(timezone.utc) + row = row_factory(now) + row.spend = 0.0 + mock_prisma_client.data[table] = [row] + + asyncio.run(run_phase(reset_budget_job)) + + writes = _batch_writes(mock_prisma_client, table) + assert len(writes) == 1 + assert writes[0]["data"]["spend"] == {"decrement": 0.0} + assert writes[0]["data"]["budget_reset_at"] > now + assert _apply_spend_payload(db_spend=0.4, spend_field=writes[0]["data"]["spend"]) == pytest.approx(0.4) + + +def test_reset_deletes_spend_counter_instead_of_seeding(reset_budget_job, mock_prisma_client, monkeypatch): + """A reset deletes the counter so the next read reseeds from the committed row.""" + counter_cache = _make_counter_invalidation_job(monkeypatch) + now = datetime.now(timezone.utc) + mock_prisma_client.data["user"] = [ + type( + "User", + (), + {"spend": 5.0, "budget_duration": "7d", "budget_reset_at": now, "id": "user-r", "user_id": "carol"}, + ) + ] + + asyncio.run(reset_budget_job.reset_budget_for_litellm_users()) + + counter_cache.in_memory_cache.delete_cache.assert_any_call(key="spend:user:carol") + counter_cache.redis_cache.async_delete_cache.assert_any_await(key="spend:user:carol") + counter_cache.in_memory_cache.set_cache.assert_not_called() + counter_cache.redis_cache.async_set_cache.assert_not_awaited() diff --git a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py index 3bc7e1f02f8..6f7ea56db51 100644 --- a/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py +++ b/tests/test_litellm/proxy/db/test_routing_prisma_wrapper.py @@ -145,6 +145,18 @@ def test_writer_pinned_client_yields_to_routed_reads_when_writer_down(): assert pinned.db.litellm_proxymodeltable.find_many is reader_inner.litellm_proxymodeltable.find_many +def test_writer_wrapper_keeps_raw_sql_on_the_writer_while_writer_flagged_down(): + from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper, writer_wrapper + + writer, writer_inner, reader, reader_inner = _make_wrappers() + routing = RoutingPrismaWrapper(writer=writer, reader=reader) + routing._writer_unavailable = True + + assert writer_wrapper(routing).query_raw is writer_inner.query_raw + assert writer_wrapper(routing).query_raw is not reader_inner.query_raw + assert writer_wrapper(writer) is writer + + @pytest.mark.asyncio async def test_connect_invokes_both_clients(): from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper diff --git a/tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py b/tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py new file mode 100644 index 00000000000..919e9c79828 --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_batch_rate_limiter.py @@ -0,0 +1,259 @@ +""" +Tests for `tpd_limit` (tokens per day) enforcement on batch submissions. + +A batch's rows are scheduled by the provider, so a caller cannot keep a large +batch under a per-minute RPM/TPM budget. Scopes that configure `tpd_limit` +are charged against a 24h token window instead of their minute counters. +""" + +from datetime import datetime + +import pytest +from fastapi import HTTPException + +from litellm import DualCache +from litellm.constants import BATCH_TPD_WINDOW_SECONDS +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.hooks.batch_rate_limiter import BatchFileUsage +from litellm.proxy.hooks.parallel_request_limiter_v3 import ( + _PROXY_MaxParallelRequestsHandler_v3, +) +from litellm.proxy.utils import InternalUsageCache, hash_token + + +class _Clock: + def __init__(self, start: datetime): + self.now = start + + def __call__(self) -> datetime: + return self.now + + +def _make_limiters(clock: _Clock | None = None): + internal_usage_cache = InternalUsageCache(dual_cache=DualCache()) + rate_limiter = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=internal_usage_cache, time_provider=clock) + batch_limiter = rate_limiter._get_batch_rate_limiter() + assert batch_limiter is not None + return internal_usage_cache, rate_limiter, batch_limiter + + +async def _counter(internal_usage_cache, rate_limiter, descriptor_key, value, rate_limit_type): + cache_key = rate_limiter.create_rate_limit_keys(descriptor_key, value, rate_limit_type) + raw = await internal_usage_cache.async_get_cache(key=cache_key, litellm_parent_otel_span=None, local_only=True) + return int(raw or 0) + + +@pytest.mark.asyncio +async def test_batch_over_rpm_and_tpm_but_under_tpd_is_accepted(): + internal_usage_cache, rate_limiter, batch_limiter = _make_limiters() + api_key = hash_token("tpd-key") + user_api_key_dict = UserAPIKeyAuth(api_key=api_key, rpm_limit=1, tpm_limit=10, tpd_limit=1000) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=500, request_count=50), + ) + + assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 500 + assert await _counter(internal_usage_cache, rate_limiter, "api_key", api_key, "requests") == 0 + assert await _counter(internal_usage_cache, rate_limiter, "api_key", api_key, "tokens") == 0 + + +@pytest.mark.asyncio +async def test_cumulative_batch_tokens_over_tpd_returns_429_with_remaining_daily_window(): + window_start = datetime(2026, 9, 13, 8, 0, 0) + clock = _Clock(window_start) + _internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters(clock) + user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("tpd-key-2"), rpm_limit=1, tpd_limit=1000) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=600, request_count=6), + ) + clock.now = datetime(2026, 9, 13, 11, 0, 0) + with pytest.raises(HTTPException) as exc: + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=600, request_count=6), + ) + + assert exc.value.status_code == 429 + assert "api_key_tpd" in str(exc.value.detail) + assert "600 tokens but only 400 tokens remaining out of 1000 TPD limit" in str(exc.value.detail) + assert exc.value.headers["retry-after"] == str(BATCH_TPD_WINDOW_SECONDS - 3 * 3600) + assert exc.value.headers["reset_at"] == "2026-09-14 08:00:00 UTC" + + +@pytest.mark.asyncio +async def test_failed_batch_submission_refunds_tpd_tokens(): + internal_usage_cache, rate_limiter, batch_limiter = _make_limiters() + api_key = hash_token("tpd-refund-key") + user_api_key_dict = UserAPIKeyAuth(api_key=api_key, rpm_limit=1, tpd_limit=1000) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=600, request_count=6), + ) + await rate_limiter.async_post_call_failure_hook( + request_data={}, + original_exception=RuntimeError("provider rejected the file"), + user_api_key_dict=user_api_key_dict, + ) + + assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 0 + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=1000, request_count=10), + ) + assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 1000 + + +@pytest.mark.asyncio +async def test_tpd_refund_applies_once_and_only_to_daily_counters(): + internal_usage_cache, rate_limiter, batch_limiter = _make_limiters() + team_key = UserAPIKeyAuth( + api_key=hash_token("tpd-refund-team-key"), + rpm_limit=100, + tpm_limit=10_000, + team_id="team-r", + team_tpd_limit=5000, + ) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=team_key, + data={}, + batch_usage=BatchFileUsage(total_tokens=800, request_count=8), + ) + await rate_limiter.async_post_call_failure_hook( + request_data={}, original_exception=RuntimeError("boom"), user_api_key_dict=team_key + ) + await rate_limiter.async_post_call_failure_hook( + request_data={}, original_exception=RuntimeError("boom"), user_api_key_dict=team_key + ) + + assert await _counter(internal_usage_cache, rate_limiter, "team_tpd", "team-r", "tokens") == 0 + assert await _counter(internal_usage_cache, rate_limiter, "api_key", team_key.api_key, "tokens") == 800 + assert await _counter(internal_usage_cache, rate_limiter, "api_key", team_key.api_key, "requests") == 8 + + +@pytest.mark.asyncio +async def test_rejected_batch_leaves_nothing_to_refund(): + internal_usage_cache, rate_limiter, batch_limiter = _make_limiters() + api_key = hash_token("tpd-rejected-key") + user_api_key_dict = UserAPIKeyAuth(api_key=api_key, tpd_limit=100) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, data={}, batch_usage=BatchFileUsage(total_tokens=90, request_count=9) + ) + with pytest.raises(HTTPException): + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, data={}, batch_usage=BatchFileUsage(total_tokens=20, request_count=2) + ) + await rate_limiter.async_post_call_failure_hook( + request_data={}, original_exception=RuntimeError("429 bubbled up"), user_api_key_dict=user_api_key_dict + ) + + assert await _counter(internal_usage_cache, rate_limiter, "api_key_tpd", api_key, "tokens") == 90 + + +@pytest.mark.asyncio +async def test_batch_without_tpd_still_enforces_minute_rpm(): + _internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters() + user_api_key_dict = UserAPIKeyAuth(api_key=hash_token("rpm-only-key"), rpm_limit=1, tpm_limit=1000) + + with pytest.raises(HTTPException) as exc: + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=user_api_key_dict, + data={}, + batch_usage=BatchFileUsage(total_tokens=50, request_count=5), + ) + + assert exc.value.status_code == 429 + assert "RPM limit" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_team_tpd_replaces_team_minute_limits_but_key_minute_limits_still_apply(): + internal_usage_cache, rate_limiter, batch_limiter = _make_limiters() + team_key = UserAPIKeyAuth( + api_key=hash_token("team-key"), + team_id="team-1", + team_rpm_limit=1, + team_tpm_limit=10, + team_tpd_limit=5000, + ) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=team_key, + data={}, + batch_usage=BatchFileUsage(total_tokens=800, request_count=8), + ) + assert await _counter(internal_usage_cache, rate_limiter, "team_tpd", "team-1", "tokens") == 800 + assert await _counter(internal_usage_cache, rate_limiter, "team", "team-1", "requests") == 0 + + key_rpm_in_team_with_tpd = UserAPIKeyAuth( + api_key=hash_token("team-key-2"), + rpm_limit=1, + team_id="team-1", + team_tpd_limit=5000, + ) + with pytest.raises(HTTPException) as exc: + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=key_rpm_in_team_with_tpd, + data={}, + batch_usage=BatchFileUsage(total_tokens=10, request_count=2), + ) + assert exc.value.status_code == 429 + assert "api_key:" in str(exc.value.detail) + assert "RPM limit" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_end_user_tpd_is_enforced_per_end_user(): + _internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters() + first_customer = UserAPIKeyAuth( + api_key=hash_token("shared-key"), end_user_id="customer-a", end_user_rpm_limit=1, end_user_tpd_limit=100 + ) + second_customer = UserAPIKeyAuth( + api_key=hash_token("shared-key"), end_user_id="customer-b", end_user_rpm_limit=1, end_user_tpd_limit=100 + ) + + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=first_customer, data={}, batch_usage=BatchFileUsage(total_tokens=90, request_count=9) + ) + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=second_customer, data={}, batch_usage=BatchFileUsage(total_tokens=90, request_count=9) + ) + with pytest.raises(HTTPException) as exc: + await batch_limiter._check_and_increment_batch_counters( + user_api_key_dict=first_customer, data={}, batch_usage=BatchFileUsage(total_tokens=20, request_count=2) + ) + assert exc.value.status_code == 429 + assert "end_user_tpd: customer-a" in str(exc.value.detail) + + +def test_tpd_only_key_is_not_skipped_as_having_no_limits(): + _internal_usage_cache, _rate_limiter, batch_limiter = _make_limiters() + descriptors = batch_limiter._create_batch_rate_limit_descriptors( + user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("tpd-only"), tpd_limit=100), + data={}, + ) + assert batch_limiter._has_applicable_batch_rate_limits(descriptors) is True + + +def test_online_descriptors_ignore_tpd_limit(): + _internal_usage_cache, rate_limiter, _batch_limiter = _make_limiters() + api_key = hash_token("online-key") + descriptors = rate_limiter._create_rate_limit_descriptors( + user_api_key_dict=UserAPIKeyAuth(api_key=api_key, rpm_limit=5, tpd_limit=100, team_id="t", team_tpd_limit=9), + data={"model": "gpt-4o"}, + rpm_limit_type=None, + tpm_limit_type=None, + model_has_failures=False, + ) + assert [(d["key"], d["rate_limit"]["window_size"]) for d in descriptors] == [("api_key", rate_limiter.window_size)] diff --git a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py index 395ce68ec54..dfc95db3e14 100644 --- a/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/test_litellm/proxy/hooks/test_proxy_track_cost_callback.py @@ -678,6 +678,149 @@ async def test_update_database_and_spend_counters_preserves_counter_exception_wh proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once() +@pytest.mark.asyncio +async def test_update_database_and_spend_counters_reconciles_reservation_before_db_update(): + call_order: list[str] = [] + proxy_logging_obj = MagicMock() + + async def _update_database(**kwargs): + call_order.append("update_database") + return True + + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=_update_database) + increment_spend_counters = AsyncMock() + budget_reservation = {"reserved_cost": 0.5, "entries": []} + + async def _reconcile(**kwargs): + call_order.append("reconcile") + + with patch( # test-quality-ok: the helper imports reconcile_budget_reservation in its body, no injection seam + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + new_callable=AsyncMock, + side_effect=_reconcile, + ) as mock_reconcile_budget_reservation: + charged = await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=budget_reservation, + ) + + assert charged is True + assert call_order == ["reconcile", "update_database"] + mock_reconcile_budget_reservation.assert_awaited_once_with( + budget_reservation=budget_reservation, + actual_cost=0.2, + finalize=False, + ) + increment_spend_counters.assert_awaited_once() + assert increment_spend_counters.await_args.kwargs["budget_reservation"] is budget_reservation + + +@pytest.mark.asyncio +async def test_update_database_and_spend_counters_releases_reservation_when_db_update_fails_after_early_reconcile(): + proxy_logging_obj = MagicMock() + db_exception = RuntimeError("db unavailable") + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(side_effect=db_exception) + increment_spend_counters = AsyncMock() + budget_reservation = {"reserved_cost": 0.5, "entries": []} + + with ( + patch( # test-quality-ok: the helper imports reconcile_budget_reservation in its body, no injection seam + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + new_callable=AsyncMock, + ) as mock_reconcile_budget_reservation, + patch( # test-quality-ok: _release_budget_reservation imports the release in its body, no injection seam + "litellm.proxy.spend_tracking.budget_reservation.release_budget_reservation", + new_callable=AsyncMock, + ) as mock_release_budget_reservation, + ): + with pytest.raises(RuntimeError) as exc_info: + await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=budget_reservation, + ) + + assert exc_info.value is db_exception + mock_reconcile_budget_reservation.assert_awaited_once_with( + budget_reservation=budget_reservation, + actual_cost=0.2, + finalize=False, + ) + mock_release_budget_reservation.assert_awaited_once_with( + budget_reservation=budget_reservation, + ) + + increment_spend_counters.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_update_database_and_spend_counters_invalidates_reservation_when_early_reconcile_fails(): + proxy_logging_obj = MagicMock() + proxy_logging_obj.db_spend_update_writer.update_database = AsyncMock(return_value=True) + increment_spend_counters = AsyncMock() + budget_reservation = { + "reserved_cost": 0.5, + "entries": [{"counter_key": "spend:key:test_api_key"}], + } + + with ( + patch( # test-quality-ok: the helper imports reconcile_budget_reservation in its body, no injection seam + "litellm.proxy.spend_tracking.budget_reservation.reconcile_budget_reservation", + new_callable=AsyncMock, + side_effect=RuntimeError("redis unavailable"), + ) as mock_reconcile_budget_reservation, + patch( # test-quality-ok: _invalidate_budget_reservation_counters imports it in its body, no injection seam + "litellm.proxy.spend_tracking.budget_reservation.invalidate_budget_reservation_counters", + new_callable=AsyncMock, + ) as mock_invalidate_budget_reservation_counters, + ): + charged = await _update_database_and_spend_counters( + proxy_logging_obj=proxy_logging_obj, + increment_spend_counters=increment_spend_counters, + user_api_key="test_api_key", + user_id="test_user_id", + end_user_id=None, + team_id="test_team_id", + org_id="test_org_id", + kwargs={}, + completion_response=None, + start_time=datetime.now(), + end_time=datetime.now(), + response_cost=0.2, + budget_reservation=budget_reservation, + ) + + assert charged is True + mock_reconcile_budget_reservation.assert_awaited_once() + mock_invalidate_budget_reservation_counters.assert_awaited_once_with( + budget_reservation=budget_reservation, + ) + assert budget_reservation["finalized"] is True + proxy_logging_obj.db_spend_update_writer.update_database.assert_awaited_once() + increment_spend_counters.assert_awaited_once() + + @pytest.mark.asyncio async def test_track_cost_callback_skips_when_no_standard_logging_object(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/management_v1/test_budgets.py b/tests/test_litellm/proxy/management_endpoints/management_v1/test_budgets.py index add2126ac7b..2b438a9d370 100644 --- a/tests/test_litellm/proxy/management_endpoints/management_v1/test_budgets.py +++ b/tests/test_litellm/proxy/management_endpoints/management_v1/test_budgets.py @@ -52,7 +52,7 @@ app.include_router(router) client = TestClient(app) BUDGETS_PATH = f"{MANAGEMENT_V1_PREFIX}/budgets" -SORTABLE = ["budget_id", "created_at", "max_budget", "rpm_limit", "tpm_limit"] +SORTABLE = ["budget_id", "created_at", "max_budget", "rpm_limit", "tpd_limit", "tpm_limit"] def _row(budget_id: str, **overrides: Any) -> dict[str, Any]: @@ -62,6 +62,7 @@ def _row(budget_id: str, **overrides: Any) -> dict[str, Any]: "soft_budget": None, "tpm_limit": None, "rpm_limit": None, + "tpd_limit": None, "budget_duration": "30d", "budget_reset_at": None, "created_at": "2026-07-20T12:00:00+00:00", @@ -123,7 +124,7 @@ def test_returns_flat_rows_in_the_control_plane_envelope(query_raw, as_proxy_adm def test_serves_the_columns_the_budgets_page_renders(query_raw, as_proxy_admin): - _serve(query_raw, [_row("b-1", soft_budget=5.0, budget_reset_at="2026-08-01T00:00:00+00:00")]) + _serve(query_raw, [_row("b-1", soft_budget=5.0, tpd_limit=250000, budget_reset_at="2026-08-01T00:00:00+00:00")]) row = _get().json()["data"][0] @@ -133,12 +134,14 @@ def test_serves_the_columns_the_budgets_page_renders(query_raw, as_proxy_admin): "soft_budget", "tpm_limit", "rpm_limit", + "tpd_limit", "budget_duration", "budget_reset_at", "created_at", "updated_at", } assert row["soft_budget"] == 5.0 + assert row["tpd_limit"] == 250000 assert row["budget_reset_at"].startswith("2026-08-01T00:00:00") diff --git a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py index 4b6815d7552..2f3be61d00f 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_budget_endpoints.py @@ -136,6 +136,21 @@ async def test_update_budget_success(client_and_mocks, monkeypatch): assert body["updated_by"] == "test_user" +@pytest.mark.asyncio +async def test_new_and_update_budget_persist_tpd_limit(client_and_mocks): + client, _, mock_table = client_and_mocks + + resp = client.post("/budget/new", json={"budget_id": "budget_tpd", "tpd_limit": 250000}) + assert resp.status_code == 200, resp.text + assert resp.json()["tpd_limit"] == 250000 + assert mock_table.create.await_args.kwargs["data"]["tpd_limit"] == 250000 + + resp = client.post("/budget/update", json={"budget_id": "budget_tpd", "tpd_limit": 500000}) + assert resp.status_code == 200, resp.text + assert resp.json()["tpd_limit"] == 500000 + assert mock_table.update.await_args.kwargs["data"]["tpd_limit"] == 500000 + + @pytest.mark.asyncio async def test_update_budget_missing_id(client_and_mocks, monkeypatch): client, mock_prisma, mock_table = client_and_mocks diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index bd59a82cbd2..9ce3a6fb4c2 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -743,6 +743,7 @@ _EXPECTED_CUSTOMER = { "max_parallel_requests": None, "tpm_limit": None, "rpm_limit": None, + "tpd_limit": None, "model_max_budget": None, "budget_duration": "30d", "allowed_models": [], diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index c590a24203c..63055872aa1 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -474,6 +474,28 @@ async def test_key_expiration_exact_duration_hours(monkeypatch): ), f"Expected expiration to be approximately 12 hours from creation, got {hours_diff} hours" +@pytest.mark.asyncio +async def test_generate_key_persists_tpd_limit(monkeypatch): + mock_prisma_client = AsyncMock() + mock_prisma_client.insert_data = AsyncMock( + return_value=MagicMock(token="hashed_token_123", litellm_budget_table=None) + ) + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=None) + mock_prisma_client.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[]) + mock_prisma_client.db.litellm_verificationtoken.count = AsyncMock(return_value=0) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client) + + data_json = GenerateKeyRequest(tpd_limit=250000, rpm_limit=5).model_dump(exclude_none=True) + response = await generate_key_helper_fn(request_type="key", **data_json, table_name="key") + + assert response["tpd_limit"] == 250000 + key_insert = mock_prisma_client.insert_data.await_args_list[-1].kwargs + assert key_insert["table_name"] == "key" + assert key_insert["data"]["tpd_limit"] == 250000 + assert key_insert["data"]["rpm_limit"] == 5 + + @pytest.mark.asyncio async def test_key_generation_with_object_permission(monkeypatch): """Ensure /key/generate correctly handles `object_permission` input by @@ -1823,6 +1845,18 @@ async def test_update_key_enable_prompt_caching_folds_into_metadata(flag_value): assert "enable_prompt_caching" not in {k for k in updated if k != "metadata"} +@pytest.mark.asyncio +@pytest.mark.parametrize("tpd_limit", [250000, None]) +async def test_update_key_writes_tpd_limit_as_a_column(tpd_limit): + data = UpdateKeyRequest(key="sk-1", tpd_limit=tpd_limit) + existing_key = LiteLLM_VerificationToken(token="hashed", tpd_limit=1) + + updated = await prepare_key_update_data(data=data, existing_key_row=existing_key) + + assert updated["tpd_limit"] == tpd_limit + assert "rpm_limit" not in updated + + @pytest.mark.asyncio async def test_update_preserves_service_account_id_when_metadata_replaced(): """ diff --git a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py index 1f17bc29f3b..ebbedc6541e 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_team_endpoints.py @@ -675,6 +675,42 @@ async def test_new_team_with_object_permission(mock_db_client, mock_admin_auth): assert "object_permission" not in team_data +@pytest.mark.asyncio +async def test_new_team_persists_tpd_limit(mock_db_client, mock_admin_auth): + mock_db_client.jsonify_team_object = lambda db_data: db_data + mock_db_client.get_data = AsyncMock(return_value=None) + mock_db_client.update_data = AsyncMock(return_value=MagicMock()) + mock_db_client.db = MagicMock() + mock_db_client.db.litellm_modeltable = MagicMock() + mock_db_client.db.litellm_modeltable.create = AsyncMock(return_value=MagicMock(id="model123")) + + team_create_result = MagicMock(team_id="team-tpd") + team_create_result.model_dump.return_value = {"team_id": "team-tpd", "tpd_limit": 250000} + mock_team_create = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_teamtable = MagicMock() + mock_db_client.db.litellm_teamtable.create = mock_team_create + _wire_team_create_tx(mock_db_client) + mock_db_client.db.litellm_teamtable.count = AsyncMock(return_value=0) + mock_db_client.db.litellm_teamtable.update = AsyncMock(return_value=team_create_result) + mock_db_client.db.litellm_usertable = MagicMock() + mock_db_client.db.litellm_usertable.update = AsyncMock(return_value=MagicMock()) + + from fastapi import Request + + from litellm.proxy._types import NewTeamRequest + from litellm.proxy.management_endpoints.team_endpoints import new_team + + await new_team( + data=NewTeamRequest(team_alias="tpd-team", rpm_limit=5, tpd_limit=250000), + http_request=MagicMock(spec=Request), + user_api_key_dict=mock_admin_auth, + ) + + team_data = mock_team_create.call_args.kwargs["data"] + assert team_data["tpd_limit"] == 250000 + assert team_data["rpm_limit"] == 5 + + @pytest.mark.asyncio async def test_new_team_with_mcp_tool_permissions(mock_db_client, mock_admin_auth): """ @@ -7596,6 +7632,48 @@ async def test_update_team_rpm_limit_not_gated_by_user_limit( assert result is not None +@pytest.mark.asyncio +async def test_update_team_persists_tpd_limit(disable_audit_logging_for_mocked_team): + from fastapi import Request + + from litellm.proxy._types import UpdateTeamRequest, UserAPIKeyAuth + from litellm.proxy.management_endpoints.team_endpoints import update_team + + with ( + patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point + "litellm.proxy.proxy_server.prisma_client" + ) as mock_prisma, + patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point + "litellm.proxy.proxy_server.user_api_key_cache" + ) as mock_cache, + patch( # test-quality-ok: proxy_server module global is the endpoint's only injection point + "litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin" + ), + patch( # test-quality-ok: stubs the audit write so the test observes only the team column written + "litellm.proxy.proxy_server.create_audit_log_for_update", new=AsyncMock() + ), + ): + existing_team = MagicMock(team_id="team-tpd", organization_id=None, model_id=None, tpd_limit=None) + existing_team.model_dump.return_value = {"team_id": "team-tpd", "organization_id": None} + mock_prisma.db.litellm_teamtable.find_unique = AsyncMock(return_value=existing_team) + mock_prisma.jsonify_team_object = lambda db_data: db_data + mock_cache.async_get_cache = AsyncMock(return_value=None) + mock_cache.async_set_cache = AsyncMock() + updated_team = MagicMock(team_id="team-tpd", organization_id=None, litellm_model_table=None) + updated_team.model_dump.return_value = {"team_id": "team-tpd", "organization_id": None, "tpd_limit": 250000} + mock_prisma.db.litellm_teamtable.update = AsyncMock(return_value=updated_team) + + await update_team( + data=UpdateTeamRequest(team_id="team-tpd", tpd_limit=250000), + http_request=MagicMock(spec=Request), + user_api_key_dict=UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN, user_id="admin"), + ) + + written = mock_prisma.db.litellm_teamtable.update.call_args.kwargs["data"] + assert written["tpd_limit"] == 250000 + assert "rpm_limit" not in written + + @pytest.mark.asyncio async def test_new_team_org_scoped_tpm_exceeds_org_limit(): """ diff --git a/tests/test_litellm/proxy/management_helpers/test_access_group_key_sync.py b/tests/test_litellm/proxy/management_helpers/test_access_group_key_sync.py index 60c36e33e09..9b379dbe330 100644 --- a/tests/test_litellm/proxy/management_helpers/test_access_group_key_sync.py +++ b/tests/test_litellm/proxy/management_helpers/test_access_group_key_sync.py @@ -11,14 +11,15 @@ from litellm.proxy.management_helpers.access_group_key_sync import ( ) -def _routed_prisma_client(): +def _routed_prisma_client(writer_unavailable: bool = False): writer_inner = MagicMock(name="writer_prisma") reader_inner = MagicMock(name="reader_prisma") writer_inner.query_raw = AsyncMock(return_value=[]) - reader_inner.query_raw = AsyncMock(return_value=[]) + reader_inner.query_raw = AsyncMock(side_effect=RuntimeError("cannot execute UPDATE in a read-only transaction")) writer = PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False) reader = PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False) routing = RoutingPrismaWrapper(writer=writer, reader=reader) + routing._writer_unavailable = writer_unavailable return SimpleNamespace(db=routing), writer_inner, reader_inner @@ -39,6 +40,39 @@ async def test_regeneration_repoint_update_runs_on_the_writer(): reader_inner.query_raw.assert_not_awaited() +@pytest.mark.asyncio +async def test_regeneration_repoint_update_stays_on_the_writer_while_writer_flagged_unavailable(): + prisma_client, writer_inner, reader_inner = _routed_prisma_client(writer_unavailable=True) + + await sync_key_regeneration_access_group_membership( + prisma_client=prisma_client, + previous_key_token="old-token", + new_key_token="new-token", + data=None, + existing_key_row=MagicMock(), + ) + + writer_inner.query_raw.assert_awaited_once() + assert writer_inner.query_raw.await_args.args[0].startswith('UPDATE "LiteLLM_AccessGroupTable"') + assert writer_inner.query_raw.await_args.args[1:] == ("old-token", "new-token") + reader_inner.query_raw.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_membership_attach_and_detach_updates_stay_on_the_writer_while_writer_flagged_unavailable(): + prisma_client, writer_inner, reader_inner = _routed_prisma_client(writer_unavailable=True) + + await sync_key_access_group_membership( + prisma_client=prisma_client, + key_token="token", + previous_access_group_ids=["ag-old"], + updated_access_group_ids=["ag-new"], + ) + + assert writer_inner.query_raw.await_count == 2 + reader_inner.query_raw.assert_not_awaited() + + @pytest.mark.asyncio async def test_membership_attach_and_detach_updates_run_on_the_writer(): prisma_client, writer_inner, reader_inner = _routed_prisma_client() diff --git a/tests/test_litellm/proxy/management_helpers/test_access_group_model_sync.py b/tests/test_litellm/proxy/management_helpers/test_access_group_model_sync.py index 65ef2d55cb8..c7ce97894d4 100644 --- a/tests/test_litellm/proxy/management_helpers/test_access_group_model_sync.py +++ b/tests/test_litellm/proxy/management_helpers/test_access_group_model_sync.py @@ -13,7 +13,7 @@ from litellm.proxy.management_helpers.access_group_model_sync import ( _INVALIDATE = "litellm.proxy.management_helpers.access_group_model_sync.invalidate_access_group_caches" -def _routed_prisma_client(deployment_count: int): +def _routed_prisma_client(deployment_count: int, writer_unavailable: bool = False): async def query_raw(sql, *params): if sql.startswith("SELECT COUNT(*)"): return [{"deployment_count": deployment_count}] @@ -26,6 +26,7 @@ def _routed_prisma_client(deployment_count: int): writer = PrismaWrapper(original_prisma=writer_inner, iam_token_db_auth=False) reader = PrismaWrapper(original_prisma=reader_inner, iam_token_db_auth=False) routing = RoutingPrismaWrapper(writer=writer, reader=reader) + routing._writer_unavailable = writer_unavailable return SimpleNamespace(db=routing), writer_inner, reader_inner @@ -53,6 +54,20 @@ async def test_rename_replaces_the_old_name_when_no_other_deployment_carries_it( reader_inner.query_raw.assert_not_awaited() +@pytest.mark.asyncio +async def test_rename_update_stays_on_the_writer_while_writer_flagged_unavailable(): + prisma_client, writer_inner, reader_inner = _routed_prisma_client(deployment_count=0, writer_unavailable=True) + + with patch(_INVALIDATE, new=AsyncMock()): + await sync_access_groups_for_renamed_model( + prisma_client, model_id="m-1", old_name="gpt-5.6", new_name="gpt-5.6-eu", llm_router=None + ) + + (update_call,) = _access_group_updates(writer_inner) + assert update_call.args[1:] == ("gpt-5.6", "gpt-5.6-eu") + reader_inner.query_raw.assert_not_awaited() + + @pytest.mark.asyncio async def test_rename_appends_the_new_name_when_a_sibling_row_keeps_the_old_one(): prisma_client, writer_inner, _ = _routed_prisma_client(deployment_count=1) @@ -168,3 +183,17 @@ async def test_delete_keeps_the_name_while_a_sibling_row_still_backs_it(): assert _access_group_updates(writer_inner) == [] invalidate.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_delete_counts_backing_rows_on_the_writer_not_a_lagging_replica_while_writer_flagged_unavailable(): + prisma_client, writer_inner, reader_inner = _routed_prisma_client(deployment_count=0, writer_unavailable=True) + reader_inner.query_raw = AsyncMock(return_value=[{"deployment_count": 1}]) + + with patch(_INVALIDATE, new=AsyncMock()) as invalidate: + await sync_access_groups_for_deleted_model(prisma_client, model_id="m-1", model_name="gpt-5.6", llm_router=None) + + (update_call,) = _access_group_updates(writer_inner) + assert "array_remove" in update_call.args[0] + invalidate.assert_awaited_once_with(("ag-1", "ag-2")) + reader_inner.query_raw.assert_not_awaited() diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py index dd9fbd9161f..e91b7ef970c 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler.py @@ -7,6 +7,7 @@ import pytest import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.proxy.pass_through_endpoints.llm_provider_handlers.vertex_passthrough_logging_handler import ( VertexPassthroughLoggingHandler, ) @@ -128,3 +129,104 @@ def test_vertex_generate_content_payload_prices_gemini_urls_at_gemini_rates(): assert result["kwargs"]["response_cost"] == pytest.approx(GEMINI_COST) assert logging_obj.model_call_details["custom_llm_provider"] == "gemini" + + +def _interrupted_anthropic_stream(model: str, output_text: str) -> list[bytes]: + def sse(event: str, data: dict) -> bytes: + return f"event: {event}\ndata: {json.dumps(data)}\n\n".encode() + + message_start = { + "type": "message_start", + "message": { + "id": "msg_interrupted", + "type": "message", + "role": "assistant", + "model": model, + "content": [], + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 29, "output_tokens": 2}, + }, + } + block_start = {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}} + delta = {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": output_text}} + return [ + sse("message_start", message_start), + sse("content_block_start", block_start), + sse("content_block_delta", delta), + ] + + +@pytest.mark.asyncio +async def test_interrupted_anthropic_stream_recovers_output_tokens_off_the_event_loop(): + from unittest.mock import AsyncMock + + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + model = "claude-fable-5" + warm_tokenizer(model) + logging_obj = _logging_obj() + logging_obj.model_call_details = {"model": model, "stream": True} + logging_obj.litellm_params = {} + logging_obj.get_router_model_id.return_value = None + logging_obj.dispatch_success_handlers = AsyncMock() + + _, took, lags = await timed_with_loop_lags( + lambda: PassThroughStreamingHandler._route_streaming_logging_to_handler( + litellm_logging_obj=logging_obj, + passthrough_success_handler_obj=PassThroughEndpointLogging(), + url_route="/anthropic/v1/messages", + request_body={"model": model, "stream": True}, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + raw_bytes=_interrupted_anthropic_stream(model, text * 100), + end_time=datetime.now(), + model=model, + ) + ) + + logging_obj.dispatch_success_handlers.assert_awaited_once() + logged_usage = logging_obj.dispatch_success_handlers.await_args.kwargs["result"].usage + assert logged_usage.completion_tokens > 100_000 + assert_loop_stayed_free(took, lags) + + +@pytest.mark.asyncio +async def test_failed_anthropic_stream_records_partial_usage_off_the_event_loop(): + from unittest.mock import AsyncMock + + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + model = "claude-fable-5" + warm_tokenizer(model) + logging_obj = _logging_obj() + logging_obj.model_call_details = {"model": model, "stream": True} + logging_obj.litellm_params = {} + logging_obj.get_router_model_id.return_value = None + logging_obj.dispatch_failure_handlers = AsyncMock() + + _, took, lags = await timed_with_loop_lags( + lambda: PassThroughStreamingHandler.schedule_stream_failure_logging( + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + request_body={"model": model, "stream": True}, + raw_bytes=_interrupted_anthropic_stream(model, text * 100), + exception=RuntimeError("upstream closed the stream"), + ) + ) + await GLOBAL_LOGGING_WORKER.flush() + + logging_obj.dispatch_failure_handlers.assert_awaited_once() + partial_usage = logging_obj.record_partial_usage_for_failure.call_args.kwargs["usage"] + assert partial_usage.completion_tokens > 100_000 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py index 68462065393..0731c233fef 100644 --- a/tests/test_litellm/proxy/proxy_server/test_spend_counters.py +++ b/tests/test_litellm/proxy/proxy_server/test_spend_counters.py @@ -921,6 +921,30 @@ async def test_reconcile_budget_reservation_for_counter_update_failure_invalidat assert fake_invalidate.called is True +@pytest.mark.asyncio +async def test_reconcile_budget_reservation_for_counter_update_finalized_reservation_falls_back_to_direct_increment( + monkeypatch, +): + """A reservation already finalized before the counter update (the pre-persist + reconcile failed and dropped its counters) must not shield its keys from the + direct increment, or the settled cost is never added back after the drop.""" + import litellm.proxy.spend_tracking.budget_reservation as br + + fake_reconcile = AsyncMock() + monkeypatch.setattr(br, "reconcile_budget_reservation", fake_reconcile) + + result = await ps._reconcile_budget_reservation_for_counter_update( + budget_reservation={ + "finalized": True, + "entries": [{"counter_key": "spend:key:abc"}], + }, + response_cost=1.0, + ) + + assert result == set() + fake_reconcile.assert_not_awaited() + + # --------------------------------------------------------------------------- # _prepare_end_user_and_tag_spend_increments # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 8283ee8395a..60bff50f000 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -675,6 +675,7 @@ ignored_keys = [ "metadata.user_api_key_team_alias", "metadata.spend_logs_metadata", "metadata.requester_ip_address", + "metadata.user_agent", "metadata.status", "metadata.proxy_server_request", "metadata.error_information", diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index a72b4e28143..8b105e94d19 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2935,6 +2935,16 @@ def test_get_spend_logs_metadata_keeps_master_key_alias_readable(): assert meta["user_api_key"] == LITELLM_PROXY_MASTER_KEY_ALIAS +def test_get_spend_logs_metadata_keeps_user_agent(): + """`add_litellm_data_to_request` stamps the caller's User-Agent next to its IP, but + the spend log metadata dropped it, so an abusive client could not be identified + from the Logs page.""" + meta = _get_spend_logs_metadata({"requester_ip_address": "203.0.113.9", "user_agent": "abusive-client/9.9"}) + assert meta["requester_ip_address"] == "203.0.113.9" + assert meta["user_agent"] == "abusive-client/9.9" + assert _get_spend_logs_metadata(None)["user_agent"] is None + + def test_redact_logged_api_key_bearer_only_returns_none(): # "bearer " with nothing after stripping is equivalent to no key assert _redact_logged_api_key("bearer ") is None diff --git a/tests/test_litellm/proxy/test_budget_reservation.py b/tests/test_litellm/proxy/test_budget_reservation.py index 40ebc03781c..032722d3259 100644 --- a/tests/test_litellm/proxy/test_budget_reservation.py +++ b/tests/test_litellm/proxy/test_budget_reservation.py @@ -2230,6 +2230,17 @@ class _ExpiringRedisCache: return None +class _TeamMembershipFloorDb: + """Stands in for `prisma_client.db`: only the team-membership row exists and its spend is the DB floor.""" + + def __init__(self, spend: float) -> None: + self.spend = spend + + def __getattr__(self, table_name: str) -> SimpleNamespace: + row = SimpleNamespace(spend=self.spend) if table_name == "litellm_teammembership" else None + return SimpleNamespace(find_unique=AsyncMock(return_value=row)) + + @pytest.mark.asyncio async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( spend_counter_state, @@ -2275,6 +2286,59 @@ async def test_reconcile_after_redis_counter_expiry_keeps_request_cost_enforced( assert reservation["finalized"] is True +@pytest.mark.asyncio +async def test_reconcile_before_db_update_does_not_double_count_when_flush_lands_between_passes( + spend_counter_state, +): + """The early reconcile (before the spend row is enqueued) reseeds from a DB + floor that cannot yet include this request. When the periodic flush commits + the row before increment_spend_counters runs its second reconcile, the + applied_adjustment early-return must keep the counter from adding the cost + a second time.""" + import litellm.proxy.proxy_server as ps + from litellm.proxy.spend_tracking.budget_reservation import reconcile_budget_reservation + + counter_cache, _ = spend_counter_state + counter_key = "spend:team_member:user-flush:team-flush" + redis_cache = _ExpiringRedisCache() + counter_cache.redis_cache = redis_cache + counter_cache.in_memory_cache.set_cache(key=counter_key, value=0.6) + db_floor = _TeamMembershipFloorDb(spend=0.3) + ps.prisma_client = SimpleNamespace(db=db_floor) + + reservation = { + "reserved_cost": 0.6, + "entries": [ + { + "counter_key": counter_key, + "entity_type": "TeamMember", + "entity_id": "user-flush:team-flush", + "reserved_cost": 0.6, + "applied_adjustment": 0.0, + } + ], + "finalized": False, + } + + await reconcile_budget_reservation(budget_reservation=reservation, actual_cost=0.05, finalize=False) + + assert redis_cache.store[counter_key] == pytest.approx(0.35) + assert reservation["entries"][0]["applied_adjustment"] == pytest.approx(-0.55) + assert reservation["finalized"] is False + + db_floor.spend = 0.35 + await ps.increment_spend_counters( + token="key-flush", + team_id="team-flush", + user_id="user-flush", + response_cost=0.05, + budget_reservation=reservation, + ) + + assert redis_cache.store[counter_key] == pytest.approx(0.35) + assert reservation["finalized"] is True + + @pytest.mark.asyncio async def test_should_invalidate_reserved_counters_after_persisted_spend_failure( spend_counter_state, diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py index 8b2b6b4c6ca..d7a6124dd97 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_post_call_failure_hook.py @@ -4,13 +4,14 @@ and ``_handle_logging_proxy_only_error``.""" from __future__ import annotations import asyncio -from typing import Any +from datetime import datetime from unittest.mock import AsyncMock, MagicMock import pytest from fastapi import HTTPException import litellm +from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_logger import CustomLogger from litellm.proxy._types import AlertType, ProxyErrorTypes from litellm.proxy.utils import ProxyLogging @@ -47,12 +48,17 @@ def test_is_proxy_only_llm_api_truth_table(proxy_logging): error_type=ProxyErrorTypes.auth_error, route="/chat/completions", ), + "guardrail_raised_on_llm_route": proxy_logging._is_proxy_only_llm_api_error( + original_exception=GuardrailRaisedException(guardrail_name="g", message="blocked"), + route="/chat/completions", + ), } assert snapshot == { "no_route": False, "non_llm_route": False, "http_on_llm_route": True, "auth_short_circuit": True, + "guardrail_raised_on_llm_route": True, } @@ -318,3 +324,50 @@ async def test_post_call_failure_hook_keeps_the_route_for_multi_operation_routes route=route, ) assert request_data["call_type"] == route + + +@pytest.mark.asyncio +async def test_post_call_failure_hook_guardrail_block_fires_failure_callback( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """A ``GuardrailRaisedException`` on an LLM route must reach the logging + object's ``async_failure_handler`` so custom loggers see a ``failure`` + status - without this, guardrail blocks produce only + ``post_call_failure_hook`` and no failure logging event.""" + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + + recorded: list[object] = [] + + class _StatusRecorder(CustomLogger): + async def async_log_failure_event( + self, kwargs: dict[str, object], response_obj: object, start_time: datetime, end_time: datetime + ) -> None: + standard_logging_object = kwargs.get("standard_logging_object") + recorded.append(standard_logging_object.get("status") if isinstance(standard_logging_object, dict) else None) + + monkeypatch.setattr(litellm, "_async_failure_callback", [_StatusRecorder()]) + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id="test_guardrail_block_failure_cb", + function_id="test_guardrail_block_failure_cb", + ) + request_data = { + "litellm_logging_obj": logging_obj, + "litellm_call_id": "test_guardrail_block_failure_cb", + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hi"}], + "metadata": {}, + } + proxy_logging.alert_types = [] + await proxy_logging.post_call_failure_hook( + request_data=request_data, + original_exception=GuardrailRaisedException(guardrail_name="g", message="blocked"), + user_api_key_dict=make_user_api_key_auth(request_route="/chat/completions"), + ) + await asyncio.sleep(0) + await asyncio.sleep(0) + assert recorded == ["failure"] diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py index ec5b994f147..ebc831b4102 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_streaming_hooks.py @@ -10,6 +10,7 @@ Covers ``_wrap_streaming_iterator_with_enrichment``, from __future__ import annotations import asyncio +from collections.abc import AsyncGenerator, AsyncIterator from datetime import datetime from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock @@ -18,12 +19,15 @@ import pytest from fastapi import HTTPException import litellm +from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.experimental_pass_through.messages.streaming_iterator import ( BaseAnthropicMessagesStreamingIterator, ) +from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ProxyLogging +from litellm.types.utils import Usage @pytest.fixture(autouse=True) @@ -479,6 +483,135 @@ async def test_native_messages_stream_logging_fires_when_guardrail_blocks_after_ assert logging_obj._deferred_stream_complete_args is None +def _armed_chat_stream( + test_name: str, request_data: dict[str, object], events: list[str] +) -> tuple[LiteLLMLoggingObj, AsyncIterator[dict[str, object]]]: + """A /chat/completions stream whose CSW shape parks ``(assembled ModelResponse, cache_hit)`` + at upstream exhaustion, with the deferred dispatch recording into ``events``.""" + logging_obj = LiteLLMLoggingObj( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + stream=True, + call_type="acompletion", + start_time=datetime.now(), + litellm_call_id=test_name, + function_id=test_name, + ) + logging_obj.optional_params = {} + logging_obj.litellm_params = {} + logging_obj.standard_built_in_tools_params = None + + async def _dispatch_deferred_logging(*args: object) -> None: + events.append("success_dispatched") + + logging_obj._on_deferred_stream_complete = _dispatch_deferred_logging + request_data["litellm_logging_obj"] = logging_obj + + assembled = litellm.ModelResponse( + model="gpt-4o-mini", + choices=[{"index": 0, "message": {"role": "assistant", "content": "BANANA"}}], + usage=Usage(prompt_tokens=3, completion_tokens=5, total_tokens=8), + ) + + async def _upstream() -> AsyncIterator[dict[str, object]]: + yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "BAN"}}]} + yield {"id": "c1", "choices": [{"index": 0, "delta": {"content": "ANA"}}]} + logging_obj._deferred_stream_complete_args = (assembled, False) + + return logging_obj, _upstream() + + +def _raising_at_end_of_stream(error: Exception) -> CustomLogger: + class _EndOfStreamRaiser(CustomLogger): + async def async_post_call_streaming_iterator_hook( + self, user_api_key_dict: UserAPIKeyAuth, response: AsyncIterator[object], request_data: dict[str, object] + ) -> AsyncGenerator[object, None]: + async for chunk in response: + yield chunk + raise error + + return _EndOfStreamRaiser() + + +@pytest.mark.asyncio +async def test_chat_stream_guardrail_block_after_stream_end_logs_failure_not_success( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """ + A guardrail that raises ``GuardrailRaisedException`` at end of a + /chat/completions stream must NOT dispatch the parked success logging: + the request is logged via the failure path instead, with the consumed + usage carried over so the failure row bills correctly. + """ + events: list[str] = [] + request_data: dict[str, object] = {"metadata": {}} + logging_obj, upstream = _armed_chat_stream("test_chat_stream_guardrail_block", request_data, events) + monkeypatch.setattr( + litellm, + "callbacks", + [_raising_at_end_of_stream(GuardrailRaisedException(guardrail_name="g", message="blocked"))], + ) + + with pytest.raises(GuardrailRaisedException): + async for _ in proxy_logging.async_post_call_streaming_iterator_hook( + response=upstream, + user_api_key_dict=make_user_api_key_auth(), + request_data=request_data, + ): + pass + await asyncio.sleep(0) + await asyncio.sleep(0) + + snapshot = { + "events": events, + "callback_cleared": logging_obj._on_deferred_stream_complete is None, + "args_cleared": logging_obj._deferred_stream_complete_args is None, + "combined_usage_total_tokens": logging_obj.model_call_details["combined_usage_object"].total_tokens, + "response_cost_positive": logging_obj.model_call_details["response_cost"] > 0, + } + assert snapshot == { + "events": [], + "callback_cleared": True, + "args_cleared": True, + "combined_usage_total_tokens": 8, + "response_cost_positive": True, + } + + +@pytest.mark.asyncio +async def test_chat_stream_generic_callback_error_after_stream_end_still_flushes_success_logging( + proxy_logging, make_user_api_key_auth, monkeypatch +): + """ + ``post_call_failure_hook`` only routes proxy-level errors (HTTPException, + ProxyException, GuardrailRaisedException) through failure logging. A + callback that dies with any other exception after the stream completed + must keep flushing the parked success dispatch, or the request ends with + no terminal log at all. + """ + events: list[str] = [] + request_data: dict[str, object] = {"metadata": {}} + logging_obj, upstream = _armed_chat_stream("test_chat_stream_generic_callback_error", request_data, events) + monkeypatch.setattr(litellm, "callbacks", [_raising_at_end_of_stream(RuntimeError("callback crashed"))]) + + with pytest.raises(RuntimeError): + async for _ in proxy_logging.async_post_call_streaming_iterator_hook( + response=upstream, + user_api_key_dict=make_user_api_key_auth(), + request_data=request_data, + ): + pass + await asyncio.sleep(0) + await asyncio.sleep(0) + + snapshot = { + "events": events, + "args_cleared": logging_obj._deferred_stream_complete_args is None, + "failure_usage_recorded": "combined_usage_object" in logging_obj.model_call_details, + } + assert snapshot == {"events": ["success_dispatched"], "args_cleared": True, "failure_usage_recorded": False} + + # --------------------------------------------------------------------------- # _fire_deferred_stream_logging # --------------------------------------------------------------------------- diff --git a/tests/test_litellm/repositories/test_unit_of_work.py b/tests/test_litellm/repositories/test_unit_of_work.py index 1a76b537e95..b52b8ced31e 100644 --- a/tests/test_litellm/repositories/test_unit_of_work.py +++ b/tests/test_litellm/repositories/test_unit_of_work.py @@ -46,16 +46,16 @@ async def test_updates_across_tables_share_one_batch_and_commit_once(): reset_at = datetime(2026, 8, 3, 12, 0, tzinfo=timezone.utc) async with spend_reset_unit_of_work(lambda: batch) as uow: - uow.keys.queue_spend_reset(token="tok-1", budget_reset_at=reset_at) - uow.users.queue_spend_reset(user_id="user-1", budget_reset_at=reset_at) - uow.teams.queue_spend_reset(team_id="team-1", budget_reset_at=None) + uow.keys.queue_spend_reset(token="tok-1", budget_reset_at=reset_at, spend_decrement=1.5) + uow.users.queue_spend_reset(user_id="user-1", budget_reset_at=reset_at, spend_decrement=2.5) + uow.teams.queue_spend_reset(team_id="team-1", budget_reset_at=None, spend_decrement=0.0) assert batch.commit_count == 0 assert batch.commit_count == 1 assert batch.calls == [ - ("litellm_verificationtoken", {"token": "tok-1"}, {"spend": 0, "budget_reset_at": reset_at}), - ("litellm_usertable", {"user_id": "user-1"}, {"spend": 0, "budget_reset_at": reset_at}), - ("litellm_teamtable", {"team_id": "team-1"}, {"spend": 0, "budget_reset_at": None}), + ("litellm_verificationtoken", {"token": "tok-1"}, {"spend": {"decrement": 1.5}, "budget_reset_at": reset_at}), + ("litellm_usertable", {"user_id": "user-1"}, {"spend": {"decrement": 2.5}, "budget_reset_at": reset_at}), + ("litellm_teamtable", {"team_id": "team-1"}, {"spend": {"decrement": 0.0}, "budget_reset_at": None}), ] @@ -64,7 +64,7 @@ async def test_raising_inside_block_skips_commit(): async def _blow_up_mid_transaction(): async with spend_reset_unit_of_work(lambda: batch) as uow: - uow.keys.queue_spend_reset(token="tok-1", budget_reset_at=None) + uow.keys.queue_spend_reset(token="tok-1", budget_reset_at=None, spend_decrement=0.0) raise RuntimeError("boom") with pytest.raises(RuntimeError, match="boom"): diff --git a/tests/test_litellm/router_strategy/test_llm_v2.py b/tests/test_litellm/router_strategy/test_llm_v2.py index a92db26c77c..de23987a242 100644 --- a/tests/test_litellm/router_strategy/test_llm_v2.py +++ b/tests/test_litellm/router_strategy/test_llm_v2.py @@ -4,6 +4,7 @@ from typing import Final from unittest.mock import AsyncMock, MagicMock import pytest +import litellm from pydantic import ValidationError from litellm import ModelResponse, Router @@ -11,6 +12,7 @@ from litellm.caching.dual_cache import DualCache from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig, ComplexityTier from litellm.router_strategy.complexity_router.llm_v2 import ( + LLM_V2_PROMPT_VERSION, LLMV2Calibration, LLMV2Config, LLMV2ProbabilityCalibration, @@ -39,6 +41,31 @@ def test_learned_policy_overrides_raw_gap_without_rewriting_raw_probabilities() assert "selective:target=rescue" in decision.signals +@pytest.mark.asyncio +async def test_learned_policy_metadata_does_not_claim_unused_raw_gap_threshold() -> None: + base: Final = _config().llm_v2_config + assert base is not None + policy: Final = SelectivePolicy( + version="test-rescue", + feature_schema="v2-v1", + target="rescue", + threshold=0.05, + heads=(SelectiveHead(constant=1),), + ) + config: Final = _config(llm_v2_config={**base.model_dump(), "selective_policy": policy.model_dump()}) + router, _ = _router(_verdict().model_dump_json(), config) + result: Final = await router.async_pre_routing_hook( + model="v2-router", messages=[{"role": "user", "content": "Fix nested behavior"}], request_kwargs={} + ) + assert result is not None and result.model == "capable" + decision: Final = result.routing_decision + assert decision is not None + assert decision["classifier_efficient_p_solve"] == 0.9 + assert decision["classifier_capable_p_solve"] == 0.92 + assert "classifier_max_quality_gap" not in decision + assert "selective:target=rescue" in decision["signals"] + + def test_learned_policy_rejects_another_classifier_feature_schema() -> None: config: Final = _config().llm_v2_config assert config is not None @@ -220,6 +247,7 @@ def test_verdict_rejects_invalid_probabilities(probability: object) -> None: ({"classifier_type": "heuristic"}, "requires classifier_type llm_v2"), ({"classifier_llm_config": None}, "classifier_llm_config is required"), ({"adaptive": True}, "adaptive=false"), + ({"classifier_fallback": "default_model", "default_model": "efficient"}, "fails closed"), ({"tiers": {"SIMPLE": ["same"], "REASONING": ["same"]}}, "distinct model"), ({"tiers": {"SIMPLE": ["a", "b"], "REASONING": ["c"]}}, "one distinct model"), ({"tiers": {"SIMPLE": ["a"], "MEDIUM": ["b"], "REASONING": ["c"]}}, "exactly"), @@ -300,13 +328,100 @@ async def test_json_object_mode_supplies_schema_in_prompt() -> None: @pytest.mark.asyncio -@pytest.mark.parametrize("content", ["", "not json", '{"tier":"SIMPLE"}', '{"forecasts":{}}']) +@pytest.mark.parametrize("mode", ("json_schema", "json_object")) +@pytest.mark.parametrize("fence", ("```json", "```")) +async def test_fenced_forecast_routes_by_validated_probabilities(mode: str, fence: str) -> None: + base: Final = _config().llm_v2_config + assert base is not None + config: Final = _config(llm_v2_config={**base.model_dump(), "response_format": mode}) + content: Final = f" {fence}\n{_verdict().model_dump_json()}\n``` " + router, client = _router(content, config) + result: Final = await router.async_pre_routing_hook( + model="v2-router", messages=[{"role": "user", "content": "Fix nested behavior"}], request_kwargs={} + ) + assert result is not None and result.model == "efficient" + assert result.routing_decision is not None + assert result.routing_decision["cause"] == "llm_v2_classifier" + assert result.routing_decision["classifier_efficient_p_solve"] == 0.9 + assert result.routing_decision["classifier_capable_p_solve"] == 0.92 + assert result.routing_decision["classifier_cost"] == 0.001 + client.acompletion.assert_awaited_once() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("user_agent", ("claude-cli/2.1.233", "curl/8.7.1")) +@pytest.mark.parametrize("metadata_key", ("metadata", "litellm_metadata")) +async def test_caller_constraints_respect_claude_code_prompt_policy(user_agent: str, metadata_key: str) -> None: + router, client = _router(_verdict().model_dump_json()) + outcome: Final = await router.aclassify( + "Fix nested behavior", "Caller system context", request_kwargs={metadata_key: {"user_agent": user_agent}} + ) + assert outcome.cause == "llm_v2_classifier" + call: Final = client.acompletion.call_args.kwargs + payload: Final = json.loads(call["messages"][1]["content"]) + assert payload["caller_constraints"] == (None if user_agent.startswith("claude") else "Caller system context") + assert payload["task_and_follow_ups"] == ["Fix nested behavior"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("calibrated", (False, True)) +async def test_routing_metadata_preserves_exact_forecasts_and_redaction( + calibrated: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + base: Final = _config().llm_v2_config + assert base is not None + calibration: Final = LLMV2Calibration( + version="test-pair-v1", + prompt_version=LLM_V2_PROMPT_VERSION, + efficient=LLMV2ProbabilityCalibration(slope=0.2, intercept=-1.0), + capable=LLMV2ProbabilityCalibration(slope=1.0, intercept=0.0), + ) + policy: Final = base.model_copy(update={"calibration": calibration if calibrated else None}) + verdict: Final = _verdict(0.900000123, 0.920000321) + router, _ = _router(verdict.model_dump_json(), _config(llm_v2_config=policy.model_dump())) + result: Final = await router.async_pre_routing_hook( + model="v2-router", messages=[{"role": "user", "content": "Fix nested behavior"}], request_kwargs={} + ) + assert result is not None + assert result.model == ("capable" if calibrated else "efficient") + decision: Final = result.routing_decision + assert decision is not None + monkeypatch.setattr(litellm, "turn_off_message_logging", True) + redacted: Final = Router._redact_prompt_text_if_needed(request_kwargs={}, routing_decision=decision) + assert redacted is not None + assert "signals" not in redacted + for record in (decision, redacted): + assert record["classifier_efficient_p_solve"] == 0.900000123 + assert record["classifier_capable_p_solve"] == 0.920000321 + assert record["classifier_max_quality_gap"] == 0.05 + assert record["classifier_prompt_version"] == LLM_V2_PROMPT_VERSION + if calibrated: + assert record["classifier_calibration_version"] == "test-pair-v1" + assert record["classifier_calibrated_efficient_p_solve"] == calibration.efficient.calibrate(0.900000123) + assert record["classifier_calibrated_capable_p_solve"] == calibration.capable.calibrate(0.920000321) + else: + assert "classifier_calibration_version" not in record + assert "classifier_calibrated_efficient_p_solve" not in record + assert "classifier_calibrated_capable_p_solve" not in record + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "content", ["", "not json", '{"tier":"SIMPLE"}', '{"forecasts":{}}', '```json\n{"forecasts":{}}\n```'] +) async def test_invalid_output_falls_back_to_capable_and_preserves_paid_call_cost(content: str) -> None: router, client = _router(content) - outcome: Final = await router.aclassify("hi") - assert outcome.tier == ComplexityTier.REASONING - assert outcome.cause == "llm_v2_fallback" - assert outcome.classifier_cost == 0.001 + result: Final = await router.async_pre_routing_hook( + model="v2-router", messages=[{"role": "user", "content": "hi"}], request_kwargs={} + ) + assert result is not None and result.model == "capable" + decision: Final = result.routing_decision + assert decision is not None + assert decision["cause"] == "llm_v2_fallback" + assert decision["classifier_cost"] == 0.001 + assert "classifier_efficient_p_solve" not in decision + assert "classifier_capable_p_solve" not in decision + assert "classifier_prompt_version" not in decision client.acompletion.assert_awaited_once() @@ -323,6 +438,17 @@ async def test_timeout_falls_back_to_capable_and_opens_shared_breaker() -> None: client.acompletion.assert_awaited_once() +@pytest.mark.asyncio +async def test_provider_failure_redacts_prompt_text_from_warning(caplog: pytest.LogCaptureFixture) -> None: + router, client = _router("") + client.acompletion.side_effect = ValueError("private task text from provider") + outcome: Final = await router.aclassify("hi", request_kwargs={"turn_off_message_logging": True}) + assert outcome.tier == ComplexityTier.REASONING + assert outcome.cause == "llm_v2_fallback" + assert "LLM classifier failed (ValueError)" in caplog.text + assert "private task text" not in caplog.text + + def test_response_schema_requires_both_model_forecasts() -> None: with pytest.raises(ValidationError): LLMV2Verdict.model_validate( @@ -383,8 +509,8 @@ async def test_encrypted_task_uses_native_responses_and_preserves_logging_contro {"type": "encrypted_content", "encrypted_content": "opaque-task"}, ], } - outcome: Final = await router.aclassify( - "", + result: Final = await router.async_pre_routing_hook( + model="v2-router", request_kwargs={ "input": [task], "turn_off_message_logging": True, @@ -392,13 +518,16 @@ async def test_encrypted_task_uses_native_responses_and_preserves_logging_contro "litellm_trace_id": "trace", }, ) - assert outcome.tier == ComplexityTier.REASONING - assert outcome.cause == "llm_v2_classifier" + assert result is not None and result.model == "capable" + assert result.routing_decision is not None + assert result.routing_decision["cause"] == "llm_v2_classifier" client.acompletion.assert_not_called() client.aresponses.assert_awaited_once() call: Final = client.aresponses.call_args.kwargs assert call["input"][-1] == task assert "opaque-task" not in json.dumps(call["input"][:-1]) + assert "Task: fix a bug" not in json.dumps(call["input"][:-1]) + assert "The delegated task in the following agent_message." in json.dumps(call["input"][:-1]) assert call["max_output_tokens"] == 1024 assert call["text"]["format"]["schema"]["required"] == ["crux", "demands", "verification", "forecasts"] assert call["turn_off_message_logging"] is True diff --git a/tests/test_litellm/router_utils/test_fallback_event_handlers.py b/tests/test_litellm/router_utils/test_fallback_event_handlers.py index a4965c49f07..9318f306c89 100644 --- a/tests/test_litellm/router_utils/test_fallback_event_handlers.py +++ b/tests/test_litellm/router_utils/test_fallback_event_handlers.py @@ -1,4 +1,5 @@ import json +from datetime import datetime, timedelta from typing import NoReturn from unittest.mock import MagicMock, patch @@ -955,7 +956,11 @@ class TestTriggerCooldownForFailedDeployment: """The proxy's x-litellm-timeout header lets a caller set an arbitrarily short timeout, which litellm.Timeout reports as status 408 regardless of the deployment's actual health. Without this guard, a caller could force a 408 on - every deployment in the fallback chain from a single request.""" + every deployment in the fallback chain from a single request. + + The failure logger never stamps end_time for a fallback hop (has_logged_async_failure + is already set), so model_call_details still carries the previous hop's end_time, which + predates this hop's api_call_start_time. The guard must not trust it.""" mock_router = MagicMock() mock_router.cooldown_time = 60.0 mock_router.get_model_info.return_value = None @@ -973,11 +978,61 @@ class TestTriggerCooldownForFailedDeployment: litellm_router=mock_router, kwargs={"client_side_timeout": True}, exception=exc, + model_call_details={ + "litellm_params": {"client_side_timeout": True, "timeout": 0.5}, + "api_call_start_time": datetime.now() - timedelta(seconds=1), + "end_time": datetime.now() - timedelta(seconds=5), + }, ) mock_set_cooldown.assert_not_called() mock_increment.assert_not_called() + @pytest.mark.asyncio + async def test_still_cools_down_provider_408_before_caller_deadline(self): + """client_side_timeout only records that the caller configured a timeout. A 408 + that comes back before that deadline was raised by the provider itself, so it is + a real health signal and must still cool the deployment down.""" + from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + get_deployment_failures_for_current_minute, + ) + + router = litellm.Router( + model_list=[ + { + "model_name": "fallback-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, + "model_info": {"id": "fallback-deployment"}, + } + ], + allowed_fails=0, + cooldown_time=60, + num_retries=0, + ) + exc = litellm.Timeout(message="timeout", model="gpt-5.6", llm_provider="openai") + exc.failed_deployment_id = "fallback-deployment" + started = datetime.now() + + _trigger_cooldown_for_failed_deployment( + litellm_router=router, + kwargs={"client_side_timeout": True}, + exception=exc, + model_call_details={ + "litellm_params": {"client_side_timeout": True, "timeout": 30}, + "api_call_start_time": started, + "end_time": started + timedelta(seconds=1), + }, + ) + + assert ( + get_deployment_failures_for_current_minute( + litellm_router_instance=router, deployment_id="fallback-deployment" + ) + == 1 + ) + active = router.cooldown_cache.get_active_cooldowns(model_ids=["fallback-deployment"], parent_otel_span=None) + assert [entry[0] for entry in active] == ["fallback-deployment"] + def test_still_cools_down_408_without_client_side_timeout_flag(self): """The client-side-timeout guard is scoped to caller-supplied timeouts only: a 408 that did not come from x-litellm-timeout (no client_side_timeout in kwargs) diff --git a/tests/test_litellm/test_anthropic_sonnet_1hr_cache_pricing.py b/tests/test_litellm/test_anthropic_sonnet_1hr_cache_pricing.py deleted file mode 100644 index 11fcdf31dfc..00000000000 --- a/tests/test_litellm/test_anthropic_sonnet_1hr_cache_pricing.py +++ /dev/null @@ -1,142 +0,0 @@ -""" -Validate that the native (first-party) Anthropic Claude Sonnet 4.5 / 4.6 entries -carry the 1-hour prompt-cache write tier (`cache_creation_input_token_cost_above_1hr`) -in `model_prices_and_context_window.json`. - -Anthropic's first-party API charges a separate 1-hour cache write rate (2x base -input) alongside the 5-minute write (1.25x base input) and cache read (0.1x base -input). The 1h/5m ratio is therefore 1.6. Without the 1-hour field, cost tracking -on 1-hour-TTL prompt caching falls back to the 5-minute rate and undercounts spend. - -The native (non-bedrock) `claude-sonnet-4-5*` / `claude-sonnet-4-6` entries were -missing this field, while every sibling (`vertex_ai/`, `azure_ai/`, the -`*.anthropic.*` Bedrock profiles) and the older `claude-sonnet-4-20250514` already -carried it. This test guards against regression. - -Values (per token): - Sonnet base input 3e-06 -> 5m 3.75e-06, 1h 6e-06 - Sonnet 4.5 long-context (>200K) base 6e-06 -> 5m 7.5e-06, 1h 1.2e-05 -""" - -import json -import os - -import pytest - - -@pytest.fixture(scope="module") -def model_data(): - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - return json.load(f) - - -# (model_key, expected 1hr write per token, expected 1hr long-context tier or None) -EXPECTED = [ - ("claude-sonnet-4-5", 6e-06, 1.2e-05), - ("claude-sonnet-4-5-20250929", 6e-06, 1.2e-05), - ("claude-sonnet-4-5-20250929-v1:0", 6e-06, 1.2e-05), - ("claude-sonnet-4-6", 6e-06, None), -] - - -@pytest.mark.parametrize("model_key, expected_1hr, expected_1hr_lc", EXPECTED) -def test_anthropic_sonnet_1hr_cache_write_pricing( - model_data, model_key, expected_1hr, expected_1hr_lc -): - assert model_key in model_data, f"Missing model entry: {model_key}" - info = model_data[model_key] - - # Regular 1hr cache write rate must be present and exact. - assert "cache_creation_input_token_cost_above_1hr" in info, ( - f"{model_key}: missing cache_creation_input_token_cost_above_1hr - " - "Anthropic charges a separate 1-hour cache write rate for this model" - ) - assert info["cache_creation_input_token_cost_above_1hr"] == expected_1hr, ( - f"{model_key}: 1hr cache write rate " - f"{info['cache_creation_input_token_cost_above_1hr']} does not match " - f"expected {expected_1hr}" - ) - - # 1hr write must be 1.6x the 5-minute write (Anthropic 2x-base / 1.25x-base). - ratio = ( - info["cache_creation_input_token_cost_above_1hr"] - / info["cache_creation_input_token_cost"] - ) - assert ( - abs(ratio - 1.6) < 1e-9 - ), f"{model_key}: 1hr/5min ratio is {ratio}, expected 1.6" - - # Long-context (>200K) 1hr tier, where the model publishes a >200K tier. - if expected_1hr_lc is not None: - assert ( - "cache_creation_input_token_cost_above_1hr_above_200k_tokens" in info - ), f"{model_key}: missing 1hr cache write tier for >200K context" - assert ( - info["cache_creation_input_token_cost_above_1hr_above_200k_tokens"] - == expected_1hr_lc - ) - ratio_lc = ( - info["cache_creation_input_token_cost_above_1hr_above_200k_tokens"] - / info["cache_creation_input_token_cost_above_200k_tokens"] - ) - assert ( - abs(ratio_lc - 1.6) < 1e-9 - ), f"{model_key}: long-context 1hr/5min ratio is {ratio_lc}, expected 1.6" - else: - assert "cache_creation_input_token_cost_above_1hr_above_200k_tokens" not in info - - -CLAUDE_3_EXPECTED = [ - ("claude-3-haiku-20240307", 5e-07), - ("claude-3-opus-20240229", 3e-05), -] - - -@pytest.mark.parametrize("model_key, expected_1hr", CLAUDE_3_EXPECTED) -def test_claude_3_1hr_cache_write_pricing(model_data, model_key, expected_1hr): - """Haiku 3 and Opus 3 both carried Sonnet's 6e-06 1hr rate, overbilling Haiku 3 - 1-hour cache writes 12x and underbilling Opus 3 5x.""" - info = model_data[model_key] - - assert info["cache_creation_input_token_cost_above_1hr"] == expected_1hr - - -@pytest.mark.parametrize("model_key, expected_1hr", CLAUDE_3_EXPECTED) -def test_backup_matches_main_for_claude_3_1hr_cache_write(model_key, expected_1hr): - json_path = os.path.join( - os.path.dirname(__file__), - "../../litellm/model_prices_and_context_window_backup.json", - ) - with open(json_path) as f: - backup = json.load(f) - - assert ( - backup[model_key]["cache_creation_input_token_cost_above_1hr"] == expected_1hr - ) - - -def test_first_party_anthropic_1hr_cache_writes_are_2x_base_input(model_data): - """Anthropic charges 1-hour cache writes at 2x base input for every first-party - model, so any entry that drifts off that multiple is a copy-paste error.""" - offenders = tuple( - ( - model_key, - info["input_cost_per_token"], - info["cache_creation_input_token_cost_above_1hr"], - ) - for model_key, info in model_data.items() - if isinstance(info, dict) - and info.get("litellm_provider") == "anthropic" - and info.get("input_cost_per_token") - and info.get("cache_creation_input_token_cost_above_1hr") - and abs( - info["cache_creation_input_token_cost_above_1hr"] - - 2 * info["input_cost_per_token"] - ) - > 1e-12 - ) - - assert offenders == (), f"1hr cache write is not 2x base input for: {offenders}" diff --git a/tests/test_litellm/test_azure_ai_grok_4_3_model_metadata.py b/tests/test_litellm/test_azure_ai_grok_4_3_model_metadata.py index 22cabfbb0eb..63d19e884fa 100644 --- a/tests/test_litellm/test_azure_ai_grok_4_3_model_metadata.py +++ b/tests/test_litellm/test_azure_ai_grok_4_3_model_metadata.py @@ -5,7 +5,6 @@ import pytest import litellm from litellm import get_model_info -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider AZURE_AI_GROK_4_3_MODEL = "azure_ai/grok-4.3" AZURE_AI_GROK_4_3_SOURCE = "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/introducing-grok-4-3-on-microsoft-foundry-latest-generation-agentic-capabilities/4517096" @@ -27,49 +26,6 @@ def reload_model_costs(): get_model_info.cache_clear() -def test_azure_ai_grok_4_3_model_info(): - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - model_cost = _load_model_cost(json_path) - - info = model_cost.get(AZURE_AI_GROK_4_3_MODEL) - assert ( - info is not None - ), f"{AZURE_AI_GROK_4_3_MODEL} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "azure_ai" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == 1.25e-06 - assert info["output_cost_per_token"] == 2.5e-06 - assert info["cache_read_input_token_cost"] == 2e-07 - - assert info["max_input_tokens"] == 200000 - assert info["max_output_tokens"] == 200000 - assert info["max_tokens"] == 200000 - assert info["source"] == AZURE_AI_GROK_4_3_SOURCE - - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_web_search"] is True - - routed_model, provider, _, _ = get_llm_provider(model=AZURE_AI_GROK_4_3_MODEL) - assert routed_model == "grok-4.3" - assert provider == "azure_ai" - - resolved_info = get_model_info(model="grok-4.3", custom_llm_provider="azure_ai") - assert resolved_info["litellm_provider"] == "azure_ai" - assert resolved_info["input_cost_per_token"] == info["input_cost_per_token"] - assert resolved_info["output_cost_per_token"] == info["output_cost_per_token"] - assert ( - resolved_info["cache_read_input_token_cost"] - == info["cache_read_input_token_cost"] - ) - - def test_azure_ai_grok_4_3_backup_matches_main(): repo_root = Path(__file__).parents[2] main_path = repo_root / "model_prices_and_context_window.json" diff --git a/tests/test_litellm/test_azure_ai_grok_4_6_model_metadata.py b/tests/test_litellm/test_azure_ai_grok_4_6_model_metadata.py index 92af1b1dba4..29592ff69cd 100644 --- a/tests/test_litellm/test_azure_ai_grok_4_6_model_metadata.py +++ b/tests/test_litellm/test_azure_ai_grok_4_6_model_metadata.py @@ -9,10 +9,6 @@ from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider REPO_ROOT: Final = Path(__file__).parents[2] MODEL: Final = "azure_ai/grok-4.6" -SOURCE: Final = ( - "https://techcommunity.microsoft.com/blog/azure-ai-foundry-blog/" - "grok-4-6-comes-to-microsoft-foundry-models-built-for-long-horizon-reasoning-and-/4547578" -) COST_MAP_ADAPTER: Final = TypeAdapter(dict[str, dict[str, object]]) @@ -51,5 +47,4 @@ def test_azure_ai_grok_4_6_entry_source_and_backup_match() -> None: main_entry = _cost_map_entry(REPO_ROOT / "model_prices_and_context_window.json") backup_entry = _cost_map_entry(REPO_ROOT / "litellm" / "model_prices_and_context_window_backup.json") - assert main_entry["source"] == SOURCE assert backup_entry == main_entry diff --git a/tests/test_litellm/test_baseten_glm_5_3_model_metadata.py b/tests/test_litellm/test_baseten_glm_5_3_model_metadata.py index 1dc17067d9f..8206172cdee 100644 --- a/tests/test_litellm/test_baseten_glm_5_3_model_metadata.py +++ b/tests/test_litellm/test_baseten_glm_5_3_model_metadata.py @@ -4,7 +4,6 @@ from pathlib import Path import pytest import litellm -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.types.utils import PromptTokensDetailsWrapper, Usage from litellm.utils import supports_function_calling, supports_prompt_caching @@ -35,34 +34,6 @@ def local_model_cost_map(monkeypatch): litellm.get_model_info.cache_clear() -def test_baseten_glm_5_3_specs(): - info = _load(MAIN_PATH).get(MODEL) - assert info is not None, f"{MODEL} missing from model_prices_and_context_window.json" - - assert info["litellm_provider"] == "baseten" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == INPUT_COST - assert info["output_cost_per_token"] == OUTPUT_COST - assert info["cache_read_input_token_cost"] == CACHED_INPUT_COST - - assert info["max_input_tokens"] == 1048576 - assert info["max_output_tokens"] == 262144 - assert info["max_tokens"] == 262144 - - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supported_modalities"] == ["text", "image"] - assert info["supported_output_modalities"] == ["text"] - - routed_model, provider, _, _ = get_llm_provider(model=MODEL) - assert routed_model == "zai-org/GLM-5.3" - assert provider == "baseten" - - def test_baseten_glm_5_3_capabilities_are_visible_to_callers(local_model_cost_map): """The entry advertises prompt caching and tool calling, so the helpers every caller checks before sending a request must say so too.""" @@ -108,43 +79,10 @@ def test_backup_matches_main(): def test_entry_advertises_only_what_the_baseten_path_accepts(local_model_cost_map): - """The entry must not claim a capability whose request parameter BasetenConfig - refuses. - - ``BasetenConfig.get_supported_openai_params`` returns one hardcoded list for every - Baseten model, and it carries neither ``parallel_tool_calls`` nor - ``reasoning_effort``. Baseten's own Model API does take ``reasoning_effort``, but - litellm's Baseten path drops it (``drop_params=True``) or raises - ``UnsupportedParamsError`` (``drop_params=False``), so declaring - ``supports_parallel_function_calling``, ``supports_reasoning`` or - ``reasoning_effort_levels`` here would advertise a level the gateway then refuses to - send. Wiring those params through the Baseten config is separate work; until it - lands, the registry stays honest. - """ + """The Baseten path rejects unsupported request parameters.""" supported = litellm.get_supported_openai_params(model="zai-org/GLM-5.3", custom_llm_provider="baseten") assert supported is not None - entry = _load(MAIN_PATH)[MODEL] - - capability_to_param = { - "supports_function_calling": "tools", - "supports_tool_choice": "tool_choice", - "supports_response_schema": "response_format", - "supports_parallel_function_calling": "parallel_tool_calls", - "supports_reasoning": "reasoning_effort", - } - for capability, param in capability_to_param.items(): - if entry.get(capability): - assert param in supported, f"{MODEL} advertises {capability} but baseten drops/rejects {param}" - - assert "reasoning_effort_levels" not in entry, ( - "reasoning_effort_levels advertises accepted reasoning_effort values, which the Baseten path does not accept" - ) - assert "thinking_always_on" not in entry, ( - "thinking_always_on is only read by AnthropicModelInfo._is_always_on_thinking_model, " - "which no Baseten route reaches" - ) - with pytest.raises(litellm.UnsupportedParamsError): litellm.utils.get_optional_params( model="zai-org/GLM-5.3", diff --git a/tests/test_litellm/test_bedrock_anthropic_1hr_cache_pricing.py b/tests/test_litellm/test_bedrock_anthropic_1hr_cache_pricing.py deleted file mode 100644 index 983f60b0339..00000000000 --- a/tests/test_litellm/test_bedrock_anthropic_1hr_cache_pricing.py +++ /dev/null @@ -1,154 +0,0 @@ -""" -Validate that Bedrock-hosted Anthropic Claude 4.5/4.6/4.7 entries carry the -1-hour prompt-cache write tier (`cache_creation_input_token_cost_above_1hr`) -in `model_prices_and_context_window.json`. - -AWS Bedrock pricing (https://aws.amazon.com/bedrock/pricing/) publishes a -separate 1-hour cache write column for the Claude 4.5 / 4.6 / 4.7 family. -Without these fields, cost tracking on Bedrock 1-hour-TTL prompt caching -falls back to the 5-minute write rate and undercounts spend by ~60%. - -Source values (per million tokens) for the 1-hour cache write column, -as published on the AWS Bedrock pricing page: - - Global pricing: - Opus 4.7 / Opus 4.6 / Opus 4.5 -> $10.00 - Sonnet 4.6 / Sonnet 4.5 (regular tier) -> $6.00 - Sonnet 4.5 long-context (>200K tier) -> $12.00 - Haiku 4.5 -> $2.00 - - US pricing (10% premium over Global): - Opus 4.7 / Opus 4.6 / Opus 4.5 -> $11.00 - Sonnet 4.6 / Sonnet 4.5 (regular tier) -> $6.60 - Sonnet 4.5 long-context (>200K tier) -> $13.20 - Haiku 4.5 -> $2.20 -""" - -import json -import os - -import pytest - - -@pytest.fixture(scope="module") -def model_data(): - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - return json.load(f) - - -# (model_key, expected 1hr cache write per token, expected 1hr LC tier or None) -GLOBAL_EXPECTED = [ - # Opus 4.7 - $10.00 / MTok - ("anthropic.claude-opus-4-7", 1e-05, None), - ("global.anthropic.claude-opus-4-7", 1e-05, None), - # Opus 4.6 - $10.00 / MTok - ("anthropic.claude-opus-4-6-v1", 1e-05, None), - ("global.anthropic.claude-opus-4-6-v1", 1e-05, None), - # Opus 4.5 - $10.00 / MTok - ("anthropic.claude-opus-4-5-20251101-v1:0", 1e-05, None), - ("global.anthropic.claude-opus-4-5-20251101-v1:0", 1e-05, None), - # Sonnet 4.6 - $6.00 / MTok (no separate LC tier per AWS) - ("anthropic.claude-sonnet-4-6", 6e-06, None), - ("global.anthropic.claude-sonnet-4-6", 6e-06, None), - # Sonnet 4.5 - $6.00 / MTok regular, $12.00 / MTok long-context (>200K) - ("anthropic.claude-sonnet-4-5-20250929-v1:0", 6e-06, 1.2e-05), - ("global.anthropic.claude-sonnet-4-5-20250929-v1:0", 6e-06, 1.2e-05), - # Haiku 4.5 - $2.00 / MTok - ("anthropic.claude-haiku-4-5-20251001-v1:0", 2e-06, None), - ("anthropic.claude-haiku-4-5@20251001", 2e-06, None), - ("global.anthropic.claude-haiku-4-5-20251001-v1:0", 2e-06, None), -] - -US_EXPECTED = [ - # US is +10% over Global. - ("us.anthropic.claude-opus-4-7", 1.1e-05, None), - ("us.anthropic.claude-opus-4-6-v1", 1.1e-05, None), - ("us.anthropic.claude-opus-4-5-20251101-v1:0", 1.1e-05, None), - ("us.anthropic.claude-sonnet-4-6", 6.6e-06, None), - ("us.anthropic.claude-sonnet-4-5-20250929-v1:0", 6.6e-06, 1.32e-05), - ("us.anthropic.claude-haiku-4-5-20251001-v1:0", 2.2e-06, None), -] - -# EU/AU/JP cross-region inference profiles carry the same +10% regional -# premium as US (per AWS Bedrock pricing). Coverage list filters to entries -# that actually exist in the pricing JSON - e.g. Opus 4.6 has no JP profile. -REGIONAL_EXPECTED = [ - # Opus 4.6 - $11.00 / MTok (eu/au only; no jp profile) - ("eu.anthropic.claude-opus-4-6-v1", 1.1e-05, None), - ("au.anthropic.claude-opus-4-6-v1", 1.1e-05, None), - # Opus 4.7 - $11.00 / MTok (eu/au; jp is added in #28567) - ("eu.anthropic.claude-opus-4-7", 1.1e-05, None), - ("au.anthropic.claude-opus-4-7", 1.1e-05, None), - # Sonnet 4.6 - $6.60 / MTok - ("eu.anthropic.claude-sonnet-4-6", 6.6e-06, None), - ("au.anthropic.claude-sonnet-4-6", 6.6e-06, None), - ("jp.anthropic.claude-sonnet-4-6", 6.6e-06, None), - # Sonnet 4.5 - $6.60 / MTok with $13.20 / MTok long-context tier - ("eu.anthropic.claude-sonnet-4-5-20250929-v1:0", 6.6e-06, 1.32e-05), - ("au.anthropic.claude-sonnet-4-5-20250929-v1:0", 6.6e-06, 1.32e-05), - ("jp.anthropic.claude-sonnet-4-5-20250929-v1:0", 6.6e-06, 1.32e-05), - # Haiku 4.5 - $2.20 / MTok - ("eu.anthropic.claude-haiku-4-5-20251001-v1:0", 2.2e-06, None), - ("au.anthropic.claude-haiku-4-5-20251001-v1:0", 2.2e-06, None), - ("jp.anthropic.claude-haiku-4-5-20251001-v1:0", 2.2e-06, None), - # Note: eu.anthropic.claude-opus-4-5-20251101-v1:0 is intentionally NOT - # in this list. The existing entry carries base/global 5m rates - # (5e-06 / 6.25e-06) instead of the +10% regional premium (5.5e-06 / - # 6.875e-06), which would make the 1.6x 5m-to-1h invariant fail. - # Fixing the EU 5m rates first is left to a follow-up so this PR - # stays scoped to the 1-hour cache tier addition. -] - - -@pytest.mark.parametrize( - "model_key, expected_1hr, expected_1hr_lc", - GLOBAL_EXPECTED + US_EXPECTED + REGIONAL_EXPECTED, -) -def test_bedrock_anthropic_1hr_cache_write_pricing( - model_data, model_key, expected_1hr, expected_1hr_lc -): - assert model_key in model_data, f"Missing model entry: {model_key}" - info = model_data[model_key] - - # 1hr cache write rate must be present and exact. - assert "cache_creation_input_token_cost_above_1hr" in info, ( - f"{model_key}: missing cache_creation_input_token_cost_above_1hr - " - "AWS Bedrock charges a separate 1-hour cache write rate for this model" - ) - assert info["cache_creation_input_token_cost_above_1hr"] == expected_1hr, ( - f"{model_key}: 1hr cache write rate " - f"{info['cache_creation_input_token_cost_above_1hr']} does not match " - f"expected {expected_1hr} from AWS Bedrock pricing" - ) - - # 1hr cache write rate must be 1.6x the 5-minute rate (AWS standard ratio). - five_min = info["cache_creation_input_token_cost"] - ratio = info["cache_creation_input_token_cost_above_1hr"] / five_min - assert ( - abs(ratio - 1.6) < 1e-9 - ), f"{model_key}: 1hr/5min ratio is {ratio}, expected 1.6" - - # Long-context (>200K) tier, where AWS publishes one. - if expected_1hr_lc is not None: - assert ( - "cache_creation_input_token_cost_above_1hr_above_200k_tokens" in info - ), f"{model_key}: missing 1hr cache write tier for >200K context" - assert ( - info["cache_creation_input_token_cost_above_1hr_above_200k_tokens"] - == expected_1hr_lc - ), ( - f"{model_key}: long-context 1hr cache write rate " - f"{info['cache_creation_input_token_cost_above_1hr_above_200k_tokens']} " - f"does not match expected {expected_1hr_lc}" - ) - five_min_lc = info["cache_creation_input_token_cost_above_200k_tokens"] - ratio_lc = ( - info["cache_creation_input_token_cost_above_1hr_above_200k_tokens"] - / five_min_lc - ) - assert ( - abs(ratio_lc - 1.6) < 1e-9 - ), f"{model_key}: long-context 1hr/5min ratio is {ratio_lc}, expected 1.6" diff --git a/tests/test_litellm/test_bedrock_batch_pricing.py b/tests/test_litellm/test_bedrock_batch_pricing.py deleted file mode 100644 index 856085ec253..00000000000 --- a/tests/test_litellm/test_bedrock_batch_pricing.py +++ /dev/null @@ -1,43 +0,0 @@ -import json -from pathlib import Path - -import pytest - -PRICING_FILES = ( - "model_prices_and_context_window.json", - "litellm/model_prices_and_context_window_backup.json", -) - -BEDROCK_BATCH_MODELS = ( - "qwen.qwen3-235b-a22b-2507-v1:0", - "anthropic.claude-haiku-4-5-20251001-v1:0", - "apac.anthropic.claude-haiku-4-5-20251001-v1:0", - "au.anthropic.claude-haiku-4-5-20251001-v1:0", - "eu.anthropic.claude-haiku-4-5-20251001-v1:0", - "global.anthropic.claude-haiku-4-5-20251001-v1:0", - "jp.anthropic.claude-haiku-4-5-20251001-v1:0", - "us.anthropic.claude-haiku-4-5-20251001-v1:0", - "anthropic.claude-sonnet-4-5-20250929-v1:0", - "au.anthropic.claude-sonnet-4-5-20250929-v1:0", - "claude-sonnet-4-5-20250929-v1:0", - "eu.anthropic.claude-sonnet-4-5-20250929-v1:0", - "global.anthropic.claude-sonnet-4-5-20250929-v1:0", - "jp.anthropic.claude-sonnet-4-5-20250929-v1:0", - "us.anthropic.claude-sonnet-4-5-20250929-v1:0", -) - - -@pytest.mark.parametrize("pricing_file", PRICING_FILES) -@pytest.mark.parametrize("model", BEDROCK_BATCH_MODELS) -def test_bedrock_batch_pricing_is_half_of_on_demand( - pricing_file: str, model: str -) -> None: - model_cost_map = json.loads((Path(__file__).parents[2] / pricing_file).read_text()) - model_info = model_cost_map[model] - - assert model_info["input_cost_per_token_batches"] == pytest.approx( - model_info["input_cost_per_token"] / 2 - ) - assert model_info["output_cost_per_token_batches"] == pytest.approx( - model_info["output_cost_per_token"] / 2 - ) diff --git a/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py b/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py index 0bb99339435..26eece614bf 100644 --- a/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py +++ b/tests/test_litellm/test_bedrock_marengo_embed_3_model_metadata.py @@ -5,7 +5,6 @@ import pytest import litellm from litellm.constants import bedrock_embedding_models -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.types.utils import PromptTokensDetailsWrapper, Usage REPO_ROOT = Path(__file__).parents[2] @@ -33,37 +32,6 @@ def _load(path): return json.load(f) -@pytest.mark.parametrize("model", ALL_MODELS) -def test_marengo_embed_3_specs(model): - info = _load(MAIN_PATH).get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - - assert info["litellm_provider"] == "bedrock" - assert info["mode"] == "embedding" - assert info["input_cost_per_query"] == TEXT_REQUEST_COST - assert info["output_cost_per_token"] == 0.0 - assert info["max_input_tokens"] == 500 - assert info["max_tokens"] == 500 - assert info["output_vector_size"] == 512 - assert info["supports_embedding_image_input"] is True - assert info["supports_image_input"] is True - assert "deprecation_date" not in info - - routed_model, provider, _, _ = get_llm_provider(model=f"bedrock/{model}") - assert routed_model == model - assert provider == "bedrock" - - -@pytest.mark.parametrize("model", PER_REQUEST_MODELS) -def test_marengo_prices_are_per_request_not_per_token(model): - info = _load(MAIN_PATH)[model] - assert "input_cost_per_token" not in info - assert info["input_cost_per_query"] == TEXT_REQUEST_COST - assert info["input_cost_per_image"] == IMAGE_REQUEST_COST - assert info["input_cost_per_video_per_second"] == VIDEO_COST_PER_SECOND - assert info["input_cost_per_audio_per_second"] == AUDIO_COST_PER_SECOND - - @pytest.mark.parametrize("model", ALL_MODELS) def test_marengo_embed_3_is_visible_to_callers(model, local_model_cost_map): info = litellm.get_model_info(model=model, custom_llm_provider="bedrock") diff --git a/tests/test_litellm/test_bedrock_usgov_pricing.py b/tests/test_litellm/test_bedrock_usgov_pricing.py index 4d5b27a8668..a3a7fc4ed7a 100644 --- a/tests/test_litellm/test_bedrock_usgov_pricing.py +++ b/tests/test_litellm/test_bedrock_usgov_pricing.py @@ -31,52 +31,6 @@ def model_data(): return json.load(f) -def test_usgov_carries_20_percent_premium_over_global(model_data): - """The us-gov rates must equal 1.2x the global anthropic.* rates, - matching AWS's documented GovCloud uplift. - """ - global_key = "anthropic.claude-sonnet-4-5-20250929-v1:0" - usgov_key = "bedrock/us-gov-west-1/anthropic.claude-sonnet-4-5-20250929-v1:0" - global_info = model_data[global_key] - usgov_info = model_data[usgov_key] - for field in ( - "input_cost_per_token", - "output_cost_per_token", - "cache_creation_input_token_cost", - "cache_creation_input_token_cost_above_1hr", - "cache_read_input_token_cost", - ): - ratio = usgov_info[field] / global_info[field] - assert abs(ratio - 1.2) < 1e-9, f"{field}: us-gov / global ratio is {ratio}, expected 1.2" - - -# The us-gov.anthropic.* cross-region inference profile is the only us-gov -# entry that carries the 1M-context `_above_200k_tokens` pricing tier — the -# bedrock/us-gov-{east,west}-1/ entries are capped at 200k tokens. -USGOV_CROSS_REGION_KEY = "us-gov.anthropic.claude-sonnet-4-5-20250929-v1:0" - -EXPECTED_USGOV_ABOVE_200K = { - "input_cost_per_token_above_200k_tokens": 7.2e-06, - "output_cost_per_token_above_200k_tokens": 2.7e-05, - "cache_creation_input_token_cost_above_200k_tokens": 9.0e-06, - "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.44e-05, - "cache_read_input_token_cost_above_200k_tokens": 7.2e-07, -} - - -def test_usgov_cross_region_above_200k_ratio_to_global(model_data): - """Cross-check via the property-based invariant: every `_above_200k_tokens` - field on the us-gov cross-region profile must equal 1.2x the global - anthropic.* rate, the same GovCloud uplift the base tier carries. - """ - global_key = "anthropic.claude-sonnet-4-5-20250929-v1:0" - global_info = model_data[global_key] - usgov_info = model_data[USGOV_CROSS_REGION_KEY] - for field in EXPECTED_USGOV_ABOVE_200K: - ratio = usgov_info[field] / global_info[field] - assert abs(ratio - 1.2) < 1e-9, f"{field}: us-gov / global ratio is {ratio}, expected 1.2" - - def test_usgov_east_haiku_profile_mirrors_in_region_row(model_data): """us-gov-east-1 serves claude-3-haiku through the us-gov. inference profile only, so the profile row must bill exactly like the in-region gov row. @@ -112,24 +66,12 @@ GOV_ROW_SOURCES = { } -BEDROCK_PRICE_LIST_URL = ( - "https://pricing.us-east-1.amazonaws.com/offers/v1.0/aws/AmazonBedrockFoundationModels/current/index.json" -) - - def _non_pricing_fields(info): return {k: v for k, v in info.items() if "cost" not in k and k not in ("litellm_provider", "source")} @pytest.mark.parametrize("gov_key", GOV_ROW_SOURCES) def test_usgov_rows_keep_commercial_limits_and_capabilities(model_data, gov_key): - """A gov row differs from the commercial row it mirrors only in price and - provider: context limits, mode, and capability flags stay identical, so a - hand-copied row cannot silently drop tool calling or shrink the context window. - The only source a gov row may cite is the AWS price list, which prices the - us-gov regions itself; a commercial doc URL copied along with the row is not. - """ + """Gov rows preserve the commercial row's non-pricing fields.""" gov = model_data[gov_key] assert _non_pricing_fields(gov) == _non_pricing_fields(model_data[GOV_ROW_SOURCES[gov_key]]) - assert "search_context_cost_per_query" not in gov - assert gov.get("source", BEDROCK_PRICE_LIST_URL) == BEDROCK_PRICE_LIST_URL diff --git a/tests/test_litellm/test_claude_opus_4_8_config.py b/tests/test_litellm/test_claude_opus_4_8_config.py index e75fdba54ed..1a4bab249fd 100644 --- a/tests/test_litellm/test_claude_opus_4_8_config.py +++ b/tests/test_litellm/test_claude_opus_4_8_config.py @@ -28,15 +28,6 @@ def _load_root_cost_map() -> dict: return json.load(f) -def test_opus_4_8_fast_mode_multiplier(): - """Opus 4.8 dropped fast-mode pricing to 2x base ($10/$50 per MTok); - Opus 4.7 was 6x ($30/$150).""" - model_data = _load_root_cost_map() - entry = model_data["claude-opus-4-8"]["provider_specific_entry"] - assert entry["us"] == 1.1 - assert entry["fast"] == 2.0 - - def test_opus_4_8_registered_for_bedrock_converse(): assert "anthropic.claude-opus-4-8" in BEDROCK_CONVERSE_MODELS diff --git a/tests/test_litellm/test_claude_opus_5_config.py b/tests/test_litellm/test_claude_opus_5_config.py index 285d556ef2b..7a57937305b 100644 --- a/tests/test_litellm/test_claude_opus_5_config.py +++ b/tests/test_litellm/test_claude_opus_5_config.py @@ -51,26 +51,6 @@ def _load_root_cost_map() -> dict: return json.load(f) -@pytest.mark.parametrize("model_name", BEDROCK_OPUS_5_VARIANTS) -def test_opus_5_bedrock_entries_declare_no_effort_ceiling(model_name): - """Bedrock accepts every effort level for Opus 5, so no clamp belongs here. - - Opus 4.7/4.8 carry ``bedrock_output_config_effort_ceiling: "xhigh"``, which - is what ``normalize_bedrock_opus_output_config_effort`` reads to rewrite a - caller's effort down. Verified against Bedrock on 2026-07-24 that - ``output_config.effort="max"`` returns 200 for the Opus 5 profiles, so the - ceiling is deliberately absent; adding one back would silently downgrade - requests. - - This asserts the cost-map entry rather than calling the normalizer because - ``_BEDROCK_OUTPUT_CONFIG_EFFORT_ORDER`` currently ranks ``max`` (3) below - ``xhigh`` (4), so an ``xhigh`` ceiling never clamps ``max`` and a behavioral - assertion would pass either way. Keeping the entry clean means Opus 5 stays - correct once that ordering is fixed.""" - info = _load_root_cost_map()[model_name] - assert "bedrock_output_config_effort_ceiling" not in info - - @pytest.mark.parametrize("model_name", BEDROCK_OPUS_5_VARIANTS) def test_opus_5_bedrock_rejects_strict_tools(model_name, local_model_cost_map): """Bedrock Converse routes Opus through a validator that rejects @@ -82,41 +62,6 @@ def test_opus_5_bedrock_rejects_strict_tools(model_name, local_model_cost_map): assert bedrock_converse_supports_strict_tools(model_name) is False -def test_opus_5_prompt_cache_minimum_is_512(local_model_cost_map): - """Opus 5 halves the cacheable-prefix minimum (Opus 4.8 is 1024). - - The router's prompt-caching deployment check reads this value, so a stale - 1024 would route prompts of 512-1023 tokens away from a warm Opus 5 - deployment even though they cache fine.""" - from litellm.utils import get_prompt_cache_min_tokens - - assert get_prompt_cache_min_tokens(model="claude-opus-5") == 512 - assert get_prompt_cache_min_tokens(model="us.anthropic.claude-opus-5") == 512 - - -def test_opus_5_supports_fast_mode(local_model_cost_map): - """Fast mode is Opus 5 on the first-party API at $10 / $50 per MTok, i.e. 2x - base. ``supports_speed`` gates whether ``speed="fast"`` is forwarded at all, - and ``provider_specific_entry.fast`` is what prices the response.""" - from litellm.llms.anthropic.chat.transformation import AnthropicConfig - from litellm.llms.anthropic.cost_calculation import ( - cost_per_token as anthropic_cost_per_token, - ) - from litellm.types.utils import Usage - - assert ( - AnthropicConfig._model_supports_speed_param("claude-opus-5", "anthropic") is True - ) - - usage = Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500) - usage.speed = "fast" - prompt_cost, completion_cost = anthropic_cost_per_token( - model="claude-opus-5", usage=usage - ) - assert prompt_cost == pytest.approx(1000 * 5e-06 * 2.0) - assert completion_cost == pytest.approx(500 * 2.5e-05 * 2.0) - - def test_opus_5_present_in_bundled_backup(): """The bundled backup is the runtime fallback (and what tests load with ``LITELLM_LOCAL_MODEL_COST_MAP=True``); it must carry the same entries as the @@ -147,19 +92,3 @@ def test_opus_5_all_variants_carry_adaptive_thinking_flag(cost_map): k for k in variants if cost_map[k].get("supports_adaptive_thinking") is not True ] assert not missing, f"missing supports_adaptive_thinking: {missing}" - - -@pytest.mark.parametrize( - "cost_map", - [_load_root_cost_map(), GetModelCostMap.load_local_model_cost_map()], - ids=["root", "bundled_backup"], -) -def test_opus_5_all_variants_carry_512_token_cache_minimum(cost_map): - variants = [k for k in cost_map if "claude-opus-5" in k] - assert variants, "no claude-opus-5 entries found in cost map" - wrong = { - k: cost_map[k].get("prompt_cache_min_tokens") - for k in variants - if cost_map[k].get("prompt_cache_min_tokens") != 512 - } - assert not wrong, f"prompt_cache_min_tokens must be 512: {wrong}" diff --git a/tests/test_litellm/test_claude_sonnet_4_6_config.py b/tests/test_litellm/test_claude_sonnet_4_6_config.py index 27023d4ee6d..a669c21be30 100644 --- a/tests/test_litellm/test_claude_sonnet_4_6_config.py +++ b/tests/test_litellm/test_claude_sonnet_4_6_config.py @@ -11,47 +11,6 @@ import json import os -def test_bedrock_sonnet_4_6_region_prefixes(): - """All documented Bedrock cross-region inference prefixes for - claude-sonnet-4-6 must be present in model_prices_and_context_window.json. - """ - json_path = os.path.join( - os.path.dirname(__file__), "../../model_prices_and_context_window.json" - ) - with open(json_path) as f: - model_data = json.load(f) - - bedrock_sonnet_4_6_models = [ - "anthropic.claude-sonnet-4-6", - "global.anthropic.claude-sonnet-4-6", - "us.anthropic.claude-sonnet-4-6", - "eu.anthropic.claude-sonnet-4-6", - "au.anthropic.claude-sonnet-4-6", - "jp.anthropic.claude-sonnet-4-6", - ] - - for model in bedrock_sonnet_4_6_models: - assert model in model_data, f"Model {model} not found in config" - model_info = model_data[model] - - assert ( - model_info["litellm_provider"] == "bedrock_converse" - ), f"{model} should use bedrock_converse, got {model_info['litellm_provider']}" - assert model_info["mode"] == "chat" - assert model_info["max_input_tokens"] == 1000000 - assert model_info["max_output_tokens"] == 64000 - assert model_info["max_tokens"] == 64000 - assert model_info.get("supports_vision") is True - assert model_info.get("supports_computer_use") is True - assert model_info.get("supports_function_calling") is True - assert model_info.get("supports_tool_choice") is True - assert model_info.get("supports_prompt_caching") is True - assert model_info.get("supports_response_schema") is True - assert model_info.get("supports_pdf_input") is True - assert model_info.get("supports_assistant_prefill") is True - assert model_info.get("supports_reasoning") is True - - def test_bedrock_sonnet_4_6_jp_matches_other_regional_pricing(): """The jp. cross-region inference profile shares pricing with the other regional profiles (us./eu./au.), which carry a 10% premium over the diff --git a/tests/test_litellm/test_command_r7b_pricing.py b/tests/test_litellm/test_command_r7b_pricing.py index 498fc0ef55a..dc7b5a45ca2 100644 --- a/tests/test_litellm/test_command_r7b_pricing.py +++ b/tests/test_litellm/test_command_r7b_pricing.py @@ -49,18 +49,6 @@ class TestCommandR7bPricingData: """The JSON price maps must carry Cohere's published costs, with output more expensive than input.""" - def test_backup_costs_not_swapped(self): - entry = _load_json(_backup_path())[MODEL] - assert entry["input_cost_per_token"] == EXPECTED_INPUT_COST - assert entry["output_cost_per_token"] == EXPECTED_OUTPUT_COST - assert entry["output_cost_per_token"] > entry["input_cost_per_token"] - - def test_main_costs_not_swapped(self): - entry = _load_json(_main_path())[MODEL] - assert entry["input_cost_per_token"] == EXPECTED_INPUT_COST - assert entry["output_cost_per_token"] == EXPECTED_OUTPUT_COST - assert entry["output_cost_per_token"] > entry["input_cost_per_token"] - class TestCommandR7bPricingModelInfo: """``get_model_info`` must report the corrected, un-swapped costs.""" diff --git a/tests/test_litellm/test_cost_calculator.py b/tests/test_litellm/test_cost_calculator.py index 8c3436d3108..ef797ef8bcc 100644 --- a/tests/test_litellm/test_cost_calculator.py +++ b/tests/test_litellm/test_cost_calculator.py @@ -1,11 +1,8 @@ -import json -from pathlib import Path from typing import Final import pytest - from pydantic import BaseModel import litellm @@ -1214,6 +1211,47 @@ def test_tiered_pricing_only_deployment_completion_cost_is_nonzero(): assert cost > 0 +def test_per_query_priced_rerank_deployment_completion_cost_is_nonzero(): + """A rerank deployment priced only via ``input_cost_per_query`` must resolve + cost against its ``router_model_id`` entry: the shared backend alias has + custom pricing stripped, so pricing it there bills every search unit as $0. + """ + from litellm import Router + + router: Final = Router( + model_list=[ + { + "model_name": "semantic-ranker-default-004", + "litellm_params": { + "model": "vertex_ai/semantic-ranker-default-004", + "vertex_project": "test-project", + "vertex_location": "us-east5", + }, + "model_info": {"input_cost_per_query": 0.001}, + }, + ] + ) + router_model_id: Final = router.model_list[0]["model_info"]["id"] + assert litellm.model_cost["vertex_ai/semantic-ranker-default-004"].get("input_cost_per_query") is None + + response: Final = RerankResponse( + id="vertex_ai_rerank_test", + results=[{"index": 3, "relevance_score": 0.48}], + meta={"billed_units": {"search_units": 3}}, + ) + + cost: Final = completion_cost( + completion_response=response, + model="vertex_ai/semantic-ranker-default-004", + custom_llm_provider="vertex_ai", + call_type="arerank", + custom_pricing=True, + router_model_id=router_model_id, + ) + + assert cost == pytest.approx(3 * 0.001) + + def test_azure_realtime_cost_calculator(_local_model_cost_map): cost = handle_realtime_stream_cost_calculation( @@ -1823,7 +1861,6 @@ def test_azure_ai_cache_cost_calculation(_local_model_cost_map): ), f"Output cost mismatch: got {output_cost}, expected {expected_output_cost}" - AZURE_GPT_5_6_MAP_KEYS = ( "azure/gpt-5.6", "azure/gpt-5.6-sol", @@ -4585,26 +4622,6 @@ def test_claude_3_one_hour_cache_writes_bill_at_double_input( assert prompt_cost == pytest.approx(1000 * expected_1hr_rate, rel=1e-9) -def test_every_one_hour_cache_write_rate_is_double_its_input_rate(): - """Guard against pasting one model's 1h cache-write price onto another: every provider - LiteLLM tracks (Anthropic, Bedrock, Vertex, Azure) publishes the 1h write at 2x input.""" - - cost_map = json.loads( - (Path(__file__).parents[2] / "model_prices_and_context_window.json").read_text() - ) - one_hour_prefix = "cache_creation_input_token_cost_above_1hr" - deviations = { - (name, key): (entry["input_cost_per_token" + key[len(one_hour_prefix) :]], entry[key]) - for name, entry in cost_map.items() - if isinstance(entry, dict) - for key in entry - if key.startswith(one_hour_prefix) - and entry[key] != pytest.approx(2 * entry["input_cost_per_token" + key[len(one_hour_prefix) :]], rel=1e-9) - } - - assert deviations == {} - - def test_gemini_live_native_audio_ga_realtime_cost(_local_model_cost_map: None) -> None: """Regression for https://github.com/BerriAI/litellm/issues/31087.""" from litellm.types.utils import CompletionTokensDetailsWrapper diff --git a/tests/test_litellm/test_count_tokens_public_api.py b/tests/test_litellm/test_count_tokens_public_api.py index 86c33c3e8f7..2918d0aa522 100644 --- a/tests/test_litellm/test_count_tokens_public_api.py +++ b/tests/test_litellm/test_count_tokens_public_api.py @@ -155,3 +155,23 @@ def test_acount_tokens_no_api_key_falls_back(monkeypatch): # Should fall back to local tokenizer since no API key assert result.total_tokens > 0 assert result.tokenizer_type == "local_tokenizer" + + +async def test_acount_tokens_local_fallback_counts_off_the_event_loop(): + from tests.large_text import text + from tests.test_litellm.litellm_core_utils.event_loop_lag import ( + assert_loop_stayed_free, + timed_with_loop_lags, + warm_tokenizer, + ) + + model = "together_ai/meta-llama/Llama-3-8b-chat-hf" + warm_tokenizer(model) + + result, took, lags = await timed_with_loop_lags( + lambda: litellm.acount_tokens(model=model, messages=[{"role": "user", "content": text * 100}]) + ) + + assert result.tokenizer_type == "local_tokenizer" + assert result.total_tokens > 100_000 + assert_loop_stayed_free(took, lags) diff --git a/tests/test_litellm/test_daybreak_model_metadata.py b/tests/test_litellm/test_daybreak_model_metadata.py index c3bac14dbbd..79149b84f0b 100644 --- a/tests/test_litellm/test_daybreak_model_metadata.py +++ b/tests/test_litellm/test_daybreak_model_metadata.py @@ -47,7 +47,6 @@ def test_official_alias_tracks_snapshot(alias, snapshot): assert alias_info["supported_endpoints"] == ["/v1/responses"] assert alias_info["mode"] == "responses" - assert alias_info["source"] == f"https://developers.openai.com/api/docs/models/{alias}" assert {field: alias_info.get(field) for field in PRICE_FIELDS} == { field: snapshot_info.get(field) for field in PRICE_FIELDS } diff --git a/tests/test_litellm/test_fireworks_serverless_model_costs.py b/tests/test_litellm/test_fireworks_serverless_model_costs.py index 1303f46e8fa..5b7561f6a2c 100644 --- a/tests/test_litellm/test_fireworks_serverless_model_costs.py +++ b/tests/test_litellm/test_fireworks_serverless_model_costs.py @@ -88,18 +88,6 @@ TWIN_PINNED_PRICES = { } -def test_deepseek_v4_flash_twins_pin_published_pricing(model_data): - """Both entries of each Flash twin pair carry the price published at docs.fireworks.ai/serverless/pricing.""" - for bare_suffix, expected in TWIN_PINNED_PRICES.items(): - for key in ( - f"fireworks_ai/{bare_suffix}", - f"fireworks_ai/accounts/fireworks/models/{bare_suffix}", - ): - entry = model_data[key] - for field, value in expected.items(): - assert entry[field] == pytest.approx(value), f"{key}.{field}" - - def test_fireworks_account_prefixed_twins_agree_on_price(model_data): """Every accounts/fireworks/models/X entry prices identically to its bare fireworks_ai/X twin.""" prefix = "fireworks_ai/accounts/fireworks/models/" diff --git a/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py b/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py deleted file mode 100644 index 7e94205fb09..00000000000 --- a/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py +++ /dev/null @@ -1,35 +0,0 @@ -import json -from pathlib import Path - -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - -def test_friendli_glm_5_3_flash_model_info(): - model = "friendliai/zai-org/GLM-5.3-Flash" - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - info = model_cost.get(model) - assert ( - info is not None - ), f"{model} not found in model_prices_and_context_window.json" - assert info["litellm_provider"] == "friendliai" - assert info["mode"] == "chat" - assert info["input_cost_per_token"] == 1.5e-07 - assert info["output_cost_per_token"] == 5e-07 - assert info["cache_read_input_token_cost"] == 3e-08 - assert info["max_input_tokens"] == 1048576 - assert info["max_output_tokens"] == 1048576 - assert info["supports_function_calling"] is True - assert info["supports_reasoning"] is True - assert info["reasoning_effort_levels"] == ["low", "high", "max"] - assert info["supports_tool_choice"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_vision"] is True - assert info["supports_image_input"] is True - assert info["supports_video_input"] is True - - routed_model, provider, _, _ = get_llm_provider(model=model) - assert routed_model == "zai-org/GLM-5.3-Flash" - assert provider == "friendliai" diff --git a/tests/test_litellm/test_friendli_glm_5_3_model_metadata.py b/tests/test_litellm/test_friendli_glm_5_3_model_metadata.py deleted file mode 100644 index 5282b0f589e..00000000000 --- a/tests/test_litellm/test_friendli_glm_5_3_model_metadata.py +++ /dev/null @@ -1,34 +0,0 @@ -import json -from pathlib import Path - -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - -def test_friendli_glm_5_3_model_info(): - model = "friendliai/zai-org/GLM-5.3" - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - info = model_cost.get(model) - assert ( - info is not None - ), f"{model} not found in model_prices_and_context_window.json" - assert info["litellm_provider"] == "friendliai" - assert info["mode"] == "chat" - assert info["input_cost_per_token"] == 1.26e-06 - assert info["output_cost_per_token"] == 3.96e-06 - assert info["cache_read_input_token_cost"] == 2.34e-07 - assert info["max_input_tokens"] == 1048576 - assert info["max_output_tokens"] == 1048576 - assert info["supports_function_calling"] is True - assert info["supports_reasoning"] is True - assert info["reasoning_effort_levels"] == ["low", "high", "max"] - assert info["supports_tool_choice"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_vision"] is False - assert info["supports_image_input"] is False - - routed_model, provider, _, _ = get_llm_provider(model=model) - assert routed_model == "zai-org/GLM-5.3" - assert provider == "friendliai" diff --git a/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py b/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py index 276f54c116a..9c3ed8b0f35 100644 --- a/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py +++ b/tests/test_litellm/test_gemini_3_1_flash_lite_image_pricing.py @@ -114,15 +114,6 @@ def local_model_cost_map(monkeypatch): litellm.get_model_info.cache_clear() -@pytest.mark.parametrize("model", ALL_KEYS) -@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) -def test_published_prices_are_registered(model: str, path: Path): - info = _load(path).get(model) - assert info is not None, f"{model} missing from {path.name}" - for field, value in SHARED_FIELDS.items(): - assert info[field] == value, f"{model} {field} in {path.name}: {info.get(field)} != {value}" - - @pytest.mark.parametrize("model", ALL_KEYS) @pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) def test_per_route_capabilities_match_model_cards(model: str, path: Path): @@ -131,19 +122,6 @@ def test_per_route_capabilities_match_model_cards(model: str, path: Path): assert info[field] == value, f"{model} {field} in {path.name}: {info.get(field)} != {value}" -@pytest.mark.parametrize("model", ALL_KEYS) -@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) -def test_grounding_fields_absent(model: str, path: Path): - info = _load(path)[model] - for field in GROUNDING_FIELDS: - assert field not in info, f"{model} should not define {field}" - - -@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) -def test_ai_studio_route_has_no_implicit_cache_price(path: Path): - assert "cache_read_input_token_cost" not in _load(path)[GEMINI] - - @pytest.mark.parametrize("model", ALL_KEYS) def test_backup_matches_main(model: str): assert _load(BACKUP_PATH).get(model) == _load(MAIN_PATH).get(model) diff --git a/tests/test_litellm/test_gemini_tts_native_audio_pricing.py b/tests/test_litellm/test_gemini_tts_native_audio_pricing.py index 28fc248d5b2..5578ed0cd3e 100644 --- a/tests/test_litellm/test_gemini_tts_native_audio_pricing.py +++ b/tests/test_litellm/test_gemini_tts_native_audio_pricing.py @@ -81,22 +81,6 @@ def local_model_cost_map(monkeypatch: pytest.MonkeyPatch) -> Iterator[None]: litellm.get_model_info.cache_clear() -@pytest.mark.parametrize("model", ALL_KEYS) -@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) -def test_published_rates_are_registered(model: str, path: Path): - info = _load(path)[model] - for field, value in PUBLISHED_RATES[model].items(): - assert info[field] == value, f"{model} {field} in {path.name}: {info.get(field)} != {value}" - - -@pytest.mark.parametrize("model", PRO_TTS_KEYS) -@pytest.mark.parametrize("path", (MAIN_PATH, BACKUP_PATH), ids=("main", "backup")) -def test_pro_tts_has_no_long_context_tier(model: str, path: Path): - info = _load(path)[model] - for field in LONG_CONTEXT_TIER_FIELDS: - assert field not in info, f"{model} has {field} but Google publishes one flat TTS rate" - - @pytest.mark.parametrize("model", ALL_KEYS) def test_backup_matches_main(model: str): assert _load(BACKUP_PATH)[model] == _load(MAIN_PATH)[model] diff --git a/tests/test_litellm/test_gpt_5_4_model_metadata.py b/tests/test_litellm/test_gpt_5_4_model_metadata.py index f93e6187dcb..294d0757069 100644 --- a/tests/test_litellm/test_gpt_5_4_model_metadata.py +++ b/tests/test_litellm/test_gpt_5_4_model_metadata.py @@ -37,43 +37,6 @@ def _pricing_key(model: str) -> str: return "gpt-5.4-nano" if "nano" in model else "gpt-5.4-mini" -@pytest.mark.parametrize("model", SMALL_MODELS) -def test_gpt_5_4_small_models_use_documented_token_limits(model: str) -> None: - """gpt-5.4-mini/nano are 400K-window models: 272K in, 128K out, not gpt-5.4's 1.05M window.""" - info = _load(MAIN_PATH).get(model) - assert info is not None, f"{model} not found in model_prices_and_context_window.json" - - assert info["max_input_tokens"] == DOCUMENTED_MAX_INPUT_TOKENS - assert info["max_output_tokens"] == DOCUMENTED_MAX_OUTPUT_TOKENS - assert info["max_tokens"] == DOCUMENTED_MAX_OUTPUT_TOKENS - - -@pytest.mark.parametrize("model", SMALL_MODELS) -def test_gpt_5_4_small_models_have_no_long_context_surcharge(model: str) -> None: - """OpenAI prices prompts above 272K at 2x input / 1.5x output for the 1.05M-window models only.""" - info = _load(MAIN_PATH)[model] - assert [key for key in info if "above_272k" in key] == [] - - -@pytest.mark.parametrize("model", SMALL_MODELS) -def test_gpt_5_4_small_models_standard_pricing(model: str) -> None: - info = _load(MAIN_PATH)[model] - input_cost, output_cost, cache_read_cost = STANDARD_PRICING[_pricing_key(model)] - - assert info["input_cost_per_token"] == input_cost - assert info["output_cost_per_token"] == output_cost - assert info["cache_read_input_token_cost"] == cache_read_cost - - -@pytest.mark.parametrize("model", LONG_CONTEXT_MODELS) -def test_gpt_5_4_long_context_models_keep_surcharge(model: str) -> None: - """The mini/nano correction must leave gpt-5.4 and gpt-5.4-pro tiered pricing intact.""" - info = _load(MAIN_PATH)[model] - - assert info["input_cost_per_token_above_272k_tokens"] == pytest.approx(info["input_cost_per_token"] * 2) - assert info["output_cost_per_token_above_272k_tokens"] == pytest.approx(info["output_cost_per_token"] * 1.5) - - @pytest.mark.parametrize("model", SMALL_MODELS) def test_gpt_5_4_small_models_backup_matches_main(model: str) -> None: assert _load(BACKUP_PATH).get(model) == _load(MAIN_PATH).get(model), ( diff --git a/tests/test_litellm/test_mistral_zai_glm_5_2_model_metadata.py b/tests/test_litellm/test_mistral_zai_glm_5_2_model_metadata.py index ad1f3b06e15..29576eb0119 100644 --- a/tests/test_litellm/test_mistral_zai_glm_5_2_model_metadata.py +++ b/tests/test_litellm/test_mistral_zai_glm_5_2_model_metadata.py @@ -4,7 +4,6 @@ from pathlib import Path import pytest import litellm -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider from litellm.types.utils import PromptTokensDetailsWrapper, Usage from litellm.utils import supports_prompt_caching, supports_reasoning @@ -35,34 +34,6 @@ def local_model_cost_map(monkeypatch): litellm.get_model_info.cache_clear() -@pytest.mark.parametrize("model", GLM_5_2_MODELS) -def test_zai_glm_5_2_specs(model): - info = _load(MAIN_PATH).get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - - assert info["litellm_provider"] == "mistral" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == INPUT_COST - assert info["output_cost_per_token"] == OUTPUT_COST - assert info["cache_read_input_token_cost"] == CACHED_INPUT_COST - - assert info["max_input_tokens"] == 1048576 - assert info["max_output_tokens"] == 131072 - assert info["max_tokens"] == 131072 - - assert info["supports_assistant_prefill"] is True - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - - routed_model, provider, _, _ = get_llm_provider(model=model) - assert routed_model == model.split("/", 1)[1] - assert provider == "mistral" - - @pytest.mark.parametrize("model", GLM_5_2_MODELS) def test_zai_glm_5_2_capabilities_are_visible_to_callers(local_model_cost_map, model): """Mistral advertises reasoning and prompt caching on this model, so the helpers diff --git a/tests/test_litellm/test_muse_spark_1_1_model_metadata.py b/tests/test_litellm/test_muse_spark_1_1_model_metadata.py index 540b97884dc..f55266a78d7 100644 --- a/tests/test_litellm/test_muse_spark_1_1_model_metadata.py +++ b/tests/test_litellm/test_muse_spark_1_1_model_metadata.py @@ -7,40 +7,6 @@ MUSE_SPARK_MODEL = "meta/muse-spark-1.1" def test_muse_spark_1_1_model_info(): - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - info = model_cost.get(MUSE_SPARK_MODEL) - assert info is not None, f"{MUSE_SPARK_MODEL} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "meta" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == 1.25e-06 - assert info["output_cost_per_token"] == 4.25e-06 - assert info["cache_read_input_token_cost"] == 1.5e-07 - - assert info["max_input_tokens"] == 1048576 - assert info["max_output_tokens"] == 131072 - assert info["max_tokens"] == 131072 - - assert info["supports_function_calling"] is True - assert info["supports_parallel_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_pdf_input"] is True - assert info["supports_web_search"] is True - assert info["supports_minimal_reasoning_effort"] is True - assert info["supports_xhigh_reasoning_effort"] is True - - assert info["supported_endpoints"] == ["/v1/chat/completions", "/v1/responses", "/v1/messages"] - assert info["supported_modalities"] == ["text", "image", "video"] - assert info["supported_output_modalities"] == ["text"] - routed_model, provider, _, api_base = get_llm_provider(model=MUSE_SPARK_MODEL, api_key="sk-test") assert routed_model == "muse-spark-1.1" assert provider == "meta" diff --git a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py index bdb2dc26813..8027d64d1ed 100644 --- a/tests/test_litellm/test_openai_service_tier_long_context_pricing.py +++ b/tests/test_litellm/test_openai_service_tier_long_context_pricing.py @@ -78,38 +78,6 @@ def _load(path: Path) -> dict[str, dict[str, object]]: return json.load(f) -@pytest.mark.parametrize("path", [MAIN_PATH, BACKUP_PATH], ids=["main", "backup"]) -@pytest.mark.parametrize("model", sorted(EXPECTED)) -def test_service_tier_long_context_rates_are_published(model: str, path: Path) -> None: - """Each tier must carry its own above-272K rates, in both price files.""" - info = _load(path).get(model) - assert info is not None, f"{model} not found in {path.name}" - for key, expected in EXPECTED[model].items(): - assert info.get(key) == pytest.approx(expected), f"{model}.{key} is {info.get(key)!r}, expected {expected!r}" - - -@pytest.mark.parametrize("model", sorted(EXPECTED)) -def test_tier_long_context_rate_is_half_or_double_the_standard(model: str) -> None: - """Flex is half the standard long-context rate; priority is double it.""" - info = _load(MAIN_PATH)[model] - tier = "flex" if model in FLEX_LONG_CONTEXT else "priority" - ratio = 0.5 if tier == "flex" else 2.0 - for base in ("input_cost_per_token", "output_cost_per_token"): - standard = info[f"{base}_above_272k_tokens"] - tiered = info[f"{base}_above_272k_tokens_{tier}"] - assert tiered == pytest.approx(standard * ratio), ( - f"{model}.{base}_above_272k_tokens_{tier} is {tiered!r}, " - f"expected {ratio}x the standard long-context rate {standard!r}" - ) - - -@pytest.mark.parametrize("model", NO_PUBLISHED_PRIORITY_LONG_CONTEXT) -def test_no_priority_long_context_rates_where_openai_publishes_none(model: str) -> None: - """Guard against back-filling a rate OpenAI does not publish.""" - info = _load(MAIN_PATH)[model] - assert "input_cost_per_token_above_272k_tokens_priority" not in info - - LONG_CONTEXT_PROMPT_TOKENS = 300_000 COMPLETION_TOKENS = 1_000 diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index c46a080976c..fc682145aca 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -7,7 +7,7 @@ import os import sys import threading from collections.abc import Awaitable, Callable, Mapping -from datetime import datetime +from datetime import datetime, timedelta from types import SimpleNamespace from typing import Final, Literal from unittest.mock import AsyncMock, MagicMock, patch @@ -1360,6 +1360,56 @@ def test_add_invalid_provider_to_router(): assert router.pattern_router.patterns == {} +@pytest.fixture +def registered_custom_provider(monkeypatch: pytest.MonkeyPatch) -> str: + from litellm import CustomLLM + from litellm.types.utils import ModelResponse + + class OnPremLLM(CustomLLM): + def completion(self, *args, **kwargs) -> ModelResponse: + return litellm.completion( + model="gpt-5.6", messages=[{"role": "user", "content": "hi"}], mock_response="served by onprem handler" + ) + + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "test-onprem-llm", "custom_handler": OnPremLLM()}]) + monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list)) + monkeypatch.setattr(litellm, "_custom_providers", list(litellm._custom_providers)) + return "test-onprem-llm" + + +def test_router_init_accepts_custom_provider_map_prefix_before_first_completion(registered_custom_provider: str): + assert registered_custom_provider not in litellm.provider_list + + router = litellm.Router( + model_list=[ + {"model_name": "onprem", "litellm_params": {"model": f"{registered_custom_provider}/my-model"}}, + ], + ) + + assert router.get_model_list(model_name="onprem")[0]["litellm_params"]["model"] == ( + f"{registered_custom_provider}/my-model" + ) + response = router.completion(model="onprem", messages=[{"role": "user", "content": "hi"}]) + assert response.choices[0].message.content == "served by onprem handler" + + +def test_router_add_deployment_accepts_explicit_custom_provider_from_custom_provider_map( + registered_custom_provider: str, +): + from litellm.types.router import Deployment + + router = litellm.Router(model_list=[]) + + router.add_deployment( + Deployment( + model_name="onprem", + litellm_params={"model": "my-model", "custom_llm_provider": registered_custom_provider}, + ) + ) + + assert router.get_model_list(model_name="onprem")[0]["litellm_params"]["model"] == "my-model" + + @pytest.mark.asyncio async def test_router_ageneric_api_call_with_fallbacks_helper(): """ @@ -8520,6 +8570,106 @@ class TestAdvisorSubCallCooldown: assert "dep-1" not in self._cooled_down_ids(router) +class TestCallerTimeoutCooldown: + """A timeout the caller set (the proxy's `timeout` body field or x-litellm-timeout + header) comes back as a 408 whatever the deployment's health, so it must neither + count toward allowed_fails nor bench the deployment. A 408 without that marker, or + one that arrives before the caller's deadline could have fired, is the provider's + and keeps cooling the deployment down.""" + + def _router(self): + return litellm.Router( + model_list=[ + { + "model_name": "slow-model", + "litellm_params": {"model": "openai/gpt-5.6", "api_key": "sk-fake"}, + "model_info": {"id": "dep-1"}, + } + ], + allowed_fails=0, + cooldown_time=120, + num_retries=0, + ) + + def _kwargs(self, marker, started=None, ended=None): + exception = litellm.Timeout(message="Request timed out", model="gpt-5.6", llm_provider="openai") + return { + "exception": exception, + "api_call_start_time": started, + "end_time": ended, + "litellm_params": {"model_info": {"id": "dep-1"}, "metadata": {}, **marker}, + } + + def _fail_count(self, router): + from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + get_deployment_failures_for_current_minute, + ) + + return get_deployment_failures_for_current_minute(litellm_router_instance=router, deployment_id="dep-1") + + def _cooled_down_ids(self, router): + active = router.cooldown_cache.get_active_cooldowns(model_ids=["dep-1"], parent_otel_span=None) + return [entry[0] for entry in active] + + @pytest.mark.asyncio + async def test_caller_timeout_408_leaves_failure_counter_and_cooldown_untouched(self): + router = self._router() + started = datetime.now() + ended = started + timedelta(seconds=2.05) + kwargs = self._kwargs({"client_side_timeout": True, "timeout": 2}, started=started, ended=ended) + assert router.deployment_callback_on_failure(kwargs, None, started, ended) is False + assert self._fail_count(router) == 0 + assert self._cooled_down_ids(router) == [] + + @pytest.mark.asyncio + async def test_provider_timeout_408_still_counts_and_cools_down(self): + router = self._router() + now = datetime.now() + assert router.deployment_callback_on_failure(self._kwargs({}), None, now, now) is True + assert self._fail_count(router) == 1 + assert self._cooled_down_ids(router) == ["dep-1"] + + @pytest.mark.asyncio + async def test_provider_408_before_caller_deadline_still_counts_and_cools_down(self): + """The marker only says the caller configured a timeout. A 408 that comes back + well before that deadline was raised by the provider, so it is a real health + signal and must not hide behind the caller's timeout.""" + router = self._router() + started = datetime.now() + ended = started + timedelta(seconds=0.4) + kwargs = self._kwargs({"client_side_timeout": True, "timeout": 30}, started=started, ended=ended) + assert router.deployment_callback_on_failure(kwargs, None, started, ended) is True + assert self._fail_count(router) == 1 + assert self._cooled_down_ids(router) == ["dep-1"] + + @pytest.mark.asyncio + async def test_caller_timeout_marker_reaches_failure_callback_end_to_end(self): + router = self._router() + seen = [] + recorded = threading.Event() + + def record(kwargs, completion_response, start_time, end_time): + seen.append(kwargs) + recorded.set() + + litellm.failure_callback.append(record) + try: + with pytest.raises(litellm.Timeout): + await router.acompletion( + model="slow-model", + messages=[{"role": "user", "content": "hello"}], + mock_timeout=True, + timeout=0.001, + client_side_timeout=True, + ) + assert await asyncio.to_thread(recorded.wait, 5) + finally: + litellm.failure_callback.remove(record) + assert seen[0]["litellm_params"]["client_side_timeout"] is True + assert self._fail_count(router) == 0 + assert self._cooled_down_ids(router) == [] + + def test_stream_chunks_have_generated_content_detects_text_and_non_text(): from litellm.router import _stream_chunks_have_generated_content from litellm.types.utils import ( diff --git a/tests/test_litellm/test_sambanova_model_metadata.py b/tests/test_litellm/test_sambanova_model_metadata.py index 972ddb4deef..20f34f9f3cc 100644 --- a/tests/test_litellm/test_sambanova_model_metadata.py +++ b/tests/test_litellm/test_sambanova_model_metadata.py @@ -11,15 +11,11 @@ def test_sambanova_minimax_m27_model_info(): model_cost = json.load(f) info = model_cost.get(model) - assert ( - info is not None - ), f"{model} not found in model_prices_and_context_window.json" + assert info is not None, f"{model} not found in model_prices_and_context_window.json" assert info["litellm_provider"] == "sambanova" assert info["mode"] == "chat" assert info["input_cost_per_token"] > 0 assert info["output_cost_per_token"] > 0 - assert info["max_input_tokens"] == 196608 - assert info["max_output_tokens"] == 131072 assert info["supports_function_calling"] is True assert info["supports_reasoning"] is True assert info["supports_tool_choice"] is True diff --git a/tests/test_litellm/test_together_ai_model_metadata.py b/tests/test_litellm/test_together_ai_model_metadata.py index b9764eca2f8..99e93ae2865 100644 --- a/tests/test_litellm/test_together_ai_model_metadata.py +++ b/tests/test_litellm/test_together_ai_model_metadata.py @@ -88,13 +88,6 @@ def test_together_chat_entries_never_carry_context_length_as_output_ceiling(cost assert inflated == [] -@pytest.mark.parametrize("model", sorted(DEPRECATED_MODELS)) -def test_together_deprecated_model_carries_deprecation_date(cost_map: CostMap, model: str): - info = cost_map.get(model) - assert info is not None, f"{model} missing from model_prices_and_context_window.json" - assert info.get("deprecation_date") == DEPRECATED_MODELS[model] - - def _successor(info: dict[str, object]) -> str | None: metadata = info.get("metadata") if not isinstance(metadata, dict): diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index 8bf8489fc52..02196a9cd26 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -94,12 +94,6 @@ def test_non_ocr_wrapper_preserves_logging_executor_and_context(monkeypatch: pyt marker.reset(token) -def test_cloudflare_model_info_includes_rpm(local_model_cost_map: None) -> None: - assert litellm.get_model_info("cloudflare/@cf/meta/llama-3.1-8b-instruct-fp8")["rpm"] == 300 - assert litellm.get_model_info("cloudflare/@cf/moonshotai/kimi-k2.6")["rpm"] == 20 - assert litellm.get_model_info("cloudflare/@cf/openai/whisper-large-v3-turbo")["rpm"] == 720 - - def test_get_utc_datetime_returns_current_aware_utc_time() -> None: before: Final = datetime.now(timezone.utc) result: Final = litellm.utils.get_utc_datetime() @@ -160,7 +154,6 @@ def test_prompt_tokens_details_cache_write_creation_stay_in_sync_on_assignment() assert details.cache_write_tokens == details.cache_creation_tokens == 375 - def test_get_model_info_surfaces_supports_adaptive_thinking(local_model_cost_map): """supports_adaptive_thinking must flow through get_model_info like every other capability flag: both from an explicit cost-map entry and from a @@ -177,7 +170,6 @@ def test_get_model_info_surfaces_supports_adaptive_thinking(local_model_cost_map assert generalized["supports_adaptive_thinking"] is True - def test_get_model_info_surfaces_supports_parallel_function_calling(local_model_cost_map): """A registry entry's supports_parallel_function_calling must read back through get_model_info and litellm.supports_parallel_function_calling. Regression: the key was never copied into @@ -493,64 +485,6 @@ def test_gpt_image_provider_detection_covers_existing_family(): assert custom_llm_provider == "openai" -def test_gpt_image_2_provider_and_model_info(local_model_cost_map): - - model, custom_llm_provider, _, _ = litellm.get_llm_provider(model="gpt-image-2") - - assert model == "gpt-image-2" - assert custom_llm_provider == "openai" - - model_info = litellm.get_model_info(model="gpt-image-2") - assert model_info["litellm_provider"] == "openai" - assert model_info["mode"] == "image_generation" - assert model_info["input_cost_per_token"] == 5e-06 - assert model_info["input_cost_per_image_token"] == 8e-06 - assert model_info["output_cost_per_token"] == 0 - assert model_info["output_cost_per_image_token"] == 3e-05 - assert ( - "/v1/images/generations" - in litellm.model_cost["gpt-image-2"]["supported_endpoints"] - ) - assert ( - "/v1/images/edits" in litellm.model_cost["gpt-image-2"]["supported_endpoints"] - ) - assert model_info["supports_vision"] is True - assert model_info["supports_pdf_input"] is True - - -def test_gpt_image_2_snapshot_model_info(local_model_cost_map): - model, custom_llm_provider, _, _ = litellm.get_llm_provider( - model="gpt-image-2-2026-04-21" - ) - - assert model == "gpt-image-2-2026-04-21" - assert custom_llm_provider == "openai" - - model_info = litellm.get_model_info(model="gpt-image-2-2026-04-21") - assert model_info["litellm_provider"] == "openai" - assert model_info["mode"] == "image_generation" - assert model_info["output_cost_per_image_token"] == 3e-05 - - -def test_azure_gpt_image_2_model_info(local_model_cost_map): - model, custom_llm_provider, _, _ = litellm.get_llm_provider( - model="azure/gpt-image-2" - ) - - assert model == "gpt-image-2" - assert custom_llm_provider == "azure" - - model_info = litellm.get_model_info( - model="gpt-image-2", custom_llm_provider="azure" - ) - assert model_info["litellm_provider"] == "azure" - assert model_info["mode"] == "image_generation" - assert model_info["input_cost_per_token"] == 5e-06 - assert model_info["input_cost_per_image_token"] == 8e-06 - assert model_info["output_cost_per_token"] == 0 - assert model_info["output_cost_per_image_token"] == 3e-05 - - def test_all_model_configs(): from litellm.llms.vertex_ai.vertex_ai_partner_models.ai21.transformation import ( VertexAIAi21Config, @@ -2907,158 +2841,6 @@ def test_model_info_for_vertex_ai_deepseek_model(): print("vertex deepseek model info", model_info) -def test_model_info_for_openrouter_kimi_k2_5(): - """ - Test that openrouter/moonshotai/kimi-k2.5 model info is correctly configured - in model_prices_and_context_window.json. - - Model properties from OpenRouter API: - - context_length: 262144 - - pricing: prompt=$0.00000045, completion=$0.00000225, input_cache_read=$0.00000007 - - modality: text+image->text (supports vision) - - supports: tool_choice, tools (function calling) - """ - import json - from pathlib import Path - - # Load directly from the local JSON file - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - model_info = model_cost.get("openrouter/moonshotai/kimi-k2.5") - assert ( - model_info is not None - ), "Model not found in model_prices_and_context_window.json" - assert model_info["litellm_provider"] == "openrouter" - assert model_info["mode"] == "chat" - - # Verify context window - assert model_info["max_input_tokens"] == 262144 - assert model_info["max_output_tokens"] == 262144 - assert model_info["max_tokens"] == 262144 - - # Verify pricing - assert model_info["input_cost_per_token"] == 4.5e-07 - assert model_info["output_cost_per_token"] == 2.25e-06 - assert model_info["cache_read_input_token_cost"] == 7e-08 - - # Verify capabilities - assert model_info["supports_vision"] is True - assert model_info["supports_function_calling"] is True - assert model_info["supports_tool_choice"] is True - - print("openrouter kimi-k2.5 model info", model_info) - - -def test_gemini_embedding_2_ga_in_cost_map(): - """GA and Vertex preview gemini-embedding-2 entries align with multimodal token pricing.""" - import json - from pathlib import Path - - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - for key, provider in ( - ("gemini/gemini-embedding-2", "gemini"), - ("vertex_ai/gemini-embedding-2", "vertex_ai"), - ("vertex_ai/gemini-embedding-2-preview", "vertex_ai"), - ("gemini-embedding-2", "vertex_ai-embedding-models"), - ): - info = model_cost.get(key) - assert ( - info is not None - ), f"{key} missing from model_prices_and_context_window.json" - assert info["litellm_provider"] == provider - assert info.get("mode") == "embedding" - assert info.get("supports_multimodal") is True - assert info.get("input_cost_per_token") == 2e-07 - assert info.get("input_cost_per_audio_token") == 6.5e-06 - assert info.get("input_cost_per_image_token") == 4.5e-07 - assert info.get("input_cost_per_video_token") == 1.2e-05 - assert info.get("input_cost_per_audio_token_batches") == 3.25e-06 - assert info.get("input_cost_per_image_token_batches") == 2.25e-07 - assert info.get("input_cost_per_video_token_batches") == 6e-06 - assert "input_cost_per_image" not in info - assert "input_cost_per_audio_per_second" not in info - assert "input_cost_per_video_per_second" not in info - if provider in ("vertex_ai-embedding-models", "vertex_ai"): - assert ( - info.get("uses_embed_content") is True - ), f"{key} must have uses_embed_content=true for correct Vertex AI routing" - - -def test_gemini_lyria_3_preview_models_in_cost_map(): - import json - from pathlib import Path - - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - clip = model_cost.get("gemini/lyria-3-clip-preview") - pro = model_cost.get("gemini/lyria-3-pro-preview") - assert clip is not None and pro is not None - assert clip["litellm_provider"] == "gemini" and pro["litellm_provider"] == "gemini" - assert clip["max_input_tokens"] == 131072 == pro["max_input_tokens"] - assert clip["output_cost_per_image"] == 0.04 - - -def test_vertex_ai_lyria_models_in_cost_map(): - import json - from pathlib import Path - - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - lyria_2 = model_cost.get("vertex_ai/lyria-002") - clip = model_cost.get("vertex_ai/lyria-3-clip-preview") - pro = model_cost.get("vertex_ai/lyria-3-pro-preview") - - assert lyria_2 is not None - assert clip is not None - assert pro is not None - assert lyria_2["litellm_provider"] == "vertex_ai" - assert clip["litellm_provider"] == "vertex_ai" - assert pro["litellm_provider"] == "vertex_ai" - assert lyria_2["mode"] == "audio_speech" - assert clip["mode"] == "audio_speech" - assert pro["mode"] == "audio_speech" - assert lyria_2["output_cost_per_image"] == 0.06 - assert lyria_2["supported_modalities"] == ["text"] - assert lyria_2["supported_output_modalities"] == ["audio"] - assert lyria_2["supports_audio_output"] is True - assert lyria_2["supported_audio_formats"] == ["wav"] - assert lyria_2["vertex_ai_audio_api"] == "lyria_predict" - assert lyria_2["supported_endpoints"] == ["/v1/audio/speech"] - assert clip["output_cost_per_image"] == 0.04 - assert pro["output_cost_per_image"] == 0.08 - assert clip["supported_audio_formats"] == ["mp3"] - assert pro["supported_audio_formats"] == ["mp3", "wav"] - assert clip["vertex_ai_audio_api"] == "lyria_interactions" - assert pro["vertex_ai_audio_api"] == "lyria_interactions" - assert clip["supported_endpoints"] == [ - "/v1beta/interactions", - "/v1/audio/speech", - ] - assert pro["supported_endpoints"] == [ - "/v1beta/interactions", - "/v1/audio/speech", - ] - assert clip["supported_modalities"] == ["text"] - assert pro["supported_modalities"] == ["text"] - assert clip["supports_vision"] is False - assert pro["supports_vision"] is False - assert "supports_image_input" not in clip - assert "supports_image_input" not in pro - assert clip["supported_regions"] == ["global"] - assert pro["supported_regions"] == ["global"] - assert clip["supports_audio_output"] is True - assert pro["supports_audio_output"] is True - - def test_model_info_for_fireworks_short_form_models(): """ Test that fireworks_ai short-form model entries (fireworks_ai/) @@ -4180,114 +3962,6 @@ class TestValidateAndFixThinkingParam: assert validate_and_fix_thinking_param(thinking=False) is None -def test_deepseek_v4_models_in_cost_map(): - """ - Test that deepseek-v4-flash and deepseek-v4-pro entries are correctly - configured in model_prices_and_context_window.json. - - Prices sourced from https://api-docs.deepseek.com/quick_start/pricing: - - deepseek-v4-flash: $0.30/M input, $1.20/M output - - deepseek-v4-pro: $1.32/M input, $3.96/M output - - Closes https://github.com/BerriAI/litellm/issues/26709 - """ - import json - from pathlib import Path - - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - # --- bare model names --- - for key, expected_input, expected_output, expected_cache, expected_vision in [ - ("deepseek-v4-flash", 3e-07, 1.2e-06, 6e-09, True), - ("deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08, False), - ]: - info = model_cost.get(key) - assert info is not None, f"{key} missing from model_prices_and_context_window.json" - assert info["litellm_provider"] == "deepseek" - assert info["mode"] == "chat" - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - assert info["cache_read_input_token_cost"] == expected_cache - assert info["max_input_tokens"] == 1_000_000 - assert info["supports_function_calling"] is True - assert info["supports_tool_choice"] is True - assert info.get("supports_vision", False) is expected_vision - - # --- provider-prefixed names --- - for key, expected_input, expected_output, expected_cache, expected_vision in [ - ("deepseek/deepseek-v4-flash", 3e-07, 1.2e-06, 6e-09, True), - ("deepseek/deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08, False), - ]: - info = model_cost.get(key) - assert info is not None, f"{key} missing from model_prices_and_context_window.json" - assert info["litellm_provider"] == "deepseek" - assert info["mode"] == "chat" - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - assert info["cache_read_input_token_cost"] == expected_cache - assert info["supports_function_calling"] is True - assert info["supports_tool_choice"] is True - assert info.get("supports_vision", False) is expected_vision - - -def test_deepseek_v4_models_in_backup_cost_map(): - """ - Test that deepseek-v4-flash and deepseek-v4-pro entries are correctly - configured in litellm/model_prices_and_context_window_backup.json. - """ - import json - from pathlib import Path - - json_path = Path(__file__).parents[2] / "litellm" / "model_prices_and_context_window_backup.json" - with open(json_path) as f: - model_cost = json.load(f) - - # --- bare model names --- - for key, expected_input, expected_output, expected_cache, expected_vision in [ - ("deepseek-v4-flash", 3e-07, 1.2e-06, 6e-09, True), - ("deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08, False), - ]: - info = model_cost.get(key) - assert info is not None, f"{key} missing from backup JSON" - assert info["litellm_provider"] == "deepseek" - assert info["mode"] == "chat" - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - assert info["cache_read_input_token_cost"] == expected_cache - assert info["max_input_tokens"] == 1_000_000 - assert info.get("supports_vision", False) is expected_vision - - # --- provider-prefixed names --- - for key, expected_input, expected_output, expected_cache, expected_vision in [ - ("deepseek/deepseek-v4-flash", 3e-07, 1.2e-06, 6e-09, True), - ("deepseek/deepseek-v4-pro", 1.32e-06, 3.96e-06, 4.4e-08, False), - ]: - info = model_cost.get(key) - assert info is not None, f"{key} missing from backup JSON" - assert info["litellm_provider"] == "deepseek" - assert info["mode"] == "chat" - assert info["input_cost_per_token"] == expected_input - assert info["output_cost_per_token"] == expected_output - assert info["cache_read_input_token_cost"] == expected_cache - assert info.get("supports_vision", False) is expected_vision - - -def test_deprecation_dates_for_retired_xai_and_groq_models(): - import json - from pathlib import Path - - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - assert model_cost["xai/grok-imagine-image-quality"]["deprecation_date"] == "2026-11-02" - assert model_cost["xai/grok-imagine-image-quality-latest"]["deprecation_date"] == "2026-11-02" - assert model_cost["xai/grok-imagine-image-quality-20260403"]["deprecation_date"] == "2026-11-02" - assert model_cost["groq/gemma-7b-it"]["deprecation_date"] == "2024-12-18" - - @pytest.mark.usefixtures("local_model_cost_map") def test_deepseek_flash_completion_cost(): from litellm.types.utils import ModelResponse @@ -4979,25 +4653,6 @@ def test_anthropic_reexport_entries_carry_explicit_prompt_cache_min_tokens(local assert not wrong, f"(cost-map value, resolved value) diverge from Anthropic's published minimums: {wrong}" -def test_anthropic_reexport_cache_minimums_present_in_root_cost_map() -> None: - """The root map ships to the CDN independently of the bundled backup, so both must carry the - minimum or proxies reading one of them regress to the 1024 default.""" - root_map_path: Final = os.path.join(os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json") - with open(root_map_path) as f: - root_map: Final = json.load(f) - wrong: Final = { - model: root_map[model].get("prompt_cache_min_tokens") - for model, expected in ANTHROPIC_REEXPORT_CACHE_MIN.items() - if root_map[model].get("prompt_cache_min_tokens") != expected - } - fable_5_wrong: Final = { - model: info.get("prompt_cache_min_tokens") - for model, info in root_map.items() - if "fable-5" in model and info.get("supports_prompt_caching") and info.get("prompt_cache_min_tokens") != 512 - } - assert not wrong and not fable_5_wrong, f"root cost map diverges: {wrong | fable_5_wrong}" - - GEMINI_4096_CACHE_MIN_MODELS: Final = tuple( prefix + base for base in ( @@ -5024,20 +4679,6 @@ def test_gemini_3_flash_and_31_pro_preview_resolve_4096_cache_minimum(local_mode assert not wrong, f"prompt_cache_min_tokens must be 4096: {wrong}" -def test_gemini_4096_cache_minimum_present_in_root_cost_map() -> None: - """The root map ships to the CDN independently of the bundled backup, so both must carry the - minimum or proxies reading one of them regress to the 1024 default.""" - root_map_path: Final = os.path.join(os.path.dirname(__file__), "..", "..", "model_prices_and_context_window.json") - with open(root_map_path) as f: - root_map: Final = json.load(f) - wrong: Final = { - model: root_map[model].get("prompt_cache_min_tokens") - for model in GEMINI_4096_CACHE_MIN_MODELS - if root_map[model].get("prompt_cache_min_tokens") != 4096 - } - assert not wrong, f"prompt_cache_min_tokens must be 4096: {wrong}" - - def test_get_prompt_cache_min_tokens_unmapped_model_falls_back_to_default(local_model_cost_map: None) -> None: """get_model_info raises for a model it has no entry for. The resolver must swallow that and fall back to the default, otherwise the raise reaches callers that would read it as @@ -6508,7 +6149,6 @@ async def test_async_mock_completion_streaming_obj_raises_mock_exception_before_ await _async_mock_stream_snapshots(mock_exception, 51234) - @contextlib.contextmanager def _recording_hidden_params_at_submit(submit_target: str) -> "Iterator[queue.SimpleQueue[dict[str, object]]]": seen: Final = queue.SimpleQueue() diff --git a/tests/test_litellm/test_xai_grok_4_3_model_metadata.py b/tests/test_litellm/test_xai_grok_4_3_model_metadata.py index 81e7f4adf1f..e6e4eada1b6 100644 --- a/tests/test_litellm/test_xai_grok_4_3_model_metadata.py +++ b/tests/test_litellm/test_xai_grok_4_3_model_metadata.py @@ -1,49 +1,6 @@ import json from pathlib import Path -import pytest - -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider - - -@pytest.mark.parametrize("model", ["xai/grok-4.3", "xai/grok-4.3-latest"]) -def test_xai_grok_4_3_model_info(model): - json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" - with open(json_path) as f: - model_cost = json.load(f) - - info = model_cost.get(model) - assert ( - info is not None - ), f"{model} not found in model_prices_and_context_window.json" - - assert info["litellm_provider"] == "xai" - assert info["mode"] == "chat" - - assert info["input_cost_per_token"] == 1.25e-06 - assert info["output_cost_per_token"] == 2.5e-06 - assert info["cache_read_input_token_cost"] == 2e-07 - - assert info["input_cost_per_token_above_200k_tokens"] == 2.5e-06 - assert info["output_cost_per_token_above_200k_tokens"] == 5e-06 - assert info["cache_read_input_token_cost_above_200k_tokens"] == 4e-07 - - assert info["max_input_tokens"] == 1000000 - assert info["max_output_tokens"] == 1000000 - assert info["max_tokens"] == 1000000 - - assert info["supports_function_calling"] is True - assert info["supports_prompt_caching"] is True - assert info["supports_reasoning"] is True - assert info["supports_response_schema"] is True - assert info["supports_tool_choice"] is True - assert info["supports_vision"] is True - assert info["supports_web_search"] is True - - routed_model, provider, _, _ = get_llm_provider(model=model) - assert routed_model == model.split("/", 1)[1] - assert provider == "xai" - def test_xai_grok_4_3_backup_matches_main(): """Ensure the bundled model cost map stays in sync with the canonical file.""" diff --git a/tests/test_litellm/test_xai_responses_auto_routing.py b/tests/test_litellm/test_xai_responses_auto_routing.py index 5b1944dcb8b..fbf2453d7fb 100644 --- a/tests/test_litellm/test_xai_responses_auto_routing.py +++ b/tests/test_litellm/test_xai_responses_auto_routing.py @@ -204,6 +204,17 @@ class TestXAIResponsesAutoRouting: assert model_info.get("mode") == "responses" assert updated_model == model + def test_responses_api_bridge_check_with_web_search_options_on_unmapped_model(self): + """web search must reach /responses even for a model missing from the cost map, chat returns 410""" + model_info, updated_model = responses_api_bridge_check( + model="grok-not-in-cost-map", + custom_llm_provider="xai", + web_search_options={"search_context_size": "medium"}, + ) + + assert model_info.get("mode") == "responses" + assert updated_model == "grok-not-in-cost-map" + @patch("litellm.completion_extras.responses_api_bridge.completion") def test_completion_with_tools_routes_to_responses_api( self, mock_responses_completion diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx index c425e766f2d..ae898645de4 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTable.test.tsx @@ -129,7 +129,7 @@ describe("BudgetTable", () => { const user = userEvent.setup(); renderWithProviders(); await showColumn(user, "created_at"); - for (const field of ["budget_id", "max_budget", "tpm_limit", "rpm_limit", "created_at"]) { + for (const field of ["budget_id", "max_budget", "tpm_limit", "rpm_limit", "tpd_limit", "created_at"]) { expect(screen.getByTestId(`sort-header-${field}`)).toBeInTheDocument(); } }); @@ -152,9 +152,10 @@ describe("BudgetTable", () => { }); it("should show n/a for missing rate limits and Unlimited for a missing max budget", () => { - const list = makeList({ rows: [makeBudget({ max_budget: null, tpm_limit: null, rpm_limit: null })] }); + const noLimits = { max_budget: null, tpm_limit: null, rpm_limit: null, tpd_limit: null }; + const list = makeList({ rows: [makeBudget(noLimits)] }); renderWithProviders(); - expect(screen.getAllByText("n/a")).toHaveLength(2); + expect(screen.getAllByText("n/a")).toHaveLength(3); expect(screen.getByText("Unlimited")).toBeInTheDocument(); }); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx index e5cd9043492..fb894322208 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/BudgetTableColumns.tsx @@ -126,6 +126,14 @@ export const getBudgetTableColumns = ({ size: 100, cell: ({ row }) => , }, + { + id: "tpd_limit", + accessorKey: "tpd_limit", + meta: { title: "TPD (batch)", numeric: true }, + header: ({ column }) => , + size: 110, + cell: ({ row }) => , + }, { id: "budget_duration", accessorKey: "budget_duration", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx index 492a6b5c630..5068cbed453 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_modal.tsx @@ -17,6 +17,7 @@ const budgetShape = { budget_id: z.string().min(1, "Please input a human-friendly name for the budget"), tpm_limit: z.number().nullish(), rpm_limit: z.number().nullish(), + tpd_limit: z.number().nullish(), max_budget: z.number().nullish(), budget_duration: z.string().nullish(), }; @@ -112,6 +113,23 @@ const BudgetModal: React.FC = ({ isModalVisible, setIsModalVis /> )} + + {({ ref, value, onChange, ...field }) => ( + onChange(event.target.value === "" ? null : event.target.valueAsNumber)} + /> + )} + diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx index 25344c52847..7455c252e26 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/budget_panel.tsx @@ -133,6 +133,7 @@ const BudgetPanel: React.FC = ({ accessToken }) => { { label: "Max Budget", value: selectedBudget?.max_budget }, { label: "TPM", value: selectedBudget?.tpm_limit }, { label: "RPM", value: selectedBudget?.rpm_limit }, + { label: "TPD (batch)", value: selectedBudget?.tpd_limit }, ]} onCancel={handleDeleteCancel} onOk={handleDeleteConfirm} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx index 71fce2de836..1931a88f096 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/budgets/_components/edit_budget_modal.tsx @@ -15,13 +15,14 @@ import { Dialog, DialogContent, DialogHeader, DialogTitle } from "@/components/u type EditBudgetFormValues = Pick< budgetItem, - "budget_id" | "tpm_limit" | "rpm_limit" | "max_budget" | "budget_duration" + "budget_id" | "tpm_limit" | "rpm_limit" | "tpd_limit" | "max_budget" | "budget_duration" >; const toFormValues = (budget: budgetItem): EditBudgetFormValues => ({ budget_id: budget.budget_id, tpm_limit: budget.tpm_limit, rpm_limit: budget.rpm_limit, + tpd_limit: budget.tpd_limit, max_budget: budget.max_budget, budget_duration: budget.budget_duration, }); @@ -118,6 +119,23 @@ const EditBudgetModal: React.FC = ({ isModalVisible, setIs /> )} + + {({ ref, value, onChange, ...field }) => ( + onChange(event.target.value === "" ? null : event.target.valueAsNumber)} + /> + )} + diff --git a/ui/litellm-dashboard/src/components/Teams.test.tsx b/ui/litellm-dashboard/src/components/Teams.test.tsx index 70c30596c75..851c9e6d487 100644 --- a/ui/litellm-dashboard/src/components/Teams.test.tsx +++ b/ui/litellm-dashboard/src/components/Teams.test.tsx @@ -1187,6 +1187,7 @@ describe("Teams - which fields reach the create payload depends on the open sect "organization_id", "rpm_limit", "team_alias", + "tpd_limit", "tpm_limit", ]); expect(payload.team_alias).toBe("Closed Sections Team"); @@ -1314,6 +1315,7 @@ describe("Teams - the exact bytes the create call sends", () => { budget_duration: undefined, tpm_limit: undefined, rpm_limit: undefined, + tpd_limit: undefined, metadata: undefined, }); expect(wireBody(payload)).toStrictEqual({ @@ -1341,6 +1343,7 @@ describe("Teams - the exact bytes the create call sends", () => { budget_duration: undefined, tpm_limit: undefined, rpm_limit: undefined, + tpd_limit: undefined, metadata: undefined, team_id: undefined, team_member_budget: undefined, @@ -1513,6 +1516,7 @@ describe("Teams - the exact bytes the create call sends", () => { budget_duration: undefined, tpm_limit: undefined, rpm_limit: undefined, + tpd_limit: undefined, metadata: undefined, team_id: undefined, team_member_budget: undefined, diff --git a/ui/litellm-dashboard/src/components/Teams.tsx b/ui/litellm-dashboard/src/components/Teams.tsx index dc531ea5dad..4f3367d8b98 100644 --- a/ui/litellm-dashboard/src/components/Teams.tsx +++ b/ui/litellm-dashboard/src/components/Teams.tsx @@ -77,6 +77,7 @@ const teamCreateFieldsSchema = z.object({ budget_duration: z.string().nullish(), tpm_limit: numericInputSchema, rpm_limit: numericInputSchema, + tpd_limit: numericInputSchema, metadata: metadataPairsSchema.optional(), team_id: z.string().optional(), team_member_budget: z.number().optional(), @@ -113,6 +114,7 @@ const EMPTY_TEAM_CREATE_VALUES: TeamCreateFormValues = { budget_duration: undefined, tpm_limit: undefined, rpm_limit: undefined, + tpd_limit: undefined, metadata: [], team_id: undefined, team_member_budget: undefined, @@ -821,6 +823,18 @@ const Teams: React.FC = ({ accessToken, userID, userRole, premiumUser )} + + {({ ref, value, ...field }) => ( + + )} + Metadata - (["llm", "llm_v2", "heuristic_first", "hybrid", "capability"] as const).some((type) => type === classifierType); + (["llm", "heuristic_first", "hybrid", "capability", "llm_v2"] as const).some((type) => type === classifierType); export type ClassifierFallback = "heuristic" | "default_model"; @@ -183,7 +183,7 @@ export const heuristicScoringRoleFor = ( classifierType: ClassifierType, classifierFallback: ClassifierFallback | undefined, ): HeuristicScoringRole => { - if (classifierType === "heuristic_v2" || classifierType === "llm_v2" || classifierType === "capability") + if (classifierType === "heuristic_v2" || classifierType === "capability" || classifierType === "llm_v2") return "never"; if (classifierType === "heuristic" || classifierType === "heuristic_first" || classifierType === "hybrid") return "decides"; diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index 56ea29a043f..eadbca87140 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -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 | null; budget_reset_at?: string | null; @@ -50,6 +51,7 @@ export interface KeyResponse { metadata: Record; tpm_limit: number; rpm_limit: number; + tpd_limit?: number | null; duration: string; budget_duration: string; budget_reset_at: string; diff --git a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts index fef94cc3c2b..7c6d5def8da 100644 --- a/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts +++ b/ui/litellm-dashboard/src/components/organisms/createKeyPayload.test.ts @@ -45,6 +45,7 @@ const DROPPED_AT_SERIALISATION = [ "rpm_limit", "tags", "throttle_on_budget_exceeded", + "tpd_limit", "tpm_limit", ]; @@ -64,6 +65,7 @@ const OPTIONAL_SETTINGS_VALUES = { tpm_limit_type: "key", rpm_limit: undefined, rpm_limit_type: "key", + tpd_limit: undefined, throttle_on_budget_exceeded: undefined, enable_prompt_caching: undefined, guardrails: undefined, @@ -456,6 +458,18 @@ describe("budget duration", () => { }); }); +describe("tpd_limit", () => { + it("forwards the daily batch token budget alongside the minute limits", () => { + expect(payloadOf(build({ key_alias: "my-key", tpd_limit: 250000, rpm_limit: 5 }))).toStrictEqual( + aliasOnly({ tpd_limit: 250000, rpm_limit: 5 }), + ); + }); + + it("keeps a zero tpd_limit rather than treating it as unset", () => { + expect(payloadOf(build({ key_alias: "my-key", tpd_limit: 0 }))).toStrictEqual(aliasOnly({ tpd_limit: 0 })); + }); +}); + describe("purity", () => { it("leaves the submitted form values untouched", () => { const values = { @@ -499,9 +513,9 @@ describe("serialised wire shape", () => { expect(payloadOf(build({ ...CLOSED_SECTIONS_VALUES, team_id: "team-1" })).team_id).toBe("team-1"); }); - it("adds fifteen keys to the object and only the two limit types to the wire when Optional Settings opens", () => { + it("adds sixteen keys to the object and only the two limit types to the wire when Optional Settings opens", () => { const payload = payloadOf(build(OPTIONAL_SETTINGS_VALUES)); - expect(Object.keys(payload)).toHaveLength(23); + expect(Object.keys(payload)).toHaveLength(24); expect(wireKeys(payload)).toStrictEqual([ "team_id", "key_alias", diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx index 0d5d9f5ec8d..3e3c29e330d 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.integration.test.tsx @@ -143,6 +143,7 @@ const OPTIONAL_OPEN_PAYLOAD = { tpm_limit_type: null, rpm_limit: undefined, rpm_limit_type: null, + tpd_limit: undefined, throttle_on_budget_exceeded: undefined, enable_prompt_caching: undefined, guardrails: undefined, @@ -395,6 +396,7 @@ describe("CreateKey", () => { it.each([ ["Tokens per minute Limit (TPM)", "tpm_limit"], ["Requests per minute Limit (RPM)", "rpm_limit"], + ["Tokens per day Limit (TPD)", "tpd_limit"], ])("routes a typed %s into the %s payload key", async (label, key) => { await openModal(); await nameTheKey(); diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index b5789101f77..b8ea8de7f59 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1150,6 +1150,32 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp /> )} + + Tokens per day Limit (TPD){" "} + + + + + } + name="tpd_limit" + help={`TPD cannot exceed team TPD limit: ${team?.tpd_limit !== null && team?.tpd_limit !== undefined ? team?.tpd_limit : "unlimited"}`} + rules={ceilingRule( + team?.tpd_limit, + (limit) => `TPD limit cannot exceed team TPD limit: ${limit}`, + )} + > + {(control) => ( + + )} + @@ -1760,6 +1786,7 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp "budget_duration", "tpm_limit", "rpm_limit", + "tpd_limit", ...(disableCustomApiKeys ? ["key"] : []), ]} /> diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx index a25ca28651a..eb912ffa3cc 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.test.tsx @@ -1968,6 +1968,7 @@ describe("TeamInfoView - the exact bytes the update call sends", () => { models: ["gpt-4"], tpm_limit: 1000, rpm_limit: 1000, + tpd_limit: null, model_tpm_limit: {}, model_rpm_limit: {}, max_budget: 100, diff --git a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx index ffc83d0165e..ce705008678 100644 --- a/ui/litellm-dashboard/src/components/team/TeamInfo.tsx +++ b/ui/litellm-dashboard/src/components/team/TeamInfo.tsx @@ -264,6 +264,7 @@ export interface TeamData { metadata: Record; tpm_limit: number | null; rpm_limit: number | null; + tpd_limit?: number | null; max_budget: number | null; soft_budget?: number | null; budget_duration: string | null; @@ -330,6 +331,7 @@ const teamUpdateFieldsSchema = z.object({ budget_duration: z.string().nullish(), tpm_limit: numericInputSchema, rpm_limit: numericInputSchema, + tpd_limit: numericInputSchema, modelLimits: z .array( z.object({ @@ -411,6 +413,7 @@ const EMPTY_TEAM_UPDATE_VALUES: TeamUpdateFormValues = { budget_duration: undefined, tpm_limit: undefined, rpm_limit: undefined, + tpd_limit: undefined, modelLimits: [], default_estimated_output_tokens: undefined, default_estimated_output_tokens_per_model: "", @@ -460,6 +463,7 @@ const toTeamFormValues = (info: TeamInfoRecord, effectiveGuardrails: string[]): budget_duration: info.budget_duration, tpm_limit: info.tpm_limit, rpm_limit: info.rpm_limit, + tpd_limit: info.tpd_limit, modelLimits: Array.from( new Set([ ...Object.keys(info.metadata?.model_tpm_limit ?? {}), @@ -918,6 +922,7 @@ const TeamInfoView: React.FC = ({ models: normalizeTeamModelSelection(values.models), tpm_limit: sanitizeNumeric(values.tpm_limit), rpm_limit: sanitizeNumeric(values.rpm_limit), + tpd_limit: sanitizeNumeric(values.tpd_limit), model_tpm_limit: modelTpmLimit, model_rpm_limit: modelRpmLimit, max_budget: values.max_budget, @@ -1168,6 +1173,7 @@ const TeamInfoView: React.FC = ({

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

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

+

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

{info.max_parallel_requests &&

Max Parallel Requests: {info.max_parallel_requests}

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

Rate Limits

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

+

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

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

Throttle on budget exceeded: Yes

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

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

+

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

Max Parallel Requests:{" "} {currentKeyData.max_parallel_requests !== null diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 20a87fce98f..785a8c35058 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -1843,6 +1843,7 @@ export interface paths { * - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. * - tpm_limit: Optional[int] - The tokens per minute limit for the budget. * - rpm_limit: Optional[int] - The requests per minute limit for the budget. + * - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit. * - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} * - budget_reset_at: Optional[datetime] - Datetime when the initial budget is reset. Default is now. */ @@ -1899,6 +1900,7 @@ export interface paths { * - max_parallel_requests: Optional[int] - The max number of parallel requests for the budget. * - tpm_limit: Optional[int] - The tokens per minute limit for the budget. * - rpm_limit: Optional[int] - The requests per minute limit for the budget. + * - tpd_limit: Optional[int] - The tokens per day limit for the budget. Charged by batch submissions instead of tpm_limit/rpm_limit. * - model_max_budget: Optional[dict] - Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d", "tpm_limit": 100000, "rpm_limit": 100000}} * - budget_reset_at: Optional[datetime] - Update the Datetime when the budget was last reset. */ @@ -3980,6 +3982,7 @@ export interface paths { * - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). * - tpm_limit: Optional[int] - [Not Implemented Yet] Specify tpm limit for a given customer (Tokens per minute) * - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute) + * - tpd_limit: Optional[int] - Specify tpd limit for a given customer (Tokens per day). Batch submissions are charged against it instead of tpm_limit/rpm_limit * - model_max_budget: Optional[dict] - [Not Implemented Yet] Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}} * - max_parallel_requests: Optional[int] - [Not Implemented Yet] Specify max parallel requests for a given customer. * - soft_budget: Optional[float] - [Not Implemented Yet] Get alerts when customer crosses given budget, doesn't block requests. @@ -4514,6 +4517,7 @@ export interface paths { * - budget_duration: Optional[str] - Budget is reset at the end of specified duration. If not set, budget is never reset. You can set duration as seconds ("30s"), minutes ("30m"), hours ("30h"), days ("30d"). * - tpm_limit: Optional[int] - [Not Implemented Yet] Specify tpm limit for a given customer (Tokens per minute) * - rpm_limit: Optional[int] - [Not Implemented Yet] Specify rpm limit for a given customer (Requests per minute) + * - tpd_limit: Optional[int] - Specify tpd limit for a given customer (Tokens per day). Batch submissions are charged against it instead of tpm_limit/rpm_limit * - model_max_budget: Optional[dict] - [Not Implemented Yet] Specify max budget for a given model. Example: {"openai/gpt-4o-mini": {"max_budget": 100.0, "budget_duration": "1d"}} * - max_parallel_requests: Optional[int] - [Not Implemented Yet] Specify max parallel requests for a given customer. * - soft_budget: Optional[float] - [Not Implemented Yet] Get alerts when customer crosses given budget, doesn't block requests. @@ -7736,6 +7740,7 @@ export interface paths { * - blocked: Optional[bool] - Whether the key is blocked. * - rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute) * - tpm_limit: Optional[int] - Specify tpm limit for a given key (Tokens per minute) + * - tpd_limit: Optional[int] - Specify tpd limit for a given key (Tokens per day). Charged by batch submissions instead of tpm_limit/rpm_limit. * - soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached. * - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). * - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. @@ -8049,6 +8054,7 @@ export interface paths { * - blocked: Optional[bool] - Whether the key is blocked. * - rpm_limit: Optional[int] - Specify rpm limit for a given key (Requests per minute) * - tpm_limit: Optional[int] - Specify tpm limit for a given key (Tokens per minute) + * - tpd_limit: Optional[int] - Specify tpd limit for a given key (Tokens per day). Charged by batch submissions instead of tpm_limit/rpm_limit. * - soft_budget: Optional[float] - Specify soft budget for a given key. Will trigger a slack alert when this soft budget is reached. * - tags: Optional[List[str]] - Tags for [tracking spend](https://litellm.vercel.app/docs/proxy/enterprise#tracking-spend-for-custom-tags) and/or doing [tag-based routing](https://litellm.vercel.app/docs/proxy/tag_routing). * - enforced_params: Optional[List[str]] - List of enforced params for the key (Enterprise only). [Docs](https://docs.litellm.ai/docs/proxy/enterprise#enforce-required-params-for-llm-requests) @@ -8175,6 +8181,7 @@ export interface paths { * - metadata: Optional[dict] - Metadata for key. Example {"team": "core-infra", "app": "app2"} * - tpm_limit: Optional[int] - Tokens per minute limit * - rpm_limit: Optional[int] - Requests per minute limit + * - tpd_limit: Optional[int] - Tokens per day limit, charged by batch submissions instead of tpm_limit/rpm_limit * - model_rpm_limit: Optional[dict] - Model-specific RPM limits {"gpt-4": 100, "claude-v1": 200} * - mcp_rpm_limit: Optional[dict] - Per-MCP-server RPM limits, keyed by MCP server name {"github": 100, "slack": 200} * - tag_rpm_limit: Optional[dict] - Per-request-tag RPM limits, keyed by request tag {"cell-1": 1000, "cell-2": 500}. Each tag gets an independent counter; absent tags fall back to the key-level rpm limit. @@ -8424,7 +8431,7 @@ export interface paths { * way to page, sort or filter it. * * `sort` takes a comma-separated list of `budget_id`, `max_budget`, `tpm_limit`, - * `rpm_limit` or `created_at`, each optionally prefixed with `-` for descending, + * `rpm_limit`, `tpd_limit` or `created_at`, each optionally prefixed with `-` for descending, * and defaults to `-created_at`. `budget_id` is appended to every sort as the * tiebreaker. `q` is a case-insensitive substring match on `budget_id`. * `page_size` defaults to 50 and is capped at 100. Filters are @@ -10619,6 +10626,7 @@ export interface paths { * - max_budget: *Optional[float]* - Max budget for org * - tpm_limit: *Optional[int]* - Max tpm limit for org * - rpm_limit: *Optional[int]* - Max rpm limit for org + * - tpd_limit: *Optional[int]* - Max tokens per day stored on the org budget. Batch submissions enforce tpd_limit at the key, team and end user scopes only. * - model_rpm_limit: *Optional[Dict[str, int]]* - The RPM (Requests Per Minute) limit per model for this organization. * - model_tpm_limit: *Optional[Dict[str, int]]* - The TPM (Tokens Per Minute) limit per model for this organization. * - max_parallel_requests: *Optional[int]* - [Not Implemented Yet] Max parallel requests for org @@ -15605,6 +15613,7 @@ export interface paths { * - mcp_rpm_limit: Optional[Dict[str, int]] - Per-MCP-server RPM limit for this team, keyed by MCP server name (alias if set, else the configured name). Example: {"github": 100, "slack": 200}. Applied across all keys for this team. * - tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for this team - all keys with this team_id will have at max this TPM limit * - rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit + * - tpd_limit: Optional[int] - The TPD (Tokens Per Day) limit for this team. Batch submissions are charged against it instead of tpm_limit/rpm_limit * - rpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of RPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating RPM, or "best_effort_throughput" for best effort enforcement. * - tpm_limit_type: Optional[Literal["guaranteed_throughput", "best_effort_throughput"]] - The type of TPM limit enforcement. Use "guaranteed_throughput" to raise an error if overallocating TPM, or "best_effort_throughput" for best effort enforcement. * - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget @@ -15833,6 +15842,7 @@ export interface paths { * - metadata: Optional[dict] - Metadata for team, store information for team. Example metadata = {"team": "core-infra", "app": "app2", "email": "ishaan@berri.ai" } * - tpm_limit: Optional[int] - The TPM (Tokens Per Minute) limit for this team - all keys with this team_id will have at max this TPM limit * - rpm_limit: Optional[int] - The RPM (Requests Per Minute) limit for this team - all keys associated with this team_id will have at max this RPM limit + * - tpd_limit: Optional[int] - The TPD (Tokens Per Day) limit for this team. Batch submissions are charged against it instead of tpm_limit/rpm_limit * - max_budget: Optional[float] - The maximum budget allocated to the team - all keys for this team_id will have at max this max_budget * - soft_budget: Optional[float] - The soft budget threshold for the team. If max_budget is set (either in the request or existing), soft_budget must be strictly lower than max_budget. Can be set independently if max_budget is not set. * - budget_duration: Optional[str] - The duration of the budget for the team. Doc [here](https://docs.litellm.ai/docs/proxy/team_budgets) @@ -24582,6 +24592,8 @@ export interface components { rpm_limit?: number | null; /** Soft Budget */ soft_budget?: number | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** @@ -24634,6 +24646,11 @@ export interface components { * @description Requests will NOT fail if this is exceeded. Will fire alerting though. */ soft_budget?: number | null; + /** + * Tpd Limit + * @description Max tokens per day, charged by batch submissions, allowed for this budget id. + */ + tpd_limit?: number | null; /** * Tpm Limit * @description Max tokens per minute, allowed for this budget id. @@ -28335,6 +28352,8 @@ export interface components { team_id?: string | null; /** Throttle On Budget Exceeded */ throttle_on_budget_exceeded?: boolean | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Tpm Limit Type */ @@ -28495,6 +28514,8 @@ export interface components { token?: string | null; /** Token Id */ token_id?: string | null; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; /** Tpm Limit Type */ @@ -29299,6 +29320,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; }; @@ -29332,6 +29355,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; }; @@ -29440,6 +29465,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 */ @@ -29607,6 +29634,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 */ @@ -30800,6 +30829,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 */ @@ -31176,6 +31207,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 */ @@ -32825,6 +32858,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. @@ -33040,6 +33078,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; }; @@ -33161,6 +33201,8 @@ export interface components { tags?: string[] | null; /** Team Id */ team_id: string; + /** Tpd Limit */ + tpd_limit?: number | null; /** Tpm Limit */ tpm_limit?: number | null; }; @@ -33345,6 +33387,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 */ @@ -33638,6 +33682,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 */ @@ -34092,6 +34138,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; }; @@ -35501,6 +35549,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 */ @@ -35730,7 +35780,7 @@ export interface components { classifier_plugin_timeout_ms: number; /** * Classifier Type - * @description Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, an LLM tier-selection call, a Switchyard-compatible capability forecast, a joint task-demand and capability forecast, a custom classifier plugin, 'heuristic_first', which scores locally and only pays for the LLM classifier when the local scorer does not confidently land a cheap tier, or 'hybrid', which trusts the local scorer everywhere except when its score lands near a tier boundary + * @description Classification strategy: local regex/keyword scoring, the bundled trained four-tier heuristic, an LLM tier-selection call, a Switchyard-compatible capability forecast, a joint Fuse V2 forecast, a custom classifier plugin, 'heuristic_first', which scores locally and only pays for the LLM classifier when the local scorer does not confidently land a cheap tier, or 'hybrid', which trusts the local scorer everywhere except when its score lands near a tier boundary * @default heuristic * @enum {string} */ @@ -37135,22 +37185,34 @@ export interface components { * @enum {string} */ cause?: "heuristic_scorer" | "heuristic_v2" | "reasoning_override" | "llm_classifier" | "capability_classifier" | "llm_v2_classifier" | "llm_v2_fallback" | "heuristic_first_short_circuit" | "hybrid_short_circuit" | "classifier_plugin" | "classifier_fallback" | "capability_classifier_fallback" | "default_model_fallback" | "literal_keyword_match" | "semantic_keyword_match" | "plan_mode" | "housekeeping" | "modality_escalation" | "modality_pin_override" | "health_failover" | "health_default_fallback" | "session_affinity_pin" | "session_affinity_escalation" | "user_turn_continuation" | "default_fallback" | "keyword" | "quality_tier" | "bandit"; + /** Classifier Calibrated Capable P Solve */ + classifier_calibrated_capable_p_solve?: number; + /** Classifier Calibrated Efficient P Solve */ + classifier_calibrated_efficient_p_solve?: number; /** Classifier Calibrated P Solve */ classifier_calibrated_p_solve?: number; /** Classifier Calibration Version */ classifier_calibration_version?: string; /** Classifier Capability Boundary */ classifier_capability_boundary?: string; + /** Classifier Capable P Solve */ + classifier_capable_p_solve?: number; /** Classifier Cost */ classifier_cost?: number; /** Classifier Crux */ classifier_crux?: string; + /** Classifier Efficient P Solve */ + classifier_efficient_p_solve?: number; + /** Classifier Max Quality Gap */ + classifier_max_quality_gap?: number; /** Classifier Model */ classifier_model?: string; /** Classifier P Solve */ classifier_p_solve?: number; /** Classifier Primary Rule */ classifier_primary_rule?: string; + /** Classifier Prompt Version */ + classifier_prompt_version?: string; /** Classifier Threshold */ classifier_threshold?: number; /** Context Escalated */ @@ -37545,6 +37607,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 */ @@ -37685,6 +37749,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 */ @@ -38791,6 +38857,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 */ @@ -39055,6 +39123,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; }; @@ -39254,6 +39324,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; }; @@ -39764,6 +39836,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 */ @@ -39932,10 +40006,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 */