Update vector.py

vector.py

@@ -191,7 +191,7 @@ class VectorDatabase:
     
     # ==================== ENHANCED STORAGE WITH CACHE INVALIDATION ====================
     
-    def store_session_document(self, text: str, filename: str, user_id: str, chat_id: str) -> bool:
+    def store_session_document(self, text: str, filename: str, user_id: str, chat_id: str, file_id: str = None) -> bool:
         """Store extracted file content with enhanced chunking and cache invalidation"""
         if not text or len(text) < 10 or not user_id:
             logger.warning(f"Invalid input for {filename}")
@@ -233,6 +233,7 @@ class VectorDatabase:
             final_meta.append({
                 "text": chunk["text"],
                 "source": filename,
+                "file_id": file_id,
                 "type": "file",
                 "subtype": chunk.get("type", "general"),
                 "name": chunk.get("name", "unknown"),
@@ -249,6 +250,7 @@ class VectorDatabase:
             "text": whole_file_text,
             "actual_content": text,
             "source": filename,
+            "file_id": file_id,
             "type": "file",
             "subtype": "whole_file",
             "is_whole_file": True,
@@ -295,6 +297,52 @@ class VectorDatabase:
                     self._rollback_partial_storage(user_id, chat_id)
             return False
     
+    def delete_file(self, user_id: str, chat_id: str, file_id: str) -> bool:
+        """Surgical Strike: Remove chunks belonging to a specific file ID"""
+        with self.memory_lock:
+            new_metadata = []
+            removed_count = 0
+            
+            # Filter loop: Keep everything that DOESN'T match our file_id
+            for meta in self.metadata:
+                # Check matches: Must match User + Chat + FileID
+                if (meta.get("user_id") == user_id and 
+                    meta.get("chat_id") == chat_id and 
+                    meta.get("file_id") == file_id):
+                    removed_count += 1
+                else:
+                    new_metadata.append(meta)
+            
+            if removed_count == 0:
+                logger.info(f"ℹ️ No vectors found for file_id {file_id}")
+                return False
+
+            logger.info(f"🧹 Surgically removing {removed_count} vectors for file {file_id}...")
+
+            # Rebuild Index (Standard Faiss Pattern)
+            if not new_metadata:
+                self.index = faiss.IndexFlatIP(384)
+            else:
+                surviving_texts = [m["text"] for m in new_metadata]
+                try:
+                    embeddings = self.embedder.encode(surviving_texts, show_progress_bar=False)
+                    faiss.normalize_L2(embeddings)
+                    
+                    new_index = faiss.IndexFlatIP(384)
+                    new_index.add(np.array(embeddings).astype('float32'))
+                    self.index = new_index
+                except Exception as e:
+                    logger.error(f"❌ Rebuild failed during file deletion: {e}")
+                    return False
+
+            self.metadata = new_metadata
+            self._save_index()
+            
+            # Invalidate Cache
+            self._invalidate_bm25_cache(user_id, chat_id)
+            
+            logger.info(f"✅ Successfully deleted file {file_id}")
+            return True
     # ==================== UPDATED BM25 SEARCH WITH LAZY LOADING ====================
     
     def bm25_search(self, query: str, user_id: str, chat_id: str,