diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 06800cb452..25680ab122 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -377,6 +377,23 @@ async def retrieve_batch( response = await proxy_logging_obj.post_call_success_hook( data=data, user_api_key_dict=user_api_key_dict, response=response ) + + # Resolve raw input_file_id to unified ID + if ( + unified_batch_id + and hasattr(response, "input_file_id") + and response.input_file_id + and not _is_base64_encoded_unified_file_id(response.input_file_id) + and prisma_client + ): + try: + _managed_file = await prisma_client.db.litellm_managedfiletable.find_first( + where={"flat_model_file_ids": {"has": response.input_file_id}} + ) + if _managed_file: + response.input_file_id = _managed_file.unified_file_id + except Exception: + pass asyncio.create_task( proxy_logging_obj.update_request_status( @@ -479,6 +496,23 @@ async def retrieve_batch( data=data, user_api_key_dict=user_api_key_dict, response=response ) + # Resolve raw input_file_id to unified ID + if ( + unified_batch_id + and hasattr(response, "input_file_id") + and response.input_file_id + and not _is_base64_encoded_unified_file_id(response.input_file_id) + and prisma_client + ): + try: + _managed_file = await prisma_client.db.litellm_managedfiletable.find_first( + where={"flat_model_file_ids": {"has": response.input_file_id}} + ) + if _managed_file: + response.input_file_id = _managed_file.unified_file_id + except Exception: + pass + ### ALERTING ### asyncio.create_task( proxy_logging_obj.update_request_status( diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index f67dc5e2aa..15c5bdfabb 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -687,6 +687,16 @@ async def get_batch_from_database( batch_data = json.loads(db_batch_object.file_object) if isinstance(db_batch_object.file_object, str) else db_batch_object.file_object response = LiteLLMBatch(**batch_data) response.id = batch_id + + if response.input_file_id and not _is_base64_encoded_unified_file_id(response.input_file_id): + try: + managed_file = await prisma_client.db.litellm_managedfiletable.find_first( + where={"flat_model_file_ids": {"has": response.input_file_id}} + ) + if managed_file: + response.input_file_id = managed_file.unified_file_id + except Exception: + pass verbose_proxy_logger.debug( f"Retrieved batch {batch_id} from ManagedObjectTable with status={response.status}" diff --git a/tests/test_litellm/enterprise/proxy/test_batch_retrieve_input_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_retrieve_input_file_id.py new file mode 100644 index 0000000000..6e9c3c0354 --- /dev/null +++ b/tests/test_litellm/enterprise/proxy/test_batch_retrieve_input_file_id.py @@ -0,0 +1,75 @@ +""" +Test that batch retrieve endpoint resolves raw input_file_id to the +unified managed file ID before returning. + +Bug: After batch completion, batches.retrieve returns the raw provider +input_file_id instead of the LiteLLM unified ID. +""" + +import base64 +import json + +import pytest +from unittest.mock import AsyncMock, MagicMock, patch + +from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, +) + + +DECODED_UNIFIED_INPUT_FILE_ID = "litellm_proxy:application/octet-stream;unified_id,test-uuid;target_model_names,azure-gpt-4" +B64_UNIFIED_INPUT_FILE_ID = base64.urlsafe_b64encode(DECODED_UNIFIED_INPUT_FILE_ID.encode()).decode().rstrip("=") +RAW_INPUT_FILE_ID = "file-raw-provider-abc123" + +DECODED_UNIFIED_BATCH_ID = "litellm_proxy;model_id:model-xyz;llm_batch_id:batch-123" +B64_UNIFIED_BATCH_ID = base64.urlsafe_b64encode(DECODED_UNIFIED_BATCH_ID.encode()).decode().rstrip("=") + + +@pytest.mark.asyncio +async def test_should_resolve_raw_input_file_id_to_unified(): + """ + When a completed batch has a raw input_file_id and the managed file table + contains a record for that raw ID, the retrieve endpoint should resolve + it to the unified file ID. + """ + unified_batch_id = _is_base64_encoded_unified_file_id(B64_UNIFIED_BATCH_ID) + assert unified_batch_id, "Test setup: batch_id should decode as unified" + + from litellm.types.utils import LiteLLMBatch + + batch_data = { + "id": B64_UNIFIED_BATCH_ID, + "completion_window": "24h", + "created_at": 1700000000, + "endpoint": "/v1/chat/completions", + "input_file_id": RAW_INPUT_FILE_ID, + "object": "batch", + "status": "completed", + "output_file_id": "file-output-xyz", + } + + mock_db_object = MagicMock() + mock_db_object.file_object = json.dumps(batch_data) + + mock_managed_file = MagicMock() + mock_managed_file.unified_file_id = B64_UNIFIED_INPUT_FILE_ID + + mock_prisma = MagicMock() + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=mock_db_object) + mock_prisma.db.litellm_managedfiletable.find_first = AsyncMock(return_value=mock_managed_file) + + from litellm.proxy.openai_files_endpoints.common_utils import get_batch_from_database + + _, response = await get_batch_from_database( + batch_id=B64_UNIFIED_BATCH_ID, + unified_batch_id=unified_batch_id, + managed_files_obj=MagicMock(), + prisma_client=mock_prisma, + verbose_proxy_logger=MagicMock(), + ) + + assert response is not None, "Batch should be found in DB" + assert response.input_file_id == B64_UNIFIED_INPUT_FILE_ID, ( + f"input_file_id should be unified '{B64_UNIFIED_INPUT_FILE_ID}', " + f"got raw '{response.input_file_id}'" + ) diff --git a/tests/test_litellm/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py b/tests/test_litellm/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py new file mode 100644 index 0000000000..420f5f9789 --- /dev/null +++ b/tests/test_litellm/enterprise/proxy/test_batch_retrieve_returns_unified_input_file_id.py @@ -0,0 +1,124 @@ +""" +Test that get_batch_from_database resolves raw input_file_id to the +unified/managed file ID when reading a batch from the database. + +Bug: The batch retrieve path stores the raw provider input_file_id in the +DB (via async_post_call_success_hook on the retrieve endpoint). When the +batch is later read from DB, get_batch_from_database returns the raw ID +without resolving it to the unified ID. +""" + +import json +import pytest +from typing import Optional +from unittest.mock import AsyncMock, MagicMock + +from litellm.proxy.openai_files_endpoints.common_utils import get_batch_from_database + + +def _mock_prisma(batch_json: str, managed_file_record=None): + """Create a mock prisma client with canned responses.""" + prisma = MagicMock() + + batch_db_record = MagicMock() + batch_db_record.file_object = batch_json + + prisma.db.litellm_managedobjecttable.find_first = AsyncMock( + return_value=batch_db_record + ) + + prisma.db.litellm_managedfiletable.find_first = AsyncMock( + return_value=managed_file_record + ) + + return prisma + + +@pytest.mark.asyncio +async def test_should_resolve_raw_input_file_id_to_unified_id(): + """ + When input_file_id in the stored batch is a raw provider ID, + get_batch_from_database must look up the unified ID from the + managed files table. + """ + unified_batch_id = "bGl0ZWxsbV9wcm94eTpiYXRjaF9pZA" + unified_input_file_id = "bGl0ZWxsbV9wcm94eTp1bmlmaWVkX2lucHV0" + raw_input_file_id = "file-abc123-raw" + + batch_data = { + "id": "batch-raw-123", + "completion_window": "24h", + "created_at": 1700000000, + "endpoint": "/v1/chat/completions", + "input_file_id": raw_input_file_id, + "object": "batch", + "status": "completed", + "output_file_id": "file-output-raw", + } + + managed_file_record = MagicMock() + managed_file_record.unified_file_id = unified_input_file_id + + prisma = _mock_prisma( + batch_json=json.dumps(batch_data), + managed_file_record=managed_file_record, + ) + + _, response = await get_batch_from_database( + batch_id=unified_batch_id, + unified_batch_id="decoded_unified_batch_id", + managed_files_obj=MagicMock(), + prisma_client=prisma, + verbose_proxy_logger=MagicMock(), + ) + + assert response is not None + assert response.input_file_id == unified_input_file_id, ( + f"input_file_id should be resolved to '{unified_input_file_id}', " + f"got raw: '{response.input_file_id}'" + ) + + prisma.db.litellm_managedfiletable.find_first.assert_called_once_with( + where={"flat_model_file_ids": {"has": raw_input_file_id}} + ) + + +@pytest.mark.asyncio +async def test_should_preserve_already_managed_input_file_id(): + """ + When input_file_id is already a managed/unified ID, it should + not be modified. + """ + import base64 + + unified_batch_id = "bGl0ZWxsbV9wcm94eTpiYXRjaF9pZA" + decoded_unified = "litellm_proxy:application/octet-stream;unified_id,test-123" + base64_input_file_id = base64.urlsafe_b64encode(decoded_unified.encode()).decode().rstrip("=") + + batch_data = { + "id": "batch-raw-123", + "completion_window": "24h", + "created_at": 1700000000, + "endpoint": "/v1/chat/completions", + "input_file_id": base64_input_file_id, + "object": "batch", + "status": "completed", + } + + prisma = _mock_prisma(batch_json=json.dumps(batch_data)) + + _, response = await get_batch_from_database( + batch_id=unified_batch_id, + unified_batch_id="decoded_unified_batch_id", + managed_files_obj=MagicMock(), + prisma_client=prisma, + verbose_proxy_logger=MagicMock(), + ) + + assert response is not None + assert response.input_file_id == base64_input_file_id, ( + f"input_file_id was already managed, should be preserved as '{base64_input_file_id}', " + f"got: '{response.input_file_id}'" + ) + + prisma.db.litellm_managedfiletable.find_first.assert_not_called()