From 88280d9cca4b7fbd28fa25cf24881a639d2990af Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Mon, 26 Jan 2026 12:02:35 -0500 Subject: [PATCH] Fix batch creation to return the input file's expires_at attribute --- .../proxy/hooks/managed_files.py | 40 ++++++----------- .../proxy/hooks/test_managed_files.py | 43 ++++++++++++++++++- 2 files changed, 55 insertions(+), 28 deletions(-) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 776706bdd1b..dd0613f87ff 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -251,14 +251,22 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): after: Optional[str] = None, provider: Optional[str] = None, target_model_names: Optional[str] = None, + llm_router: Optional[Router] = None, ) -> Dict[str, Any]: # Provider filtering is not supported for managed batches # This is because the encoded object ids stored in the managed objects table do not contain the provider information # To support provider filtering, we would need to store the provider information in the encoded object ids if provider: raise Exception( - "Filtering by 'provider' is not supported when using managed batches. " - "Use 'target_model_names' to filter by specific model names instead." + "Filtering by 'provider' is not supported when using managed batches." + ) + + # Model name filtering is not supported for managed batches + # This is because the encoded object ids stored in the managed objects table do not contain the model name + # A hash of the model name + litellm_params for the model name is encoded as the model id. This is not sufficient to reliably map the target model names to the model ids. + if target_model_names: + raise Exception( + "Filtering by 'target_model_names' is not supported when using managed batches." ) where_clause: Dict[str, Any] = {"file_purpose": "batch"} @@ -281,12 +289,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): take=fetch_limit, order={"created_at": "desc"}, ) - - # Parse target_model_names filter - target_models_filter: List[str] = [] - if target_model_names: - target_models_filter = [m.strip() for m in target_model_names.split(",") if m.strip()] - + batch_objects: List[LiteLLMBatch] = [] for batch in batches: try: @@ -297,26 +300,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): batch_data = json.loads(batch.file_object) if isinstance(batch.file_object, str) else batch.file_object batch_obj = LiteLLMBatch(**batch_data) batch_obj.id = batch.unified_object_id + batch_objects.append(batch_obj) - # If no target_model_names filter, add the batch to the list - if not target_models_filter: - batch_objects.append(batch_obj) - continue - - # Filter by target_model_names - decoded_id = _is_base64_encoded_unified_file_id(batch.unified_object_id) - model_id = None - if decoded_id: - model_id = get_model_id_from_unified_batch_id(decoded_id) - - # Skip batches without decodable IDs if filtering is requested - if not model_id: - continue - - if any(target.lower() in model_id.lower() for target in target_models_filter): - batch_objects.append(batch_obj) - continue - except Exception as e: verbose_logger.warning( f"Failed to parse batch object {batch.unified_object_id}: {e}" @@ -760,6 +745,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): bytes=file_objects[0].bytes, filename=file_objects[0].filename, status="uploaded", + expires_at=file_objects[0].expires_at, ) return response diff --git a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py index e70e0640aa6..4fa16066e4b 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -1044,4 +1044,45 @@ async def test_list_batches_from_managed_objects_table_filters_by_created_by(): where={"file_purpose": "batch", "created_by": "user2"}, take=10, order={"created_at": "desc"}, - ) \ No newline at end of file + ) + + +@pytest.mark.asyncio +async def test_return_unified_file_id_includes_expires_at(): + from litellm.types.llms.openai import OpenAIFileObject + + # Create a mock file object with expires_at set + file_object = OpenAIFileObject( + id="file-abc123", + object="file", + bytes=1234, + created_at=1234567890, + filename="test.jsonl", + purpose="batch", + status="uploaded", + expires_at=1234657890, + ) + file_object._hidden_params = {"model_id": "test-model-id"} + + create_file_request = { + "file": ("test.jsonl", b"test content", "application/jsonl"), + "purpose": "batch", + } + + internal_usage_cache = MagicMock() + + result = await _PROXY_LiteLLMManagedFiles.return_unified_file_id( + file_objects=[file_object], + create_file_request=create_file_request, + internal_usage_cache=internal_usage_cache, + litellm_parent_otel_span=None, + target_model_names_list=["gpt-4o"], + ) + + # Verify expires_at is passed through + assert result.expires_at == 1234657890 + assert result.purpose == "batch" + assert result.filename == "test.jsonl" + assert result.bytes == 1234 + assert result.created_at == 1234567890 + assert _is_base64_encoded_unified_file_id(result.id) \ No newline at end of file