From 27a33565e78ec86d0ef3143ac3a2174c56599b2b Mon Sep 17 00:00:00 2001 From: Ephrim Stanley Date: Mon, 23 Feb 2026 07:52:58 -0500 Subject: [PATCH] State management fixes for CheckBatchCost - Address greptile comments --- .../proxy/common_utils/check_batch_cost.py | 5 +++-- enterprise/litellm_enterprise/proxy/hooks/managed_files.py | 7 ++----- 2 files changed, 5 insertions(+), 7 deletions(-) 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 fbc56d37745..4dcabb9c58b 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -74,7 +74,8 @@ class CheckBatchCost: jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( where={ "file_purpose": "batch", - "batch_processed" : False + "batch_processed" : False, + "status": {"not_in": ["failed", "expired", "cancelled"]} } ) completed_jobs = [] @@ -111,7 +112,7 @@ class CheckBatchCost: model=model_id, batch_id=batch_id, litellm_metadata={ - "user_api_key_user_id": job.created_by, + "user_api_key_user_id": job.created_by or "default-user-id", "batch_ignore_default_logging": True, }, ) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 35a44d18cfe..4fa050a84aa 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1087,11 +1087,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): ) -> List[Dict[str, Any]]: """ Find batches that reference this file and still need cost tracking. - - Blocks on: validating, in_progress, finalizing (in-flight) - Blocks on: completed AND batch_processed=False (cost not yet computed) - Allows: completed AND batch_processed=True, failed, expired, cancelled - + Find batches that are in non-terminal state and have not yet been processed by CheckBatchCost. Args: file_id: The unified file ID to check @@ -1123,6 +1119,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): where={ "file_purpose": "batch", "batch_processed": False, + "status": {"not_in": ["failed", "expired", "cancelled"]} }, take=MAX_MATCHES_TO_RETURN, order={"created_at": "desc"},