feat: add datadog cost management support and fix startup callback issue

docs/my-website/docs/observability/datadog.md

@@ -7,6 +7,7 @@ import TabItem from '@theme/TabItem';
 LiteLLM Supports logging to the following Datdog Integrations:
 - `datadog` [Datadog Logs](https://docs.datadoghq.com/logs/)
 - `datadog_llm_observability` [Datadog LLM Observability](https://www.datadoghq.com/product/llm-observability/)
+- `datadog_cost_management` [Datadog Cloud Cost Management](#datadog-cloud-cost-management)
 - `ddtrace-run` [Datadog Tracing](#datadog-tracing)
 
 ## Datadog Logs
@@ -161,6 +162,50 @@ On the Datadog LLM Observability page, you should see that both input messages a
 
 
 
+<Image img={require('../../img/dd_llm_obs.png')} />
+
+
+## Datadog Cloud Cost Management
+
+| Feature | Details |
+|---------|---------|
+| **What is logged** | Aggregated LLM Costs (FOCUS format) |
+| **Events** | Periodic Uploads of Aggregated Cost Data |
+| **Product Link** | [Datadog Cloud Cost Management](https://docs.datadoghq.com/cost_management/) |
+
+We will use the `--config` to set `litellm.callbacks = ["datadog_cost_management"]`. This will periodically upload aggregated LLM cost data to Datadog.
+
+**Step 1**: Create a `config.yaml` file and set `litellm_settings`: `success_callback`
+
+```yaml
+model_list:
+ - model_name: gpt-3.5-turbo
+    litellm_params:
+      model: gpt-3.5-turbo
+litellm_settings:
+  callbacks: ["datadog_cost_management"]
+```
+
+**Step 2**: Set Required env variables
+
+```shell
+DD_API_KEY="your-api-key"
+DD_APP_KEY="your-app-key" # REQUIRED for Cost Management
+DD_SITE="us5.datadoghq.com"
+```
+
+**Step 3**: Start the proxy
+
+```shell
+litellm --config config.yaml
+```
+
+**How it works**
+* LiteLLM aggregates costs in-memory by Provider, Model, Date, and Tags.
+* Requires `DD_APP_KEY` for the Custom Costs API.
+* Costs are uploaded periodically (flushed).
+
+
 ### Datadog Tracing
 
 Use `ddtrace-run` to enable [Datadog Tracing](https://ddtrace.readthedocs.io/en/stable/installation_quickstart.html) on litellm proxy

litellm/integrations/callback_configs.json

@@ -83,6 +83,33 @@
     },
     "description": "Datadog Logging Integration"
   },
+  {
+    "id": "datadog_cost_management",
+    "displayName": "Datadog Cost Management",
+    "logo": "datadog.png",
+    "supports_key_team_logging": false,
+    "dynamic_params": {
+      "dd_api_key": {
+        "type": "password",
+        "ui_name": "API Key",
+        "description": "Datadog API key for authentication",
+        "required": true
+      },
+      "dd_app_key": {
+        "type": "password",
+        "ui_name": "App Key",
+        "description": "Datadog Application Key for Cloud Cost Management",
+        "required": true
+      },
+      "dd_site": {
+        "type": "text",
+        "ui_name": "Site",
+        "description": "Datadog site URL (e.g., us5.datadoghq.com)",
+        "required": true
+      }
+    },
+    "description": "Datadog Cloud Cost Management Integration"
+  },
   {
     "id": "lago",
     "displayName": "Lago",
@@ -407,4 +434,4 @@
     },
     "description": "SQS Queue (AWS) Logging Integration"
   }
-]
+]

litellm/integrations/datadog/datadog_cost_management.py

@@ -0,0 +1,202 @@
+import asyncio
+import os
+import time
+from datetime import datetime
+from typing import Dict, List, Optional
+
+from litellm._logging import verbose_logger
+from litellm.integrations.custom_batch_logger import CustomBatchLogger
+from litellm.llms.custom_httpx.http_handler import (
+    get_async_httpx_client,
+    httpxSpecialProvider,
+)
+from litellm.types.integrations.datadog_cost_management import (
+    DatadogFOCUSCostEntry,
+)
+from litellm.types.utils import StandardLoggingPayload
+
+
+class DatadogCostManagementLogger(CustomBatchLogger):
+    def __init__(self, **kwargs):
+        self.dd_api_key = os.getenv("DD_API_KEY")
+        self.dd_app_key = os.getenv("DD_APP_KEY")
+        self.dd_site = os.getenv("DD_SITE", "datadoghq.com")
+
+        if not self.dd_api_key or not self.dd_app_key:
+            verbose_logger.warning(
+                "Datadog Cost Management: DD_API_KEY and DD_APP_KEY are required. Integration will not work."
+            )
+
+        self.upload_url = f"https://api.{self.dd_site}/api/v2/cost/custom_costs"
+
+        self.async_client = get_async_httpx_client(
+            llm_provider=httpxSpecialProvider.LoggingCallback
+        )
+
+        # Initialize lock and start periodic flush task
+        self.flush_lock = asyncio.Lock()
+        asyncio.create_task(self.periodic_flush())
+
+        # Check if flush_lock is already in kwargs to avoid double passing (unlikely but safe)
+        if "flush_lock" not in kwargs:
+            kwargs["flush_lock"] = self.flush_lock
+
+        super().__init__(**kwargs)
+
+    async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
+        try:
+            standard_logging_object: Optional[StandardLoggingPayload] = kwargs.get(
+                "standard_logging_object", None
+            )
+
+            if standard_logging_object is None:
+                return
+
+            # Only log if there is a cost associated
+            if standard_logging_object.get("response_cost", 0) > 0:
+                self.log_queue.append(standard_logging_object)
+
+                if len(self.log_queue) >= self.batch_size:
+                    await self.async_send_batch()
+
+        except Exception as e:
+            verbose_logger.exception(
+                f"Datadog Cost Management: Error in async_log_success_event: {str(e)}"
+            )
+
+    async def async_send_batch(self):
+        if not self.log_queue:
+            return
+
+        try:
+            # Aggregate costs from the batch
+            aggregated_entries = self._aggregate_costs(self.log_queue)
+
+            if not aggregated_entries:
+                return
+
+            # Send to Datadog
+            await self._upload_to_datadog(aggregated_entries)
+
+            # Clear queue only on success (or if we decide to drop on failure)
+            # CustomBatchLogger clears queue in flush_queue, so we just process here
+
+        except Exception as e:
+            verbose_logger.exception(
+                f"Datadog Cost Management: Error in async_send_batch: {str(e)}"
+            )
+
+    def _aggregate_costs(
+        self, logs: List[StandardLoggingPayload]
+    ) -> List[DatadogFOCUSCostEntry]:
+        """
+        Aggregates costs by Provider, Model, and Date.
+        Returns a list of DatadogFOCUSCostEntry.
+        """
+        aggregator: Dict[str, DatadogFOCUSCostEntry] = {}
+
+        for log in logs:
+            try:
+                # Extract keys for aggregation
+                provider = log.get("custom_llm_provider") or "unknown"
+                model = log.get("model") or "unknown"
+                cost = log.get("response_cost", 0)
+
+                if cost == 0:
+                    continue
+
+                # Get date strings (FOCUS format requires specific keys, but for aggregation we group by Day)
+                # UTC date
+                # We interpret "ChargePeriod" as the day of the request.
+                ts = log.get("startTime") or time.time()
+                dt = datetime.fromtimestamp(ts)
+                date_str = dt.strftime("%Y-%m-%d")
+
+                # ChargePeriodStart and End
+                # If we want daily granularity, end date is usually same day or next day?
+                # Datadog Custom Costs usually expects periods.
+                # "ChargePeriodStart": "2023-01-01", "ChargePeriodEnd": "2023-12-31" in example.
+                # If we send daily, we can say Start=Date, End=Date.
+
+                # Grouping Key: Provider + Model + Date + Tags?
+                # For simplicity, let's aggregate by Provider + Model + Date first.
+                # If we handle tags, we need to include them in the key.
+
+                tags = self._extract_tags(log)
+                tags_key = tuple(sorted(tags.items())) if tags else ()
+
+                key = (provider, model, date_str, tags_key)
+
+                if key not in aggregator:
+                    aggregator[key] = {
+                        "ProviderName": provider,
+                        "ChargeDescription": f"LLM Usage for {model}",
+                        "ChargePeriodStart": date_str,
+                        "ChargePeriodEnd": date_str,
+                        "BilledCost": 0.0,
+                        "BillingCurrency": "USD",
+                        "Tags": tags if tags else None,
+                    }
+
+                aggregator[key]["BilledCost"] += cost
+
+            except Exception as e:
+                verbose_logger.warning(
+                    f"Error processing log for cost aggregation: {e}"
+                )
+                continue
+
+        return list(aggregator.values())
+
+    def _extract_tags(self, log: StandardLoggingPayload) -> Dict[str, str]:
+        from litellm.integrations.datadog.datadog_handler import (
+            get_datadog_env,
+            get_datadog_hostname,
+            get_datadog_pod_name,
+            get_datadog_service,
+        )
+
+        tags = {
+            "env": get_datadog_env(),
+            "service": get_datadog_service(),
+            "host": get_datadog_hostname(),
+            "pod_name": get_datadog_pod_name(),
+        }
+
+        # Add metadata as tags
+        metadata = log.get("metadata", {})
+        if metadata:
+            # Add user info
+            if "user_api_key_alias" in metadata:
+                tags["user"] = str(metadata["user_api_key_alias"])
+            if "user_api_key_team_alias" in metadata:
+                tags["team"] = str(metadata["user_api_key_team_alias"])
+            if "model_group" in metadata:
+                tags["model_group"] = str(metadata["model_group"])
+
+        return tags
+
+    async def _upload_to_datadog(self, payload: List[Dict]):
+        if not self.dd_api_key or not self.dd_app_key:
+            return
+
+        headers = {
+            "Content-Type": "application/json",
+            "DD-API-KEY": self.dd_api_key,
+            "DD-APPLICATION-KEY": self.dd_app_key,
+        }
+
+        # The API endpoint expects a list of objects directly in the body (file content behavior)
+        from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
+
+        data_json = safe_dumps(payload)
+
+        response = await self.async_client.put(
+            self.upload_url, content=data_json, headers=headers
+        )
+
+        response.raise_for_status()
+
+        verbose_logger.debug(
+            f"Datadog Cost Management: Uploaded {len(payload)} cost entries. Status: {response.status_code}"
+        )

litellm/proxy/common_utils/callback_utils.py

@@ -274,11 +274,20 @@ def initialize_callbacks_on_proxy(  # noqa: PLR0915
                     WebSearchInterceptionLogger,
                 )
 
-                websearch_interception_obj = WebSearchInterceptionLogger.initialize_from_proxy_config(
-                    litellm_settings=litellm_settings,
-                    callback_specific_params=callback_specific_params,
+                websearch_interception_obj = (
+                    WebSearchInterceptionLogger.initialize_from_proxy_config(
+                        litellm_settings=litellm_settings,
+                        callback_specific_params=callback_specific_params,
+                    )
                 )
                 imported_list.append(websearch_interception_obj)
+            elif isinstance(callback, str) and callback == "datadog_cost_management":
+                from litellm.integrations.datadog.datadog_cost_management import (
+                    DatadogCostManagementLogger,
+                )
+
+                datadog_cost_management_obj = DatadogCostManagementLogger()
+                imported_list.append(datadog_cost_management_obj)
             elif isinstance(callback, CustomLogger):
                 imported_list.append(callback)
             else:
@@ -353,17 +362,17 @@ def get_remaining_tokens_and_requests_from_request_data(data: Dict) -> Dict[str,
     remaining_requests_variable_name = f"litellm-key-remaining-requests-{model_group}"
     remaining_requests = _metadata.get(remaining_requests_variable_name, None)
     if remaining_requests:
-        headers[f"x-litellm-key-remaining-requests-{h11_model_group_name}"] = (
-            remaining_requests
-        )
+        headers[
+            f"x-litellm-key-remaining-requests-{h11_model_group_name}"
+        ] = remaining_requests
 
     # Remaining Tokens
     remaining_tokens_variable_name = f"litellm-key-remaining-tokens-{model_group}"
     remaining_tokens = _metadata.get(remaining_tokens_variable_name, None)
     if remaining_tokens:
-        headers[f"x-litellm-key-remaining-tokens-{h11_model_group_name}"] = (
-            remaining_tokens
-        )
+        headers[
+            f"x-litellm-key-remaining-tokens-{h11_model_group_name}"
+        ] = remaining_tokens
 
     return headers
 
@@ -412,9 +421,9 @@ def add_guardrail_response_to_standard_logging_object(
 ):
     if litellm_logging_obj is None:
         return
-    standard_logging_object: Optional[StandardLoggingPayload] = (
-        litellm_logging_obj.model_call_details.get("standard_logging_object")
-    )
+    standard_logging_object: Optional[
+        StandardLoggingPayload
+    ] = litellm_logging_obj.model_call_details.get("standard_logging_object")
     if standard_logging_object is None:
         return
     guardrail_information = standard_logging_object.get("guardrail_information", [])
@@ -443,7 +452,9 @@ def get_metadata_variable_name_from_kwargs(
     return "litellm_metadata" if "litellm_metadata" in kwargs else "metadata"
 
 
-def process_callback(_callback: str, callback_type: str, environment_variables: dict) -> dict:
+def process_callback(
+    _callback: str, callback_type: str, environment_variables: dict
+) -> dict:
     """Process a single callback and return its data with environment variables"""
     env_vars = CustomLogger.get_callback_env_vars(_callback)
 
@@ -455,11 +466,9 @@ def process_callback(_callback: str, callback_type: str, environment_variables:
         else:
             env_vars_dict[_var] = env_variable
 
-    return {
-        "name": _callback,
-        "variables": env_vars_dict,
-        "type": callback_type
-    }
+    return {"name": _callback, "variables": env_vars_dict, "type": callback_type}
+
+
 def normalize_callback_names(callbacks: Iterable[Any]) -> List[Any]:
     if callbacks is None:
         return []