diff --git a/litellm/batches/batch_utils.py b/litellm/batches/batch_utils.py index f80eae20f3..16a467e00c 100644 --- a/litellm/batches/batch_utils.py +++ b/litellm/batches/batch_utils.py @@ -39,11 +39,19 @@ async def _handle_completed_batch( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"], model_name: Optional[str] = None, + litellm_params: Optional[dict] = None, ) -> Tuple[float, Usage, List[str]]: - """Helper function to process a completed batch and handle logging""" + """Helper function to process a completed batch and handle logging + + Args: + batch: The batch object + custom_llm_provider: The LLM provider + model_name: Optional model name + litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.) + """ # Get batch results file_content_dictionary = await _get_batch_output_file_content_as_dictionary( - batch, custom_llm_provider + batch, custom_llm_provider, litellm_params=litellm_params ) # Calculate costs and usage @@ -187,9 +195,16 @@ def calculate_vertex_ai_batch_cost_and_usage( async def _get_batch_output_file_content_as_dictionary( batch: Batch, custom_llm_provider: Literal["openai", "azure", "vertex_ai", "hosted_vllm", "anthropic"] = "openai", + litellm_params: Optional[dict] = None, ) -> List[dict]: """ Get the batch output file content as a list of dictionaries + + Args: + batch: The batch object + custom_llm_provider: The LLM provider + litellm_params: Optional litellm parameters containing credentials (api_key, api_base, etc.) + Required for Azure and other providers that need authentication """ from litellm.files.main import afile_content from litellm.proxy.openai_files_endpoints.common_utils import ( @@ -211,13 +226,50 @@ async def _get_batch_output_file_content_as_dictionary( except (IndexError, AttributeError) as e: verbose_logger.error(f"Failed to extract LLM output file ID from unified file ID: {batch.output_file_id}, error: {e}") - _file_content = await afile_content( - file_id=file_id, - custom_llm_provider=custom_llm_provider, - ) + # Build kwargs for afile_content with credentials from litellm_params + file_content_kwargs = { + "file_id": file_id, + "custom_llm_provider": custom_llm_provider, + } + + # Extract and add credentials for file access + credentials = _extract_file_access_credentials(litellm_params) + file_content_kwargs.update(credentials) + + _file_content = await afile_content(**file_content_kwargs) return _get_file_content_as_dictionary(_file_content.content) +def _extract_file_access_credentials(litellm_params: Optional[dict]) -> dict: + """ + Extract credentials from litellm_params for file access operations. + + This method extracts relevant authentication and configuration parameters + needed for accessing files across different providers (Azure, Vertex AI, etc.). + + Args: + litellm_params: Dictionary containing litellm parameters with credentials + + Returns: + Dictionary containing only the credentials needed for file access + """ + credentials = {} + + if litellm_params: + # List of credential keys that should be passed to file operations + credential_keys = [ + "api_key", "api_base", "api_version", "organization", + "azure_ad_token", "azure_ad_token_provider", + "vertex_project", "vertex_location", "vertex_credentials", + "timeout", "max_retries" + ] + for key in credential_keys: + if key in litellm_params: + credentials[key] = litellm_params[key] + + return credentials + + def _get_file_content_as_dictionary(file_content: bytes) -> List[dict]: """ Get the file content as a list of dictionaries from JSON Lines format diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 8907b14a9d..bdbbc7579b 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2369,6 +2369,7 @@ class Logging(LiteLLMLoggingBaseClass): ) = await _handle_completed_batch( batch=result, custom_llm_provider=self.custom_llm_provider, + litellm_params=self.litellm_params, ) result._hidden_params["response_cost"] = response_cost diff --git a/litellm/proxy/hooks/batch_rate_limiter.py b/litellm/proxy/hooks/batch_rate_limiter.py index 78a371ad66..45b1bd8653 100644 --- a/litellm/proxy/hooks/batch_rate_limiter.py +++ b/litellm/proxy/hooks/batch_rate_limiter.py @@ -255,12 +255,24 @@ class _PROXY_BatchRateLimiter(CustomLogger): BatchFileUsage with total_tokens and request_count """ try: - # Read file content - file_content = await litellm.afile_content( - file_id=file_id, - custom_llm_provider=custom_llm_provider, - user_api_key_dict=user_api_key_dict, + # Check if this is a managed file (base64 encoded unified file ID) + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, ) + is_managed_file = _is_base64_encoded_unified_file_id(file_id) + if is_managed_file and user_api_key_dict is not None: + # For managed files, use the managed files hook directly + file_content = await self._fetch_managed_file_content( + file_id=file_id, + user_api_key_dict=user_api_key_dict, + ) + else: + # For non-managed files, use the standard litellm.afile_content + 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( file_content.content @@ -282,6 +294,67 @@ class _PROXY_BatchRateLimiter(CustomLogger): ) raise + async def _fetch_managed_file_content( + self, + file_id: str, + user_api_key_dict: UserAPIKeyAuth, + ) -> Any: + """ + Fetch file content from managed files hook. + + This is needed for managed files because they require proper user context + to verify file ownership and access permissions. + + Args: + file_id: The managed file ID (base64 encoded) + user_api_key_dict: User authentication information + + Returns: + HttpxBinaryResponseContent with the file content + """ + from litellm.llms.base_llm.files.transformation import BaseFileEndpoints + + # Import proxy_server dependencies at runtime to avoid circular imports + try: + from litellm.proxy.proxy_server import llm_router, proxy_logging_obj + except ImportError as e: + raise ValueError( + f"Cannot import proxy_server dependencies: {str(e)}. " + "Managed files require proxy_server to be initialized." + ) + + # Get the managed files hook + if proxy_logging_obj is None: + raise ValueError( + "proxy_logging_obj not available. Cannot access managed files hook." + ) + + managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") + if managed_files_obj is None: + raise ValueError( + "Managed files hook not found. Cannot access managed file." + ) + + if not isinstance(managed_files_obj, BaseFileEndpoints): + raise ValueError( + "Managed files hook is not a BaseFileEndpoints instance." + ) + + if llm_router is None: + raise ValueError( + "llm_router not available. Cannot access managed files." + ) + + # Use the managed files hook to get file content + # This properly handles user permissions and file ownership + file_content = await managed_files_obj.afile_content( + file_id=file_id, + litellm_parent_otel_span=user_api_key_dict.parent_otel_span, + llm_router=llm_router, + ) + + return file_content + async def async_pre_call_hook( self, user_api_key_dict: UserAPIKeyAuth, diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index dab5fb1bfd..37b79e6d06 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -202,8 +202,8 @@ class _ProxyDBLogger(CustomLogger): max_budget=end_user_max_budget, ) else: - if kwargs["stream"] is not True or ( - kwargs["stream"] is True and "complete_streaming_response" in kwargs + if kwargs.get("stream") is not True or ( + kwargs.get("stream") is True and "complete_streaming_response" in kwargs ): if sl_object is not None: cost_tracking_failure_debug_info: Union[dict, str] = ( diff --git a/tests/batches_tests/test_batch_rate_limits.py b/tests/batches_tests/test_batch_rate_limits.py index 13241e94d5..6bff3b82e5 100644 --- a/tests/batches_tests/test_batch_rate_limits.py +++ b/tests/batches_tests/test_batch_rate_limits.py @@ -619,3 +619,442 @@ async def test_batch_rate_limiter_without_user_context(): finally: os.unlink(file_path) + + +@pytest.mark.asyncio() +async def test_batch_rate_limiter_managed_files_regression(): + """ + Regression test for GEN-2166: Batch Rate Limiter Cannot Access User Files + + This test ensures that the batch rate limiter can properly access managed files + by verifying that: + 1. Managed files are detected correctly (base64 encoded unified file IDs) + 2. The _fetch_managed_file_content method uses the managed files hook + 3. User context (user_api_key_dict) is properly passed through + 4. No 403 errors occur when accessing files owned by the user + 5. The fix doesn't break non-managed file access + + This is a unit test that doesn't require external API calls. + """ + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.llms.base_llm.files.transformation import BaseFileEndpoints + from litellm.types.llms.openai import HttpxBinaryResponseContent + import httpx + + print("\n=== Regression Test: GEN-2166 Batch Rate Limiter Managed Files ===") + + # Setup: Create batch rate limiter + dual_cache = DualCache() + internal_usage_cache = InternalUsageCache(dual_cache=dual_cache) + rate_limiter = _PROXY_MaxParallelRequestsHandler_v3( + internal_usage_cache=internal_usage_cache + ) + batch_limiter = rate_limiter._get_batch_rate_limiter() + assert batch_limiter is not None + + # Setup: Create user API key dict + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key-regression", + user_id="test-user-regression", + tpm_limit=1000, + rpm_limit=10, + ) + + # Setup: Create mock file content (batch input file) + batch_content = b'{"custom_id": "request-1", "method": "POST", "url": "/v1/chat/completions", "body": {"model": "gpt-3.5-turbo", "messages": [{"role": "user", "content": "Test message for regression"}]}}' + + # Mock managed file ID (base64 encoded unified file ID format) + managed_file_id = "bGl0ZWxsbV9wcm94eTphcHBsaWNhdGlvbi9vY3RldC1zdHJlYW07dW5pZmllZF9pZCxyZWdyZXNzaW9uLXRlc3QtZmlsZQ==" + + # Test 1: Verify managed file detection + print("\n1. Verifying managed file detection...") + from litellm.proxy.openai_files_endpoints.common_utils import ( + _is_base64_encoded_unified_file_id, + ) + is_managed = _is_base64_encoded_unified_file_id(managed_file_id) + assert is_managed, "Managed file should be detected correctly" + print(" ✓ Managed file detected") + + # Test 2: Verify _fetch_managed_file_content uses managed files hook + print("\n2. Verifying managed files hook integration...") + + # Create mock managed files hook + class MockManagedFiles(BaseFileEndpoints): + def __init__(self): + self._afile_content_called = False + self._last_call_args = None + + async def acreate_file(self, *args, **kwargs): + pass + + async def afile_content(self, *args, **kwargs): + self._afile_content_called = True + self._last_call_args = kwargs + # Return mock file content + mock_response = httpx.Response( + status_code=200, + content=batch_content, + headers={"content-type": "application/octet-stream"}, + ) + return HttpxBinaryResponseContent(response=mock_response) + + async def afile_delete(self, *args, **kwargs): + pass + + async def afile_list(self, *args, **kwargs): + pass + + async def afile_retrieve(self, *args, **kwargs): + pass + + mock_managed_files = MockManagedFiles() + mock_llm_router = MagicMock() + mock_proxy_logging_obj = MagicMock() + mock_proxy_logging_obj.get_proxy_hook.return_value = mock_managed_files + + # Patch proxy_server imports + with patch.dict('sys.modules', { + 'litellm.proxy.proxy_server': MagicMock( + llm_router=mock_llm_router, + proxy_logging_obj=mock_proxy_logging_obj, + ) + }): + # Call _fetch_managed_file_content + result = await batch_limiter._fetch_managed_file_content( + file_id=managed_file_id, + user_api_key_dict=user_api_key_dict, + ) + + # Verify managed files hook was called + assert mock_managed_files._afile_content_called, \ + "REGRESSION: managed_files_obj.afile_content was not called! Bug GEN-2166 has returned." + + # Verify user context was passed + assert mock_managed_files._last_call_args is not None, \ + "REGRESSION: No arguments passed to afile_content" + assert 'file_id' in mock_managed_files._last_call_args, \ + "REGRESSION: file_id not passed to managed files hook" + assert mock_managed_files._last_call_args['file_id'] == managed_file_id, \ + "REGRESSION: Incorrect file_id passed" + assert 'llm_router' in mock_managed_files._last_call_args, \ + "REGRESSION: llm_router not passed to managed files hook" + + print(" ✓ Managed files hook called correctly") + print(" ✓ User context passed correctly") + + # Test 3: Verify count_input_file_usage uses managed files path + print("\n3. Verifying count_input_file_usage integration...") + + with patch.object(batch_limiter, '_fetch_managed_file_content') as mock_fetch: + mock_response = httpx.Response( + status_code=200, + content=batch_content, + headers={"content-type": "application/octet-stream"}, + ) + mock_fetch.return_value = HttpxBinaryResponseContent(response=mock_response) + + # Call count_input_file_usage with managed file + usage = await batch_limiter.count_input_file_usage( + file_id=managed_file_id, + custom_llm_provider="openai", + user_api_key_dict=user_api_key_dict, + ) + + # Verify _fetch_managed_file_content was called + assert mock_fetch.called, \ + "REGRESSION: _fetch_managed_file_content not called for managed files! Bug GEN-2166 has returned." + + # Verify correct parameters were passed + call_kwargs = mock_fetch.call_args.kwargs + assert call_kwargs['file_id'] == managed_file_id, \ + "REGRESSION: Incorrect file_id passed to _fetch_managed_file_content" + assert call_kwargs['user_api_key_dict'] == user_api_key_dict, \ + "REGRESSION: user_api_key_dict not passed! Bug GEN-2166 has returned." + + # Verify usage was calculated + assert usage.total_tokens > 0, "Token count should be greater than 0" + assert usage.request_count == 1, "Request count should be 1" + + print(" ✓ Managed file path used") + print(f" ✓ Token count: {usage.total_tokens}") + print(f" ✓ Request count: {usage.request_count}") + + # Test 4: Verify non-managed files still work + print("\n4. Verifying non-managed files still work...") + + non_managed_file_id = "file-abc123" # Standard OpenAI file ID + + with patch('litellm.afile_content') as mock_afile_content: + mock_response = httpx.Response( + status_code=200, + content=batch_content, + headers={"content-type": "application/octet-stream"}, + ) + mock_afile_content.return_value = HttpxBinaryResponseContent(response=mock_response) + + # Call count_input_file_usage with non-managed file + usage = await batch_limiter.count_input_file_usage( + file_id=non_managed_file_id, + custom_llm_provider="openai", + user_api_key_dict=user_api_key_dict, + ) + + # Verify litellm.afile_content was called + assert mock_afile_content.called, \ + "REGRESSION: litellm.afile_content not called for non-managed files" + + print(" ✓ Standard file path used") + print(f" ✓ Token count: {usage.total_tokens}") + + # Test 5: Verify the fix prevents 403 errors + print("\n5. Verifying 403 error prevention...") + + # Simulate the bug scenario: managed files hook not being used + with patch.object(batch_limiter, '_fetch_managed_file_content') as mock_fetch: + # If this is NOT called for managed files, the bug has returned + mock_fetch.side_effect = Exception("Should not be called if bug exists") + + # This should call _fetch_managed_file_content + try: + with patch('litellm.afile_content') as mock_afile_content: + # If litellm.afile_content is called for managed files, bug exists + mock_afile_content.side_effect = Exception( + "Error code: 403 - User does not have access to the file" + ) + + # Reset mock_fetch to return valid content + mock_response = httpx.Response( + status_code=200, + content=batch_content, + headers={"content-type": "application/octet-stream"}, + ) + mock_fetch.side_effect = None + mock_fetch.return_value = HttpxBinaryResponseContent(response=mock_response) + + # This should use _fetch_managed_file_content, not litellm.afile_content + usage = await batch_limiter.count_input_file_usage( + file_id=managed_file_id, + custom_llm_provider="openai", + user_api_key_dict=user_api_key_dict, + ) + + # Verify managed files path was used (not standard path that causes 403) + assert mock_fetch.called, \ + "REGRESSION: Managed files path not used! This would cause 403 errors." + assert not mock_afile_content.called, \ + "REGRESSION: Standard path used for managed files! This causes 403 errors." + + print(" ✓ 403 error prevention verified") + + except Exception as e: + if "403" in str(e): + pytest.fail( + f"REGRESSION: 403 error occurred! Bug GEN-2166 has returned. Error: {str(e)}" + ) + raise + + print("\n=== Regression Test Passed ===") + print("✓ Bug GEN-2166 is fixed and protected against regression") + print("✓ Managed files are properly accessed via managed files hook") + print("✓ User context is correctly passed through") + print("✓ No 403 errors occur") + print("✓ Non-managed files still work correctly\n") + + +@pytest.mark.asyncio() +async def test_batch_logging_azure_credentials_regression(): + """ + Regression test: LoggingWorker Missing Azure Credentials When Fetching Batch Output + + This test ensures that Azure credentials are properly passed when fetching batch + output files during logging, preventing "Missing credentials" errors. + + Bug: The LoggingWorker failed when processing completed Azure batches because + it attempted to fetch batch output file content without Azure credentials. + + Fix: Pass litellm_params (containing credentials) from the logging object + through to the file content retrieval functions. + """ + from unittest.mock import AsyncMock, MagicMock, patch + from litellm.batches.batch_utils import ( + _extract_file_access_credentials, + _get_batch_output_file_content_as_dictionary, + _handle_completed_batch, + ) + from litellm.types.llms.openai import Batch, HttpxBinaryResponseContent + import httpx + + print("\n=== Regression Test: Azure Batch Logging Credentials ===") + + # Setup: Create mock batch with output file + mock_batch = Batch( + id="batch-azure-test", + object="batch", + endpoint="/v1/chat/completions", + errors=None, + input_file_id="file-input-azure", + completion_window="24h", + status="completed", + output_file_id="file-output-azure", + error_file_id=None, + created_at=1234567890, + in_progress_at=1234567900, + expires_at=1234654290, + finalizing_at=1234568000, + completed_at=1234568100, + failed_at=None, + expired_at=None, + cancelling_at=None, + cancelled_at=None, + request_counts=None, + metadata=None, + ) + + # Setup: Azure credentials (as they would be in litellm_params) + azure_credentials = { + "api_key": "test-azure-key-regression", + "api_base": "https://test-regression.openai.azure.com", + "api_version": "2024-02-15-preview", + "organization": "test-org", + "timeout": 600, + } + + # Setup: Mock batch output content + batch_output = b'{"id": "batch_req_1", "custom_id": "request-1", "response": {"status_code": 200, "body": {"id": "chatcmpl-azure", "object": "chat.completion", "model": "gpt-4", "usage": {"prompt_tokens": 15, "completion_tokens": 25, "total_tokens": 40}}}}\n' + + # Test 1: Verify _extract_file_access_credentials works correctly + print("\n1. Testing credential extraction...") + + extracted_creds = _extract_file_access_credentials(azure_credentials) + assert "api_key" in extracted_creds, "api_key should be extracted" + assert extracted_creds["api_key"] == "test-azure-key-regression", "Incorrect api_key" + assert "api_base" in extracted_creds, "api_base should be extracted" + assert "api_version" in extracted_creds, "api_version should be extracted" + assert "timeout" in extracted_creds, "timeout should be extracted" + + print(" ✓ Credentials extracted correctly") + print(f" ✓ Extracted keys: {list(extracted_creds.keys())}") + + # Test 2: Verify credentials are passed to afile_content + print("\n2. Testing credentials passed to afile_content...") + + credentials_received = {"value": False, "params": None} + + async def mock_afile_content_tracker(**kwargs): + # Track if Azure credentials were passed + if "api_key" in kwargs and "api_base" in kwargs and "api_version" in kwargs: + credentials_received["value"] = True + credentials_received["params"] = { + "api_key": kwargs.get("api_key"), + "api_base": kwargs.get("api_base"), + "api_version": kwargs.get("api_version"), + } + mock_response = httpx.Response( + status_code=200, + content=batch_output, + headers={"content-type": "application/octet-stream"}, + ) + return HttpxBinaryResponseContent(response=mock_response) + + with patch('litellm.files.main.afile_content', side_effect=mock_afile_content_tracker): + result = await _get_batch_output_file_content_as_dictionary( + batch=mock_batch, + custom_llm_provider="azure", + litellm_params=azure_credentials, + ) + + # Verify credentials were passed + assert credentials_received["value"], \ + "REGRESSION: Azure credentials not passed to afile_content! This causes 'Missing credentials' error." + assert credentials_received["params"]["api_key"] == "test-azure-key-regression", \ + "REGRESSION: Incorrect api_key" + assert credentials_received["params"]["api_base"] == "https://test-regression.openai.azure.com", \ + "REGRESSION: Incorrect api_base" + + print(" ✓ Credentials passed to afile_content") + print(f" ✓ api_key: {credentials_received['params']['api_key']}") + print(f" ✓ api_base: {credentials_received['params']['api_base']}") + + # Test 3: Verify full flow through _handle_completed_batch + print("\n3. Testing full logging flow...") + + credentials_received["value"] = False + credentials_received["params"] = None + + with patch('litellm.files.main.afile_content', side_effect=mock_afile_content_tracker): + cost, usage, models = await _handle_completed_batch( + batch=mock_batch, + custom_llm_provider="azure", + litellm_params=azure_credentials, + ) + + # Verify credentials were passed through the entire flow + assert credentials_received["value"], \ + "REGRESSION: Credentials not passed through _handle_completed_batch" + + # Verify cost and usage were calculated + assert cost > 0, "Cost should be calculated" + assert usage.total_tokens == 40, "Usage should be calculated correctly" + + print(" ✓ Credentials passed through full flow") + print(f" ✓ Cost: {cost}") + print(f" ✓ Usage: {usage.total_tokens} tokens") + print(f" ✓ Models: {models}") + + # Test 4: Verify error prevention + print("\n4. Testing 'Missing credentials' error prevention...") + + # Simulate the bug: if credentials are NOT passed, Azure would fail + with patch('litellm.files.main.afile_content') as mock_afile_content_fail: + # This is what would happen without the fix + mock_afile_content_fail.side_effect = Exception( + "Missing credentials. Please pass one of `api_key`, `azure_ad_token`, " + "`azure_ad_token_provider`, or the `AZURE_OPENAI_API_KEY` or " + "`AZURE_OPENAI_AD_TOKEN` environment variables." + ) + + # Now test with the fix - should NOT raise the error + with patch('litellm.files.main.afile_content', side_effect=mock_afile_content_tracker): + try: + cost, usage, models = await _handle_completed_batch( + batch=mock_batch, + custom_llm_provider="azure", + litellm_params=azure_credentials, + ) + print(" ✓ No 'Missing credentials' error with fix") + except Exception as e: + if "Missing credentials" in str(e): + pytest.fail( + f"REGRESSION: 'Missing credentials' error occurred! " + f"Credentials not being passed. Error: {str(e)}" + ) + raise + + # Test 5: Verify backwards compatibility (works without credentials for OpenAI) + print("\n5. Testing backwards compatibility...") + + with patch('litellm.files.main.afile_content') as mock_afile_content: + mock_response = httpx.Response( + status_code=200, + content=batch_output, + headers={"content-type": "application/octet-stream"}, + ) + mock_afile_content.return_value = HttpxBinaryResponseContent(response=mock_response) + + # Call without litellm_params (should still work for OpenAI) + result = await _get_batch_output_file_content_as_dictionary( + batch=mock_batch, + custom_llm_provider="openai", + litellm_params=None, + ) + + assert len(result) > 0, "Should return file content" + print(" ✓ Backwards compatibility maintained") + print(" ✓ Works without litellm_params for OpenAI") + + print("\n=== Regression Test Passed ===") + print("✓ Azure credentials properly passed from logging to file retrieval") + print("✓ 'Missing credentials' error prevented") + print("✓ Batch output files can be fetched with Azure credentials") + print("✓ Cost and usage tracking works for Azure batches") + print("✓ Backwards compatibility maintained\n")