fix(proxy): rewrite only alias rows and authorize them against the alias at batch create

litellm/litellm_core_utils/batch_utils.py

@@ -155,7 +155,12 @@ def should_replace_model_in_jsonl(
     return purpose == "batch" and not passthrough
 
 
-def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str, keep_line_numbers: bool = False) -> FileTypes:
+def replace_model_in_jsonl(
+    file_content: FileTypes,
+    new_model_name: str,
+    keep_line_numbers: bool = False,
+    only_rows_naming: str | None = None,
+) -> FileTypes:
     try:
         ## if pathlike, return the original file content
         if isinstance(file_content, PathLike):
@@ -209,7 +214,11 @@ def replace_model_in_jsonl(file_content: FileTypes, new_model_name: str, keep_li
                 json_object = _DECODED_JSONL_ROW.validate_python(json.loads(stripped))
             except json.JSONDecodeError:
                 continue  # object not complete yet; keep accumulating
-            if isinstance(json_object, dict) and isinstance(json_object.get("body"), dict):
+            if (
+                isinstance(json_object, dict)
+                and isinstance(json_object.get("body"), dict)
+                and (only_rows_naming is None or json_object["body"].get("model") == only_rows_naming)
+            ):
                 json_object["body"]["model"] = new_model_name
             output.write(
                 (

litellm/proxy/hooks/batch_rate_limiter.py

@@ -223,6 +223,33 @@ class _PROXY_BatchRateLimiter(CustomLogger):
 
         return None
 
+    def _routing_alias_by_row_model(self, routing_alias: str | None) -> Mapping[str, str]:
+        """Map the model a model-routed upload writes into its rows back to the alias it routed by.
+
+        A key's allowlist and project quotas name the proxy alias, while the uploaded
+        rows name the deployment's model. A row that maps here runs on the alias's own
+        deployment, so the alias is what the key has to be granted and charged for.
+        """
+        from litellm.proxy.openai_files_endpoints.common_utils import (
+            batch_row_model_for_deployment,
+            get_credentials_for_model,
+        )
+        from litellm.proxy.proxy_server import llm_router
+
+        if routing_alias is None or llm_router is None:
+            return MappingProxyType({})
+        try:
+            credentials: Final = get_credentials_for_model(
+                llm_router=llm_router,
+                model_id=routing_alias,
+                operation_context="batch input file read (rate limiting)",
+            )
+        except HTTPException:
+            return MappingProxyType({})
+        return MappingProxyType(
+            {batch_row_model_for_deployment(credentials["model"], credentials["custom_llm_provider"]): routing_alias}
+        )
+
     def _resolve_batch_provider(self, batch_model: str | None) -> str | None:
         """Resolve the provider from the deployment that serves ``batch_model``.
 
@@ -942,6 +969,13 @@ class _PROXY_BatchRateLimiter(CustomLogger):
             # For managed files the unified file id encodes the proxy model
             # alias(es) the file was uploaded for; auth validates against those.
             target_model_names: Final = get_models_from_unified_file_id(is_managed_file) if is_managed_file else []
+            routing_alias_by_row_model: Final = (
+                MappingProxyType({})
+                if is_managed_file
+                else self._routing_alias_by_row_model(
+                    (self._get_file_bound_batch_model(data) or _batch_create_routing_model(data)) if data else None
+                )
+            )
             if is_managed_file and user_api_key_dict is not None:
                 file_content = await self._fetch_managed_file_content(
                     file_id=file_id,
@@ -1001,7 +1035,8 @@ class _PROXY_BatchRateLimiter(CustomLogger):
                     output_tokens += entry_output_tokens
                     continue
 
-                model: str | None = (entry.get("body") or {}).get("model") if isinstance(entry, dict) else None
+                row_model: str | None = (entry.get("body") or {}).get("model") if isinstance(entry, dict) else None
+                model: str | None = routing_alias_by_row_model.get(row_model, row_model) if row_model else None
                 if model:
                     models.add(model)
 

litellm/proxy/openai_files_endpoints/common_utils.py

@@ -22,6 +22,7 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict
 
 from litellm.batches.batch_utils import batch_cost_is_final
 from litellm.constants import MAX_FILE_LIST_LIMIT
+from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
 from litellm.proxy._types import ProxyException
 from litellm.repositories.managed_file_repository import ManagedFileRepository
 from litellm.repositories.table_repositories import (
@@ -477,6 +478,10 @@ async def get_authorized_credentials_for_model(
     )
 
 
+def batch_row_model_for_deployment(deployment_model: str, custom_llm_provider: str) -> str:
+    return get_llm_provider(model=deployment_model, custom_llm_provider=custom_llm_provider)[0]
+
+
 def deployment_drop_params(
     llm_router: Optional["Router"], model_id: str, request: Mapping[str, object]
 ) -> Mapping[str, bool | str]:

litellm/proxy/openai_files_endpoints/files_endpoints.py

@@ -38,7 +38,6 @@ from litellm.litellm_core_utils.cloud_storage_security import (
     is_managed_cloud_storage_uri,
 )
 from litellm.litellm_core_utils.core_helpers import get_or_create_metadata_bucket
-from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider
 from litellm.llms.base_llm.files.litellm_db_storage_backend import LITELLM_DB_STORAGE_BACKEND_NAME
 from litellm.llms.base_llm.files.transformation import BaseFileEndpoints, FileListModelScope
 from litellm.llms.base_llm.managed_resources.isolation import build_list_page
@@ -83,6 +82,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import (  # noqa: F401  #
     add_internal_model_credentials,
     apply_team_provider_credentials,
     authorize_model_for_key,
+    batch_row_model_for_deployment,
     deployment_drop_params,
     encode_file_id_with_model,
     extract_file_creation_params,
@@ -383,12 +383,14 @@ def get_model_from_json_obj(json_object: dict) -> str | None:
     return body.get("model") if isinstance(body, dict) else None
 
 
-def _rewrite_batch_rows_to_deployment_model(file: FileTypes, deployment_model: str, provider: str) -> FileTypes:
-    stripped_model: Final = get_llm_provider(model=deployment_model, custom_llm_provider=provider)[0]
+def _rewrite_alias_rows_to_deployment_model(
+    file: FileTypes, alias: str, deployment_model: str, provider: str
+) -> FileTypes:
     rewritten: Final = replace_model_in_jsonl(
         file_content=file,
-        new_model_name=stripped_model,
+        new_model_name=batch_row_model_for_deployment(deployment_model, provider),
         keep_line_numbers=provider == LlmProviders.VERTEX_AI.value,
+        only_rows_naming=alias,
     )
     source: Final = file[1] if isinstance(file, tuple) else file
     if rewritten is file:
@@ -493,8 +495,9 @@ async def route_create_file(
             {
                 **_create_file_request,
                 "file": await asyncio.to_thread(
-                    _rewrite_batch_rows_to_deployment_model,
+                    _rewrite_alias_rows_to_deployment_model,
                     _create_file_request["file"],
+                    model,
                     credentials["model"],
                     credentials["custom_llm_provider"],
                 ),

tests/unit/proxy/hooks/test_batch_rate_limiter.py

@@ -306,6 +306,7 @@ async def test_batch_rate_limit_error_reports_reset_time_in_utc_on_a_non_utc_pro
 
 
 _BATCH_MODEL: Final = "gpt-3.5-turbo"
+_MANAGED_DEPLOYMENT_ALIAS: Final = "managed-deployment-alias"
 _MANAGED_FILE_ID: Final = (
     "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9vY3RldC1zdHJlYW07dW5pZmllZF9pZCxyZWdyZXNzaW9uLXRlc3QtZmlsZQ=="
 )
@@ -666,7 +667,12 @@ async def _install_managed_file_hook(
                 "model_name": _BATCH_MODEL,
                 "litellm_params": {"model": f"openai/{_BATCH_MODEL}", "api_key": "sk-test"},
                 "model_info": {"id": "deployment-1"},
-            }
+            },
+            {
+                "model_name": _MANAGED_DEPLOYMENT_ALIAS,
+                "litellm_params": {"model": f"openai/{_BATCH_MODEL}", "api_key": "sk-test"},
+                "model_info": {"id": "deployment-1-alias"},
+            },
         ]
     )
     monkeypatch.setattr(proxy_server, "proxy_logging_obj", proxy_logging)
@@ -720,3 +726,145 @@ async def test_managed_file_of_another_tenant_is_denied_before_it_is_read(
     assert denied.value.status_code == 403
     assert not provider_route.called
     assert not unrouted_managed_read.called
+
+
+_BATCH_ALIAS: Final = "nano-batch-alias"
+_DEPLOYMENT_MODEL: Final = "nano-deployment-model"
+
+
+def _row_naming(model: str) -> _BatchLine:
+    return _BatchLine(
+        custom_id=model,
+        method="POST",
+        url="/v1/chat/completions",
+        body=_BatchBody(model=model, messages=(ChatCompletionUserMessage(role="user", content="hi"),)),
+    )
+
+
+@pytest.fixture
+def alias_routed_batch(openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> respx.MockRouter:
+    from litellm import Router
+    from litellm.proxy import proxy_server
+
+    monkeypatch.setattr(
+        proxy_server,
+        "llm_router",
+        Router(
+            model_list=[
+                {
+                    "model_name": _BATCH_ALIAS,
+                    "litellm_params": {
+                        "model": f"openai/{_DEPLOYMENT_MODEL}",
+                        "api_key": "sk-deployment",
+                        "api_base": "https://routed.example/v1",
+                    },
+                    "model_info": {"id": "alias-deployment"},
+                }
+            ]
+        ),
+    )
+    return openai_files
+
+
+def _alias_routing(route: Literal["file id", "body"]) -> dict[str, str]:
+    from litellm.proxy.openai_files_endpoints.common_utils import encode_file_id_with_model
+
+    if route == "file id":
+        return {"input_file_id": encode_file_id_with_model("file-abc123", _BATCH_ALIAS)}
+    return {"input_file_id": "file-abc123", "model": _BATCH_ALIAS}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("route", ["file id", "body"])
+async def test_key_granted_only_the_alias_can_batch_rows_naming_its_deployment_model(
+    alias_routed_batch: respx.MockRouter, route: Literal["file id", "body"]
+):
+    alias_routed_batch.get(_ROUTED_FILE_URL).mock(
+        side_effect=_serve_file_to_the_deployment_key((_row_naming(_DEPLOYMENT_MODEL),))
+    )
+    data: Final = _alias_routing(route)
+    _, _, batch_limiter = _make_limiters()
+
+    usage: Final = await batch_limiter.count_input_file_usage(
+        file_id=data["input_file_id"],
+        user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-batch"), models=[_BATCH_ALIAS]),
+        data=dict(data),
+    )
+
+    assert set(usage.per_model_usage) == {_BATCH_ALIAS}
+    assert usage.per_model_usage[_BATCH_ALIAS]["request_count"] == 1
+
+
+@pytest.mark.asyncio
+async def test_key_granted_the_alias_cannot_smuggle_another_model_into_the_rows(
+    alias_routed_batch: respx.MockRouter,
+):
+    alias_routed_batch.get(_ROUTED_FILE_URL).mock(
+        side_effect=_serve_file_to_the_deployment_key((_row_naming(_DEPLOYMENT_MODEL), _row_naming("another-model")))
+    )
+    data: Final = _alias_routing("body")
+    _, _, batch_limiter = _make_limiters()
+
+    with pytest.raises(HTTPException) as denied:
+        await batch_limiter.count_input_file_usage(
+            file_id=data["input_file_id"],
+            user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-batch"), models=[_BATCH_ALIAS]),
+            data=dict(data),
+        )
+
+    assert denied.value.status_code == 403
+    assert "another-model" in str(denied.value.detail)
+
+
+@pytest.mark.asyncio
+async def test_key_without_the_alias_cannot_batch_rows_naming_its_deployment_model(
+    alias_routed_batch: respx.MockRouter,
+):
+    alias_routed_batch.get(_ROUTED_FILE_URL).mock(
+        side_effect=_serve_file_to_the_deployment_key((_row_naming(_DEPLOYMENT_MODEL),))
+    )
+    data: Final = _alias_routing("body")
+    _, _, batch_limiter = _make_limiters()
+
+    with pytest.raises(HTTPException) as denied:
+        await batch_limiter.count_input_file_usage(
+            file_id=data["input_file_id"],
+            user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-batch"), models=["a-different-model"]),
+            data=dict(data),
+        )
+
+    assert denied.value.status_code == 403
+
+
+@pytest.mark.asyncio
+async def test_managed_file_usage_stays_keyed_by_the_row_model(
+    openai_files: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
+):
+    rows: Final = _batch_rows(("Test message for regression",))
+    await _install_managed_file_hook(openai_files, monkeypatch, rows)
+    _, _, batch_limiter = _make_limiters()
+
+    usage: Final = await batch_limiter.count_input_file_usage(
+        file_id=_MANAGED_FILE_ID,
+        custom_llm_provider="openai",
+        user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user-regression", team_id="team-owner"),
+        data={"input_file_id": _MANAGED_FILE_ID, "model": _MANAGED_DEPLOYMENT_ALIAS},
+    )
+
+    assert set(usage.per_model_usage) == {_BATCH_MODEL}
+
+
+@pytest.mark.asyncio
+async def test_batch_routed_by_a_model_missing_from_the_proxy_is_counted_under_the_row_model(
+    alias_routed_batch: respx.MockRouter,
+):
+    _serve_file(alias_routed_batch, "file-abc123", (_row_naming(_DEPLOYMENT_MODEL),))
+    _, _, batch_limiter = _make_limiters()
+
+    usage: Final = await batch_limiter.count_input_file_usage(
+        file_id="file-abc123",
+        user_api_key_dict=UserAPIKeyAuth(api_key=hash_token("sk-batch"), models=[_DEPLOYMENT_MODEL]),
+        data={"input_file_id": "file-abc123", "model": "not-on-this-proxy"},
+    )
+
+    assert set(usage.per_model_usage) == {_DEPLOYMENT_MODEL}

tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py

@@ -1036,10 +1036,10 @@ def _models_named_by_rows(received: ExtractedFileData) -> list[str]:
     [({"x-litellm-model": "gpt-5.4-nano-batch"}, {}), ({}, {"model": "gpt-5.4-nano-batch"})],
     ids=["x-litellm-model header", "model form field"],
 )
-def test_model_routed_batch_upload_names_the_deployment_model_on_every_row(
+def test_model_routed_batch_upload_names_the_deployment_model_on_rows_naming_the_alias(
     files_the_provider_received: list[ExtractedFileData], headers: dict[str, str], form: dict[str, str]
 ):
-    content = _batch_row("alias", "gpt-5.4-nano-batch") + b"\n" + _batch_row("other", "a-model-the-key-cannot-reach")
+    content = _batch_row("first", "gpt-5.4-nano-batch") + b"\n" + _batch_row("second", "gpt-5.4-nano-batch")
 
     response = _upload_batch_file(headers, form, content)
 
@@ -1047,6 +1047,17 @@ def test_model_routed_batch_upload_names_the_deployment_model_on_every_row(
     assert _models_named_by_rows(files_the_provider_received[0]) == ["gpt-5.4-nano", "gpt-5.4-nano"]
 
 
+def test_model_routed_batch_upload_leaves_a_row_naming_another_model_for_the_limiter_to_check(
+    files_the_provider_received: list[ExtractedFileData],
+):
+    content = _batch_row("alias", "gpt-5.4-nano-batch") + b"\n" + _batch_row("other", "another-model")
+
+    response = _upload_batch_file({"x-litellm-model": "gpt-5.4-nano-batch"}, {}, content)
+
+    assert response.status_code == 200, response.text
+    assert _models_named_by_rows(files_the_provider_received[0]) == ["gpt-5.4-nano", "another-model"]
+
+
 def test_model_routed_batch_upload_strips_the_provider_prefix_from_the_deployment_model(
     files_the_provider_received: list[ExtractedFileData],
 ):
@@ -1080,7 +1091,7 @@ def test_wildcard_routed_batch_upload_names_the_requested_model(
         ),
     )
 
-    response = _upload_batch_file({"x-litellm-model": requested_model}, {}, _batch_row("r", "a-row-naming-another-model"))
+    response = _upload_batch_file({"x-litellm-model": requested_model}, {}, _batch_row("r", requested_model))
 
     assert response.status_code == 200, response.text
     assert _models_named_by_rows(files_the_provider_received[0]) == ["gpt-x"]