From 6586718b610cf1e6853cc4dac6808916a71b6fd3 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 26 Nov 2024 17:06:59 -0800
Subject: [PATCH] use type for DatadogProxyFailureHookJsonMessage

---
 litellm/integrations/datadog/datadog.py | 18 ++++++++----------
 litellm/integrations/datadog/types.py   |  7 +++++++
 2 files changed, 15 insertions(+), 10 deletions(-)

diff --git a/litellm/integrations/datadog/datadog.py b/litellm/integrations/datadog/datadog.py
index d0dee5369..52e040d45 100644
--- a/litellm/integrations/datadog/datadog.py
+++ b/litellm/integrations/datadog/datadog.py
@@ -35,7 +35,7 @@ from litellm.llms.custom_httpx.http_handler import (
 from litellm.proxy._types import UserAPIKeyAuth
 from litellm.types.services import ServiceLoggerPayload
 
-from .types import DD_ERRORS, DatadogPayload, DataDogStatus
+from .types import *
 from .utils import make_json_serializable
 
 DD_MAX_BATCH_SIZE = 1000  # max number of logs DD API can accept
@@ -410,11 +410,13 @@ class DataDogLogger(CustomBatchLogger):
                 "Datadog: Logging - Enters failure logging function for model %s",
                 request_data,
             )
-            _json_message = {
-                "exception": str(original_exception),
-                "request_data": request_data,
-                "user_api_key_dict": user_api_key_dict.model_dump(),
-            }
+            _json_message = DatadogProxyFailureHookJsonMessage(
+                exception=str(original_exception),
+                traceback=traceback.format_exc(),
+                request_data=request_data,
+                user_api_key_dict=user_api_key_dict.model_dump(),
+            )
+
             dd_payload = DatadogPayload(
                 ddsource=DD_SOURCE_NAME,
                 ddtags="",
@@ -428,10 +430,6 @@ class DataDogLogger(CustomBatchLogger):
             verbose_logger.debug(
                 f"Datadog, failure event added to queue. Will flush in {self.flush_interval} seconds..."
             )
-
-            if len(self.log_queue) >= self.batch_size:
-                await self.async_send_batch()
-
         except Exception as e:
             verbose_logger.exception(
                 f"Datadog Layer Error - {str(e)}\n{traceback.format_exc()}"
diff --git a/litellm/integrations/datadog/types.py b/litellm/integrations/datadog/types.py
index 87aa3ce17..4d11daa75 100644
--- a/litellm/integrations/datadog/types.py
+++ b/litellm/integrations/datadog/types.py
@@ -19,3 +19,10 @@ class DatadogPayload(TypedDict, total=False):
 
 class DD_ERRORS(Enum):
     DATADOG_413_ERROR = "Datadog API Error - Payload too large (batch is above 5MB uncompressed). If you want this logged either disable request/response logging or set `DD_BATCH_SIZE=50`"
+
+
+class DatadogProxyFailureHookJsonMessage(TypedDict, total=False):
+    exception: str
+    traceback: str
+    request_data: dict
+    user_api_key_dict: dict