(feat) Allow enabling logging message / response for specific virtual keys (#7071)

* redact_message_input_output_from_logging

* initialize_standard_callback_dynamic_params

* allow dynamically opting out of redaction

* test_redact_msgs_from_logs_with_dynamic_params

* fix AddTeamCallback

* _get_turn_off_message_logging_from_dynamic_params

* test_global_redaction_with_dynamic_params

* test_dynamic_turn_off_message_logging

* docs Disable/Enable Message redaction

* fix doe qual check

* _get_turn_off_message_logging_from_dynamic_params

docs/my-website/docs/proxy/team_logging.md

@@ -386,3 +386,79 @@ A key is **unhealthy** when the logging callbacks are not setup correctly.
 
 </TabItem>
 </Tabs>
+
+### Disable/Enable Message redaction
+
+Use this to enable prompt logging for specific keys when you have globally disabled it
+
+Example config.yaml with globally disabled prompt logging (message redaction)
+```yaml
+model_list:
+ - model_name: gpt-4o
+    litellm_params:
+      model: gpt-4o
+litellm_settings:
+  callbacks: ["datadog"]
+  turn_off_message_logging: True # 👈 Globally logging prompt / response is disabled
+```
+
+**Enable prompt logging for key**
+
+Set `turn_off_message_logging` to `false` for the key you want to enable prompt logging for. This will override the global `turn_off_message_logging` setting.
+
+```shell
+curl -X POST 'http://0.0.0.0:4000/key/generate' \
+-H 'Authorization: Bearer sk-1234' \
+-H 'Content-Type: application/json' \
+-d '{
+    "metadata": {
+        "logging": [{
+            "callback_name": "datadog",
+            "callback_vars": {
+                "turn_off_message_logging": false # 👈 Enable prompt logging
+            }
+        }]
+    }
+}'
+```
+
+Response from `/key/generate`
+
+```json
+{
+    "key_alias": null,
+    "key": "sk-9v6I-jf9-eYtg_PwM8OKgQ",
+    "metadata": {
+        "logging": [
+            {
+                "callback_name": "datadog",
+                "callback_vars": {
+                    "turn_off_message_logging": false
+                }
+            }
+        ]
+    },
+    "token_id": "a53a33db8c3cf832ceb28565dbb034f19f0acd69ee7f03b7bf6752f9f804081e"
+}
+```
+
+Use key for `/chat/completions` request
+
+This key will log the prompt to the callback specified in the request
+
+```shell
+curl -i http://localhost:4000/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -H "Authorization: Bearer sk-9v6I-jf9-eYtg_PwM8OKgQ" \
+  -d '{
+    "model": "gpt-4o",
+    "messages": [
+      {"role": "user", "content": "hi my name is ishaan what key alias is this"}
+    ]
+  }'
+```
+
+
+
+
+

litellm/litellm_core_utils/litellm_logging.py

@@ -374,7 +374,11 @@ class Logging:
             for param in _supported_callback_params:
                 if param in kwargs:
                     _param_value = kwargs.pop(param)
-                    if _param_value is not None and "os.environ/" in _param_value:
+                    if (
+                        _param_value is not None
+                        and isinstance(_param_value, str)
+                        and "os.environ/" in _param_value
+                    ):
                         _param_value = get_secret_str(secret_name=_param_value)
                     standard_callback_dynamic_params[param] = _param_value  # type: ignore
         return standard_callback_dynamic_params

litellm/litellm_core_utils/redact_messages.py

@@ -12,6 +12,8 @@ from typing import TYPE_CHECKING, Any, Optional
 
 import litellm
 from litellm.integrations.custom_logger import CustomLogger
+from litellm.secret_managers.main import str_to_bool
+from litellm.types.utils import StandardCallbackDynamicParams
 
 if TYPE_CHECKING:
     from litellm.litellm_core_utils.litellm_logging import (
@@ -88,6 +90,8 @@ def redact_message_input_output_from_logging(
     if (
         litellm.turn_off_message_logging is not True
         and request_headers.get("litellm-enable-message-redaction", False) is not True
+        and _get_turn_off_message_logging_from_dynamic_params(model_call_details)
+        is not True
     ):
         return result
 
@@ -96,9 +100,35 @@ def redact_message_input_output_from_logging(
     ):
         return result
 
+    # user has OPTED OUT of message redaction
+    if _get_turn_off_message_logging_from_dynamic_params(model_call_details) is False:
+        return result
+
     return perform_redaction(model_call_details, result)
 
 
+def _get_turn_off_message_logging_from_dynamic_params(
+    model_call_details: dict,
+) -> Optional[bool]:
+    """
+    gets the value of `turn_off_message_logging` from the dynamic params, if it exists.
+
+    handles boolean and string values of `turn_off_message_logging`
+    """
+    standard_callback_dynamic_params: Optional[StandardCallbackDynamicParams] = (
+        model_call_details.get("standard_callback_dynamic_params", None)
+    )
+    if standard_callback_dynamic_params:
+        _turn_off_message_logging = standard_callback_dynamic_params.get(
+            "turn_off_message_logging"
+        )
+        if isinstance(_turn_off_message_logging, bool):
+            return _turn_off_message_logging
+        elif isinstance(_turn_off_message_logging, str):
+            return str_to_bool(_turn_off_message_logging)
+    return None
+
+
 def redact_user_api_key_info(metadata: dict) -> dict:
     """
     removes any user_api_key_info before passing to logging object, if flag set

litellm/proxy/_types.py

@@ -1012,7 +1012,9 @@ class BlockKeyRequest(LiteLLMBase):
 
 class AddTeamCallback(LiteLLMBase):
     callback_name: str
-    callback_type: Literal["success", "failure", "success_and_failure"]
+    callback_type: Optional[Literal["success", "failure", "success_and_failure"]] = (
+        "success_and_failure"
+    )
     callback_vars: Dict[str, str]
 
     @model_validator(mode="before")
@@ -1020,11 +1022,13 @@ class AddTeamCallback(LiteLLMBase):
     def validate_callback_vars(cls, values):
         callback_vars = values.get("callback_vars", {})
         valid_keys = set(StandardCallbackDynamicParams.__annotations__.keys())
-        for key in callback_vars:
+        for key, value in callback_vars.items():
             if key not in valid_keys:
                 raise ValueError(
                     f"Invalid callback variable: {key}. Must be one of {valid_keys}"
                 )
+            if not isinstance(value, str):
+                callback_vars[key] = str(value)
         return values
 
 

litellm/proxy/proxy_config.yaml

@@ -13,6 +13,8 @@ model_list:
 litellm_settings:
   callbacks: ["datadog"] 
 
+  turn_off_message_logging: True
+
 
 router_settings:
   provider_budget_config: 
@@ -35,4 +37,5 @@ router_settings:
   # OPTIONAL: Set Redis Host, Port, and Password if using multiple instance of LiteLLM
   redis_host: os.environ/REDIS_HOST
   redis_port: os.environ/REDIS_PORT
-  redis_password: os.environ/REDIS_PASSWORD
+  redis_password: os.environ/REDIS_PASSWORD
+

litellm/types/utils.py

@@ -1607,6 +1607,9 @@ class StandardCallbackDynamicParams(TypedDict, total=False):
     langsmith_project: Optional[str]
     langsmith_base_url: Optional[str]
 
+    # Logging settings
+    turn_off_message_logging: Optional[bool]  # when true will not log messages
+
 
 class KeyGenerationConfig(TypedDict, total=False):
     required_params: List[

tests/local_testing/test_utils.py

@@ -6,6 +6,8 @@ from unittest import mock
 
 from dotenv import load_dotenv
 
+from litellm.types.utils import StandardCallbackDynamicParams
+
 load_dotenv()
 import os
 
@@ -545,6 +547,92 @@ def test_redact_msgs_from_logs():
     print("Test passed")
 
 
+def test_redact_msgs_from_logs_with_dynamic_params():
+    """
+    Tests redaction behavior based on standard_callback_dynamic_params setting:
+    In all tests litellm.turn_off_message_logging is True
+
+
+    1. When standard_callback_dynamic_params.turn_off_message_logging is False (or not set): No redaction should occur. User has opted out of redaction.
+    2. When standard_callback_dynamic_params.turn_off_message_logging is True: Redaction should occur. User has opted in to redaction.
+    3. standard_callback_dynamic_params.turn_off_message_logging not set, litellm.turn_off_message_logging is True: Redaction should occur.
+    """
+    from litellm.litellm_core_utils.litellm_logging import Logging
+    from litellm.litellm_core_utils.redact_messages import (
+        redact_message_input_output_from_logging,
+    )
+
+    litellm.turn_off_message_logging = True
+    test_content = "I'm LLaMA, an AI assistant developed by Meta AI that can understand and respond to human input in a conversational manner."
+    response_obj = litellm.ModelResponse(
+        choices=[
+            {
+                "finish_reason": "stop",
+                "index": 0,
+                "message": {
+                    "content": test_content,
+                    "role": "assistant",
+                },
+            }
+        ]
+    )
+
+    litellm_logging_obj = Logging(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "hi"}],
+        stream=False,
+        call_type="acompletion",
+        litellm_call_id="1234",
+        start_time=datetime.now(),
+        function_id="1234",
+    )
+
+    # Test Case 1: standard_callback_dynamic_params = False (or not set)
+    standard_callback_dynamic_params = StandardCallbackDynamicParams(
+        turn_off_message_logging=False
+    )
+    litellm_logging_obj.model_call_details["standard_callback_dynamic_params"] = (
+        standard_callback_dynamic_params
+    )
+    _redacted_response_obj = redact_message_input_output_from_logging(
+        result=response_obj,
+        model_call_details=litellm_logging_obj.model_call_details,
+    )
+    # Assert no redaction occurred
+    assert _redacted_response_obj.choices[0].message.content == test_content
+
+    # Test Case 2: standard_callback_dynamic_params = True
+    standard_callback_dynamic_params = StandardCallbackDynamicParams(
+        turn_off_message_logging=True
+    )
+    litellm_logging_obj.model_call_details["standard_callback_dynamic_params"] = (
+        standard_callback_dynamic_params
+    )
+    _redacted_response_obj = redact_message_input_output_from_logging(
+        result=response_obj,
+        model_call_details=litellm_logging_obj.model_call_details,
+    )
+    # Assert redaction occurred
+    assert _redacted_response_obj.choices[0].message.content == "redacted-by-litellm"
+
+    # Test Case 3: standard_callback_dynamic_params does not override litellm.turn_off_message_logging
+    # since litellm.turn_off_message_logging is True redaction should occur
+    standard_callback_dynamic_params = StandardCallbackDynamicParams()
+    litellm_logging_obj.model_call_details["standard_callback_dynamic_params"] = (
+        standard_callback_dynamic_params
+    )
+    _redacted_response_obj = redact_message_input_output_from_logging(
+        result=response_obj,
+        model_call_details=litellm_logging_obj.model_call_details,
+    )
+    # Assert no redaction occurred
+    assert _redacted_response_obj.choices[0].message.content == "redacted-by-litellm"
+
+    # Reset settings
+    litellm.turn_off_message_logging = False
+    print("Test passed")
+
+
 @pytest.mark.parametrize(
     "duration, unit",
     [("7s", "s"), ("7m", "m"), ("7h", "h"), ("7d", "d"), ("7mo", "mo")],

tests/logging_callback_tests/test_logging_redaction_e2e_test.py

@@ -0,0 +1,120 @@
+import io
+import os
+import sys
+
+from typing import Optional
+
+sys.path.insert(0, os.path.abspath("../.."))
+
+import asyncio
+import gzip
+import json
+import logging
+import time
+from unittest.mock import AsyncMock, patch
+
+import pytest
+
+import litellm
+from litellm._logging import verbose_logger
+from litellm.integrations.custom_logger import CustomLogger
+from litellm.types.utils import StandardLoggingPayload
+
+
+class TestCustomLogger(CustomLogger):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.logged_standard_logging_payload: Optional[StandardLoggingPayload] = None
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        standard_logging_payload = kwargs.get("standard_logging_object", None)
+        self.logged_standard_logging_payload = standard_logging_payload
+
+
+@pytest.mark.asyncio
+async def test_global_redaction_on():
+    litellm.turn_off_message_logging = True
+    test_custom_logger = TestCustomLogger()
+    litellm.callbacks = [test_custom_logger]
+    response = await litellm.acompletion(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "hi"}],
+        mock_response="hello",
+    )
+
+    await asyncio.sleep(1)
+    standard_logging_payload = test_custom_logger.logged_standard_logging_payload
+    assert standard_logging_payload is not None
+    assert standard_logging_payload["response"] == "redacted-by-litellm"
+    assert standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm"
+    print(
+        "logged standard logging payload",
+        json.dumps(standard_logging_payload, indent=2),
+    )
+
+
+@pytest.mark.parametrize("turn_off_message_logging", [True, False])
+@pytest.mark.asyncio
+async def test_global_redaction_with_dynamic_params(turn_off_message_logging):
+    litellm.turn_off_message_logging = True
+    test_custom_logger = TestCustomLogger()
+    litellm.callbacks = [test_custom_logger]
+    response = await litellm.acompletion(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "hi"}],
+        turn_off_message_logging=turn_off_message_logging,
+        mock_response="hello",
+    )
+
+    await asyncio.sleep(1)
+    standard_logging_payload = test_custom_logger.logged_standard_logging_payload
+    assert standard_logging_payload is not None
+    print(
+        "logged standard logging payload",
+        json.dumps(standard_logging_payload, indent=2),
+    )
+
+    if turn_off_message_logging is True:
+        assert standard_logging_payload["response"] == "redacted-by-litellm"
+        assert (
+            standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm"
+        )
+    else:
+        assert (
+            standard_logging_payload["response"]["choices"][0]["message"]["content"]
+            == "hello"
+        )
+        assert standard_logging_payload["messages"][0]["content"] == "hi"
+
+
+@pytest.mark.parametrize("turn_off_message_logging", [True, False])
+@pytest.mark.asyncio
+async def test_global_redaction_off_with_dynamic_params(turn_off_message_logging):
+    litellm.turn_off_message_logging = False
+    test_custom_logger = TestCustomLogger()
+    litellm.callbacks = [test_custom_logger]
+    response = await litellm.acompletion(
+        model="gpt-3.5-turbo",
+        messages=[{"role": "user", "content": "hi"}],
+        turn_off_message_logging=turn_off_message_logging,
+        mock_response="hello",
+    )
+
+    await asyncio.sleep(1)
+    standard_logging_payload = test_custom_logger.logged_standard_logging_payload
+    assert standard_logging_payload is not None
+    print(
+        "logged standard logging payload",
+        json.dumps(standard_logging_payload, indent=2),
+    )
+    if turn_off_message_logging is True:
+        assert standard_logging_payload["response"] == "redacted-by-litellm"
+        assert (
+            standard_logging_payload["messages"][0]["content"] == "redacted-by-litellm"
+        )
+    else:
+        assert (
+            standard_logging_payload["response"]["choices"][0]["message"]["content"]
+            == "hello"
+        )
+        assert standard_logging_payload["messages"][0]["content"] == "hi"

tests/proxy_unit_tests/test_proxy_utils.py

@@ -296,6 +296,83 @@ def test_dynamic_logging_metadata_key_and_team_metadata(callback_vars):
         assert "os.environ" not in var
 
 
+@pytest.mark.parametrize(
+    "callback_vars",
+    [
+        {
+            "turn_off_message_logging": True,
+        },
+        {
+            "turn_off_message_logging": False,
+        },
+    ],
+)
+def test_dynamic_turn_off_message_logging(callback_vars):
+    user_api_key_dict = UserAPIKeyAuth(
+        token="6f8688eaff1d37555bb9e9a6390b6d7032b3ab2526ba0152da87128eab956432",
+        key_name="sk-...63Fg",
+        key_alias=None,
+        spend=0.000111,
+        max_budget=None,
+        expires=None,
+        models=[],
+        aliases={},
+        config={},
+        user_id=None,
+        team_id="ishaan-special-team_e02dd54f-f790-4755-9f93-73734f415898",
+        max_parallel_requests=None,
+        metadata={
+            "logging": [
+                {
+                    "callback_name": "datadog",
+                    "callback_vars": callback_vars,
+                }
+            ]
+        },
+        tpm_limit=None,
+        rpm_limit=None,
+        budget_duration=None,
+        budget_reset_at=None,
+        allowed_cache_controls=[],
+        permissions={},
+        model_spend={},
+        model_max_budget={},
+        soft_budget_cooldown=False,
+        litellm_budget_table=None,
+        org_id=None,
+        team_spend=0.000132,
+        team_alias=None,
+        team_tpm_limit=None,
+        team_rpm_limit=None,
+        team_max_budget=None,
+        team_models=[],
+        team_blocked=False,
+        soft_budget=None,
+        team_model_aliases=None,
+        team_member_spend=None,
+        team_member=None,
+        team_metadata={},
+        end_user_id=None,
+        end_user_tpm_limit=None,
+        end_user_rpm_limit=None,
+        end_user_max_budget=None,
+        last_refreshed_at=1726101560.967527,
+        api_key="7c305cc48fe72272700dc0d67dc691c2d1f2807490ef5eb2ee1d3a3ca86e12b1",
+        user_role=LitellmUserRoles.INTERNAL_USER,
+        allowed_model_region=None,
+        parent_otel_span=None,
+        rpm_limit_per_model=None,
+        tpm_limit_per_model=None,
+    )
+    callbacks = _get_dynamic_logging_metadata(user_api_key_dict=user_api_key_dict)
+
+    assert callbacks is not None
+    assert (
+        callbacks.callback_vars["turn_off_message_logging"]
+        == callback_vars["turn_off_message_logging"]
+    )
+
+
 @pytest.mark.parametrize(
     "allow_client_side_credentials, expect_error", [(True, False), (False, True)]
 )