diff options
Diffstat (limited to 'cogs/music/search_utils.py')
| -rw-r--r-- | cogs/music/search_utils.py | 287 |
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 | |||
| 6 | import json | ||
| 7 | import logging | ||
| 8 | import re | ||
| 9 | from pathlib import Path | ||
| 10 | import aiofiles | ||
| 11 | from rapidfuzz import fuzz | ||
| 12 | |||
| 13 | logger = logging.getLogger(__name__) | ||
| 14 | |||
| 15 | class 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 | |||
| 176 | class 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) | ||
