Update search_utils.py

search_utils.py

@@ -36,29 +36,46 @@ class MetadataManager:
             self.total_docs = max(self.total_docs, end + 1)
 
     def get_metadata(self, global_indices):
-        """Retrieve metadata for global indices"""
-        results = []
-        shard_groups = {}
+        """Retrieve metadata with validation"""
+        if not global_indices:
+            return pd.DataFrame(columns=["title", "summary", "source", "similarity"])
         
-        # Organize indices by their respective shards
-        for idx in global_indices:
+        # Filter valid indices
+        valid_indices = [idx for idx in global_indices if 0 <= idx < self.total_docs]
+        if not valid_indices:
+            return pd.DataFrame(columns=["title", "summary", "source", "similarity"])
+    
+        # Group indices by shard with boundary check
+        shard_groups = {}
+        for idx in valid_indices:
+            found = False
             for (start, end), shard in self.shard_map.items():
                 if start <= idx <= end:
                     if shard not in shard_groups:
                         shard_groups[shard] = []
                     shard_groups[shard].append(idx - start)
+                    found = True
                     break
-        
-        # Load and process required shards
+            if not found:
+                st.warning(f"Index {idx} out of shard range (0-{self.total_docs-1})")
+    
+        # Load and process shards
+        results = []
         for shard, local_indices in shard_groups.items():
-            if shard not in self.loaded_shards:
-                self.loaded_shards[shard] = pd.read_parquet(
-                    self.shard_dir / shard,
-                    columns=["title", "summary", "source"]
-                )
-            results.append(self.loaded_shards[shard].iloc[local_indices])
-        
-        return pd.concat(results).reset_index(drop=True)
+            try:
+                if shard not in self.loaded_shards:
+                    self.loaded_shards[shard] = pd.read_parquet(
+                        self.shard_dir / shard,
+                        columns=["title", "summary", "source"]
+                    )
+                
+                if local_indices:
+                    results.append(self.loaded_shards[shard].iloc[local_indices])
+            except Exception as e:
+                st.error(f"Error loading shard {shard}: {str(e)}")
+                continue
+    
+        return pd.concat(results).reset_index(drop=True) if results else pd.DataFrame()
 
 class SemanticSearch:
     def __init__(self):
@@ -89,19 +106,39 @@ class SemanticSearch:
         return sum(self.shard_sizes[:shard_idx]) + local_idx
 
     def search(self, query, top_k=5):
-        """Main search functionality"""
-        query_embedding = self.model.encode([query], convert_to_numpy=True)
+        """Search with validation"""
+        if not query or not self.index_shards:
+            return pd.DataFrame()
+        
+        try:
+            query_embedding = self.model.encode([query], convert_to_numpy=True)
+        except Exception as e:
+            st.error(f"Query encoding failed: {str(e)}")
+            return pd.DataFrame()
+    
         all_distances = []
         all_global_indices = []
-
-        # Search across all shards
+    
+        # Search with index validation
         for shard_idx, index in enumerate(self.index_shards):
-            distances, indices = index.search(query_embedding, top_k)
-            global_indices = [self._global_index(shard_idx, idx) for idx in indices[0]]
-            all_distances.extend(distances[0])
-            all_global_indices.extend(global_indices)
-
-        # Process and format results
-        results = self.metadata_mgr.get_metadata(all_global_indices)
-        results['similarity'] = 1 - (np.array(all_distances) / 2)  # Convert L2 to cosine
-        return results.sort_values('similarity', ascending=False).head(top_k)
+            if index.ntotal == 0:
+                continue
+                
+            try:
+                distances, indices = index.search(query_embedding, top_k)
+                valid_indices = [idx for idx in indices[0] if 0 <= idx < index.ntotal]
+                global_indices = [self._global_index(shard_idx, idx) for idx in valid_indices]
+                
+                all_distances.extend(distances[0][:len(valid_indices)])
+                all_global_indices.extend(global_indices)
+            except Exception as e:
+                st.error(f"Search failed in shard {shard_idx}: {str(e)}")
+                continue
+    
+        # Ensure equal array lengths
+        min_length = min(len(all_distances), len(all_global_indices))
+        return self._process_results(
+            np.array(all_distances[:min_length]), 
+            np.array(all_global_indices[:min_length]), 
+            top_k
+        )