fix(batches): count non-chat tokens, validate batch-file model access (VERIA-39)

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

litellm/proxy/hooks/batch_rate_limiter.py

@@ -27,6 +27,7 @@
 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
@@ -288,6 +289,17 @@ async def count_input_file_usage(
 
             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,
@@ -298,12 +310,69 @@ async def count_input_file_usage(
                 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,