diff --git a/docs/my-website/docs/proxy/team_logging.md b/docs/my-website/docs/proxy/team_logging.md index 25b3679946..779a6516b4 100644 --- a/docs/my-website/docs/proxy/team_logging.md +++ b/docs/my-website/docs/proxy/team_logging.md @@ -386,3 +386,79 @@ A key is **unhealthy** when the logging callbacks are not setup correctly. + +### 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"} + ] + }' +``` + + + + + diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index b05e572dd1..fed1cc2863 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/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 diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 8dad714393..3be27c44df 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/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 diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 3e991edb6b..6bb48bf8e6 100644 --- a/litellm/proxy/_types.py +++ b/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 diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml index d93f6bb202..d93c7bc1e5 100644 --- a/litellm/proxy/proxy_config.yaml +++ b/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 \ No newline at end of file + redis_password: os.environ/REDIS_PASSWORD + diff --git a/litellm/types/utils.py b/litellm/types/utils.py index d4fb3a97e7..168af87386 100644 --- a/litellm/types/utils.py +++ b/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[ diff --git a/tests/local_testing/test_utils.py b/tests/local_testing/test_utils.py index 54c2e8c6c3..1682235254 100644 --- a/tests/local_testing/test_utils.py +++ b/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")], diff --git a/tests/logging_callback_tests/test_logging_redaction_e2e_test.py b/tests/logging_callback_tests/test_logging_redaction_e2e_test.py new file mode 100644 index 0000000000..ea821e7880 --- /dev/null +++ b/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" diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 7910cc7ea9..5d2979de13 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/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)] )