[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 <cursoragent@cursor.com>
This commit is contained in:
Ishaan Jaff
GitHub
parent
02a095d4db
commit
8c5fb6f539
194
docs/my-website/docs/proxy/dynamic_logging.md
Normal file
194
docs/my-website/docs/proxy/dynamic_logging.md
Normal file
@ -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
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="disable-single" label="Disable a single callback">
|
||||
|
||||
```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"
|
||||
}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="disable-multiple" label="Disable multiple 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,datadog' \
|
||||
--data '{
|
||||
"model": "claude-sonnet-4-20250514",
|
||||
"messages": [
|
||||
{
|
||||
"role": "user",
|
||||
"content": "what llm are you"
|
||||
}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## 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.
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="Curl" label="Curl Request">
|
||||
|
||||
```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"
|
||||
}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="OpenAI" label="OpenAI Python SDK">
|
||||
|
||||
```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)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
## 3. Disable Multiple Callbacks
|
||||
|
||||
You can disable multiple callbacks by providing a comma-separated list in the header.
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="Curl" label="Curl Request">
|
||||
|
||||
```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"
|
||||
}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="OpenAI" label="OpenAI Python SDK">
|
||||
|
||||
```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)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
@ -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' \
|
||||
<Image img={require('../../img/message_redaction_spend_logs.png')} />
|
||||
|
||||
|
||||
### 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: <callback_name>` in the request headers.
|
||||
|
||||
Send the list of callbacks to disable in the request header `x-litellm-disable-callbacks`.
|
||||
|
||||
<Tabs>
|
||||
<TabItem value="Curl" label="Curl Request">
|
||||
|
||||
```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"
|
||||
}
|
||||
]
|
||||
}'
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
<TabItem value="OpenAI" label="OpenAI Python SDK">
|
||||
|
||||
```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)
|
||||
```
|
||||
|
||||
</TabItem>
|
||||
</Tabs>
|
||||
|
||||
|
||||
### ✨ 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?
|
||||
|
||||
|
||||
@ -198,7 +198,8 @@ const sidebars = {
|
||||
items: [
|
||||
"proxy/logging",
|
||||
"proxy/logging_spec",
|
||||
"proxy/team_logging"
|
||||
"proxy/team_logging",
|
||||
"proxy/dynamic_logging"
|
||||
],
|
||||
},
|
||||
|
||||
|
||||
@ -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
|
||||
@ -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(
|
||||
|
||||
133
litellm/litellm_core_utils/custom_logger_registry.py
Normal file
133
litellm/litellm_core_utils/custom_logger_registry.py
Normal file
@ -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
|
||||
@ -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,
|
||||
|
||||
@ -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
|
||||
@ -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)
|
||||
|
||||
@ -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
|
||||
Loading…
Reference in New Issue
Block a user