diff --git a/litellm/proxy/openai_files_endpoints/files_endpoints.py b/litellm/proxy/openai_files_endpoints/files_endpoints.py index 7e3f582081..2c6b378ae3 100644 --- a/litellm/proxy/openai_files_endpoints/files_endpoints.py +++ b/litellm/proxy/openai_files_endpoints/files_endpoints.py @@ -146,8 +146,8 @@ async def route_create_file( Priority: 1. If target_storage is specified and not "default" -> use storage backend 2. If model parameter provided -> use model credentials and encode ID - 3. If enable_loadbalancing_on_batch_endpoints -> deprecated loadbalancing - 4. If target_model_names_list -> managed files (requires DB) + 3. If target_model_names_list -> managed files (requires DB, supports loadbalancing) + 4. If enable_loadbalancing_on_batch_endpoints -> deprecated loadbalancing 5. Else -> use custom_llm_provider with files_settings """ @@ -202,18 +202,9 @@ async def route_create_file( return response - # EXISTING: Deprecated loadbalancing approach - if ( - litellm.enable_loadbalancing_on_batch_endpoints is True - and is_router_model - and router_model is not None - ): - response = await _deprecated_loadbalanced_create_file( - llm_router=llm_router, - router_model=router_model, - _create_file_request=_create_file_request, - ) - elif target_model_names_list: + # Handle managed files (supports loadbalancing via llm_router.acreate_file) + # Priority: Check for managed files BEFORE deprecated loadbalancing + if target_model_names_list: managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") if managed_files_obj is None: raise ProxyException( @@ -236,6 +227,7 @@ async def route_create_file( param="None", code=500, ) + # Managed files internally calls llm_router.acreate_file() which includes loadbalancing response = await managed_files_obj.acreate_file( llm_router=llm_router, create_file_request=_create_file_request, @@ -243,6 +235,17 @@ async def route_create_file( litellm_parent_otel_span=user_api_key_dict.parent_otel_span, user_api_key_dict=user_api_key_dict, ) + # EXISTING: Deprecated loadbalancing approach (for backwards compatibility when not using managed files) + elif ( + litellm.enable_loadbalancing_on_batch_endpoints is True + and is_router_model + and router_model is not None + ): + response = await _deprecated_loadbalanced_create_file( + llm_router=llm_router, + router_model=router_model, + _create_file_request=_create_file_request, + ) else: # get configs for custom_llm_provider llm_provider_config = get_files_provider_config( diff --git a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py index 4651bf59b4..b86f927ea0 100644 --- a/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py +++ b/tests/test_litellm/proxy/openai_files_endpoint/test_files_endpoint.py @@ -856,3 +856,97 @@ def test_create_file_without_expires_after(mocker: MockerFixture, monkeypatch, l result = response.json() assert result["id"] == "file-abc123" assert result["purpose"] == "fine-tune" + + +def test_managed_files_with_loadbalancing(mocker: MockerFixture, monkeypatch, llm_router: Router): + """ + Test that managed files work with loadbalancing when both target_model_names + and enable_loadbalancing_on_batch_endpoints are enabled. + + This ensures that the priority order is correct: + - managed files should take precedence over deprecated loadbalancing + - managed files internally use llm_router.acreate_file() which provides loadbalancing + """ + from litellm.llms.base_llm.files.transformation import BaseFileEndpoints + from litellm.types.llms.openai import OpenAIFileObject + + # Enable loadbalancing on batch endpoints + monkeypatch.setattr("litellm.enable_loadbalancing_on_batch_endpoints", True) + + proxy_logging_obj = ProxyLogging( + user_api_key_cache=DualCache(default_in_memory_ttl=1) + ) + proxy_logging_obj._add_proxy_hooks(llm_router) + + # Track calls to verify loadbalancing through router + router_acreate_file_calls = [] + + class ManagedFilesWithLoadbalancing(BaseFileEndpoints): + async def acreate_file(self, llm_router, create_file_request, target_model_names_list, litellm_parent_otel_span, user_api_key_dict): + # Verify we receive the target model names + assert len(target_model_names_list) > 0, "Should have target_model_names_list" + + # Simulate what managed files does - call llm_router.acreate_file for each model + # This is where loadbalancing happens internally + for model in target_model_names_list: + router_acreate_file_calls.append({ + "model": model, + "via_router": True + }) + + # Return a managed file ID (base64 encoded) + return OpenAIFileObject( + id="litellm_managed_file_abc123", + object="file", + bytes=100, + created_at=1234567890, + filename="batch_data.jsonl", + purpose="batch", + status="uploaded", + ) + + async def afile_retrieve(self, file_id, litellm_parent_otel_span, llm_router): + raise NotImplementedError("Not implemented for test") + + async def afile_list(self, purpose, litellm_parent_otel_span): + raise NotImplementedError("Not implemented for test") + + async def afile_delete(self, file_id, litellm_parent_otel_span, llm_router, **data): + raise NotImplementedError("Not implemented for test") + + async def afile_content(self, file_id, litellm_parent_otel_span, llm_router, **data): + raise NotImplementedError("Not implemented for test") + + proxy_logging_obj.proxy_hook_mapping["managed_files"] = ManagedFilesWithLoadbalancing() + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", llm_router) + monkeypatch.setattr( + "litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj + ) + + # Create batch file content + test_file_content = b'{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello"}]}}' + test_file = ("batch_data.jsonl", test_file_content, "application/jsonl") + + # Make request with both target_model_names AND enable_loadbalancing_on_batch_endpoints + response = client.post( + "/v1/files", + files={"file": test_file}, + data={ + "purpose": "batch", + "target_model_names": "azure-gpt-3-5-turbo,gpt-3.5-turbo", # Multiple models + }, + headers={"Authorization": "Bearer test-key"}, + ) + + # Verify success + assert response.status_code == 200 + result = response.json() + assert result["id"] == "litellm_managed_file_abc123" + assert result["purpose"] == "batch" + + # Verify that managed files was called (via router for loadbalancing) + # This proves that managed files took precedence over deprecated loadbalancing + assert len(router_acreate_file_calls) == 2, "Should have called router for both models" + assert router_acreate_file_calls[0]["model"] == "azure-gpt-3-5-turbo" + assert router_acreate_file_calls[1]["model"] == "gpt-3.5-turbo" + assert all(call["via_router"] for call in router_acreate_file_calls), "All calls should go through router"