diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 445d2b242b4..776706bdd1b 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -244,6 +244,93 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): return managed_object.created_by == user_id return True # don't raise error if managed object is not found + async def list_user_batches( + self, + user_api_key_dict: UserAPIKeyAuth, + limit: Optional[int] = None, + after: Optional[str] = None, + provider: Optional[str] = None, + target_model_names: Optional[str] = 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." + ) + + where_clause: Dict[str, Any] = {"file_purpose": "batch"} + + # Filter by user who created the batch + if user_api_key_dict.user_id: + where_clause["created_by"] = user_api_key_dict.user_id + + if after: + where_clause["id"] = {"gt": after} + + # Fetch more than needed to allow for post-fetch filtering + fetch_limit = limit or 20 + if target_model_names: + # Fetch extra to account for filtering + fetch_limit = max(fetch_limit * 3, 100) + + batches = await self.prisma_client.db.litellm_managedobjecttable.find_many( + where=where_clause, + 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: + # Stop once we have enough after filtering + if len(batch_objects) >= (limit or 20): + break + + 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 + + # 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}" + ) + continue + + return { + "object": "list", + "data": batch_objects, + "first_id": batch_objects[0].id if batch_objects else None, + "last_id": batch_objects[-1].id if batch_objects else None, + "has_more": len(batch_objects) == (limit or 20), + } + async def get_user_created_file_ids( self, user_api_key_dict: UserAPIKeyAuth, model_object_ids: List[str] ) -> List[OpenAIFileObject]: diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 086105042e8..078e21f9bb4 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -542,14 +542,26 @@ async def list_batches( route_type="alist_batches", ) - model_param = ( + # Try to use managed objects table for listing batches (returns encoded IDs) + managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") + if managed_files_obj is not None and hasattr(managed_files_obj, "list_user_batches"): + verbose_proxy_logger.debug( + "Using managed objects table for batch listing" + ) + response = await managed_files_obj.list_user_batches( + user_api_key_dict=user_api_key_dict, + limit=limit, + after=after, + provider=provider, + target_model_names=target_model_names, + llm_router=llm_router, + ) + elif (model_param := ( data.get("model") or request.query_params.get("model") or request.headers.get("x-litellm-model") - ) - - # SCENARIO 2: Use model-based routing from header/query/body - if model_param: + )): + # SCENARIO 2: Use model-based routing from header/query/body credentials = get_credentials_for_model( llm_router=llm_router, model_id=model_param, 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 9a6e153a22b..e70e0640aa6 100644 --- a/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py +++ b/tests/enterprise/litellm_enterprise/proxy/hooks/test_managed_files.py @@ -827,3 +827,221 @@ async def test_afile_retrieve_raises_error_for_non_managed_file(): ) assert "not found" in str(exc_info.value) + + +@pytest.mark.asyncio +async def test_list_batches_from_managed_objects_table(): + from litellm.proxy._types import UserAPIKeyAuth + from openai.types.batch import BatchRequestCounts + + prisma_client = AsyncMock() + + batch_record_1 = MagicMock() + batch_record_1.unified_object_id = "unified-batch-id-1" + batch_record_1.file_object = json.dumps({ + "id": "batch_abc123", + "object": "batch", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "status": "completed", + "created_at": 1234567890, + "input_file_id": "file-input-1", + "request_counts": {"total": 1, "completed": 1, "failed": 0}, + }) + + batch_record_2 = MagicMock() + batch_record_2.unified_object_id = "unified-batch-id-2" + batch_record_2.file_object = json.dumps({ + "id": "batch_xyz789", + "object": "batch", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "status": "in_progress", + "created_at": 1234567891, + "input_file_id": "file-input-2", + "request_counts": {"total": 5, "completed": 2, "failed": 0}, + }) + + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [ + batch_record_1, + batch_record_2, + ] + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), + limit=10, + ) + + assert result["object"] == "list" + assert len(result["data"]) == 2 + assert result["data"][0].id == "unified-batch-id-1" + assert result["data"][1].id == "unified-batch-id-2" + assert result["first_id"] == "unified-batch-id-1" + assert result["last_id"] == "unified-batch-id-2" + + # Should filter by user_id (created_by) + prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with( + where={"file_purpose": "batch", "created_by": "test-user"}, + take=10, + order={"created_at": "desc"}, + ) + + +@pytest.mark.asyncio +async def test_list_batches_from_managed_objects_table_empty_list(): + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [] + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + result = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), + ) + + assert result["object"] == "list" + assert len(result["data"]) == 0 + assert result["first_id"] is None + assert result["last_id"] is None + assert result["has_more"] is False + + # Verify where clause includes created_by filter + # Default take is 20 when no limit is provided + prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with( + where={"file_purpose": "batch", "created_by": "test-user"}, + take=20, + order={"created_at": "desc"}, + ) + + +def _create_unified_batch_id(model_id: str, batch_id: str) -> str: + import base64 + unified_str = f"litellm_proxy;model_id:{model_id};llm_batch_id:{batch_id}" + return base64.urlsafe_b64encode(unified_str.encode()).decode().rstrip("=") + + +@pytest.mark.asyncio +async def test_list_batches_from_managed_objects_table_provider_filter_raises_exception(): + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + # Filtering by provider should raise Exception + with pytest.raises(Exception) as exc_info: + await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), + limit=10, + provider="openai", + ) + + assert str(exc_info.value) == ( + "Filtering by 'provider' is not supported when using managed batches." + ) + + # Verify find_many was NOT called since exception is raised before database query + prisma_client.db.litellm_managedobjecttable.find_many.assert_not_called() + + +@pytest.mark.asyncio +async def test_list_batches_from_managed_objects_table_target_model_name_filter_raises_exception(): + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + # Filtering by provider should raise Exception + with pytest.raises(Exception) as exc_info: + await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="test-user"), + limit=10, + target_model_names="gpt-4o,gpt-3.5", + ) + + assert str(exc_info.value) == ( + "Filtering by 'target_model_names' is not supported when using managed batches." + ) + + # Verify find_many was NOT called since exception is raised before database query + prisma_client.db.litellm_managedobjecttable.find_many.assert_not_called() + +@pytest.mark.asyncio +async def test_list_batches_from_managed_objects_table_filters_by_created_by(): + from litellm.proxy._types import UserAPIKeyAuth + + prisma_client = AsyncMock() + + # Create batch for user1 + batch_user1 = MagicMock() + batch_user1.unified_object_id = "unified-batch-user1" + batch_user1.file_object = json.dumps({ + "id": "batch_user1_abc", + "object": "batch", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "status": "completed", + "created_at": 1234567890, + "input_file_id": "file-input-user1", + "request_counts": {"total": 1, "completed": 1, "failed": 0}, + }) + + # Create batch for user2 + batch_user2 = MagicMock() + batch_user2.unified_object_id = "unified-batch-user2" + batch_user2.file_object = json.dumps({ + "id": "batch_user2_xyz", + "object": "batch", + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "status": "completed", + "created_at": 1234567891, + "input_file_id": "file-input-user2", + "request_counts": {"total": 2, "completed": 2, "failed": 0}, + }) + + proxy_managed_files = _PROXY_LiteLLMManagedFiles( + DualCache(), prisma_client=prisma_client + ) + + # Query with user1's API key - should only return user1's batch + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [batch_user1] + result_user1 = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="user1"), + limit=10, + ) + + assert len(result_user1["data"]) == 1 + assert result_user1["data"][0].id == "unified-batch-user1" + prisma_client.db.litellm_managedobjecttable.find_many.assert_called_with( + where={"file_purpose": "batch", "created_by": "user1"}, + take=10, + order={"created_at": "desc"}, + ) + + # Query with user2's API key - should only return user2's batch + prisma_client.db.litellm_managedobjecttable.find_many.return_value = [batch_user2] + result_user2 = await proxy_managed_files.list_user_batches( + user_api_key_dict=UserAPIKeyAuth(user_id="user2"), + limit=10, + ) + + assert len(result_user2["data"]) == 1 + assert result_user2["data"][0].id == "unified-batch-user2" + prisma_client.db.litellm_managedobjecttable.find_many.assert_called_with( + where={"file_purpose": "batch", "created_by": "user2"}, + take=10, + order={"created_at": "desc"}, + ) \ No newline at end of file