support for later embeddings

litellm/integrations/llmonitor.py

@@ -12,8 +12,10 @@ dotenv.load_dotenv()  # Loading env variables using dotenv
 # convert to {completion: xx, tokens: xx}
 def parse_usage(usage):
     return {
-        "completion": usage["completion_tokens"],
+        "completion":
-        "prompt": usage["prompt_tokens"],
+        usage["completion_tokens"] if "completion_tokens" in usage else 0,
+        "prompt":
+        usage["prompt_tokens"] if "prompt_tokens" in usage else 0,
     }
 
 
@@ -23,6 +25,10 @@ def parse_messages(input):
         return None
 
     def clean_message(message):
+        #if is strin, return as is
+        if isinstance(message, str):
+            return message
+
         if "message" in message:
             return clean_message(message["message"])
 
@@ -51,10 +57,11 @@ class LLMonitorLogger:
     def log_event(
             self,
             type,
+            event,
             run_id,
             model,
             print_verbose,
-            messages=None,
+            input=None,
             user_id=None,
             response_obj=None,
             start_time=datetime.datetime.now(),
@@ -67,8 +74,11 @@ class LLMonitorLogger:
                 f"LLMonitor Logging - Logging request for model {model}")
 
             if response_obj:
-                usage = parse_usage(response_obj['usage'])
+                print(response_obj)
-                output = response_obj['choices']
+                usage = parse_usage(
+                    response_obj['usage']) if 'usage' in response_obj else None
+                output = response_obj[
+                    'choices'] if 'choices' in response_obj else None
             else:
                 usage = None
                 output = None
@@ -80,19 +90,19 @@ class LLMonitorLogger:
                 error_obj = None
 
             data = [{
-                "type": "llm",
+                "type": type,
                 "name": model,
                 "runId": run_id,
                 "app": self.app_id,
                 'event': 'start',
                 "timestamp": start_time.isoformat(),
                 "userId": user_id,
-                "input": parse_messages(messages),
+                "input": parse_messages(input),
             }, {
-                "type": "llm",
+                "type": type,
                 "runId": run_id,
                 "app": self.app_id,
-                "event": type,
+                "event": event,
                 "error": error_obj,
                 "timestamp": end_time.isoformat(),
                 "userId": user_id,
@@ -100,6 +110,8 @@ class LLMonitorLogger:
                 "tokensUsage": usage,
             }]
 
+            print(data)
+
             # print_verbose(f"LLMonitor Logging - final data object: {data}")
 
             response = requests.post(

litellm/tests/test_llmonitor_integration.py

@@ -6,7 +6,7 @@ import os
 
 sys.path.insert(0, os.path.abspath('../..'))
 
-from litellm import completion
+from litellm import completion, embedding
 import litellm
 
 litellm.success_callback = ["llmonitor"]
@@ -15,10 +15,23 @@ litellm.failure_callback = ["llmonitor"]
 litellm.set_verbose = True
 
 # openai call
-first_success_test = completion(model="gpt-3.5-turbo",
+# first_success_test = completion(model="gpt-3.5-turbo",
-                                messages=[{
+#                                 messages=[{
-                                    "role": "user",
+#                                     "role": "user",
-                                    "content": "Hi 👋 - i'm openai"
+#                                     "content": "Hi 👋 - i'm openai"
-                                }])
+#                                 }])
 
-print(first_success_test)
+# print(first_success_test)
+
+
+def test_embedding_openai():
+    try:
+        response = embedding(model="text-embedding-ada-002", input=['test'])
+        # Add any assertions here to check the response
+        print(f"response: {str(response)[:50]}")
+    except Exception as e:
+        print(e)
+        # pytest.fail(f"Error occurred: {e}")
+
+
+test_embedding_openai()

litellm/utils.py

@@ -842,14 +842,21 @@ def handle_failure(exception, traceback_exception, start_time, end_time, args,
                         print_verbose=print_verbose,
                     )
                 elif callback == "llmonitor":
-                    print_verbose("reaches llmonitor for logging!")
+                    print_verbose("reaches llmonitor for logging error!")
+
                     model = args[0] if len(args) > 0 else kwargs["model"]
-                    messages = args[1] if len(args) > 1 else kwargs["messages"]
+
+                    input = args[1] if len(args) > 1 else kwargs.get(
+                        "messages", kwargs.get("input", None))
+
+                    type = 'embed' if 'input' in kwargs else 'llm'
 
                     llmonitorLogger.log_event(
-                        type="error",
+                        type=type,
+                        event="error",
                         user_id=litellm._thread_context.user,
                         model=model,
+                        input=input,
                         error=traceback_exception,
                         run_id=kwargs["litellm_call_id"],
                         start_time=start_time,
@@ -969,11 +976,18 @@ def handle_success(args, kwargs, result, start_time, end_time):
                 elif callback == "llmonitor":
                     print_verbose("reaches llmonitor for logging!")
                     model = args[0] if len(args) > 0 else kwargs["model"]
-                    messages = args[1] if len(args) > 1 else kwargs["messages"]
+
+                    input = args[1] if len(args) > 1 else kwargs.get(
+                        "messages", kwargs.get("input", None))
+
+                    #if contains input, it's 'embedding', otherwise 'llm'
+                    type = 'embed' if 'input' in kwargs else 'llm'
+
                     llmonitorLogger.log_event(
-                        type="end",
+                        type=type,
+                        event="end",
                         model=model,
-                        messages=messages,
+                        input=input,
                         user_id=litellm._thread_context.user,
                         response_obj=result,
                         start_time=start_time,