diff --git a/tests/test_litellm/enterprise/test_responses_background_cost.py b/tests/test_litellm/enterprise/test_responses_background_cost.py new file mode 100644 index 0000000000..9e82aceffd --- /dev/null +++ b/tests/test_litellm/enterprise/test_responses_background_cost.py @@ -0,0 +1,494 @@ +""" +Integration tests for responses API background cost tracking +""" + +import os +import sys +from datetime import datetime +from unittest.mock import AsyncMock, MagicMock, Mock, patch + +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) + +from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse +from litellm_enterprise.proxy.common_utils.check_responses_cost import ( + CheckResponsesCost, +) + + +class TestResponsesBackgroundCostTracking: + """Integration tests for responses API background cost tracking""" + + @pytest.fixture + def mock_managed_files_obj(self): + """Create a mock managed files object""" + managed_files = MagicMock() + managed_files.store_unified_object_id = AsyncMock() + return managed_files + + @pytest.fixture + def mock_proxy_logging_obj(self, mock_managed_files_obj): + """Create a mock proxy logging object""" + logging_obj = MagicMock() + logging_obj.get_proxy_hook = MagicMock(return_value=mock_managed_files_obj) + return logging_obj + + @pytest.fixture + def mock_llm_router(self): + """Create a mock LLM router""" + router = MagicMock() + return router + + @pytest.mark.asyncio + async def test_store_response_in_managed_objects_table( + self, mock_managed_files_obj, mock_proxy_logging_obj, mock_llm_router + ): + """Test that background responses are stored in managed objects table""" + # Create a mock response with queued status and hidden params + response = ResponsesAPIResponse( + id="resp_bGl0ZWxsbTpjdXN0b21fbGxtX3Byb3ZpZGVyOm9wZW5haTttb2RlbF9pZDpncHQtNDtsbGxfcmVzcG9uc2VfaWQ6cmVzcF8xMjM", + object="response", + status="queued", + created_at=int(datetime.now().timestamp()), + output=[], + usage=None, + ) + + # Add hidden params with model_id (simulating what base_process_llm_request does) + response._hidden_params = { + "model_id": "model-deployment-id-123" + } + + # Mock request data + data = { + "model": "gpt-4", + "input": "Test input", + "background": True, + } + + # Mock user_api_key_dict + user_api_key_dict = MagicMock() + user_api_key_dict.user_id = "test-user" + + # Simulate the storage logic from endpoints.py + if data.get("background") and isinstance(response, ResponsesAPIResponse): + if response.status in ["queued", "in_progress"]: + # Get model_id from hidden params + hidden_params = getattr(response, "_hidden_params", {}) or {} + model_id = hidden_params.get("model_id", None) + + if model_id: + # Store in managed objects table using response.id directly + await mock_managed_files_obj.store_unified_object_id( + unified_object_id=response.id, + file_object=response, + litellm_parent_otel_span=None, + model_object_id=response.id, + file_purpose="response", + user_api_key_dict=user_api_key_dict, + ) + + # Verify store_unified_object_id was called + mock_managed_files_obj.store_unified_object_id.assert_called_once() + call_args = mock_managed_files_obj.store_unified_object_id.call_args + + # Verify the arguments - unified_object_id should be response.id + assert call_args[1]["unified_object_id"] == response.id + assert call_args[1]["model_object_id"] == response.id + assert call_args[1]["file_purpose"] == "response" + assert call_args[1]["user_api_key_dict"] == user_api_key_dict + + @pytest.mark.asyncio + async def test_no_storage_for_non_background_requests( + self, mock_managed_files_obj, mock_proxy_logging_obj + ): + """Test that non-background requests are not stored""" + # Create a mock response + response = ResponsesAPIResponse( + id="resp_456", + object="response", + status="completed", + created_at=int(datetime.now().timestamp()), + output=[], + usage=ResponseAPIUsage( + input_tokens=100, + output_tokens=50, + total_tokens=150, + ), + ) + + # Mock request data without background flag + data = { + "model": "gpt-4", + "input": "Test input", + "background": False, + } + + # Simulate the storage logic + if data.get("background") and isinstance(response, ResponsesAPIResponse): + if response.status in ["queued", "in_progress"]: + await mock_managed_files_obj.store_unified_object_id() + + # Verify store_unified_object_id was NOT called + mock_managed_files_obj.store_unified_object_id.assert_not_called() + + @pytest.mark.asyncio + async def test_no_storage_for_completed_responses( + self, mock_managed_files_obj, mock_proxy_logging_obj + ): + """Test that completed responses are not stored""" + # Create a mock response with completed status + response = ResponsesAPIResponse( + id="resp_789", + object="response", + status="completed", + created_at=int(datetime.now().timestamp()), + output=[], + usage=ResponseAPIUsage( + input_tokens=100, + output_tokens=50, + total_tokens=150, + ), + ) + + # Mock request data with background flag + data = { + "model": "gpt-4", + "input": "Test input", + "background": True, + } + + # Simulate the storage logic + if data.get("background") and isinstance(response, ResponsesAPIResponse): + if response.status in ["queued", "in_progress"]: + await mock_managed_files_obj.store_unified_object_id() + + # Verify store_unified_object_id was NOT called (status is completed) + mock_managed_files_obj.store_unified_object_id.assert_not_called() + + @pytest.mark.asyncio + async def test_no_storage_without_model_id( + self, mock_managed_files_obj, mock_proxy_logging_obj + ): + """Test that responses without model_id in hidden params are not stored""" + # Create a mock response without hidden params + response = ResponsesAPIResponse( + id="resp_no_model", + object="response", + status="queued", + created_at=int(datetime.now().timestamp()), + output=[], + usage=None, + ) + + # Mock request data with background flag + data = { + "model": "gpt-4", + "input": "Test input", + "background": True, + } + + user_api_key_dict = MagicMock() + + # Simulate the storage logic + if data.get("background") and isinstance(response, ResponsesAPIResponse): + if response.status in ["queued", "in_progress"]: + hidden_params = getattr(response, "_hidden_params", {}) or {} + model_id = hidden_params.get("model_id", None) + + if model_id: # This will be False + await mock_managed_files_obj.store_unified_object_id( + unified_object_id=response.id, + file_object=response, + litellm_parent_otel_span=None, + model_object_id=response.id, + file_purpose="response", + user_api_key_dict=user_api_key_dict, + ) + + # Verify store_unified_object_id was NOT called (no model_id) + mock_managed_files_obj.store_unified_object_id.assert_not_called() + + @pytest.mark.asyncio + async def test_error_handling_in_storage( + self, mock_managed_files_obj, mock_proxy_logging_obj + ): + """Test that errors during storage are handled gracefully""" + # Mock store_unified_object_id to raise an exception + mock_managed_files_obj.store_unified_object_id = AsyncMock( + side_effect=Exception("Database error") + ) + + response = ResponsesAPIResponse( + id="resp_error", + object="response", + status="queued", + created_at=int(datetime.now().timestamp()), + output=[], + usage=None, + ) + response._hidden_params = {"model_id": "test-model-id"} + + data = { + "model": "gpt-4", + "input": "Test input", + "background": True, + } + + user_api_key_dict = MagicMock() + user_api_key_dict.user_id = "test-user" + + # Try to store - should not raise (error is caught in endpoints.py) + try: + if data.get("background") and isinstance(response, ResponsesAPIResponse): + if response.status in ["queued", "in_progress"]: + hidden_params = getattr(response, "_hidden_params", {}) or {} + model_id = hidden_params.get("model_id", None) + + if model_id: + await mock_managed_files_obj.store_unified_object_id( + unified_object_id=response.id, + file_object=response, + litellm_parent_otel_span=None, + model_object_id=response.id, + file_purpose="response", + user_api_key_dict=user_api_key_dict, + ) + except Exception: + # Exception should be caught and logged, not raised + pass + + # Verify the method was called (even though it raised) + assert mock_managed_files_obj.store_unified_object_id.called + + +class TestCheckResponsesCost: + """Tests for the CheckResponsesCost polling class""" + + @pytest.fixture + def mock_prisma_client(self): + """Create a mock Prisma client""" + client = MagicMock() + client.db = MagicMock() + client.db.litellm_managedobjecttable = MagicMock() + return client + + @pytest.fixture + def mock_proxy_logging_obj(self): + """Create a mock proxy logging object""" + return MagicMock() + + @pytest.fixture + def mock_llm_router(self): + """Create a mock LLM router""" + return MagicMock() + + @pytest.mark.asyncio + async def test_check_responses_cost_initialization( + self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router + ): + """Test CheckResponsesCost initialization""" + checker = CheckResponsesCost( + proxy_logging_obj=mock_proxy_logging_obj, + prisma_client=mock_prisma_client, + llm_router=mock_llm_router, + ) + + assert checker.proxy_logging_obj == mock_proxy_logging_obj + assert checker.prisma_client == mock_prisma_client + assert checker.llm_router == mock_llm_router + + @pytest.mark.asyncio + async def test_check_responses_cost_no_jobs( + self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router + ): + """Test polling when there are no jobs""" + # Mock find_many to return empty list + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[] + ) + + checker = CheckResponsesCost( + proxy_logging_obj=mock_proxy_logging_obj, + prisma_client=mock_prisma_client, + llm_router=mock_llm_router, + ) + + # Should not raise any errors + await checker.check_responses_cost() + + # Verify find_many was called with correct parameters + mock_prisma_client.db.litellm_managedobjecttable.find_many.assert_called_once_with( + where={ + "status": {"in": ["queued", "in_progress"]}, + "file_purpose": "response", + } + ) + + @pytest.mark.asyncio + async def test_check_responses_cost_with_completed_job( + self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router + ): + """Test polling with a completed job""" + # Create a mock job + mock_job = MagicMock() + mock_job.id = "job-123" + mock_job.unified_object_id = "resp_test_id" + mock_job.created_by = "test-user" + + # Mock find_many to return the job + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + + # Mock update_many + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock() + + # Create a completed response + completed_response = ResponsesAPIResponse( + id="resp_test_id", + object="response", + status="completed", + created_at=int(datetime.now().timestamp()), + output=[], + usage=ResponseAPIUsage( + input_tokens=100, + output_tokens=50, + total_tokens=150, + ), + ) + + checker = CheckResponsesCost( + proxy_logging_obj=mock_proxy_logging_obj, + prisma_client=mock_prisma_client, + llm_router=mock_llm_router, + ) + + # Mock litellm.aget_responses to return completed response + with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget: + mock_aget.return_value = completed_response + + await checker.check_responses_cost() + + # Verify update_many was called to mark job as completed + mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_called_once() + call_args = ( + mock_prisma_client.db.litellm_managedobjecttable.update_many.call_args + ) + assert call_args[1]["where"]["id"]["in"] == ["job-123"] + assert call_args[1]["data"]["status"] == "completed" + + @pytest.mark.asyncio + async def test_check_responses_cost_with_failed_job( + self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router + ): + """Test polling with a failed job""" + # Create a mock job + mock_job = MagicMock() + mock_job.id = "job-456" + mock_job.unified_object_id = "resp_failed" + mock_job.created_by = "test-user" + + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock() + + # Create a failed response + failed_response = ResponsesAPIResponse( + id="resp_failed", + object="response", + status="failed", + created_at=int(datetime.now().timestamp()), + output=[], + usage=None, + ) + + checker = CheckResponsesCost( + proxy_logging_obj=mock_proxy_logging_obj, + prisma_client=mock_prisma_client, + llm_router=mock_llm_router, + ) + + with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget: + mock_aget.return_value = failed_response + + await checker.check_responses_cost() + + # Verify job was marked as completed even though it failed + mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_called_once() + + @pytest.mark.asyncio + async def test_check_responses_cost_with_in_progress_job( + self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router + ): + """Test polling with a job still in progress""" + # Create a mock job + mock_job = MagicMock() + mock_job.id = "job-789" + mock_job.unified_object_id = "resp_in_progress" + mock_job.created_by = "test-user" + + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock() + + # Create an in-progress response + in_progress_response = ResponsesAPIResponse( + id="resp_in_progress", + object="response", + status="in_progress", + created_at=int(datetime.now().timestamp()), + output=[], + usage=None, + ) + + checker = CheckResponsesCost( + proxy_logging_obj=mock_proxy_logging_obj, + prisma_client=mock_prisma_client, + llm_router=mock_llm_router, + ) + + with patch("litellm.aget_responses", new_callable=AsyncMock) as mock_aget: + mock_aget.return_value = in_progress_response + + await checker.check_responses_cost() + + # Verify update_many was NOT called (job still in progress) + mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_not_called() + + @pytest.mark.asyncio + async def test_check_responses_cost_error_handling( + self, mock_proxy_logging_obj, mock_prisma_client, mock_llm_router + ): + """Test that errors when querying responses are handled gracefully""" + # Create a mock job + mock_job = MagicMock() + mock_job.id = "job-error" + mock_job.unified_object_id = "resp_error" + mock_job.created_by = "test-user" + + mock_prisma_client.db.litellm_managedobjecttable.find_many = AsyncMock( + return_value=[mock_job] + ) + mock_prisma_client.db.litellm_managedobjecttable.update_many = AsyncMock() + + checker = CheckResponsesCost( + proxy_logging_obj=mock_proxy_logging_obj, + prisma_client=mock_prisma_client, + llm_router=mock_llm_router, + ) + + # Mock litellm.aget_responses to raise an exception + with patch( + "litellm.aget_responses", + new_callable=AsyncMock, + side_effect=Exception("API error"), + ): + # Should not raise - errors are caught and logged + await checker.check_responses_cost() + + # Verify update_many was NOT called (error occurred) + mock_prisma_client.db.litellm_managedobjecttable.update_many.assert_not_called()