diff --git a/docs/my-website/docs/proxy/logging.md b/docs/my-website/docs/proxy/logging.md index 20fc6103883..cf36963b7e1 100644 --- a/docs/my-website/docs/proxy/logging.md +++ b/docs/my-website/docs/proxy/logging.md @@ -1366,14 +1366,12 @@ Your logs should be available on the specified s3 Bucket ### Team Alias Prefix in Object Key -**This is a preview feature** - -You can add the team alias to the object key by setting the `team_alias` in the `config.yaml` file. This will prefix the object key with the team alias. +You can add the team alias to the object key by setting the `team_alias` in the `config.yaml` file. +This will prefix the object key with the team alias. ```yaml litellm_settings: callbacks: ["s3_v2"] - enable_preview_features: true s3_callback_params: s3_bucket_name: logs-bucket-litellm s3_region_name: us-west-2 @@ -1386,6 +1384,28 @@ litellm_settings: On s3 bucket, you will see the object key as `my-test-path/my-team-alias/...` +### Key Alias Prefix in Object Key + +You can add the user api key alias to the s3 object key by enabling s3_use_key_prefix. + +```yaml +litellm_settings: + callbacks: ["s3_v2"] + s3_callback_params: + s3_bucket_name: logs-bucket-litellm + s3_region_name: us-west-2 + s3_aws_access_key_id: os.environ/AWS_ACCESS_KEY_ID + s3_aws_secret_access_key: os.environ/AWS_SECRET_ACCESS_KEY + s3_path: my-test-path + s3_endpoint_url: https://s3.amazonaws.com + s3_use_key_prefix: true +``` + +On s3 bucket, you will see the object key as `my-test-path/my-key-alias/...` + +if both team alias and key alias are enabled then the path becomes +`my-test-path/my-team-alias/my-key-alias/...` + ## AWS SQS @@ -1432,9 +1452,13 @@ litellm_settings: # AWS Region for your SQS queue (e.g., us-east-1, eu-central-1, etc.) # --- Logging Controls --- - sqs_strip_base64_files: true + sqs_strip_base64_files: false # If true, LiteLLM will remove or redact base64-encoded binary data (e.g., PDFs, images, audio) # from logged messages to avoid large payloads. SQS has a 1 MB payload size limit. + s3_use_team_prefix: false + # If true, Litellm will add the team alias prefix to s3 path + s3_use_key_prefix: false + # If true, Litellm will add the key alias prefix to s3 path ``` diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 6c613b6c142..fd8ab2bad9d 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -635,6 +635,38 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac f"[CustomLogger] Completed base64 strip; retained {total_items} content items" ) return payload + + def _strip_base64_from_messages_sync( + self, payload: "StandardLoggingPayload", max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER + ) -> "StandardLoggingPayload": + """ + Removes or redacts base64-encoded file data (e.g., PDFs, images, audio) + from messages and responses before sending to SQS. + + Behavior: + • Drop entries with a 'file' key. + • Drop entries with type == 'file' or any non-text type. + • Keep untyped or text content. + • Recursively redact inline base64 blobs in *any* string field, at any depth. + """ + raw_messages: Any = payload.get("messages", []) + messages: List[Any] = raw_messages if isinstance(raw_messages, list) else [] + verbose_logger.debug(f"[CustomLogger] Stripping base64 from {len(messages)} messages") + + if messages: + payload["messages"] = self._process_messages(messages=messages, max_depth=max_depth) + + total_items = 0 + for m in payload.get("messages", []) or []: + if isinstance(m, dict): + content = m.get("content", []) + if isinstance(content, list): + total_items += len(content) + + verbose_logger.debug( + f"[CustomLogger] Completed base64 strip; retained {total_items} content items" + ) + return payload def _redact_base64(self, value: Any, depth: int = 0, max_depth: int = DEFAULT_MAX_RECURSE_DEPTH_SENSITIVE_DATA_MASKER) -> Any: diff --git a/litellm/integrations/s3.py b/litellm/integrations/s3.py index 53caeb0d198..2e70b1d6519 100644 --- a/litellm/integrations/s3.py +++ b/litellm/integrations/s3.py @@ -181,13 +181,13 @@ class S3Logger: def get_s3_object_key( s3_path: str, - team_alias_prefix: str, + prefix: str, start_time: datetime, s3_file_name: str, ) -> str: s3_object_key = ( (s3_path.rstrip("/") + "/" if s3_path else "") - + team_alias_prefix + + prefix + start_time.strftime("%Y-%m-%d") + "/" + s3_file_name diff --git a/litellm/integrations/s3_v2.py b/litellm/integrations/s3_v2.py index fcfdd112a5c..534b85e4752 100644 --- a/litellm/integrations/s3_v2.py +++ b/litellm/integrations/s3_v2.py @@ -50,6 +50,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_config=None, s3_use_team_prefix: bool = False, s3_strip_base64_files: bool = False, + s3_use_key_prefix: bool = False, **kwargs, ): try: @@ -76,7 +77,8 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_config=s3_config, s3_path=s3_path, s3_use_team_prefix=s3_use_team_prefix, - s3_strip_base64_files=s3_strip_base64_files + s3_strip_base64_files=s3_strip_base64_files, + s3_use_key_prefix=s3_use_key_prefix ) verbose_logger.debug(f"s3 logger using endpoint url {s3_endpoint_url}") @@ -132,6 +134,7 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): s3_path: Optional[str] = None, s3_use_team_prefix: bool = False, s3_strip_base64_files: bool = False, + s3_use_key_prefix: bool = False, ): """ Initialize the s3 params for this logging callback @@ -204,6 +207,11 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): or s3_use_team_prefix ) + self.s3_use_key_prefix = ( + bool(litellm.s3_callback_params.get("s3_use_key_prefix", False)) + or s3_use_key_prefix + ) + self.s3_strip_base64_files = ( bool(litellm.s3_callback_params.get("s3_strip_base64_files", False)) or s3_strip_base64_files @@ -384,36 +392,36 @@ class S3Logger(CustomBatchLogger, BaseAWSLLM): return None if self.s3_strip_base64_files: - import asyncio - standard_logging_payload = asyncio.run(self._strip_base64_from_messages(standard_logging_payload)) + standard_logging_payload = self._strip_base64_from_messages_sync(standard_logging_payload) - team_alias = standard_logging_payload["metadata"].get("user_api_key_team_alias") + # Base prefix (default empty) + prefix_components = [] + if self.s3_use_team_prefix: + team_alias = standard_logging_payload.get("metadata", {}).get("user_api_key_team_alias", None) + if team_alias: + prefix_components.append(team_alias) + if self.s3_use_key_prefix: + user_api_key_alias = standard_logging_payload.get("metadata", {}).get("user_api_key_alias", None) + if user_api_key_alias: + prefix_components.append(user_api_key_alias) - team_alias_prefix = "" - if ( - litellm.enable_preview_features - and self.s3_use_team_prefix - and team_alias is not None - ): - team_alias_prefix = f"{team_alias}/" + + # Construct full prefix path + prefix_path = "/".join(prefix_components) + if prefix_path: + prefix_path += "/" s3_file_name = ( litellm.utils.get_logging_id(start_time, standard_logging_payload) or "" ) + verbose_logger.debug(f"Creating s3 file with prefix_components={prefix_components},prefix_path={prefix_path} and {s3_file_name}") s3_object_key = get_s3_object_key( s3_path=cast(Optional[str], self.s3_path) or "", - team_alias_prefix=team_alias_prefix, + prefix=prefix_path, start_time=start_time, s3_file_name=s3_file_name, ) - - s3_object_download_filename = ( - "time-" - + start_time.strftime("%Y-%m-%dT%H-%M-%S-%f") - + "_" - + standard_logging_payload["id"] - + ".json" - ) + verbose_logger.debug(f"s3_object_key={s3_object_key}") s3_object_download_filename = f"time-{start_time.strftime('%Y-%m-%dT%H-%M-%S-%f')}_{standard_logging_payload['id']}.json" diff --git a/tests/test_litellm/integrations/test_s3_v2.py b/tests/test_litellm/integrations/test_s3_v2.py index 417153e24cc..0a3523699a9 100644 --- a/tests/test_litellm/integrations/test_s3_v2.py +++ b/tests/test_litellm/integrations/test_s3_v2.py @@ -478,3 +478,141 @@ async def test_strip_base64_recursive_redaction(): s = json.dumps(c).lower() # "[base64_redacted]" is fine, but raw base64 is not assert "base64," not in s, f"Found real base64 blob in: {s}" + + + +# -------------------------------------------------------------- +# Shared fixture that silences asyncio.create_task during tests +# -------------------------------------------------------------- +@pytest.fixture(autouse=True) +def patch_asyncio_create_task(): + """Prevent 'no running event loop' errors when S3Logger calls asyncio.create_task().""" + with patch("asyncio.create_task"): + yield + + +# -------------------------------------------------------------- +# Parametrized prefix combination test +# -------------------------------------------------------------- +@pytest.mark.parametrize( + "use_team_prefix,use_key_prefix,team_alias,key_alias,expected_prefix", + [ + (False, False, "teamA", "keyA", ""), + (True, False, "teamA", "keyA", "teamA/"), + (False, True, "teamA", "keyA", "keyA/"), + (True, True, "teamA", "keyA", "teamA/keyA/"), + (True, True, None, "keyA", "keyA/"), + (True, True, "teamA", None, "teamA/"), + (True, True, None, None, ""), + ], +) +def test_s3_object_key_prefix_combinations( + use_team_prefix, use_key_prefix, team_alias, key_alias, expected_prefix +): + """ + Validate correct S3 prefix composition for team alias + key alias combinations. + """ + with patch("litellm.integrations.s3_v2.get_s3_object_key") as mock_get_key: + mock_get_key.return_value = "mocked/s3/object/key.json" + + logger = S3Logger( + s3_bucket_name="test-bucket", + s3_region_name="us-east-1", + s3_use_team_prefix=use_team_prefix, + s3_use_key_prefix=use_key_prefix, + ) + + payload = StandardLoggingPayload( + id="abc123", + metadata={ + "user_api_key_team_alias": team_alias, + "user_api_key_alias": key_alias, + }, + messages=[{"role": "user", "content": [{"type": "text", "text": "hi"}]}], + ) + + result = logger.create_s3_batch_logging_element(datetime.utcnow(), payload) + assert result is not None + mock_get_key.assert_called_once() + + prefix_arg = mock_get_key.call_args.kwargs.get("prefix") + assert prefix_arg == expected_prefix, ( + f"Expected prefix '{expected_prefix}', got '{prefix_arg}' " + f"for team={team_alias}, key={key_alias}, " + f"use_team_prefix={use_team_prefix}, use_key_prefix={use_key_prefix}" + ) + + +# -------------------------------------------------------------- +# Test prefix priority and concatenation +# -------------------------------------------------------------- +def test_prefix_priority_and_path_construction(): + """ + Validate that prefix components are ordered and joined with '/' only once. + """ + with patch("litellm.integrations.s3_v2.get_s3_object_key") as mock_get_key: + mock_get_key.return_value = "mocked/key" + + logger = S3Logger(s3_use_team_prefix=True, s3_use_key_prefix=True) + payload = StandardLoggingPayload( + id="xyz999", + metadata={ + "user_api_key_team_alias": "Team-Alpha", + "user_api_key_alias": "API-12345", + }, + messages=[], + ) + + logger.create_s3_batch_logging_element(datetime.utcnow(), payload) + prefix_arg = mock_get_key.call_args.kwargs.get("prefix", "") + + assert prefix_arg == "Team-Alpha/API-12345/" + assert "//" not in prefix_arg + + +# -------------------------------------------------------------- +# Test when prefixes are disabled +# -------------------------------------------------------------- +def test_prefix_absent_when_flags_disabled(): + """ + Verify prefix is omitted entirely when prefix flags are False. + """ + with patch("litellm.integrations.s3_v2.get_s3_object_key") as mock_get_key: + mock_get_key.return_value = "mocked/key" + + logger = S3Logger(s3_use_team_prefix=False, s3_use_key_prefix=False) + payload = StandardLoggingPayload( + id="no-prefix", + metadata={ + "user_api_key_team_alias": "team-x", + "user_api_key_alias": "key-x", + }, + messages=[], + ) + + logger.create_s3_batch_logging_element(datetime.utcnow(), payload) + prefix_arg = mock_get_key.call_args.kwargs.get("prefix", None) + assert prefix_arg == "", f"Expected empty prefix, got {prefix_arg}" + + +# -------------------------------------------------------------- +# Integration-style test (asyncio fixture will patch create_task) +# -------------------------------------------------------------- +@pytest.mark.asyncio +async def test_combined_prefix_reflects_in_s3_object_key(): + """ + Integration-style test ensuring final s3_object_key includes both prefixes correctly. + """ + logger = S3Logger(s3_use_team_prefix=True, s3_use_key_prefix=True) + payload = StandardLoggingPayload( + id="int-test", + metadata={ + "user_api_key_team_alias": "myteam", + "user_api_key_alias": "apikey", + }, + messages=[], + ) + + result = logger.create_s3_batch_logging_element(datetime.utcnow(), payload) + key = result.s3_object_key + assert "myteam/apikey/" in key, f"Expected both prefixes in key: {key}"