From bca4f84c24fa54ee3b50e48a8cece4a96792148c Mon Sep 17 00:00:00 2001
From: Marc Klingen <git@marcklingen.com>
Date: Mon, 6 Jan 2025 11:19:10 +0100
Subject: [PATCH] fix: include assistant message as trace.output in langfuse

---
 examples/filters/langfuse_filter_pipeline.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/examples/filters/langfuse_filter_pipeline.py b/examples/filters/langfuse_filter_pipeline.py
index 8576ef2..cb0135e 100644
--- a/examples/filters/langfuse_filter_pipeline.py
+++ b/examples/filters/langfuse_filter_pipeline.py
@@ -44,6 +44,7 @@ class Pipeline:
             }
         )
         self.langfuse = None
+        self.chat_traces = {}
         self.chat_generations = {}
 
     async def on_startup(self):
@@ -107,17 +108,18 @@ class Pipeline:
             metadata={"interface": "open-webui"},
         )
 
+        self.chat_traces[body["chat_id"]] = trace
         self.chat_generations[body["chat_id"]] = generation
-        print(trace.get_trace_url())
 
         return body
 
     async def outlet(self, body: dict, user: Optional[dict] = None) -> dict:
         print(f"outlet:{__name__}")
         print(f"Received body: {body}")
-        if body["chat_id"] not in self.chat_generations:
+        if body["chat_id"] not in self.chat_generations or body["chat_id"] not in self.chat_traces:
             return body
 
+        trace = self.chat_traces[body["chat_id"]]
         generation = self.chat_generations[body["chat_id"]]
         assistant_message = get_last_assistant_message(body["messages"])
 
@@ -138,6 +140,9 @@ class Pipeline:
                     }
 
         # Update generation
+        trace.update(
+            output=assistant_message,
+        )
         generation.end(
             output=assistant_message,
             metadata={"interface": "open-webui"},
@@ -145,6 +150,7 @@ class Pipeline:
         )
 
         # Clean up the chat_generations dictionary
+        del self.chat_traces[body["chat_id"]]
         del self.chat_generations[body["chat_id"]]
 
         return body