State management fixes for CheckBatchCost - Address greptile comments

This commit is contained in:
Ephrim Stanley 2026-02-23 07:52:58 -05:00
parent 7b5dc3fb9c
commit 27a33565e7
2 changed files with 5 additions and 7 deletions

View file

@ -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,
},
)

View file

@ -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"},