improve searching algorithm

1 files changed, 104 insertions, 57 deletions

diff --git a/cogs/music/play_utils.py b/cogs/music/play_utils.py
index bc20d7b..ad2a2e1 100644
--- a/cogs/music/play_utils.py
+++ b/cogs/music/play_utils.py
@@ -1,31 +1,59 @@
 import json
 import logging
+import math
+import os
+import asyncio
+from pathlib import Path
 
+import discord
 import mutagen
-from youtubesearchpython import VideosSearch
 import yt_dlp
-import discord
-import asyncio
-import os
+from rapidfuzz import fuzz
+from youtubesearchpython import VideosSearch
 
 logger = logging.getLogger(__name__)
 
 class Utils:
     @staticmethod
-    def levenshtein_distance(s1, s2):
-        if len(s1) > len(s2):
-            s1, s2 = s2, s1
-        distances = range(len(s1) + 1)
-        for i2, c2 in enumerate(s2):
-            distances_ = [i2+1]
-            for i1, c1 in enumerate(s1):
-                if c1 == c2:
-                    distances_.append(distances[i1])
-                else:
-                    distances_.append(
-                        min(distances[i1], distances[i1 + 1], distances_[-1]) + 1)
-            distances = distances_
-        return distances[-1]
+    def normalize_name(p: Path, keep_ext=False):
+        """Preprocess filenames: lower, replace separators, strip extensions optionally."""
+        name = str(p.name if keep_ext else p.stem)
+        return name.lower().replace('_', ' ').replace('-', ' ')
+    
+    @staticmethod
+    def length_penalty(query, candidate):
+        """Penalize tiny filenames when query is long; returns multiplier [0.0,1.0]."""
+        ql = len(query)
+        cl = len(candidate)
+        if ql == 0: 
+            return 1.0
+        ratio = cl / ql
+        # Prefer candidate lengths near query length; gentle falloff otherwise
+        return math.exp(-0.5 * ((math.log(max(ratio, 1e-6)))**2))
+    
+    @staticmethod
+    def combined_score(query, candidate):
+        """Calculate a combined fuzzy matching score using multiple strategies."""
+        # Use several scorers and combine with weights
+        r_ratio = fuzz.ratio(query, candidate)            # overall edit similarity
+        r_partial = fuzz.partial_ratio(query, candidate)  # substring allowance
+        r_token_set = fuzz.token_set_ratio(query, candidate)  # word overlap ignoring extras
+        
+        # Calculate prefix score using partial_ratio of the start of the strings
+        min_len = min(len(query), len(candidate))
+        if min_len > 0:
+            r_prefix = fuzz.ratio(query, candidate[:min_len])  # prefix matching
+        else:
+            r_prefix = 0
+
+        # Weighted sum (tune weights as needed)
+        score = (0.35 * r_token_set +  # Increased weight for token set
+                0.30 * r_partial +    # Slightly increased weight for partial matches
+                0.25 * r_ratio +      # Standard ratio
+                0.10 * r_prefix)      # Reduced weight for prefix matches
+                
+        score *= Utils.length_penalty(query, candidate)  # apply length penalty
+        return min(100, score)  # Ensure score doesn't exceed 100
 
 class PlayUtils:
     def __init__(self, bot):
@@ -168,65 +196,84 @@ class PlayUtils:
     # -- LOCAL PLAY --
 
     @staticmethod
-    async def search_local(query):
+    async def search_local(query, max_results=5, min_score=30):
+        """Search for audio files using fuzzy matching on both filename and path.
+        
+        Args:
+            query: Search query string
+            max_results: Maximum number of results to return
+            min_score: Minimum score threshold for matches (0-100)
+            
+        Returns:
+            List of matching file paths, sorted by relevance
+        """
         logger.info(f"Starting local search for query: '{query}'")
         try:
             with open('settings.json') as f:
                 data = json.load(f)
-                paths = data.get('audio_paths', [])
-                logger.info(f"Searching in paths: {paths}")
+                search_paths = data.get('audio_paths', [])
+                logger.info(f"Searching in paths: {search_paths}")
+                
                 candidates = []
                 files_checked = 0
