Update search_utils.py

search_utils.py

@@ -12,7 +12,6 @@ from urllib.parse import quote
 import requests
 import shutil
 import concurrent.futures
-# Optional: Uncomment if you want to use lru_cache for instance methods
 from functools import lru_cache
 
 # Configure logging
@@ -144,22 +143,28 @@ 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","authors", "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", "source", "authors"])
+                    # Load with explicit dtype for source column
+                    self.loaded_shards[shard] = pd.read_parquet(
+                        shard_path, 
+                        columns=["title", "summary", "source", "authors"],
+                        dtype={'source': 'str'}
+                    )
+                    # Convert source strings to lists
+                    self.loaded_shards[shard]['source'] = self.loaded_shards[shard]['source'].apply(
+                        lambda x: x.split("; ") if isinstance(x, str) else []
+                    )
+                    # Handle missing summaries
+                    self.loaded_shards[shard]['summary'] = self.loaded_shards[shard]['summary'].fillna("")
                     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:
-                        schema = pd.read_parquet(shard_path, engine='pyarrow').dtypes
-                        logger.info(f"Parquet schema: {schema}")
-                    except Exception:
-                        pass
                     return pd.DataFrame(columns=["title", "summary", "similarity", "source", "authors"])
             df = self.loaded_shards[shard]
             df_len = len(df)
@@ -220,8 +225,8 @@ class MetadataManager:
         else:
             logger.warning("No metadata records retrieved")
             return pd.DataFrame(columns=["title", "summary", "similarity", "source"])
-    
-    
+
+
 class SemanticSearch:
     def __init__(self):
         self.shard_dir = Path("compressed_shards")
@@ -310,7 +315,6 @@ class SemanticSearch:
         
         all_distances = []
         all_global_indices = []
-        # Run shard searches in parallel
         with concurrent.futures.ThreadPoolExecutor() as executor:
             futures = {
                 executor.submit(self._search_shard, shard_idx, index, query_embedding, top_k): shard_idx
@@ -351,7 +355,7 @@ class SemanticSearch:
             return None
 
     def _process_results(self, distances, global_indices, top_k):
-        """Process raw search results: retrieve metadata, calculate similarity, and deduplicate."""
+        """Process raw search results with correct similarity calculation."""
         process_start = time.time()
         if global_indices.size == 0 or distances.size == 0:
             self.logger.warning("No search results to process")
@@ -367,33 +371,41 @@ class SemanticSearch:
                 self.logger.warning("No metadata found for indices")
                 return pd.DataFrame(columns=["title", "summary", "source", "authors", "similarity"])
     
+            # Handle distance-results alignment
             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)
-    
-            # Ensure all required columns
-            results['source'] = results["source"]
-    
+                min_len = min(len(results), len(distances))
+                results = results.iloc[:min_len]
+                distances = distances[:min_len]
+
+            # Calculate similarity (cosine similarity = inner product for normalized embeddings)
+            results['similarity'] = distances
+            
+            # Ensure URL lists are properly formatted
+            results['source'] = results['source'].apply(
+                lambda x: [
+                    url.strip().rstrip(')')  # Clean trailing parentheses and whitespace
+                    for url in str(x).split(';')  # Split on semicolons
+                    if url.strip()  # Remove empty strings
+                ] if isinstance(x, (str, list)) else []
+            )
+
+            # Deduplicate and sort
             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"])
+                .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")
+
+            self.logger.info(f"Results processed in {time.time() - process_start:.2f}s")
             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"])