Update rag.py

rag.py

@@ -96,6 +96,84 @@ class SessionRAG:
             
         return embeddings.cpu().numpy()[0]
 
+    def load_existing_session_data(self, chunks_from_db: List[Dict[str, Any]]):
+        """Load pre-existing chunks with embeddings from database"""
+        print(f"Loading existing session data for {self.session_id}: {len(chunks_from_db)} chunks...")
+        
+        # Process chunks from MongoDB format
+        self.chunks_data = self.process_db_chunks(chunks_from_db)
+        
+        # Rebuild indices from existing embeddings (don't recreate embeddings)
+        self.rebuild_indices_from_existing_embeddings()
+        
+        print(f"Session {self.session_id} loaded with existing embeddings!")
+    
+    def rebuild_indices_from_existing_embeddings(self):
+        """Rebuild search indices using existing embeddings from database"""
+        if not self.chunks_data:
+            raise ValueError("No chunks data available")
+        
+        print(f"Rebuilding indices from existing embeddings...")
+        
+        # Extract existing embeddings
+        embeddings = []
+        for chunk in self.chunks_data:
+            if 'embedding' in chunk and chunk['embedding'] is not None:
+                embeddings.append(chunk['embedding'])
+            else:
+                raise ValueError(f"Missing embedding for chunk {chunk.get('id', 'unknown')}")
+        
+        # Build FAISS index from existing embeddings
+        embeddings_matrix = np.vstack(embeddings)
+        self.dense_index = faiss.IndexFlatIP(embeddings_matrix.shape[1])
+        self.dense_index.add(embeddings_matrix.astype('float32'))
+        
+        # Build other indices
+        tokenized_corpus = [chunk['text'].lower().split() for chunk in self.chunks_data]
+        self.bm25_index = BM25Okapi(tokenized_corpus)
+        
+        # 3. ColBERT-style token index
+        self.token_to_chunks = defaultdict(set)
+        for i, chunk in enumerate(self.chunks_data):
+            tokens = chunk['text'].lower().split()
+            for token in tokens:
+                self.token_to_chunks[token].add(i)
+        
+        # 4. Legal concept graph
+        self.concept_graph = nx.Graph()
+        for i, chunk in enumerate(self.chunks_data):
+            self.concept_graph.add_node(i, text=chunk['text'][:200], importance=chunk['importance_score'])
+            
+            for j, other_chunk in enumerate(self.chunks_data[i+1:], i+1):
+                shared_entities = set(e['text'] for e in chunk['entities']) & \
+                                set(e['text'] for e in other_chunk['entities'])
+                if shared_entities:
+                    self.concept_graph.add_edge(i, j, weight=len(shared_entities))
+        
+        print(f"All indices rebuilt from existing embeddings for session {self.session_id}!")
+
+    def process_db_chunks(self, chunks_from_db: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Convert MongoDB chunk format to internal format"""
+        processed_chunks = []
+        for chunk in chunks_from_db:
+            # Convert embedding from list to numpy array if needed
+            embedding = chunk.get('embedding')
+            if embedding and isinstance(embedding, list):
+                embedding = np.array(embedding)
+            
+            processed_chunk = {
+                'id': chunk.get('chunk_id', chunk.get('id')),
+                'text': chunk.get('content', chunk.get('text', '')),
+                'title': chunk.get('title', 'Document'),
+                'section_type': chunk.get('section_type', 'general'),
+                'importance_score': chunk.get('importance_score', 1.0),
+                'entities': chunk.get('entities', []),
+                'embedding': embedding
+            }
+            processed_chunks.append(processed_chunk)
+        
+        return processed_chunks
+
     def extract_legal_entities(self, text: str) -> List[Dict[str, Any]]:
         """Extract legal entities from text"""
         entities = []