fix: batch retrieve returns unified input_file_id

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(

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}"

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}'"
+    )

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()