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 bf8bc46f72..4dcabb9c58 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -13,6 +13,9 @@ if TYPE_CHECKING: from litellm.router import Router +CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost" + + class CheckBatchCost: def __init__( self, @@ -27,6 +30,25 @@ class CheckBatchCost: self.prisma_client: PrismaClient = prisma_client self.llm_router: Router = llm_router + async def _get_user_info(self, batch_id, user_id) -> dict: + """ + Look up user email and key alias by user_id for enriching the S3 callback metadata. + Returns a dict with user_api_key_user_email and user_api_key_alias (both may be None). + """ + try: + user_row = await self.prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_id} + ) + if user_row is None: + return {} + return { + "user_api_key_user_email": getattr(user_row, "user_email", None), + "user_api_key_alias": getattr(user_row, "user_alias", None), + } + except Exception as e: + verbose_proxy_logger.error(f"CheckBatchCost: could not look up user {user_id} for batch {batch_id}: {e}") + return {} + async def check_batch_cost(self): """ Check if the batch JOB has been tracked. @@ -48,10 +70,12 @@ class CheckBatchCost: get_model_id_from_unified_batch_id, ) + # Look for all batches that have not yet been processed by CheckBatchCost jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( where={ - "status": {"in": ["validating", "in_progress", "finalizing"]}, "file_purpose": "batch", + "batch_processed" : False, + "status": {"not_in": ["failed", "expired", "cancelled"]} } ) completed_jobs = [] @@ -107,6 +131,21 @@ class CheckBatchCost: f"Batch ID: {batch_id} is complete, tracking cost and usage" ) + # aretrieve_batch is called with the raw provider batch ID, so response.id + # is the raw provider value (e.g. "batch_20260223-0518.234"). We need the + # unified base64 ID in the S3 log so downstream consumers can correlate it + # back to the batch they submitted via the proxy. + # + # CheckBatchCost builds its own LiteLLMLogging object (logging_obj below) and + # calls async_success_handler(result=response) directly. That handler calls + # _build_standard_logging_payload(response, ...) which reads response.id at + # that point — so setting response.id here is sufficient. + # + # The HTTP endpoint does this substitution via the managed files hook + # (async_post_call_success_hook). CheckBatchCost bypasses that hook entirely, + # so we do it explicitly here. + response.id = job.unified_object_id + # This background job runs as default_user_id, so going through the HTTP endpoint # would trigger check_managed_file_id_access and get 403. Instead, extract the raw # provider file ID and call afile_content directly with deployment credentials. @@ -171,11 +210,21 @@ class CheckBatchCost: function_id=str(uuid.uuid4()), ) + creator_user_id = job.created_by + user_info = await self._get_user_info(batch_id, job.created_by) + logging_obj.update_environment_variables( litellm_params={ + # set the user-agent header so that S3 callback consumers can easily identify CheckBatchCost callbacks + "proxy_server_request": { + "headers": { + "user-agent": CHECK_BATCH_COST_USER_AGENT, + } + }, "metadata": { - "user_api_key_user_id": job.created_by or "default-user-id", - } + "user_api_key_user_id": creator_user_id, + **user_info, + }, }, optional_params={}, ) @@ -191,8 +240,7 @@ class CheckBatchCost: completed_jobs.append(job) if len(completed_jobs) > 0: - # mark the jobs as complete await self.prisma_client.db.litellm_managedobjecttable.update_many( where={"id": {"in": [job.id for job in completed_jobs]}}, - data={"status": "complete"}, + data={"batch_processed": True, "status": "complete"}, ) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index bda20e2f74..4fa050a84a 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -1086,11 +1086,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): self, file_id: str ) -> List[Dict[str, Any]]: """ - Find batches in non-terminal states that reference this file. - - Non-terminal states: validating, in_progress, finalizing - Terminal states: completed, complete, failed, expired, cancelled - + Find batches that reference this file and still need cost tracking. + 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 @@ -1121,7 +1118,8 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): batches = await self.prisma_client.db.litellm_managedobjecttable.find_many( where={ "file_purpose": "batch", - "status": {"in": ["validating", "in_progress", "finalizing"]}, + "batch_processed": False, + "status": {"not_in": ["failed", "expired", "cancelled"]} }, take=MAX_MATCHES_TO_RETURN, order={"created_at": "desc"}, @@ -1205,7 +1203,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): error_message += ( f"To delete this file before complete cost tracking, please delete or cancel the referencing batch(es) first. " - f"Alternatively, wait for all batches to complete processing." + f"Alternatively, wait for all batches to complete and for cost to be computed (batch_processed=true)." ) raise HTTPException( diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260222000000_add_batch_processed_to_managed_object_table/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260222000000_add_batch_processed_to_managed_object_table/migration.sql new file mode 100644 index 0000000000..ac390d164d --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260222000000_add_batch_processed_to_managed_object_table/migration.sql @@ -0,0 +1,3 @@ +-- Add batch_processed column to LiteLLM_ManagedObjectTable +-- Set to true by CheckBatchCost after cost has been computed for a completed batch +ALTER TABLE "LiteLLM_ManagedObjectTable" ADD COLUMN "batch_processed" BOOLEAN NOT NULL DEFAULT false; diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 5d2cad6da5..6d2f1e4236 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -813,6 +813,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t file_object Json // Stores the OpenAIFileObject file_purpose String // either 'batch' or 'fine-tune' status String? // check if batch cost has been tracked + batch_processed Boolean @default(false) // set to true by CheckBatchCost after cost is computed created_at DateTime @default(now()) created_by String? updated_at DateTime @updatedAt diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 5d2cad6da5..6d2f1e4236 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -813,6 +813,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t file_object Json // Stores the OpenAIFileObject file_purpose String // either 'batch' or 'fine-tune' status String? // check if batch cost has been tracked + batch_processed Boolean @default(false) // set to true by CheckBatchCost after cost is computed created_at DateTime @default(now()) created_by String? updated_at DateTime @updatedAt diff --git a/schema.prisma b/schema.prisma index 5d2cad6da5..6d2f1e4236 100644 --- a/schema.prisma +++ b/schema.prisma @@ -813,6 +813,7 @@ model LiteLLM_ManagedObjectTable { // for batches or finetuning jobs which use t file_object Json // Stores the OpenAIFileObject file_purpose String // either 'batch' or 'fine-tune' status String? // check if batch cost has been tracked + batch_processed Boolean @default(false) // set to true by CheckBatchCost after cost is computed created_at DateTime @default(now()) created_by String? updated_at DateTime @updatedAt