fix(langfuse.py): langfuse success logging fix

BerriAI · Feb 9, 2024 · 64fd1f7 · 64fd1f7
1 parent bc23a92
commit 64fd1f7
Showing 1 changed file with 19 additions and 14 deletions.
diff --git a/litellm/integrations/langfuse.py b/litellm/integrations/langfuse.py
@@ -275,19 +275,24 @@ def _log_langfuse_v2(
                     "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,
-            )
+
+            generation_params = {
+                "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,
+            }
+
+            if output is not None and isinstance(output, str) and level == "ERROR":
+                generation_params["statusMessage"] = output
+
+            trace.generation(**generation_params)
         except Exception as e:
             print(f"Langfuse Layer Error - {traceback.format_exc()}")