From 4e74582b6fc9c1a9c20f2c0075ddbb45720beecd Mon Sep 17 00:00:00 2001
From: Vince Loewe <vince@lunary.ai>
Date: Tue, 27 Feb 2024 22:30:32 -0800
Subject: [PATCH] fix timestamps and user

---
 litellm/integrations/lunary.py | 49 ++++++++++++++++++++--------------
 litellm/tests/test_lunary.py   |  5 ++--
 litellm/utils.py               |  3 ++-
 3 files changed, 34 insertions(+), 23 deletions(-)

diff --git a/litellm/integrations/lunary.py b/litellm/integrations/lunary.py
index fca8927d6..2cd65fada 100644
--- a/litellm/integrations/lunary.py
+++ b/litellm/integrations/lunary.py
@@ -1,15 +1,12 @@
 #### What this does ####
-#    On success + failure, log events to aispend.io
-import datetime
+#    On success + failure, log events to lunary.ai
+from datetime import datetime, timezone
 import traceback
 import dotenv
-import os
-import requests
+import subprocess
+import sys
 
-dotenv.load_dotenv()  # Loading env variables using dotenv
-import traceback
-import datetime, subprocess, sys
-import litellm
+dotenv.load_dotenv() 
 
 # convert to {completion: xx, tokens: xx}
 def parse_usage(usage):
@@ -86,8 +83,8 @@ class LunaryLogger:
         input=None,
         user_id=None,
         response_obj=None,
-        start_time=datetime.datetime.now(),
-        end_time=datetime.datetime.now(),
+        start_time=datetime.now(timezone.utc),
+        end_time=datetime.now(timezone.utc),
         error=None,
     ):
         # Method definition
@@ -102,23 +99,38 @@ class LunaryLogger:
 
             tags = litellm_params.pop("tags", None) or []
 
-            template_id = extra.pop("templateId", None),
+            if extra:
+                template_id = extra.pop("templateId", None)
+                extra.pop("extra_body", None)
+                extra.pop("user", None)
 
+            # keep only serializable types
             for param, value in extra.items():
                 if not isinstance(value, (str, int, bool, float)):
                     try:
                         extra[param] = str(value)
                     except:
-                        # if casting value to str fails don't block logging
                         pass
 
+            for param, value in metadata.items():
+                if not isinstance(value, (str, int, bool, float)):
+                    try:
+                        extra[param] = str(value)
+                    except:
+                        pass
+
+
+
             if response_obj:
                 usage = (
                     parse_usage(response_obj["usage"])
                     if "usage" in response_obj
                     else None
                 )
+
                 output = response_obj["choices"] if "choices" in response_obj else None
+
+                print(output)
             else:
                 usage = None
                 output = None
@@ -128,7 +140,7 @@ class LunaryLogger:
             else:
                 error_obj = None
 
-            print(start_time.isoformat())
+
 
             self.lunary_client.track_event(
                 type,
@@ -137,8 +149,8 @@ class LunaryLogger:
                 user_id=user_id,
                 name=model,
                 input=parse_messages(input),
-                timestamp=start_time.isoformat(),
-                # template_id=template_id,
+                timestamp=start_time.astimezone(timezone.utc).isoformat(),
+                template_id=template_id,
                 metadata=metadata,
                 runtime="litellm",
                 tags=tags,
@@ -150,14 +162,11 @@ class LunaryLogger:
                 type,
                 event,
                 run_id,
-                timestamp=end_time.isoformat(),
+                timestamp=end_time.astimezone(timezone.utc).isoformat(),
                 runtime="litellm",
                 error=error_obj,
                 output=parse_messages(output),
-                token_usage={
-                    "prompt": usage.get("prompt_tokens"),
-                    "completion": usage.get("completion_tokens"),
-                }
+                token_usage=usage
             )
 
 
diff --git a/litellm/tests/test_lunary.py b/litellm/tests/test_lunary.py
index 34b4ef8df..964843bfd 100644
--- a/litellm/tests/test_lunary.py
+++ b/litellm/tests/test_lunary.py
@@ -9,7 +9,7 @@ import litellm
 
 litellm.success_callback = ["lunary"]
 litellm.set_verbose = True
-import time
+
 
 
 def test_lunary_logging():
@@ -19,6 +19,7 @@ def test_lunary_logging():
             messages=[{"role": "user", "content": "what llm are u"}],
             max_tokens=10,
             temperature=0.2,
+            user="test-user",
         )
         print(response)
     except Exception as e:
@@ -67,4 +68,4 @@ def test_lunary_logging_with_streaming_and_metadata():
         print(e)
 
 
-# test_lunary_logging_with_streaming_and_metadata()
+test_lunary_logging_with_streaming_and_metadata()
diff --git a/litellm/utils.py b/litellm/utils.py
index 8c731d065..4bbffe0d2 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -1347,7 +1347,8 @@ class Logging:
                             event="end",
                             model=model,
                             input=input,
-                            user_id=self.model_call_details.get("user", "default"),
+                            user_id=self.model_call_details.get("user_id", self.model_call_details.get("user", None)),
+                            #user_props=self.model_call_details.get("user_props", None),
                             extra=self.model_call_details.get("optional_params", {}),
                             response_obj=result,
                             start_time=start_time,