From 939e1c9b19a8723908a5003d60fe20fe40afbbe0 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Thu, 16 Jan 2025 22:02:24 -0800
Subject: [PATCH] (datadog llm observability) - fixes + improvements for using
 `datadog llm observability` logging integration  (#7824)

* dd llm obs fixes

* _ensure_string_content

* fix _get_dd_llm_obs_payload_metadata
---
 .../integrations/datadog/datadog_llm_obs.py   | 42 +++++++++++++++++--
 litellm/proxy/proxy_config.yaml               |  2 +-
 litellm/types/integrations/datadog_llm_obs.py |  3 +-
 3 files changed, 42 insertions(+), 5 deletions(-)

diff --git a/litellm/integrations/datadog/datadog_llm_obs.py b/litellm/integrations/datadog/datadog_llm_obs.py
index 6b7aa435465..58d16aff3e2 100644
--- a/litellm/integrations/datadog/datadog_llm_obs.py
+++ b/litellm/integrations/datadog/datadog_llm_obs.py
@@ -7,10 +7,11 @@ API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=examp
 """
 
 import asyncio
+import json
 import os
 import uuid
 from datetime import datetime
-from typing import Any, Dict, List, Optional
+from typing import Any, Dict, List, Optional, Union
 
 import litellm
 from litellm._logging import verbose_logger
@@ -97,7 +98,7 @@ class DataDogLLMObsLogger(CustomBatchLogger):
                     ),
                 ),
             }
-
+            verbose_logger.debug("payload", json.dumps(payload, indent=4))
             response = await self.async_client.post(
                 url=self.intake_url,
                 json=payload,
@@ -130,12 +131,19 @@ class DataDogLLMObsLogger(CustomBatchLogger):
             raise Exception("DataDogLLMObs: standard_logging_object is not set")
 
         messages = standard_logging_payload["messages"]
+        messages = self._ensure_string_content(messages=messages)
+
         metadata = kwargs.get("litellm_params", {}).get("metadata", {})
 
         input_meta = InputMeta(messages=messages)  # type: ignore
         output_meta = OutputMeta(messages=self._get_response_messages(response_obj))
 
-        meta = Meta(kind="llm", input=input_meta, output=output_meta)
+        meta = Meta(
+            kind="llm",
+            input=input_meta,
+            output=output_meta,
+            metadata=self._get_dd_llm_obs_payload_metadata(standard_logging_payload),
+        )
 
         # Calculate metrics (you may need to adjust these based on available data)
         metrics = LLMMetrics(
@@ -164,3 +172,31 @@ class DataDogLLMObsLogger(CustomBatchLogger):
         if isinstance(response_obj, litellm.ModelResponse):
             return [response_obj["choices"][0]["message"].json()]
         return []
+
+    def _ensure_string_content(
+        self, messages: Optional[Union[str, List[Any], Dict[Any, Any]]]
+    ) -> List[Any]:
+        if messages is None:
+            return []
+        if isinstance(messages, str):
+            return [messages]
+        elif isinstance(messages, list):
+            return [message for message in messages]
+        elif isinstance(messages, dict):
+            return [str(messages.get("content", ""))]
+        return []
+
+    def _get_dd_llm_obs_payload_metadata(
+        self, standard_logging_payload: StandardLoggingPayload
+    ) -> Dict:
+        _metadata = {
+            "model_name": standard_logging_payload.get("model", "unknown"),
+            "model_provider": standard_logging_payload.get(
+                "custom_llm_provider", "unknown"
+            ),
+        }
+        _standard_logging_metadata: dict = (
+            dict(standard_logging_payload.get("metadata", {})) or {}
+        )
+        _metadata.update(_standard_logging_metadata)
+        return _metadata
diff --git a/litellm/proxy/proxy_config.yaml b/litellm/proxy/proxy_config.yaml
index b2edcffab9e..04c4e3f07e6 100644
--- a/litellm/proxy/proxy_config.yaml
+++ b/litellm/proxy/proxy_config.yaml
@@ -13,4 +13,4 @@ model_list:
       health_check_model: anthropic/claude-3-5-sonnet-20240620
 
 litellm_settings:
-  callbacks: ["prometheus"]
+  callbacks: ["datadog_llm_observability"]
diff --git a/litellm/types/integrations/datadog_llm_obs.py b/litellm/types/integrations/datadog_llm_obs.py
index 119d8ecc7a4..91ea8c25758 100644
--- a/litellm/types/integrations/datadog_llm_obs.py
+++ b/litellm/types/integrations/datadog_llm_obs.py
@@ -4,7 +4,7 @@ Payloads for Datadog LLM Observability Service (LLMObs)
 API Reference: https://docs.datadoghq.com/llm_observability/setup/api/?tab=example#api-standards
 """
 
-from typing import Any, List, Literal, Optional, TypedDict
+from typing import Any, Dict, List, Literal, Optional, TypedDict
 
 
 class InputMeta(TypedDict):
@@ -20,6 +20,7 @@ class Meta(TypedDict):
     kind: Literal["llm", "tool", "task", "embedding", "retrieval"]
     input: InputMeta  # The span’s input information.
     output: OutputMeta  # The span’s output information.
+    metadata: Dict[str, Any]
 
 
 class LLMMetrics(TypedDict, total=False):