From 06804bc70adf1ac723b776f82ee7830434d45d17 Mon Sep 17 00:00:00 2001
From: Ishaan Jaff <ishaanjaffer0324@gmail.com>
Date: Tue, 30 Apr 2024 12:48:17 -0700
Subject: [PATCH] fix - working exception writing

---
 litellm/proxy/proxy_server.py | 65 +++++++++++++++++++++++++++++++++++
 1 file changed, 65 insertions(+)

diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py
index 29f3c41db..81ea8961a 100644
--- a/litellm/proxy/proxy_server.py
+++ b/litellm/proxy/proxy_server.py
@@ -1217,6 +1217,59 @@ def cost_tracking():
                 litellm.success_callback.append(_PROXY_track_cost_callback)  # type: ignore
 
 
+async def _PROXY_failure_handler(
+    kwargs,  # kwargs to completion
+    completion_response: litellm.ModelResponse,  # response from completion
+    start_time=None,
+    end_time=None,  # start/end time for completion
+):
+    global prisma_client
+    if prisma_client is not None:
+        verbose_proxy_logger.debug(
+            "inside _PROXY_failure_handler kwargs=", extra=kwargs
+        )
+
+        _exception = kwargs.get("exception")
+        traceback = kwargs.get("traceback")
+        _exception_type = _exception.__class__.__name__
+
+        _model = kwargs.get("model", None)
+        _status_code = _exception.status_code
+
+        _litellm_params = kwargs.get("litellm_params", {}) or {}
+        _metadata = _litellm_params.get("metadata", {}) or {}
+        _model_id = _metadata.get("model_info", {}).get("id", None)
+        verbose_proxy_logger.debug(
+            "\nexception_type",
+            _exception_type,
+            "\nrequest_model",
+            _model,
+            "\nmodel_id",
+            _model_id,
+            "\nexception",
+            _exception,
+            "\ntraceback",
+            traceback,
+        )
+        error_log = LiteLLM_ErrorLogs(
+            model_name=_model,
+            model_id=_model_id,
+            exception_type=_exception_type,
+            status_code=_status_code,
+            exception_string=str(_exception),
+        )
+
+        # helper function to convert to dict on pydantic v2 & v1
+        error_log_dict = _get_pydantic_json_dict(error_log)
+        error_log_dict["request_kwargs"] = json.dumps(error_log_dict["request_kwargs"])
+
+        await prisma_client.db.litellm_errorlogs.create(
+            data=error_log_dict  # type: ignore
+        )
+
+    pass
+
+
 async def _PROXY_track_cost_callback(
     kwargs,  # kwargs to completion
     completion_response: litellm.ModelResponse,  # response from completion
@@ -1302,6 +1355,15 @@ async def _PROXY_track_cost_callback(
         verbose_proxy_logger.debug("error in tracking cost callback - %s", e)
 
 
+def error_tracking():
+    global prisma_client, custom_db_client
+    if prisma_client is not None or custom_db_client is not None:
+        if isinstance(litellm.failure_callback, list):
+            verbose_proxy_logger.debug("setting litellm failure callback to track cost")
+            if (_PROXY_failure_handler) not in litellm.failure_callback:  # type: ignore
+                litellm.failure_callback.append(_PROXY_failure_handler)  # type: ignore
+
+
 def _set_spend_logs_payload(
     payload: dict, prisma_client: PrismaClient, spend_logs_url: Optional[str] = None
 ):
@@ -3194,6 +3256,9 @@ async def startup_event():
     ## COST TRACKING ##
     cost_tracking()
 
+    ## Error Tracking ##
+    error_tracking()
+
     db_writer_client = HTTPHandler()
 
     proxy_logging_obj._init_litellm_callbacks()  # INITIALIZE LITELLM CALLBACKS ON SERVER STARTUP <- do this to catch any logging errors on startup, not when calls are being made