From 70a523eee741b203e95be579beb5d4dfd9cab2f1 Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Fri, 23 Feb 2024 10:04:12 -0800
Subject: [PATCH] fix(langfuse.py): support time to first token logging on
 langfuse

---
 litellm/integrations/langfuse.py |  8 ++++++++
 litellm/utils.py                 | 10 +++++++++-
 2 files changed, 17 insertions(+), 1 deletion(-)

diff --git a/litellm/integrations/langfuse.py b/litellm/integrations/langfuse.py
index f7f2634de3..bb8b62b507 100644
--- a/litellm/integrations/langfuse.py
+++ b/litellm/integrations/langfuse.py
@@ -235,6 +235,9 @@ class LangFuseLogger:
             supports_tags = Version(langfuse.version.__version__) >= Version("2.6.3")
             supports_prompt = Version(langfuse.version.__version__) >= Version("2.7.3")
             supports_costs = Version(langfuse.version.__version__) >= Version("2.7.3")
+            supports_completion_start_time = Version(
+                langfuse.version.__version__
+            ) >= Version("2.7.3")
 
             print_verbose(f"Langfuse Layer Logging - logging to langfuse v2 ")
 
@@ -308,6 +311,11 @@ class LangFuseLogger:
             if output is not None and isinstance(output, str) and level == "ERROR":
                 generation_params["statusMessage"] = output
 
+            if supports_completion_start_time:
+                generation_params["completion_start_time"] = kwargs.get(
+                    "completion_start_time", None
+                )
+
             trace.generation(**generation_params)
         except Exception as e:
             print(f"Langfuse Layer Error - {traceback.format_exc()}")
diff --git a/litellm/utils.py b/litellm/utils.py
index 8b469294cb..a2139c1430 100644
--- a/litellm/utils.py
+++ b/litellm/utils.py
@@ -820,6 +820,8 @@ class Logging:
         ## DYNAMIC LANGFUSE KEYS ##
         self.langfuse_public_key = langfuse_public_key
         self.langfuse_secret = langfuse_secret
+        ## TIME TO FIRST TOKEN LOGGING ##
+        self.completion_start_time: Optional[datetime.datetime] = None
 
     def update_environment_variables(
         self, model, user, optional_params, litellm_params, **additional_params
@@ -840,6 +842,7 @@ class Logging:
             "user": user,
             "call_type": str(self.call_type),
             "litellm_call_id": self.litellm_call_id,
+            "completion_start_time": self.completion_start_time,
             **self.optional_params,
             **additional_params,
         }
@@ -1069,6 +1072,11 @@ class Logging:
                 start_time = self.start_time
             if end_time is None:
                 end_time = datetime.datetime.now()
+            if self.completion_start_time is None:
+                self.completion_start_time = end_time
+                self.model_call_details["completion_start_time"] = (
+                    self.completion_start_time
+                )
             self.model_call_details["log_event_type"] = "successful_api_call"
             self.model_call_details["end_time"] = end_time
             self.model_call_details["cache_hit"] = cache_hit
@@ -1358,7 +1366,7 @@ class Logging:
                                 f"is complete_streaming_response in kwargs: {kwargs.get('complete_streaming_response', None)}"
                             )
                             if complete_streaming_response is None:
-                                break
+                                continue
                             else:
                                 print_verbose("reaches langfuse for streaming logging!")
                                 result = kwargs["complete_streaming_response"]