aboutsummaryrefslogtreecommitdiffstats
path: root/cogs/music/play_utils.py
diff options
context:
space:
mode:
Diffstat (limited to 'cogs/music/play_utils.py')
-rw-r--r--cogs/music/play_utils.py168
1 files changed, 135 insertions, 33 deletions
diff --git a/cogs/music/play_utils.py b/cogs/music/play_utils.py
index b3d81a8..8de7543 100644
--- a/cogs/music/play_utils.py
+++ b/cogs/music/play_utils.py
@@ -3,6 +3,7 @@ import logging
3import math 3import math
4import os 4import os
5import asyncio 5import asyncio
6import re
6from pathlib import Path 7from pathlib import Path
7 8
8import discord 9import discord
@@ -33,27 +34,90 @@ class Utils:
33 34
34 @staticmethod 35 @staticmethod
35 def combined_score(query, candidate): 36 def combined_score(query, candidate):
36 """Calculate a combined fuzzy matching score using multiple strategies.""" 37 """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 38
42 # Calculate prefix score using partial_ratio of the start of the strings 39 Args:
43 min_len = min(len(query), len(candidate)) 40 query: The search query
44 if min_len > 0: 41 candidate: The candidate string to score against
45 r_prefix = fuzz.ratio(query, candidate[:min_len]) # prefix matching 42
46 else: 43 Returns:
47 r_prefix = 0 44 float: Score between 0-100, where higher means better match
48 45 """
49 # Weighted sum (tune weights as needed) 46 if not query or not candidate:
50 score = (0.35 * r_token_set + # Increased weight for token set 47 return 0.0
51 0.30 * r_partial + # Slightly increased weight for partial matches 48
52 0.25 * r_ratio + # Standard ratio 49 query = query.lower().strip()
53 0.10 * r_prefix) # Reduced weight for prefix matches 50 candidate = candidate.lower()
51
52 # First check for direct substring match (fast path)
53 if query in candidate:
54 # Base score of 100 for exact match, with bonus for being at word boundary
55 base_score = 100
56 # Check if match is at word boundary
57 if re.search(rf'\b{re.escape(query)}\b', candidate):
58 base_score = 110 # Slight bonus for word boundary match
59 return min(100, base_score)
60
61 # For very short queries (2-3 chars), be more lenient
62 if len(query) <= 3:
63 # Check for partial matches at word boundaries
64 if re.search(rf'\b{re.escape(query)}', candidate):
65 return 90 # High score for short query matching word start
66 if query in candidate:
67 return 80 # Still good for any substring match with short query
68
69 # For longer queries, check word presence
70 query_words = set(word for word in re.split(r'\W+', query) if len(word) > 1) # Reduced min word length to 1
71 if query_words:
72 candidate_words = set(re.split(r'\W+', candidate))
73 # Check if all query words are in candidate (with some leniency)
74 matching_words = sum(1 for word in query_words if word in candidate_words)
75 word_ratio = matching_words / len(query_words)
76
77 # If no words match at all, return 0
78 if matching_words == 0:
79 return 0.0
54 80
55 score *= Utils.length_penalty(query, candidate) # apply length penalty 81 # Calculate different similarity metrics
56 return min(100, score) # Ensure score doesn't exceed 100 82 r_ratio = fuzz.ratio(query, candidate) # Overall edit similarity
83 r_partial = fuzz.partial_ratio(query, candidate) # Best substring match
84 r_token_set = fuzz.token_set_ratio(query, candidate) # Word overlap
85
86 # Calculate word-based scores
87 query_tokens = set(query.split())
88 candidate_tokens = set(candidate.split())
89 word_intersection = len(query_tokens & candidate_tokens)
90 word_union = len(query_tokens | candidate_tokens)
91
92 # Jaccard similarity for word overlap
93 jaccard_sim = word_intersection / word_union if word_union > 0 else 0
94
95 # Calculate prefix score (bonus for matching start of words)
96 prefix_bonus = 0
97 for q_word in query_words:
98 for c_word in candidate_words:
99 if c_word.startswith(q_word):
100 prefix_bonus += (len(q_word) / max(len(c_word), 1)) * 20 # Increased weight
101
102 # Weighted sum with adjusted weights
103 score = (
104 0.40 * r_token_set + # Increased weight for token set
105 0.25 * r_ratio + # Overall similarity
106 0.15 * r_partial + # Substring matches
107 0.15 * (jaccard_sim * 100) + # Word overlap ratio
108 0.05 * min(prefix_bonus, 100) # Prefix bonus (capped at 100)
109 )
110
111 # Apply length penalty (less severe)
112 score *= min(1.0, 0.9 + (0.1 * Utils.length_penalty(query, candidate)))
113
114 # Boost score based on word ratio
115 score = score * (0.7 + (0.3 * word_ratio))
116
117 # Ensure score is within bounds
118 return min(100, max(0, score))
119
120 return 0.0 # No match found
57 121
58class PlayUtils: 122class PlayUtils:
59 def __init__(self, bot): 123 def __init__(self, bot):
@@ -235,18 +299,24 @@ class PlayUtils:
235 299
236 # -- LOCAL PLAY -- 300 # -- LOCAL PLAY --
237 301
238 @staticmethod 302 @classmethod
239 async def search_local(query, max_results=5): 303 async def search_local(cls, query: str = None, max_results: int = 25, autocomplete: bool = False):
240 """Search for audio files that contain the query in their name (case-insensitive). 304 """Search for audio files with fuzzy matching and optional autocomplete support.
241 305
242 Args: 306 Args:
243 query: Search query string 307 query: Search query string (can be None for autocomplete)
244 max_results: Maximum number of results to return 308 max_results: Maximum number of results to return
309 autocomplete: If True, returns results even with empty query
245 310
246 Returns: 311 Returns:
247 List of matching file paths 312 List of matching file paths with scores
248 """ 313 """
314 if not query and not autocomplete:
315 return []
316
317 query = (query or '').lower().strip()
249 logger.info(f"Starting local search for query: '{query}'") 318 logger.info(f"Starting local search for query: '{query}'")
319
250 try: 320 try:
251 with open('settings.json') as f: 321 with open('settings.json') as f:
252 data = json.load(f) 322 data = json.load(f)
@@ -254,7 +324,7 @@ class PlayUtils:
254 logger.info(f"Searching in paths: {search_paths}") 324 logger.info(f"Searching in paths: {search_paths}")
255 325
256 matches = [] 326 matches = []
257 q = query.lower() 327 file_count = 0
258 328
259 for path in search_paths: 329 for path in search_paths:
260 path = Path(path) 330 path = Path(path)
@@ -273,19 +343,51 @@ class PlayUtils:
273 if file_path.suffix.lower() not in {'.mp4', '.mp3', '.wav', '.opus', '.m4a'}: 343 if file_path.suffix.lower() not in {'.mp4', '.mp3', '.wav', '.opus', '.m4a'}:
274 continue 344 continue
275 345
276 # Check if query is in filename (case-insensitive) 346 file_count += 1
277 if q in file_path.name.lower(): 347 file_name = file_path.name.lower()
278 matches.append(str(file_path.absolute())) 348
279 if len(matches) >= max_results: 349 # For autocomplete, return all files if query is empty
280 logger.info(f"Found {len(matches)} matches, stopping search") 350 if not query and autocomplete:
281 return matches 351 matches.append((0, str(file_path.absolute())))
352 # For regular search, calculate fuzzy match score
353 elif query and len(query) >= 2:
354 # First do a simple substring check for debugging
355 if query in file_name:
356 logger.debug(f"Found direct match: {file_name} (contains '{query}')")
357
358 score = Utils.combined_score(query, file_name)
359 logger.debug(f"File: {file_name} - Score: {score}")
360
361 if score >= 30: # Minimum score threshold
362 matches.append((score, str(file_path.absolute())))
363
364 # Early return if we have enough high-scoring matches
365 if len(matches) >= max_results * 5: # Get more candidates for sorting
366 break
367
368 if len(matches) >= max_results * 5: # Limit memory usage
369 break
370
371 logger.info(f"Scanned {file_count} files total")
282 372
283 logger.info(f"Search complete. Found {len(matches)} matches.") 373 # Sort by score (descending) and take top results
284 return matches 374 matches.sort(reverse=True, key=lambda x: x[0])
375 result = [path for score, path in matches[:max_results]]
376
377 logger.info(f"Search complete. Found {len(result)} matches out of {file_count} files.")
378 if not result and query:
379 logger.warning(f"No matches found for query: '{query}'. Try a different search term.")
380 return result
285 381
286 except Exception as e: 382 except Exception as e:
287 logger.error(f"Error during local search: {str(e)}", exc_info=True) 383 logger.error(f"Error during local search: {str(e)}", exc_info=True)
288 return [] 384 return []
385
386 @classmethod
387 async def search_local_autocomplete(cls, query: str = None, max_results: int = 10):
388 """Wrapper for autocomplete searches that returns more results and allows empty query."""
389 # For autocomplete, we want to show more options and allow empty queries
390 return await cls.search_local(query, max_results=max_results, autocomplete=True)
289 391
290 392
291 393