From 8c5fb6f539b3fee5ea5eac3892fe7908902db4dc Mon Sep 17 00:00:00 2001 From: Ishaan Jaff Date: Mon, 23 Jun 2025 14:32:05 -0700 Subject: [PATCH] [Feat] Enterprise - Allow dynamically disabling callbacks in request headers (#11985) * Add support for disabling callbacks via x-litellm-disable-callbacks header * add _is_callback_disabled_via_headers * add get_proxy_server_request_headers * _is_callback_disabled_via_headers * X_LITELLM_DISABLE_CALLBACKS * add EnterpriseCallbackControls * use EnterpriseCallbackControls * use CustomLoggerRegistry * use CustomLoggerRegistry * CustomLoggerRegistry * EnterpriseCallbackControls * TestEnterpriseCallbackControls * docs clean up * docs dynamic callbacks * doc fixes * fix code qa checks * fix CustomLoggerRegistry --------- Co-authored-by: Cursor Agent --- docs/my-website/docs/proxy/dynamic_logging.md | 194 ++++++++++++++++++ docs/my-website/docs/proxy/logging.md | 108 ++++++++-- docs/my-website/sidebars.js | 3 +- .../enterprise_callbacks/callback_controls.py | 66 ++++++ litellm/constants.py | 3 + .../custom_logger_registry.py | 133 ++++++++++++ litellm/litellm_core_utils/litellm_logging.py | 28 +++ .../litellm_core_utils/llm_request_utils.py | 15 ++ .../test_unit_tests_init_callbacks.py | 84 +------- .../test_callback_controls.py | 176 ++++++++++++++++ 10 files changed, 709 insertions(+), 101 deletions(-) create mode 100644 docs/my-website/docs/proxy/dynamic_logging.md create mode 100644 enterprise/litellm_enterprise/enterprise_callbacks/callback_controls.py create mode 100644 litellm/litellm_core_utils/custom_logger_registry.py create mode 100644 tests/test_litellm/enterprise/enterprise_callbacks/test_callback_controls.py diff --git a/docs/my-website/docs/proxy/dynamic_logging.md b/docs/my-website/docs/proxy/dynamic_logging.md new file mode 100644 index 00000000000..484c0c1eacb --- /dev/null +++ b/docs/my-website/docs/proxy/dynamic_logging.md @@ -0,0 +1,194 @@ +import Image from '@theme/IdealImage'; +import Tabs from '@theme/Tabs'; +import TabItem from '@theme/TabItem'; + + +# Dynamic Callback Management + +:::info + +This is an enterprise feature. + +[Get started with LiteLLM Enterprise](https://www.litellm.ai/enterprise) + +::: + +LiteLLM's dynamic callback management enables teams to control logging behavior on a per-request basis without requiring central infrastructure changes. This is essential for organizations managing large-scale service ecosystems where: + +- **Teams manage their own compliance** - Services can handle sensitive data appropriately without central oversight +- **Decentralized responsibility** - Each team controls their data handling while using shared infrastructure + +You can disable callbacks by passing the `x-litellm-disable-callbacks` header with your requests, giving teams granular control over where their data is logged. + +## Quick Start + + + + +```bash +curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Content-Type: application/json' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'x-litellm-disable-callbacks: langfuse' \ + --data '{ + "model": "claude-sonnet-4-20250514", + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ] +}' +``` + + + + +```bash +curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Content-Type: application/json' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'x-litellm-disable-callbacks: langfuse,datadog' \ + --data '{ + "model": "claude-sonnet-4-20250514", + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ] +}' +``` + + + + +## 1. View Active Logging Callbacks + +Before disabling callbacks, you can view all currently enabled callbacks on your proxy. + +### Request + +```bash +curl --location 'http://0.0.0.0:4000/callbacks/list' \ + --header 'Authorization: Bearer sk-1234' +``` + +### Response + +```json +{ + "callbacks": [ + "langfuse", + "datadog", + "prometheus", + "slack_alerting" + ] +} +``` + +## 2. Disable a Single Callback + +Use the `x-litellm-disable-callbacks` header to disable specific callbacks for individual requests. + + + + +```bash +curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Content-Type: application/json' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'x-litellm-disable-callbacks: langfuse' \ + --data '{ + "model": "claude-sonnet-4-20250514", + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ] +}' +``` + + + + +```python +import openai + +client = openai.OpenAI( + api_key="sk-1234", + base_url="http://0.0.0.0:4000" +) + +response = client.chat.completions.create( + model="claude-sonnet-4-20250514", + messages=[ + { + "role": "user", + "content": "what llm are you" + } + ], + extra_headers={ + "x-litellm-disable-callbacks": "langfuse" + } +) + +print(response) +``` + + + + +## 3. Disable Multiple Callbacks + +You can disable multiple callbacks by providing a comma-separated list in the header. + + + + +```bash +curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Content-Type: application/json' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'x-litellm-disable-callbacks: langfuse,datadog,prometheus' \ + --data '{ + "model": "claude-sonnet-4-20250514", + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ] +}' +``` + + + + +```python +import openai + +client = openai.OpenAI( + api_key="sk-1234", + base_url="http://0.0.0.0:4000" +) + +response = client.chat.completions.create( + model="claude-sonnet-4-20250514", + messages=[ + { + "role": "user", + "content": "what llm are you" + } + ], + extra_headers={ + "x-litellm-disable-callbacks": "langfuse,datadog,prometheus" + } +) + +print(response) +``` + + + diff --git a/docs/my-website/docs/proxy/logging.md b/docs/my-website/docs/proxy/logging.md index 3bf5ed12300..7ec9080dfdd 100644 --- a/docs/my-website/docs/proxy/logging.md +++ b/docs/my-website/docs/proxy/logging.md @@ -56,27 +56,6 @@ components in your system, including in logging tools. ## Logging Features -### Conditional Logging by Virtual Keys, Teams - -Use this to: -1. Conditionally enable logging for some virtual keys/teams -2. Set different logging providers for different virtual keys/teams - -[👉 **Get Started** - Team/Key Based Logging](team_logging) - - -### 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 -``` - ### Redact Messages, Response Content @@ -172,6 +151,18 @@ curl -L -X POST 'http://0.0.0.0:4000/v1/chat/completions' \ +### 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 +``` + ### Disable Message Redaction If you have `litellm.turn_on_message_logging` turned on, you can override it for specific requests by @@ -269,6 +260,81 @@ print(response) LiteLLM.Info: "no-log request, skipping logging" ``` +### ✨ Dynamically Disable specific callbacks + +:::info + +This is an enterprise feature. + +[Proceed with LiteLLM Enterprise](https://www.litellm.ai/enterprise) + +::: + +For some use cases, you may want to disable specific callbacks for a request. You can do this by passing `x-litellm-disable-callbacks: ` in the request headers. + +Send the list of callbacks to disable in the request header `x-litellm-disable-callbacks`. + + + + +```bash +curl --location 'http://0.0.0.0:4000/chat/completions' \ + --header 'Content-Type: application/json' \ + --header 'Authorization: Bearer sk-1234' \ + --header 'x-litellm-disable-callbacks: langfuse' \ + --data '{ + "model": "claude-sonnet-4-20250514", + "messages": [ + { + "role": "user", + "content": "what llm are you" + } + ] +}' +``` + + + + +```python +import openai + +client = openai.OpenAI( + api_key="sk-1234", + base_url="http://0.0.0.0:4000" +) + +response = client.chat.completions.create( + model="claude-sonnet-4-20250514", + messages=[ + { + "role": "user", + "content": "what llm are you" + } + ], + extra_headers={ + "x-litellm-disable-callbacks": "langfuse" + } +) + +print(response) +``` + + + + + +### ✨ Conditional Logging by Virtual Keys, Teams + +Use this to: +1. Conditionally enable logging for some virtual keys/teams +2. Set different logging providers for different virtual keys/teams + +[👉 **Get Started** - Team/Key Based Logging](team_logging) + + + + ## What gets logged? diff --git a/docs/my-website/sidebars.js b/docs/my-website/sidebars.js index ece07984d34..0ab4a6b3e4f 100644 --- a/docs/my-website/sidebars.js +++ b/docs/my-website/sidebars.js @@ -198,7 +198,8 @@ const sidebars = { items: [ "proxy/logging", "proxy/logging_spec", - "proxy/team_logging" + "proxy/team_logging", + "proxy/dynamic_logging" ], }, diff --git a/enterprise/litellm_enterprise/enterprise_callbacks/callback_controls.py b/enterprise/litellm_enterprise/enterprise_callbacks/callback_controls.py new file mode 100644 index 00000000000..c4316012738 --- /dev/null +++ b/enterprise/litellm_enterprise/enterprise_callbacks/callback_controls.py @@ -0,0 +1,66 @@ +import litellm +from litellm._logging import verbose_logger +from litellm.constants import X_LITELLM_DISABLE_CALLBACKS +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.llm_request_utils import ( + get_proxy_server_request_headers, +) +from litellm.proxy._types import CommonProxyErrors + + +class EnterpriseCallbackControls: + @staticmethod + def is_callback_disabled_via_headers( + callback: litellm.CALLBACK_TYPES, litellm_params: dict + ) -> bool: + """ + Check if a callback is disabled via the x-litellm-disable-callbacks header. + + Args: + callback: The callback to check (can be string, CustomLogger instance, or callable) + litellm_params: Parameters containing proxy server request info + + Returns: + bool: True if the callback should be disabled, False otherwise + """ + from litellm.litellm_core_utils.custom_logger_registry import ( + CustomLoggerRegistry, + ) + + try: + request_headers = get_proxy_server_request_headers(litellm_params) + disabled_callbacks = request_headers.get(X_LITELLM_DISABLE_CALLBACKS, None) + verbose_logger.debug(f"Dynamically disabled callbacks from {X_LITELLM_DISABLE_CALLBACKS}: {disabled_callbacks}") + verbose_logger.debug(f"Checking if {callback} is disabled via headers. Disable callbacks from headers: {disabled_callbacks}") + if disabled_callbacks is not None: + ######################################################### + # premium user check + ######################################################### + if not EnterpriseCallbackControls._premium_user_check(): + return False + ######################################################### + disabled_callbacks = set([cb.strip().lower() for cb in disabled_callbacks.split(",")]) + if isinstance(callback, str): + if callback.lower() in disabled_callbacks: + verbose_logger.debug(f"Not logging to {callback} because it is disabled via {X_LITELLM_DISABLE_CALLBACKS}") + return True + elif isinstance(callback, CustomLogger): + # get the string name of the callback + callback_str = CustomLoggerRegistry.get_callback_str_from_class_type(callback.__class__) + if callback_str is not None and callback_str.lower() in disabled_callbacks: + verbose_logger.debug(f"Not logging to {callback_str} because it is disabled via {X_LITELLM_DISABLE_CALLBACKS}") + return True + return False + except Exception as e: + verbose_logger.debug( + f"Error checking disabled callbacks header: {str(e)}" + ) + return False + + @staticmethod + def _premium_user_check(): + from litellm.proxy.proxy_server import premium_user + if premium_user: + return True + verbose_logger.warning(f"Disabling callbacks using request headers is an enterprise feature. {CommonProxyErrors.not_premium_user.value}") + return False \ No newline at end of file diff --git a/litellm/constants.py b/litellm/constants.py index c96e1dfaba8..61eeb615069 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -693,6 +693,9 @@ PROMETHEUS_BUDGET_METRICS_REFRESH_INTERVAL_MINUTES = int( MCP_TOOL_NAME_PREFIX = "mcp_tool" MAXIMUM_TRACEBACK_LINES_TO_LOG = int(os.getenv("MAXIMUM_TRACEBACK_LINES_TO_LOG", 100)) +# Headers to control callbacks +X_LITELLM_DISABLE_CALLBACKS = "x-litellm-disable-callbacks" + ########################### LiteLLM Proxy Specific Constants ########################### ######################################################################################## MAX_SPENDLOG_ROWS_TO_QUERY = int( diff --git a/litellm/litellm_core_utils/custom_logger_registry.py b/litellm/litellm_core_utils/custom_logger_registry.py new file mode 100644 index 00000000000..20dc2b0c903 --- /dev/null +++ b/litellm/litellm_core_utils/custom_logger_registry.py @@ -0,0 +1,133 @@ +""" +Registry mapping the callback class string to the class type. + +This is used to get the class type from the callback class string. + +Example: + "datadog" -> DataDogLogger + "prometheus" -> PrometheusLogger +""" + +from litellm.integrations.agentops import AgentOps +from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook +from litellm.integrations.argilla import ArgillaLogger +from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger +from litellm.integrations.braintrust_logging import BraintrustLogger +from litellm.integrations.datadog.datadog import DataDogLogger +from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger +from litellm.integrations.deepeval import DeepEvalLogger +from litellm.integrations.galileo import GalileoObserve +from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger +from litellm.integrations.gcs_pubsub.pub_sub import GcsPubSubLogger +from litellm.integrations.humanloop import HumanloopLogger +from litellm.integrations.lago import LagoLogger +from litellm.integrations.langfuse.langfuse_prompt_management import ( + LangfusePromptManagement, +) +from litellm.integrations.langsmith import LangsmithLogger +from litellm.integrations.literal_ai import LiteralAILogger +from litellm.integrations.mlflow import MlflowLogger +from litellm.integrations.openmeter import OpenMeterLogger +from litellm.integrations.opentelemetry import OpenTelemetry +from litellm.integrations.opik.opik import OpikLogger +from litellm.integrations.prometheus import PrometheusLogger +from litellm.integrations.s3_v2 import S3Logger +from litellm.integrations.vector_stores.bedrock_vector_store import BedrockVectorStore +from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHandler + + +class CustomLoggerRegistry: + """ + Registry mapping the callback class string to the class type. + """ + CALLBACK_CLASS_STR_TO_CLASS_TYPE = { + "lago": LagoLogger, + "openmeter": OpenMeterLogger, + "braintrust": BraintrustLogger, + "galileo": GalileoObserve, + "langsmith": LangsmithLogger, + "literalai": LiteralAILogger, + "prometheus": PrometheusLogger, + "datadog": DataDogLogger, + "datadog_llm_observability": DataDogLLMObsLogger, + "gcs_bucket": GCSBucketLogger, + "opik": OpikLogger, + "argilla": ArgillaLogger, + "opentelemetry": OpenTelemetry, + "azure_storage": AzureBlobStorageLogger, + "humanloop": HumanloopLogger, + # OTEL compatible loggers + "logfire": OpenTelemetry, + "arize": OpenTelemetry, + "langfuse_otel": OpenTelemetry, + "arize_phoenix": OpenTelemetry, + "langtrace": OpenTelemetry, + "mlflow": MlflowLogger, + "langfuse": LangfusePromptManagement, + "otel": OpenTelemetry, + "gcs_pubsub": GcsPubSubLogger, + "anthropic_cache_control_hook": AnthropicCacheControlHook, + "agentops": AgentOps, + "bedrock_vector_store": BedrockVectorStore, + "deepeval": DeepEvalLogger, + "s3_v2": S3Logger, + "dynamic_rate_limiter": _PROXY_DynamicRateLimitHandler, + } + + try: + from litellm_enterprise.enterprise_callbacks.generic_api_callback import ( + GenericAPILogger, + ) + from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import ( + PagerDutyAlerting, + ) + from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ( + ResendEmailLogger, + ) + from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import ( + SMTPEmailLogger, + ) + + enterprise_loggers = { + "pagerduty": PagerDutyAlerting, + "generic_api": GenericAPILogger, + "resend_email": ResendEmailLogger, + "smtp_email": SMTPEmailLogger, + } + CALLBACK_CLASS_STR_TO_CLASS_TYPE.update(enterprise_loggers) + except ImportError: + pass # enterprise not installed + + @classmethod + def get_callback_str_from_class_type(cls, class_type: type) -> str | None: + """ + Get the callback string from the class type. + + Args: + class_type: The class type to find the string for + + Returns: + str: The callback string, or None if not found + """ + for callback_str, callback_class in cls.CALLBACK_CLASS_STR_TO_CLASS_TYPE.items(): + if callback_class == class_type: + return callback_str + return None + + @classmethod + def get_all_callback_strs_from_class_type(cls, class_type: type) -> list[str]: + """ + Get all callback strings that map to the same class type. + Some class types (like OpenTelemetry) have multiple string mappings. + + Args: + class_type: The class type to find all strings for + + Returns: + list: List of callback strings that map to the class type + """ + callback_strs: list[str] = [] + for callback_str, callback_class in cls.CALLBACK_CLASS_STR_TO_CLASS_TYPE.items(): + if callback_class == class_type: + callback_strs.append(callback_str) + return callback_strs \ No newline at end of file diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 4523342a091..72589d57e34 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -147,6 +147,9 @@ from .initialize_dynamic_callback_params import ( from .specialty_caches.dynamic_logging_cache import DynamicLoggingCache try: + from litellm_enterprise.enterprise_callbacks.callback_controls import ( + EnterpriseCallbackControls, + ) from litellm_enterprise.enterprise_callbacks.generic_api_callback import ( GenericAPILogger, ) @@ -174,6 +177,7 @@ except Exception as e: ResendEmailLogger = CustomLogger # type: ignore SMTPEmailLogger = CustomLogger # type: ignore PagerDutyAlerting = CustomLogger # type: ignore + EnterpriseCallbackControls = None # type: ignore EnterpriseStandardLoggingPayloadSetupVAR = None _in_memory_loggers: List[Any] = [] @@ -1217,6 +1221,14 @@ class Logging(LiteLLMLoggingBaseClass): f"no-log request, skipping logging for {event_hook} event" ) return False + + # Check for dynamically disabled callbacks via headers + if EnterpriseCallbackControls is not None and EnterpriseCallbackControls.is_callback_disabled_via_headers(callback, litellm_params): + verbose_logger.debug( + f"Callback {callback} disabled via x-litellm-disable-callbacks header for {event_hook} event" + ) + return False + return True def _update_completion_start_time(self, completion_start_time: datetime.datetime): @@ -2246,6 +2258,14 @@ class Logging(LiteLLMLoggingBaseClass): self.has_run_logging(event_type="sync_failure") for callback in callbacks: try: + litellm_params = self.model_call_details.get("litellm_params", {}) + should_run = self.should_run_callback( + callback=callback, + litellm_params=litellm_params, + event_hook="failure_handler", + ) + if not should_run: + continue if callback == "lunary" and lunaryLogger is not None: print_verbose("reaches lunary for logging error!") @@ -2427,6 +2447,14 @@ class Logging(LiteLLMLoggingBaseClass): self.has_run_logging(event_type="async_failure") for callback in callbacks: try: + litellm_params = self.model_call_details.get("litellm_params", {}) + should_run = self.should_run_callback( + callback=callback, + litellm_params=litellm_params, + event_hook="async_failure_handler", + ) + if not should_run: + continue if isinstance(callback, CustomLogger): # custom logger class await callback.async_log_failure_event( kwargs=self.model_call_details, diff --git a/litellm/litellm_core_utils/llm_request_utils.py b/litellm/litellm_core_utils/llm_request_utils.py index 50dbdc5536e..89f5728979f 100644 --- a/litellm/litellm_core_utils/llm_request_utils.py +++ b/litellm/litellm_core_utils/llm_request_utils.py @@ -66,3 +66,18 @@ def pick_cheapest_chat_models_from_llm_provider(custom_llm_provider: str, n=1): # Return the top n cheapest models return [model for model, _ in model_costs[:n]] + +def get_proxy_server_request_headers(litellm_params: Optional[dict]) -> dict: + """ + Get the `proxy_server_request` headers from the litellm_params.\ + + Use this if you want to access the request headers made to LiteLLM proxy server. + """ + if litellm_params is None: + return {} + + proxy_request_headers = ( + litellm_params.get("proxy_server_request", {}).get("headers", {}) or {} + ) + + return proxy_request_headers \ No newline at end of file diff --git a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py index 74165ccdb7d..aff622201fb 100644 --- a/tests/logging_callback_tests/test_unit_tests_init_callbacks.py +++ b/tests/logging_callback_tests/test_unit_tests_init_callbacks.py @@ -16,48 +16,10 @@ import asyncio import logging from litellm._logging import verbose_logger from prometheus_client import REGISTRY, CollectorRegistry - -from litellm.integrations.lago import LagoLogger -from litellm.integrations.deepeval import DeepEvalLogger -from litellm.integrations.openmeter import OpenMeterLogger -from litellm.integrations.braintrust_logging import BraintrustLogger -from litellm.integrations.galileo import GalileoObserve -from litellm.integrations.langsmith import LangsmithLogger -from litellm.integrations.literal_ai import LiteralAILogger -from litellm.integrations.prometheus import PrometheusLogger -from litellm.integrations.datadog.datadog import DataDogLogger -from litellm.integrations.datadog.datadog_llm_obs import DataDogLLMObsLogger -from litellm.integrations.gcs_bucket.gcs_bucket import GCSBucketLogger -from litellm.integrations.gcs_pubsub.pub_sub import GcsPubSubLogger -from litellm.integrations.opik.opik import OpikLogger -from litellm.integrations.opentelemetry import OpenTelemetry -from litellm.integrations.mlflow import MlflowLogger -from litellm.integrations.argilla import ArgillaLogger -from litellm.integrations.deepeval.deepeval import DeepEvalLogger -from litellm.integrations.s3_v2 import S3Logger -from litellm.integrations.langfuse.langfuse_otel import LangfuseOtelLogger -from litellm.integrations.anthropic_cache_control_hook import AnthropicCacheControlHook -from litellm.integrations.vector_stores.bedrock_vector_store import BedrockVectorStore -from litellm.integrations.langfuse.langfuse_prompt_management import ( - LangfusePromptManagement, -) -from litellm.integrations.azure_storage.azure_storage import AzureBlobStorageLogger -from litellm.integrations.agentops import AgentOps -from litellm.integrations.humanloop import HumanloopLogger -from litellm.proxy.hooks.dynamic_rate_limiter import _PROXY_DynamicRateLimitHandler -from litellm_enterprise.enterprise_callbacks.generic_api_callback import ( - GenericAPILogger, -) -from litellm_enterprise.enterprise_callbacks.send_emails.resend_email import ( - ResendEmailLogger, -) -from litellm_enterprise.enterprise_callbacks.send_emails.smtp_email import ( - SMTPEmailLogger, -) -from litellm_enterprise.enterprise_callbacks.pagerduty.pagerduty import ( - PagerDutyAlerting, -) from unittest.mock import patch +from litellm.litellm_core_utils.custom_logger_registry import ( + CustomLoggerRegistry, +) # clear prometheus collectors / registry collectors = list(REGISTRY._collector_to_names.keys()) @@ -65,43 +27,7 @@ for collector in collectors: REGISTRY.unregister(collector) ###################################### -callback_class_str_to_classType = { - "lago": LagoLogger, - "openmeter": OpenMeterLogger, - "braintrust": BraintrustLogger, - "galileo": GalileoObserve, - "langsmith": LangsmithLogger, - "literalai": LiteralAILogger, - "prometheus": PrometheusLogger, - "datadog": DataDogLogger, - "datadog_llm_observability": DataDogLLMObsLogger, - "gcs_bucket": GCSBucketLogger, - "opik": OpikLogger, - "argilla": ArgillaLogger, - "opentelemetry": OpenTelemetry, - "azure_storage": AzureBlobStorageLogger, - "humanloop": HumanloopLogger, - # OTEL compatible loggers - "logfire": OpenTelemetry, - "arize": OpenTelemetry, - "langfuse_otel": OpenTelemetry, - "arize_phoenix": OpenTelemetry, - "langtrace": OpenTelemetry, - "mlflow": MlflowLogger, - "langfuse": LangfusePromptManagement, - "otel": OpenTelemetry, - "pagerduty": PagerDutyAlerting, - "gcs_pubsub": GcsPubSubLogger, - "anthropic_cache_control_hook": AnthropicCacheControlHook, - "agentops": AgentOps, - "bedrock_vector_store": BedrockVectorStore, - "generic_api": GenericAPILogger, - "resend_email": ResendEmailLogger, - "smtp_email": SMTPEmailLogger, - "deepeval": DeepEvalLogger, - "s3_v2": S3Logger, - "langfuse_otel": OpenTelemetry, -} + expected_env_vars = { "LAGO_API_KEY": "api_key", @@ -215,7 +141,7 @@ async def use_callback_in_llm_call( await asyncio.sleep(0.5) - expected_class = callback_class_str_to_classType[callback] + expected_class = CustomLoggerRegistry.CALLBACK_CLASS_STR_TO_CLASS_TYPE[callback] if used_in == "callbacks": assert isinstance(litellm._async_success_callback[0], expected_class) diff --git a/tests/test_litellm/enterprise/enterprise_callbacks/test_callback_controls.py b/tests/test_litellm/enterprise/enterprise_callbacks/test_callback_controls.py new file mode 100644 index 00000000000..a06911170c0 --- /dev/null +++ b/tests/test_litellm/enterprise/enterprise_callbacks/test_callback_controls.py @@ -0,0 +1,176 @@ +import unittest.mock as mock +from unittest.mock import MagicMock, patch + +import pytest + +from enterprise.litellm_enterprise.enterprise_callbacks.callback_controls import ( + EnterpriseCallbackControls, +) +from litellm.constants import X_LITELLM_DISABLE_CALLBACKS +from litellm.integrations.custom_logger import CustomLogger +from litellm.integrations.datadog.datadog import DataDogLogger +from litellm.integrations.langfuse.langfuse_prompt_management import ( + LangfusePromptManagement, +) +from litellm.integrations.s3_v2 import S3Logger + + +class TestEnterpriseCallbackControls: + + @pytest.fixture + def mock_premium_user(self): + """Fixture to mock premium user check as True""" + with patch.object(EnterpriseCallbackControls, '_premium_user_check', return_value=True): + yield + + @pytest.fixture + def mock_non_premium_user(self): + """Fixture to mock premium user check as False""" + with patch.object(EnterpriseCallbackControls, '_premium_user_check', return_value=False): + yield + + @pytest.fixture + def mock_request_headers(self): + """Fixture to mock get_proxy_server_request_headers""" + with patch('enterprise.litellm_enterprise.enterprise_callbacks.callback_controls.get_proxy_server_request_headers') as mock_headers: + yield mock_headers + + def test_callback_disabled_langfuse_string(self, mock_premium_user, mock_request_headers): + """Test that 'langfuse' string callback is disabled when specified in headers""" + mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "langfuse"} + litellm_params = {"proxy_server_request": {"url": "test"}} + + result = EnterpriseCallbackControls.is_callback_disabled_via_headers("langfuse", litellm_params) + assert result is True + + def test_callback_disabled_langfuse_customlogger(self, mock_premium_user, mock_request_headers): + """Test that LangfusePromptManagement CustomLogger instance is disabled when 'langfuse' specified in headers""" + mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "langfuse"} + litellm_params = {"proxy_server_request": {"url": "test"}} + + langfuse_logger = LangfusePromptManagement() + result = EnterpriseCallbackControls.is_callback_disabled_via_headers(langfuse_logger, litellm_params) + assert result is True + + def test_callback_disabled_s3_v2_string(self, mock_premium_user, mock_request_headers): + """Test that 's3_v2' string callback is disabled when specified in headers""" + mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "s3_v2"} + litellm_params = {"proxy_server_request": {"url": "test"}} + + result = EnterpriseCallbackControls.is_callback_disabled_via_headers("s3_v2", litellm_params) + assert result is True + + def test_callback_disabled_s3_v2_customlogger(self, mock_premium_user, mock_request_headers): + """Test that S3Logger CustomLogger instance is disabled when 's3_v2' specified in headers""" + mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "s3_v2"} + litellm_params = {"proxy_server_request": {"url": "test"}} + + # Mock S3Logger to avoid async initialization issues + with patch('litellm.integrations.s3_v2.S3Logger.__init__', return_value=None): + s3_logger = S3Logger() + result = EnterpriseCallbackControls.is_callback_disabled_via_headers(s3_logger, litellm_params) + assert result is True + + def test_callback_disabled_datadog_string(self, mock_premium_user, mock_request_headers): + """Test that 'datadog' string callback is disabled when specified in headers""" + mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "datadog"} + litellm_params = {"proxy_server_request": {"url": "test"}} + + result = EnterpriseCallbackControls.is_callback_disabled_via_headers("datadog", litellm_params) + assert result is True + + def test_callback_disabled_datadog_customlogger(self, mock_premium_user, mock_request_headers): + """Test that DataDogLogger CustomLogger instance is disabled when 'datadog' specified in headers""" + mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "datadog"} + litellm_params = {"proxy_server_request": {"url": "test"}} + + # Mock DataDogLogger to avoid async initialization issues + with patch('litellm.integrations.datadog.datadog.DataDogLogger.__init__', return_value=None): + datadog_logger = DataDogLogger() + result = EnterpriseCallbackControls.is_callback_disabled_via_headers(datadog_logger, litellm_params) + assert result is True + + def test_multiple_callbacks_disabled(self, mock_premium_user, mock_request_headers): + """Test that multiple callbacks can be disabled with comma-separated list""" + mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "langfuse,datadog,s3_v2"} + litellm_params = {"proxy_server_request": {"url": "test"}} + + # Test each callback is disabled + assert EnterpriseCallbackControls.is_callback_disabled_via_headers("langfuse", litellm_params) is True + assert EnterpriseCallbackControls.is_callback_disabled_via_headers("datadog", litellm_params) is True + assert EnterpriseCallbackControls.is_callback_disabled_via_headers("s3_v2", litellm_params) is True + + # Test non-disabled callback is not disabled + assert EnterpriseCallbackControls.is_callback_disabled_via_headers("prometheus", litellm_params) is False + + def test_callback_not_disabled_when_not_in_list(self, mock_premium_user, mock_request_headers): + """Test that callbacks not in the disabled list are not disabled""" + mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "langfuse"} + litellm_params = {"proxy_server_request": {"url": "test"}} + + result = EnterpriseCallbackControls.is_callback_disabled_via_headers("datadog", litellm_params) + assert result is False + + def test_callback_not_disabled_when_no_header(self, mock_premium_user, mock_request_headers): + """Test that callbacks are not disabled when the header is not present""" + mock_request_headers.return_value = {} + litellm_params = {"proxy_server_request": {"url": "test"}} + + result = EnterpriseCallbackControls.is_callback_disabled_via_headers("langfuse", litellm_params) + assert result is False + + def test_callback_not_disabled_when_header_none(self, mock_premium_user, mock_request_headers): + """Test that callbacks are not disabled when the header value is None""" + mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: None} + litellm_params = {"proxy_server_request": {"url": "test"}} + + result = EnterpriseCallbackControls.is_callback_disabled_via_headers("langfuse", litellm_params) + assert result is False + + def test_non_premium_user_cannot_disable_callbacks(self, mock_non_premium_user, mock_request_headers): + """Test that non-premium users cannot disable callbacks even with the header""" + mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "langfuse"} + litellm_params = {"proxy_server_request": {"url": "test"}} + + result = EnterpriseCallbackControls.is_callback_disabled_via_headers("langfuse", litellm_params) + assert result is False + + def test_case_insensitive_callback_matching(self, mock_premium_user, mock_request_headers): + """Test that callback matching is case insensitive""" + mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "LANGFUSE,DataDog"} + litellm_params = {"proxy_server_request": {"url": "test"}} + + # Test lowercase callbacks are disabled + assert EnterpriseCallbackControls.is_callback_disabled_via_headers("langfuse", litellm_params) is True + assert EnterpriseCallbackControls.is_callback_disabled_via_headers("datadog", litellm_params) is True + + def test_whitespace_handling_in_disabled_callbacks(self, mock_premium_user, mock_request_headers): + """Test that whitespace around callback names is handled correctly""" + mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: " langfuse , datadog , s3_v2 "} + litellm_params = {"proxy_server_request": {"url": "test"}} + + assert EnterpriseCallbackControls.is_callback_disabled_via_headers("langfuse", litellm_params) is True + assert EnterpriseCallbackControls.is_callback_disabled_via_headers("datadog", litellm_params) is True + assert EnterpriseCallbackControls.is_callback_disabled_via_headers("s3_v2", litellm_params) is True + + def test_custom_logger_not_in_registry(self, mock_premium_user, mock_request_headers): + """Test that CustomLogger not in registry is not disabled""" + mock_request_headers.return_value = {X_LITELLM_DISABLE_CALLBACKS: "unknown_logger"} + litellm_params = {"proxy_server_request": {"url": "test"}} + + # Create a mock CustomLogger that's not in the registry + class UnknownLogger(CustomLogger): + pass + + unknown_logger = UnknownLogger() + result = EnterpriseCallbackControls.is_callback_disabled_via_headers(unknown_logger, litellm_params) + assert result is False + + def test_exception_handling(self, mock_premium_user, mock_request_headers): + """Test that exceptions are handled gracefully and return False""" + # Make get_proxy_server_request_headers raise an exception + mock_request_headers.side_effect = Exception("Test exception") + litellm_params = {"proxy_server_request": {"url": "test"}} + + result = EnterpriseCallbackControls.is_callback_disabled_via_headers("langfuse", litellm_params) + assert result is False