Merge pull request #2578 from BerriAI/litellm_datadog_int

[FEAT] DataDog Logging Provider

litellm/integrations/datadog.py

@@ -0,0 +1,143 @@
+#### What this does ####
+#    On success + failure, log events to Supabase
+
+import dotenv, os
+import requests
+
+dotenv.load_dotenv()  # Loading env variables using dotenv
+import traceback
+import datetime, subprocess, sys
+import litellm, uuid
+from litellm._logging import print_verbose, verbose_logger
+
+
+class DataDogLogger:
+    # Class variables or attributes
+    def __init__(
+        self,
+        **kwargs,
+    ):
+        from datadog_api_client import ApiClient, Configuration
+
+        # check if the correct env variables are set
+        if os.getenv("DD_API_KEY", None) is None:
+            raise Exception("DD_API_KEY is not set, set 'DD_API_KEY=<>")
+        if os.getenv("DD_SITE", None) is None:
+            raise Exception("DD_SITE is not set in .env, set 'DD_SITE=<>")
+        self.configuration = Configuration()
+
+        try:
+            verbose_logger.debug(f"in init datadog logger")
+            pass
+
+        except Exception as e:
+            print_verbose(f"Got exception on init s3 client {str(e)}")
+            raise e
+
+    async def _async_log_event(
+        self, kwargs, response_obj, start_time, end_time, print_verbose, user_id
+    ):
+        self.log_event(kwargs, response_obj, start_time, end_time, print_verbose)
+
+    def log_event(
+        self, kwargs, response_obj, start_time, end_time, user_id, print_verbose
+    ):
+        try:
+            # Define DataDog client
+            from datadog_api_client.v2.api.logs_api import LogsApi
+            from datadog_api_client.v2 import ApiClient
+            from datadog_api_client.v2.models import HTTPLogItem, HTTPLog
+
+            verbose_logger.debug(
+                f"datadog Logging - Enters logging function for model {kwargs}"
+            )
+            litellm_params = kwargs.get("litellm_params", {})
+            metadata = (
+                litellm_params.get("metadata", {}) or {}
+            )  # if litellm_params['metadata'] == None
+            messages = kwargs.get("messages")
+            optional_params = kwargs.get("optional_params", {})
+            call_type = kwargs.get("call_type", "litellm.completion")
+            cache_hit = kwargs.get("cache_hit", False)
+            usage = response_obj["usage"]
+            id = response_obj.get("id", str(uuid.uuid4()))
+            usage = dict(usage)
+            try:
+                response_time = (end_time - start_time).total_seconds()
+            except:
+                response_time = None
+
+            try:
+                response_obj = dict(response_obj)
+            except:
+                response_obj = response_obj
+
+            # Clean Metadata before logging - never log raw metadata
+            # the raw metadata can contain circular references which leads to infinite recursion
+            # we clean out all extra litellm metadata params before logging
+            clean_metadata = {}
+            if isinstance(metadata, dict):
+                for key, value in metadata.items():
+                    # clean litellm metadata before logging
+                    if key in [
+                        "endpoint",
+                        "caching_groups",
+                        "previous_models",
+                    ]:
+                        continue
+                    else:
+                        clean_metadata[key] = value
+
+            # Build the initial payload
+            payload = {
+                "id": id,
+                "call_type": call_type,
+                "cache_hit": cache_hit,
+                "startTime": start_time,
+                "endTime": end_time,
+                "responseTime (seconds)": response_time,
+                "model": kwargs.get("model", ""),
+                "user": kwargs.get("user", ""),
+                "modelParameters": optional_params,
+                "spend": kwargs.get("response_cost", 0),
+                "messages": messages,
+                "response": response_obj,
+                "usage": usage,
+                "metadata": clean_metadata,
+            }
+
+            # Ensure everything in the payload is converted to str
+            for key, value in payload.items():
+                try:
+                    payload[key] = str(value)
+                except:
+                    # non blocking if it can't cast to a str
+                    pass
+            import json
+
+            payload = json.dumps(payload)
+
+            print_verbose(f"\ndd Logger - Logging payload = {payload}")
+
+            with ApiClient(self.configuration) as api_client:
+                api_instance = LogsApi(api_client)
+                body = HTTPLog(
+                    [
+                        HTTPLogItem(
+                            ddsource="litellm",
+                            message=payload,
+                            service="litellm-server",
+                        ),
+                    ]
+                )
+                response = api_instance.submit_log(body)
+
+            print_verbose(
+                f"Datadog Layer Logging - final response object: {response_obj}"
+            )
+        except Exception as e:
+            traceback.print_exc()
+            verbose_logger.debug(
+                f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}"
+            )
+            pass

litellm/tests/test_datadog.py

