diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index f90515ac81f..e89470152d9 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -2732,9 +2732,10 @@ class Logging(LiteLLMLoggingBaseClass): def record_native_bedrock_guardrail(self, result: object) -> None: """Bill the guardrail trace a Bedrock response carries when the deployment's ``guardrailConfig`` ran it, by - recording it as guardrail information on this call, the same entry the proxy-side Bedrock guardrail writes.""" + recording it as guardrail information on this call, the same entry the proxy-side Bedrock guardrail writes. + A cache hit carries the original response's trace but made no Bedrock call, so it bills nothing.""" trace: Final = bedrock_guardrail_trace(result) - if trace is None: + if trace is None or served_from_cache(result) or self.model_call_details.get("cache_hit") is True: return optional_params: Final = object_mapping(getattr(self, "optional_params", None)) or {} guardrail_config: Final = object_mapping(optional_params.get("guardrailConfig")) or {} diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 7f8a38cc8ec..25e4482ab4c 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -14,7 +14,7 @@ from collections.abc import Mapping, Sequence from dataclasses import replace from typing import TYPE_CHECKING, Any, Final -from pydantic import TypeAdapter +from pydantic import JsonValue, TypeAdapter, ValidationError import litellm from litellm.constants import REDACTED_BY_LITELLM @@ -120,11 +120,52 @@ def _redacted_streaming_response_copy(streaming_response: object) -> object: return _redact_streaming_response(redacted_response) +BEDROCK_GUARDRAIL_TRACE_FIELDS: Final = ("trace", "amazon-bedrock-trace") +_MATCHED_TEXT_KEYS: Final = frozenset({"match", "regex"}) +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) + + +def _without_matched_text(node: JsonValue) -> JsonValue: + if isinstance(node, dict): + return { + key: "[REDACTED]" if key in _MATCHED_TEXT_KEYS and isinstance(value, str) else _without_matched_text(value) + for key, value in node.items() + } + if isinstance(node, list): + return [_without_matched_text(item) for item in node] + return node + + +def _redacted_bedrock_trace(trace: object) -> object: + try: + return _without_matched_text(_JSON_OBJECT.validate_python(trace)) + except ValidationError: + return trace + + +def _redact_bedrock_guardrail_trace(response: object) -> None: + """A Bedrock guardrail trace echoes the matched text of its sensitive-information and word findings back.""" + mapping: Final = _as_mapping(response) + if isinstance(response, dict): + traces: Final = { + field: _redacted_bedrock_trace(value) + for field, value in (mapping or {}).items() + if field in BEDROCK_GUARDRAIL_TRACE_FIELDS + } + response.update(traces) + return + for field, value in ((name, getattr(response, name, None)) for name in BEDROCK_GUARDRAIL_TRACE_FIELDS): + if value is not None: + setattr(response, field, _redacted_bedrock_trace(value)) + + def _redact_streaming_response(streaming_response): if hasattr(streaming_response, "choices"): for choice in streaming_response.choices: _redact_choice_content(choice) redact_vertex_ai_metadata_from_logged_object(streaming_response) + if isinstance(streaming_response, litellm.ModelResponse): + _redact_bedrock_guardrail_trace(streaming_response) elif isinstance(streaming_response, litellm.ResponsesAPIResponse): return perform_redaction({}, streaming_response) elif hasattr(streaming_response, "output"): @@ -351,6 +392,7 @@ def _redact_standard_logging_object( if isinstance(response.get("choices"), list): _redact_model_response_dict_choices(response["choices"], redacted_str) redact_vertex_ai_metadata_from_logged_object(response) + _redact_bedrock_guardrail_trace(response) elif isinstance(response, str): standard_logging_object["response"] = redacted_str else: @@ -523,11 +565,13 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons for choice in _result.choices: _redact_choice_content(choice) redact_vertex_ai_metadata_from_logged_object(_result) + _redact_bedrock_guardrail_trace(_result) elif isinstance(_result, dict) and "choices" in _result: # Handle dict representation of ModelResponse (e.g., from model_dump()) if _result.get("choices") is not None: _redact_model_response_dict_choices(_result["choices"], REDACTED_BY_LITELLM) redact_vertex_ai_metadata_from_logged_object(_result) + _redact_bedrock_guardrail_trace(_result) elif isinstance(_result, dict) and "output" in _result: return _redacted_responses_api_response(_result, redact_instructions=redact_streaming_responses) elif isinstance(_result, litellm.ResponsesAPIResponse): diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index f44bb1a9050..c5938ae8ecb 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -13666,3 +13666,30 @@ def test_record_native_bedrock_guardrail_bills_when_the_metadata_is_none() -> No (entry,) = logging_obj.model_call_details["litellm_params"]["metadata"]["standard_logging_guardrail_information"] assert entry["guardrail_usage"] == {"contentPolicyUnits": 2} assert entry["guardrail_cost"] > 0 + + +def test_record_native_bedrock_guardrail_bills_nothing_for_a_response_served_from_cache() -> None: + import datetime + + from litellm.litellm_core_utils.litellm_logging import Logging + + logging_obj = Logging( + model="bedrock/anthropic.claude-sonnet-4-5-20250929-v1:0", + messages=[{"role": "user", "content": "hello"}], + stream=False, + call_type="completion", + start_time=datetime.datetime.now(), + litellm_call_id="call-cache-hit", + function_id="fn", + ) + logging_obj.model_call_details["litellm_params"] = {"metadata": {}} + trace = {"guardrail": {"inputAssessment": {"gid": {"invocationMetrics": {"usage": {"contentPolicyUnits": 2}}}}}} + cached = litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "hi"}}], model="x") + setattr(cached, "trace", trace) + cached._hidden_params = {"cache_hit": True} + + logging_obj.record_native_bedrock_guardrail(cached) + logging_obj.model_call_details["cache_hit"] = True + logging_obj.record_native_bedrock_guardrail({"trace": trace}) + + assert "standard_logging_guardrail_information" not in logging_obj.model_call_details["litellm_params"]["metadata"] diff --git a/tests/unit/litellm_core_utils/test_redact_messages.py b/tests/unit/litellm_core_utils/test_redact_messages.py index 6a734ca3b70..5b64e42d8a8 100644 --- a/tests/unit/litellm_core_utils/test_redact_messages.py +++ b/tests/unit/litellm_core_utils/test_redact_messages.py @@ -26,15 +26,16 @@ import litellm from litellm.constants import LOGGING_WORKER_MAX_TIME_PER_COROUTINE from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER -from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.litellm_core_utils.redact_messages import ( _redact_responses_api_output, _redacted_responses_api_response, perform_redaction, redacted_model_call_details_for_custom_logger, + redacted_response_obj_for_custom_logger, redacted_standard_logging_payload, should_redact_message_logging, ) +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.responses.main import mock_responses_api_response from litellm.types.router import BaselineRouteStamp from litellm.types.utils import ( @@ -176,7 +177,11 @@ def test_redaction_hides_a_compaction_summary_in_responses_api_output(shape: str object="response", status="completed", output=[ - {"type": "compaction", "id": "cmp_1", "encrypted_content": json.dumps({"type": "compaction", "content": SECRET})} + { + "type": "compaction", + "id": "cmp_1", + "encrypted_content": json.dumps({"type": "compaction", "content": SECRET}), + } ], ) result: Final = response if shape == "object" else response.model_dump() @@ -2041,3 +2046,32 @@ async def test_redaction_with_metadata_completion_api(): response = standard_logging_payload["response"] assert response["choices"][0]["message"]["content"] == "redacted-by-litellm" assert standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm" + + +def test_perform_redaction_strips_the_matched_text_from_the_bedrock_guardrail_traces() -> None: + finding = { + "guardrail": { + "inputAssessment": { + "gid": { + "sensitiveInformationPolicy": { + "piiEntities": [{"type": "US_SSN", "match": "123-45-6789", "action": "ANONYMIZED"}] + } + } + } + } + } + response = litellm.ModelResponse(choices=[{"message": {"role": "assistant", "content": "hi"}}], model="x") + setattr(response, "trace", copy.deepcopy(finding)) + setattr(response, "amazon-bedrock-trace", copy.deepcopy(finding)) + opted_out: Final = CustomLogger(message_logging=False) + + redacted: Final = redacted_response_obj_for_custom_logger(response, opted_out) + redacted_dict: Final = redacted_response_obj_for_custom_logger( + {"choices": [], "trace": copy.deepcopy(finding)}, opted_out + ) + + assert isinstance(redacted, litellm.ModelResponse) + assert "123-45-6789" not in redacted.model_dump_json() + assert "US_SSN" in redacted.model_dump_json(), "the finding stays, only its matched text goes" + assert "123-45-6789" not in str(redacted_dict) + assert "123-45-6789" in response.model_dump_json(), "the caller's own response is not redacted"