diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index be37f648397..a70d85cf59c 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -45,6 +45,7 @@ from litellm.proxy._types import ( UserAPIKeyAuth, ) from litellm.proxy.openai_files_endpoints.common_utils import ( + FILE_LIST_CONTINUATION_CHUNK_SIZE, MAX_FILE_LIST_LIMIT, _is_base64_encoded_unified_file_id, apply_unified_file_ids, @@ -1390,7 +1391,10 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): yield fewer matches than the page holds. Successive chunks are read until the page is full or the caller's rows run out, which keeps ``data`` non-empty while matches remain and its last id usable as the - next cursor. + next cursor. A first chunk that fills the page costs one query; once a + scan has to continue past it, the chunk widens to + ``FILE_LIST_CONTINUATION_CHUNK_SIZE`` so a page whose matches sit far + behind the newest rows cannot degenerate into thousands of queries. """ validate_file_list_limit(limit) @@ -1412,9 +1416,9 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) page_size: Final = min(limit or MAX_FILE_LIST_LIMIT, MAX_FILE_LIST_LIMIT) - chunk_size: Final = page_size + 1 matches: Final[List[OpenAIFileObject]] = [] cursor_id = after + chunk_size = page_size + 1 while len(matches) <= page_size: cursor_args: _CursorPageArgs = {"cursor": {"unified_file_id": cursor_id}, "skip": 1} if cursor_id else {} @@ -1433,6 +1437,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if len(chunk) < chunk_size: break cursor_id = chunk[-1].unified_file_id + chunk_size = max(chunk_size, FILE_LIST_CONTINUATION_CHUNK_SIZE) return build_list_page(matches[:page_size], has_more=len(matches) > page_size) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 605da435848..837b4c43652 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -25,6 +25,8 @@ if TYPE_CHECKING: MAX_FILE_LIST_LIMIT: Final = 10000 +FILE_LIST_CONTINUATION_CHUNK_SIZE: Final = 500 + def validate_file_list_limit(limit: int | None) -> None: """Reject a ``limit`` outside the range OpenAI documents for GET /v1/files.""" diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py index 40abcac1008..0b81c1d23e3 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_hook.py @@ -441,6 +441,69 @@ async def test_afile_list_fills_a_page_past_rows_that_do_not_parse(): assert page["has_more"] is False +_DEEP_SCAN_ROW_COUNT = 2000 +_DEEP_SCAN_QUERY_BUDGET = 10 + + +@pytest.mark.asyncio +async def test_afile_list_bounds_the_queries_a_deep_purpose_match_costs(): + """A tiny limit over rows the filter drops must not turn one request into thousands of queries.""" + managed_files, table = _make_managed_files_over_rows( + [_make_managed_file_row(f"unified-{index:05d}") for index in range(_DEEP_SCAN_ROW_COUNT)] + + [_make_managed_file_row("unified-match", purpose="batch")] + ) + + page = await managed_files.afile_list( + purpose="batch", + litellm_parent_otel_span=None, + user_api_key_dict=_make_user_api_key_dict(), + limit=1, + ) + + assert [file.id for file in page["data"]] == ["unified-match"] + assert page["has_more"] is False + assert len(table.find_many_calls) <= _DEEP_SCAN_QUERY_BUDGET + + +@pytest.mark.asyncio +async def test_afile_list_bounds_the_queries_a_deep_unparseable_run_costs(): + """Rows that will not parse drop out like a filter does, so they get the same bound.""" + managed_files, table = _make_managed_files_over_rows( + [_make_unparseable_managed_file_row(f"unified-{index:05d}") for index in range(_DEEP_SCAN_ROW_COUNT)] + + [_make_managed_file_row("unified-parses")] + ) + + page = await managed_files.afile_list( + purpose=None, + litellm_parent_otel_span=None, + user_api_key_dict=_make_user_api_key_dict(), + limit=1, + ) + + assert [file.id for file in page["data"]] == ["unified-parses"] + assert page["has_more"] is False + assert len(table.find_many_calls) <= _DEEP_SCAN_QUERY_BUDGET + + +@pytest.mark.asyncio +async def test_afile_list_reads_one_chunk_when_the_first_one_fills_the_page(): + """The widened chunk must stay off the common path, where the newest rows already fill the page.""" + managed_files, table = _make_managed_files_over_rows( + [_make_managed_file_row(f"unified-{index:05d}") for index in range(_DEEP_SCAN_ROW_COUNT)] + ) + + page = await managed_files.afile_list( + purpose=None, + litellm_parent_otel_span=None, + user_api_key_dict=_make_user_api_key_dict(), + limit=2, + ) + + assert [file.id for file in page["data"]] == ["unified-00000", "unified-00001"] + assert page["has_more"] is True + assert [call["take"] for call in table.find_many_calls] == [3] + + @pytest.mark.asyncio async def test_afile_list_reports_no_more_pages_when_nothing_matches(): managed_files, _ = _make_managed_files_over_rows(