Merge pull request #14534 from hula-la/fix/response-api-bugfixes

fix: improve response api handling and cold storage configuration

Author: krrishdholakia

litellm/litellm_core_utils/litellm_logging.py

@@ -4150,15 +4150,28 @@ def _generate_cold_storage_object_key(
         from litellm.integrations.s3 import get_s3_object_key
 
         # Only generate object key if cold storage is configured
-        if litellm.configured_cold_storage_logger is None:
+        configured_cold_storage_logger = litellm.configured_cold_storage_logger
+        if configured_cold_storage_logger is None:
             return None
 
         try:
             # Generate file name in same format as litellm.utils.get_logging_id
             s3_file_name = f"time-{start_time.strftime('%H-%M-%S-%f')}_{response_id}"
 
+            # Get the actual s3_path from the configured cold storage logger instance
+            s3_path = ""  # default value
+            
+            # Try to get the actual logger instance from the logger name
+            try:
+                custom_logger = litellm.logging_callback_manager.get_active_custom_logger_for_callback_name(configured_cold_storage_logger)
+                if custom_logger and hasattr(custom_logger, 's3_path') and custom_logger.s3_path:
+                    s3_path = custom_logger.s3_path
+            except Exception:
+                # If any error occurs in getting the logger instance, use default empty s3_path
+                pass
+            
             s3_object_key = get_s3_object_key(
-                s3_path="",  # Use empty path as default
+                s3_path=s3_path,  # Use actual s3_path from logger configuration
                 team_alias_prefix="",  # Don't split by team alias for cold storage
                 start_time=start_time,
                 s3_file_name=s3_file_name,

litellm/responses/litellm_completion_transformation/session_handler.py

@@ -141,7 +141,7 @@ async def extend_chat_completion_message_with_spend_log_payload(
         # Add Output messages for this Spend Log
         ############################################################
         _response_output = spend_log.get("response", "{}")
-        if isinstance(_response_output, dict):
+        if isinstance(_response_output, dict) and _response_output and _response_output != {}:
             # transform `ChatCompletion Response` to `ResponsesAPIResponse`
             model_response = ModelResponse(**_response_output)
             for choice in model_response.choices:

tests/test_litellm/litellm_core_utils/test_litellm_logging.py

@@ -438,6 +438,97 @@ async def test_e2e_generate_cold_storage_object_key_successful():
         assert isinstance(result, str)
 
 
+@pytest.mark.asyncio
+async def test_e2e_generate_cold_storage_object_key_with_custom_logger_s3_path():
+    """
+    Test that _generate_cold_storage_object_key uses s3_path from custom logger instance.
+    """
+    from datetime import datetime, timezone
+    from unittest.mock import MagicMock, patch
+
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    # Create test data
+    start_time = datetime(2025, 1, 15, 10, 30, 45, 123456, timezone.utc)
+    response_id = "chatcmpl-test-12345"
+    
+    # Create mock custom logger with s3_path
+    mock_custom_logger = MagicMock()
+    mock_custom_logger.s3_path = "storage"
+    
+    with patch("litellm.configured_cold_storage_logger", "s3_v2"), \
+         patch("litellm.logging_callback_manager.get_active_custom_logger_for_callback_name") as mock_get_logger, \
+         patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key:
+        
+        # Setup mocks
+        mock_get_logger.return_value = mock_custom_logger
+        mock_get_s3_key.return_value = "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
+        
+        # Call the function
+        result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
+            start_time=start_time,
+            response_id=response_id
+        )
+        
+        # Verify logger was queried correctly
+        mock_get_logger.assert_called_once_with("s3_v2")
+        
+        # Verify the S3 function was called with the custom logger's s3_path
+        mock_get_s3_key.assert_called_once_with(
+            s3_path="storage",  # Should use custom logger's s3_path
+            team_alias_prefix="",
+            start_time=start_time,
+            s3_file_name="time-10-30-45-123456_chatcmpl-test-12345"
+        )
+        
+        # Verify the result
+        assert result == "storage/2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
+
+
+@pytest.mark.asyncio
+async def test_e2e_generate_cold_storage_object_key_with_logger_no_s3_path():
+    """
+    Test that _generate_cold_storage_object_key falls back to empty s3_path when logger has no s3_path.
+    """
+    from datetime import datetime, timezone
+    from unittest.mock import MagicMock, patch
+
+    from litellm.litellm_core_utils.litellm_logging import StandardLoggingPayloadSetup
+
+    # Create test data
+    start_time = datetime(2025, 1, 15, 10, 30, 45, 123456, timezone.utc)
+    response_id = "chatcmpl-test-12345"
+    
+    # Create mock custom logger without s3_path
+    mock_custom_logger = MagicMock()
+    mock_custom_logger.s3_path = None  # or could be missing attribute
+    
+    with patch("litellm.configured_cold_storage_logger", "s3_v2"), \
+         patch("litellm.logging_callback_manager.get_active_custom_logger_for_callback_name") as mock_get_logger, \
+         patch("litellm.integrations.s3.get_s3_object_key") as mock_get_s3_key:
+        
+        # Setup mocks
+        mock_get_logger.return_value = mock_custom_logger
+        mock_get_s3_key.return_value = "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
+        
+        # Call the function
+        result = StandardLoggingPayloadSetup._generate_cold_storage_object_key(
+            start_time=start_time,
+            response_id=response_id
+        )
+        
+        # Verify the S3 function was called with empty s3_path (fallback)
+        mock_get_s3_key.assert_called_once_with(
+            s3_path="",  # Should fall back to empty string
+            team_alias_prefix="",
+            start_time=start_time,
+            s3_file_name="time-10-30-45-123456_chatcmpl-test-12345"
+        )
+        
+        # Verify the result
+        assert result == "2025-01-15/time-10-30-45-123456_chatcmpl-test-12345.json"
+
+
 @pytest.mark.asyncio
 async def test_e2e_generate_cold_storage_object_key_not_configured():
     """

tests/test_litellm/responses/litellm_completion_transformation/test_session_handler.py

@@ -364,3 +364,53 @@ async def test_should_check_cold_storage_for_full_payload():
     with patch.object(litellm, 'configured_cold_storage_logger', None):
         result5 = ResponsesSessionHandler._should_check_cold_storage_for_full_payload(proxy_request_with_truncated_pdf)
         assert result5 == False, "Should return False when cold storage is not configured, even with truncated content"
+
+
+@pytest.mark.asyncio
+async def test_get_chat_completion_message_history_empty_response_dict():
+    """
+    Test that empty response dict is handled correctly without processing.
+    This tests the fix for response validation to check for empty dict responses.
+    """
+    from unittest.mock import AsyncMock, patch
+    
+    # Mock spend logs with empty response dict
+    mock_spend_logs = [
+        {
+            "request_id": "chatcmpl-test-empty-response",
+            "call_type": "aresponses", 
+            "api_key": "test_key",
+            "spend": 0.001,
+            "total_tokens": 0,
+            "prompt_tokens": 0,
+            "completion_tokens": 0,
+            "startTime": "2025-01-15T10:30:00.000+00:00",
+            "endTime": "2025-01-15T10:30:01.000+00:00",
+            "model": "gpt-4",
+            "session_id": "test-session",
+            "proxy_server_request": {
+                "input": "test input",
+                "model": "gpt-4"
+            },
+            "response": {}  # Empty dict - should not be processed
+        }
+    ]
+    
+    with patch.object(ResponsesSessionHandler, "get_all_spend_logs_for_previous_response_id") as mock_get_spend_logs:
+        mock_get_spend_logs.return_value = mock_spend_logs
+        
+        # Call the function
+        result = await ResponsesSessionHandler.get_chat_completion_message_history_for_previous_response_id(
+            "chatcmpl-test-empty-response"
+        )
+        
+        # Verify that user message was added but no assistant response
+        # Since response is empty dict, no assistant response should be processed
+        # But user input from proxy_server_request should still be included
+        messages = result["messages"]
+        assert len(messages) == 1  # Only user message, no assistant response
+        assert messages[0]["role"] == "user"
+        assert messages[0]["content"] == "test input"
+        
+        # Verify the session was still created correctly
+        assert result["litellm_session_id"] == "test-session"