Update utils/vector_store.py

utils/vector_store.py

@@ -112,47 +112,32 @@ class VectorStore:
         
         return chunks
 
-    def similarity_search(
-        self,
-        query: str,
-        k: int = 5,
-        threshold: float = 0.5,
-        filter_criteria: Dict[str, List] = None
-    ) -> List[Dict]:
-        """Enhanced similarity search with filtering and re-ranking."""
-        # Encode query
+    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 []
+
         query_vector = self.model.encode(query, convert_to_tensor=True)
-        
-        # Calculate similarities and filter results
         results = []
+
         for doc in self.vectors:
-            # Apply filters if specified
-            if filter_criteria:
-                skip = False
-                for key, values in filter_criteria.items():
-                    doc_value = self._get_nested_dict_value(doc["metadata"], key)
-                    if doc_value not in values:
-                        skip = True
-                        break
-                if skip:
-                    continue
-            
-            # Calculate similarity
             similarity = util.pytorch_cos_sim(query_vector, doc["vector"]).item()
-            if similarity >= threshold:
-                results.append({
-                    **doc,
-                    "score": similarity
-                })
-        
-        # Sort by similarity score
+            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)
-        
-        # Re-rank results based on chunk position and metadata
-        reranked_results = self._rerank_results(results[:k*2], query)
-        
-        return reranked_results[:k]
+        return results[:k]
 
+    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: