Merge pull request #20838 from BerriAI/litellm_managed_error_file

Add support managed error file
2026-08-02 12:21:10 +00:00 · 2026-02-10 18:20:26 +05:30
parent 0b4c511576 63eedc26b0
commit cdab87dec0
4 changed files with 370 additions and 53 deletions
@@ -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):
            """
@@ -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(
@@ -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)
@@ -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"