fix(langfuse.py): support time to first token logging on langfuse

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()}")

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"]