Update utils/vector_store.py

utils/vector_store.py

@@ -5,12 +5,13 @@ from sentence_transformers import SentenceTransformer, util
 import numpy as np
 from datetime import datetime
 
+
 class VectorStore:
     def __init__(self, storage_path: str = "data/vector_store", model_name: str = 'all-MiniLM-L6-v2'):
         """Initialize VectorStore with improved chunk handling."""
         self.storage_path = storage_path
         os.makedirs(storage_path, exist_ok=True)
-        
+
         self.model = SentenceTransformer(model_name)
         self.vectors = self._load_vectors()
         self.chunk_size = 512  # Optimal size for most transformer models
@@ -32,11 +33,11 @@ class VectorStore:
         """Save vectors with backup and atomic write."""
         vector_file = os.path.join(self.storage_path, "vectors.pkl")
         backup_file = vector_file + ".backup"
-        
+
         # Create backup of existing vectors
         if os.path.exists(vector_file):
             os.replace(vector_file, backup_file)
-        
+
         try:
             with open(vector_file, "wb") as f:
                 pickle.dump(self.vectors, f)
@@ -53,7 +54,7 @@ class VectorStore:
         """Add document with improved chunking and metadata."""
         # Create chunks with overlap
         chunks = self._create_chunks(text)
-        
+
         # Add timestamp and chunk info to metadata
         base_metadata = {
             **metadata,
@@ -69,10 +70,10 @@ class VectorStore:
                 "chunk_idx": chunk_idx,
                 "chunk_text": chunk[:200]  # Store preview of chunk text
             }
-            
+
             # Encode chunk
             vector = self.model.encode(chunk, convert_to_tensor=True)
-            
+
             # Store chunk with metadata
             self.vectors.append({
                 "doc_id": f"{doc_id}_chunk_{chunk_idx}",
@@ -80,7 +81,7 @@ class VectorStore:
                 "text": chunk,
                 "metadata": chunk_metadata
             })
-        
+
         self._save_vectors()
 
     def _create_chunks(self, text: str) -> List[str]:
@@ -90,10 +91,10 @@ class VectorStore:
         chunks = []
         current_chunk = []
         current_size = 0
-        
+
         for sentence in sentences:
             sentence_size = len(sentence.split())
-            
+
             if current_size + sentence_size > self.chunk_size:
                 # Save current chunk
                 if current_chunk:
@@ -105,39 +106,39 @@ class VectorStore:
             else:
                 current_chunk.append(sentence)
                 current_size += sentence_size
-        
+
         # Add final chunk
         if current_chunk:
             chunks.append(' '.join(current_chunk))
-        
+
         return chunks
 
     def similarity_search(self, query: str, k: int = 3) -> List[Dict]:
-    """Perform similarity search with error handling."""
-    try:
-        # If no vectors are stored yet, return empty list
-        if not self.vectors:
-            return []
+        """Perform similarity search with error handling."""
+        try:
+            # If no vectors are stored yet, return empty list
+            if not self.vectors:
+                return []
 
-        query_vector = self.model.encode(query, convert_to_tensor=True)
-        results = []
+            query_vector = self.model.encode(query, convert_to_tensor=True)
+            results = []
 
-        for doc in self.vectors:
-            similarity = util.pytorch_cos_sim(query_vector, doc["vector"]).item()
-            results.append({
-                "text": doc["text"],
-                "metadata": doc["metadata"],
-                "score": similarity
-            })
+            for doc in self.vectors:
+                similarity = util.pytorch_cos_sim(query_vector, doc["vector"]).item()
+                results.append({
+                    "text": doc["text"],
+                    "metadata": doc["metadata"],
+                    "score": similarity
+                })
 
-        # Sort by similarity and return top k
-        results.sort(key=lambda x: x["score"], reverse=True)
-        return results[:k]
+            # Sort by similarity and return top k
+            results.sort(key=lambda x: x["score"], reverse=True)
+            return results[:k]
+
+        except Exception as e:
+            print(f"Error in similarity search: {str(e)}")
+            return []
 
-    except Exception as e:
-        st.error(f"Error in similarity search: {str(e)}")
-        return []
-        
     def _rerank_results(self, results: List[Dict], query: str) -> List[Dict]:
         """Re-rank results considering chunk position and metadata relevance."""
         for result in results:
@@ -145,24 +146,24 @@ class VectorStore:
             chunk_idx = result["metadata"].get("chunk_idx", 0)
             total_chunks = result["metadata"].get("total_chunks", 1)
             position_score = 1 - (chunk_idx / total_chunks)  # Favor earlier chunks
-            
+
             # Adjust score based on metadata relevance
             metadata_score = self._calculate_metadata_relevance(result["metadata"], query)
-            
+
             # Combine scores
             result["final_score"] = (
                 result["score"] * 0.6 +  # Base similarity
                 position_score * 0.2 +    # Position importance
                 metadata_score * 0.2      # Metadata relevance
             )
-        
+
         return sorted(results, key=lambda x: x["final_score"], reverse=True)
 
     def _calculate_metadata_relevance(self, metadata: Dict, query: str) -> float:
         """Calculate relevance score based on metadata matching."""
         relevance_score = 0.0
         query_lower = query.lower()
-        
+
         # Check for metadata field matches
         for key, value in metadata.items():
             if isinstance(value, str):
@@ -170,19 +171,8 @@ class VectorStore:
                     relevance_score += 0.2
                 elif query_lower in value.lower():
                     relevance_score += 0.1
-        
-        return min(1.0, relevance_score)  # Normalize to [0,1]
 
-    def _get_nested_dict_value(self, d: Dict, key_path: str):
-        """Get value from nested dictionary using dot notation."""
-        keys = key_path.split('.')
-        value = d
-        for key in keys:
-            if isinstance(value, dict):
-                value = value.get(key)
-            else:
-                return None
-        return value
+        return min(1.0, relevance_score)  # Normalize to [0,1]
 
     def get_document_embeddings(self, doc_id: str) -> List[Dict]:
         """Retrieve all embeddings for a specific document."""
@@ -190,8 +180,7 @@ class VectorStore:
 
     def delete_document(self, doc_id: str):
         """Delete all chunks associated with a document."""
-        self.vectors = [doc for doc in self.vectors 
-                       if doc["metadata"]["doc_id"] != doc_id]
+        self.vectors = [doc for doc in self.vectors if doc["metadata"]["doc_id"] != doc_id]
         self._save_vectors()
 
     def update_metadata(self, doc_id: str, metadata_updates: Dict):
@@ -199,4 +188,4 @@ class VectorStore:
         for doc in self.vectors:
             if doc["metadata"]["doc_id"] == doc_id:
                 doc["metadata"].update(metadata_updates)
-        self._save_vectors()
+        self._save_vectors()