tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-06 06:24:16 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #21092 from BerriAI/litellm_azure_batches_issues

Browse Source 
Fix azure batches issues
This commit is contained in:
Sameer Kankute
2026-02-13 22:06:21 +05:30
committed by GitHub
parent fedfec51ab e0c98d62d4
commit a653a8ddde
5 changed files with 578 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/batches/batch_utils.py
+58 -6
View File
@@ -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
litellm/litellm_core_utils/litellm_logging.py
+1
View File
@@ -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
litellm/proxy/hooks/batch_rate_limiter.py
+78 -5
View File
@@ -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,
litellm/proxy/hooks/proxy_track_cost_callback.py
+2 -2
View File
@@ -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] = (
tests/batches_tests/test_batch_rate_limits.py
+439
View File
@@ -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")