From 2e1ee63d5830aa7a95f3649c4d1ad8ba4b4f46d8 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Wed, 23 Aug 2023 10:38:46 -0700
Subject: [PATCH] log embedding success/failure as well

---
 litellm/integrations/litedebugger.py | 20 ++++++++++++++++++++
 litellm/utils.py                     | 25 +++++++------------------
 pyproject.toml                       |  2 +-
 3 files changed, 28 insertions(+), 19 deletions(-)

diff --git a/litellm/integrations/litedebugger.py b/litellm/integrations/litedebugger.py
index 19b2b17772..ca695641bb 100644
--- a/litellm/integrations/litedebugger.py
+++ b/litellm/integrations/litedebugger.py
@@ -88,6 +88,26 @@ class LiteDebugger:
                     headers={"content-type": "application/json"},
                     data=json.dumps(litellm_data_obj),
                 )
+            elif "embedding" in response_obj:
+                litellm_data_obj = {
+                    "response_time": response_time,
+                    "model": response_obj["model"],
+                    "total_cost": total_cost,
+                    "messages": messages,
+                    "response": response_obj["embedding"][:5],
+                    "end_user": end_user,
+                    "litellm_call_id": litellm_call_id,
+                    "status": "success",
+                    "user_email": self.user_email,
+                }
+                print_verbose(
+                    f"LiteDebugger: Logging - final data object: {litellm_data_obj}"
+                )
+                response = requests.post(
+                    url=self.api_url,
+                    headers={"content-type": "application/json"},
+                    data=json.dumps(litellm_data_obj),
+                )
             elif "error" in response_obj:
                 if "Unable to map your input to a model." in response_obj["error"]:
                     total_cost = 0
diff --git a/litellm/utils.py b/litellm/utils.py
index 7d38a38090..10022268cd 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -293,8 +293,8 @@ def exception_logging(
 # make it easy to log if completion/embedding runs succeeded or failed + see what happened | Non-Blocking
 def client(original_function):
     global liteDebuggerClient, get_all_keys
-
-      def function_setup(
+    
+    def function_setup(
         *args, **kwargs
     ):  # just run once to check if user wants to send their data anywhere - PostHog/Sentry/Slack/etc.
         try:
@@ -899,7 +899,7 @@ def handle_failure(exception, traceback_exception, start_time, end_time, args,
                     print_verbose("reaches lite_debugger for logging!")
                     print_verbose(f"liteDebuggerClient: {liteDebuggerClient}")
                     model = args[0] if len(args) > 0 else kwargs["model"]
-                    messages = args[1] if len(args) > 1 else kwargs["messages"]
+                    messages = args[1] if len(args) > 1 else kwargs.get("messages", {"role": "user", "content": kwargs.get("input", "")})
                     result = {
                         "model": model,
                         "created": time.time(),
@@ -943,6 +943,8 @@ def handle_failure(exception, traceback_exception, start_time, end_time, args,
 def handle_success(args, kwargs, result, start_time, end_time):
     global heliconeLogger, aispendLogger, supabaseClient, liteDebuggerClient, llmonitorLogger
     try:
+        model = args[0] if len(args) > 0 else kwargs["model"]
+        input = args[1] if len(args) > 1 else kwargs.get("messages", kwargs.get("input", None))
         success_handler = additional_details.pop("success_handler", None)
         failure_handler = additional_details.pop("failure_handler", None)
         additional_details["Event_Name"] = additional_details.pop(
@@ -1011,22 +1013,10 @@ def handle_success(args, kwargs, result, start_time, end_time):
                         end_time=end_time,
                         print_verbose=print_verbose,
                     )
-                elif callback == "berrispend":
-                    print_verbose("reaches berrispend for logging!")
-                    model = args[0] if len(args) > 0 else kwargs["model"]
-                    messages = args[1] if len(args) > 1 else kwargs["messages"]
-                    berrispendLogger.log_event(
-                        model=model,
-                        messages=messages,
-                        response_obj=result,
-                        start_time=start_time,
-                        end_time=end_time,
-                        print_verbose=print_verbose,
-                    )
                 elif callback == "supabase":
                     print_verbose("reaches supabase for logging!")
                     model = args[0] if len(args) > 0 else kwargs["model"]
-                    messages = args[1] if len(args) > 1 else kwargs["messages"]
+                    messages = args[1] if len(args) > 1 else kwargs.get("messages", {"role": "user", "content": ""})
                     print(f"supabaseClient: {supabaseClient}")
                     supabaseClient.log_event(
                         model=model,
@@ -1040,9 +1030,8 @@ def handle_success(args, kwargs, result, start_time, end_time):
                     )
                 elif callback == "lite_debugger":
                     print_verbose("reaches lite_debugger for logging!")
-                    model = args[0] if len(args) > 0 else kwargs["model"]
-                    messages = args[1] if len(args) > 1 else kwargs["messages"]
                     print_verbose(f"liteDebuggerClient: {liteDebuggerClient}")
+                    messages = args[1] if len(args) > 1 else kwargs.get("messages", {"role": "user", "content": kwargs.get("input")}) 
                     liteDebuggerClient.log_event(
                         model=model,
                         messages=messages,
diff --git a/pyproject.toml b/pyproject.toml
index 94a7262f0d..1a444d0629 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "litellm"
-version = "0.1.458"
+version = "0.1.459"
 description = "Library to easily interface with LLM API providers"
 authors = ["BerriAI"]
 license = "MIT License"