From e2f6fb2d7c64e3657fe5412c18bf864d3f3107ff Mon Sep 17 00:00:00 2001 From: Krish Dholakia Date: Wed, 25 Jun 2025 21:27:06 -0700 Subject: [PATCH] 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 | 46 +++++++++-- litellm/proxy/_new_secret_config.yaml | 17 ++++ litellm/proxy/proxy_server.py | 2 + litellm/router.py | 18 ++++- litellm/router_utils/common_utils.py | 26 ++++++ .../enterprise_hooks/test_managed_files.py | 54 +++++++++++++ tests/test_litellm/proxy/test_proxy_server.py | 81 +++++++++++++++++++ tests/test_litellm/test_router.py | 39 +++++++++ 8 files changed, 272 insertions(+), 11 deletions(-) diff --git a/enterprise/enterprise_hooks/managed_files.py b/enterprise/enterprise_hooks/managed_files.py index c752395ac6..e06e83e272 100644 --- a/enterprise/enterprise_hooks/managed_files.py +++ b/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, diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 40243ae668..21c8a6148c 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/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 \ No newline at end of file diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index e0004bebc9..299d244c8e 100644 --- a/litellm/proxy/proxy_server.py +++ b/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() diff --git a/litellm/router.py b/litellm/router.py index 2d9796b62f..d301fe8d5e 100644 --- a/litellm/router.py +++ b/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" diff --git a/litellm/router_utils/common_utils.py b/litellm/router_utils/common_utils.py index 6e90943d49..a7163545d3 100644 --- a/litellm/router_utils/common_utils.py +++ b/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 diff --git a/tests/enterprise/enterprise_hooks/test_managed_files.py b/tests/enterprise/enterprise_hooks/test_managed_files.py index 81e27d941f..76ebcb9fb9 100644 --- a/tests/enterprise/enterprise_hooks/test_managed_files.py +++ b/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) diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index deaf69afe6..8688d08ee4 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/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}" diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index eb6c740818..921cb17808 100644 --- a/tests/test_litellm/test_router.py +++ b/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"