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 d4ee4042b1a..bb25e4f0626 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -53,7 +53,7 @@ class CheckBatchCost: jobs = await self.prisma_client.db.litellm_managedobjecttable.find_many( where={ - "status": "validating", + "status": {"in": ["validating", "in_progress", "finalizing"]}, "file_purpose": "batch", } ) diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 5ee3372cca7..fe476c0ba2f 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -166,7 +166,11 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "updated_by": user_api_key_dict.user_id, "status": file_object.status, }, - "update": {}, # don't do anything if it already exists + "update": { + "file_object": file_object.model_dump_json(), + "status": file_object.status, + "updated_by": user_api_key_dict.user_id, + }, # FIX: Update status and file_object on every operation to keep state in sync }, ) @@ -460,8 +464,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): if retrieve_object_id else False ) - print(f"🔥potential_llm_object_id: {potential_llm_object_id}") - print(f"🔥retrieve_object_id: {retrieve_object_id}") if potential_llm_object_id and retrieve_object_id: ## VALIDATE USER HAS ACCESS TO THE OBJECT ## if not await self.can_user_call_unified_object_id( diff --git a/litellm/proxy/batches_endpoints/endpoints.py b/litellm/proxy/batches_endpoints/endpoints.py index f47e2e1667b..06800cb4524 100644 --- a/litellm/proxy/batches_endpoints/endpoints.py +++ b/litellm/proxy/batches_endpoints/endpoints.py @@ -24,10 +24,12 @@ from litellm.proxy.openai_files_endpoints.common_utils import ( _is_base64_encoded_unified_file_id, decode_model_from_file_id, encode_file_id_with_model, + get_batch_from_database, get_credentials_for_model, get_models_from_unified_file_id, get_original_file_id, prepare_data_with_credentials, + update_batch_in_database, ) from litellm.proxy.utils import handle_exception_on_proxy, is_known_model from litellm.types.llms.openai import LiteLLMBatchCreateRequest @@ -357,6 +359,57 @@ async def retrieve_batch( route_type="aretrieve_batch", ) + # FIX: First, try to read from ManagedObjectTable for consistent state + managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") + from litellm.proxy.proxy_server import prisma_client + + db_batch_object, response = await get_batch_from_database( + batch_id=batch_id, + unified_batch_id=unified_batch_id, + managed_files_obj=managed_files_obj, + prisma_client=prisma_client, + verbose_proxy_logger=verbose_proxy_logger, + ) + + # If batch is in a terminal state, return immediately + if response is not None and response.status in ["completed", "failed", "cancelled", "expired"]: + # Call hooks and return + response = await proxy_logging_obj.post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + + asyncio.create_task( + proxy_logging_obj.update_request_status( + litellm_call_id=data.get("litellm_call_id", ""), status="success" + ) + ) + + hidden_params = getattr(response, "_hidden_params", {}) or {} + model_id = hidden_params.get("model_id", None) or "" + cache_key = hidden_params.get("cache_key", None) or "" + api_base = hidden_params.get("api_base", None) or "" + + fastapi_response.headers.update( + ProxyBaseLLMRequestProcessing.get_custom_headers( + user_api_key_dict=user_api_key_dict, + model_id=model_id, + cache_key=cache_key, + api_base=api_base, + version=version, + model_region=getattr(user_api_key_dict, "allowed_model_region", ""), + request_data=data, + ) + ) + + return response + + # If batch is still processing, sync with provider to get latest state + if response is not None: + verbose_proxy_logger.debug( + f"Batch {batch_id} is in non-terminal state {response.status}, syncing with provider" + ) + + # Retrieve from provider (for non-terminal states or if DB lookup failed) # SCENARIO 1: Batch ID is encoded with model info if model_from_id is not None: credentials = get_credentials_for_model( @@ -408,6 +461,18 @@ async def retrieve_batch( response = await litellm.aretrieve_batch( custom_llm_provider=custom_llm_provider, **data # type: ignore ) + + # 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, + managed_files_obj=managed_files_obj, + prisma_client=prisma_client, + verbose_proxy_logger=verbose_proxy_logger, + db_batch_object=db_batch_object, + operation="retrieve", + ) ### CALL HOOKS ### - modify outgoing data response = await proxy_logging_obj.post_call_success_hook( @@ -769,6 +834,20 @@ async def cancel_batch( **_cancel_batch_data, ) + # FIX: Update the database with the new cancelled state + managed_files_obj = proxy_logging_obj.get_proxy_hook("managed_files") + from litellm.proxy.proxy_server import prisma_client + + await update_batch_in_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, + operation="cancel", + ) + ### 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 diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 2ff1183579f..f67dc5e2aaa 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -637,3 +637,127 @@ def _extract_model_param(request: "Request", request_body: dict) -> Optional[str or request.query_params.get("model") or request.headers.get("x-litellm-model") ) + + +# ============================================================================ +# BATCH DATABASE OPERATIONS +# ============================================================================ + + +async def get_batch_from_database( + batch_id: str, + unified_batch_id: Union[str, Literal[False]], + managed_files_obj, + prisma_client, + verbose_proxy_logger, +): + """ + Try to retrieve batch object from ManagedObjectTable for consistent state. + + Args: + batch_id: The batch ID (may be unified/encoded) + unified_batch_id: Result from _is_base64_encoded_unified_file_id() + managed_files_obj: The managed_files proxy hook object + prisma_client: Prisma database client + verbose_proxy_logger: Logger instance + + Returns: + Tuple of (db_batch_object, response_batch) + - db_batch_object: Raw database object (or None) + - response_batch: Parsed LiteLLMBatch object (or None) + """ + import json + from litellm.types.utils import LiteLLMBatch + + if managed_files_obj is None or not unified_batch_id: + return None, None + + try: + if not prisma_client: + return None, None + + db_batch_object = await prisma_client.db.litellm_managedobjecttable.find_first( + where={"unified_object_id": batch_id} + ) + + if not db_batch_object or not db_batch_object.file_object: + return None, None + + # Parse the batch object from database + batch_data = json.loads(db_batch_object.file_object) if isinstance(db_batch_object.file_object, str) else db_batch_object.file_object + response = LiteLLMBatch(**batch_data) + response.id = batch_id + + verbose_proxy_logger.debug( + f"Retrieved batch {batch_id} from ManagedObjectTable with status={response.status}" + ) + + return db_batch_object, response + + except Exception as e: + verbose_proxy_logger.warning( + f"Failed to retrieve batch from ManagedObjectTable: {e}, falling back to provider" + ) + return None, None + + +async def update_batch_in_database( + batch_id: str, + unified_batch_id: Union[str, Literal[False]], + response, + managed_files_obj, + prisma_client, + verbose_proxy_logger, + db_batch_object=None, + operation: str = "update", +): + """ + Update batch status and object in ManagedObjectTable. + + Args: + batch_id: The batch ID (unified/encoded) + unified_batch_id: Result from _is_base64_encoded_unified_file_id() + response: The batch response object with updated state + managed_files_obj: The managed_files proxy hook object + prisma_client: Prisma database client + verbose_proxy_logger: Logger instance + db_batch_object: Optional existing database object (for comparison) + operation: Description of operation ("update", "cancel", etc.) + """ + import litellm.utils + + if managed_files_obj is None or not unified_batch_id: + return + + try: + if not prisma_client: + return + + # Only update if status has changed (when db_batch_object is provided) + if db_batch_object and response.status == db_batch_object.status: + return + + if db_batch_object: + verbose_proxy_logger.info( + f"Updating batch {batch_id} status from {db_batch_object.status} to {response.status}" + ) + else: + verbose_proxy_logger.info( + f"Updating batch {batch_id} status to {response.status} after {operation}" + ) + + # Normalize status for database storage + db_status = response.status if response.status != "completed" else "complete" + + await prisma_client.db.litellm_managedobjecttable.update( + where={"unified_object_id": batch_id}, + data={ + "status": db_status, + "file_object": response.model_dump_json(), + "updated_at": litellm.utils.get_utc_datetime(), + }, + ) + except Exception as e: + verbose_proxy_logger.error( + f"Failed to update batch status in ManagedObjectTable: {e}" + ) diff --git a/tests/openai_endpoints_tests/test_openai_batches_endpoint.py b/tests/openai_endpoints_tests/test_openai_batches_endpoint.py index ecc0e3b370f..215ac0874f2 100644 --- a/tests/openai_endpoints_tests/test_openai_batches_endpoint.py +++ b/tests/openai_endpoints_tests/test_openai_batches_endpoint.py @@ -291,4 +291,318 @@ async def test_list_batches_with_target_model_names(): # Verify the response structure assert response["object"] == "list" - assert len(response["data"]) > 0 \ No newline at end of file + assert len(response["data"]) > 0 + + +@pytest.mark.asyncio +async def test_batch_status_sync_from_provider_to_database(): + """ + Test that when batch status changes at the provider, + it gets synced to the ManagedObjectTable database. + + This tests the new refactored utility functions: + - get_batch_from_database() + - update_batch_in_database() + """ + from unittest.mock import MagicMock, AsyncMock + from litellm.proxy.openai_files_endpoints.common_utils import ( + get_batch_from_database, + update_batch_in_database, + ) + from litellm.types.utils import LiteLLMBatch + import json + + # Setup: Create mock objects + batch_id = "batch_test123" + unified_batch_id = "litellm_proxy:test_unified_batch" + + # Mock database batch object with "validating" status + mock_db_batch = MagicMock() + mock_db_batch.unified_object_id = batch_id + mock_db_batch.status = "validating" + mock_db_batch.file_object = json.dumps({ + "id": batch_id, + "object": "batch", + "status": "validating", + "endpoint": "/v1/chat/completions", + "input_file_id": "file-test123", + "completion_window": "24h", + "created_at": 1234567890, + }) + + # Mock prisma client + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( + return_value=mock_db_batch + ) + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + + # Mock managed_files_obj + mock_managed_files = MagicMock() + + # Mock logger + mock_logger = MagicMock() + mock_logger.debug = MagicMock() + mock_logger.info = MagicMock() + mock_logger.warning = MagicMock() + mock_logger.error = MagicMock() + + # Test 1: Retrieve batch from database (initial state) + db_batch_object, response_batch = await get_batch_from_database( + batch_id=batch_id, + unified_batch_id=unified_batch_id, + managed_files_obj=mock_managed_files, + prisma_client=mock_prisma_client, + verbose_proxy_logger=mock_logger, + ) + + # Verify database was queried + mock_prisma_client.db.litellm_managedobjecttable.find_first.assert_called_once_with( + where={"unified_object_id": batch_id} + ) + + # Verify batch was retrieved correctly + assert db_batch_object is not None + assert response_batch is not None + assert response_batch.id == batch_id + assert response_batch.status == "validating" + + # Test 2: Simulate provider returning updated status + updated_batch_response = LiteLLMBatch( + id=batch_id, + object="batch", + status="completed", # Status changed from "validating" to "completed" + endpoint="/v1/chat/completions", + input_file_id="file-test123", + completion_window="24h", + created_at=1234567890, + output_file_id="file-output123", + ) + + # Test 3: Update database with new status from provider + await update_batch_in_database( + batch_id=batch_id, + unified_batch_id=unified_batch_id, + response=updated_batch_response, + managed_files_obj=mock_managed_files, + prisma_client=mock_prisma_client, + verbose_proxy_logger=mock_logger, + db_batch_object=db_batch_object, + operation="retrieve", + ) + + # Verify database was updated + mock_prisma_client.db.litellm_managedobjecttable.update.assert_called_once() + update_call_args = mock_prisma_client.db.litellm_managedobjecttable.update.call_args + + # Verify the update call had correct parameters + assert update_call_args.kwargs["where"]["unified_object_id"] == batch_id + assert update_call_args.kwargs["data"]["status"] == "complete" # "completed" normalized to "complete" + assert "file_object" in update_call_args.kwargs["data"] + assert "updated_at" in update_call_args.kwargs["data"] + + # Verify logger was called with status change message + mock_logger.info.assert_called() + log_message = mock_logger.info.call_args[0][0] + assert "validating" in log_message + assert "completed" in log_message + + print("✅ Test passed: Batch status synced from provider to database") + + +@pytest.mark.asyncio +async def test_batch_cancel_updates_database(): + """ + Test that canceling a batch updates the database status. + """ + from unittest.mock import MagicMock, AsyncMock + from litellm.proxy.openai_files_endpoints.common_utils import ( + update_batch_in_database, + ) + from litellm.types.utils import LiteLLMBatch + + # Setup + batch_id = "batch_cancel_test" + unified_batch_id = "litellm_proxy:cancel_test" + + # Mock cancelled batch response from provider + cancelled_batch_response = LiteLLMBatch( + id=batch_id, + object="batch", + status="cancelled", + endpoint="/v1/chat/completions", + input_file_id="file-test123", + completion_window="24h", + created_at=1234567890, + cancelled_at=1234567999, + ) + + # Mock prisma client + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + + # Mock managed_files_obj + mock_managed_files = MagicMock() + + # Mock logger + mock_logger = MagicMock() + mock_logger.info = MagicMock() + mock_logger.error = MagicMock() + + # Call update_batch_in_database for cancel operation + await update_batch_in_database( + batch_id=batch_id, + unified_batch_id=unified_batch_id, + response=cancelled_batch_response, + managed_files_obj=mock_managed_files, + prisma_client=mock_prisma_client, + verbose_proxy_logger=mock_logger, + operation="cancel", + ) + + # Verify database was updated + mock_prisma_client.db.litellm_managedobjecttable.update.assert_called_once() + update_call_args = mock_prisma_client.db.litellm_managedobjecttable.update.call_args + + # Verify the update call had correct parameters + assert update_call_args.kwargs["where"]["unified_object_id"] == batch_id + assert update_call_args.kwargs["data"]["status"] == "cancelled" + assert "file_object" in update_call_args.kwargs["data"] + + # Verify logger was called + mock_logger.info.assert_called() + log_message = mock_logger.info.call_args[0][0] + assert "cancel" in log_message.lower() + assert "cancelled" in log_message + + print("✅ Test passed: Batch cancel updates database") + + +@pytest.mark.asyncio +async def test_batch_terminal_state_skip_provider_call(): + """ + Test that when a batch is in a terminal state (completed, failed, cancelled, expired), + it returns immediately from database without calling the provider. + """ + from unittest.mock import MagicMock, AsyncMock + from litellm.proxy.openai_files_endpoints.common_utils import ( + get_batch_from_database, + ) + from litellm.types.utils import LiteLLMBatch + import json + + # Setup: Create mock objects for a completed batch + batch_id = "batch_completed_test" + unified_batch_id = "litellm_proxy:completed_test" + + # Mock database batch object with "completed" status + mock_db_batch = MagicMock() + mock_db_batch.unified_object_id = batch_id + mock_db_batch.status = "complete" + mock_db_batch.file_object = json.dumps({ + "id": batch_id, + "object": "batch", + "status": "completed", + "endpoint": "/v1/chat/completions", + "input_file_id": "file-test123", + "output_file_id": "file-output123", + "completion_window": "24h", + "created_at": 1234567890, + "completed_at": 1234567999, + }) + + # Mock prisma client + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_managedobjecttable.find_first = AsyncMock( + return_value=mock_db_batch + ) + + # Mock managed_files_obj + mock_managed_files = MagicMock() + + # Mock logger + mock_logger = MagicMock() + mock_logger.debug = MagicMock() + + # Retrieve batch from database + db_batch_object, response_batch = await get_batch_from_database( + batch_id=batch_id, + unified_batch_id=unified_batch_id, + managed_files_obj=mock_managed_files, + prisma_client=mock_prisma_client, + verbose_proxy_logger=mock_logger, + ) + + # Verify batch was retrieved + assert db_batch_object is not None + assert response_batch is not None + assert response_batch.status == "completed" + + # In the actual endpoint, when status is in terminal states, + # it should return immediately without calling the provider + # This test verifies the database retrieval works correctly + assert response_batch.status in ["completed", "failed", "cancelled", "expired"] + + print("✅ Test passed: Terminal state batch retrieved from database") + + +@pytest.mark.asyncio +async def test_batch_no_status_change_skip_update(): + """ + Test that when batch status hasn't changed, database update is skipped. + """ + from unittest.mock import MagicMock, AsyncMock + from litellm.proxy.openai_files_endpoints.common_utils import ( + update_batch_in_database, + ) + from litellm.types.utils import LiteLLMBatch + + # Setup + batch_id = "batch_no_change_test" + unified_batch_id = "litellm_proxy:no_change_test" + + # Mock database batch object with "validating" status + mock_db_batch = MagicMock() + mock_db_batch.status = "validating" + + # Mock batch response from provider with same status + batch_response = LiteLLMBatch( + id=batch_id, + object="batch", + status="validating", # Same status as in database + endpoint="/v1/chat/completions", + input_file_id="file-test123", + completion_window="24h", + created_at=1234567890, + ) + + # Mock prisma client + mock_prisma_client = MagicMock() + mock_prisma_client.db.litellm_managedobjecttable.update = AsyncMock() + + # Mock managed_files_obj + mock_managed_files = MagicMock() + + # Mock logger + mock_logger = MagicMock() + mock_logger.info = MagicMock() + + # Call update_batch_in_database + await update_batch_in_database( + batch_id=batch_id, + unified_batch_id=unified_batch_id, + response=batch_response, + managed_files_obj=mock_managed_files, + prisma_client=mock_prisma_client, + verbose_proxy_logger=mock_logger, + db_batch_object=mock_db_batch, + operation="retrieve", + ) + + # Verify database update was NOT called (status hasn't changed) + mock_prisma_client.db.litellm_managedobjecttable.update.assert_not_called() + + # Verify logger info was NOT called (no status change to log) + mock_logger.info.assert_not_called() + + print("✅ Test passed: Database update skipped when status unchanged") \ No newline at end of file