mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
feat(datadog): add team-scoped Datadog callback support
Enable teams to configure their own Datadog credentials via
POST /team/{team_id}/callback, following the same pattern as Langfuse.
This commit is contained in:
parent
35f6961526
commit
9c049daa1b
6 changed files with 414 additions and 20 deletions
|
|
@ -75,12 +75,22 @@ class DataDogLogger(
|
||||||
# Class variables or attributes
|
# Class variables or attributes
|
||||||
def __init__(
|
def __init__(
|
||||||
self,
|
self,
|
||||||
|
dd_api_key: Optional[str] = None,
|
||||||
|
dd_site: Optional[str] = None,
|
||||||
|
dd_agent_host: Optional[str] = None,
|
||||||
|
dd_agent_port: Optional[str] = None,
|
||||||
**kwargs,
|
**kwargs,
|
||||||
):
|
):
|
||||||
"""
|
"""
|
||||||
Initializes the datadog logger, checks if the correct env variables are set
|
Initializes the datadog logger, checks if the correct env variables are set
|
||||||
|
|
||||||
Required environment variables (Direct API):
|
Args:
|
||||||
|
dd_api_key: Datadog API key. Falls back to DD_API_KEY env var.
|
||||||
|
dd_site: Datadog site (e.g. "us5.datadoghq.com"). Falls back to DD_SITE env var.
|
||||||
|
dd_agent_host: Hostname or IP of DataDog agent. Falls back to LITELLM_DD_AGENT_HOST env var.
|
||||||
|
dd_agent_port: Port of DataDog agent (default: 10518). Falls back to LITELLM_DD_AGENT_PORT env var.
|
||||||
|
|
||||||
|
Required environment variables (Direct API) when kwargs not provided:
|
||||||
`DD_API_KEY` - your datadog api key
|
`DD_API_KEY` - your datadog api key
|
||||||
`DD_SITE` - your datadog site, example = `"us5.datadoghq.com"`
|
`DD_SITE` - your datadog site, example = `"us5.datadoghq.com"`
|
||||||
|
|
||||||
|
|
@ -113,12 +123,19 @@ class DataDogLogger(
|
||||||
)
|
)
|
||||||
|
|
||||||
# Configure DataDog endpoint (Agent or Direct API)
|
# Configure DataDog endpoint (Agent or Direct API)
|
||||||
# Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST
|
# Prefer explicit kwargs, then fall back to env vars
|
||||||
dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST")
|
resolved_agent_host = dd_agent_host or os.getenv("LITELLM_DD_AGENT_HOST")
|
||||||
if dd_agent_host:
|
if resolved_agent_host:
|
||||||
self._configure_dd_agent(dd_agent_host=dd_agent_host)
|
self._configure_dd_agent(
|
||||||
|
dd_agent_host=resolved_agent_host,
|
||||||
|
dd_agent_port=dd_agent_port,
|
||||||
|
dd_api_key=dd_api_key,
|
||||||
|
)
|
||||||
else:
|
else:
|
||||||
self._configure_dd_direct_api()
|
self._configure_dd_direct_api(
|
||||||
|
dd_api_key=dd_api_key,
|
||||||
|
dd_site=dd_site,
|
||||||
|
)
|
||||||
|
|
||||||
# Optional override for testing
|
# Optional override for testing
|
||||||
dd_base_url = get_datadog_base_url_from_env()
|
dd_base_url = get_datadog_base_url_from_env()
|
||||||
|
|
@ -153,34 +170,54 @@ class DataDogLogger(
|
||||||
).model_dump()
|
).model_dump()
|
||||||
return dict_datadog_params
|
return dict_datadog_params
|
||||||
|
|
||||||
def _configure_dd_agent(self, dd_agent_host: str) -> None:
|
def _configure_dd_agent(
|
||||||
|
self,
|
||||||
|
dd_agent_host: str,
|
||||||
|
dd_agent_port: Optional[str] = None,
|
||||||
|
dd_api_key: Optional[str] = None,
|
||||||
|
) -> None:
|
||||||
"""
|
"""
|
||||||
Configure DataDog Agent for log forwarding
|
Configure DataDog Agent for log forwarding
|
||||||
|
|
||||||
Args:
|
Args:
|
||||||
dd_agent_host: Hostname or IP of DataDog agent
|
dd_agent_host: Hostname or IP of DataDog agent
|
||||||
|
dd_agent_port: Port of DataDog agent. Falls back to LITELLM_DD_AGENT_PORT env var (default: 10518).
|
||||||
|
dd_api_key: Datadog API key. Falls back to DD_API_KEY env var. Optional when using agent.
|
||||||
"""
|
"""
|
||||||
dd_agent_port = os.getenv(
|
resolved_port = dd_agent_port or os.getenv(
|
||||||
"LITELLM_DD_AGENT_PORT", "10518"
|
"LITELLM_DD_AGENT_PORT", "10518"
|
||||||
) # default port for logs
|
) # default port for logs
|
||||||
self.intake_url = f"http://{dd_agent_host}:{dd_agent_port}/api/v2/logs"
|
self.intake_url = f"http://{dd_agent_host}:{resolved_port}/api/v2/logs"
|
||||||
self.DD_API_KEY = os.getenv("DD_API_KEY") # Optional when using agent
|
self.DD_API_KEY = dd_api_key or os.getenv(
|
||||||
|
"DD_API_KEY"
|
||||||
|
) # Optional when using agent
|
||||||
verbose_logger.debug(f"Datadog: Using DD Agent at {self.intake_url}")
|
verbose_logger.debug(f"Datadog: Using DD Agent at {self.intake_url}")
|
||||||
|
|
||||||
def _configure_dd_direct_api(self) -> None:
|
def _configure_dd_direct_api(
|
||||||
|
self,
|
||||||
|
dd_api_key: Optional[str] = None,
|
||||||
|
dd_site: Optional[str] = None,
|
||||||
|
) -> None:
|
||||||
"""
|
"""
|
||||||
Configure direct DataDog API connection
|
Configure direct DataDog API connection
|
||||||
|
|
||||||
|
Args:
|
||||||
|
dd_api_key: Datadog API key. Falls back to DD_API_KEY env var.
|
||||||
|
dd_site: Datadog site. Falls back to DD_SITE env var.
|
||||||
|
|
||||||
Raises:
|
Raises:
|
||||||
Exception: If required environment variables are not set
|
Exception: If required credentials are not provided via args or env vars
|
||||||
"""
|
"""
|
||||||
if os.getenv("DD_API_KEY", None) is None:
|
resolved_api_key = dd_api_key or os.getenv("DD_API_KEY")
|
||||||
|
resolved_site = dd_site or os.getenv("DD_SITE")
|
||||||
|
|
||||||
|
if resolved_api_key is None:
|
||||||
raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>")
|
raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>")
|
||||||
if os.getenv("DD_SITE", None) is None:
|
if resolved_site is None:
|
||||||
raise Exception("DD_SITE is not set in .env, set 'DD_SITE=<>")
|
raise Exception("DD_SITE is not set in .env, set 'DD_SITE=<>")
|
||||||
|
|
||||||
self.DD_API_KEY = os.getenv("DD_API_KEY")
|
self.DD_API_KEY = resolved_api_key
|
||||||
self.intake_url = f"https://http-intake.logs.{os.getenv('DD_SITE')}/api/v2/logs"
|
self.intake_url = f"https://http-intake.logs.{resolved_site}/api/v2/logs"
|
||||||
|
|
||||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||||
"""
|
"""
|
||||||
|
|
|
||||||
117
litellm/integrations/datadog/datadog_team_handler.py
Normal file
117
litellm/integrations/datadog/datadog_team_handler.py
Normal file
|
|
@ -0,0 +1,117 @@
|
||||||
|
"""
|
||||||
|
DataDog Team Handler
|
||||||
|
|
||||||
|
Used to get the DataDogLogger for a given request.
|
||||||
|
Handles Key/Team Based Datadog Logging, following the same pattern as LangFuseHandler.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING, Any, Dict, Optional, TypedDict
|
||||||
|
|
||||||
|
from litellm._logging import verbose_logger
|
||||||
|
from litellm.litellm_core_utils.litellm_logging import StandardCallbackDynamicParams
|
||||||
|
|
||||||
|
from .datadog import DataDogLogger
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from litellm.litellm_core_utils.litellm_logging import DynamicLoggingCache
|
||||||
|
else:
|
||||||
|
DynamicLoggingCache = Any
|
||||||
|
|
||||||
|
|
||||||
|
class DatadogLoggingConfig(TypedDict):
|
||||||
|
dd_api_key: Optional[str]
|
||||||
|
dd_site: Optional[str]
|
||||||
|
dd_agent_host: Optional[str]
|
||||||
|
dd_agent_port: Optional[str]
|
||||||
|
|
||||||
|
|
||||||
|
class DataDogHandler:
|
||||||
|
@staticmethod
|
||||||
|
def get_datadog_logger_for_request(
|
||||||
|
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||||
|
in_memory_dynamic_logger_cache: DynamicLoggingCache,
|
||||||
|
) -> DataDogLogger:
|
||||||
|
"""
|
||||||
|
Get a team-scoped DataDogLogger for a given request.
|
||||||
|
|
||||||
|
Resolves and caches per-team DataDogLogger instances using DynamicLoggingCache,
|
||||||
|
keyed by the team's DD credentials. Each unique set of credentials gets its own
|
||||||
|
logger instance with its own batch/flush loop.
|
||||||
|
|
||||||
|
Note: This handler is only called when team-scoped DD credentials are present.
|
||||||
|
The global (env-var based) DataDogLogger is managed separately by
|
||||||
|
_init_custom_logger_compatible_class via _in_memory_loggers.
|
||||||
|
"""
|
||||||
|
_credentials = DataDogHandler.get_dynamic_datadog_logging_config(
|
||||||
|
standard_callback_dynamic_params=standard_callback_dynamic_params,
|
||||||
|
)
|
||||||
|
credentials_dict = dict(_credentials)
|
||||||
|
|
||||||
|
# check if datadog logger is already cached
|
||||||
|
temp_datadog_logger = in_memory_dynamic_logger_cache.get_cache(
|
||||||
|
credentials=credentials_dict, service_name="datadog"
|
||||||
|
)
|
||||||
|
|
||||||
|
# if not cached, create a new datadog logger and cache it
|
||||||
|
if temp_datadog_logger is None:
|
||||||
|
temp_datadog_logger = (
|
||||||
|
DataDogHandler._create_datadog_logger_from_credentials(
|
||||||
|
credentials=credentials_dict,
|
||||||
|
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
|
||||||
|
)
|
||||||
|
)
|
||||||
|
|
||||||
|
return temp_datadog_logger
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _create_datadog_logger_from_credentials(
|
||||||
|
credentials: Dict,
|
||||||
|
in_memory_dynamic_logger_cache: DynamicLoggingCache,
|
||||||
|
) -> DataDogLogger:
|
||||||
|
"""
|
||||||
|
Create a DataDogLogger from the credentials and cache it.
|
||||||
|
"""
|
||||||
|
datadog_logger = DataDogLogger(
|
||||||
|
dd_api_key=credentials.get("dd_api_key"),
|
||||||
|
dd_site=credentials.get("dd_site"),
|
||||||
|
dd_agent_host=credentials.get("dd_agent_host"),
|
||||||
|
dd_agent_port=credentials.get("dd_agent_port"),
|
||||||
|
)
|
||||||
|
in_memory_dynamic_logger_cache.set_cache(
|
||||||
|
credentials=credentials,
|
||||||
|
service_name="datadog",
|
||||||
|
logging_obj=datadog_logger,
|
||||||
|
)
|
||||||
|
verbose_logger.debug(
|
||||||
|
"Datadog: Created and cached new DataDogLogger for team-scoped credentials"
|
||||||
|
)
|
||||||
|
return datadog_logger
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def get_dynamic_datadog_logging_config(
|
||||||
|
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||||
|
) -> DatadogLoggingConfig:
|
||||||
|
"""
|
||||||
|
Get the Datadog logging config for a given request from dynamic params.
|
||||||
|
"""
|
||||||
|
return DatadogLoggingConfig(
|
||||||
|
dd_api_key=standard_callback_dynamic_params.get("dd_api_key"),
|
||||||
|
dd_site=standard_callback_dynamic_params.get("dd_site"),
|
||||||
|
dd_agent_host=standard_callback_dynamic_params.get("dd_agent_host"),
|
||||||
|
dd_agent_port=standard_callback_dynamic_params.get("dd_agent_port"),
|
||||||
|
)
|
||||||
|
|
||||||
|
@staticmethod
|
||||||
|
def _dynamic_datadog_credentials_are_passed(
|
||||||
|
standard_callback_dynamic_params: StandardCallbackDynamicParams,
|
||||||
|
) -> bool:
|
||||||
|
"""
|
||||||
|
Check if dynamic Datadog credentials are passed in standard_callback_dynamic_params.
|
||||||
|
"""
|
||||||
|
if (
|
||||||
|
standard_callback_dynamic_params.get("dd_api_key") is not None
|
||||||
|
or standard_callback_dynamic_params.get("dd_site") is not None
|
||||||
|
or standard_callback_dynamic_params.get("dd_agent_host") is not None
|
||||||
|
):
|
||||||
|
return True
|
||||||
|
return False
|
||||||
|
|
@ -53,11 +53,19 @@ _supported_callback_params = [
|
||||||
"braintrust_host",
|
"braintrust_host",
|
||||||
"slack_webhook_url",
|
"slack_webhook_url",
|
||||||
"lunary_public_key",
|
"lunary_public_key",
|
||||||
|
"dd_api_key",
|
||||||
|
"dd_site",
|
||||||
|
"dd_agent_host",
|
||||||
|
"dd_agent_port",
|
||||||
]
|
]
|
||||||
|
|
||||||
_request_blocked_callback_params = {
|
_request_blocked_callback_params = {
|
||||||
"gcs_bucket_name",
|
"gcs_bucket_name",
|
||||||
"gcs_path_service_account",
|
"gcs_path_service_account",
|
||||||
|
"dd_api_key",
|
||||||
|
"dd_site",
|
||||||
|
"dd_agent_host",
|
||||||
|
"dd_agent_port",
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|
||||||
|
|
|
||||||
|
|
@ -376,13 +376,14 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||||
List[Union[str, Callable, CustomLogger]]
|
List[Union[str, Callable, CustomLogger]]
|
||||||
] = dynamic_async_failure_callbacks
|
] = dynamic_async_failure_callbacks
|
||||||
|
|
||||||
# Process dynamic callbacks
|
|
||||||
self.process_dynamic_callbacks()
|
|
||||||
|
|
||||||
## DYNAMIC LANGFUSE / GCS / logging callback KEYS ##
|
## DYNAMIC LANGFUSE / GCS / logging callback KEYS ##
|
||||||
self.standard_callback_dynamic_params: StandardCallbackDynamicParams = (
|
self.standard_callback_dynamic_params: StandardCallbackDynamicParams = (
|
||||||
self.initialize_standard_callback_dynamic_params(kwargs)
|
self.initialize_standard_callback_dynamic_params(kwargs)
|
||||||
)
|
)
|
||||||
|
|
||||||
|
# Process dynamic callbacks (after standard_callback_dynamic_params is initialized,
|
||||||
|
# so team-scoped credentials are available for callback initialization)
|
||||||
|
self.process_dynamic_callbacks()
|
||||||
self.standard_built_in_tools_params: StandardBuiltInToolsParams = (
|
self.standard_built_in_tools_params: StandardBuiltInToolsParams = (
|
||||||
self.initialize_standard_built_in_tools_params(kwargs)
|
self.initialize_standard_built_in_tools_params(kwargs)
|
||||||
)
|
)
|
||||||
|
|
@ -477,8 +478,21 @@ class Logging(LiteLLMLoggingBaseClass):
|
||||||
isinstance(callback, str)
|
isinstance(callback, str)
|
||||||
and callback in litellm._known_custom_logger_compatible_callbacks
|
and callback in litellm._known_custom_logger_compatible_callbacks
|
||||||
):
|
):
|
||||||
|
# For callbacks that support team-scoped credentials (e.g. datadog),
|
||||||
|
# pass only the relevant dynamic params as custom_logger_init_args.
|
||||||
|
_custom_logger_init_args: Optional[dict] = None
|
||||||
|
if callback == "datadog":
|
||||||
|
_custom_logger_init_args = {
|
||||||
|
k: v
|
||||||
|
for k, v in self.standard_callback_dynamic_params.items()
|
||||||
|
if k.startswith("dd_")
|
||||||
|
}
|
||||||
|
|
||||||
callback_class = _init_custom_logger_compatible_class(
|
callback_class = _init_custom_logger_compatible_class(
|
||||||
callback, internal_usage_cache=None, llm_router=None # type: ignore
|
callback, # type: ignore[arg-type]
|
||||||
|
internal_usage_cache=None,
|
||||||
|
llm_router=None, # type: ignore
|
||||||
|
custom_logger_init_args=_custom_logger_init_args,
|
||||||
)
|
)
|
||||||
if callback_class is not None:
|
if callback_class is not None:
|
||||||
processed_list.append(callback_class)
|
processed_list.append(callback_class)
|
||||||
|
|
@ -3797,6 +3811,24 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
|
||||||
_in_memory_loggers.append(_prometheus_logger)
|
_in_memory_loggers.append(_prometheus_logger)
|
||||||
return _prometheus_logger # type: ignore
|
return _prometheus_logger # type: ignore
|
||||||
elif logging_integration == "datadog":
|
elif logging_integration == "datadog":
|
||||||
|
# Check if team-scoped credentials are provided
|
||||||
|
_dd_api_key = custom_logger_init_args.get("dd_api_key")
|
||||||
|
_dd_site = custom_logger_init_args.get("dd_site")
|
||||||
|
_dd_agent_host = custom_logger_init_args.get("dd_agent_host")
|
||||||
|
_dd_agent_port = custom_logger_init_args.get("dd_agent_port")
|
||||||
|
|
||||||
|
if _dd_api_key or _dd_site or _dd_agent_host:
|
||||||
|
# Team-scoped credentials: use DynamicLoggingCache for per-credential isolation
|
||||||
|
from litellm.integrations.datadog.datadog_team_handler import (
|
||||||
|
DataDogHandler,
|
||||||
|
)
|
||||||
|
|
||||||
|
return DataDogHandler.get_datadog_logger_for_request(
|
||||||
|
standard_callback_dynamic_params=custom_logger_init_args, # type: ignore
|
||||||
|
in_memory_dynamic_logger_cache=in_memory_dynamic_logger_cache,
|
||||||
|
)
|
||||||
|
|
||||||
|
# Global (env-var based): reuse cached instance
|
||||||
for callback in _in_memory_loggers:
|
for callback in _in_memory_loggers:
|
||||||
if isinstance(callback, DataDogLogger):
|
if isinstance(callback, DataDogLogger):
|
||||||
return callback # type: ignore
|
return callback # type: ignore
|
||||||
|
|
|
||||||
|
|
@ -3010,6 +3010,12 @@ class StandardCallbackDynamicParams(TypedDict, total=False):
|
||||||
wandb_api_key: Optional[str]
|
wandb_api_key: Optional[str]
|
||||||
weave_project_id: Optional[str]
|
weave_project_id: Optional[str]
|
||||||
|
|
||||||
|
# Datadog dynamic params
|
||||||
|
dd_api_key: Optional[str]
|
||||||
|
dd_site: Optional[str]
|
||||||
|
dd_agent_host: Optional[str]
|
||||||
|
dd_agent_port: Optional[str]
|
||||||
|
|
||||||
# Logging settings
|
# Logging settings
|
||||||
turn_off_message_logging: Optional[bool] # when true will not log messages
|
turn_off_message_logging: Optional[bool] # when true will not log messages
|
||||||
litellm_disabled_callbacks: Optional[List[str]]
|
litellm_disabled_callbacks: Optional[List[str]]
|
||||||
|
|
|
||||||
|
|
@ -0,0 +1,194 @@
|
||||||
|
"""
|
||||||
|
Tests for team-scoped Datadog callback support.
|
||||||
|
|
||||||
|
Verifies that DataDogLogger can be instantiated with per-team credentials
|
||||||
|
(dd_api_key, dd_site) instead of relying solely on environment variables,
|
||||||
|
and that the DataDogHandler correctly resolves and caches per-team loggers.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from unittest.mock import patch
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
from litellm.integrations.datadog.datadog import DataDogLogger
|
||||||
|
from litellm.integrations.datadog.datadog_team_handler import (
|
||||||
|
DataDogHandler,
|
||||||
|
DatadogLoggingConfig,
|
||||||
|
)
|
||||||
|
from litellm.litellm_core_utils.specialty_caches.dynamic_logging_cache import (
|
||||||
|
DynamicLoggingCache,
|
||||||
|
)
|
||||||
|
from litellm.types.utils import StandardCallbackDynamicParams
|
||||||
|
|
||||||
|
|
||||||
|
@pytest.fixture
|
||||||
|
def datadog_env(monkeypatch):
|
||||||
|
"""Set global DD env vars for the default/global logger."""
|
||||||
|
monkeypatch.setenv("DD_API_KEY", "global_api_key")
|
||||||
|
monkeypatch.setenv("DD_SITE", "us1.datadoghq.com")
|
||||||
|
|
||||||
|
|
||||||
|
class TestDataDogLoggerCredentialKwargs:
|
||||||
|
"""Test that DataDogLogger accepts credentials as kwargs."""
|
||||||
|
|
||||||
|
def test_init_with_explicit_credentials(self):
|
||||||
|
"""Logger should use explicit kwargs instead of env vars."""
|
||||||
|
with patch("asyncio.create_task"):
|
||||||
|
logger = DataDogLogger(
|
||||||
|
dd_api_key="team_api_key",
|
||||||
|
dd_site="eu1.datadoghq.com",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert logger.DD_API_KEY == "team_api_key"
|
||||||
|
assert "eu1.datadoghq.com" in logger.intake_url
|
||||||
|
|
||||||
|
def test_init_falls_back_to_env_vars(self, datadog_env):
|
||||||
|
"""Logger should fall back to env vars when no kwargs provided."""
|
||||||
|
with patch("asyncio.create_task"):
|
||||||
|
logger = DataDogLogger()
|
||||||
|
|
||||||
|
assert logger.DD_API_KEY == "global_api_key"
|
||||||
|
assert "us1.datadoghq.com" in logger.intake_url
|
||||||
|
|
||||||
|
def test_init_kwargs_override_env_vars(self, datadog_env):
|
||||||
|
"""Explicit kwargs should take precedence over env vars."""
|
||||||
|
with patch("asyncio.create_task"):
|
||||||
|
logger = DataDogLogger(
|
||||||
|
dd_api_key="override_key",
|
||||||
|
dd_site="ap1.datadoghq.com",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert logger.DD_API_KEY == "override_key"
|
||||||
|
assert "ap1.datadoghq.com" in logger.intake_url
|
||||||
|
|
||||||
|
def test_init_with_agent_credentials(self):
|
||||||
|
"""Logger should use agent mode when dd_agent_host is provided."""
|
||||||
|
with patch("asyncio.create_task"):
|
||||||
|
logger = DataDogLogger(
|
||||||
|
dd_agent_host="dd-agent.local",
|
||||||
|
dd_agent_port="8125",
|
||||||
|
dd_api_key="agent_api_key",
|
||||||
|
)
|
||||||
|
|
||||||
|
assert "dd-agent.local:8125" in logger.intake_url
|
||||||
|
assert logger.DD_API_KEY == "agent_api_key"
|
||||||
|
|
||||||
|
def test_init_raises_without_credentials(self, monkeypatch):
|
||||||
|
"""Logger should raise if no credentials are available."""
|
||||||
|
monkeypatch.delenv("DD_API_KEY", raising=False)
|
||||||
|
monkeypatch.delenv("DD_SITE", raising=False)
|
||||||
|
monkeypatch.delenv("LITELLM_DD_AGENT_HOST", raising=False)
|
||||||
|
|
||||||
|
with pytest.raises(Exception, match="DD_API_KEY"):
|
||||||
|
with patch("asyncio.create_task"):
|
||||||
|
DataDogLogger()
|
||||||
|
|
||||||
|
|
||||||
|
class TestDataDogHandler:
|
||||||
|
"""Test that DataDogHandler resolves the correct logger per team."""
|
||||||
|
|
||||||
|
def test_creates_team_logger_with_dynamic_credentials(self, datadog_env):
|
||||||
|
"""Should create a new logger when team credentials are provided."""
|
||||||
|
cache = DynamicLoggingCache()
|
||||||
|
params = StandardCallbackDynamicParams(
|
||||||
|
dd_api_key="team_a_key",
|
||||||
|
dd_site="eu1.datadoghq.com",
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch("asyncio.create_task"):
|
||||||
|
result = DataDogHandler.get_datadog_logger_for_request(
|
||||||
|
standard_callback_dynamic_params=params,
|
||||||
|
in_memory_dynamic_logger_cache=cache,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result.DD_API_KEY == "team_a_key"
|
||||||
|
assert "eu1.datadoghq.com" in result.intake_url
|
||||||
|
|
||||||
|
def test_caches_team_logger(self, datadog_env):
|
||||||
|
"""Same team credentials should return the same cached logger instance."""
|
||||||
|
cache = DynamicLoggingCache()
|
||||||
|
params = StandardCallbackDynamicParams(
|
||||||
|
dd_api_key="team_b_key",
|
||||||
|
dd_site="us5.datadoghq.com",
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch("asyncio.create_task"):
|
||||||
|
result1 = DataDogHandler.get_datadog_logger_for_request(
|
||||||
|
standard_callback_dynamic_params=params,
|
||||||
|
in_memory_dynamic_logger_cache=cache,
|
||||||
|
)
|
||||||
|
result2 = DataDogHandler.get_datadog_logger_for_request(
|
||||||
|
standard_callback_dynamic_params=params,
|
||||||
|
in_memory_dynamic_logger_cache=cache,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result1 is result2
|
||||||
|
|
||||||
|
def test_different_teams_get_different_loggers(self, datadog_env):
|
||||||
|
"""Different team credentials should create separate logger instances."""
|
||||||
|
cache = DynamicLoggingCache()
|
||||||
|
|
||||||
|
params_a = StandardCallbackDynamicParams(
|
||||||
|
dd_api_key="team_a_key",
|
||||||
|
dd_site="us1.datadoghq.com",
|
||||||
|
)
|
||||||
|
params_b = StandardCallbackDynamicParams(
|
||||||
|
dd_api_key="team_b_key",
|
||||||
|
dd_site="eu1.datadoghq.com",
|
||||||
|
)
|
||||||
|
|
||||||
|
with patch("asyncio.create_task"):
|
||||||
|
result_a = DataDogHandler.get_datadog_logger_for_request(
|
||||||
|
standard_callback_dynamic_params=params_a,
|
||||||
|
in_memory_dynamic_logger_cache=cache,
|
||||||
|
)
|
||||||
|
result_b = DataDogHandler.get_datadog_logger_for_request(
|
||||||
|
standard_callback_dynamic_params=params_b,
|
||||||
|
in_memory_dynamic_logger_cache=cache,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert result_a is not result_b
|
||||||
|
assert result_a.DD_API_KEY == "team_a_key"
|
||||||
|
assert result_b.DD_API_KEY == "team_b_key"
|
||||||
|
|
||||||
|
def test_request_blocked_callback_params_includes_dd(self):
|
||||||
|
"""DD params should be blocked from request-level metadata (security)."""
|
||||||
|
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
|
||||||
|
_request_blocked_callback_params,
|
||||||
|
)
|
||||||
|
|
||||||
|
assert "dd_api_key" in _request_blocked_callback_params
|
||||||
|
assert "dd_site" in _request_blocked_callback_params
|
||||||
|
assert "dd_agent_host" in _request_blocked_callback_params
|
||||||
|
assert "dd_agent_port" in _request_blocked_callback_params
|
||||||
|
|
||||||
|
|
||||||
|
class TestDynamicCredentialDetection:
|
||||||
|
"""Test that _dynamic_datadog_credentials_are_passed works correctly."""
|
||||||
|
|
||||||
|
def test_no_credentials(self):
|
||||||
|
params = StandardCallbackDynamicParams()
|
||||||
|
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is False
|
||||||
|
|
||||||
|
def test_dd_api_key_only(self):
|
||||||
|
params = StandardCallbackDynamicParams(dd_api_key="key")
|
||||||
|
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True
|
||||||
|
|
||||||
|
def test_dd_site_only(self):
|
||||||
|
params = StandardCallbackDynamicParams(dd_site="site")
|
||||||
|
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True
|
||||||
|
|
||||||
|
def test_dd_agent_host_only(self):
|
||||||
|
params = StandardCallbackDynamicParams(dd_agent_host="host")
|
||||||
|
assert DataDogHandler._dynamic_datadog_credentials_are_passed(params) is True
|
||||||
|
|
||||||
|
|
||||||
|
class TestStandardCallbackDynamicParamsIncludesDatadog:
|
||||||
|
"""Verify that Datadog params are in the allow-list."""
|
||||||
|
|
||||||
|
def test_dd_params_in_annotations(self):
|
||||||
|
annotations = StandardCallbackDynamicParams.__annotations__
|
||||||
|
assert "dd_api_key" in annotations
|
||||||
|
assert "dd_site" in annotations
|
||||||
|
assert "dd_agent_host" in annotations
|
||||||
|
assert "dd_agent_port" in annotations
|
||||||
Loading…
Add table
Reference in a new issue