Managed Files + Batches - filter deployments to only those where file was written + save all model file id mappings in DB (prev just 1st one) (#12048)

* test(test_router.py): initial unit test confirming router.afile_content uses dynamic api key / api base

* fix(managed_files.py): filter deployments for only those within file id mapping

ensure call works - only route to models where the file was written

* fix(proxy_server.py): fix loading in model ids from config, if config id is int

* fix(router.py): return all model file id mappings on create_file

if multiple deployments - this ensures all the file id mappings are bubbled up

Fixes issue when trying to use loadbalanced deployments - only 1 file id mapping was being stored

enterprise/enterprise_hooks/managed_files.py

@@ -88,6 +88,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
             litellm_parent_otel_span=litellm_parent_otel_span,
         )
 
+        ## STORE MODEL MAPPINGS IN DB
+
         await self.prisma_client.db.litellm_managedfiletable.create(
             data={
                 "unified_file_id": file_id,
@@ -367,6 +369,36 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
 
         return data
 
+    async def async_filter_deployments(
+        self,
+        model: str,
+        healthy_deployments: List,
+        messages: Optional[List[AllMessageValues]],
+        request_kwargs: Optional[Dict] = None,
+        parent_otel_span: Optional[Span] = None,
+    ) -> List[Dict]:
+        if request_kwargs is None:
+            return healthy_deployments
+
+        input_file_id = cast(Optional[str], request_kwargs.get("input_file_id"))
+        model_file_id_mapping = cast(
+            Optional[Dict[str, Dict[str, str]]],
+            request_kwargs.get("model_file_id_mapping"),
+        )
+        allowed_model_ids = []
+        if input_file_id and model_file_id_mapping:
+            model_id_dict = model_file_id_mapping.get(input_file_id, {})
+            allowed_model_ids = list(model_id_dict.keys())
+
+        if len(allowed_model_ids) == 0:
+            return healthy_deployments
+
+        return [
+            deployment
+            for deployment in healthy_deployments
+            if deployment.get("model_info", {}).get("id") in allowed_model_ids
+        ]
+
     async def async_pre_call_deployment_hook(
         self, kwargs: Dict[str, Any], call_type: Optional[CallTypes]
     ) -> Optional[dict]:
@@ -500,15 +532,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
 
         ## STORE MODEL MAPPINGS IN DB
         model_mappings: Dict[str, str] = {}
+
         for file_object in responses:
-            model_id = file_object._hidden_params.get("model_id")
-            if model_id is None:
-                verbose_logger.warning(
-                    f"Skipping file_object: {file_object} because model_id in hidden_params={file_object._hidden_params} is None"
-                )
-                continue
-            file_id = file_object.id
-            model_mappings[model_id] = file_id
+            model_file_id_mapping = file_object._hidden_params.get(
+                "model_file_id_mapping"
+            )
+            if model_file_id_mapping and isinstance(model_file_id_mapping, dict):
+                model_mappings.update(model_file_id_mapping)
 
         await self.store_unified_file_id(
             file_id=response.id,

litellm/proxy/_new_secret_config.yaml

@@ -2,3 +2,20 @@ model_list:
   - model_name: gemini-2.5-pro
     litellm_params:
       model: gemini/gemini-2.5-pro
+  - model_name: azure-batches
+    litellm_params:
+      model: azure/gpt-4o
+      api_key: os.environ/AZURE_API_KEY_HIDDEN
+      api_base: os.environ/AZURE_API_BASE_HIDDEN
+  - model_name: openai-gpt-4o
+    litellm_params:
+      model: gpt-4o
+      api_key: os.environ/OPENAI_API_KEY_TEST
+    model_info:
+      id: 12345678
+  - model_name: openai-gpt-4o
+    litellm_params:
+      model: gpt-4o
+      api_key: os.environ/OPENAI_API_KEY_TEST_2
+    model_info:
+      id: 12345679

litellm/proxy/proxy_server.py

@@ -2279,6 +2279,8 @@ class ProxyConfig:
                         model_group=model["model_name"],
                         litellm_params=model["litellm_params"],
                     )
+                else:
+                    model_id = str(model_id)
                 combined_id_list.append(model_id)  # ADD CONFIG MODEL TO COMBINED LIST
 
         router_model_ids = llm_router.get_model_ids()

litellm/router.py

@@ -743,6 +743,7 @@ class Router:
         self.afile_delete = self.factory_function(
             litellm.afile_delete, call_type="afile_delete"
         )
+
         self.afile_content = self.factory_function(
             litellm.afile_content, call_type="afile_content"
         )
@@ -2480,9 +2481,9 @@ class Router:
         self._update_kwargs_before_fallbacks(
             model=model,
             kwargs=kwargs,
-            metadata_variable_name = _get_router_metadata_variable_name(
+            metadata_variable_name=_get_router_metadata_variable_name(
                 function_name=function_name
-            )
+            ),
         )
         try:
             verbose_router_logger.debug(
@@ -2812,6 +2813,8 @@ class Router:
         **kwargs,
     ) -> OpenAIFileObject:
         try:
+            from litellm.router_utils.common_utils import add_model_file_id_mappings
+
             verbose_router_logger.debug(
                 f"Inside _atext_completion()- model: {model}; kwargs: {kwargs}"
             )
@@ -2906,6 +2909,7 @@ class Router:
                 return response
 
             tasks = []
+
             if isinstance(healthy_deployments, dict):
                 tasks.append(create_file_for_deployment(healthy_deployments))
             else:
@@ -2916,7 +2920,15 @@ class Router:
 
             if len(responses) == 0:
                 raise Exception("No healthy deployments found.")
-            return responses[0]
+
+            model_file_id_mapping = add_model_file_id_mappings(
+                healthy_deployments=healthy_deployments, responses=responses
+            )
+            returned_response = cast(OpenAIFileObject, responses[0])
+            returned_response._hidden_params["model_file_id_mapping"] = (
+                model_file_id_mapping
+            )
+            return returned_response
         except Exception as e:
             verbose_router_logger.exception(
                 f"litellm.acreate_file(model={model}, {kwargs})\033[31m Exception {str(e)}\033[0m"

litellm/router_utils/common_utils.py

@@ -1,5 +1,9 @@
 import hashlib
 import json
+from typing import TYPE_CHECKING, Dict, List, Union
+
+if TYPE_CHECKING:
+    from litellm.types.llms.openai import OpenAIFileObject
 
 from litellm.types.router import CredentialLiteLLMParams
 
@@ -12,3 +16,25 @@ def get_litellm_params_sensitive_credential_hash(litellm_params: dict) -> str:
     return hashlib.sha256(
         json.dumps(sensitive_params.model_dump()).encode()
     ).hexdigest()
+
+
+def add_model_file_id_mappings(
+    healthy_deployments: Union[List[Dict], Dict], responses: List["OpenAIFileObject"]
+) -> dict:
+    """
+    Create a mapping of model name to file id
+    {
+        "model_id": "file_id",
+        "model_id": "file_id",
+    }
+    """
+    model_file_id_mapping = {}
+    if isinstance(healthy_deployments, list):
+        for deployment, response in zip(healthy_deployments, responses):
+            model_file_id_mapping[deployment.get("model_info", {}).get("id")] = (
+                response.id
+            )
+    elif isinstance(healthy_deployments, dict):
+        for model_id, file_id in healthy_deployments.items():
+            model_file_id_mapping[model_id] = file_id
+    return model_file_id_mapping

tests/enterprise/enterprise_hooks/test_managed_files.py

@@ -269,3 +269,57 @@ async def test_can_user_call_unified_file_id(call_type):
             data={"file_id": unified_file_id},
             call_type=call_type,
         )
+
+
+@pytest.mark.asyncio
+async def test_router_acreate_batch_only_selects_from_file_id_mapping(monkeypatch):
+    """
+    Test that router.acreate_batch only selects model_id from the file_id_mapping
+    """
+    import litellm
+
+    prisma_client = AsyncMock()
+    return_value = MagicMock()
+    return_value.created_by = "123"
+    prisma_client.db.litellm_managedobjecttable.find_first.return_value = return_value
+    proxy_managed_files = _PROXY_LiteLLMManagedFiles(
+        DualCache(), prisma_client=prisma_client
+    )
+
+    monkeypatch.setattr(
+        litellm,
+        "callbacks",
+        [proxy_managed_files],
+    )
+
+    router = litellm.Router(
+        model_list=[
+            {
+                "model_name": "gpt-3.5-turbo",
+                "litellm_params": {"model": "gpt-3.5-turbo"},
+                "model_info": {"id": "1234"},
+            },
+            {
+                "model_name": "gpt-3.5-turbo",
+                "litellm_params": {"model": "gpt-3.5-turbo"},
+                "model_info": {"id": "5678"},
+            },
+        ],
+    )
+
+    file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9vY3RldC1zdHJlYW07dW5pZmllZF9pZCw2YmQ4ZjhhYS02NmEzLTRmY2MtOTIxZS1lMTYwYzIzZWZjNzU7dGFyZ2V0X21vZGVsX25hbWVzLGdwdC00bztsbG1fb3V0cHV0X2ZpbGVfaWQsZmlsZS1MTENVRkI1MnVUTWE5aE5ZanRldzlWO2xsbV9vdXRwdXRfZmlsZV9tb2RlbF9pZCxmMzJlNWQ0OC05YWZmLTQ5YjMtOWE1Ny0zYzJhN2JjN2NjMmE"
+
+    model_file_id_mapping = {file_id: {"5678": "file-LLCUFB52uTMa9hNYjtew9V"}}
+
+    with patch.object(
+        litellm, "acreate_batch", return_value=AsyncMock()
+    ) as mock_acreate_batch:
+        for _ in range(1000):
+            response = await router.acreate_batch(
+                model="gpt-3.5-turbo",
+                input_file_id=file_id,
+                model_file_id_mapping=model_file_id_mapping,
+            )
+
+            mock_acreate_batch.assert_called()
+            assert "5678" in json.dumps(mock_acreate_batch.call_args.kwargs)

tests/test_litellm/proxy/test_proxy_server.py

@@ -458,3 +458,84 @@ def test_add_team_models_to_all_models():
         llm_router=llm_router,
     )
     assert result == {"gpt-4-model-2": {"team1"}}
+
+
+@pytest.mark.asyncio
+async def test_delete_deployment_type_mismatch():
+    """
+    Test that the _delete_deployment function handles type mismatches correctly.
+    Specifically test that models 12345678 and 12345679 are NOT deleted when
+    they exist in both combined_id_list (as integers) and router_model_ids (as strings).
+
+    This test reproduces the bug where type mismatch causes valid models to be deleted.
+    """
+    from unittest.mock import MagicMock, patch
+
+    from litellm.proxy.proxy_server import ProxyConfig
+
+    # Create mock ProxyConfig instance
+    pc = ProxyConfig()
+
+    pc.get_config = MagicMock(
+        return_value={
+            "model_list": [
+                {
+                    "model_name": "openai-gpt-4o",
+                    "litellm_params": {"model": "gpt-4o"},
+                    "model_info": {"id": 12345678},
+                },
+                {
+                    "model_name": "openai-gpt-4o",
+                    "litellm_params": {"model": "gpt-4o"},
+                    "model_info": {"id": 12345679},
+                },
+            ]
+        }
+    )
+
+    # Mock llm_router with string IDs (this is the source of the type mismatch)
+    mock_llm_router = MagicMock()
+    mock_llm_router.get_model_ids.return_value = [
+        "a96e12e76b36a57cfae57a41288eb41567629cac89b4828c6f7074afc3534695",
+        "a40186dd0fdb9b7282380277d7f57044d29de95bfbfcd7f4322b3493702d5cd3",
+        "12345678",  # String ID
+        "12345679",  # String ID
+    ]
+
+    # Track which deployments were deleted
+    deleted_ids = []
+
+    def mock_delete_deployment(id):
+        deleted_ids.append(id)
+        return True  # Simulate successful deletion
+
+    mock_llm_router.delete_deployment = MagicMock(side_effect=mock_delete_deployment)
+
+    # Mock get_config to return empty config (no config models)
+    async def mock_get_config(config_file_path):
+        return {}
+
+    pc.get_config = MagicMock(side_effect=mock_get_config)
+
+    # Patch the global llm_router
+    with patch("litellm.proxy.proxy_server.llm_router", mock_llm_router), patch(
+        "litellm.proxy.proxy_server.user_config_file_path", "test_config.yaml"
+    ):
+
+        # Call the function under test
+        deleted_count = await pc._delete_deployment(db_models=[])
+
+        # Assertions: Models 12345678 and 12345679 should NOT be deleted
+        # because they exist in combined_id_list (as integers) even though
+        # router has them as strings
+
+        # The function should delete the other 2 models that are not in combined_id_list
+        assert deleted_count == 0, f"Expected 0 deletions, got {deleted_count}"
+
+        # Verify that 12345678 and 12345679 were NOT deleted
+        assert (
+            "12345678" not in deleted_ids
+        ), f"Model 12345678 should NOT be deleted. Deleted IDs: {deleted_ids}"
+        assert (
+            "12345679" not in deleted_ids
+        ), f"Model 12345679 should NOT be deleted. Deleted IDs: {deleted_ids}"

tests/test_litellm/test_router.py

@@ -384,3 +384,42 @@ async def test_router_aretrieve_batch():
         print(mock_aretrieve_batch.call_args.kwargs)
         assert mock_aretrieve_batch.call_args.kwargs["api_key"] == "my-custom-key"
         assert mock_aretrieve_batch.call_args.kwargs["api_base"] == "my-custom-base"
+
+
+@pytest.mark.asyncio
+async def test_router_aretrieve_file_content():
+    """
+    Test that router.acreate_file with JSONL file returns the correct response
+    """
+
+    with patch.object(
+        litellm, "afile_content", return_value=AsyncMock()
+    ) as mock_afile_content:
+        router = litellm.Router(
+            model_list=[
+                {
+                    "model_name": "gpt-3.5-turbo",
+                    "litellm_params": {
+                        "model": "gpt-3.5-turbo",
+                        "custom_llm_provider": "azure",
+                        "api_key": "my-custom-key",
+                        "api_base": "my-custom-base",
+                    },
+                }
+            ],
+        )
+        try:
+            response = await router.afile_content(
+                **{
+                    "model": "gpt-3.5-turbo",
+                    "file_id": "my-unique-file-id",
+                }
+            )  # type: ignore
+        except Exception as e:
+            print(f"Error: {e}")
+
+        mock_afile_content.assert_called_once()
+
+        print(mock_afile_content.call_args.kwargs)
+        assert mock_afile_content.call_args.kwargs["api_key"] == "my-custom-key"
+        assert mock_afile_content.call_args.kwargs["api_base"] == "my-custom-base"