From 5ad0d0367181af342e281e11b8b6c3afc9055e08 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 10:26:03 -0300 Subject: [PATCH 1/8] fix(proxy): encode batch IDs with model info when x-litellm-model header is used When create_batch routes via x-litellm-model header, the response batch_id was returned raw without model routing info. This meant retrieve_batch could not determine which provider/credentials to use, defaulting to "openai" instead of the correct provider (e.g., VLLM). Now encodes batch_id, output_file_id, and error_file_id with model info (same pattern as the model-embedded file_id flow in Scenario 1), so retrieve_batch can decode and route back to the correct provider. --- litellm/proxy/batches_endpoints/endpoints.py | 21 +- .../test_batch_x_litellm_model_encoding.py | 362 ++++++++++++++++++ 2 files changed, 382 insertions(+), 1 deletion(-) create mode 100644 tests/litellm/proxy/test_batch_x_litellm_model_encoding.py diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 1c9ba6cb248..7352e6e2085 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -242,7 +242,26 @@ async def create_batch( # noqa: PLR0915 custom_llm_provider=credentials["custom_llm_provider"], **_create_batch_data # type: ignore ) - + + # Encode response IDs with model info so retrieve_batch + # can route back to the correct provider/credentials. + if response and hasattr(response, "id") and response.id: + response.id = encode_file_id_with_model( + file_id=response.id, + model=model_param, + id_type="batch", + ) + + if hasattr(response, "output_file_id") and response.output_file_id: + response.output_file_id = encode_file_id_with_model( + file_id=response.output_file_id, model=model_param + ) + + if hasattr(response, "error_file_id") and response.error_file_id: + response.error_file_id = encode_file_id_with_model( + file_id=response.error_file_id, model=model_param + ) + verbose_proxy_logger.debug(f"Created batch using model: {model_param}") else: # SCENARIO 3: Fallback to custom_llm_provider (uses env variables) diff --git a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py new file mode 100644 index 00000000000..01d6a8aaacd --- /dev/null +++ b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -0,0 +1,362 @@ +""" +Unit tests for batch ID encoding when x-litellm-model header is used. + +Verifies that create_batch encodes response IDs with model info so that +retrieve_batch can route back to the correct provider/credentials. +""" + +import json +from unittest.mock import AsyncMock, MagicMock, patch + +import pytest + +import litellm +from litellm.proxy.openai_files_endpoints.common_utils import ( + decode_model_from_file_id, + get_original_file_id, +) +from litellm.types.utils import LiteLLMBatch + + +def _make_mock_request(headers: dict) -> MagicMock: + """Create a mock FastAPI Request with the given headers.""" + mock_request = MagicMock() + mock_request.headers = headers + mock_request.query_params = {} + mock_request.url = MagicMock() + mock_request.url.port = 4000 + mock_request.method = "POST" + mock_request.url.path = "/v1/batches" + return mock_request + + +def _make_batch_response( + batch_id: str = "batch_abc123", + input_file_id: str = "file-input456", + output_file_id: str = None, + error_file_id: str = None, + status: str = "validating", +) -> LiteLLMBatch: + """Create a mock LiteLLMBatch response from a provider.""" + return LiteLLMBatch( + id=batch_id, + object="batch", + status=status, + endpoint="/v1/chat/completions", + input_file_id=input_file_id, + completion_window="24h", + created_at=1234567890, + output_file_id=output_file_id, + error_file_id=error_file_id, + ) + + +@pytest.mark.asyncio +async def test_create_batch_with_x_litellm_model_encodes_batch_id(): + """ + When x-litellm-model header is provided, create_batch should encode the + response batch_id with model info so retrieve_batch can route correctly. + """ + from litellm.proxy.batches_endpoints.endpoints import create_batch + + model_name = "my-vllm-model" + raw_batch_id = "batch_abc123" + + mock_response = _make_batch_response(batch_id=raw_batch_id) + mock_request = _make_mock_request(headers={"x-litellm-model": model_name}) + mock_fastapi_response = MagicMock() + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.parent_otel_span = None + mock_user_api_key_dict.user_id = "test_user" + + mock_credentials = { + "api_key": "sk-test", + "api_base": "http://vllm:8000", + "custom_llm_provider": "openai", + } + + with ( + patch( + "litellm.proxy.batches_endpoints.endpoints._read_request_body", + new=AsyncMock( + return_value={"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"} + ), + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processor_cls, + patch( + "litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model", + return_value=mock_credentials, + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials", + ), + patch( + "litellm.acreate_batch", + new=AsyncMock(return_value=mock_response), + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.is_known_model", + return_value=False, + ), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_config", MagicMock()), + patch("litellm.proxy.proxy_server.version", "1.0.0"), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock( + post_call_success_hook=AsyncMock(return_value=mock_response), + update_request_status=AsyncMock(), + ), + ), + ): + # Setup the mock processor to return data and logging obj + mock_processor = MagicMock() + mock_processor.common_processing_pre_call_logic = AsyncMock( + return_value=( + {"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + MagicMock(), + ) + ) + mock_processor_cls.return_value = mock_processor + + response = await create_batch( + request=mock_request, + fastapi_response=mock_fastapi_response, + provider=None, + user_api_key_dict=mock_user_api_key_dict, + ) + + # The batch_id should be encoded with model info + assert response.id != raw_batch_id, ( + f"Expected batch_id to be encoded, but got raw ID: {response.id}" + ) + assert response.id.startswith("batch_"), ( + f"Encoded batch_id should keep batch_ prefix, got: {response.id}" + ) + + # Should be decodable back to the original + decoded_model = decode_model_from_file_id(response.id) + assert decoded_model == model_name, ( + f"Expected model '{model_name}' from decoded batch_id, got: {decoded_model}" + ) + + original_id = get_original_file_id(response.id) + assert original_id == raw_batch_id, ( + f"Expected original ID '{raw_batch_id}', got: {original_id}" + ) + + +@pytest.mark.asyncio +async def test_create_batch_with_x_litellm_model_encodes_output_and_error_file_ids(): + """ + When a completed batch is returned with output_file_id and error_file_id, + these should also be encoded with model info. + """ + from litellm.proxy.batches_endpoints.endpoints import create_batch + + model_name = "my-vllm-model" + raw_output_file = "file-output789" + raw_error_file = "file-error012" + + mock_response = _make_batch_response( + batch_id="batch_abc123", + output_file_id=raw_output_file, + error_file_id=raw_error_file, + status="completed", + ) + mock_request = _make_mock_request(headers={"x-litellm-model": model_name}) + mock_fastapi_response = MagicMock() + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.parent_otel_span = None + mock_user_api_key_dict.user_id = "test_user" + + mock_credentials = { + "api_key": "sk-test", + "api_base": "http://vllm:8000", + "custom_llm_provider": "openai", + } + + with ( + patch( + "litellm.proxy.batches_endpoints.endpoints._read_request_body", + new=AsyncMock( + return_value={"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"} + ), + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processor_cls, + patch( + "litellm.proxy.batches_endpoints.endpoints.get_credentials_for_model", + return_value=mock_credentials, + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.prepare_data_with_credentials", + ), + patch( + "litellm.acreate_batch", + new=AsyncMock(return_value=mock_response), + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.is_known_model", + return_value=False, + ), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.llm_router", MagicMock()), + patch("litellm.proxy.proxy_server.proxy_config", MagicMock()), + patch("litellm.proxy.proxy_server.version", "1.0.0"), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock( + post_call_success_hook=AsyncMock(return_value=mock_response), + update_request_status=AsyncMock(), + ), + ), + ): + mock_processor = MagicMock() + mock_processor.common_processing_pre_call_logic = AsyncMock( + return_value=( + {"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + MagicMock(), + ) + ) + mock_processor_cls.return_value = mock_processor + + response = await create_batch( + request=mock_request, + fastapi_response=mock_fastapi_response, + provider=None, + user_api_key_dict=mock_user_api_key_dict, + ) + + # output_file_id should be encoded + assert decode_model_from_file_id(response.output_file_id) == model_name + assert get_original_file_id(response.output_file_id) == raw_output_file + + # error_file_id should be encoded + assert decode_model_from_file_id(response.error_file_id) == model_name + assert get_original_file_id(response.error_file_id) == raw_error_file + + +@pytest.mark.asyncio +async def test_create_batch_without_x_litellm_model_returns_raw_ids(): + """ + Without x-litellm-model header, create_batch should NOT encode batch IDs + (falls through to Scenario 3 / custom_llm_provider fallback). + """ + from litellm.proxy.batches_endpoints.endpoints import create_batch + + raw_batch_id = "batch_abc123" + mock_response = _make_batch_response(batch_id=raw_batch_id) + mock_request = _make_mock_request(headers={}) + mock_fastapi_response = MagicMock() + mock_user_api_key_dict = MagicMock() + mock_user_api_key_dict.parent_otel_span = None + mock_user_api_key_dict.user_id = "test_user" + + with ( + patch( + "litellm.proxy.batches_endpoints.endpoints._read_request_body", + new=AsyncMock( + return_value={"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"} + ), + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.ProxyBaseLLMRequestProcessing" + ) as mock_processor_cls, + patch( + "litellm.acreate_batch", + new=AsyncMock(return_value=mock_response), + ), + patch( + "litellm.proxy.batches_endpoints.endpoints.is_known_model", + return_value=False, + ), + patch("litellm.proxy.proxy_server.general_settings", {}), + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.proxy_config", MagicMock()), + patch("litellm.proxy.proxy_server.version", "1.0.0"), + patch( + "litellm.proxy.proxy_server.proxy_logging_obj", + MagicMock( + post_call_success_hook=AsyncMock(return_value=mock_response), + update_request_status=AsyncMock(), + ), + ), + ): + mock_processor = MagicMock() + mock_processor.common_processing_pre_call_logic = AsyncMock( + return_value=( + {"input_file_id": "file-input456", "endpoint": "/v1/chat/completions", "completion_window": "24h"}, + MagicMock(), + ) + ) + mock_processor_cls.return_value = mock_processor + + response = await create_batch( + request=mock_request, + fastapi_response=mock_fastapi_response, + provider=None, + user_api_key_dict=mock_user_api_key_dict, + ) + + # Without x-litellm-model, the batch_id should remain raw + assert response.id == raw_batch_id + assert decode_model_from_file_id(response.id) is None + + +class TestBatchIdRoundTripWithRetrieve: + """ + Tests that batch IDs encoded during create_batch can be decoded + correctly during retrieve_batch (Scenario 1: model_from_id). + """ + + def test_encoded_batch_id_is_decoded_for_retrieve(self): + """ + Simulates the full round-trip: create encodes the ID, + retrieve decodes it to get the model and original batch_id. + """ + from litellm.proxy.openai_files_endpoints.common_utils import ( + encode_file_id_with_model, + ) + + model_name = "my-vllm-model" + raw_batch_id = "batch_vllm_12345" + + # What create_batch does: + encoded_id = encode_file_id_with_model( + file_id=raw_batch_id, model=model_name, id_type="batch" + ) + + # What retrieve_batch does: + decoded_model = decode_model_from_file_id(encoded_id) + original_id = get_original_file_id(encoded_id) + + assert decoded_model == model_name + assert original_id == raw_batch_id + + def test_vllm_style_batch_id_roundtrip(self): + """ + VLLM may return batch IDs in various formats. + Verify round-trip works for common patterns. + """ + from litellm.proxy.openai_files_endpoints.common_utils import ( + encode_file_id_with_model, + ) + + test_cases = [ + ("batch_abc123", "vllm-llama3"), + ("batch_67890", "openai/llama-3-8b"), + ("batch_some-uuid-here", "my-custom-vllm"), + ] + + for raw_id, model in test_cases: + encoded = encode_file_id_with_model( + file_id=raw_id, model=model, id_type="batch" + ) + assert encoded.startswith("batch_") + assert decode_model_from_file_id(encoded) == model + assert get_original_file_id(encoded) == raw_id From 3426b905cedd66d59506174efb04d8c53566b53b Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Tue, 3 Mar 2026 10:44:47 -0300 Subject: [PATCH 2/8] Update tests/litellm/proxy/test_batch_x_litellm_model_encoding.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/litellm/proxy/test_batch_x_litellm_model_encoding.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py index 01d6a8aaacd..062443d9425 100644 --- a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -33,8 +33,8 @@ def _make_mock_request(headers: dict) -> MagicMock: def _make_batch_response( batch_id: str = "batch_abc123", input_file_id: str = "file-input456", - output_file_id: str = None, - error_file_id: str = None, + output_file_id: Optional[str] = None, + error_file_id: Optional[str] = None, status: str = "validating", ) -> LiteLLMBatch: """Create a mock LiteLLMBatch response from a provider.""" From 7d664f0c096a7f85befc17b22d80c8cbf9097a93 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Tue, 3 Mar 2026 10:45:00 -0300 Subject: [PATCH 3/8] Update tests/litellm/proxy/test_batch_x_litellm_model_encoding.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- tests/litellm/proxy/test_batch_x_litellm_model_encoding.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py index 062443d9425..5150b57568b 100644 --- a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -5,7 +5,6 @@ Verifies that create_batch encodes response IDs with model info so that retrieve_batch can route back to the correct provider/credentials. """ -import json from unittest.mock import AsyncMock, MagicMock, patch import pytest From 096edface5b7772eccf398113a57b26941c00457 Mon Sep 17 00:00:00 2001 From: Cesar Garcia <128240629+Chesars@users.noreply.github.com> Date: Tue, 3 Mar 2026 10:46:14 -0300 Subject: [PATCH 4/8] Update litellm/proxy/batches_endpoints/endpoints.py Co-authored-by: greptile-apps[bot] <165735046+greptile-apps[bot]@users.noreply.github.com> --- litellm/proxy/batches_endpoints/endpoints.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 7352e6e2085..58c8e2d4d0c 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -257,11 +257,16 @@ async def create_batch( # noqa: PLR0915 file_id=response.output_file_id, model=model_param ) - if hasattr(response, "error_file_id") and response.error_file_id: + if hasattr(response, "error_file_id") and response.error_file_id: response.error_file_id = encode_file_id_with_model( file_id=response.error_file_id, model=model_param ) + if hasattr(response, "input_file_id") and response.input_file_id: + response.input_file_id = encode_file_id_with_model( + file_id=response.input_file_id, model=model_param + ) + verbose_proxy_logger.debug(f"Created batch using model: {model_param}") else: # SCENARIO 3: Fallback to custom_llm_provider (uses env variables) From 9463de0c6629354dbc616b1a1beb394f27a45b80 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 10:51:48 -0300 Subject: [PATCH 5/8] fix: correct indentation from commit suggestions and add missing Optional import --- litellm/proxy/batches_endpoints/endpoints.py | 2 +- tests/litellm/proxy/test_batch_x_litellm_model_encoding.py | 1 + 2 files changed, 2 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 58c8e2d4d0c..cdee69f2b30 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -257,7 +257,7 @@ async def create_batch( # noqa: PLR0915 file_id=response.output_file_id, model=model_param ) - if hasattr(response, "error_file_id") and response.error_file_id: + if hasattr(response, "error_file_id") and response.error_file_id: response.error_file_id = encode_file_id_with_model( file_id=response.error_file_id, model=model_param ) diff --git a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py index 5150b57568b..521a3632dcb 100644 --- a/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py +++ b/tests/litellm/proxy/test_batch_x_litellm_model_encoding.py @@ -5,6 +5,7 @@ Verifies that create_batch encodes response IDs with model info so that retrieve_batch can route back to the correct provider/credentials. """ +from typing import Optional from unittest.mock import AsyncMock, MagicMock, patch import pytest From 9fd4c00b064e6358fdb5202daa417b3cfb482796 Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 11:06:38 -0300 Subject: [PATCH 6/8] fix(proxy): re-encode response IDs in retrieve_batch for model-based routing The provider returns raw IDs in the retrieve response (output_file_id, error_file_id). These need to be encoded with model info so the client can use them for subsequent file download calls through the proxy. --- litellm/proxy/batches_endpoints/endpoints.py | 26 ++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index cdee69f2b30..d3e983c5ad5 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -464,8 +464,30 @@ async def retrieve_batch( # noqa: PLR0915 custom_llm_provider=credentials["custom_llm_provider"], **data # type: ignore ) - - + + # Re-encode response IDs so the client always sees encoded IDs. + # The provider returns raw IDs (e.g. output_file_id, error_file_id) + # which the client needs encoded to route future file downloads. + if response and hasattr(response, "id") and response.id: + response.id = encode_file_id_with_model( + file_id=response.id, model=model_from_id, id_type="batch", + ) + + if hasattr(response, "output_file_id") and response.output_file_id: + response.output_file_id = encode_file_id_with_model( + file_id=response.output_file_id, model=model_from_id + ) + + if hasattr(response, "error_file_id") and response.error_file_id: + response.error_file_id = encode_file_id_with_model( + file_id=response.error_file_id, model=model_from_id + ) + + if hasattr(response, "input_file_id") and response.input_file_id: + response.input_file_id = encode_file_id_with_model( + file_id=response.input_file_id, model=model_from_id + ) + verbose_proxy_logger.debug( f"Retrieved batch using model: {model_from_id}, original_id: {original_batch_id}" ) From 7506fd0426a70c0387fd811d44fdd06c1fc0006b Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 11:22:43 -0300 Subject: [PATCH 7/8] fix(proxy): re-encode response IDs in cancel_batch for model-based routing --- litellm/proxy/batches_endpoints/endpoints.py | 23 +++++++++++++++++++- 1 file changed, 22 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index d3e983c5ad5..7ab50e321ac 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -855,7 +855,28 @@ async def cancel_batch( custom_llm_provider=credentials["custom_llm_provider"], **data # type: ignore ) - + + # Re-encode response IDs so the client always sees encoded IDs. + if response and hasattr(response, "id") and response.id: + response.id = encode_file_id_with_model( + file_id=response.id, model=model_from_id, id_type="batch", + ) + + if hasattr(response, "output_file_id") and response.output_file_id: + response.output_file_id = encode_file_id_with_model( + file_id=response.output_file_id, model=model_from_id + ) + + if hasattr(response, "error_file_id") and response.error_file_id: + response.error_file_id = encode_file_id_with_model( + file_id=response.error_file_id, model=model_from_id + ) + + if hasattr(response, "input_file_id") and response.input_file_id: + response.input_file_id = encode_file_id_with_model( + file_id=response.input_file_id, model=model_from_id + ) + verbose_proxy_logger.debug( f"Cancelled batch using model: {model_from_id}, original_id: {original_batch_id}" ) From 59bde4a81a92c9e7b8f48ca2e32311a735a51a7d Mon Sep 17 00:00:00 2001 From: Chesars Date: Tue, 3 Mar 2026 11:38:50 -0300 Subject: [PATCH 8/8] refactor(proxy): extract encode_batch_response_ids helper and fix list_batches encoding Extract duplicated batch ID encoding logic into a shared helper encode_batch_response_ids() in common_utils.py. Use it in create_batch, retrieve_batch, and cancel_batch. Also add encoding to list_batches when x-litellm-model is used. --- litellm/proxy/batches_endpoints/endpoints.py | 77 +++---------------- .../openai_files_endpoints/common_utils.py | 16 ++++ 2 files changed, 27 insertions(+), 66 deletions(-) diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 7ab50e321ac..ae99d59c631 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -23,6 +23,7 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, decode_model_from_file_id, + encode_batch_response_ids, encode_file_id_with_model, get_batch_from_database, get_credentials_for_model, @@ -243,29 +244,7 @@ async def create_batch( # noqa: PLR0915 **_create_batch_data # type: ignore ) - # Encode response IDs with model info so retrieve_batch - # can route back to the correct provider/credentials. - if response and hasattr(response, "id") and response.id: - response.id = encode_file_id_with_model( - file_id=response.id, - model=model_param, - id_type="batch", - ) - - if hasattr(response, "output_file_id") and response.output_file_id: - response.output_file_id = encode_file_id_with_model( - file_id=response.output_file_id, model=model_param - ) - - if hasattr(response, "error_file_id") and response.error_file_id: - response.error_file_id = encode_file_id_with_model( - file_id=response.error_file_id, model=model_param - ) - - if hasattr(response, "input_file_id") and response.input_file_id: - response.input_file_id = encode_file_id_with_model( - file_id=response.input_file_id, model=model_param - ) + encode_batch_response_ids(response, model=model_param) verbose_proxy_logger.debug(f"Created batch using model: {model_param}") else: @@ -465,28 +444,7 @@ async def retrieve_batch( # noqa: PLR0915 **data # type: ignore ) - # Re-encode response IDs so the client always sees encoded IDs. - # The provider returns raw IDs (e.g. output_file_id, error_file_id) - # which the client needs encoded to route future file downloads. - if response and hasattr(response, "id") and response.id: - response.id = encode_file_id_with_model( - file_id=response.id, model=model_from_id, id_type="batch", - ) - - if hasattr(response, "output_file_id") and response.output_file_id: - response.output_file_id = encode_file_id_with_model( - file_id=response.output_file_id, model=model_from_id - ) - - if hasattr(response, "error_file_id") and response.error_file_id: - response.error_file_id = encode_file_id_with_model( - file_id=response.error_file_id, model=model_from_id - ) - - if hasattr(response, "input_file_id") and response.input_file_id: - response.input_file_id = encode_file_id_with_model( - file_id=response.input_file_id, model=model_from_id - ) + encode_batch_response_ids(response, model=model_from_id) verbose_proxy_logger.debug( f"Retrieved batch using model: {model_from_id}, original_id: {original_batch_id}" @@ -679,7 +637,13 @@ async def list_batches( limit=limit, **data # type: ignore ) - + + # Encode batch IDs in the list response so clients can use + # them for retrieve/cancel/file downloads through the proxy. + if response and hasattr(response, "data") and response.data: + for batch in response.data: + encode_batch_response_ids(batch, model=model_param) + verbose_proxy_logger.debug(f"Listed batches using model: {model_param}") # SCENARIO 2 (alternative): target_model_names based routing @@ -856,26 +820,7 @@ async def cancel_batch( **data # type: ignore ) - # Re-encode response IDs so the client always sees encoded IDs. - if response and hasattr(response, "id") and response.id: - response.id = encode_file_id_with_model( - file_id=response.id, model=model_from_id, id_type="batch", - ) - - if hasattr(response, "output_file_id") and response.output_file_id: - response.output_file_id = encode_file_id_with_model( - file_id=response.output_file_id, model=model_from_id - ) - - if hasattr(response, "error_file_id") and response.error_file_id: - response.error_file_id = encode_file_id_with_model( - file_id=response.error_file_id, model=model_from_id - ) - - if hasattr(response, "input_file_id") and response.input_file_id: - response.input_file_id = encode_file_id_with_model( - file_id=response.input_file_id, model=model_from_id - ) + encode_batch_response_ids(response, model=model_from_id) verbose_proxy_logger.debug( f"Cancelled batch using model: {model_from_id}, original_id: {original_batch_id}" diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index ceaf3c7550e..343ea119672 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -129,6 +129,22 @@ def encode_file_id_with_model( return f"{prefix}{encoded_b64}" +def encode_batch_response_ids(response, model: str) -> None: + """Encode all IDs in a batch response with model routing info (in-place).""" + if not response or not hasattr(response, "id") or not response.id: + return + response.id = encode_file_id_with_model( + file_id=response.id, model=model, id_type="batch" + ) + for attr in ("output_file_id", "error_file_id", "input_file_id"): + if hasattr(response, attr) and getattr(response, attr): + setattr( + response, + attr, + encode_file_id_with_model(file_id=getattr(response, attr), model=model), + ) + + def decode_model_from_file_id(encoded_id: str) -> Optional[str]: """ Extract model name from an encoded file/batch ID.