diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index 4b965d4e63..aaf083e75d 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -387,6 +387,27 @@ def _get_batch_job_total_usage_from_file_content( ) +def _get_models_from_batch_input_file_content( + file_content_dictionary: List[dict], +) -> List[str]: + """Extract the distinct ``body.model`` values from a batch *input* file. + + Used by the proxy's batch pre-call hook to enforce that the caller is + authorized for every model named inside the JSONL — not just the one + on the outer request — so the proxy's per-key model allowlist isn't + bypassed by smuggling expensive models into the batch file. + """ + models: List[str] = [] + seen: set = set() + for _item in file_content_dictionary: + body = _item.get("body") or {} + model = body.get("model") + if model and model not in seen: + seen.add(model) + models.append(model) + return models + + def _get_batch_job_input_file_usage( file_content_dictionary: List[dict], custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", @@ -403,11 +424,25 @@ def _get_batch_job_input_file_usage( for _item in file_content_dictionary: body = _item.get("body", {}) model = body.get("model", model_name or "") - messages = body.get("messages", []) + # Chat completion payloads. + messages = body.get("messages") if messages: - item_tokens = token_counter(model=model, messages=messages) - prompt_tokens += item_tokens + prompt_tokens += token_counter(model=model, messages=messages) + continue + + # Text completion payloads (`prompt`). + prompt = body.get("prompt") + if prompt: + prompt_tokens += _count_prompt_or_input_tokens(model=model, value=prompt) + continue + + # Embedding payloads (`input`). + input_data = body.get("input") + if input_data: + prompt_tokens += _count_prompt_or_input_tokens( + model=model, value=input_data + ) return Usage( total_tokens=prompt_tokens + completion_tokens, @@ -416,6 +451,43 @@ def _get_batch_job_input_file_usage( ) +def _count_prompt_or_input_tokens(model: str, value: Any) -> int: + """Token-count a ``prompt`` / ``input`` field that the OpenAI batch + schema allows in four shapes: + + - ``str``: a single text prompt. + - ``list[str]``: multiple text prompts. + - ``list[int]``: a pre-tokenized prompt (each int counts as 1 token). + - ``list[list[int]]``: multiple pre-tokenized prompts. + + Pre-fix only the string shapes were counted, so a caller could send + a large ``list[list[int]]`` payload and slip past TPM rate limits + with a recorded cost of zero tokens. + """ + if isinstance(value, str): + return token_counter(model=model, text=value) + if isinstance(value, list): + total = 0 + for chunk in value: + if isinstance(chunk, str): + total += token_counter(model=model, text=chunk) + elif isinstance(chunk, int): + # Single pre-tokenized prompt at the top level: each + # int counts as one token. + total += 1 + elif isinstance(chunk, list): + # Nested pre-tokenized prompt: every int contributes a + # token. Mixed string/int items still count. + total += sum(1 if isinstance(t, int) else 0 for t in chunk) + total += sum( + token_counter(model=model, text=t) + for t in chunk + if isinstance(t, str) + ) + return total + return 0 + + def _get_batch_job_usage_from_response_body(response_body: dict) -> Usage: """ Get the tokens of a batch job from the response body diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 2ee0588f19..f740d5dd40 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -27,6 +27,7 @@ from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import ( _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 @@ -246,6 +247,17 @@ class _PROXY_BatchRateLimiter(CustomLogger): file_content_as_dict = _get_file_content_as_dictionary(file_content.content) + # Validate every model named in the batch JSONL against the + # caller's per-key model allowlist. Without this, a caller + # could smuggle restricted/expensive models inside the file + # and the upstream provider would execute the batch under + # the proxy's shared API key. + if user_api_key_dict is not None: + await self._enforce_batch_file_model_access( + user_api_key_dict=user_api_key_dict, + file_content_as_dict=file_content_as_dict, + ) + input_file_usage = _get_batch_job_input_file_usage( file_content_dictionary=file_content_as_dict, custom_llm_provider=custom_llm_provider, @@ -256,12 +268,69 @@ class _PROXY_BatchRateLimiter(CustomLogger): request_count=request_count, ) + except HTTPException as e: + # Distinguish intentional 403s from `_enforce_batch_file_model_access` + # from genuine I/O failures so security-relevant rejections show up + # in the access log instead of getting buried in error noise. + if e.status_code == 403: + verbose_proxy_logger.warning( + f"Batch rejected: caller not authorized for a model named in {file_id}: {e.detail}" + ) + else: + verbose_proxy_logger.error( + f"Batch input file rejected for {file_id}: status={e.status_code} detail={e.detail}" + ) + raise except Exception as e: verbose_proxy_logger.error( f"Error counting input file usage for {file_id}: {str(e)}" ) raise + async def _enforce_batch_file_model_access( + self, + user_api_key_dict: UserAPIKeyAuth, + file_content_as_dict: List[dict], + ) -> None: + """Reject the batch if the caller is not authorized for every + ``body.model`` named inside the JSONL. + + Reuses ``can_key_call_model`` so the same allowlist semantics + (wildcards, access groups, ``all-proxy-models``, team aliases) + the proxy enforces on `/chat/completions` apply here. + """ + from litellm.proxy.auth.auth_checks import can_key_call_model + from litellm.proxy.proxy_server import llm_router + + models = _get_models_from_batch_input_file_content(file_content_as_dict) + if not models: + return + + llm_model_list = llm_router.model_list if llm_router is not None else None + for model in models: + try: + await can_key_call_model( + model=model, + llm_model_list=llm_model_list, + valid_token=user_api_key_dict, + llm_router=llm_router, + ) + except HTTPException: + raise + except Exception as e: + # `can_key_call_model` raises ProxyException on denial; + # re-shape to a 403 so the batch endpoint returns a + # consistent rejection without leaking internal types. + raise HTTPException( + status_code=403, + detail={ + "error": ( + "Batch input file references a model the caller is " + f"not authorized to use: model={model}, reason={str(e)}" + ) + }, + ) + async def _fetch_managed_file_content( self, file_id: str, diff --git a/tests/test_litellm/proxy/hooks/test_batch_file_validation.py b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py new file mode 100644 index 0000000000..7f1006543b --- /dev/null +++ b/tests/test_litellm/proxy/hooks/test_batch_file_validation.py @@ -0,0 +1,285 @@ +""" +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 + + +# --------------------------------------------------------------------------- +# Token counter — covers all three batch payload shapes +# --------------------------------------------------------------------------- + + +def test_token_counter_counts_chat_messages(): + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "gpt-4o-mini", + "messages": [{"role": "user", "content": "hello"}], + } + } + ] + ) + assert usage.prompt_tokens > 0 + + +def test_token_counter_counts_text_completion_prompt(): + """Pre-fix this returned 0 tokens (the function only inspected + `messages`), letting `prompt`-style batches slip past TPM limits.""" + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + {"body": {"model": "gpt-3.5-turbo-instruct", "prompt": "hello world"}} + ] + ) + assert usage.prompt_tokens > 0 + + +def test_token_counter_counts_embedding_input_string(): + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + {"body": {"model": "text-embedding-3-small", "input": "hello world"}} + ] + ) + assert usage.prompt_tokens > 0 + + +def test_token_counter_counts_embedding_input_list(): + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "text-embedding-3-small", + "input": ["hello", "world"], + } + } + ] + ) + assert usage.prompt_tokens > 0 + + +def test_token_counter_counts_text_completion_prompt_list(): + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": ["alpha", "beta"], + } + } + ] + ) + assert usage.prompt_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 _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": [1, 2, 3, 4, 5], + } + } + ] + ) + assert usage.prompt_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 _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "gpt-3.5-turbo-instruct", + "prompt": [[1] * 250, [2] * 250, [3] * 500], + } + } + ] + ) + assert usage.prompt_tokens == 1000 + + +def test_token_counter_counts_pre_tokenized_input_for_embeddings(): + """Same shape applies to embeddings (`input`).""" + from litellm.batches.batch_utils import _get_batch_job_input_file_usage + + usage = _get_batch_job_input_file_usage( + file_content_dictionary=[ + { + "body": { + "model": "text-embedding-3-small", + "input": [[1, 2, 3], [4, 5, 6]], + } + } + ] + ) + assert usage.prompt_tokens == 6 + + +# --------------------------------------------------------------------------- +# Model extractor +# --------------------------------------------------------------------------- + + +def test_model_extractor_returns_distinct_models(): + from litellm.batches.batch_utils import _get_models_from_batch_input_file_content + + models = _get_models_from_batch_input_file_content( + [ + {"body": {"model": "gpt-4o", "messages": []}}, + {"body": {"model": "gpt-4o", "messages": []}}, # duplicate + {"body": {"model": "gpt-4o-mini", "messages": []}}, + {"body": {}}, # missing model + ] + ) + assert models == ["gpt-4o", "gpt-4o-mini"] + + +# --------------------------------------------------------------------------- +# 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, + file_content_as_dict=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_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, + file_content_as_dict=file_dict, + ) + + +@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, + file_content_as_dict=[], + ) + await rate_limiter._enforce_batch_file_model_access( + user_api_key_dict=user, + file_content_as_dict=[{"body": {}}], + )