Update opik_filter_pipeline.py

examples/filters/opik_filter_pipeline.py

@@ -2,9 +2,9 @@
 title: Opik Filter Pipeline
 author: open-webui
 date: 2025-03-12
-version: 1.0
+version: 1.1
 license: MIT
-description: A filter pipeline that uses Opik for LLM observability.
+description: A filter pipeline that uses Opik for LLM observability with improved error handling.
 requirements: opik
 """
 
@@ -12,6 +12,7 @@ from typing import List, Optional
 import os
 import uuid
 import json
+import time
 
 from pydantic import BaseModel
 from opik import Opik
@@ -97,11 +98,50 @@ class Pipeline:
                 f"Opik error: {e} Please re-enter your Opik credentials in the pipeline settings."
             )
 
+    def cleanup_existing_trace(self, chat_id: str):
+        """Safely cleanup existing trace and span for a chat_id"""
+        try:
+            if chat_id in self.chat_spans:
+                existing_span = self.chat_spans[chat_id]
+                try:
+                    existing_span.end(output={"status": "interrupted", "reason": "new_request_received"})
+                    self.log(f"Ended existing span for chat_id: {chat_id}")
+                except Exception as e:
+                    self.log(f"Warning: Could not end existing span: {e}")
+            
+            if chat_id in self.chat_traces:
+                existing_trace = self.chat_traces[chat_id]
+                try:
+                    existing_trace.end(output={"status": "interrupted", "reason": "new_request_received"})
+                    self.log(f"Ended existing trace for chat_id: {chat_id}")
+                except Exception as e:
+                    self.log(f"Warning: Could not end existing trace: {e}")
+            
+            # Clean up the dictionaries
+            self.chat_traces.pop(chat_id, None)
+            self.chat_spans.pop(chat_id, None)
+            self.log(f"Cleaned up existing trace/span for chat_id: {chat_id}")
+            
+        except Exception as e:
+            self.log(f"Error during cleanup for chat_id {chat_id}: {e}")
+            # Force cleanup even if there are errors
+            self.chat_traces.pop(chat_id, None)
+            self.chat_spans.pop(chat_id, None)
+
+    def cleanup_stale_traces(self, max_count: int = 100):
+        """Clean up traces if we have too many active ones"""
+        if len(self.chat_traces) > max_count:
+            self.log(f"Too many active traces ({len(self.chat_traces)}), cleaning up oldest ones")
+            # Clean up oldest traces (simple FIFO approach)
+            chat_ids_to_remove = list(self.chat_traces.keys())[:len(self.chat_traces) - max_count + 10]
+            for chat_id in chat_ids_to_remove:
+                self.cleanup_existing_trace(chat_id)
+
     async def inlet(self, body: dict, user: Optional[dict] = None) -> dict:
         """
         Inlet handles the incoming request (usually a user message).
         - If no trace exists yet for this chat_id, we create a new trace.
-        - If a trace does exist, we simply create a new span for the new user message.
+        - If a trace does exist, we clean it up and create a new one.
         """
         if self.valves.debug:
             print(f"[DEBUG] Received request: {json.dumps(body, indent=2)}")
@@ -117,7 +157,8 @@ class Pipeline:
             return body
 
         if "chat_id" not in metadata:
-            chat_id = str(uuid.uuid4())  # Regular chat messages
+            # Generate unique chat_id with timestamp for extra uniqueness
+            chat_id = f"{uuid.uuid4()}-{int(time.time() * 1000)}"
             self.log(f"Assigned normal chat_id: {chat_id}")
 
             metadata["chat_id"] = chat_id
@@ -136,13 +177,18 @@ class Pipeline:
 
         user_email = user.get("email") if user else None
 
-        assert chat_id not in self.chat_traces, (
+        # FIXED: Instead of asserting, clean up any existing trace/span
-            f"There shouldn't be a trace already exists for chat_id {chat_id}"
+        if chat_id in self.chat_traces:
-        )
+            self.log(f"Found existing trace for chat_id {chat_id}, cleaning up...")
+            self.cleanup_existing_trace(chat_id)
+
+        # Periodic cleanup to prevent memory leaks
+        self.cleanup_stale_traces()
 
         # Create a new trace and span
         self.log(f"Creating new chat trace for chat_id: {chat_id}")
 
+        try:
             # Body copy for traces and span
             trace_body = body.copy()
             span_body = body.copy()
@@ -196,6 +242,14 @@ class Pipeline:
             self.chat_spans[chat_id] = span
             self.log(f"Trace and span objects successfully created for chat_id: {chat_id}")
 
+        except Exception as e:
+            self.log(f"Error creating Opik trace/span for chat_id {chat_id}: {e}")
+            # Clean up on error
+            self.chat_traces.pop(chat_id, None)
+            self.chat_spans.pop(chat_id, None)
+            # Don't fail the request, just log the error
+            self.log(f"Continuing without Opik logging for this request")
+
         return body
 
     async def outlet(self, body: dict, user: Optional[dict] = None) -> dict:
@@ -217,6 +271,7 @@ class Pipeline:
         trace = self.chat_traces[chat_id]
         span = self.chat_spans[chat_id]
 
+        try:
             # Body copy for traces and span
             trace_body = body.copy()
             span_body = body.copy()
@@ -262,13 +317,28 @@ class Pipeline:
             self.log(f"span ended for chat_id: {chat_id}")
 
             # Chat_id is already logged as trace thread
-        span_body.pop("chat_id", None)
+            trace_body.pop("chat_id", None)
 
             # Optionally update the trace with the final assistant output
             trace.end(output=trace_body)
+            self.log(f"trace ended for chat_id: {chat_id}")
 
-        # Force the creation of a new trace and span for the next chat even if they are part of the same thread
+        except Exception as e:
-        del self.chat_traces[chat_id]
+            self.log(f"Error ending Opik trace/span for chat_id {chat_id}: {e}")
-        del self.chat_spans[chat_id]
+            # Try to end gracefully even if there are errors
+            try:
+                span.end(output={"status": "error", "error": str(e)})
+            except:
+                pass
+            try:
+                trace.end(output={"status": "error", "error": str(e)})
+            except:
+                pass
+
+        finally:
+            # Always clean up the dictionaries, even if ending the trace/span failed
+            self.chat_traces.pop(chat_id, None)
+            self.chat_spans.pop(chat_id, None)
+            self.log(f"Cleaned up trace/span references for chat_id: {chat_id}")
 
         return body