aboutsummaryrefslogtreecommitdiffstats
path: root/cogs/music/search_utils.py
diff options
context:
space:
mode:
Diffstat (limited to 'cogs/music/search_utils.py')
-rw-r--r--cogs/music/search_utils.py287
1 files changed, 0 insertions, 287 deletions
diff --git a/cogs/music/search_utils.py b/cogs/music/search_utils.py
deleted file mode 100644
index bc9e335..0000000
--- a/cogs/music/search_utils.py
+++ /dev/null
@@ -1,287 +0,0 @@
1# This file is part of Senator Sperma
2# Licensed under the GPL-3.0 License
3# Copyright (C) 2026 Salmonik @ salmonik@tuta.io
4# https://codeberg.org/salmonik/senator-sperma
5
6import json
7import logging
8import re
9from pathlib import Path
10import aiofiles
11from rapidfuzz import fuzz
12
13logger = logging.getLogger(__name__)
14
15class SearchUtils:
16 """Utilities for searching and scoring audio files."""
17
18 @staticmethod
19 def normalize_name(p: Path, keep_ext=False):
20 """Preprocess filenames: lower, replace separators, strip extensions optionally."""
21 name = str(p.name if keep_ext else p.stem)
22 return name.lower().replace('_', ' ').replace('-', ' ')
23
24 @staticmethod
25 def length_penalty(query, candidate):
26 """Penalize tiny filenames when query is long; returns multiplier [0.0,1.0]."""
27 ql = len(query)
28 cl = len(candidate)
29 if ql == 0:
30 return 1.0
31 ratio = cl / ql
32 # Prefer candidate lengths near query length; gentle falloff otherwise
33 return math.exp(-0.5 * ((math.log(max(ratio, 1e-6)))**2))
34
35 @staticmethod
36 def combined_score(query, candidate):
37 """Calculate a combined fuzzy matching score using multiple strategies.
38
39 Args:
40 query: The search query
41 candidate: The candidate string to score against
42
43 Returns:
44 float: Score between 0-100, where higher means better match
45 """
46 if not query or not candidate:
47 return 0.0
48
49 query = query.lower().strip()
50 candidate = candidate.lower()
51
52 # Check for direct matches first
53 direct_match_score = SearchUtils._check_direct_matches(query, candidate)
54 if direct_match_score is not None:
55 return direct_match_score
56
57 # Handle short queries specially
58 short_query_score = SearchUtils._handle_short_queries(query, candidate)
59 if short_query_score is not None:
60 return short_query_score
61
62 # For longer queries, use complex scoring
63 return SearchUtils._calculate_complex_score(query, candidate)
64
65 @staticmethod
66 def _check_direct_matches(query: str, candidate: str) -> float:
67 """Check for direct substring matches."""
68 if query in candidate:
69 base_score = 100
70 if re.search(rf'\b{re.escape(query)}\b', candidate):
71 base_score = 110
72 return min(100, base_score)
73 return 0
74
75 @staticmethod
76 def _handle_short_queries(query: str, candidate: str) -> float:
77 """Handle scoring for short queries (2-3 chars)."""
78 if len(query) <= 3:
79 if re.search(rf'\b{re.escape(query)}', candidate):
80 return 90
81 if query in candidate:
82 return 80
83 return 0
84
85 @staticmethod
86 def _calculate_complex_score(query: str, candidate: str) -> float:
87 """Calculate complex score for longer queries."""
88 query_words = {word for word in re.split(r'\W+', query) if len(word) > 1}
89 if not query_words:
90 return 0.0
91
92 candidate_words = set(re.split(r'\W+', candidate))
93 word_ratio = SearchUtils._calculate_word_ratio(query_words, candidate_words)
94
95 if word_ratio == 0:
96 return 0.0
97
98 # Calculate various similarity metrics
99 similarity_scores = SearchUtils._get_similarity_scores(query, candidate)
100 word_scores = SearchUtils._get_word_scores(query, candidate, query_words, candidate_words)
101
102 # Calculate final weighted score
103 final_score = SearchUtils._calculate_weighted_score(similarity_scores, word_scores, word_ratio)
104
105 # Apply penalties and boosts
106 final_score = SearchUtils._apply_adjustments(final_score, query, candidate, word_ratio)
107
108 return min(100, max(0, int(final_score)))
109
110 @staticmethod
111 def _calculate_word_ratio(query_words: set, candidate_words: set) -> float:
112 """Calculate the ratio of matching words."""
113 matching_words = sum(1 for word in query_words if word in candidate_words)
114 return matching_words / len(query_words)
115
116 @staticmethod
117 def _get_similarity_scores(query: str, candidate: str) -> dict:
118 """Get various fuzzy similarity scores."""
119 return {
120 'ratio': fuzz.ratio(query, candidate),
121 'partial': fuzz.partial_ratio(query, candidate),
122 'token_set': fuzz.token_set_ratio(query, candidate)
123 }
124
125 @staticmethod
126 def _get_word_scores(query: str, candidate: str, query_words: set, candidate_words: set) -> dict:
127 """Calculate word-based scores."""
128 query_tokens = set(query.split())
129 candidate_tokens = set(candidate.split())
130
131 word_intersection = len(query_tokens & candidate_tokens)
132 word_union = len(query_tokens | candidate_tokens)
133 jaccard_sim = word_intersection / word_union if word_union > 0 else 0
134
135 prefix_bonus = SearchUtils._calculate_prefix_bonus(query_words, candidate_words)
136
137 return {
138 'jaccard': jaccard_sim,
139 'prefix': min(prefix_bonus, 100)
140 }
141
142 @staticmethod
143 def _calculate_prefix_bonus(query_words: set, candidate_words: set) -> float:
144 """Calculate prefix matching bonus."""
145 prefix_bonus = 0
146 for q_word in query_words:
147 for c_word in candidate_words:
148 if c_word.startswith(q_word):
149 prefix_bonus += (len(q_word) / max(len(c_word), 1)) * 20
150 return prefix_bonus
151
152 @staticmethod
153 def _calculate_weighted_score(similarity_scores: dict, word_scores: dict, word_ratio: float) -> float:
154 """Calculate the final weighted score."""
155 return (
156 0.40 * similarity_scores['token_set'] +
157 0.25 * similarity_scores['ratio'] +
158 0.15 * similarity_scores['partial'] +
159 0.15 * (word_scores['jaccard'] * 100) +
160 0.05 * word_scores['prefix']
161 )
162
163 @staticmethod
164 def _apply_adjustments(score: float, query: str, candidate: str, word_ratio: float) -> float:
165 """Apply length penalty and word ratio boost."""
166 import math
167
168 # Apply length penalty
169 score *= min(1.0, 0.9 + (0.1 * SearchUtils.length_penalty(query, candidate)))
170
171 # Apply word ratio boost
172 score = score * (0.7 + (0.3 * word_ratio))
173
174 return score
175
176class AudioSearcher:
177 """Handles searching for audio files in configured paths."""
178
179 @staticmethod
180 def _is_valid_audio_file(file_path: Path) -> bool:
181 """Check if file is a valid audio file."""
182 valid_extensions = {'.mp4', '.mp3', '.wav', '.opus', '.m4a'}
183 return file_path.is_file() and file_path.suffix.lower() in valid_extensions
184
185 @staticmethod
186 def _calculate_match_score(query: str, file_name: str, autocomplete: bool) -> tuple[int, str]:
187 """Calculate match score for a file."""
188 if not query and autocomplete:
189 return 0, file_name
190
191 if query and len(query) >= 2:
192 if query in file_name:
193 logger.debug(f"Found direct match: {file_name} (contains '{query}')")
194
195 score = SearchUtils.combined_score(query, file_name)
196 logger.debug(f"File: {file_name} - Score: {score}")
197
198 if score >= 30:
199 return score, file_name
200
201 return -1, file_name
202
203 @staticmethod
204 async def _search_single_path(path: Path, query: str, max_results: int, autocomplete: bool) -> list[tuple[int, str]]:
205 """Search for audio files in a single path."""
206 matches = []
207
208 logger.debug(f"Searching in directory: {path}")
209
210 for file_path in path.rglob('*'):
211 if not AudioSearcher._is_valid_audio_file(file_path):
212 continue
213
214 file_name = file_path.name.lower()
215 score, _ = AudioSearcher._calculate_match_score(query, file_name, autocomplete)
216
217 if score >= 0:
218 matches.append((score, str(file_path.absolute())))
219
220 if len(matches) >= max_results * 5:
221 break
222
223 return matches
224
225 @classmethod
226 async def search_local(cls, query: str = None, max_results: int = 25, autocomplete: bool = False):
227 """Search for audio files with fuzzy matching and optional autocomplete support.
228
229 Args:
230 query: Search query string (can be None for autocomplete)
231 max_results: Maximum number of results to return
232 autocomplete: If True, returns results even with empty query
233
234 Returns:
235 List of matching file paths with scores
236 """
237 if not query and not autocomplete:
238 return []
239
240 query = (query or '').lower().strip()
241 logger.info(f"Starting local search for query: '{query}'")
242
243 try:
244 async with aiofiles.open('settings.json') as f:
245 content = await f.read()
246 data = json.loads(content)
247 search_paths = data.get('audio_paths', [])
248 logger.info(f"Searching in paths: {search_paths}")
249
250 matches = []
251 file_count = 0
252
253 for path_str in search_paths:
254 path = Path(path_str)
255 if not path.exists():
256 logger.warning(f"Path does not exist: {path}")
257 continue
258
259 path_matches = await cls._search_single_path(path, query, max_results, autocomplete)
260 matches.extend(path_matches)
261 file_count += len(path_matches)
262
263 if len(matches) >= max_results * 5:
264 break
265
266 logger.info(f"Scanned {file_count} files total")
267
268 matches.sort(reverse=True, key=lambda x: x[0])
269 result = [path for score, path in matches[:max_results]]
270
271 logger.info(f"Search complete. Found {len(result)} matches out of {file_count} files.")
272 if not result and query:
273 logger.warning(f"No matches found for query: '{query}'. Try a different search term.")
274 return result
275
276 except Exception as e:
277 logger.error(f"Error during local search: {str(e)}", exc_info=True)
278 return []
279
280 @classmethod
281 async def search_local_autocomplete(cls, query: str = None, max_results: int = 10):
282 """Wrapper for autocomplete searches that returns more results and allows empty query."""
283 # For autocomplete, we want to show more options and allow empty queries
284 if not query:
285 return []
286
287 return await cls.search_local(query, max_results=max_results, autocomplete=True)