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 47229466021..3178f57facd 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -662,6 +662,11 @@ class CheckBatchCost: @staticmethod def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]: + stored_batch: Final = CheckBatchCost._stored_batch(job) + return stored_batch.input_file_id if stored_batch is not None else None + + @staticmethod + def _stored_batch(job: "ManagedObjectRow") -> Optional["LiteLLMBatch"]: from litellm.types.utils import LiteLLMBatch file_object: Final = job.file_object @@ -674,7 +679,7 @@ class CheckBatchCost: if not isinstance(parsed_file_object, dict): return None try: - return LiteLLMBatch.model_validate(parsed_file_object).input_file_id + return LiteLLMBatch.model_validate(parsed_file_object) except Exception: return None @@ -1138,6 +1143,8 @@ class CheckBatchCost: - if not, return False - if so, return True """ + from litellm.proxy.openai_files_endpoints.common_utils import with_created_batch_endpoint + try: from litellm.integrations.prometheus import PrometheusLogger @@ -1199,7 +1206,7 @@ class CheckBatchCost: ) try: - response = await self.llm_router.aretrieve_batch( + provider_response = await self.llm_router.aretrieve_batch( model=model_id, batch_id=batch_id, litellm_metadata={ @@ -1218,6 +1225,7 @@ class CheckBatchCost: else: await self._expire_if_old_and_unfinished(job) continue + response = with_created_batch_endpoint(provider_response, self._stored_batch(job)) if self._is_completed_with_output_id_pending(response): verbose_proxy_logger.info( f"CheckBatchCost: batch {batch_id} is completed but its output file id " diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index 763737191ef..0dcad0ee9f1 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -67,6 +67,7 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( # noqa: F401 # prepare_data_with_credentials, update_batch_in_database, validate_managed_id_requirement, + with_created_batch_endpoint, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import request_tags_from_request_data from litellm.proxy.route_llm_request import raise_if_required_body_param_missing @@ -644,6 +645,7 @@ async def retrieve_batch( proxy_logging_obj=proxy_logging_obj, user_api_key_dict=user_api_key_dict, ) + stored_batch: Final = response if executed_batch and response is None: raise batch_error(404, f"No batch found with id '{batch_id}'.") @@ -788,11 +790,12 @@ async def retrieve_batch( **data, ) + retrieved_batch: Final = with_created_batch_endpoint(response, stored_batch) # FIX: Update the database with the latest state from provider await update_batch_in_database( batch_id=batch_id, unified_batch_id=unified_batch_id, - response=response, + response=retrieved_batch, managed_files_obj=managed_files_obj, prisma_client=prisma_client, verbose_proxy_logger=verbose_proxy_logger, @@ -803,7 +806,7 @@ async def retrieve_batch( ### CALL HOOKS ### - modify outgoing data response = await proxy_logging_obj.post_call_success_hook( - data=data, user_api_key_dict=user_api_key_dict, response=response + data=data, user_api_key_dict=user_api_key_dict, response=retrieved_batch ) # Fix: bug_feb14_batch_retrieve_returns_raw_input_file_id @@ -1230,20 +1233,29 @@ async def cancel_batch( managed_files_obj: Final = proxy_logging_obj.get_proxy_hook("managed_files") from litellm.proxy.proxy_server import prisma_client - await update_batch_in_database( + db_batch_object, stored_batch = await get_batch_from_database( batch_id=batch_id, unified_batch_id=unified_batch_id, - response=response, managed_files_obj=managed_files_obj, prisma_client=prisma_client, verbose_proxy_logger=verbose_proxy_logger, + ) + cancelled_batch: Final = with_created_batch_endpoint(response, stored_batch) + await update_batch_in_database( + batch_id=batch_id, + unified_batch_id=unified_batch_id, + response=cancelled_batch, + managed_files_obj=managed_files_obj, + prisma_client=prisma_client, + verbose_proxy_logger=verbose_proxy_logger, + db_batch_object=db_batch_object, operation="cancel", user_api_key_dict=user_api_key_dict, ) ### CALL HOOKS ### - modify outgoing data response = await proxy_logging_obj.post_call_success_hook( - data=data, user_api_key_dict=user_api_key_dict, response=response + data=data, user_api_key_dict=user_api_key_dict, response=cancelled_batch ) ### ALERTING ### diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index f7289445d3a..dec9574c885 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1593,6 +1593,12 @@ def batch_cost_poller_is_active() -> bool: return False +def with_created_batch_endpoint(provider_batch: "LiteLLMBatch", stored_batch: "LiteLLMBatch | None") -> "LiteLLMBatch": + if stored_batch is None or not stored_batch.endpoint or stored_batch.endpoint == provider_batch.endpoint: + return provider_batch + return provider_batch.model_copy(update={"endpoint": stored_batch.endpoint}) + + def _completed_batch_safe_to_retire(response: "LiteLLMBatch") -> bool: """Whether a "completed" batch may be retired from cost recovery. diff --git a/tests/unit/proxy/batches_endpoints/test_endpoints.py b/tests/unit/proxy/batches_endpoints/test_endpoints.py index 37d64a1fc79..353584da92c 100644 --- a/tests/unit/proxy/batches_endpoints/test_endpoints.py +++ b/tests/unit/proxy/batches_endpoints/test_endpoints.py @@ -1878,6 +1878,20 @@ async def test_retrieve__db_non_terminal_state_syncs_with_provider(retrieve_harn assert retrieve_harness.update_batch_in_db.call_count == 1 +@pytest.mark.asyncio +async def test_retrieve__provider_sync_keeps_the_endpoint_the_batch_was_created_for(retrieve_harness): + stored = make_batch(id="batch-unified-blob", status="validating").model_copy(update={"endpoint": "/v1/embeddings"}) + retrieve_harness.get_batch_from_db.return_value = (MagicMock(status="validating"), stored) + retrieve_harness.router.aretrieve_batch.return_value = make_batch(id="batch-raw-xyz", status="in_progress") + + with patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID): + resp = await call_retrieve(retrieve_harness, "batch-unified-blob") + + assert (resp.endpoint, resp.status) == ("/v1/embeddings", "in_progress") + assert retrieve_harness.update_batch_in_db.call_args.kwargs["response"].endpoint == "/v1/embeddings" + assert resp._hidden_params["unified_batch_id"] == UNIFIED_BATCH_ID + + @pytest.mark.asyncio @pytest.mark.parametrize("status", ["validating", "in_progress", "finalizing", "cancelling"]) async def test_retrieve__executed_batch_served_from_db_in_every_status(retrieve_harness, status): @@ -2632,6 +2646,23 @@ async def test_cancel__db_write_receives_caller_auth(cancel_harness): assert cancel_harness.update_batch_in_db.call_args.kwargs["user_api_key_dict"] is caller +@pytest.mark.asyncio +async def test_cancel__keeps_the_endpoint_the_batch_was_created_for(cancel_harness): + stored_row = MagicMock(status="in_progress") + stored = make_batch(id="batch-unified-blob", status="in_progress").model_copy(update={"endpoint": "/v1/embeddings"}) + cancel_harness.router.acancel_batch.return_value = make_batch(id="batch-raw-xyz", status="cancelling") + + with ( + patch.object(endpoints, "is_base64_encoded_unified_file_id", return_value=UNIFIED_BATCH_ID), + patch.object(endpoints, "get_batch_from_database", AsyncMock(return_value=(stored_row, stored))), + ): + resp = await call_cancel(cancel_harness, "batch-unified-blob") + + assert (resp.endpoint, resp.status) == ("/v1/embeddings", "cancelling") + written = cancel_harness.update_batch_in_db.call_args.kwargs + assert (written["response"].endpoint, written["db_batch_object"]) == ("/v1/embeddings", stored_row) + + @pytest.mark.asyncio async def test_cancel__unified_missing_model_id_400(cancel_harness): # unified id with no model_id segment -> get_model_id returns None -> 400. diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 5ff159f80d3..1e51482eac1 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -3784,6 +3784,20 @@ class TestMultiPodBatchCostClaim: assert row.batch_processed is True await self._run_deletion_guard(prisma, _CLAIM_OUTPUT_FILE_ID) + @pytest.mark.asyncio + @pytest.mark.parametrize("has_batch_processed_column", (True, False)) + async def test_billing_a_batch_keeps_the_endpoint_it_was_created_for(self, has_batch_processed_column: bool) -> None: + row: Final = _FakeManagedObjectRow() + row.file_object = json.dumps({**json.loads(row.file_object), "endpoint": "/v1/embeddings"}) + instance: Final = self._instance(self._prisma(row, []), self._router()) + instance._has_batch_processed_column = has_batch_processed_column + + with self._billing_patches([]): + await instance.check_batch_cost() + + assert row.status == "complete" + assert json.loads(row.file_object)["endpoint"] == "/v1/embeddings" + @pytest.mark.asyncio async def test_schema_without_batch_processed_still_bills(self): """Older schemas have no column to claim, so they keep the pre-fix behavior instead