From 99d0d13e7e13b0f94033b5aa69eda9ca481ac27f Mon Sep 17 00:00:00 2001 From: salmonik Date: Fri, 26 Sep 2025 10:13:50 +0200 Subject: improve searching algorithm --- cogs/music/play_utils.py | 161 ++++++++++++++++++++++++++++++----------------- 1 file changed, 104 insertions(+), 57 deletions(-) (limited to 'cogs') 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] -- cgit v1.2.3