ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fixes based on greptile reviews

Browse Source
This commit is contained in:
Sameer Kankute 2026-02-18 12:19:11 +05:30
parent 9f5580fddd
commit 03f5717456
2 changed files with 15 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
enterprise/litellm_enterprise/proxy/hooks/managed_files.py
View File

@ -1086,7 +1086,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
self, file_id: str
) -> List[Dict[str, Any]]:
"""
Find all batches in non-terminal states that reference this file.
Find batches in non-terminal states that reference this file.
Non-terminal states: validating, in_progress, finalizing
Terminal states: completed, complete, failed, expired, cancelled
@ -1096,7 +1096,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
Returns:
List of batch objects referencing this file in non-terminal state
(limited to first 10 matches for error message display)
(max 10 for error message display)
"""
# Prepare list of file IDs to check (both unified and provider IDs)
file_ids_to_check = [file_id]
@ -1116,8 +1116,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
f"Could not get model file ID mapping for {file_id}: {e}. "
f"Will only check unified file ID."
)
MAX_BATCHES_TO_CHECK = 500
MAX_MATCHES_TO_RETURN = 10
batches = await self.prisma_client.db.litellm_managedobjecttable.find_many(
@ -1125,19 +1123,12 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints):
"file_purpose": "batch",
"status": {"in": ["validating", "in_progress", "finalizing"]},
},
take=MAX_BATCHES_TO_CHECK,
take=MAX_MATCHES_TO_RETURN,
order={"created_at": "desc"},
)
referencing_batches = []
for batch in batches:
# Early exit if we have enough matches for error message
if len(referencing_batches) >= MAX_MATCHES_TO_RETURN:
verbose_logger.debug(
f"Found {MAX_MATCHES_TO_RETURN}+ batches referencing file {file_id}, "
)
break
try:
# Parse the batch file_object to check for file references
batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object

22
tests/test_litellm/enterprise/proxy/test_file_deletion_blocking.py
View File

@ -438,16 +438,16 @@ async def test_afile_delete_calls_check_deletion_allowed():
@pytest.mark.asyncio
async def test_early_exit_after_max_matches():
async def test_database_limit_respected():
"""
Test that we stop checking batches once we find enough matches.
This is a performance optimization to avoid parsing all batches.
Test that we only fetch 10 batches from DB (not 500).
This is a performance optimization - we only fetch what we need.
"""
unified_file_id = _make_unified_file_id("file-shared")
# Create more batches than MAX_MATCHES_TO_RETURN (10)
many_batches = []
for i in range(15):
# Create exactly 10 batches (what DB will return with take=10)
ten_batches = []
for i in range(10):
batch = _make_batch_db_record(
unified_object_id=_make_unified_batch_id(f"batch-{i}"),
status="validating",
@ -457,19 +457,20 @@ async def test_early_exit_after_max_matches():
"status": "validating"
},
)
many_batches.append(batch)
ten_batches.append(batch)
# Mock will return only 10 batches (as DB would with take=10)
managed_files = _make_managed_files_instance_with_batches(
file_id=unified_file_id,
batches=many_batches,
batches=ten_batches,
)
referencing_batches = await managed_files._get_batches_referencing_file(unified_file_id)
# Should return exactly 10 (MAX_MATCHES_TO_RETURN)
# Should return all 10 that reference the file
assert len(referencing_batches) == 10
# Verify error message handles "10+" case
# Verify error message handles "10+" case (since we got exactly 10, might be more in DB)
mock_scheduler = MagicMock()
mock_scheduler.get_job.return_value = MagicMock()
@ -478,6 +479,7 @@ async def test_early_exit_after_max_matches():
await managed_files._check_file_deletion_allowed(unified_file_id)
error_detail = exc_info.value.detail
# When we get exactly 10 matches, show "10+" to indicate there might be more
assert "10+ batch(es)" in error_detail