aboutsummaryrefslogtreecommitdiffstats
path: root/cogs
diff options
context:
space:
mode:
Diffstat (limited to '')
-rw-r--r--cogs/music/play_utils.py161
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 @@
1import json 1import json
2import logging 2import logging
3import math
4import os
5import asyncio
6from pathlib import Path
3 7
8import discord
4import mutagen 9import mutagen
5from youtubesearchpython import VideosSearch
6import yt_dlp 10import yt_dlp
7import discord 11from rapidfuzz import fuzz
8import asyncio 12from youtubesearchpython import VideosSearch
9import os
10 13
11logger = logging.getLogger(__name__) 14logger = logging.getLogger(__name__)
12 15
13class Utils: 16class Utils:
14 @staticmethod 17 @staticmethod
15 def levenshtein_distance(s1, s2): 18 def normalize_name(p: Path, keep_ext=False):
16 if len(s1) > len(s2): 19 """Preprocess filenames: lower, replace separators, strip extensions optionally."""
17 s1, s2 = s2, s1 20 name = str(p.name if keep_ext else p.stem)
18 distances = range(len(s1) + 1) 21 return name.lower().replace('_', ' ').replace('-', ' ')
19 for i2, c2 in enumerate(s2): 22
20 distances_ = [i2+1] 23 @staticmethod
21 for i1, c1 in enumerate(s1): 24 def length_penalty(query, candidate):
22 if c1 == c2: 25 """Penalize tiny filenames when query is long; returns multiplier [0.0,1.0]."""
23 distances_.append(distances[i1]) 26 ql = len(query)
24 else: 27 cl = len(candidate)
25 distances_.append( 28 if ql == 0:
26 min(distances[i1], distances[i1 + 1], distances_[-1]) + 1) 29 return 1.0
27 distances = distances_ 30 ratio = cl / ql
28 return distances[-1] 31 # Prefer candidate lengths near query length; gentle falloff otherwise
32 return math.exp(-0.5 * ((math.log(max(ratio, 1e-6)))**2))
33
34 @staticmethod
35 def combined_score(query, candidate):
36 """Calculate a combined fuzzy matching score using multiple strategies."""
37 # Use several scorers and combine with weights
38 r_ratio = fuzz.ratio(query, candidate) # overall edit similarity
39 r_partial = fuzz.partial_ratio(query, candidate) # substring allowance
40 r_token_set = fuzz.token_set_ratio(query, candidate) # word overlap ignoring extras
41
42 # Calculate prefix score using partial_ratio of the start of the strings
43 min_len = min(len(query), len(candidate))
44 if min_len > 0:
45 r_prefix = fuzz.ratio(query, candidate[:min_len]) # prefix matching
46 else:
47 r_prefix = 0
48
49 # Weighted sum (tune weights as needed)
50 score = (0.35 * r_token_set + # Increased weight for token set
51 0.30 * r_partial + # Slightly increased weight for partial matches
52 0.25 * r_ratio + # Standard ratio
53 0.10 * r_prefix) # Reduced weight for prefix matches
54
55 score *= Utils.length_penalty(query, candidate) # apply length penalty
56 return min(100, score) # Ensure score doesn't exceed 100
29 57
30class PlayUtils: 58class PlayUtils:
31 def __init__(self, bot): 59 def __init__(self, bot):
@@ -168,65 +196,84 @@ class PlayUtils:
168 # -- LOCAL PLAY -- 196 # -- LOCAL PLAY --
169 197
170 @staticmethod 198 @staticmethod
171 async def search_local(query): 199 async def search_local(query, max_results=5, min_score=30):
200 """Search for audio files using fuzzy matching on both filename and path.
201
202 Args:
203 query: Search query string
204 max_results: Maximum number of results to return
205 min_score: Minimum score threshold for matches (0-100)
206
207 Returns:
208 List of matching file paths, sorted by relevance
209 """
172 logger.info(f"Starting local search for query: '{query}'") 210 logger.info(f"Starting local search for query: '{query}'")
173 try: 211 try:
174 with open('settings.json') as f: 212 with open('settings.json') as f:
175 data = json.load(f) 213 data = json.load(f)
176 paths = data.get('audio_paths', []) 214 search_paths = data.get('audio_paths', [])
177 logger.info(f"Searching in paths: {paths}") 215 logger.info(f"Searching in paths: {search_paths}")
216
178 candidates = [] 217 candidates = []
179 files_checked = 0 218 files_checked = 0
219 q = query.lower().replace('_', ' ').replace('-', ' ')
180 220
181 for path in paths: 221 for path in search_paths:
182 if not os.path.exists(path): 222 path = Path(path)
223 if not path.exists():
183 logger.warning(f"Path does not exist: {path}") 224 logger.warning(f"Path does not exist: {path}")
184 continue 225 continue
185 226
186 logger.debug(f"Searching in directory: {path}") 227 logger.debug(f"Searching in directory: {path}")
228
187 # Walk through all subdirectories 229 # Walk through all subdirectories
188 for root, dirs, files in os.walk(path): 230 for file_path in path.rglob('*'):
189 logger.debug(f"Scanning directory: {root} (contains {len(files)} files)") 231 if not file_path.is_file():
190 for file in files: 232 continue
191 files_checked += 1 233
192 # Check if file has a valid extension 234 files_checked += 1
193 if file.lower().endswith(('.mp4', '.mp3', '.wav', '.opus', '.m4a')): 235
194 # Get relative path from the search path 236 # Check if file has a valid extension
195 rel_path = os.path.relpath(root, path) 237 if file_path.suffix.lower() not in {'.mp4', '.mp3', '.wav', '.opus', '.m4a'}:
196 238 continue
197 # Get filename without extension 239
198 filename_without_ext = os.path.splitext(file)[0] 240 # Get relative path from search root
199 241 rel_path = file_path.relative_to(path).parent
200 # Create searchable string with both filename and relative path 242
201 searchable_name = f"{os.path.join(rel_path, filename_without_ext)}".lower() 243 # Create searchable strings
202 244 filename = Utils.normalize_name(file_path, keep_ext=False)
203 # Calculate Levenshtein distance for the searchable string 245 full_path_str = str(rel_path / file_path.stem).lower()
204 distance = Utils.levenshtein_distance(query.lower(), searchable_name) 246
205 247 # Calculate scores
206 # Also check just the filename for better matching of partial queries 248 score_filename = Utils.combined_score(q, filename)
207 filename_distance = Utils.levenshtein_distance(query.lower(), filename_without_ext.lower()) 249 score_path = Utils.combined_score(q, full_path_str)
208 250
209 # Use the better (lower) distance score 251 # Use the better score
210 final_distance = min(distance, filename_distance) 252 score = max(score_filename, score_path)
211 253
212 # Add to candidates with distance and full path 254 # Apply bonuses
213 file_path = os.path.abspath(os.path.join(root, file)) 255 depth_bonus = max(0, 1.0 - (len(rel_path.parts) / 10.0)) # Prefer shallower paths
214 candidates.append((final_distance, file_path)) 256 ext_bonus = 1.1 if file_path.suffix.lower() in q.split('.')[-1].lower() else 1.0
215 logger.debug(f"Found match: {file_path} (distance: {distance})") 257
258 final_score = score * (0.9 + depth_bonus * 0.1) * ext_bonus
259
260 if final_score >= min_score:
261 candidates.append((final_score, str(file_path.absolute())))
262 logger.debug(f"Match: {file_path} (score: {final_score:.1f})")
216 263
217 logger.info(f"Search complete. Checked {files_checked} files, found {len(candidates)} audio files.") 264 logger.info(f"Search complete. Checked {files_checked} files, found {len(candidates)} matches.")
218 265
219 if not candidates: 266 if not candidates:
220 logger.warning("No matching audio files found.") 267 logger.warning("No matching audio files found.")
221 return [] 268 return []
222 269
223 # Sort by distance and get top 5 matches 270 # Sort by score (descending) and get top results
224 sorted_candidates = sorted(candidates, key=lambda x: x[0]) 271 candidates.sort(key=lambda x: x[0], reverse=True)
225 top_matches = sorted_candidates[:5] 272 top_matches = candidates[:max_results]
226 273
227 logger.info(f"Top {len(top_matches)} matches:") 274 logger.info(f"Top {len(top_matches)} matches:")
228 for i, (distance, path) in enumerate(top_matches, 1): 275 for i, (score, path) in enumerate(top_matches, 1):
229 logger.info(f" {i}. {path} (distance: {distance})") 276 logger.info(f" {i}. {path} (score: {score:.1f})")
230 277
231 return [c[1] for c in top_matches] 278 return [c[1] for c in top_matches]
232 279