From bc23a9266eac24346164fceeef3a3451a2af376f Mon Sep 17 00:00:00 2001
From: Krrish Dholakia <krrishdholakia@gmail.com>
Date: Thu, 8 Feb 2024 16:37:33 -0800
Subject: [PATCH] fix(langfuse.py): support passing input params for langfuse
 errors

---
 litellm/integrations/langfuse.py | 55 +++++++++++++-------------------
 1 file changed, 23 insertions(+), 32 deletions(-)

diff --git a/litellm/integrations/langfuse.py b/litellm/integrations/langfuse.py
index bb49cb048af..bed3d32591a 100644
--- a/litellm/integrations/langfuse.py
+++ b/litellm/integrations/langfuse.py
@@ -266,37 +266,28 @@ class LangFuseLogger:
 
             trace = self.Langfuse.trace(**trace_params)
 
-            if level == "ERROR":
-                trace.generation(
-                    level="ERROR",  # can be any of DEBUG, DEFAULT, WARNING or ERROR
-                    status_message=output,  # can be any string (e.g. stringified stack trace or error body)
-                )
-                print(f"SUCCESSFULLY LOGGED ERROR")
-            else:
-                # get generation_id
-                generation_id = None
-                if (
-                    response_obj is not None
-                    and response_obj.get("id", None) is not None
-                ):
-                    generation_id = litellm.utils.get_logging_id(
-                        start_time, response_obj
-                    )
-                trace.generation(
-                    name=generation_name,
-                    id=metadata.get("generation_id", generation_id),
-                    startTime=start_time,
-                    endTime=end_time,
-                    model=kwargs["model"],
-                    modelParameters=optional_params,
-                    input=input,
-                    output=output,
-                    usage={
-                        "prompt_tokens": response_obj["usage"]["prompt_tokens"],
-                        "completion_tokens": response_obj["usage"]["completion_tokens"],
-                        "total_cost": cost if supports_costs else None,
-                    },
-                    metadata=metadata,
-                )
+            generation_id = None
+            usage = None
+            if response_obj is not None and response_obj.get("id", None) is not None:
+                generation_id = litellm.utils.get_logging_id(start_time, response_obj)
+                usage = {
+                    "prompt_tokens": response_obj["usage"]["prompt_tokens"],
+                    "completion_tokens": response_obj["usage"]["completion_tokens"],
+                    "total_cost": cost if supports_costs else None,
+                }
+            trace.generation(
+                name=generation_name,
+                id=metadata.get("generation_id", generation_id),
+                startTime=start_time,
+                endTime=end_time,
+                model=kwargs["model"],
+                modelParameters=optional_params,
+                input=input,
+                output=output,
+                usage=usage,
+                metadata=metadata,
+                level=level,
+                status_message=output,
+            )
         except Exception as e:
             print(f"Langfuse Layer Error - {traceback.format_exc()}")