Merge pull request #387 from marcklingen/fix-output-tracking-in-langf…

…use-pipe

fix: include assistant message as trace.output in langfuse

examples/filters/langfuse_filter_pipeline.py

@@ -44,6 +44,7 @@ def __init__(self):
             }
         )
         self.langfuse = None
+        self.chat_traces = {}
         self.chat_generations = {}
 
     async def on_startup(self):
@@ -107,17 +108,18 @@ async def inlet(self, body: dict, user: Optional[dict] = None) -> dict:
             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,13 +140,17 @@ async def outlet(self, body: dict, user: Optional[dict] = None) -> dict:
                     }
 
         # Update generation
+        trace.update(
+            output=assistant_message,
+        )
         generation.end(
             output=assistant_message,
             metadata={"interface": "open-webui"},
             usage=usage,
         )
 
         # Clean up the chat_generations dictionary
+        del self.chat_traces[body["chat_id"]]
         del self.chat_generations[body["chat_id"]]
 
         return body