Skip to content
69 changes: 53 additions & 16 deletions litellm/integrations/datadog/datadog.py
Original file line number Diff line number Diff line change
Expand Up @@ -92,12 +92,22 @@ class DataDogLogger(
# Class variables or attributes
def __init__(
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,
):
"""
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_SITE` - your datadog site, example = `"us5.datadoghq.com"`

Expand Down Expand Up @@ -130,12 +140,19 @@ def __init__(
)

# Configure DataDog endpoint (Agent or Direct API)
# Use LITELLM_DD_AGENT_HOST to avoid conflicts with ddtrace's DD_AGENT_HOST
dd_agent_host = os.getenv("LITELLM_DD_AGENT_HOST")
if dd_agent_host:
self._configure_dd_agent(dd_agent_host=dd_agent_host)
# Prefer explicit kwargs, then fall back to env vars
resolved_agent_host = dd_agent_host or os.getenv("LITELLM_DD_AGENT_HOST")
if resolved_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:
self._configure_dd_direct_api()
self._configure_dd_direct_api(
dd_api_key=dd_api_key,
dd_site=dd_site,
)

# Optional override for testing
dd_base_url = get_datadog_base_url_from_env()
Expand Down Expand Up @@ -172,34 +189,54 @@ def _get_datadog_params(self) -> Dict:
).model_dump()
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

Args:
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"
) # default port for logs
self.intake_url = f"http://{dd_agent_host}:{dd_agent_port}/api/v2/logs"
self.DD_API_KEY = os.getenv("DD_API_KEY") # Optional when using agent
self.intake_url = f"http://{dd_agent_host}:{resolved_port}/api/v2/logs"
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}")

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

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:
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=<>")
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=<>")

self.DD_API_KEY = os.getenv("DD_API_KEY")
self.intake_url = f"https://http-intake.logs.{os.getenv('DD_SITE')}/api/v2/logs"
self.DD_API_KEY = resolved_api_key
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):
"""
Expand Down
117 changes: 117 additions & 0 deletions litellm/integrations/datadog/datadog_team_handler.py
Original file line number Diff line number Diff line change
@@ -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
Original file line number Diff line number Diff line change
Expand Up @@ -53,11 +53,19 @@ def validate_no_callback_env_reference(
"braintrust_host",
"slack_webhook_url",
"lunary_public_key",
"dd_api_key",
"dd_site",
"dd_agent_host",
"dd_agent_port",
]

_request_blocked_callback_params = {
"gcs_bucket_name",
"gcs_path_service_account",
"dd_api_key",
"dd_site",
"dd_agent_host",
"dd_agent_port",
}


Expand Down
40 changes: 36 additions & 4 deletions litellm/litellm_core_utils/litellm_logging.py
Original file line number Diff line number Diff line change
Expand Up @@ -376,13 +376,14 @@ def __init__(
List[Union[str, Callable, CustomLogger]]
] = dynamic_async_failure_callbacks

# Process dynamic callbacks
self.process_dynamic_callbacks()

## DYNAMIC LANGFUSE / GCS / logging callback KEYS ##
self.standard_callback_dynamic_params: StandardCallbackDynamicParams = (
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.initialize_standard_built_in_tools_params(kwargs)
)
Expand Down Expand Up @@ -477,8 +478,21 @@ def _process_dynamic_callback_list(
isinstance(callback, str)
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, 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:
processed_list.append(callback_class)
Expand Down Expand Up @@ -3890,6 +3904,24 @@ def _init_custom_logger_compatible_class( # noqa: PLR0915
_in_memory_loggers.append(_prometheus_logger)
return _prometheus_logger # type: ignore
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:
if isinstance(callback, DataDogLogger):
return callback # type: ignore
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@

import litellm
from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS
from litellm.litellm_core_utils.llm_cost_calc.utils import _get_web_search_requests
from litellm.types.llms.openai import (
FileSearchTool,
ResponsesAPIResponse,
Expand Down Expand Up @@ -339,8 +340,7 @@ def response_object_includes_web_search_call(
# and _handle_web_search_cost() is never called.
if (
hasattr(usage, "server_tool_use")
and usage.server_tool_use is not None
and usage.server_tool_use.web_search_requests is not None
and _get_web_search_requests(usage.server_tool_use) is not None
):
return True
return False
Expand All @@ -352,8 +352,7 @@ def response_object_includes_web_search_call(
elif usage is not None:
if (
hasattr(usage, "server_tool_use")
and usage.server_tool_use is not None
and usage.server_tool_use.web_search_requests is not None
and _get_web_search_requests(usage.server_tool_use) is not None
):
return True
elif (
Expand Down
22 changes: 21 additions & 1 deletion litellm/litellm_core_utils/llm_cost_calc/utils.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# What is this?
## Helper utilities for cost_per_token()

from typing import Literal, Optional, Tuple, TypedDict, cast
from typing import Any, Literal, Optional, Tuple, TypedDict, cast

import litellm
from litellm._logging import verbose_logger
Expand Down Expand Up @@ -42,6 +42,26 @@ def _get_token_detail_value(details: object, key: str) -> Optional[int]:
return value if isinstance(value, int) else None


def _get_web_search_requests(server_tool_use: Any) -> Optional[int]:
"""
Tolerantly read ``web_search_requests`` from a ``server_tool_use`` value
that may be ``None``, a ``dict``, a ``ServerToolUse`` pydantic instance,
or any other object supporting attribute access.

Returns ``None`` when the value cannot be resolved — callers can
distinguish "absent" from "zero" using ``is None``.

See https://github.com/BerriAI/litellm/issues/26153 — ``stream_chunk_builder``
historically left this as a plain ``dict``, which broke direct attribute
access in cost calculation.
"""
if server_tool_use is None:
return None
if isinstance(server_tool_use, dict):
return server_tool_use.get("web_search_requests")
return getattr(server_tool_use, "web_search_requests", None)


def _is_above_128k(tokens: float) -> bool:
if tokens > 128000:
return True
Expand Down
Loading
Loading