Update core/models/knowledge_base.py

core/models/knowledge_base.py

@@ -1,15 +1,17 @@
-
-import faiss
-import pickle
 import logging
-import numpy as np
+import pickle
 from pathlib import Path
 from typing import List, Dict, Any
+
+import faiss
+import numpy as np
 from sentence_transformers import SentenceTransformer
 
 logger = logging.getLogger(__name__)
 
 class OptimizedGazaKnowledgeBase:
+    """Optimized knowledge base that loads pre-made FAISS index and assets"""
+    
     def __init__(self, vector_store_dir: str = "./vector_store"):
         self.vector_store_dir = Path(vector_store_dir)
         self.faiss_index = None
@@ -17,60 +19,121 @@ class OptimizedGazaKnowledgeBase:
         self.chunks = []
         self.metadata = []
         self.is_initialized = False
-
+        
     def initialize(self):
-        logger.info("🔄 Loading FAISS index...")
-        index_path = self.vector_store_dir / "index.faiss"
-        chunks_path = self.vector_store_dir / "chunks.txt"
-        metadata_path = self.vector_store_dir / "metadata.pkl"
-
-        if not index_path.exists() or not chunks_path.exists():
-            raise FileNotFoundError("Missing FAISS index or chunks.txt")
-
-        self.faiss_index = faiss.read_index(str(index_path))
-        self.embedding_model = SentenceTransformer('sentence-transformers/all-mpnet-base-v2')
-
-        with open(chunks_path, 'r', encoding='utf-8') as f:
-            self.chunks = [line.strip() for line in f if line.strip()]
-
-        if metadata_path.exists():
-            with open(metadata_path, 'rb') as f:
-                metadata_dict = pickle.load(f)
-                self.metadata = metadata_dict.get('metadata', [{}] * len(self.chunks))
-        else:
-            self.metadata = [{}] * len(self.chunks)
-
-        if len(self.chunks) != self.faiss_index.ntotal:
-            self.chunks = self.chunks[:self.faiss_index.ntotal]
-            self.metadata = self.metadata[:self.faiss_index.ntotal]
-
-        self.is_initialized = True
-        logger.info("✅ Knowledge base initialized.")
-
+        """Load pre-made FAISS index and associated data"""
+        try:
+            logger.info("🔄 Loading pre-made FAISS index and assets...")
+            
+            # 1. Load FAISS index
+            index_path = self.vector_store_dir / "index.faiss"
+            if not index_path.exists():
+                raise FileNotFoundError(f"FAISS index not found at {index_path}")
+            
+            self.faiss_index = faiss.read_index(str(index_path))
+            logger.info(f"✅ Loaded FAISS index: {self.faiss_index.ntotal} vectors, {self.faiss_index.d} dimensions")
+            
+            # 2. Load chunks
+            chunks_path = self.vector_store_dir / "chunks.txt"
+            if not chunks_path.exists():
+                raise FileNotFoundError(f"Chunks file not found at {chunks_path}")
+            
+            with open(chunks_path, \'r\', encoding=\'utf-8\') as f:
+                lines = f.readlines()
+            
+            # Parse chunks from the formatted file
+            current_chunk = ""
+            for line in lines:
+                line = line.strip()
+                if line.startswith("=== Chunk") and current_chunk:
+                    self.chunks.append(current_chunk.strip())
+                    current_chunk = ""
+                elif not line.startswith("===") and not line.startswith("Source:") and not line.startswith("Length:"):
+                    current_chunk += line + " "
+            
+            # Add the last chunk
+            if current_chunk:
+                self.chunks.append(current_chunk.strip())
+            
+            logger.info(f"✅ Loaded {len(self.chunks)} text chunks")
+            
+            # 3. Load metadata
+            metadata_path = self.vector_store_dir / "metadata.pkl"
+            if metadata_path.exists():
+                with open(metadata_path, \'rb\') as f:
+                    metadata_dict = pickle.load(f)
+                
+                if isinstance(metadata_dict, dict) and \'metadata\' in metadata_dict:
+                    self.metadata = metadata_dict[\'metadata\']
+                    logger.info(f"✅ Loaded {len(self.metadata)} metadata entries")
+                else:
+                    logger.warning("⚠️ Metadata format not recognized, using empty metadata")
+                    self.metadata = [{}] * len(self.chunks)
+            else:
+                logger.warning("⚠️ No metadata file found, using empty metadata")
+                self.metadata = [{}] * len(self.chunks)
+            
+            # 4. Initialize embedding model for query encoding
+            logger.info("🔄 Loading embedding model for queries...")
+            self.embedding_model = SentenceTransformer(\'sentence-transformers/all-mpnet-base-v2\')
+            logger.info("✅ Embedding model loaded")
+            
+            # 5. Verify data consistency
+            if len(self.chunks) != self.faiss_index.ntotal:
+                logger.warning(f"⚠️ Mismatch: {len(self.chunks)} chunks vs {self.faiss_index.ntotal} vectors")
+                # Trim chunks to match index size
+                self.chunks = self.chunks[:self.faiss_index.ntotal]
+                self.metadata = self.metadata[:self.faiss_index.ntotal]
+                logger.info(f"✅ Trimmed to {len(self.chunks)} chunks to match index")
+            
+            self.is_initialized = True
+            logger.info("🎉 Knowledge base initialization complete!")
+            
+        except Exception as e:
+            logger.error(f"❌ Failed to initialize knowledge base: {e}")
+            raise
+    
     def search(self, query: str, k: int = 5) -> List[Dict[str, Any]]:
+        """Search using pre-made FAISS index"""
         if not self.is_initialized:
-            raise RuntimeError("Knowledge base not initialized.")
-
-        query_vector = np.array(self.embedding_model.encode([query]), dtype=np.float32)
-        distances, indices = self.faiss_index.search(query_vector, k)
-
-        results = []
-        for distance, idx in zip(distances[0], indices[0]):
-            if 0 <= idx < len(self.chunks):
-                results.append({
-                    "text": self.chunks[idx],
-                    "score": float(1.0 / (1.0 + distance)),
-                    "chunk_index": idx,
-                    "distance": float(distance),
-                    "metadata": self.metadata[idx],
-                    "source": self.metadata[idx].get("source", "unknown")
-                })
-
-        return results
-
+            raise RuntimeError("Knowledge base not initialized")
+        
+        try:
+            # 1. Encode query
+            query_embedding = self.embedding_model.encode([query])
+            query_vector = np.array(query_embedding, dtype=np.float32)
+            
+            # 2. Search FAISS index
+            distances, indices = self.faiss_index.search(query_vector, k)
+            
+            # 3. Prepare results
+            results = []
+            for i, (distance, idx) in enumerate(zip(distances[0], indices[0])):
+                if idx >= 0 and idx < len(self.chunks):  # Valid index
+                    chunk_metadata = self.metadata[idx] if idx < len(self.metadata) else {}
+                    
+                    result = {
+                        "text": self.chunks[idx],
+                        "score": float(1.0 / (1.0 + distance)),  # Convert distance to similarity score
+                        "source": chunk_metadata.get("source", "unknown"),
+                        "chunk_index": int(idx),
+                        "distance": float(distance),
+                        "metadata": chunk_metadata
+                    }
+                    results.append(result)
+            
+            logger.info(f"🔍 Search for \'{query}...\' returned {len(results)} results")
+            return results
+            
+        except Exception as e:
+            logger.error(f"❌ Search error: {e}")
+            return []
+    
     def get_stats(self) -> Dict[str, Any]:
+        """Get knowledge base statistics"""
         if not self.is_initialized:
             return {"status": "not_initialized"}
+        
         return {
             "status": "initialized",
             "total_chunks": len(self.chunks),
@@ -79,3 +142,4 @@ class OptimizedGazaKnowledgeBase:
             "index_type": type(self.faiss_index).__name__,
             "sources": list(set(meta.get("source", "unknown") for meta in self.metadata))
         }
+