diff --git a/litellm/integrations/custom_logger.py b/litellm/integrations/custom_logger.py index 473bf223105..df3237bff07 100644 --- a/litellm/integrations/custom_logger.py +++ b/litellm/integrations/custom_logger.py @@ -897,10 +897,9 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac This is useful for logging payloads that contain sensitive information. """ - from copy import copy - import litellm from litellm import Choices, Message, ModelResponse + from litellm.litellm_core_utils.classifier_logging import CLASSIFIER_AUDIT_FIELDS, without_classifier_audit turn_off_message_logging: Final[bool] = getattr(self, "turn_off_message_logging", False) excluded_fields: Final[list[str] | None] = getattr(litellm, "standard_logging_payload_excluded_fields", None) @@ -909,41 +908,25 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac if turn_off_message_logging is False and not excluded_fields: return model_call_details - # Only make a shallow copy of the top-level dict to avoid deepcopy issues - # with complex objects like AuthenticationError that may be present - model_call_details_copy: Final = copy(model_call_details) standard_logging_object: Final = model_call_details.get("standard_logging_object") if standard_logging_object is None: - return model_call_details_copy + return model_call_details.copy() # Make a copy of just the standard_logging_object to avoid modifying the original - standard_logging_object_copy: Final = copy(standard_logging_object) - - # Handle excluded fields - remove them entirely from the payload - if excluded_fields: - for field in excluded_fields: - if field in standard_logging_object_copy: - del standard_logging_object_copy[field] + standard_logging_object_copy: Final = { + key: value + for key, value in standard_logging_object.items() + if key not in (excluded_fields or ()) and not (turn_off_message_logging and key in CLASSIFIER_AUDIT_FIELDS) + } # Handle turn_off_message_logging - redact messages and responses (if not already excluded) if turn_off_message_logging: - from litellm.litellm_core_utils.classifier_logging import CLASSIFIER_AUDIT_FIELDS, without_classifier_audit - - for field in CLASSIFIER_AUDIT_FIELDS: - standard_logging_object_copy.pop(field, None) - params: Final = model_call_details_copy.get("litellm_params") - request: Final = params.get("proxy_server_request") if isinstance(params, dict) else None - if isinstance(params, dict) and isinstance(request, dict): - model_call_details_copy["litellm_params"] = { - **params, - "proxy_server_request": without_classifier_audit(request), - } redacted_str: Final = "redacted-by-litellm" - if "messages" not in (excluded_fields or []) and standard_logging_object_copy.get("messages") is not None: + if "messages" not in (excluded_fields or ()) and standard_logging_object_copy.get("messages") is not None: standard_logging_object_copy["messages"] = [Message(content=redacted_str).model_dump()] - if "response" not in (excluded_fields or []) and standard_logging_object_copy.get("response") is not None: + if "response" not in (excluded_fields or ()) and standard_logging_object_copy.get("response") is not None: response: Final = standard_logging_object_copy["response"] # Check if this is a ResponsesAPIResponse (has "output" field) if isinstance(response, dict) and "output" in response: @@ -967,8 +950,15 @@ class CustomLogger: # https://docs.litellm.ai/docs/observability/custom_callbac model_response_dict: Final = model_response.model_dump() standard_logging_object_copy["response"] = model_response_dict - model_call_details_copy["standard_logging_object"] = standard_logging_object_copy - return model_call_details_copy + redacted_details: Final = {**model_call_details, "standard_logging_object": standard_logging_object_copy} + params: Final = model_call_details.get("litellm_params") + request: Final = params.get("proxy_server_request") if isinstance(params, dict) else None + if turn_off_message_logging and isinstance(params, dict) and isinstance(request, dict): + return { + **redacted_details, + "litellm_params": {**params, "proxy_server_request": without_classifier_audit(request)}, + } + return redacted_details async def get_proxy_server_request_from_cold_storage_with_object_key( self, diff --git a/litellm/litellm_core_utils/classifier_logging.py b/litellm/litellm_core_utils/classifier_logging.py index c0c93161099..4cbfdb8cd49 100644 --- a/litellm/litellm_core_utils/classifier_logging.py +++ b/litellm/litellm_core_utils/classifier_logging.py @@ -38,7 +38,7 @@ def is_classifier_call(call_type: str, params: Mapping[str, object]) -> bool: def masked_originating_request(request_kwargs: Mapping[str, object] | None) -> Mapping[str, JsonValue] | None: - request: Final = (request_kwargs or {}).get("proxy_server_request") + request: Final = request_kwargs.get("proxy_server_request") if request_kwargs is not None else None body: Final = request.get("body") if isinstance(request, Mapping) else None if not isinstance(body, Mapping): return None diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index c97566c21a8..9d22a5ddef5 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -16,7 +16,7 @@ from typing import TYPE_CHECKING, Any, Final import litellm from litellm.constants import REDACTED_BY_LITELLM from litellm.integrations.custom_logger import CustomLogger -from litellm.litellm_core_utils.classifier_logging import CLASSIFIER_AUDIT_FIELDS, without_classifier_audit +from litellm.litellm_core_utils.classifier_logging import without_classifier_audit from litellm.litellm_core_utils.core_helpers import ( get_metadata_variable_name_from_kwargs, ) @@ -171,22 +171,13 @@ def redacted_standard_logging_payload(payload: Mapping[str, object]) -> Mapping[ payload, but the failure path does not, so a callback that batches both has to redact the ones it is handed. """ - redacted: Final = copy.deepcopy(dict(payload)) # mutable-ok: redacted in place below - _redact_standard_logging_object({"standard_logging_object": redacted}) # mutable-ok: the callee's shape - return redacted + return _redact_standard_logging_object(payload) -def _redact_standard_logging_object(model_call_details: dict): - """Redact messages and response inside standard_logging_object if present.""" - standard_logging_object: Final = model_call_details.get("standard_logging_object") - if standard_logging_object is None: - return - +def _redact_standard_logging_object(payload: Mapping[str, object]) -> dict[str, object]: + standard_logging_object: Final = copy.deepcopy(without_classifier_audit(payload)) redacted_str: Final = REDACTED_BY_LITELLM - for field in CLASSIFIER_AUDIT_FIELDS: - standard_logging_object.pop(field, None) - if standard_logging_object.get("messages") is not None: standard_logging_object["messages"] = [{"role": "user", "content": redacted_str}] @@ -207,6 +198,7 @@ def _redact_standard_logging_object(model_call_details: dict): else: # For other formats (empty dict, None, etc.), use simple text format standard_logging_object["response"] = {"text": redacted_str} + return standard_logging_object def _redact_tool_calls_dict(message: Mapping[str, object]) -> None: @@ -258,8 +250,6 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons copy via redact_streaming_responses_for_custom_logger instead. """ # Redact model_call_details - for field in CLASSIFIER_AUDIT_FIELDS: - model_call_details.pop(field, None) params: Final = model_call_details.get("litellm_params") request: Final = params.get("proxy_server_request") if isinstance(params, dict) else None if isinstance(params, dict) and isinstance(request, Mapping): @@ -267,7 +257,9 @@ def perform_redaction(model_call_details: dict, result, redact_streaming_respons model_call_details["messages"] = [{"role": "user", "content": REDACTED_BY_LITELLM}] model_call_details["prompt"] = "" model_call_details["input"] = "" - _redact_standard_logging_object(model_call_details) + standard_logging_object: Final = model_call_details.get("standard_logging_object") + if isinstance(standard_logging_object, Mapping): + model_call_details["standard_logging_object"] = _redact_standard_logging_object(standard_logging_object) redact_vertex_ai_metadata_from_litellm_params(model_call_details) # Redact streaming response diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index da55d1cb153..360b24575bf 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -22,7 +22,7 @@ from litellm.constants import ( from litellm.constants import ( MAX_STRING_LENGTH_PROMPT_IN_DB as DEFAULT_MAX_STRING_LENGTH_PROMPT_IN_DB, ) -from litellm.litellm_core_utils.classifier_logging import classifier_audit_fields +from litellm.litellm_core_utils.classifier_logging import classifier_audit_fields, without_classifier_audit from litellm.litellm_core_utils.core_helpers import ( get_litellm_metadata_from_kwargs, reconstruct_model_name, @@ -1277,7 +1277,7 @@ def _get_proxy_server_request_for_spend_logs_payload( # If redaction is enabled, convert to serializable dict before redacting if should_redact_message_logging(model_call_details=model_call_details): - _request_body = _convert_mapping_to_json_serializable(_request_body) + _request_body = _convert_mapping_to_json_serializable(without_classifier_audit(_request_body)) perform_redaction(model_call_details=_request_body, result=None) _request_body = _sanitize_request_body_for_spend_logs_payload(_request_body) diff --git a/tests/test_litellm/litellm_core_utils/test_redact_messages.py b/tests/test_litellm/litellm_core_utils/test_redact_messages.py index 1ba9901aab0..76092e96307 100644 --- a/tests/test_litellm/litellm_core_utils/test_redact_messages.py +++ b/tests/test_litellm/litellm_core_utils/test_redact_messages.py @@ -17,6 +17,7 @@ from litellm.litellm_core_utils.redact_messages import ( _redact_responses_api_output, perform_redaction, redact_streaming_responses_for_custom_logger, + redacted_standard_logging_payload, should_redact_message_logging, ) from litellm.responses.main import mock_responses_api_response @@ -777,29 +778,6 @@ class TestPerformRedaction: assert response_obj.choices[0].message.content == "secret content" - -@pytest.mark.parametrize("callback_only", [False, True]) -def test_classifier_audit_redaction_removes_both_fields_and_source_carrier(callback_only: bool) -> None: - audit: Final = {"classifier_input": {"system": "private rubric"}, "originating_request_masked": {"input": "private source"}} - details: Final = { - "standard_logging_object": {**audit, "messages": [], "response": {}}, - "litellm_params": {"proxy_server_request": {"body": {}, "originating_request_masked": audit["originating_request_masked"]}}, - } - logger: Final = CustomLogger() - logger.turn_off_message_logging = True - if callback_only: - redacted: Final = logger.redact_standard_logging_payload_from_model_call_details(details) - assert "classifier_input" not in redacted["standard_logging_object"] - assert "originating_request_masked" not in redacted["standard_logging_object"] - assert "originating_request_masked" not in redacted["litellm_params"]["proxy_server_request"] - assert details["standard_logging_object"]["classifier_input"] == audit["classifier_input"] - assert details["litellm_params"]["proxy_server_request"]["originating_request_masked"] == audit["originating_request_masked"] - else: - perform_redaction(details, result=None) - assert "classifier_input" not in details["standard_logging_object"] - assert "originating_request_masked" not in details["standard_logging_object"] - assert "originating_request_masked" not in details["litellm_params"]["proxy_server_request"] - def test_unredactable_result_is_not_deepcopied(self): """A result shape no branch can redact must not be deepcopied. @@ -882,3 +860,69 @@ class TestRedactStreamingResponsesForCustomLogger: assert result_details is model_call_details assert response_obj.choices[0].message.content == "secret content" + + +@pytest.mark.parametrize("callback_only", [False, True]) +def test_classifier_audit_redaction_removes_both_fields_and_source_carrier(callback_only: bool) -> None: + audit: Final = {"classifier_input": {"system": "private rubric"}, "originating_request_masked": {"input": "private source"}} + standard_payload: Final = { + **audit, + "messages": [{"role": "user", "content": "private prompt"}], + "response": {"choices": [{"message": {"content": "private answer"}}]}, + "model": "classifier", + } + details: Final = { + "standard_logging_object": standard_payload, + "litellm_params": {"proxy_server_request": {"body": {}, "originating_request_masked": audit["originating_request_masked"]}}, + } + logger: Final = CustomLogger() + logger.turn_off_message_logging = True + if callback_only: + redacted: Final = logger.redact_standard_logging_payload_from_model_call_details(details) + assert "classifier_input" not in redacted["standard_logging_object"] + assert "originating_request_masked" not in redacted["standard_logging_object"] + assert "originating_request_masked" not in redacted["litellm_params"]["proxy_server_request"] + assert details["standard_logging_object"]["classifier_input"] == audit["classifier_input"] + assert details["litellm_params"]["proxy_server_request"]["originating_request_masked"] == audit["originating_request_masked"] + else: + perform_redaction(details, result=None) + assert "classifier_input" not in details["standard_logging_object"] + assert "originating_request_masked" not in details["standard_logging_object"] + assert "originating_request_masked" not in details["litellm_params"]["proxy_server_request"] + + assert standard_payload["classifier_input"] == audit["classifier_input"] + assert standard_payload["originating_request_masked"] == audit["originating_request_masked"] + assert standard_payload["messages"][0]["content"] == "private prompt" + assert standard_payload["response"]["choices"][0]["message"]["content"] == "private answer" + + +@pytest.mark.parametrize("excluded", [False, True]) +def test_classifier_callback_redaction_preserves_exclusions(monkeypatch: pytest.MonkeyPatch, excluded: bool) -> None: + monkeypatch.setattr(litellm, "standard_logging_payload_excluded_fields", ["messages", "response"] if excluded else []) + payload: Final = { + "classifier_input": {"system": "private rubric"}, + "originating_request_masked": {"input": "private source"}, + "messages": [{"role": "user", "content": "private prompt"}], + "response": {"choices": [{"message": {"content": "private answer"}}]}, + "model": "classifier", + } + logger: Final = CustomLogger() + logger.turn_off_message_logging = True + redacted: Final = logger.redact_standard_logging_payload_from_model_call_details({"standard_logging_object": payload}) + stored: Final = redacted["standard_logging_object"] + assert "classifier_input" not in stored + assert "originating_request_masked" not in stored + assert stored["model"] == "classifier" + assert ("messages" not in stored) is excluded + assert ("response" not in stored) is excluded + if not excluded: + assert stored["messages"][0]["content"] == "redacted-by-litellm" + assert stored["response"]["choices"][0]["message"]["content"] == "redacted-by-litellm" + + failure_payload: Final = redacted_standard_logging_payload(payload) + assert "classifier_input" not in failure_payload + assert "originating_request_masked" not in failure_payload + assert failure_payload["messages"][0]["content"] == "redacted-by-litellm" + assert failure_payload["response"]["choices"][0]["message"]["content"] == "redacted-by-litellm" + assert payload["classifier_input"] == {"system": "private rubric"} + assert payload["response"]["choices"][0]["message"]["content"] == "private answer"