scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-02 02:11:58 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 101
litellm/tests/unit/proxy/hooks/test_batch_file_validation.py
devin-ai-integration[bot] 39e31958f8
test(proxy): move auth, hooks, policy_engine and client tests into tests/unit/proxy (#43998) 
* test(proxy): move auth, hooks, policy_engine and client tests into tests/unit/proxy

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): stub HIBP through respx by disabling the aiohttp transport

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): share the httpx transport fixture across proxy unit tests

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): restore proxy globals without a missing-value sentinel

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): package moved dirs and stub the login breach check at the HTTP boundary

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* test(proxy): isolate the mcp server manager per test

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yuneng <yuneng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
2026-10-01 10:11:45 -07:00

2288 lines
83 KiB
Python
Raw Blame History

"""
VERIA-39 regression tests:
- The batch input-file token counter must measure embeddings (`input`)
and text-completion (`prompt`) payloads, not only chat (`messages`).
- The batch rate-limiter pre-call hook must reject batch files that name
models the caller is not authorized to use.
"""
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fastapi import HTTPException
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
def _models(file_content_as_dict):
"""Distinct body.model values, mirroring how the rate limiter collects the
models from a streamed batch file before the access check."""
return [
entry["body"]["model"]
for entry in file_content_as_dict
if (entry.get("body") or {}).get("model")
]
# ---------------------------------------------------------------------------
# Token counter — covers all three batch payload shapes
# ---------------------------------------------------------------------------
def test_token_counter_counts_chat_messages():
from litellm.batches.batch_utils import _count_entry_tokens
tokens = _count_entry_tokens(
{
"body": {
"model": "gpt-4o-mini",
"messages": [{"role": "user", "content": "hello"}],
}
}
)
assert tokens > 0
def test_token_counter_counts_text_completion_prompt():
"""Pre-fix this returned 0 tokens (the counter only inspected
`messages`), letting `prompt`-style batches slip past TPM limits."""
from litellm.batches.batch_utils import _count_entry_tokens
tokens = _count_entry_tokens(
{"body": {"model": "gpt-3.5-turbo-instruct", "prompt": "hello world"}}
)
assert tokens > 0
def test_token_counter_counts_embedding_input_string():
from litellm.batches.batch_utils import _count_entry_tokens
tokens = _count_entry_tokens(
{"body": {"model": "text-embedding-3-small", "input": "hello world"}}
)
assert tokens > 0
def test_token_counter_counts_embedding_input_list():
from litellm.batches.batch_utils import _count_entry_tokens
tokens = _count_entry_tokens(
{
"body": {
"model": "text-embedding-3-small",
"input": ["hello", "world"],
}
}
)
assert tokens > 0
def test_token_counter_counts_text_completion_prompt_list():
from litellm.batches.batch_utils import _count_entry_tokens
tokens = _count_entry_tokens(
{
"body": {
"model": "gpt-3.5-turbo-instruct",
"prompt": ["alpha", "beta"],
}
}
)
assert tokens > 0
def test_token_counter_counts_pre_tokenized_prompt_int_list():
"""OpenAI's text-completion API accepts a single pre-tokenized prompt as
a list of ints. Each int is one token; pre-fix this shape was silently
counted as zero, leaving a TPM bypass."""
from litellm.batches.batch_utils import _count_entry_tokens
tokens = _count_entry_tokens(
{
"body": {
"model": "gpt-3.5-turbo-instruct",
"prompt": [1, 2, 3, 4, 5],
}
}
)
assert tokens == 5
def test_token_counter_counts_pre_tokenized_prompt_list_of_int_lists():
"""Multiple pre-tokenized prompts (`list[list[int]]`) — the most
important bypass shape. A 1000-token batch must report 1000 tokens,
not zero."""
from litellm.batches.batch_utils import _count_entry_tokens
tokens = _count_entry_tokens(
{
"body": {
"model": "gpt-3.5-turbo-instruct",
"prompt": [[1] * 250, [2] * 250, [3] * 500],
}
}
)
assert tokens == 1000
def test_token_counter_counts_pre_tokenized_input_for_embeddings():
"""Same shape applies to embeddings (`input`)."""
from litellm.batches.batch_utils import _count_entry_tokens
tokens = _count_entry_tokens(
{
"body": {
"model": "text-embedding-3-small",
"input": [[1, 2, 3], [4, 5, 6]],
}
}
)
assert tokens == 6
# ---------------------------------------------------------------------------
# Pre-call hook model validation
# ---------------------------------------------------------------------------
@pytest.mark.asyncio
async def test_pre_call_rejects_unauthorized_model_in_batch_file():
"""Pre-fix the hook only validated the outer `model` parameter and
forwarded the file as-is. With this fix, a model named inside the
JSONL that the caller cannot use must trigger a 403."""
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
# Simulated decoded batch file: caller is restricted to gpt-3.5
# but the JSONL points at gpt-4o.
file_dict = [
{"body": {"model": "gpt-4o", "messages": [{"role": "user", "content": "x"}]}}
]
user = UserAPIKeyAuth(
api_key="sk-restricted",
user_id="alice",
models=["gpt-3.5-turbo"],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
# `can_key_call_model` raises a ProxyException for non-allowed models.
async def _raise_unauthorized(**kwargs):
raise Exception(
f"Key not allowed to access model. This key only has access to models={kwargs['valid_token'].models}"
)
with (
patch(
"litellm.proxy.auth.auth_checks.can_key_call_model",
new=AsyncMock(side_effect=_raise_unauthorized),
),
patch("litellm.proxy.proxy_server.llm_router", None),
):
with pytest.raises(HTTPException) as exc:
await rate_limiter._enforce_batch_file_model_access(
user_api_key_dict=user,
models=_models(file_dict),
)
assert exc.value.status_code == 403
assert "gpt-4o" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_pre_call_allows_all_team_models_key_when_model_in_team_allowlist():
"""Keys with ``all-team-models`` must inherit the team allowlist when
validating models embedded in batch JSONL."""
from litellm.proxy._types import SpecialModelNames
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
proxy_alias = "openai/openai/gpt-5.5-batch"
file_dict = [
{
"body": {
"model": proxy_alias,
"messages": [{"role": "user", "content": "x"}],
}
}
]
user = UserAPIKeyAuth(
api_key="sk-team",
user_id="alice",
team_id="team-123",
models=[SpecialModelNames.all_team_models.value],
team_models=[proxy_alias],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
with patch("litellm.proxy.proxy_server.llm_router", None):
await rate_limiter._enforce_batch_file_model_access(
user_api_key_dict=user,
models=_models(file_dict),
)
@pytest.mark.asyncio
async def test_pre_call_uses_current_team_allowlist_for_all_team_models_key():
from litellm.proxy._types import LiteLLM_TeamTable, SpecialModelNames
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
stale_model = "stale-model"
current_model = "current-model"
file_dict = [
{
"body": {
"model": stale_model,
"messages": [{"role": "user", "content": "x"}],
}
}
]
user = UserAPIKeyAuth(
api_key="sk-team",
user_id="alice",
team_id="team-123",
models=[SpecialModelNames.all_team_models.value],
team_models=[stale_model],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
team_object = LiteLLM_TeamTable(
team_id="team-123",
models=[current_model],
)
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.llm_router", None),
patch(
"litellm.proxy.auth.auth_checks.get_team_object",
new=AsyncMock(return_value=team_object),
) as mock_get_team_object,
pytest.raises(HTTPException) as exc_info,
):
await rate_limiter._enforce_batch_file_model_access(
user_api_key_dict=user,
models=_models(file_dict),
)
assert exc_info.value.status_code == 403
mock_get_team_object.assert_awaited_once()
@pytest.mark.asyncio
async def test_pre_call_allows_all_team_models_key_via_current_team_object():
"""Happy path for the team_object branch: with a DB client present, an
``all-team-models`` key whose batch model is on the *current* team
allowlist must be authorized through the freshly-fetched team object,
not the cached-``team_models`` fallback."""
from litellm.proxy._types import LiteLLM_TeamTable, SpecialModelNames
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
current_model = "current-model"
file_dict = [
{
"body": {
"model": current_model,
"messages": [{"role": "user", "content": "x"}],
}
}
]
user = UserAPIKeyAuth(
api_key="sk-team",
user_id="alice",
team_id="team-123",
models=[SpecialModelNames.all_team_models.value],
team_models=["stale-model"],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
team_object = LiteLLM_TeamTable(
team_id="team-123",
models=[current_model],
)
can_key_call_model = AsyncMock(return_value=True)
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.llm_router", None),
patch(
"litellm.proxy.auth.auth_checks.get_team_object",
new=AsyncMock(return_value=team_object),
) as mock_get_team_object,
patch(
"litellm.proxy.auth.auth_checks.get_team_membership",
new=AsyncMock(return_value=None),
),
patch(
"litellm.proxy.auth.auth_checks.can_key_call_model",
new=can_key_call_model,
),
):
await rate_limiter._enforce_batch_file_model_access(
user_api_key_dict=user,
models=_models(file_dict),
)
mock_get_team_object.assert_awaited_once()
can_key_call_model.assert_not_awaited()
@pytest.mark.asyncio
async def test_pre_call_denies_all_team_models_key_via_member_scope():
"""The team_object branch must also apply the per-member model scope: a
model on the team allowlist but outside the member's ``allowed_models``
must be rejected with a 403."""
from litellm.proxy._types import (
LiteLLM_BudgetTable,
LiteLLM_TeamMembership,
LiteLLM_TeamTable,
SpecialModelNames,
)
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
team_model = "team-model"
file_dict = [
{
"body": {
"model": team_model,
"messages": [{"role": "user", "content": "x"}],
}
}
]
user = UserAPIKeyAuth(
api_key="sk-team",
user_id="alice",
team_id="team-123",
models=[SpecialModelNames.all_team_models.value],
team_models=[team_model],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
team_object = LiteLLM_TeamTable(team_id="team-123", models=[team_model])
membership = LiteLLM_TeamMembership(
user_id="alice",
team_id="team-123",
litellm_budget_table=LiteLLM_BudgetTable(allowed_models=["other-model"]),
)
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.llm_router", None),
patch(
"litellm.proxy.auth.auth_checks.get_team_object",
new=AsyncMock(return_value=team_object),
),
patch(
"litellm.proxy.auth.auth_checks.get_team_membership",
new=AsyncMock(return_value=membership),
),
pytest.raises(HTTPException) as exc_info,
):
await rate_limiter._enforce_batch_file_model_access(
user_api_key_dict=user,
models=_models(file_dict),
)
assert exc_info.value.status_code == 403
assert team_model in str(exc_info.value.detail)
@pytest.mark.parametrize(
("team_fetch_error", "expected_status"),
[
(HTTPException(status_code=404, detail="team not found"), 404),
(Exception("team fetch failed"), 403),
],
)
@pytest.mark.asyncio
async def test_pre_call_fails_closed_when_current_team_fetch_fails_for_all_team_models_key(
team_fetch_error, expected_status
):
from litellm.proxy._types import SpecialModelNames
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
stale_model = "stale-model"
file_dict = [
{
"body": {
"model": stale_model,
"messages": [{"role": "user", "content": "x"}],
}
}
]
user = UserAPIKeyAuth(
api_key="sk-team",
user_id="alice",
team_id="team-123",
models=[SpecialModelNames.all_team_models.value],
team_models=[stale_model],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
with (
patch("litellm.proxy.proxy_server.prisma_client", MagicMock()),
patch("litellm.proxy.proxy_server.llm_router", None),
patch(
"litellm.proxy.auth.auth_checks.get_team_object",
new=AsyncMock(side_effect=team_fetch_error),
) as mock_get_team_object,
patch(
"litellm.proxy.auth.auth_checks.can_key_call_model",
new=AsyncMock(return_value=True),
) as mock_can_key_call_model,
pytest.raises(HTTPException) as exc_info,
):
await rate_limiter._enforce_batch_file_model_access(
user_api_key_dict=user,
models=_models(file_dict),
)
assert exc_info.value.status_code == expected_status
mock_get_team_object.assert_awaited_once()
mock_can_key_call_model.assert_not_awaited()
@pytest.mark.asyncio
async def test_pre_call_allows_teamless_all_team_models_key():
"""A teamless key with all-team-models must be allowed to submit batch jobs
for any model (same as leaving models empty = unrestricted). Fails if
someone re-introduces a teamless denial in _resolve_key_models_for_auth_check
or adds a team_id guard that blocks the batch path."""
from litellm.proxy._types import SpecialModelNames
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
file_dict = [
{
"body": {
"model": "gpt-4o",
"messages": [{"role": "user", "content": "x"}],
}
}
]
user = UserAPIKeyAuth(
api_key="sk-orphan",
user_id="alice",
models=[SpecialModelNames.all_team_models.value],
team_models=[],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
with patch("litellm.proxy.proxy_server.llm_router", None):
await rate_limiter._enforce_batch_file_model_access(
user_api_key_dict=user,
models=_models(file_dict),
)
@pytest.mark.asyncio
async def test_pre_call_allows_authorized_model_in_batch_file():
"""If every model in the JSONL is on the caller's allowlist, the hook
must not raise."""
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
file_dict = [
{
"body": {
"model": "gpt-3.5-turbo",
"messages": [{"role": "user", "content": "x"}],
}
}
]
user = UserAPIKeyAuth(
api_key="sk-ok",
user_id="alice",
models=["gpt-3.5-turbo"],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
with (
patch(
"litellm.proxy.auth.auth_checks.can_key_call_model",
new=AsyncMock(return_value=True),
),
patch("litellm.proxy.proxy_server.llm_router", None),
):
# Should not raise
await rate_limiter._enforce_batch_file_model_access(
user_api_key_dict=user,
models=_models(file_dict),
)
@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).
Auth must check target_model_names from the unified file id, not reverse-map
the stripped id."""
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
proxy_alias = "openai/openai/gpt-5.5-batch"
file_dict = [
{"body": {"model": "gpt-5.5", "messages": [{"role": "user", "content": "x"}]}}
]
user = UserAPIKeyAuth(
api_key="sk-ok",
user_id="alice",
models=[proxy_alias],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
mock_router = MagicMock()
mock_router.model_list = []
can_key_call_model = AsyncMock(return_value=True)
with (
patch(
"litellm.proxy.auth.auth_checks.can_key_call_model",
new=can_key_call_model,
),
patch("litellm.proxy.proxy_server.llm_router", mock_router),
):
await rate_limiter._enforce_batch_file_model_access(
user_api_key_dict=user,
models=_models(file_dict),
target_model_names=[proxy_alias],
)
can_key_call_model.assert_awaited_once()
assert can_key_call_model.await_args.kwargs["model"] == proxy_alias
mock_router.resolve_model_name_from_model_id.assert_not_called()
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model_list_order",
[
[
"openai/openai/gpt-5.5",
"openai/openai/gpt-5.5-batch",
"us/azure/openai/gpt-5.5",
],
[
"us/azure/openai/gpt-5.5",
"openai/openai/gpt-5.5",
"openai/openai/gpt-5.5-batch",
],
[
"openai/openai/gpt-5.5-batch",
"us/azure/openai/gpt-5.5",
"openai/openai/gpt-5.5",
],
],
)
async def test_pre_call_uses_target_model_names_not_stripped_reverse_lookup(
model_list_order,
):
"""LIT-3593: three deployments strip to gpt-5.5; auth must use the upload
target alias from target_model_names, not first-match reverse lookup."""
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
batch_alias = "openai/openai/gpt-5.5-batch"
deployment_templates = {
"openai/openai/gpt-5.5": {
"model_name": "openai/openai/gpt-5.5",
"litellm_params": {"model": "openai/gpt-5.5"},
"model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"},
},
"openai/openai/gpt-5.5-batch": {
"model_name": "openai/openai/gpt-5.5-batch",
"litellm_params": {"model": "openai/gpt-5.5"},
"model_info": {"id": "openai/openai/gpt-5.5-batch", "mode": "batch"},
},
"us/azure/openai/gpt-5.5": {
"model_name": "us/azure/openai/gpt-5.5",
"litellm_params": {"model": "azure/gpt-5.5"},
"model_info": {"id": "openai/openai/gpt-5.5", "mode": "chat"},
},
}
mock_router = MagicMock()
mock_router.model_list = [deployment_templates[name] for name in model_list_order]
def _resolve(model_id):
for deployment in mock_router.model_list:
actual_model = deployment.get("litellm_params", {}).get("model")
if actual_model == model_id or (
actual_model and actual_model.endswith(f"/{model_id}")
):
return deployment.get("model_name")
return None
mock_router.resolve_model_name_from_model_id.side_effect = _resolve
file_dict = [
{"body": {"model": "gpt-5.5", "messages": [{"role": "user", "content": "x"}]}}
]
user = UserAPIKeyAuth(
api_key="sk-ok",
user_id="alice",
models=[batch_alias],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
can_key_call_model = AsyncMock(return_value=True)
with (
patch(
"litellm.proxy.auth.auth_checks.can_key_call_model",
new=can_key_call_model,
),
patch("litellm.proxy.proxy_server.llm_router", mock_router),
):
await rate_limiter._enforce_batch_file_model_access(
user_api_key_dict=user,
models=_models(file_dict),
target_model_names=[batch_alias],
)
can_key_call_model.assert_awaited_once()
assert can_key_call_model.await_args.kwargs["model"] == batch_alias
mock_router.resolve_model_name_from_model_id.assert_not_called()
@pytest.mark.asyncio
async def test_pre_call_skips_check_when_no_models_present():
"""Files without any `body.model` (corrupt or empty) must not 500;
the rate limiter logs a warning elsewhere and proceeds."""
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")
# Should not raise even though `can_key_call_model` is the default
# (would fail). The early-return on empty models keeps the call out
# entirely.
await rate_limiter._enforce_batch_file_model_access(
user_api_key_dict=user,
models=_models([]),
)
await rate_limiter._enforce_batch_file_model_access(
user_api_key_dict=user,
models=_models([{"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-<base64>`` 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_pre_call_enforces_project_otpm_limit_for_batch():
"""VERIA regression: ``_create_batch_rate_limit_descriptors`` only asked
for the generic key/user/team/model descriptors, so a project caller
could submit a batch that consumed none of its configured project OTPM
quota. The project OTPM descriptor must now be present and charged with
the batch's estimated *output* tokens, not its input tokens."""
from litellm import DualCache
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.utils import InternalUsageCache
local_cache = DualCache()
parallel_request_limiter = _PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=InternalUsageCache(local_cache)
)
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=InternalUsageCache(local_cache),
parallel_request_limiter=parallel_request_limiter,
)
user = UserAPIKeyAuth(
api_key="sk-project-batch-otpm",
models=["*"],
project_id="proj-mantle-batch",
project_metadata={"model_otpm_limit": {"gpt-4o-mini": 50}},
)
# Two rows each declaring max_tokens=40: 80 output tokens total, over the
# configured 50-token project OTPM limit but negligible input tokens.
mock_content = MagicMock()
mock_content.content = (
b'{"body": {"model": "gpt-4o-mini", "max_tokens": 40, '
b'"messages": [{"role": "user", "content": "hi"}]}}\n'
b'{"body": {"model": "gpt-4o-mini", "max_tokens": 40, '
b'"messages": [{"role": "user", "content": "hi"}]}}\n'
)
with (
patch("litellm.proxy.proxy_server.general_settings", {}),
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": "openai"},
),
patch("litellm.afile_content", new=AsyncMock(return_value=mock_content)),
):
with pytest.raises(HTTPException) as exc:
await rate_limiter.async_pre_call_hook(
user_api_key_dict=user,
cache=local_cache,
data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"},
call_type="acreate_batch",
)
assert exc.value.status_code == 429
assert "model_per_project_otpm" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_pre_call_enforces_project_itpm_limit_for_batch():
"""Companion to the OTPM regression above: a project's ITPM quota must
also apply to batch submissions."""
from litellm import DualCache
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.utils import InternalUsageCache
local_cache = DualCache()
parallel_request_limiter = _PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=InternalUsageCache(local_cache)
)
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=InternalUsageCache(local_cache),
parallel_request_limiter=parallel_request_limiter,
)
user = UserAPIKeyAuth(
api_key="sk-project-batch-itpm",
models=["*"],
project_id="proj-mantle-batch",
project_metadata={"model_itpm_limit": {"gpt-4o-mini": 1}},
)
mock_content = MagicMock()
mock_content.content = (
b'{"body": {"model": "gpt-4o-mini", "max_tokens": 1, '
b'"messages": [{"role": "user", "content": "well over one token of input"}]}}\n'
)
with (
patch("litellm.proxy.proxy_server.general_settings", {}),
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": "openai"},
),
patch("litellm.afile_content", new=AsyncMock(return_value=mock_content)),
):
with pytest.raises(HTTPException) as exc:
await rate_limiter.async_pre_call_hook(
user_api_key_dict=user,
cache=local_cache,
data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"},
call_type="acreate_batch",
)
assert exc.value.status_code == 429
assert "model_per_project_itpm" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_pre_call_enforces_project_otpm_limit_for_non_routing_row_model():
"""VERIA regression: project ITPM/OTPM descriptors were built only for the
file-bound/top-level routing model, so a caller could bind the batch file
to an unlimited model while a JSONL row's own `body.model` named a
different, quota-limited model. That row's tokens must still be charged
against its own model's project OTPM quota."""
from litellm import DualCache
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.utils import InternalUsageCache
local_cache = DualCache()
parallel_request_limiter = _PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=InternalUsageCache(local_cache)
)
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=InternalUsageCache(local_cache),
parallel_request_limiter=parallel_request_limiter,
)
# The routing model ("unlimited-model") has no configured quota; only
# "quota-limited-model" -- named inside the JSONL row, not the routing
# model -- has a project OTPM limit.
user = UserAPIKeyAuth(
api_key="sk-project-batch-cross-model",
models=["*"],
project_id="proj-mantle-batch",
project_metadata={"model_otpm_limit": {"quota-limited-model": 50}},
)
mock_content = MagicMock()
mock_content.content = (
b'{"body": {"model": "quota-limited-model", "max_tokens": 80, '
b'"messages": [{"role": "user", "content": "hi"}]}}\n'
)
with (
patch("litellm.proxy.proxy_server.general_settings", {}),
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": "openai"},
),
patch("litellm.afile_content", new=AsyncMock(return_value=mock_content)),
):
with pytest.raises(HTTPException) as exc:
await rate_limiter.async_pre_call_hook(
user_api_key_dict=user,
cache=local_cache,
data={"input_file_id": "file-abc123", "model": "unlimited-model"},
call_type="acreate_batch",
)
assert exc.value.status_code == 429
assert "model_per_project_otpm" in str(exc.value.detail)
assert "quota-limited-model" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_pre_call_charges_each_row_model_against_its_own_project_quota():
"""A batch whose rows target two different project-quota-limited models
must charge each row's tokens only against its own model's quota, never
the other model's or the whole batch's combined total. The under-limit
model's request must succeed even though the over-limit model's row
would fail on its own."""
from litellm import DualCache
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
PROJECT_OTPM_DESCRIPTOR_KEY,
_PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.utils import InternalUsageCache
local_cache = DualCache()
parallel_request_limiter = _PROXY_MaxParallelRequestsHandler_v3(
internal_usage_cache=InternalUsageCache(local_cache)
)
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=InternalUsageCache(local_cache),
parallel_request_limiter=parallel_request_limiter,
)
user = UserAPIKeyAuth(
api_key="sk-project-batch-two-models",
models=["*"],
project_id="proj-mantle-batch",
project_metadata={
"model_otpm_limit": {"model-a": 1000, "model-b": 10},
},
)
# model-a stays comfortably under its 1000 OTPM limit; model-b's single
# row alone exceeds its 10 OTPM limit. If the two were combined into one
# counter (the pre-fix behavior for the routing model), model-a's ample
# headroom would mask model-b's overage.
mock_content = MagicMock()
mock_content.content = (
b'{"body": {"model": "model-a", "max_tokens": 5, '
b'"messages": [{"role": "user", "content": "hi"}]}}\n'
b'{"body": {"model": "model-b", "max_tokens": 40, '
b'"messages": [{"role": "user", "content": "hi"}]}}\n'
)
with (
patch("litellm.proxy.proxy_server.general_settings", {}),
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": "openai"},
),
patch("litellm.afile_content", new=AsyncMock(return_value=mock_content)),
):
with pytest.raises(HTTPException) as exc:
await rate_limiter.async_pre_call_hook(
user_api_key_dict=user,
cache=local_cache,
data={"input_file_id": "file-abc123", "model": "model-a"},
call_type="acreate_batch",
)
assert exc.value.status_code == 429
assert "model-b" in str(exc.value.detail)
# model-a's own counter was not touched by model-b's rejection: a
# follow-up model-a-only batch well within its own limit must still pass.
model_a_status = await parallel_request_limiter.should_rate_limit(
descriptors=[
{
"key": PROJECT_OTPM_DESCRIPTOR_KEY,
"value": "proj-mantle-batch:model-a",
"rate_limit": {
"requests_per_unit": None,
"tokens_per_unit": 1000,
"window_size": parallel_request_limiter.window_size,
},
}
],
read_only=True,
)
assert model_a_status["overall_code"] == "OK"
def test_should_not_skip_when_project_has_io_limit_for_non_routing_model():
"""The no-limits skip must not fire just because the file-bound/top-level
routing model itself has no configured quota: a JSONL row can name a
different model that the project *does* quota, and that isn't knowable
without downloading and parsing the file."""
rate_limiter = _make_rate_limiter()
# No key/team/model-level limits at all -- only a project OTPM limit for a
# model unrelated to the routing model below.
rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [
{"rate_limit": {}}
]
user = UserAPIKeyAuth(
api_key="sk",
models=["*"],
project_id="proj-mantle-batch",
project_metadata={"model_otpm_limit": {"some-other-model": 50}},
)
with patch("litellm.proxy.proxy_server.general_settings", {}):
should_skip, descriptors = rate_limiter._should_skip_batch_input_file_processing(
data={"model": "unlimited-model", "input_file_id": "file-abc"},
user_api_key_dict=user,
)
assert should_skip is False
assert descriptors is not None
def test_should_skip_when_project_has_no_io_limits_and_no_other_limits():
"""Sanity check for the new project-limits carve-out: a project caller
with no ITPM/OTPM configuration anywhere must still get the fast-path
skip when no other rate limits apply, exactly as before this fix."""
rate_limiter = _make_rate_limiter()
rate_limiter.parallel_request_limiter._create_rate_limit_descriptors.return_value = [
{"rate_limit": {}}
]
user = UserAPIKeyAuth(
api_key="sk",
models=["*"],
project_id="proj-mantle-batch",
project_metadata={},
)
with patch("litellm.proxy.proxy_server.general_settings", {}):
should_skip, descriptors = rate_limiter._should_skip_batch_input_file_processing(
data={"model": "unlimited-model", "input_file_id": "file-abc"},
user_api_key_dict=user,
)
assert should_skip is True
assert descriptors is None
@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={},
)
# Streaming input counting — peak memory must not scale with a full dict list
# ---------------------------------------------------------------------------
def _make_batch_input_bytes(n_rows: int, padding: int = 200) -> bytes:
import json as _json
pad = "x" * padding
rows = []
for i in range(n_rows):
rows.append(
_json.dumps(
{
"custom_id": f"request-{i}",
"method": "POST",
"url": "/v1/chat/completions",
"body": {
"model": "gpt-4o" if i % 2 else "gpt-3.5-turbo",
"messages": [{"role": "user", "content": f"{pad} {i}"}],
},
}
)
)
return ("\n".join(rows)).encode("utf-8")
def test_iter_batch_output_entries_matches_dict_list():
from litellm.batches.batch_utils import (
_get_file_content_as_dictionary,
_iter_batch_output_entries,
)
raw = _make_batch_input_bytes(50)
streamed = list(_iter_batch_output_entries(raw))
assert streamed == _get_file_content_as_dictionary(raw)
assert streamed[0]["custom_id"] == "request-0"
# tolerant of blank lines and a missing trailing newline
assert list(_iter_batch_output_entries(raw + b"\n\n")) == streamed
def test_streaming_count_peak_below_dict_list():
import gc
import tracemalloc
from litellm.batches.batch_utils import (
_get_file_content_as_dictionary,
_iter_batch_output_entries,
)
raw = _make_batch_input_bytes(8000)
def _measure(fn):
gc.collect()
tracemalloc.start()
try:
fn()
_, peak = tracemalloc.get_traced_memory()
finally:
tracemalloc.stop()
return peak
def _stream():
count = 0
models: set = set()
for entry in _iter_batch_output_entries(raw):
count += 1
model = (entry.get("body") or {}).get("model")
if model:
models.add(model)
return count
def _build_list():
return len(_get_file_content_as_dictionary(raw))
stream_peak = _measure(_stream)
list_peak = _measure(_build_list)
assert stream_peak < list_peak * 0.5, (
f"streaming count peak {stream_peak} is not a clear win over the dict "
f"list {list_peak} (ratio {stream_peak / list_peak:.2f})"
)
@pytest.mark.asyncio
async def test_count_input_file_usage_streams_without_building_list():
"""count_input_file_usage must count requests/tokens in one streaming pass.
Mocks the download; asserts the count is correct and that the dict-list
helper is never called (a revert to the list approach would call it)."""
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
raw = _make_batch_input_bytes(10)
fake_content = MagicMock()
fake_content.content = raw
with (
patch("litellm.afile_content", new=AsyncMock(return_value=fake_content)),
patch(
"litellm.batches.batch_utils._get_file_content_as_dictionary"
) as mock_dict_list,
):
usage = await rate_limiter.count_input_file_usage(
file_id="file-not-managed",
custom_llm_provider="openai",
user_api_key_dict=None,
)
assert usage.request_count == 10
assert usage.total_tokens > 0
mock_dict_list.assert_not_called()
def _one_row_batch_bytes(model: str) -> bytes:
import json as _json
return (
_json.dumps(
{
"custom_id": "r0",
"method": "POST",
"url": "/v1/chat/completions",
"body": {
"model": model,
"messages": [{"role": "user", "content": "x"}],
},
}
)
+ "\n"
).encode("utf-8")
@pytest.mark.asyncio
async def test_count_input_file_usage_enforces_models_when_token_counting_fails():
"""Security regression: a row whose content makes token counting raise must
NOT skip the model allowlist check. async_pre_call_hook swallows non-HTTP
exceptions and submits the batch, so a raised counting error would otherwise
fail open. The access check must still run and deny the restricted model."""
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
fake_content = MagicMock()
fake_content.content = _one_row_batch_bytes("restricted-model")
user = UserAPIKeyAuth(
api_key="sk-x",
user_id="bob",
models=["only-allowed"],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
def _boom(*args, **kwargs):
raise ValueError("unsupported content part: input_audio")
deny = AsyncMock(side_effect=Exception("model not in allowlist"))
with (
patch("litellm.afile_content", new=AsyncMock(return_value=fake_content)),
patch("litellm.proxy.hooks.batch_rate_limiter._count_entry_tokens", new=_boom),
patch("litellm.proxy.auth.auth_checks.can_key_call_model", new=deny),
patch("litellm.proxy.proxy_server.llm_router", MagicMock(model_list=[])),
):
with pytest.raises(HTTPException) as exc:
await rate_limiter.count_input_file_usage(
file_id="file-not-managed",
custom_llm_provider="openai",
user_api_key_dict=user,
)
# The access check ran despite token counting failing, and denied the model.
deny.assert_awaited()
assert exc.value.status_code == 403
@pytest.mark.asyncio
async def test_count_input_file_usage_estimates_tokens_when_counting_fails_for_allowed_model():
"""A token-counting failure for an allowed model must not hard-block the batch
(the pre-streaming behavior let such batches through), but it also must not
zero the token total, which would let a caller evade the TPM limit by sending
rows the counter cannot measure. The row falls back to a conservative
size-based estimate so the batch proceeds with a non-zero count."""
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
fake_content = MagicMock()
fake_content.content = _one_row_batch_bytes("allowed-model")
user = UserAPIKeyAuth(
api_key="sk-x",
user_id="bob",
models=["allowed-model"],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
def _boom(*args, **kwargs):
raise ValueError("unsupported content part: file")
allow = AsyncMock(return_value=True)
with (
patch("litellm.afile_content", new=AsyncMock(return_value=fake_content)),
patch("litellm.proxy.hooks.batch_rate_limiter._count_entry_tokens", new=_boom),
patch("litellm.proxy.auth.auth_checks.can_key_call_model", new=allow),
patch("litellm.proxy.proxy_server.llm_router", MagicMock(model_list=[])),
):
usage = await rate_limiter.count_input_file_usage(
file_id="file-not-managed",
custom_llm_provider="openai",
user_api_key_dict=user,
)
allow.assert_awaited()
assert usage.request_count == 1
# Estimated, not zeroed: a crafted uncountable row can't evade the TPM limit.
assert usage.total_tokens > 0
@pytest.mark.asyncio
async def test_count_input_file_usage_collects_models_after_malformed_line():
"""A malformed JSONL line must not abort model collection. A restricted model
named on a row AFTER a malformed line must still be collected and denied by the
allowlist check, otherwise a caller could hide a restricted model behind a bad
row."""
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=MagicMock(),
)
fake_content = MagicMock()
fake_content.content = (
_one_row_batch_bytes("only-allowed")
+ b"{ this is not valid json\n"
+ _one_row_batch_bytes("restricted-model")
)
user = UserAPIKeyAuth(
api_key="sk-x",
user_id="bob",
models=["only-allowed"],
user_role=LitellmUserRoles.INTERNAL_USER.value,
)
async def _deny_restricted(model, **kwargs):
if model == "restricted-model":
raise Exception("model not in allowlist")
return True
deny = AsyncMock(side_effect=_deny_restricted)
with (
patch("litellm.afile_content", new=AsyncMock(return_value=fake_content)),
patch("litellm.proxy.auth.auth_checks.can_key_call_model", new=deny),
patch("litellm.proxy.proxy_server.llm_router", MagicMock(model_list=[])),
):
with pytest.raises(HTTPException) as exc:
await rate_limiter.count_input_file_usage(
file_id="file-not-managed",
custom_llm_provider="openai",
user_api_key_dict=user,
)
assert exc.value.status_code == 403
# ---------------------------------------------------------------------------
# VERIA-Low regression: Responses batch rows must not bypass project OTPM
# ---------------------------------------------------------------------------
def _output_estimator():
"""A `_PROXY_BatchRateLimiter` whose output-token floor is observable:
the no-`max_tokens` floor mock returns a distinctive sentinel so tests can
tell "floor was used" apart from "an explicit cap was read"."""
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_PROXY_MaxParallelRequestsHandler_v3,
)
limiter = MagicMock()
limiter.no_max_tokens_output_floor.return_value = 999
limiter.get_output_candidate_count = _PROXY_MaxParallelRequestsHandler_v3.get_output_candidate_count
return _PROXY_BatchRateLimiter(
internal_usage_cache=MagicMock(),
parallel_request_limiter=limiter,
)
def test_estimate_entry_output_tokens_zero_for_embeddings_url():
"""A real `/v1/embeddings` row reserves zero output tokens."""
rate_limiter = _output_estimator()
entry = {
"url": "/v1/embeddings",
"body": {"model": "text-embedding-3-small", "input": "hello world"},
}
assert rate_limiter._estimate_entry_output_tokens(entry, None) == 0
def test_estimate_entry_output_tokens_does_not_zero_responses_row_with_input():
"""Pre-fix: a `/v1/responses` row carries `body.input` with no `messages`/
`prompt`, so the old body-shape heuristic misclassified it as embeddings
and reserved zero output tokens -- a project caller could submit large
Responses generations against a quota-limited model without consuming
OTPM. The row's own `url` (not body shape) must decide this."""
rate_limiter = _output_estimator()
entry = {
"url": "/v1/responses",
"body": {"model": "gpt-4o", "input": "write me an essay"},
}
# No explicit cap on the row, so it must fall back to the no-max-tokens
# floor -- never straight to zero.
assert rate_limiter._estimate_entry_output_tokens(entry, None) == 999
rate_limiter.parallel_request_limiter.no_max_tokens_output_floor.assert_called_once_with(None)
def test_estimate_entry_output_tokens_uses_max_output_tokens_for_responses():
"""`/v1/responses` caps output with `max_output_tokens`, not `max_tokens`/
`max_completion_tokens`. Pre-fix this field was never inspected, so a
capped Responses row still fell through to the (possibly larger) floor
estimate instead of the caller's own declared cap."""
rate_limiter = _output_estimator()
entry = {
"url": "/v1/responses",
"body": {"model": "gpt-4o", "input": "hi", "max_output_tokens": 123},
}
assert rate_limiter._estimate_entry_output_tokens(entry, None) == 123
rate_limiter.parallel_request_limiter.no_max_tokens_output_floor.assert_not_called()
def test_estimate_entry_output_tokens_prefers_max_tokens_over_max_output_tokens():
"""When a row somehow carries both fields, the chat-style cap wins first --
`max_output_tokens` is only consulted once the chat-style caps are absent."""
rate_limiter = _output_estimator()
entry = {
"url": "/v1/chat/completions",
"body": {
"model": "gpt-4o",
"messages": [],
"max_tokens": 50,
"max_output_tokens": 500,
},
}
assert rate_limiter._estimate_entry_output_tokens(entry, None) == 50
@pytest.mark.parametrize(
("body_extra", "expected"),
[
({"max_tokens": 40, "n": 10}, 400),
({"max_tokens": 40, "best_of": 5}, 200),
({"max_tokens": 40, "n": 3, "best_of": 5}, 200),
({"max_tokens": 40, "n": 0}, 40),
({"max_tokens": 40, "n": -2}, 40),
({"max_tokens": 40, "n": 5.0}, 200),
({"max_tokens": 40, "n": "10"}, 400),
({"max_tokens": 40, "n": "not-a-number"}, 40),
({"max_tokens": 40, "n": 1e309}, 40),
({"max_tokens": 1e309, "n": 3}, 2997),
({"n": 3}, 2997),
],
)
def test_estimate_entry_output_tokens_multiplies_candidate_count(body_extra, expected):
"""A row generating n / best_of candidates consumes that many completions'
worth of output tokens, so the OTPM reservation must scale with the
effective candidate count. Pre-fix a `max_tokens: 40, n: 10` row consumed
up to 400 output tokens while reserving only 40."""
rate_limiter = _output_estimator()
entry = {
"url": "/v1/chat/completions",
"body": {"model": "gpt-4o", "messages": [], **body_extra},
}
assert rate_limiter._estimate_entry_output_tokens(entry, None) == expected
# ---------------------------------------------------------------------------
# LIT-5273: enqueued-token limits govern batch submission when opted in
# ---------------------------------------------------------------------------
def _enqueued_rate_limiter():
from litellm import DualCache
from litellm.proxy.hooks.batch_rate_limiter import _PROXY_BatchRateLimiter
from litellm.proxy.hooks.parallel_request_limiter_v3 import (
_PROXY_MaxParallelRequestsHandler_v3,
)
from litellm.proxy.utils import InternalUsageCache
local_cache = DualCache(default_in_memory_ttl=60)
internal_usage_cache = InternalUsageCache(local_cache)
parallel_request_limiter = _PROXY_MaxParallelRequestsHandler_v3(internal_usage_cache=internal_usage_cache)
rate_limiter = _PROXY_BatchRateLimiter(
internal_usage_cache=internal_usage_cache,
parallel_request_limiter=parallel_request_limiter,
)
return rate_limiter, local_cache
_ENQUEUED_BATCH_FILE_CONTENT = (
b'{"body": {"model": "gpt-4o-mini", "max_tokens": 40, "messages": [{"role": "user", "content": "hi"}]}}\n'
b'{"body": {"model": "gpt-4o-mini", "max_tokens": 40, "messages": [{"role": "user", "content": "hi"}]}}\n'
b'{"body": {"model": "gpt-4o-mini", "max_tokens": 40, "messages": [{"role": "user", "content": "hi"}]}}\n'
)
def _enqueued_batch_patches():
mock_content = MagicMock()
mock_content.content = _ENQUEUED_BATCH_FILE_CONTENT
afile_content_mock = AsyncMock(return_value=mock_content)
return afile_content_mock, (
patch("litellm.proxy.proxy_server.general_settings", {}),
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": "openai"},
),
)
@pytest.mark.asyncio
async def test_enqueued_limit_accepts_batch_over_per_minute_limits():
"""The headline LIT-5273 behavior: a key that opted into an enqueued-token
allowance submits a batch whose row count and token count both exceed its
per-minute RPM/TPM limits, and the batch is accepted (repeatedly) because
only the enqueued allowance governs. Without the opt-in the same key is
rejected on RPM before the batch reaches the provider."""
from litellm.proxy.hooks.parallel_request_limiter_v3 import get_request_stash
rate_limiter, local_cache = _enqueued_rate_limiter()
afile_content_mock, patches = _enqueued_batch_patches()
legacy_user = UserAPIKeyAuth(api_key="sk-legacy-rpm", models=["*"], rpm_limit=1, tpm_limit=10)
opted_in_user = UserAPIKeyAuth(
api_key="sk-enqueued-rpm",
models=["*"],
rpm_limit=1,
tpm_limit=10,
metadata={"batch_enqueued_token_limit": 100000},
)
with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock):
with pytest.raises(HTTPException) as legacy_exc:
await rate_limiter.async_pre_call_hook(
user_api_key_dict=legacy_user,
cache=local_cache,
data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"},
call_type="acreate_batch",
)
assert legacy_exc.value.status_code == 429
first_data = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"}
result = await rate_limiter.async_pre_call_hook(
user_api_key_dict=opted_in_user,
cache=local_cache,
data=first_data,
call_type="acreate_batch",
)
assert result is first_data
stash = get_request_stash()
assert stash is not None and stash.batch_enqueued_reservation is not None
assert stash.batch_enqueued_reservation.tokens == first_data["_batch_token_count"] > 0
second = await rate_limiter.async_pre_call_hook(
user_api_key_dict=opted_in_user,
cache=local_cache,
data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"},
call_type="acreate_batch",
)
assert second is not None
@pytest.mark.asyncio
async def test_enqueued_limit_rejects_when_allowance_is_exhausted():
"""Submissions are rejected pre-provider once the enqueued allowance can't
fit the batch, even for a key with no per-minute limits at all (which
previously skipped batch rate limiting entirely)."""
rate_limiter, local_cache = _enqueued_rate_limiter()
afile_content_mock, patches = _enqueued_batch_patches()
sizing_user = UserAPIKeyAuth(
api_key="sk-enqueued-sizing", models=["*"], metadata={"batch_enqueued_token_limit": 1000000}
)
with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock):
sizing_data = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"}
await rate_limiter.async_pre_call_hook(
user_api_key_dict=sizing_user,
cache=local_cache,
data=sizing_data,
call_type="acreate_batch",
)
batch_tokens = sizing_data["_batch_token_count"]
assert batch_tokens > 0
capped_user = UserAPIKeyAuth(
api_key="sk-enqueued-capped",
models=["*"],
metadata={"batch_enqueued_token_limit": batch_tokens + batch_tokens // 2},
)
await rate_limiter.async_pre_call_hook(
user_api_key_dict=capped_user,
cache=local_cache,
data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"},
call_type="acreate_batch",
)
with pytest.raises(HTTPException) as exc:
await rate_limiter.async_pre_call_hook(
user_api_key_dict=capped_user,
cache=local_cache,
data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"},
call_type="acreate_batch",
)
assert exc.value.status_code == 429
assert "Batch enqueued token limit exceeded for api_key" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_enqueued_team_limit_applies_to_batch_submission():
rate_limiter, local_cache = _enqueued_rate_limiter()
afile_content_mock, patches = _enqueued_batch_patches()
team_user = UserAPIKeyAuth(
api_key="sk-enqueued-team-key",
models=["*"],
team_id="team-enqueued-batch",
team_metadata={"batch_enqueued_token_limit": 10},
)
with patches[0], patches[1], patches[2], patch("litellm.afile_content", new=afile_content_mock):
with pytest.raises(HTTPException) as exc:
await rate_limiter.async_pre_call_hook(
user_api_key_dict=team_user,
cache=local_cache,
data={"input_file_id": "file-abc123", "model": "gpt-4o-mini"},
call_type="acreate_batch",
)
assert exc.value.status_code == 429
assert "Batch enqueued token limit exceeded for team: team-enqueued-batch" in str(exc.value.detail)
@pytest.mark.asyncio
async def test_disable_flag_still_skips_batch_processing_with_enqueued_limits():
rate_limiter, local_cache = _enqueued_rate_limiter()
afile_content_mock, _ = _enqueued_batch_patches()
opted_in_user = UserAPIKeyAuth(
api_key="sk-enqueued-disabled",
models=["*"],
metadata={"batch_enqueued_token_limit": 10},
)
with (
patch("litellm.proxy.proxy_server.general_settings", {"disable_batch_input_file_rate_limiting": True}),
patch("litellm.proxy.proxy_server.llm_router", MagicMock()),
patch("litellm.afile_content", new=afile_content_mock),
):
data = {"input_file_id": "file-abc123", "model": "gpt-4o-mini"}
result = await rate_limiter.async_pre_call_hook(
user_api_key_dict=opted_in_user,
cache=local_cache,
data=data,
call_type="acreate_batch",
)
assert result is data
afile_content_mock.assert_not_awaited()
Reference in a new issue View git blame Copy permalink