diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index 4dcabb9c58..10f7f98b71 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -78,8 +78,6 @@ class CheckBatchCost: "status": {"not_in": ["failed", "expired", "cancelled"]} } ) - completed_jobs = [] - for job in jobs: # get the model from the job unified_object_id = job.unified_object_id @@ -237,10 +235,16 @@ class CheckBatchCost: ) # mark the job as complete - completed_jobs.append(job) - - if len(completed_jobs) > 0: - await self.prisma_client.db.litellm_managedobjecttable.update_many( - where={"id": {"in": [job.id for job in completed_jobs]}}, - data={"batch_processed": True, "status": "complete"}, - ) + try: + await self.prisma_client.db.litellm_managedobjecttable.update( + where={"id": job.id}, + data={ + "batch_processed": True, + "status": "complete", + "file_object": response.model_dump_json(), + }, + ) + except Exception as db_err: + verbose_proxy_logger.error( + f"CheckBatchCost: failed to mark job {job.id} complete in DB: {db_err}" + ) diff --git a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py index 2db5a2214c..8cb642b7a4 100644 --- a/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py +++ b/tests/test_litellm/enterprise/proxy/test_managed_files_access_check.py @@ -134,7 +134,7 @@ async def test_check_batch_cost_should_call_afile_content_directly_with_credenti mock_prisma.db.litellm_managedobjecttable.find_many = AsyncMock( return_value=[mock_job] ) - mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock() + mock_prisma.db.litellm_managedobjecttable.update = AsyncMock() # Mock proxy_logging_obj — should NOT be called for file content mock_proxy_logging = MagicMock() @@ -198,3 +198,12 @@ async def test_check_batch_cost_should_call_afile_content_directly_with_credenti # managed_files_obj.afile_content should NOT have been called mock_managed_files_hook.afile_content.assert_not_called() + + # Verify the DB update writes batch_processed, status, and file_object + mock_prisma.db.litellm_managedobjecttable.update.assert_called_once() + update_call_kwargs = mock_prisma.db.litellm_managedobjecttable.update.call_args.kwargs + assert update_call_kwargs["data"]["batch_processed"] is True + assert update_call_kwargs["data"]["status"] == "complete" + assert "file_object" in update_call_kwargs["data"], ( + "file_object must be written to DB so list_batches reads updated status" + )