From 77dc0493030c1f1909f0dfe8438f2ad982f1f52f Mon Sep 17 00:00:00 2001
From: Marc Klingen <git@marcklingen.com>
Date: Wed, 18 Sep 2024 18:22:36 +0200
Subject: [PATCH] feat: track token usage in Langfuse (and fix memory leak)

---
 examples/filters/langfuse_filter_pipeline.py | 19 ++++++++++++++++---
 1 file changed, 16 insertions(+), 3 deletions(-)

diff --git a/examples/filters/langfuse_filter_pipeline.py b/examples/filters/langfuse_filter_pipeline.py
index 8130c20..342c107 100644
--- a/examples/filters/langfuse_filter_pipeline.py
+++ b/examples/filters/langfuse_filter_pipeline.py
@@ -113,13 +113,26 @@ class Pipeline:
             return body
 
         generation = self.chat_generations[body["chat_id"]]
+        assistant_message = get_last_assistant_message(body["messages"])
 
-        user_message = get_last_user_message(body["messages"])
-        generated_message = get_last_assistant_message(body["messages"])
+        # Extract usage information
+        info = assistant_message.get("info", {})
+        usage = None
+        if "prompt_tokens" in info and "completion_tokens" in info:
+            usage = {
+                "input": info["prompt_tokens"],
+                "output": info["completion_tokens"],
+                "unit": "TOKENS",
+            }
 
+        # Update generation
         generation.end(
-            output=generated_message,
+            output=assistant_message,
             metadata={"interface": "open-webui"},
+            usage=usage,
         )
 
+        # Clean up the chat_generations dictionary
+        del self.chat_generations[body["chat_id"]]
+
         return body