From ce61bf6bd5949046d15f58609f572fedcb8a073e Mon Sep 17 00:00:00 2001 From: salmonik Date: Wed, 4 Feb 2026 16:16:26 +0100 Subject: refactor music --- cogs/music/search_utils.py | 286 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 286 insertions(+) create mode 100644 cogs/music/search_utils.py (limited to 'cogs/music/search_utils.py') diff --git a/cogs/music/search_utils.py b/cogs/music/search_utils.py new file mode 100644 index 0000000..50d229c --- /dev/null +++ b/cogs/music/search_utils.py @@ -0,0 +1,286 @@ +# This file is part of Senator Sperma +# Licensed under the GPL-3.0 License +# Copyright (C) 2026 Salmonik @ salmonik@tuta.io +# https://codeberg.org/salmonik/senator-sperma + +import json +import logging +import re +from pathlib import Path +import aiofiles +from rapidfuzz import fuzz + +logger = logging.getLogger(__name__) + +class SearchUtils: + """Utilities for searching and scoring audio files.""" + + @staticmethod + 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. + + Args: + query: The search query + candidate: The candidate string to score against + + Returns: + float: Score between 0-100, where higher means better match + """ + if not query or not candidate: + return 0.0 + + query = query.lower().strip() + candidate = candidate.lower() + + # Check for direct matches first + direct_match_score = SearchUtils._check_direct_matches(query, candidate) + if direct_match_score is not None: + return direct_match_score + + # Handle short queries specially + short_query_score = SearchUtils._handle_short_queries(query, candidate) + if short_query_score is not None: + return short_query_score + + # For longer queries, use complex scoring + return SearchUtils._calculate_complex_score(query, candidate) + + @staticmethod + def _check_direct_matches(query: str, candidate: str) -> float: + """Check for direct substring matches.""" + if query in candidate: + base_score = 100 + if re.search(rf'\b{re.escape(query)}\b', candidate): + base_score = 110 + return min(100, base_score) + return 0 + + @staticmethod + def _handle_short_queries(query: str, candidate: str) -> float: + """Handle scoring for short queries (2-3 chars).""" + if len(query) <= 3: + if re.search(rf'\b{re.escape(query)}', candidate): + return 90 + if query in candidate: + return 80 + return 0 + + @staticmethod + def _calculate_complex_score(query: str, candidate: str) -> float: + """Calculate complex score for longer queries.""" + query_words = {word for word in re.split(r'\W+', query) if len(word) > 1} + if not query_words: + return 0.0 + + candidate_words = set(re.split(r'\W+', candidate)) + word_ratio = SearchUtils._calculate_word_ratio(query_words, candidate_words) + + if word_ratio == 0: + return 0.0 + + # Calculate various similarity metrics + similarity_scores = SearchUtils._get_similarity_scores(query, candidate) + word_scores = SearchUtils._get_word_scores(query, candidate, query_words, candidate_words) + + # Calculate final weighted score + final_score = SearchUtils._calculate_weighted_score(similarity_scores, word_scores, word_ratio) + + # Apply penalties and boosts + final_score = SearchUtils._apply_adjustments(final_score, query, candidate, word_ratio) + + return min(100, max(0, int(final_score))) + + @staticmethod + def _calculate_word_ratio(query_words: set, candidate_words: set) -> float: + """Calculate the ratio of matching words.""" + matching_words = sum(1 for word in query_words if word in candidate_words) + return matching_words / len(query_words) + + @staticmethod + def _get_similarity_scores(query: str, candidate: str) -> dict: + """Get various fuzzy similarity scores.""" + return { + 'ratio': fuzz.ratio(query, candidate), + 'partial': fuzz.partial_ratio(query, candidate), + 'token_set': fuzz.token_set_ratio(query, candidate) + } + + @staticmethod + def _get_word_scores(query: str, candidate: str, query_words: set, candidate_words: set) -> dict: + """Calculate word-based scores.""" + query_tokens = set(query.split()) + candidate_tokens = set(candidate.split()) + + word_intersection = len(query_tokens & candidate_tokens) + word_union = len(query_tokens | candidate_tokens) + jaccard_sim = word_intersection / word_union if word_union > 0 else 0 + + prefix_bonus = SearchUtils._calculate_prefix_bonus(query_words, candidate_words) + + return { + 'jaccard': jaccard_sim, + 'prefix': min(prefix_bonus, 100) + } + + @staticmethod + def _calculate_prefix_bonus(query_words: set, candidate_words: set) -> float: + """Calculate prefix matching bonus.""" + prefix_bonus = 0 + for q_word in query_words: + for c_word in candidate_words: + if c_word.startswith(q_word): + prefix_bonus += (len(q_word) / max(len(c_word), 1)) * 20 + return prefix_bonus + + @staticmethod + def _calculate_weighted_score(similarity_scores: dict, word_scores: dict, word_ratio: float) -> float: + """Calculate the final weighted score.""" + return ( + 0.40 * similarity_scores['token_set'] + + 0.25 * similarity_scores['ratio'] + + 0.15 * similarity_scores['partial'] + + 0.15 * (word_scores['jaccard'] * 100) + + 0.05 * word_scores['prefix'] + ) + + @staticmethod + def _apply_adjustments(score: float, query: str, candidate: str, word_ratio: float) -> float: + """Apply length penalty and word ratio boost.""" + import math + + # Apply length penalty + score *= min(1.0, 0.9 + (0.1 * SearchUtils.length_penalty(query, candidate))) + + # Apply word ratio boost + score = score * (0.7 + (0.3 * word_ratio)) + + return score + +class AudioSearcher: + """Handles searching for audio files in configured paths.""" + + @staticmethod + def _is_valid_audio_file(file_path: Path) -> bool: + """Check if file is a valid audio file.""" + valid_extensions = {'.mp4', '.mp3', '.wav', '.opus', '.m4a'} + return file_path.is_file() and file_path.suffix.lower() in valid_extensions + + @staticmethod + def _calculate_match_score(query: str, file_name: str, autocomplete: bool) -> tuple[int, str]: + """Calculate match score for a file.""" + if not query and autocomplete: + return 0, file_name + + if query and len(query) >= 2: + if query in file_name: + logger.debug(f"Found direct match: {file_name} (contains '{query}')") + + score = SearchUtils.combined_score(query, file_name) + logger.debug(f"File: {file_name} - Score: {score}") + + if score >= 30: + return score, file_name + + return -1, file_name + + @staticmethod + async def _search_single_path(path: Path, query: str, max_results: int, autocomplete: bool) -> list[tuple[int, str]]: + """Search for audio files in a single path.""" + matches = [] + + logger.debug(f"Searching in directory: {path}") + + for file_path in path.rglob('*'): + if not AudioSearcher._is_valid_audio_file(file_path): + continue + + file_name = file_path.name.lower() + score, _ = AudioSearcher._calculate_match_score(query, file_name, autocomplete) + + if score >= 0: + matches.append((score, str(file_path.absolute()))) + + if len(matches) >= max_results * 5: + break + + return matches + + @classmethod + async def search_local(cls, query: str = None, max_results: int = 25, autocomplete: bool = False): + """Search for audio files with fuzzy matching and optional autocomplete support. + + Args: + query: Search query string (can be None for autocomplete) + max_results: Maximum number of results to return + autocomplete: If True, returns results even with empty query + + Returns: + List of matching file paths with scores + """ + if not query and not autocomplete: + return [] + + query = (query or '').lower().strip() + logger.info(f"Starting local search for query: '{query}'") + + try: + async with aiofiles.open('settings.json') as f: + data = json.load(f) + search_paths = data.get('audio_paths', []) + logger.info(f"Searching in paths: {search_paths}") + + matches = [] + file_count = 0 + + for path_str in search_paths: + path = Path(path_str) + if not path.exists(): + logger.warning(f"Path does not exist: {path}") + continue + + path_matches = await cls._search_single_path(path, query, max_results, autocomplete) + matches.extend(path_matches) + file_count += len(path_matches) + + if len(matches) >= max_results * 5: + break + + logger.info(f"Scanned {file_count} files total") + + matches.sort(reverse=True, key=lambda x: x[0]) + result = [path for score, path in matches[:max_results]] + + logger.info(f"Search complete. Found {len(result)} matches out of {file_count} files.") + if not result and query: + logger.warning(f"No matches found for query: '{query}'. Try a different search term.") + return result + + except Exception as e: + logger.error(f"Error during local search: {str(e)}", exc_info=True) + return [] + + @classmethod + async def search_local_autocomplete(cls, query: str = None, max_results: int = 10): + """Wrapper for autocomplete searches that returns more results and allows empty query.""" + # For autocomplete, we want to show more options and allow empty queries + if not query: + return [] + + return await cls.search_local(query, max_results=max_results, autocomplete=True) -- cgit v1.2.3