diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 569ea17f6d..a41b3f3bf6 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -899,49 +899,49 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): batch_id=response.id, model_id=model_id ) - if ( - response.output_file_id and model_id - ): # return a file id with the model_id and output_file_id - original_output_file_id = response.output_file_id - response.output_file_id = self.get_unified_output_file_id( - output_file_id=response.output_file_id, - model_id=model_id, - model_name=model_name, - ) - - # Fetch the actual file object for the output file - file_object = None - try: - # Use litellm to retrieve the file object from the provider - from litellm import afile_retrieve - file_object = await afile_retrieve( - custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai", - file_id=original_output_file_id + # Handle both output_file_id and error_file_id + for file_attr in ["output_file_id", "error_file_id"]: + file_id_value = getattr(response, file_attr, None) + if file_id_value and model_id: + original_file_id = file_id_value + unified_file_id = self.get_unified_output_file_id( + output_file_id=original_file_id, + model_id=model_id, + model_name=model_name, ) - verbose_logger.debug( - f"Successfully retrieved file object for output_file_id={original_output_file_id}" + setattr(response, file_attr, unified_file_id) + + # Fetch the actual file object from the provider + file_object = None + try: + # Use litellm to retrieve the file object from the provider + from litellm import afile_retrieve + file_object = await afile_retrieve( + custom_llm_provider=model_name.split("/")[0] if model_name and "/" in model_name else "openai", + file_id=original_file_id + ) + verbose_logger.debug( + f"Successfully retrieved file object for {file_attr}={original_file_id}" + ) + except Exception as e: + verbose_logger.warning( + f"Failed to retrieve file object for {file_attr}={original_file_id}: {str(e)}. Storing with None and will fetch on-demand." + ) + + await self.store_unified_file_id( + file_id=unified_file_id, + file_object=file_object, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + model_mappings={model_id: original_file_id}, + user_api_key_dict=user_api_key_dict, ) - except Exception as e: - verbose_logger.warning( - f"Failed to retrieve file object for output_file_id={original_output_file_id}: {str(e)}. Storing with None and will fetch on-demand." - ) - - await self.store_unified_file_id( - file_id=response.output_file_id, - file_object=file_object, - litellm_parent_otel_span=user_api_key_dict.parent_otel_span, - model_mappings={model_id: original_output_file_id}, - user_api_key_dict=user_api_key_dict, - ) - asyncio.create_task( - self.store_unified_object_id( - unified_object_id=response.id, - file_object=response, - litellm_parent_otel_span=user_api_key_dict.parent_otel_span, - model_object_id=original_response_id, - file_purpose="batch", - user_api_key_dict=user_api_key_dict, - ) + await self.store_unified_object_id( + unified_object_id=response.id, + file_object=response, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + model_object_id=original_response_id, + file_purpose="batch", + user_api_key_dict=user_api_key_dict, ) elif isinstance(response, LiteLLMFineTuningJob): ## Check if unified_file_id is in the response @@ -958,15 +958,13 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): response.id = self.get_unified_generic_response_id( model_id=model_id, generic_response_id=response.id ) - asyncio.create_task( - self.store_unified_object_id( - unified_object_id=response.id, - file_object=response, - litellm_parent_otel_span=user_api_key_dict.parent_otel_span, - model_object_id=original_response_id, - file_purpose="fine-tune", - user_api_key_dict=user_api_key_dict, - ) + await self.store_unified_object_id( + unified_object_id=response.id, + file_object=response, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + model_object_id=original_response_id, + file_purpose="fine-tune", + user_api_key_dict=user_api_key_dict, ) elif isinstance(response, AsyncCursorPage): """ diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index ecad8bc1b1..78a371ad66 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -17,7 +17,7 @@ Quick summary: - async_log_success_event() fires on GET /v1/batches/{id} (batch completion) """ -from typing import TYPE_CHECKING, Any, Dict, List, Literal, Union +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Union from fastapi import HTTPException from pydantic import BaseModel @@ -241,6 +241,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): self, file_id: str, custom_llm_provider: Literal["openai", "azure", "vertex_ai"] = "openai", + user_api_key_dict: Optional[UserAPIKeyAuth] = None, ) -> BatchFileUsage: """ Count number of requests and tokens in a batch input file. @@ -248,6 +249,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): Args: file_id: The file ID to read custom_llm_provider: The custom LLM provider to use for token encoding + user_api_key_dict: User authentication information for file access (required for managed files) Returns: BatchFileUsage with total_tokens and request_count @@ -257,6 +259,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): file_content = await litellm.afile_content( file_id=file_id, custom_llm_provider=custom_llm_provider, + user_api_key_dict=user_api_key_dict, ) file_content_as_dict = _get_file_content_as_dictionary( @@ -336,6 +339,7 @@ class _PROXY_BatchRateLimiter(CustomLogger): batch_usage = await self.count_input_file_usage( file_id=input_file_id, custom_llm_provider=custom_llm_provider, + user_api_key_dict=user_api_key_dict, ) verbose_proxy_logger.debug( diff --git a/tests/batches_tests/test_batch_rate_limits.py b/tests/batches_tests/test_batch_rate_limits.py index 776aba438c..13241e94d5 100644 --- a/tests/batches_tests/test_batch_rate_limits.py +++ b/tests/batches_tests/test_batch_rate_limits.py @@ -389,3 +389,233 @@ async def test_batch_rate_limit_multiple_requests(): print(f" Error: {exc_info.value.detail}") finally: os.unlink(file_path_2) + + +@pytest.mark.asyncio() +@pytest.mark.skipif( + os.environ.get("OPENAI_API_KEY") is None, + reason="OPENAI_API_KEY not set - skipping integration test" +) +async def test_batch_rate_limiter_with_managed_files(): + """ + Test for GEN-2166: Verify batch rate limiter can read user files when managed files are enabled. + + This test ensures that: + 1. The batch rate limiter passes user_api_key_dict to afile_content() + 2. The managed files hook can verify file ownership correctly + 3. Rate limiting is enforced (not silently bypassed) + 4. No 403 Permission Denied errors occur for files owned by the user + """ + import tempfile + from unittest.mock import AsyncMock, MagicMock, patch + + CUSTOM_LLM_PROVIDER = "openai" + + # Setup: Create internal usage cache and rate limiter + dual_cache = DualCache() + internal_usage_cache = InternalUsageCache(dual_cache=dual_cache) + rate_limiter = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=internal_usage_cache + ) + + # Setup: Get batch rate limiter + batch_limiter = rate_limiter._get_batch_rate_limiter() + assert batch_limiter is not None, "Batch rate limiter should be available" + + # Setup: Create user API key with TPM = 500, RPM = 10 + test_user_id = "test-user-abc123" + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key-managed-files", + user_id=test_user_id, + tpm_limit=500, + rpm_limit=10, + ) + + print(f"\n=== Testing Batch Rate Limiter with Managed Files ===") + print(f"User ID: {test_user_id}") + + # Create a batch file with ~200 tokens + import json as json_lib + message = "This is a test message for batch rate limiting with managed files. " * 5 + requests = [] + for i in range(1, 4): + request_obj = { + "custom_id": f"request-{i}", + "method": "POST", + "url": "/v1/chat/completions", + "body": { + "model": "gpt-3.5-turbo", + "messages": [{"role": "user", "content": message}] + } + } + requests.append(json_lib.dumps(request_obj)) + + batch_content = "\n".join(requests) + + with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False) as f: + f.write(batch_content) + file_path = f.name + + try: + # Step 1: Upload file to OpenAI (simulating user upload) + print("\n1. Uploading batch input file...") + file_obj = await litellm.acreate_file( + file=open(file_path, "rb"), + purpose="batch", + custom_llm_provider=CUSTOM_LLM_PROVIDER, + ) + print(f" ✓ File uploaded: {file_obj.id}") + await asyncio.sleep(1) # Give API time to process + + # Step 2: Mock managed files hook to simulate file ownership check + # In a real scenario, the managed files hook would check if the user owns the file + # For this test, we'll verify that user_api_key_dict is passed correctly + print("\n2. Testing rate limiter file access with user context...") + + # Track if user_api_key_dict was passed to afile_content + original_afile_content = litellm.afile_content + user_context_passed = {"value": False} + + async def mock_afile_content(*args, **kwargs): + # Check if user_api_key_dict was passed + if "user_api_key_dict" in kwargs and kwargs["user_api_key_dict"] is not None: + user_context_passed["value"] = True + print(f" ✓ user_api_key_dict passed to afile_content") + print(f" User ID: {kwargs['user_api_key_dict'].user_id}") + else: + print(f" ✗ user_api_key_dict NOT passed to afile_content (BUG!)") + + # Call original function + return await original_afile_content(*args, **kwargs) + + # Patch afile_content to track the call + with patch('litellm.afile_content', side_effect=mock_afile_content): + data = { + "model": "gpt-3.5-turbo", + "input_file_id": file_obj.id, + "custom_llm_provider": CUSTOM_LLM_PROVIDER, + } + + # Step 3: Submit batch and verify rate limiting works + print("\n3. Submitting batch with rate limiting...") + result = await batch_limiter.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=dual_cache, + data=data, + call_type="acreate_batch", + ) + + tokens_used = result.get('_batch_token_count', 0) + requests_count = result.get('_batch_request_count', 0) + print(f" ✓ Batch submitted successfully") + print(f" Tokens counted: {tokens_used}") + print(f" Requests counted: {requests_count}") + print(f" Rate limit usage: {tokens_used}/500 TPM, {requests_count}/10 RPM") + + # Step 4: Verify user context was passed + print("\n4. Verifying fix for GEN-2166...") + assert user_context_passed["value"], ( + "FAILED: user_api_key_dict was not passed to afile_content(). " + "This means the bug GEN-2166 is not fixed!" + ) + print(" ✓ Fix verified: user_api_key_dict is correctly passed") + + # Step 5: Verify rate limiting is actually enforced (not bypassed) + print("\n5. Verifying rate limiting is enforced...") + assert tokens_used > 0, "Token count should be greater than 0" + assert requests_count > 0, "Request count should be greater than 0" + print(" ✓ Rate limiting is active (not silently bypassed)") + + print("\n=== Test Passed: GEN-2166 Fix Verified ===") + print("✓ Batch rate limiter can access user files") + print("✓ User context is correctly passed") + print("✓ Rate limiting is enforced") + print("✓ No silent failures") + + except HTTPException as e: + if e.status_code == 403: + pytest.fail( + f"FAILED: Got 403 Permission Denied error. " + f"This indicates the bug GEN-2166 is not fixed. " + f"Error: {e.detail}" + ) + else: + raise + except Exception as e: + pytest.fail(f"Unexpected error: {str(e)}") + finally: + os.unlink(file_path) + + +@pytest.mark.asyncio() +async def test_batch_rate_limiter_without_user_context(): + """ + Test that verifies the bug scenario from GEN-2166. + + When user_api_key_dict is NOT passed to count_input_file_usage(), + the function should still work for non-managed files, but would fail + for managed files (which is the bug we fixed). + + This test documents the expected behavior with and without user context. + """ + import tempfile + + CUSTOM_LLM_PROVIDER = "openai" + + # Setup + BATCH_LIMITER = _PROXY_BatchRateLimiter( + internal_usage_cache=None, + parallel_request_limiter=None, + ) + + # Create a simple batch file + batch_content = """{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Hello"}]}}""" + + with tempfile.NamedTemporaryFile(mode='w', suffix='.jsonl', delete=False) as f: + f.write(batch_content) + file_path = f.name + + try: + # Upload file + file_obj = await litellm.acreate_file( + file=open(file_path, "rb"), + purpose="batch", + custom_llm_provider=CUSTOM_LLM_PROVIDER, + ) + await asyncio.sleep(1) + + # Test 1: Without user context (old behavior - would fail with managed files) + print("\n=== Test 1: count_input_file_usage WITHOUT user context ===") + try: + usage_without_context = await BATCH_LIMITER.count_input_file_usage( + file_id=file_obj.id, + custom_llm_provider=CUSTOM_LLM_PROVIDER, + user_api_key_dict=None, # Explicitly passing None + ) + print(f"✓ Works for non-managed files (tokens: {usage_without_context.total_tokens})") + print(" Note: Would fail with 403 for managed files (GEN-2166 bug)") + except Exception as e: + print(f"✗ Failed: {str(e)}") + + # Test 2: With user context (new behavior - works with managed files) + print("\n=== Test 2: count_input_file_usage WITH user context ===") + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + user_id="test-user-123", + ) + + usage_with_context = await BATCH_LIMITER.count_input_file_usage( + file_id=file_obj.id, + custom_llm_provider=CUSTOM_LLM_PROVIDER, + user_api_key_dict=user_api_key_dict, # Passing user context + ) + print(f"✓ Works with user context (tokens: {usage_with_context.total_tokens})") + print(" Note: This fixes GEN-2166 for managed files") + + # Verify both return the same results + assert usage_with_context.total_tokens == usage_without_context.total_tokens + assert usage_with_context.request_count == usage_without_context.request_count + print("\n✓ Both methods return identical results for non-managed files") + + finally: + os.unlink(file_path) diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index 946c5ad172..58efa854e7 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -192,7 +192,7 @@ async def test_async_post_call_success_hook_for_unified_finetuning_job(): "model_id": "gpt-3.5-turbo-0613", } proxy_managed_files = _PROXY_LiteLLMManagedFiles( - DualCache(), prisma_client=MagicMock() + DualCache(), prisma_client=AsyncMock() ) data = { "user_api_key_dict": {"parent_otel_span": MagicMock()}, @@ -373,6 +373,91 @@ async def test_output_file_id_for_batch_retrieve(): assert not cast(LiteLLMBatch, response).output_file_id.startswith("file-") +@pytest.mark.asyncio +async def test_error_file_id_for_failed_batch(): + """ + Test that the error_file_id is properly managed when a batch fails + """ + from typing import cast + + from openai.types.batch import BatchRequestCounts + + from litellm.proxy._types import UserAPIKeyAuth + from litellm.types.llms.openai import OpenAIFileObject + from litellm.types.utils import LiteLLMBatch + + batch = LiteLLMBatch( + id="bGl0ZWxsbV9wcm94eTttb2RlbF9pZDoxMjM0NTY3OTtsbG1fYmF0Y2hfaWQ6YmF0Y2hfYWJjMTIz", + completion_window="24h", + created_at=1714508499, + endpoint="/v1/chat/completions", + input_file_id="file-abc123", + object="batch", + status="failed", + cancelled_at=None, + cancelling_at=None, + completed_at=None, + error_file_id="error-abc123", + errors=None, + expired_at=None, + expires_at=1714536634, + failed_at=None, + finalizing_at=None, + in_progress_at=None, + metadata=None, + output_file_id=None, + request_counts=BatchRequestCounts(completed=0, failed=0, total=0), + usage=None, + ) + + batch._hidden_params = { + "litellm_call_id": "test-call-id", + "api_base": "https://api.openai.com", + "model_id": "test-model-id", + "model_name": "gpt-4o", + "response_cost": 0.0, + "additional_headers": {}, + "litellm_model_name": "gpt-4o", + "unified_batch_id": "litellm_proxy;model_id:test-model-id;llm_batch_id:batch_abc123", + } + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=AsyncMock() + ) + + # Create a proper OpenAIFileObject for the error file + error_file_object = OpenAIFileObject( + id="error-abc123", + object="file", + bytes=1234, + created_at=1714508500, + filename="error.jsonl", + purpose="batch_output", + status="processed", + ) + + # Mock the afile_retrieve to simulate retrieving error file metadata + with patch("litellm.afile_retrieve", new_callable=AsyncMock) as mock_retrieve: + mock_retrieve.return_value = error_file_object + + user_api_key_dict = UserAPIKeyAuth( + user_id="test-user-123", + parent_otel_span=MagicMock() + ) + + response = await proxy_managed_files.async_post_call_success_hook( + data={}, + user_api_key_dict=user_api_key_dict, + response=batch, + ) + + # Verify that error_file_id was transformed to a managed file ID + assert cast(LiteLLMBatch, response).error_file_id is not None + assert not cast(LiteLLMBatch, response).error_file_id.startswith("error-") + # Verify it's a base64 encoded managed file ID + assert _is_base64_encoded_unified_file_id(cast(LiteLLMBatch, response).error_file_id) + + @pytest.mark.asyncio async def test_async_post_call_success_hook_twice_assert_no_unique_violation(): import asyncio @@ -849,7 +934,7 @@ async def test_check_file_ids_access_with_unified_file_ids(): Test that check_file_ids_access validates user access to managed file IDs. """ from litellm.proxy._types import UserAPIKeyAuth - + # Create a unified file ID unified_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9wZGY7dW5pZmllZF9pZCw2YzBiNTg5MC04OTE0LTQ4ZTAtYjhmNC0wYWU1ZWQzYzE0YTU7dGFyZ2V0X21vZGVsX25hbWVzLGdwdC00bztsbG1fb3V0cHV0X2ZpbGVfaWQsZmlsZS1FQ0JQVzdNTDlnN1hIZHdHZ1VQWmFNO2xsbV9vdXRwdXRfZmlsZV9tb2RlbF9pZCxlMjY0NTNmOWU3NmU3OTkzNjgwZDAwNjhkOThjMWY0Y2MyMDViYmFkMDk2N2EzM2M2NjQ4OTM1NjhjYTc0M2My" regular_file_id = "file-abc123"