From 548e4f53f8a759db419b23f9d8d2bb797d2ec6dd Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Mon, 22 Jul 2024 20:58:02 -0700 Subject: [PATCH] feat(redact_messages.py): allow remove sensitive key information before passing to logging integration --- docs/my-website/docs/proxy/logging.md | 14 +++++++++ litellm/__init__.py | 1 + litellm/integrations/langfuse.py | 3 ++ litellm/integrations/logfire_logger.py | 12 ++++++-- litellm/integrations/opentelemetry.py | 5 +++- litellm/litellm_core_utils/redact_messages.py | 30 +++++++++++++++++++ litellm/proxy/_new_secret_config.yaml | 4 +++ 7 files changed, 65 insertions(+), 4 deletions(-) diff --git a/docs/my-website/docs/proxy/logging.md b/docs/my-website/docs/proxy/logging.md index 680264ec7..5314395cc 100644 --- a/docs/my-website/docs/proxy/logging.md +++ b/docs/my-website/docs/proxy/logging.md @@ -48,6 +48,20 @@ A number of these headers could be useful for troubleshooting, but the `x-litellm-call-id` is the one that is most useful for tracking a request across components in your system, including in logging tools. +## Redacting UserAPIKeyInfo + +Redact information about the user api key (hashed token, user_id, team id, etc.), from logs. + +Currently supported for Langfuse, OpenTelemetry, Logfire, ArizeAI logging. + +```yaml +litellm_settings: + callbacks: ["langfuse"] + redact_user_api_key_info: true +``` + +Removes any field with `user_api_key_*` from metadata. + ## Logging Proxy Input/Output - Langfuse We will use the `--config` to set `litellm.success_callback = ["langfuse"]` this will log all successfull LLM calls to langfuse. Make sure to set `LANGFUSE_PUBLIC_KEY` and `LANGFUSE_SECRET_KEY` in your environment diff --git a/litellm/__init__.py b/litellm/__init__.py index 78043b906..9982097b1 100644 --- a/litellm/__init__.py +++ b/litellm/__init__.py @@ -76,6 +76,7 @@ post_call_rules: List[Callable] = [] turn_off_message_logging: Optional[bool] = False log_raw_request_response: bool = False redact_messages_in_exceptions: Optional[bool] = False +redact_user_api_key_info: Optional[bool] = False store_audit_logs = False # Enterprise feature, allow users to see audit logs ## end of callbacks ############# diff --git a/litellm/integrations/langfuse.py b/litellm/integrations/langfuse.py index 0647afabc..0217f7458 100644 --- a/litellm/integrations/langfuse.py +++ b/litellm/integrations/langfuse.py @@ -8,6 +8,7 @@ from packaging.version import Version import litellm from litellm._logging import verbose_logger +from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info class LangFuseLogger: @@ -382,6 +383,8 @@ class LangFuseLogger: mask_input = clean_metadata.pop("mask_input", False) mask_output = clean_metadata.pop("mask_output", False) + clean_metadata = redact_user_api_key_info(metadata=clean_metadata) + if trace_name is None and existing_trace_id is None: # just log `litellm-{call_type}` as the trace name ## DO NOT SET TRACE_NAME if trace-id set. this can lead to overwriting of past traces. diff --git a/litellm/integrations/logfire_logger.py b/litellm/integrations/logfire_logger.py index b4ab00820..fa4ab7bd5 100644 --- a/litellm/integrations/logfire_logger.py +++ b/litellm/integrations/logfire_logger.py @@ -1,17 +1,21 @@ #### What this does #### # On success + failure, log events to Logfire -import dotenv, os +import os + +import dotenv dotenv.load_dotenv() # Loading env variables using dotenv import traceback import uuid -from litellm._logging import print_verbose, verbose_logger - from enum import Enum from typing import Any, Dict, NamedTuple + from typing_extensions import LiteralString +from litellm._logging import print_verbose, verbose_logger +from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info + class SpanConfig(NamedTuple): message_template: LiteralString @@ -135,6 +139,8 @@ class LogfireLogger: else: clean_metadata[key] = value + clean_metadata = redact_user_api_key_info(metadata=clean_metadata) + # Build the initial payload payload = { "id": id, diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index bc58efad3..c47911b4f 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Dict, Optional, Union import litellm from litellm._logging import verbose_logger from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.redact_messages import redact_user_api_key_info from litellm.types.services import ServiceLoggerPayload if TYPE_CHECKING: @@ -315,7 +316,9 @@ class OpenTelemetry(CustomLogger): ############################################# metadata = litellm_params.get("metadata", {}) or {} - for key, value in metadata.items(): + clean_metadata = redact_user_api_key_info(metadata=metadata) + + for key, value in clean_metadata.items(): if self.is_primitive(value): span.set_attribute("metadata.{}".format(key), value) diff --git a/litellm/litellm_core_utils/redact_messages.py b/litellm/litellm_core_utils/redact_messages.py index 378c46ba0..7f342e271 100644 --- a/litellm/litellm_core_utils/redact_messages.py +++ b/litellm/litellm_core_utils/redact_messages.py @@ -87,3 +87,33 @@ def redact_message_input_output_from_logging( # by default return result return result + + +def redact_user_api_key_info(metadata: dict) -> dict: + """ + removes any user_api_key_info before passing to logging object, if flag set + + Usage: + + SDK + ```python + litellm.redact_user_api_key_info = True + ``` + + PROXY: + ```yaml + litellm_settings: + redact_user_api_key_info: true + ``` + """ + if litellm.redact_user_api_key_info is not True: + return metadata + + new_metadata = {} + for k, v in metadata.items(): + if isinstance(k, str) and k.startswith("user_api_key"): + pass + else: + new_metadata[k] = v + + return new_metadata diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 81244f0fa..1b5c17246 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -3,3 +3,7 @@ model_list: litellm_params: model: groq/llama3-groq-70b-8192-tool-use-preview api_key: os.environ/GROQ_API_KEY + +litellm_settings: + callbacks: ["logfire"] + redact_user_api_key_info: true \ No newline at end of file