Update search_utils.py

search_utils.py

@@ -230,6 +230,7 @@ class SemanticSearch:
         self.metadata_mgr = MetadataManager()
         self.shard_sizes = []
         self.cumulative_offsets = None
+        self.total_vectors = 0
         self.logger = logging.getLogger("SemanticSearch")
         self.logger.info("Initializing SemanticSearch")
     
@@ -270,8 +271,8 @@ class SemanticSearch:
                         self.logger.info(f"Loaded index {shard_path.name} with {size} vectors")
                 except Exception as e:
                     self.logger.error(f"Error loading index {shard_path}: {str(e)}")
-        total_vectors = sum(self.shard_sizes)
-        self.logger.info(f"Total loaded vectors: {total_vectors} across {len(self.index_shards)} shards")
+        self.total_vectors = sum(self.shard_sizes)
+        self.logger.info(f"Total loaded vectors: {aelf.total_vectors} across {len(self.index_shards)} shards")
         self.cumulative_offsets = np.cumsum([0] + self.shard_sizes)
     
     def _load_single_index(self, shard_path):
@@ -348,52 +349,51 @@ class SemanticSearch:
         except Exception as e:
             self.logger.error(f"Search failed in shard {shard_idx}: {str(e)}")
             return None
-    
+
     def _process_results(self, distances, global_indices, top_k):
         """Process raw search results: retrieve metadata, calculate similarity, and deduplicate."""
         process_start = time.time()
         if global_indices.size == 0 or distances.size == 0:
             self.logger.warning("No search results to process")
             return pd.DataFrame(columns=["title", "summary", "source", "authors", "similarity"])
+    
         try:
             self.logger.info(f"Retrieving metadata for {len(global_indices)} indices")
             metadata_start = time.time()
             results = self.metadata_mgr.get_metadata(global_indices)
             self.logger.info(f"Metadata retrieved in {time.time() - metadata_start:.2f}s, got {len(results)} records")
-            
+    
             if len(results) == 0:
                 self.logger.warning("No metadata found for indices")
                 return pd.DataFrame(columns=["title", "summary", "source", "authors", "similarity"])
+    
             if len(results) != len(distances):
                 self.logger.warning(f"Mismatch between distances ({len(distances)}) and results ({len(results)})")
                 if len(results) < len(distances):
                     distances = distances[:len(results)]
                 else:
                     distances = np.pad(distances, (0, len(results) - len(distances)), 'constant', constant_values=1.0)
-                    
+    
             self.logger.debug("Calculating similarity scores")
             results['similarity'] = 1 - (distances / 2)
-            if not results.empty:
-                self.logger.debug(f"Similarity stats: min={results['similarity'].min():.3f}, " +
-                                  f"max={results['similarity'].max():.3f}, " +
-                                  f"mean={results['similarity'].mean():.3f}")
+    
+            # Ensure all required columns
             results['source'] = results["source"]
-
-            # Ensure we have all required columns
+    
             required_columns = ["title", "summary", "authors", "source", "similarity"]
             for col in required_columns:
                 if col not in results.columns:
                     results[col] = None  # Fill missing columns with None
-            
+    
             pre_dedup = len(results)
-            results = results.drop_duplicates(subset=["title","authors", "source"]).sort_values("similarity", ascending=False).head(top_k)
-            
+            results = results.drop_duplicates(subset=["title", "authors", "source"]).sort_values("similarity", ascending=False).head(top_k)
+    
             post_dedup = len(results)
             if pre_dedup > post_dedup:
                 self.logger.info(f"Removed {pre_dedup - post_dedup} duplicate results")
+    
             self.logger.info(f"Results processed in {time.time() - process_start:.2f}s, returning {len(results)} items")
-            
-            return results.reset_index(drop=True)
+            return results[required_columns].reset_index(drop=True)
         except Exception as e:
             self.logger.error(f"Result processing failed: {str(e)}", exc_info=True)
-            return pd.DataFrame(columns=["title", "summary", "source", "similarity"])
+            return pd.DataFrame(columns=["title", "summary", "source", "similarity"])