diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index 2e70b1d6519..7f711708fd6 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -1,8 +1,8 @@ #### What this does #### -# On success + failure, log events to Supabase +# On success + failure, log events to S3 from datetime import datetime -from typing import Optional, cast +from typing import Any, Dict, Optional, cast import litellm from litellm._logging import print_verbose, verbose_logger @@ -34,6 +34,7 @@ class S3Logger: ) s3_use_team_prefix = False + s3_log_response = True # Default: log both prompts and responses if litellm.s3_callback_params is not None: # read in .env variables - example os.environ/AWS_BUCKET_NAME @@ -62,10 +63,17 @@ class S3Logger: s3_use_team_prefix = bool( litellm.s3_callback_params.get("s3_use_team_prefix", False) ) + # New param: control whether to log responses + # Set to False to log only prompts (no responses) + s3_log_response = litellm.s3_callback_params.get( + "s3_log_response", True + ) self.s3_use_team_prefix = s3_use_team_prefix + self.s3_log_response = s3_log_response self.bucket_name = s3_bucket_name self.s3_path = s3_path verbose_logger.debug(f"s3 logger using endpoint url {s3_endpoint_url}") + verbose_logger.debug(f"s3 logger s3_log_response={s3_log_response}") # Create an S3 client with custom endpoint URL self.s3_client = boto3.client( "s3", @@ -89,6 +97,33 @@ class S3Logger: ): self.log_event(kwargs, response_obj, start_time, end_time, print_verbose) + def _filter_payload_fields( + self, payload: StandardLoggingPayload + ) -> Dict[str, Any]: + """ + Filter out fields from the payload based on configuration. + + When s3_log_response is False, removes the 'response' field from the payload + to log only prompts (no responses). + + Args: + payload: The StandardLoggingPayload to filter + + Returns: + A dictionary with filtered fields + """ + # Convert TypedDict to regular dict for modification + filtered_payload: Dict[str, Any] = dict(payload) + + if not self.s3_log_response: + # Remove response field when s3_log_response is False + filtered_payload.pop("response", None) + verbose_logger.debug( + "s3 Logger - s3_log_response=False, excluding response from payload" + ) + + return filtered_payload + def log_event(self, kwargs, response_obj, start_time, end_time, print_verbose): try: verbose_logger.debug( @@ -156,7 +191,9 @@ class S3Logger: from litellm.litellm_core_utils.safe_json_dumps import safe_dumps - payload_str = safe_dumps(payload) + # Filter payload fields based on configuration + filtered_payload = self._filter_payload_fields(payload) + payload_str = safe_dumps(filtered_payload) print_verbose(f"\ns3 Logger - Logging payload = {payload_str}") diff --git a/tests/logging_callback_tests/test_s3_log_response.py b/tests/logging_callback_tests/test_s3_log_response.py new file mode 100644 index 00000000000..76b485e45fb --- /dev/null +++ b/tests/logging_callback_tests/test_s3_log_response.py @@ -0,0 +1,97 @@ +""" +Tests for S3 s3_log_response parameter. + +When s3_log_response=False, the S3 logger should log only prompts (no responses). +""" + +import pytest +from unittest.mock import MagicMock, patch +from datetime import datetime + +import litellm +from litellm.integrations.s3 import S3Logger + + +class TestS3LogResponse: + """Test the s3_log_response parameter functionality.""" + + def test_filter_payload_excludes_response_when_false(self): + """Test that response is excluded from payload when s3_log_response=False.""" + with patch("boto3.client"): + # Set up S3 callback params with s3_log_response=False + litellm.s3_callback_params = { + "s3_bucket_name": "test-bucket", + "s3_log_response": False, + } + + logger = S3Logger() + + # Create a mock payload with both messages and response + mock_payload = { + "id": "test-id", + "messages": [{"role": "user", "content": "Hello"}], + "response": "Hi there!", + "model": "gpt-4", + "metadata": {}, + } + + filtered = logger._filter_payload_fields(mock_payload) + + # Response should be removed + assert "response" not in filtered + # Messages should still be present + assert "messages" in filtered + assert filtered["messages"] == [{"role": "user", "content": "Hello"}] + + def test_filter_payload_includes_response_when_true(self): + """Test that response is included in payload when s3_log_response=True (default).""" + with patch("boto3.client"): + # Set up S3 callback params with s3_log_response=True (or not set) + litellm.s3_callback_params = { + "s3_bucket_name": "test-bucket", + "s3_log_response": True, + } + + logger = S3Logger() + + # Create a mock payload with both messages and response + mock_payload = { + "id": "test-id", + "messages": [{"role": "user", "content": "Hello"}], + "response": "Hi there!", + "model": "gpt-4", + "metadata": {}, + } + + filtered = logger._filter_payload_fields(mock_payload) + + # Both should be present + assert "response" in filtered + assert "messages" in filtered + assert filtered["response"] == "Hi there!" + + def test_filter_payload_default_includes_response(self): + """Test that response is included by default when s3_log_response is not set.""" + with patch("boto3.client"): + # Set up S3 callback params without s3_log_response + litellm.s3_callback_params = { + "s3_bucket_name": "test-bucket", + } + + logger = S3Logger() + + # Default should be True + assert logger.s3_log_response is True + + mock_payload = { + "id": "test-id", + "messages": [{"role": "user", "content": "Hello"}], + "response": "Hi there!", + "model": "gpt-4", + "metadata": {}, + } + + filtered = logger._filter_payload_fields(mock_payload) + + # Response should be present (default behavior) + assert "response" in filtered