+                q = query.lower().replace('_', ' ').replace('-', ' ')
                 
-                for path in paths:
-                    if not os.path.exists(path):
+                for path in search_paths:
+                    path = Path(path)
+                    if not path.exists():
                         logger.warning(f"Path does not exist: {path}")
                         continue
                         
                     logger.debug(f"Searching in directory: {path}")
+                    
                     # Walk through all subdirectories
-                    for root, dirs, files in os.walk(path):
-                        logger.debug(f"Scanning directory: {root} (contains {len(files)} files)")
-                        for file in files:
-                            files_checked += 1
-                            # Check if file has a valid extension
-                            if file.lower().endswith(('.mp4', '.mp3', '.wav', '.opus', '.m4a')):
-                                # Get relative path from the search path
-                                rel_path = os.path.relpath(root, path)
-                                
-                                # Get filename without extension
-                                filename_without_ext = os.path.splitext(file)[0]
-                                
-                                # Create searchable string with both filename and relative path
-                                searchable_name = f"{os.path.join(rel_path, filename_without_ext)}".lower()
-                                
-                                # Calculate Levenshtein distance for the searchable string
-                                distance = Utils.levenshtein_distance(query.lower(), searchable_name)
-                                
-                                # Also check just the filename for better matching of partial queries
-                                filename_distance = Utils.levenshtein_distance(query.lower(), filename_without_ext.lower())
-                                
-                                # Use the better (lower) distance score
-                                final_distance = min(distance, filename_distance)
-                                
-                                # Add to candidates with distance and full path
-                                file_path = os.path.abspath(os.path.join(root, file))
-                                candidates.append((final_distance, file_path))
-                                logger.debug(f"Found match: {file_path} (distance: {distance})")
+                    for file_path in path.rglob('*'):
+                        if not file_path.is_file():
+                            continue
+                            
+                        files_checked += 1
+                        
+                        # Check if file has a valid extension
+                        if file_path.suffix.lower() not in {'.mp4', '.mp3', '.wav', '.opus', '.m4a'}:
+                            continue
+                            
+                        # Get relative path from search root
+                        rel_path = file_path.relative_to(path).parent
+                        
+                        # Create searchable strings
+                        filename = Utils.normalize_name(file_path, keep_ext=False)
+                        full_path_str = str(rel_path / file_path.stem).lower()
+                        
+                        # Calculate scores
+                        score_filename = Utils.combined_score(q, filename)
+                        score_path = Utils.combined_score(q, full_path_str)
+                        
+                        # Use the better score
+                        score = max(score_filename, score_path)
+                        
+                        # Apply bonuses
+                        depth_bonus = max(0, 1.0 - (len(rel_path.parts) / 10.0))  # Prefer shallower paths
+                        ext_bonus = 1.1 if file_path.suffix.lower() in q.split('.')[-1].lower() else 1.0
+                        
+                        final_score = score * (0.9 + depth_bonus * 0.1) * ext_bonus
+                        
+                        if final_score >= min_score:
+                            candidates.append((final_score, str(file_path.absolute())))
+                            logger.debug(f"Match: {file_path} (score: {final_score:.1f})")
                 
-                logger.info(f"Search complete. Checked {files_checked} files, found {len(candidates)} audio files.")
+                logger.info(f"Search complete. Checked {files_checked} files, found {len(candidates)} matches.")
                 
                 if not candidates:
                     logger.warning("No matching audio files found.")
                     return []
                 
-                # Sort by distance and get top 5 matches
-                sorted_candidates = sorted(candidates, key=lambda x: x[0])
-                top_matches = sorted_candidates[:5]
+                # Sort by score (descending) and get top results
+                candidates.sort(key=lambda x: x[0], reverse=True)
+                top_matches = candidates[:max_results]
                 
                 logger.info(f"Top {len(top_matches)} matches:")
-                for i, (distance, path) in enumerate(top_matches, 1):
-                    logger.info(f"  {i}. {path} (distance: {distance})")
+                for i, (score, path) in enumerate(top_matches, 1):
+                    logger.info(f"  {i}. {path} (score: {score:.1f})")
                 
                 return [c[1] for c in top_matches]