Update search_utils.py

search_utils.py

@@ -36,7 +36,6 @@ class MetadataManager:
         self._ensure_directories()
         self._unzip_if_needed()
         self._build_shard_map()
-        self._init_url_resolver()
         logger.info(f"Total documents indexed: {self.total_docs}")
         logger.info(f"Total shards found: {len(self.shard_map)}")
     
@@ -145,12 +144,15 @@ class MetadataManager:
                 shard_path = self.shard_dir / shard
                 if not shard_path.exists():
                     logger.error(f"Shard file not found: {shard_path}")
-                    return pd.DataFrame(columns=["title", "summary", "similarity", "source"])
+                    return pd.DataFrame(columns=["title", "summary", "similarity","authors", "source"])
+                    
                 file_size_mb = os.path.getsize(shard_path) / (1024 * 1024)
                 logger.info(f"Loading shard file: {shard} (size: {file_size_mb:.2f} MB)")
+                
                 try:
-                    self.loaded_shards[shard] = pd.read_parquet(shard_path, columns=["title", "summary"])
+                    self.loaded_shards[shard] = pd.read_parquet(shard_path, columns=["title", "summary", "source", "authors"])
                     logger.info(f"Loaded shard {shard} with {len(self.loaded_shards[shard])} rows")
+                    
                 except Exception as e:
                     logger.error(f"Failed to read parquet file {shard}: {str(e)}")
                     try:
@@ -158,7 +160,7 @@ class MetadataManager:
                         logger.info(f"Parquet schema: {schema}")
                     except Exception:
                         pass
-                    return pd.DataFrame(columns=["title", "summary", "similarity", "source"])
+                    return pd.DataFrame(columns=["title", "summary", "similarity", "source", "authors"])
             df = self.loaded_shards[shard]
             df_len = len(df)
             valid_local_indices = [idx for idx in local_indices if 0 <= idx < df_len]
@@ -168,9 +170,10 @@ class MetadataManager:
                 chunk = df.iloc[valid_local_indices]
                 logger.info(f"Retrieved {len(chunk)} records from shard {shard}")
                 return chunk
+                
         except Exception as e:
             logger.error(f"Error processing shard {shard}: {str(e)}", exc_info=True)
-        return pd.DataFrame(columns=["title", "summary", "similarity", "source"])
+        return pd.DataFrame(columns=["title", "summary", "similarity", "source", "authors"])
     
     def get_metadata(self, global_indices):
         """Retrieve metadata for a batch of global indices using parallel shard processing."""
@@ -328,14 +331,17 @@ class SemanticSearch:
         if index.ntotal == 0:
             self.logger.warning(f"Skipping empty shard {shard_idx}")
             return None
+            
         try:
             shard_start = time.time()
             distances, indices = index.search(query_embedding, top_k)
             valid_mask = (indices[0] >= 0) & (indices[0] < index.ntotal)
             valid_indices = indices[0][valid_mask].tolist()
             valid_distances = distances[0][valid_mask].tolist()
+            
             if len(valid_indices) != top_k:
                 self.logger.debug(f"Shard {shard_idx}: Found {len(valid_indices)} valid results out of {top_k}")
+                
             global_indices = [self._global_index(shard_idx, idx) for idx in valid_indices]
             self.logger.debug(f"Shard {shard_idx} search completed in {time.time() - shard_start:.3f}s")
             return valid_distances, global_indices
@@ -348,21 +354,23 @@ class SemanticSearch:
         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", "similarity"])
+            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", "similarity"])
+                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:
@@ -370,13 +378,21 @@ class SemanticSearch:
                                   f"max={results['similarity'].max():.3f}, " +
                                   f"mean={results['similarity'].mean():.3f}")
             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", "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)
         except Exception as e:
             self.logger.error(f"Result processing failed: {str(e)}", exc_info=True)