mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-16 23:41:43 +00:00
fix(proxy): redact guardrail match fields once for spend logs
- Add redact_nested_match_and_regex_keys in core_helpers for nested match/regex. - Apply in CustomGuardrail standard logging; Bedrock forwards raw JSON to avoid double redaction. - Delegate Bedrock HTTP detail assessments and _redact_pii_matches to the same helper. - Extend unit tests (core_helpers, CustomGuardrail, Bedrock spend-log mock). Made-with: Cursor
This commit is contained in:
Milan
parent
2e5982aa96
commit
688d278748
6 changed files with 187 additions and 46 deletions
|
|
@ -13,6 +13,7 @@ from typing import (
|
|||
)
|
||||
|
||||
from litellm._logging import verbose_logger
|
||||
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
|
||||
from litellm.caching import DualCache
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.types.guardrails import (
|
||||
|
|
@ -618,6 +619,13 @@ class CustomGuardrail(CustomLogger):
|
|||
if isinstance(item, dict):
|
||||
item.pop("secret_fields", None)
|
||||
|
||||
# Default-safe behavior: never persist raw matched spans in standard
|
||||
# guardrail logging payloads (single shared implementation; Bedrock hooks pass
|
||||
# raw provider JSON so redaction is not duplicated upstream).
|
||||
clean_guardrail_response = redact_nested_match_and_regex_keys(
|
||||
clean_guardrail_response
|
||||
)
|
||||
|
||||
slg = StandardLoggingGuardrailInformation(
|
||||
guardrail_name=self.guardrail_name,
|
||||
guardrail_provider=guardrail_provider,
|
||||
|
|
|
|||
|
|
@ -1,5 +1,6 @@
|
|||
# What is this?
|
||||
## Helper utilities
|
||||
import copy
|
||||
from typing import TYPE_CHECKING, Any, Iterable, List, Literal, Optional, Union
|
||||
|
||||
import httpx
|
||||
|
|
@ -435,3 +436,38 @@ def filter_internal_params(
|
|||
|
||||
# Filter out internal parameters
|
||||
return {k: v for k, v in data.items() if k not in internal_params}
|
||||
|
||||
|
||||
def redact_nested_match_and_regex_keys(
|
||||
payload: Union[dict, List[Any], str, None],
|
||||
) -> Union[dict, List[Any], str, None]:
|
||||
"""
|
||||
Deep-copy `payload` and replace every `match` / `regex` string field with
|
||||
"[REDACTED]" anywhere in nested dict/list structures.
|
||||
|
||||
Used for guardrail spend/compliance logging so raw spans are not persisted.
|
||||
"""
|
||||
if payload is None or isinstance(payload, str):
|
||||
return payload
|
||||
try:
|
||||
redacted: Union[dict, List[Any], str, None] = copy.deepcopy(payload)
|
||||
except Exception:
|
||||
return payload
|
||||
|
||||
def _walk(node: Any) -> None:
|
||||
if isinstance(node, dict):
|
||||
if "match" in node:
|
||||
node["match"] = "[REDACTED]"
|
||||
if "regex" in node:
|
||||
node["regex"] = "[REDACTED]"
|
||||
for value in node.values():
|
||||
_walk(value)
|
||||
elif isinstance(node, list):
|
||||
for item in node:
|
||||
_walk(item)
|
||||
|
||||
try:
|
||||
_walk(redacted)
|
||||
except Exception:
|
||||
return payload
|
||||
return redacted
|
||||
|
|
|
|||
|
|
@ -5,7 +5,6 @@
|
|||
# +-------------------------------------------------------------+
|
||||
# Thank you users! We ❤️ you! - Krrish & Ishaan
|
||||
|
||||
import copy
|
||||
import os
|
||||
import sys
|
||||
|
||||
|
|
@ -34,6 +33,7 @@ from fastapi import HTTPException
|
|||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.litellm_core_utils.core_helpers import redact_nested_match_and_regex_keys
|
||||
from litellm.caching import DualCache
|
||||
from litellm.exceptions import GuardrailInterventionNormalStringError
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
|
|
@ -79,50 +79,33 @@ class GuardrailMessageFilterResult(NamedTuple):
|
|||
|
||||
|
||||
def _redact_pii_matches(response_json: dict) -> dict:
|
||||
"""
|
||||
Redact match-like fields from a Bedrock ApplyGuardrail JSON payload.
|
||||
|
||||
Delegates to :func:`redact_nested_match_and_regex_keys` (same rules as spend
|
||||
logging). Kept as a Bedrock-module entry point for existing unit tests.
|
||||
"""
|
||||
try:
|
||||
# Create a deep copy to avoid modifying the original response
|
||||
redacted_response = copy.deepcopy(response_json)
|
||||
|
||||
# Get assessments from the response
|
||||
assessments = redacted_response.get("assessments", [])
|
||||
if not assessments:
|
||||
return redacted_response
|
||||
|
||||
for assessment in assessments:
|
||||
# Redact PII entities in sensitive information policy
|
||||
sensitive_info_policy = assessment.get("sensitiveInformationPolicy")
|
||||
if sensitive_info_policy:
|
||||
pii_entities = sensitive_info_policy.get("piiEntities", [])
|
||||
for pii_entity in pii_entities:
|
||||
if "match" in pii_entity:
|
||||
pii_entity["match"] = "[REDACTED]"
|
||||
|
||||
# Redact regex matches
|
||||
regexes = sensitive_info_policy.get("regexes", [])
|
||||
for regex_match in regexes:
|
||||
if "match" in regex_match:
|
||||
regex_match["match"] = "[REDACTED]"
|
||||
|
||||
# Redact custom word matches in word policy
|
||||
word_policy = assessment.get("wordPolicy")
|
||||
if word_policy:
|
||||
custom_words = word_policy.get("customWords", [])
|
||||
for custom_word in custom_words:
|
||||
if "match" in custom_word:
|
||||
custom_word["match"] = "[REDACTED]"
|
||||
|
||||
managed_words = word_policy.get("managedWordLists", [])
|
||||
for managed_word in managed_words:
|
||||
if "match" in managed_word:
|
||||
managed_word["match"] = "[REDACTED]"
|
||||
|
||||
return redacted_response
|
||||
redacted = redact_nested_match_and_regex_keys(response_json)
|
||||
return redacted if isinstance(redacted, dict) else response_json
|
||||
except Exception as e:
|
||||
# We do not want to fail in any case so this is just a warning
|
||||
verbose_proxy_logger.warning("Guardrail log redaction failed: %s", str(e))
|
||||
return response_json
|
||||
|
||||
|
||||
def _redact_assessment_match_fields(assessments: List[dict]) -> List[dict]:
|
||||
"""
|
||||
Redact sensitive match-like fields from blocked assessment summaries.
|
||||
|
||||
This is used for customer-visible error payloads (HTTPException.detail) where
|
||||
we want to preserve policy/type/action metadata without echoing raw matched
|
||||
content.
|
||||
"""
|
||||
redacted = redact_nested_match_and_regex_keys(assessments)
|
||||
return redacted if isinstance(redacted, list) else assessments
|
||||
|
||||
|
||||
class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
||||
# During-call must use async_moderation_hook (not unified apply_guardrail), otherwise
|
||||
# OpenAI translation always passes input_type="request" and spend/UI show PRE-CALL.
|
||||
|
|
@ -521,9 +504,12 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
#########################################################
|
||||
# Add guardrail information to request trace
|
||||
#########################################################
|
||||
_json_response = httpx_response.json()
|
||||
# Raw Bedrock JSON is passed here; match/regex redaction runs once inside
|
||||
# CustomGuardrail.add_standard_logging_guardrail_information_to_request_data.
|
||||
self.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_provider=self.guardrail_provider,
|
||||
guardrail_json_response=httpx_response.json(),
|
||||
guardrail_json_response=_json_response,
|
||||
request_data=request_data or {},
|
||||
guardrail_status=self._get_bedrock_guardrail_response_status(
|
||||
response=httpx_response
|
||||
|
|
@ -536,9 +522,10 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
#########################################################
|
||||
if httpx_response.status_code == 200:
|
||||
# check if the response was flagged
|
||||
_json_response = httpx_response.json()
|
||||
redacted_response = _redact_pii_matches(_json_response)
|
||||
verbose_proxy_logger.debug("Bedrock AI response : %s", redacted_response)
|
||||
verbose_proxy_logger.debug(
|
||||
"Bedrock AI response : %s",
|
||||
redact_nested_match_and_regex_keys(_json_response),
|
||||
)
|
||||
bedrock_guardrail_response = BedrockGuardrailResponse(**_json_response)
|
||||
if self._should_raise_guardrail_blocked_exception(
|
||||
bedrock_guardrail_response
|
||||
|
|
@ -815,7 +802,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM):
|
|||
|
||||
assessments = self._extract_blocked_assessments(response)
|
||||
if assessments:
|
||||
detail["assessments"] = assessments
|
||||
detail["assessments"] = _redact_assessment_match_fields(assessments)
|
||||
|
||||
return HTTPException(status_code=400, detail=detail)
|
||||
|
||||
|
|
|
|||
|
|
@ -1047,3 +1047,50 @@ class TestTracingFieldsPopulation:
|
|||
assert slg["classification"] == classification
|
||||
assert slg["detection_method"] == "llm-judge"
|
||||
assert slg["confidence_score"] == 0.94
|
||||
|
||||
|
||||
class TestCustomGuardrailSpendLogMatchRedaction:
|
||||
"""Guardrail JSON persisted via standard_logging must not contain raw match spans."""
|
||||
|
||||
def test_add_standard_logging_redacts_nested_match(self):
|
||||
cg = CustomGuardrail(guardrail_name="test-rail")
|
||||
raw = {
|
||||
"assessments": [
|
||||
{
|
||||
"sensitiveInformationPolicy": {
|
||||
"piiEntities": [
|
||||
{"type": "NAME", "match": "GG", "action": "BLOCKED"}
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
request_data: dict = {"metadata": {}}
|
||||
cg.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response=raw,
|
||||
request_data=request_data,
|
||||
guardrail_status="guardrail_intervened",
|
||||
)
|
||||
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
||||
assert (
|
||||
slg["guardrail_response"]["assessments"][0]["sensitiveInformationPolicy"][
|
||||
"piiEntities"
|
||||
][0]["match"]
|
||||
== "[REDACTED]"
|
||||
)
|
||||
assert raw["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
|
||||
"match"
|
||||
] == "GG"
|
||||
|
||||
def test_add_standard_logging_redacts_regex_field(self):
|
||||
cg = CustomGuardrail(guardrail_name="test-rail")
|
||||
raw = {"filters": [{"regex": r"\d{3}-\d{2}-\d{4}", "action": "BLOCKED"}]}
|
||||
request_data: dict = {"metadata": {}}
|
||||
cg.add_standard_logging_guardrail_information_to_request_data(
|
||||
guardrail_json_response=raw,
|
||||
request_data=request_data,
|
||||
guardrail_status="success",
|
||||
)
|
||||
slg = request_data["metadata"]["standard_logging_guardrail_information"][0]
|
||||
assert slg["guardrail_response"]["filters"][0]["regex"] == "[REDACTED]"
|
||||
assert raw["filters"][0]["regex"] == r"\d{3}-\d{2}-\d{4}"
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@ from litellm.litellm_core_utils.core_helpers import (
|
|||
_FINISH_REASON_MAP,
|
||||
map_finish_reason,
|
||||
reconstruct_model_name,
|
||||
redact_nested_match_and_regex_keys,
|
||||
)
|
||||
|
||||
|
||||
|
|
@ -150,3 +151,37 @@ class TestFinishReasonMapOutputsAreValid:
|
|||
f"Mapped value '{openai_reason}' (from '{provider_reason}') "
|
||||
f"is not a valid OpenAI finish reason"
|
||||
)
|
||||
|
||||
|
||||
class TestRedactNestedMatchAndRegexKeys:
|
||||
def test_redacts_match_and_regex_recursively(self):
|
||||
payload = {
|
||||
"assessments": [
|
||||
{
|
||||
"sensitiveInformationPolicy": {
|
||||
"piiEntities": [
|
||||
{"type": "NAME", "match": "secret-name", "action": "BLOCKED"}
|
||||
]
|
||||
},
|
||||
"wordPolicy": {
|
||||
"customWords": [{"match": "badword", "action": "BLOCKED"}]
|
||||
},
|
||||
}
|
||||
],
|
||||
"regex": "should-redact-key-named-regex",
|
||||
}
|
||||
out = redact_nested_match_and_regex_keys(payload)
|
||||
assert out["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][0][
|
||||
"match"
|
||||
] == "[REDACTED]"
|
||||
assert out["assessments"][0]["wordPolicy"]["customWords"][0]["match"] == (
|
||||
"[REDACTED]"
|
||||
)
|
||||
assert out["regex"] == "[REDACTED]"
|
||||
assert payload["assessments"][0]["sensitiveInformationPolicy"]["piiEntities"][
|
||||
0
|
||||
]["match"] == "secret-name"
|
||||
|
||||
def test_passes_through_none_and_str(self):
|
||||
assert redact_nested_match_and_regex_keys(None) is None
|
||||
assert redact_nested_match_and_regex_keys("plain") == "plain"
|
||||
|
|
|
|||
|
|
@ -190,7 +190,7 @@ async def test__redact_pii_matches_multiple_assessments():
|
|||
|
||||
@pytest.mark.asyncio
|
||||
async def test_bedrock_guardrail_logging_uses_redacted_response():
|
||||
"""Test that the Bedrock guardrail uses redacted response for logging"""
|
||||
"""Debug logs and standard_logging payloads must not include raw match values."""
|
||||
|
||||
# Create proper mock objects
|
||||
mock_user_api_key_dict = UserAPIKeyAuth()
|
||||
|
|
@ -292,6 +292,14 @@ async def test_bedrock_guardrail_logging_uses_redacted_response():
|
|||
== "PHONE"
|
||||
)
|
||||
|
||||
slg_list = request_data["metadata"]["standard_logging_guardrail_information"]
|
||||
assert (
|
||||
slg_list[0]["guardrail_response"]["assessments"][0][
|
||||
"sensitiveInformationPolicy"
|
||||
]["piiEntities"][0]["match"]
|
||||
== "[REDACTED]"
|
||||
)
|
||||
|
||||
print("Bedrock guardrail logging redaction test passed")
|
||||
|
||||
|
||||
|
|
@ -1221,7 +1229,18 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs():
|
|||
|
||||
mock_bedrock_response = MagicMock()
|
||||
mock_bedrock_response.status_code = 200
|
||||
mock_bedrock_response.json.return_value = {"action": "NONE", "assessments": []}
|
||||
mock_bedrock_response.json.return_value = {
|
||||
"action": "NONE",
|
||||
"assessments": [
|
||||
{
|
||||
"sensitiveInformationPolicy": {
|
||||
"piiEntities": [
|
||||
{"type": "NAME", "match": "GG", "action": "BLOCKED"}
|
||||
]
|
||||
}
|
||||
}
|
||||
],
|
||||
}
|
||||
|
||||
request_data = {
|
||||
"model": "gpt-4o",
|
||||
|
|
@ -1245,6 +1264,14 @@ async def test_make_bedrock_api_request_logging_event_type_for_spend_logs():
|
|||
logging_event_type=GuardrailEventHooks.during_call,
|
||||
)
|
||||
assert mock_log.call_args.kwargs["event_type"] == GuardrailEventHooks.during_call
|
||||
# Raw Bedrock JSON is forwarded; redaction runs once in
|
||||
# CustomGuardrail.add_standard_logging_guardrail_information_to_request_data.
|
||||
assert (
|
||||
mock_log.call_args.kwargs["guardrail_json_response"]["assessments"][0][
|
||||
"sensitiveInformationPolicy"
|
||||
]["piiEntities"][0]["match"]
|
||||
== "GG"
|
||||
)
|
||||
|
||||
mock_log.reset_mock()
|
||||
|
||||
|
|
@ -1308,7 +1335,7 @@ def _make_guardrail() -> BedrockGuardrail:
|
|||
|
||||
|
||||
def test_extract_blocked_assessments_pii_entity():
|
||||
"""L3: PII entity match (BLOCKED) is surfaced with category, type, and matched term."""
|
||||
"""L3: PII entity match (BLOCKED) is surfaced with category, type, and match."""
|
||||
g = _make_guardrail()
|
||||
response = {
|
||||
"action": "GUARDRAIL_INTERVENED",
|
||||
|
|
@ -1419,6 +1446,7 @@ def test_get_http_exception_includes_assessments_and_identifier():
|
|||
assert exc.detail["guardrailVersion"] == "1"
|
||||
assert exc.detail["assessments"][0]["policy"] == "sensitiveInformationPolicy"
|
||||
assert exc.detail["assessments"][0]["matches"][0]["type"] == "NAME"
|
||||
assert exc.detail["assessments"][0]["matches"][0]["match"] == "[REDACTED]"
|
||||
|
||||
|
||||
def test_get_http_exception_no_blocked_assessments_omits_field():
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue