From 0db47c24b811837dce2ef10e9277961d3d374fc4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 30 Oct 2026 11:00:49 -0700 Subject: [PATCH] fix(proxy): rewrite only alias rows and authorize them against the alias at batch create --- litellm/litellm_core_utils/batch_utils.py | 13 +- litellm/proxy/hooks/batch_rate_limiter.py | 37 ++++- .../openai_files_endpoints/common_utils.py | 5 + .../openai_files_endpoints/files_endpoints.py | 13 +- .../proxy/hooks/test_batch_rate_limiter.py | 150 +++++++++++++++++- .../test_files_endpoint.py | 17 +- 6 files changed, 223 insertions(+), 12 deletions(-) diff --git a/litellm/litellm_core_utils/batch_utils.py b/litellm/litellm_core_utils/batch_utils.py index 5b94d4b4e5b..d7a5154b0f7 100644 --- a/litellm/litellm_core_utils/batch_utils.py +++ b/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( ( diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 747dda100bf..d0906579e00 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/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) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 907c2ca303e..aa4b96ab663 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/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]: diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 7736cbb84af..532bf440b22 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/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"], ), diff --git a/tests/unit/proxy/hooks/test_batch_rate_limiter.py b/tests/unit/proxy/hooks/test_batch_rate_limiter.py index b10cd9e7107..1246d55c9d5 100644 --- a/tests/unit/proxy/hooks/test_batch_rate_limiter.py +++ b/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} diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py index b7a3df8025b..0bd833e194f 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/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"]