@@ -0,0 +1,27 @@
+import sys
+import os
+import io
+
+sys.path.insert(0, os.path.abspath("../.."))
+
+from litellm import completion
+import litellm
+import pytest
+
+import time
+
+
+@pytest.mark.skip(reason="beta test - this is a new feature")
+def test_datadog_logging():
+    try:
+        litellm.success_callback = ["datadog"]
+        litellm.set_verbose = True
+        response = completion(
+            model="gpt-3.5-turbo",
+            messages=[{"role": "user", "content": "what llm are u"}],
+            max_tokens=10,
+            temperature=0.2,
+        )
+        print(response)
+    except Exception as e:
+        print(e)

litellm/utils.py

@@ -65,6 +65,7 @@ from .integrations.langsmith import LangsmithLogger
 from .integrations.weights_biases import WeightsBiasesLogger
 from .integrations.custom_logger import CustomLogger
 from .integrations.langfuse import LangFuseLogger
+from .integrations.datadog import DataDogLogger
 from .integrations.dynamodb import DyanmoDBLogger
 from .integrations.s3 import S3Logger
 from .integrations.clickhouse import ClickhouseLogger
@@ -121,6 +122,7 @@ langsmithLogger = None
 weightsBiasesLogger = None
 customLogger = None
 langFuseLogger = None
+dataDogLogger = None
 dynamoLogger = None
 s3Logger = None
 genericAPILogger = None
@@ -1473,6 +1475,33 @@ class Logging:
                             user_id=kwargs.get("user", None),
                             print_verbose=print_verbose,
                         )
+                    if callback == "datadog":
+                        global dataDogLogger
+                        verbose_logger.debug("reaches datadog for success logging!")
+                        kwargs = {}
+                        for k, v in self.model_call_details.items():
+                            if (
+                                k != "original_response"
+                            ):  # copy.deepcopy raises errors as this could be a coroutine
+                                kwargs[k] = v
+                        # this only logs streaming once, complete_streaming_response exists i.e when stream ends
+                        if self.stream:
+                            verbose_logger.debug(
+                                f"datadog: is complete_streaming_response in kwargs: {kwargs.get('complete_streaming_response', None)}"
+                            )
+                            if complete_streaming_response is None:
+                                continue
+                            else:
+                                print_verbose("reaches datadog for streaming logging!")
+                                result = kwargs["complete_streaming_response"]
+                        dataDogLogger.log_event(
+                            kwargs=kwargs,
+                            response_obj=result,
+                            start_time=start_time,
+                            end_time=end_time,
+                            user_id=kwargs.get("user", None),
+                            print_verbose=print_verbose,
+                        )
                     if callback == "generic":
                         global genericAPILogger
                         verbose_logger.debug("reaches langfuse for success logging!")
@@ -6082,7 +6111,7 @@ def validate_environment(model: Optional[str] = None) -> dict:
 
 
 def set_callbacks(callback_list, function_id=None):
-    global sentry_sdk_instance, capture_exception, add_breadcrumb, posthog, slack_app, alerts_channel, traceloopLogger, athinaLogger, heliconeLogger, aispendLogger, berrispendLogger, supabaseClient, liteDebuggerClient, llmonitorLogger, promptLayerLogger, langFuseLogger, customLogger, weightsBiasesLogger, langsmithLogger, dynamoLogger, s3Logger
+    global sentry_sdk_instance, capture_exception, add_breadcrumb, posthog, slack_app, alerts_channel, traceloopLogger, athinaLogger, heliconeLogger, aispendLogger, berrispendLogger, supabaseClient, liteDebuggerClient, llmonitorLogger, promptLayerLogger, langFuseLogger, customLogger, weightsBiasesLogger, langsmithLogger, dynamoLogger, s3Logger, dataDogLogger
     try:
         for callback in callback_list:
             print_verbose(f"callback: {callback}")
@@ -6148,6 +6177,8 @@ def set_callbacks(callback_list, function_id=None):
                 promptLayerLogger = PromptLayerLogger()
             elif callback == "langfuse":
                 langFuseLogger = LangFuseLogger()
+            elif callback == "datadog":
+                dataDogLogger = DataDogLogger()
             elif callback == "dynamodb":
                 dynamoLogger = DyanmoDBLogger()
             elif callback == "s3":

requirements.txt

@@ -18,7 +18,7 @@ google-generativeai==0.3.2 # for vertex ai calls
 async_generator==1.10.0 # for async ollama calls
 traceloop-sdk==0.5.3 # for open telemetry logging
 langfuse>=2.6.3 # for langfuse self-hosted logging
-clickhouse_connect==0.7.0
+datadog-api-client==2.23.0 # for datadog logging
 orjson==3.9.15 # fast /embedding responses
 apscheduler==3.10.4 # for resetting budget in background 
 fastapi-sso==0.10.0 # admin UI, SSO