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 1/2] feat: track token usage in Langfuse (and fix memory leak)

---
 examples/filters/langfuse_filter_pipeline.py | 21 ++++++++++++++++----
 1 file changed, 17 insertions(+), 4 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 @@ async def outlet(self, body: dict, user: Optional[dict] = None) -> dict:
             return body
 
         generation = self.chat_generations[body["chat_id"]]
+        assistant_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",
+            }
 
-        user_message = get_last_user_message(body["messages"])
-        generated_message = get_last_assistant_message(body["messages"])
-
+        # 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

From 234571e06e6df78f43be6576b48eeeadb38d0e3c Mon Sep 17 00:00:00 2001
From: Marc Klingen <git@marcklingen.com>
Date: Wed, 18 Sep 2024 18:24:13 +0200
Subject: [PATCH 2/2] bump version

---
 examples/filters/langfuse_filter_pipeline.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/filters/langfuse_filter_pipeline.py b/examples/filters/langfuse_filter_pipeline.py
index 342c107..ada6d88 100644
--- a/examples/filters/langfuse_filter_pipeline.py
+++ b/examples/filters/langfuse_filter_pipeline.py
@@ -2,7 +2,7 @@
 title: Langfuse Filter Pipeline
 author: open-webui
 date: 2024-05-30
-version: 1.2
+version: 1.3
 license: MIT
 description: A filter pipeline that uses Langfuse.
 requirements: langfuse