diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 1c14e7d751..435b6eea45 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -17,7 +17,7 @@ Quick summary: - async_log_success_event() fires on GET /v1/batches/{id} (batch completion) """ -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple, Union from fastapi import HTTPException from pydantic import BaseModel @@ -25,12 +25,13 @@ from pydantic import BaseModel import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import ( + _extract_file_access_credentials, _get_batch_job_input_file_usage, _get_file_content_as_dictionary, _get_models_from_batch_input_file_content, ) from litellm.integrations.custom_logger import CustomLogger -from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy._types import SpecialModelNames, UserAPIKeyAuth if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -97,6 +98,276 @@ class _PROXY_BatchRateLimiter(CustomLogger): """ self.internal_usage_cache = internal_usage_cache self.parallel_request_limiter = parallel_request_limiter + self._warned_unsupported_model_skip = False + + def _get_file_bound_batch_model(self, data: Dict) -> Optional[str]: + """Resolve the model bound to the batch input file ID. + + ``create_batch`` routes a file-bound id (model-embedded ``file-...`` or + unified managed file) on that bound model and ignores the top-level + ``model``, so this is the authoritative routing model whenever the file + binds one. The provider is then read from that deployment's trusted + credentials for the provider-level skip decision. + """ + input_file_id = data.get("input_file_id") + if not isinstance(input_file_id, str) or not input_file_id: + return None + + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + decode_model_from_file_id, + get_models_from_unified_file_id, + ) + + model_from_file_id = decode_model_from_file_id(input_file_id) + if model_from_file_id: + return model_from_file_id + + unified_file_id = _is_base64_encoded_unified_file_id(input_file_id) + if unified_file_id: + target_model_names = get_models_from_unified_file_id(unified_file_id) + if target_model_names: + return target_model_names[0] + + return None + + def _get_batch_routing_model(self, data: Dict) -> Optional[str]: + """Resolve the deployment/model used for this batch from request data. + + Mirrors ``create_batch`` routing precedence: a model bound to the input + file id wins over the top-level ``model``, because the batch endpoint + ignores the top-level model for file-bound ids. Resolving the provider + skip from the top-level model first would let a caller point ``model`` + at a skip-listed provider while the file routes a rate-limited one. + """ + file_bound_model = self._get_file_bound_batch_model(data) + if file_bound_model: + return file_bound_model + + model = data.get("model") + if isinstance(model, str) and model: + return model + + return None + + def _resolve_batch_provider(self, batch_model: Optional[str]) -> Optional[str]: + """Resolve the provider from the deployment that serves ``batch_model``. + + The provider is read from trusted router credentials rather than the + user-supplied ``custom_llm_provider`` request field, so a caller cannot + spoof a skip-listed provider to bypass batch rate limiting. + """ + if not batch_model: + return None + + from litellm.proxy.openai_files_endpoints.common_utils import ( + get_credentials_for_model, + ) + from litellm.proxy.proxy_server import llm_router + + if llm_router is None: + return None + + try: + credentials = get_credentials_for_model( + llm_router=llm_router, + model_id=batch_model, + operation_context="batch input file read (rate limiting)", + ) + except HTTPException: + return None + + provider = credentials.get("custom_llm_provider") + return provider if isinstance(provider, str) and provider else None + + def _create_batch_rate_limit_descriptors( + self, + user_api_key_dict: UserAPIKeyAuth, + data: Dict, + ) -> List["RateLimitDescriptor"]: + return 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, + ) + + def _should_skip_batch_input_file_processing( + self, + data: Dict, + user_api_key_dict: UserAPIKeyAuth, + ) -> Tuple[bool, Optional[List["RateLimitDescriptor"]]]: + """ + Skip downloading batch input files when the operator disabled batch + input-file rate limiting, when the batch runs entirely on a skip-listed + provider, or when there is nothing to enforce (no applicable rate + limits). + + A skip is only honored for keys with unrestricted model access. When + the key has a model allowlist, the JSONL must still be downloaded so + ``_enforce_batch_file_model_access`` can validate every ``body.model`` + entry, otherwise a restricted key could smuggle unauthorized models + into the file via an admin-configured skip. + + The skip is never keyed on a specific model name. The models a batch + actually runs are its JSONL ``body.model`` entries, and any model + identifier the caller can influence (the top-level ``model`` or the + unsigned model embedded in a ``file-...`` id) can be pointed at a + skip-listed deployment while the file routes a different, rate-limited + model. The provider skip is safe because the provider is read from the + routing deployment's trusted credentials and the batch is constrained + to run on that provider. + + Returns ``(should_skip, descriptors)`` where ``descriptors`` is the + rate-limit descriptor list computed for the no-limits check, so the + caller can reuse it for counter enforcement without recomputing. + """ + from litellm.proxy.proxy_server import general_settings + + self._warn_if_unsupported_model_skip_configured(general_settings) + + if self._key_requires_batch_model_access_check(user_api_key_dict): + return False, None + + if general_settings.get("disable_batch_input_file_rate_limiting") is True: + return True, None + + skip_providers = ( + general_settings.get("skip_batch_input_file_rate_limiting_for_providers") + or [] + ) + if skip_providers: + batch_provider = self._resolve_batch_provider( + self._get_batch_routing_model(data) + ) + if batch_provider and batch_provider in skip_providers: + verbose_proxy_logger.debug( + f"Skipping batch input file processing for provider={batch_provider}" + ) + return True, None + + descriptors = self._create_batch_rate_limit_descriptors( + user_api_key_dict=user_api_key_dict, + data=data, + ) + if not self._has_applicable_batch_rate_limits(descriptors): + verbose_proxy_logger.debug( + "Skipping batch input file processing: no rate limits configured" + ) + return True, None + + return False, descriptors + + def _warn_if_unsupported_model_skip_configured( + self, general_settings: Dict + ) -> None: + """Warn once that ``skip_batch_input_file_rate_limiting_for_models`` is a no-op. + + A per-model skip is intentionally not honored because the model a batch + runs on is caller-influenced and can be pointed at a skip-listed + deployment while the JSONL routes a different, rate-limited model. + """ + if self._warned_unsupported_model_skip: + return + if general_settings.get("skip_batch_input_file_rate_limiting_for_models"): + self._warned_unsupported_model_skip = True + verbose_proxy_logger.warning( + "general_settings.skip_batch_input_file_rate_limiting_for_models is not " + "supported and has no effect. Use " + "skip_batch_input_file_rate_limiting_for_providers or " + "disable_batch_input_file_rate_limiting instead." + ) + + @staticmethod + def _key_requires_batch_model_access_check( + user_api_key_dict: UserAPIKeyAuth, + ) -> bool: + """True when the key may only call a subset of models (JSONL must be checked).""" + models = user_api_key_dict.models or [] + if "*" in models: + return False + if SpecialModelNames.all_proxy_models.value in models: + return False + if user_api_key_dict.access_group_ids: + return True + if not models: + return False + return True + + @staticmethod + def _has_applicable_batch_rate_limits( + descriptors: List["RateLimitDescriptor"], + ) -> bool: + for descriptor in descriptors: + rate_limit = descriptor.get("rate_limit") or {} + if ( + rate_limit.get("requests_per_unit") is not None + or rate_limit.get("tokens_per_unit") is not None + or rate_limit.get("max_parallel_requests") is not None + ): + return True + return False + + def _resolve_batch_input_file_fetch_params( + self, + file_id: str, + custom_llm_provider: str, + data: Dict, + ) -> Tuple[str, Dict[str, Any]]: + """ + Map proxy-facing file IDs to provider file IDs and credentials. + + Model-embedded IDs (``file-``) are not unified managed-file IDs; + without decoding them, ``afile_content`` is called with the encoded ID + and the upstream provider returns 404. + """ + from litellm.proxy.openai_files_endpoints.common_utils import ( + decode_model_from_file_id, + get_credentials_for_model, + get_original_file_id, + ) + from litellm.proxy.proxy_server import llm_router + + fetch_kwargs: Dict[str, Any] = { + "custom_llm_provider": custom_llm_provider, + } + + model_from_file_id = decode_model_from_file_id(file_id) + if model_from_file_id: + if llm_router is not None: + try: + credentials = get_credentials_for_model( + llm_router=llm_router, + model_id=model_from_file_id, + operation_context="batch input file read (rate limiting)", + ) + fetch_kwargs.update(_extract_file_access_credentials(credentials)) + fetch_kwargs["model"] = model_from_file_id + provider = credentials.get("custom_llm_provider") + if provider: + fetch_kwargs["custom_llm_provider"] = provider + except HTTPException: + pass + return get_original_file_id(file_id), fetch_kwargs + + request_model = data.get("model") + if isinstance(request_model, str) and request_model and llm_router is not None: + try: + credentials = get_credentials_for_model( + llm_router=llm_router, + model_id=request_model, + operation_context="batch input file read (rate limiting)", + ) + fetch_kwargs.update(_extract_file_access_credentials(credentials)) + fetch_kwargs["model"] = request_model + provider = credentials.get("custom_llm_provider") + if provider: + fetch_kwargs["custom_llm_provider"] = provider + except HTTPException: + pass + + return file_id, fetch_kwargs def _raise_rate_limit_error( self, @@ -163,6 +434,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): user_api_key_dict: UserAPIKeyAuth, data: Dict, batch_usage: BatchFileUsage, + descriptors: Optional[List["RateLimitDescriptor"]] = None, ) -> None: """ Atomically check + increment rate-limit counters by the batch amounts. @@ -171,14 +443,15 @@ class _PROXY_BatchRateLimiter(CustomLogger): case no counter is modified. Backed by `atomic_check_and_increment_by_n` which uses a Redis Lua script when available (multi-process atomic) and falls back to a per-process asyncio.Lock + in-memory operation. + + ``descriptors`` may be passed in by the pre-call hook to reuse the list + already computed when deciding whether to skip file processing. """ - descriptors = 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, - ) + if descriptors is None: + descriptors = self._create_batch_rate_limit_descriptors( + user_api_key_dict=user_api_key_dict, + data=data, + ) increment: Dict[Literal["requests", "tokens"], int] = { "requests": batch_usage.request_count, @@ -211,6 +484,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): file_id: str, custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", user_api_key_dict: Optional[UserAPIKeyAuth] = None, + data: Optional[Dict] = None, ) -> BatchFileUsage: """ Count number of requests and tokens in a batch input file. @@ -238,14 +512,27 @@ class _PROXY_BatchRateLimiter(CustomLogger): user_api_key_dict=user_api_key_dict, ) else: + provider_file_id, fetch_kwargs = ( + self._resolve_batch_input_file_fetch_params( + file_id=file_id, + custom_llm_provider=custom_llm_provider, + data=data or {}, + ) + ) # For non-managed files, use the standard litellm.afile_content file_content = await litellm.afile_content( - file_id=file_id, - custom_llm_provider=custom_llm_provider, + file_id=provider_file_id, user_api_key_dict=user_api_key_dict, + **fetch_kwargs, ) - file_content_as_dict = _get_file_content_as_dictionary(file_content.content) + file_content_bytes = getattr(file_content, "content", None) + if not isinstance(file_content_bytes, bytes): + raise ValueError( + f"Expected bytes content from file retrieval for {file_id}, " + f"got {type(file_content_bytes)}" + ) + file_content_as_dict = _get_file_content_as_dictionary(file_content_bytes) # Validate every model named in the batch JSONL against the # caller's per-key model allowlist. Without this, a caller @@ -441,6 +728,14 @@ class _PROXY_BatchRateLimiter(CustomLogger): ) return data + should_skip, batch_rate_limit_descriptors = ( + self._should_skip_batch_input_file_processing( + data=data, user_api_key_dict=user_api_key_dict + ) + ) + if should_skip: + return data + # Get custom_llm_provider for token counting custom_llm_provider = data.get("custom_llm_provider", "openai") @@ -452,6 +747,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): file_id=input_file_id, custom_llm_provider=custom_llm_provider, user_api_key_dict=user_api_key_dict, + data=data, ) verbose_proxy_logger.debug( @@ -469,6 +765,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): user_api_key_dict=user_api_key_dict, data=data, batch_usage=batch_usage, + descriptors=batch_rate_limit_descriptors, ) verbose_proxy_logger.debug( diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py index f047d62547..af5a5cde8b 100644 --- a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -259,6 +259,188 @@ async def test_pre_call_allows_authorized_model_in_batch_file(): ) +@pytest.mark.asyncio +async def test_pre_call_skips_file_fetch_when_disabled_in_general_settings(): + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + user = UserAPIKeyAuth(api_key="sk-ok", user_id="alice", models=["*"]) + + with patch( + "litellm.proxy.proxy_server.general_settings", + {"disable_batch_input_file_rate_limiting": True}, + ): + result = await rate_limiter.async_pre_call_hook( + user_api_key_dict=user, + cache=MagicMock(), + data={"input_file_id": "file-abc123"}, + call_type="acreate_batch", + ) + + assert result == {"input_file_id": "file-abc123"} + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.assert_not_called() + + +@pytest.mark.asyncio +async def test_pre_call_skips_file_fetch_for_configured_provider(): + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + user = UserAPIKeyAuth(api_key="sk-ok", user_id="alice", models=["*"]) + data = {"input_file_id": "file-abc123", "model": "my-vllm-model"} + + with ( + patch( + "litellm.proxy.proxy_server.general_settings", + {"skip_batch_input_file_rate_limiting_for_providers": ["hosted_vllm"]}, + ), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model", + return_value={"custom_llm_provider": "hosted_vllm"}, + ), + patch("litellm.afile_content", new=AsyncMock()) as mock_afile_content, + ): + result = await rate_limiter.async_pre_call_hook( + user_api_key_dict=user, + cache=MagicMock(), + data=data, + call_type="acreate_batch", + ) + + assert result == data + # A real skip must short-circuit before any file download or rate-limit + # work — assert the skip happened rather than the hook's error-recovery + # path (which also returns data unchanged). + mock_afile_content.assert_not_awaited() + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.assert_not_called() + + +@pytest.mark.asyncio +async def test_pre_call_does_not_skip_for_spoofed_provider(): + """The provider skip is resolved from trusted deployment credentials, so a + user-supplied ``custom_llm_provider`` that is not backed by the routing + deployment must not trigger a skip: the input file must still be fetched + and the rate-limit counters incremented.""" + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + # An applicable rate limit keeps the no-limits shortcut from firing, so the + # only thing that could prevent the fetch below is the provider skip. If the + # spoofed ``custom_llm_provider`` were honored, afile_content would never be + # awaited. + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [ + {"rate_limit": {"requests_per_unit": 100}} + ] + rate_limiter.parallel_request_limiter.atomic_check_and_increment_by_n = AsyncMock( + return_value={"overall_code": "OK", "statuses": []} + ) + user = UserAPIKeyAuth(api_key="sk-ok", user_id="alice", models=["*"]) + + mock_router = MagicMock() + mock_router.model_list = [] + mock_router.resolve_model_name_from_model_id.return_value = "my-openai-model" + + mock_content = MagicMock() + mock_content.content = ( + b'{"body": {"model": "my-openai-model", ' + b'"messages": [{"role": "user", "content": "hi"}]}}\n' + ) + + with ( + patch( + "litellm.proxy.proxy_server.general_settings", + {"skip_batch_input_file_rate_limiting_for_providers": ["hosted_vllm"]}, + ), + patch("litellm.proxy.proxy_server.llm_router", mock_router), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model", + return_value={"custom_llm_provider": "openai"}, + ), + patch( + "litellm.afile_content", new=AsyncMock(return_value=mock_content) + ) as mock_afile_content, + ): + await rate_limiter.async_pre_call_hook( + user_api_key_dict=user, + cache=MagicMock(), + data={ + "input_file_id": "file-abc123", + "model": "my-openai-model", + "custom_llm_provider": "hosted_vllm", + }, + call_type="acreate_batch", + ) + + # The spoofed provider did not short-circuit the skip decision: the file was + # fetched and the counters were incremented. + mock_afile_content.assert_awaited_once() + rate_limiter.parallel_request_limiter.atomic_check_and_increment_by_n.assert_awaited_once() + + +@pytest.mark.asyncio +async def test_count_input_file_usage_decodes_model_embedded_file_id(): + import base64 + + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + original_file_id = "file-provider-xyz" + encoded_payload = ( + base64.urlsafe_b64encode( + f"litellm:{original_file_id};model,my-vllm-batch".encode() + ) + .decode() + .rstrip("=") + ) + encoded_file_id = f"file-{encoded_payload}" + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + + mock_content = MagicMock() + mock_content.content = b'{"custom_id": "1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "my-vllm-batch", "messages": [{"role": "user", "content": "hi"}]}}\n' + + with ( + patch( + "litellm.afile_content", + new=AsyncMock(return_value=mock_content), + ) as mock_afile_content, + patch( + "litellm.proxy.proxy_server.llm_router", + MagicMock(), + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model", + return_value={ + "api_key": "test-key", + "api_base": "http://vllm:8000/v1", + "custom_llm_provider": "hosted_vllm", + }, + ), + ): + await rate_limiter.count_input_file_usage( + file_id=encoded_file_id, + custom_llm_provider="openai", + user_api_key_dict=UserAPIKeyAuth(api_key="sk-ok", user_id="alice"), + data={}, + ) + + mock_afile_content.assert_awaited_once() + assert mock_afile_content.await_args.kwargs["file_id"] == original_file_id + assert mock_afile_content.await_args.kwargs["custom_llm_provider"] == "hosted_vllm" + + @pytest.mark.asyncio async def test_pre_call_allows_stripped_provider_model_when_key_has_proxy_alias(): """After replace_model_in_jsonl, body.model is the provider id (e.g. gpt-5.5). @@ -323,3 +505,524 @@ async def test_pre_call_skips_check_when_no_models_present(): user_api_key_dict=user, file_content_as_dict=[{"body": {}}], ) + + +# --------------------------------------------------------------------------- +# Skip-path helpers +# --------------------------------------------------------------------------- + + +def _make_rate_limiter(): + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + return _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + + +def test_get_batch_routing_model_uses_request_model_for_plain_file(): + rate_limiter = _make_rate_limiter() + assert ( + rate_limiter._get_batch_routing_model({"model": "gpt-4o-mini"}) == "gpt-4o-mini" + ) + + +def test_get_batch_routing_model_prefers_file_bound_over_request_model(): + """``create_batch`` routes a model-embedded file id on its bound model and + ignores the top-level ``model``. The skip decision must use the same + precedence, otherwise a caller could point ``model`` at a skip-listed + provider while the file routes a rate-limited one.""" + import base64 + + rate_limiter = _make_rate_limiter() + encoded = ( + base64.urlsafe_b64encode(b"litellm:file-xyz;model,vllm-batch") + .decode() + .rstrip("=") + ) + assert ( + rate_limiter._get_batch_routing_model( + {"input_file_id": f"file-{encoded}", "model": "gpt-4o-mini"} + ) + == "vllm-batch" + ) + + +def test_get_batch_routing_model_returns_none_without_model_or_file(): + rate_limiter = _make_rate_limiter() + assert rate_limiter._get_batch_routing_model({}) is None + assert rate_limiter._get_batch_routing_model({"input_file_id": ""}) is None + + +def test_get_batch_routing_model_decodes_model_embedded_file_id(): + import base64 + + rate_limiter = _make_rate_limiter() + encoded = ( + base64.urlsafe_b64encode(b"litellm:file-xyz;model,vllm-batch") + .decode() + .rstrip("=") + ) + assert ( + rate_limiter._get_batch_routing_model({"input_file_id": f"file-{encoded}"}) + == "vllm-batch" + ) + + +def test_get_batch_routing_model_uses_unified_file_id_target(): + rate_limiter = _make_rate_limiter() + with ( + patch( + "litellm.proxy.openai_files_endpoints.common_utils.decode_model_from_file_id", + return_value=None, + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils._is_base64_encoded_unified_file_id", + return_value="unified-id", + ), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_models_from_unified_file_id", + return_value=["model-a", "model-b"], + ), + ): + assert ( + rate_limiter._get_batch_routing_model({"input_file_id": "file-managed"}) + == "model-a" + ) + + +def test_key_requires_batch_model_access_check_branches(): + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + check = _PROXY_BatchRateLimiter._key_requires_batch_model_access_check + assert check(UserAPIKeyAuth(api_key="sk", models=["*"])) is False + assert check(UserAPIKeyAuth(api_key="sk", models=["all-proxy-models"])) is False + assert ( + check(UserAPIKeyAuth(api_key="sk", models=[], access_group_ids=["grp"])) is True + ) + assert check(UserAPIKeyAuth(api_key="sk", models=[])) is False + assert check(UserAPIKeyAuth(api_key="sk", models=["gpt-4o-mini"])) is True + # Wildcard / all-proxy-models grant access to every model, so + # can_key_call_model passes any model regardless of access groups (which + # only ever widen access). Such keys must not be forced to download and + # validate the JSONL even when access_group_ids are also present. + assert ( + check(UserAPIKeyAuth(api_key="sk", models=["*"], access_group_ids=["grp"])) + is False + ) + assert ( + check( + UserAPIKeyAuth( + api_key="sk", models=["all-proxy-models"], access_group_ids=["grp"] + ) + ) + is False + ) + # A concrete model allowlist is still a subset even with access groups. + assert ( + check( + UserAPIKeyAuth( + api_key="sk", models=["gpt-4o-mini"], access_group_ids=["grp"] + ) + ) + is True + ) + + +def test_has_applicable_batch_rate_limits(): + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + has_limits = _PROXY_BatchRateLimiter._has_applicable_batch_rate_limits + assert has_limits([{"rate_limit": {"tokens_per_unit": 100}}]) is True + assert has_limits([{"rate_limit": {"requests_per_unit": 5}}]) is True + assert has_limits([{"rate_limit": {"max_parallel_requests": 2}}]) is True + assert has_limits([{"rate_limit": {}}, {}]) is False + + +def test_should_skip_returns_false_when_key_needs_model_access_check(): + rate_limiter = _make_rate_limiter() + user = UserAPIKeyAuth(api_key="sk", models=["gpt-4o-mini"]) + should_skip, descriptors = rate_limiter._should_skip_batch_input_file_processing( + data={"input_file_id": "file-abc"}, user_api_key_dict=user + ) + assert should_skip is False + assert descriptors is None + + +def test_should_skip_ignores_client_supplied_metadata_flag(): + """A caller must not be able to bypass batch rate limits by setting + ``litellm_metadata.skip_batch_input_file_rate_limiting`` in the request + body. The skip decision is server-controlled only, so with applicable rate + limits the JSONL is still processed despite the client flag.""" + rate_limiter = _make_rate_limiter() + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [ + {"rate_limit": {"requests_per_unit": 5}} + ] + user = UserAPIKeyAuth(api_key="sk", models=["*"]) + with patch("litellm.proxy.proxy_server.general_settings", {}): + should_skip, descriptors = ( + rate_limiter._should_skip_batch_input_file_processing( + data={ + "input_file_id": "file-abc", + "litellm_metadata": {"skip_batch_input_file_rate_limiting": True}, + }, + user_api_key_dict=user, + ) + ) + assert should_skip is False + + +def test_should_not_skip_for_forged_model_embedded_file_id(): + """A ``file-`` id embeds an unsigned model name the caller fully + controls, so a caller can re-encode any accessible provider file id with a + skip-listed model while the JSONL still routes rate-limited ``body.model`` + entries. The per-model skip must therefore never fire: with applicable rate + limits, a forged skip-listed file-bound model still falls through to file + processing and counter enforcement.""" + import base64 + + rate_limiter = _make_rate_limiter() + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [ + {"rate_limit": {"requests_per_unit": 5}} + ] + user = UserAPIKeyAuth(api_key="sk", models=["*"]) + encoded = ( + base64.urlsafe_b64encode(b"litellm:file-xyz;model,gpt-4o-mini") + .decode() + .rstrip("=") + ) + with patch( + "litellm.proxy.proxy_server.general_settings", + {"skip_batch_input_file_rate_limiting_for_models": ["gpt-4o-mini"]}, + ): + should_skip, descriptors = ( + rate_limiter._should_skip_batch_input_file_processing( + data={"input_file_id": f"file-{encoded}"}, + user_api_key_dict=user, + ) + ) + assert should_skip is False + assert descriptors is not None + + +def test_should_not_skip_for_skip_listed_top_level_model(): + """A caller must not bypass batch rate limits by naming a skip-listed model + in the top-level ``model`` while routing a different model through the JSONL + ``body.model`` entries. No per-model skip exists, so a skip-listed model over + a plain file still gets processed.""" + rate_limiter = _make_rate_limiter() + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [ + {"rate_limit": {"requests_per_unit": 5}} + ] + user = UserAPIKeyAuth(api_key="sk", models=["*"]) + with patch( + "litellm.proxy.proxy_server.general_settings", + {"skip_batch_input_file_rate_limiting_for_models": ["gpt-4o-mini"]}, + ): + should_skip, descriptors = ( + rate_limiter._should_skip_batch_input_file_processing( + data={"model": "gpt-4o-mini", "input_file_id": "file-abc"}, + user_api_key_dict=user, + ) + ) + assert should_skip is False + + +def test_should_not_skip_when_file_bound_provider_is_rate_limited(): + """A caller must not bypass batch rate limits by pointing the top-level + ``model`` at a skip-listed provider while the model-embedded ``input_file_id`` + routes to a rate-limited provider. ``create_batch`` runs the batch on the + file-bound model, so the skip decision must resolve the provider from that + model and still process the file when its provider is not skip-listed.""" + import base64 + + rate_limiter = _make_rate_limiter() + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [ + {"rate_limit": {"requests_per_unit": 5}} + ] + user = UserAPIKeyAuth(api_key="sk", models=["*"]) + encoded = ( + base64.urlsafe_b64encode(b"litellm:file-orig;model,vllm-batch") + .decode() + .rstrip("=") + ) + + def _creds(model_id, **kwargs): + provider = "hosted_vllm" if model_id == "vllm-batch" else "openai" + return {"custom_llm_provider": provider} + + with ( + patch( + "litellm.proxy.proxy_server.general_settings", + {"skip_batch_input_file_rate_limiting_for_providers": ["openai"]}, + ), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model", + side_effect=_creds, + ), + ): + should_skip, descriptors = ( + rate_limiter._should_skip_batch_input_file_processing( + data={"input_file_id": f"file-{encoded}", "model": "gpt-skip"}, + user_api_key_dict=user, + ) + ) + assert should_skip is False + assert descriptors is not None + + +def test_should_skip_when_file_bound_provider_is_skip_listed(): + """The provider skip must still fire when the model the batch actually runs + on (the file-bound model) resolves to a skip-listed provider, even if the + top-level ``model`` resolves to a different, non-skipped provider.""" + import base64 + + rate_limiter = _make_rate_limiter() + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [ + {"rate_limit": {"requests_per_unit": 5}} + ] + user = UserAPIKeyAuth(api_key="sk", models=["*"]) + encoded = ( + base64.urlsafe_b64encode(b"litellm:file-orig;model,vllm-batch") + .decode() + .rstrip("=") + ) + + def _creds(model_id, **kwargs): + provider = "hosted_vllm" if model_id == "vllm-batch" else "openai" + return {"custom_llm_provider": provider} + + with ( + patch( + "litellm.proxy.proxy_server.general_settings", + {"skip_batch_input_file_rate_limiting_for_providers": ["hosted_vllm"]}, + ), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model", + side_effect=_creds, + ), + ): + should_skip, descriptors = ( + rate_limiter._should_skip_batch_input_file_processing( + data={"input_file_id": f"file-{encoded}", "model": "gpt-skip"}, + user_api_key_dict=user, + ) + ) + assert should_skip is True + + +def test_warns_once_for_unsupported_model_skip_setting(): + """Operators who set the no-op per-model skip key get a single warning so a + misconfigured deployment does not silently leave batch limits unenforced.""" + rate_limiter = _make_rate_limiter() + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [ + {"rate_limit": {"requests_per_unit": 5}} + ] + user = UserAPIKeyAuth(api_key="sk", models=["*"]) + with ( + patch( + "litellm.proxy.proxy_server.general_settings", + {"skip_batch_input_file_rate_limiting_for_models": ["gpt-4o-mini"]}, + ), + patch( + "litellm.proxy.hooks.batch_rate_limiter.verbose_proxy_logger" + ) as mock_logger, + ): + for _ in range(3): + rate_limiter._should_skip_batch_input_file_processing( + data={"model": "gpt-4o-mini", "input_file_id": "file-abc"}, + user_api_key_dict=user, + ) + assert mock_logger.warning.call_count == 1 + assert ( + "skip_batch_input_file_rate_limiting_for_models" + in mock_logger.warning.call_args[0][0] + ) + + +def test_no_warning_when_model_skip_setting_absent(): + rate_limiter = _make_rate_limiter() + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [ + {"rate_limit": {"requests_per_unit": 5}} + ] + user = UserAPIKeyAuth(api_key="sk", models=["*"]) + with ( + patch( + "litellm.proxy.proxy_server.general_settings", + {"skip_batch_input_file_rate_limiting_for_providers": ["openai"]}, + ), + patch( + "litellm.proxy.hooks.batch_rate_limiter.verbose_proxy_logger" + ) as mock_logger, + ): + rate_limiter._should_skip_batch_input_file_processing( + data={"model": "gpt-4o-mini", "input_file_id": "file-abc"}, + user_api_key_dict=user, + ) + mock_logger.warning.assert_not_called() + + +def test_should_skip_when_no_rate_limits_configured(): + rate_limiter = _make_rate_limiter() + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [ + {"rate_limit": {}} + ] + user = UserAPIKeyAuth(api_key="sk", models=["*"]) + with patch("litellm.proxy.proxy_server.general_settings", {}): + should_skip, descriptors = ( + rate_limiter._should_skip_batch_input_file_processing( + data={"model": "gpt-4o-mini", "input_file_id": "file-abc"}, + user_api_key_dict=user, + ) + ) + assert should_skip is True + assert descriptors is None + + +def test_should_not_skip_and_reuses_descriptors_when_limits_present(): + rate_limiter = _make_rate_limiter() + descriptors = [{"rate_limit": {"tokens_per_unit": 100}}] + rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = ( + descriptors + ) + user = UserAPIKeyAuth(api_key="sk", models=["*"]) + with patch("litellm.proxy.proxy_server.general_settings", {}): + should_skip, returned = rate_limiter._should_skip_batch_input_file_processing( + data={"model": "gpt-4o-mini", "input_file_id": "file-abc"}, + user_api_key_dict=user, + ) + assert should_skip is False + assert returned is descriptors + + +def test_resolve_fetch_params_uses_request_model_credentials(): + rate_limiter = _make_rate_limiter() + with ( + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model", + return_value={ + "api_key": "k", + "api_base": "http://vllm:8000/v1", + "custom_llm_provider": "hosted_vllm", + }, + ), + ): + provider_file_id, fetch_kwargs = ( + rate_limiter._resolve_batch_input_file_fetch_params( + file_id="file-plain-openai", + custom_llm_provider="openai", + data={"model": "my-vllm-batch"}, + ) + ) + assert provider_file_id == "file-plain-openai" + assert fetch_kwargs["model"] == "my-vllm-batch" + assert fetch_kwargs["custom_llm_provider"] == "hosted_vllm" + assert fetch_kwargs["api_base"] == "http://vllm:8000/v1" + + +def test_resolve_fetch_params_fails_open_on_credential_lookup_error(): + rate_limiter = _make_rate_limiter() + with ( + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model", + side_effect=HTTPException(status_code=404, detail="no creds"), + ), + ): + provider_file_id, fetch_kwargs = ( + rate_limiter._resolve_batch_input_file_fetch_params( + file_id="file-plain-openai", + custom_llm_provider="openai", + data={"model": "my-vllm-batch"}, + ) + ) + assert provider_file_id == "file-plain-openai" + assert fetch_kwargs == {"custom_llm_provider": "openai"} + + +def test_resolve_fetch_params_model_embedded_fails_open_on_credential_error(): + import base64 + + rate_limiter = _make_rate_limiter() + encoded = ( + base64.urlsafe_b64encode(b"litellm:file-orig;model,vllm-batch") + .decode() + .rstrip("=") + ) + encoded_file_id = f"file-{encoded}" + + get_credentials = MagicMock( + side_effect=HTTPException(status_code=404, detail="no creds") + ) + with ( + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch( + "litellm.proxy.openai_files_endpoints.common_utils.get_credentials_for_model", + get_credentials, + ), + ): + provider_file_id, fetch_kwargs = ( + rate_limiter._resolve_batch_input_file_fetch_params( + file_id=encoded_file_id, + custom_llm_provider="openai", + data={}, + ) + ) + get_credentials.assert_called_once() + assert provider_file_id == "file-orig" + assert fetch_kwargs == {"custom_llm_provider": "openai"} + + +@pytest.mark.asyncio +async def test_check_and_increment_computes_descriptors_when_not_passed(): + from litellm.proxy.hooks.batch_rate_limiter import ( + BatchFileUsage, + _PROXY_BatchRateLimiter, + ) + + parallel_request_limiter = MagicMock() + parallel_request_limiter._create_rate_limit_descriptors.return_value = [ + {"rate_limit": {"tokens_per_unit": 100}} + ] + parallel_request_limiter.atomic_check_and_increment_by_n = AsyncMock( + return_value={"overall_code": "OK", "statuses": []} + ) + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=parallel_request_limiter, + ) + + await rate_limiter._check_and_increment_batch_counters( + user_api_key_dict=UserAPIKeyAuth(api_key="sk", models=["*"]), + data={"model": "gpt-4o-mini"}, + batch_usage=BatchFileUsage(total_tokens=10, request_count=1), + descriptors=None, + ) + + parallel_request_limiter._create_rate_limit_descriptors.assert_called_once() + + +@pytest.mark.asyncio +async def test_count_input_file_usage_raises_on_non_bytes_content(): + from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter + + rate_limiter = _PROXY_BatchRateLimiter( + internal_usage_cache=MagicMock(), + parallel_request_limiter=MagicMock(), + ) + + bad_content = MagicMock() + bad_content.content = "not-bytes" + + with patch("litellm.afile_content", new=AsyncMock(return_value=bad_content)): + with pytest.raises(ValueError, match="Expected bytes content"): + await rate_limiter.count_input_file_usage( + file_id="file-plain", + custom_llm_provider="openai", + user_api_key_dict=UserAPIKeyAuth(api_key="sk", models=["*"]), + data={}, + )