scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-11 03:38:38 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 84

Add allowing Key based prefix to s3 path (#16237)

Browse source 
* add Key based prefix

* add Key based prefix

* update documentation

* correct doc

---------

Co-authored-by: deepanshu <deepanshu.lulla@hq.bill.com>
This commit is contained in:
Deepanshu Lulla 2025-11-05 17:43:21 -05:00 • committed by GitHub
parent 466e7d178c
commit 3ef210e5f4
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 229 additions and 27 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

34
docs/my-website/docs/proxy/logging.md
View file

@ -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
```

32
litellm/integrations/custom_logger.py
View file

@ -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:

4
litellm/integrations/s3.py
View file

@ -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

48
litellm/integrations/s3_v2.py
View file

@ -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"

138
tests/test_litellm/integrations/test_s3_v2.py
View file

@ -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}"