From b2b3713f155a68749a41741fce4db7b69f8d6f0a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Thu, 29 Oct 2026 09:48:43 -0700 Subject: [PATCH] fix(batches): read batch input files for rate limiting from the batch deployment --- litellm/proxy/hooks/batch_rate_limiter.py | 7 +- litellm/router.py | 5 -- .../proxy/hooks/test_batch_file_validation.py | 78 +++++++++++++++++++ 3 files changed, 83 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 94762232e0b..98a1b1414bd 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -542,11 +542,13 @@ class _PROXY_BatchRateLimiter(CustomLogger): and the upstream provider returns 404. """ from litellm.proxy.openai_files_endpoints.common_utils import ( + batch_deployment_id, decode_model_from_file_id, explicit_request_credentials, get_credentials_for_model, get_original_file_id, get_team_provider_credentials, + hinted_batch_deployment_id, ) from litellm.proxy.proxy_server import llm_router @@ -560,7 +562,8 @@ class _PROXY_BatchRateLimiter(CustomLogger): try: credentials = get_credentials_for_model( llm_router=llm_router, - model_id=model_from_file_id, + model_id=hinted_batch_deployment_id(llm_router, model_from_file_id, file_id) + or model_from_file_id, operation_context="batch input file read (rate limiting)", ) fetch_kwargs.update(extract_file_access_credentials(credentials)) @@ -578,7 +581,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): try: credentials = get_credentials_for_model( llm_router=llm_router, - model_id=request_model, + model_id=batch_deployment_id(llm_router, request_model) or request_model, operation_context="batch input file read (rate limiting)", ) fetch_kwargs.update(extract_file_access_credentials(credentials)) diff --git a/litellm/router.py b/litellm/router.py index 8eb31775efb..4b0f7d4a70e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -12030,11 +12030,6 @@ class Router: return Deployment(**first_usable) if isinstance(first_usable, dict) else first_usable def get_batch_deployment_id(self, model_group_name: str, preferred_deployment_id: str | None = None) -> str | None: - """ - The shared deployment of ``model_group_name`` that file and batch calls - should reach when the group mixes batch and non-batch deployments, or - None when the batch rule drops nothing and the group name resolves as before. - """ group: Final = tuple( deployment for deployment in self.deployments_named(model_group_name) or () diff --git a/tests/unit/proxy/hooks/test_batch_file_validation.py b/tests/unit/proxy/hooks/test_batch_file_validation.py index 8eeee2e6468..42d90d57619 100644 --- a/tests/unit/proxy/hooks/test_batch_file_validation.py +++ b/tests/unit/proxy/hooks/test_batch_file_validation.py @@ -1353,6 +1353,84 @@ def test_resolve_fetch_params_model_embedded_fails_open_on_credential_error(): assert fetch_kwargs == {"custom_llm_provider": "openai"} + +def _mixed_azure_chat_and_batch_router(): + from litellm import Router + + return Router( + model_list=[ + { + "model_name": "gpt-mixed", + "litellm_params": { + "model": "azure/gpt-chat", + "api_base": "https://chat-resource.openai.azure.com", + "api_key": "chat-key", + "api_version": "2025-04-01-preview", + }, + "model_info": {"id": "chat-id", "mode": "chat"}, + }, + { + "model_name": "gpt-mixed", + "litellm_params": { + "model": "azure/gpt-batch", + "api_base": "https://batch-resource.openai.azure.com", + "api_key": "batch-key", + "api_version": "2025-04-01-preview", + }, + "model_info": {"id": "batch-id", "mode": "batch"}, + }, + ] + ) + + +@pytest.mark.parametrize( + ("file_id", "data", "expected_api_base"), + [ + pytest.param( + "file-hinted", + {}, + "https://batch-resource.openai.azure.com", + id="hinted-file-id", + ), + pytest.param( + "file-legacy", + {}, + "https://chat-resource.openai.azure.com", + id="group-only-file-id", + ), + pytest.param( + "file-raw-provider-id", + {"model": "gpt-mixed"}, + "https://batch-resource.openai.azure.com", + id="request-model", + ), + ], +) +def test_resolve_fetch_params_reads_a_mixed_group_file_where_the_batch_runs( + file_id, data, expected_api_base +): + from litellm.proxy.openai_files_endpoints.common_utils import ( + encode_file_id_with_model, + ) + + proxy_file_id = { + "file-hinted": encode_file_id_with_model( + "file-orig", "gpt-mixed", deployment_id="batch-id" + ), + "file-legacy": encode_file_id_with_model("file-orig", "gpt-mixed"), + }.get(file_id, file_id) + rate_limiter = _make_rate_limiter() + with patch( + "litellm.proxy.proxy_server.llm_router", _mixed_azure_chat_and_batch_router() + ): + _, fetch_kwargs = rate_limiter._resolve_batch_input_file_fetch_params( + file_id=proxy_file_id, + custom_llm_provider="openai", + data=data, + ) + assert fetch_kwargs["api_base"] == expected_api_base + assert fetch_kwargs["model"] == "gpt-mixed" + @pytest.mark.asyncio async def test_check_and_increment_computes_descriptors_when_not_passed(): from litellm.proxy.hooks.batch_rate_limiter import (