From c634305604d0c993343a34262a3f3811d4c22e2c Mon Sep 17 00:00:00 2001 From: salmonik Date: Sun, 4 Jan 2026 13:31:49 +0100 Subject: live search with good algorithm --- cogs/music/play.py | 272 ++++++++++++++++++++++++++++++----------------- cogs/music/play_utils.py | 168 +++++++++++++++++++++++------ 2 files changed, 307 insertions(+), 133 deletions(-) (limited to 'cogs/music') diff --git a/cogs/music/play.py b/cogs/music/play.py index 8390ff5..c9abe4b 100644 --- a/cogs/music/play.py +++ b/cogs/music/play.py @@ -1,10 +1,11 @@ import os - import discord import mutagen +from typing import List, Tuple, Optional, Callable, Awaitable, Union from . import play_utils from discord.ext import commands +from discord.commands import Option, OptionChoice import logging logger = logging.getLogger(__name__) @@ -13,113 +14,184 @@ class Play(commands.Cog): def __init__(self, bot): self.bot = bot self.temp_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "temp") + self.path_mapping = {} # Maps display names to full paths - @discord.slash_command(name="grajcuj", description="grajcujmy wszyscy!") - async def play( - self, - ctx: discord.ApplicationContext, - query: str = discord.Option(str, "co zagrać??", required=True), - source: str = discord.Option(str, "co zagrać??", required=True, choices=["local", "soundcloud", "youtube"]) - - ): - await ctx.defer(ephemeral=True) - if source == "local": - candidates = await play_utils.PlayUtils(self.bot).search_local(query) - if candidates: - embed = discord.Embed(title="Wyniki wyszukiwania", description="Wybierz plik do odtworzenia:", color=discord.Color.blue()) - view = discord.ui.View(disable_on_timeout=True) - for i, candidate in enumerate(candidates, start=1): - filename = os.path.basename(candidate) - metadata = None - if os.path.splitext(filename)[1].lower() in ['.mp3', '.wav', '.opus']: - metadata = mutagen.File(candidate) + async def _create_search_embed(self, candidates: List[str]) -> Tuple[discord.Embed, discord.ui.View]: + """Create an embed and view for search results.""" + embed = discord.Embed( + title="Wyniki wyszukiwania", + description="Wybierz plik do odtworzenia:", + color=discord.Color.blue() + ) + + view = discord.ui.View(disable_on_timeout=True) + + for i, candidate in enumerate(candidates, start=1): + filename = os.path.basename(candidate) + title = None + + # Try to get metadata if it's an audio file + if os.path.splitext(filename)[1].lower() in ['.mp3', '.wav', '.opus']: + try: + metadata = mutagen.File(candidate) + if metadata: title = metadata.get('title', None) - if metadata and title: - embed.add_field(name=f"{i}. {title}", value="\u200b", inline=False) - else: - embed.add_field(name=f"{i}. {filename}", value="\u200b", inline=False) - button = discord.ui.Button(label=f"Play {i}") - async def button_callback(interaction: discord.Interaction, candidate=candidate): - await ctx.interaction.delete_original_response() - await play_utils.PlayUtils(self.bot).play(ctx, candidate) - button.callback = button_callback - view.add_item(button) - await ctx.respond(embed=embed, view=view, ephemeral=True) - else: - await ctx.respond("Nie znaleziono pliku do odtworzenia", ephemeral=True) + except Exception as e: + logger.debug(f"Could not read metadata for {filename}: {str(e)}") + + # Add field with title if available, otherwise use filename + display_name = title if title else filename + embed.add_field(name=f"{i}. {display_name}", value="\u200b", inline=False) + + return embed, view + + async def _handle_search_results( + self, + ctx: discord.ApplicationContext, + candidates: List[str], + play_callback: Callable[[discord.Interaction, str], Awaitable[None]] + ) -> None: + """Handle the display of search results and button interactions.""" + if not candidates: + await ctx.respond("Nie znaleziono pliku do odtworzenia", ephemeral=True) + return + + embed, view = await self._create_search_embed(candidates) + + # Add buttons for each candidate + for i, candidate in enumerate(candidates, start=1): + button = discord.ui.Button(label=f"Play {i}") + + async def button_callback(interaction: discord.Interaction, candidate_path=candidate): + await interaction.response.defer(ephemeral=True) + await ctx.interaction.delete_original_response() + try: + await play_callback(interaction, candidate_path) + except Exception as e: + logger.error(f"Error in play callback: {str(e)}", exc_info=True) + await interaction.followup.send(f"Wystąpił błąd podczas odtwarzania: {str(e)}", ephemeral=True) + + button.callback = button_callback + view.add_item(button) + + await ctx.respond(embed=embed, view=view, ephemeral=True) - @discord.slash_command(name="grajcuj_zkompresowane", description="grajcujmy wszyscy w wysokiej jakosci!") - async def play_compressed( + async def _search_autocomplete(self, ctx: discord.AutocompleteContext) -> List[OptionChoice]: + """Autocomplete handler for search queries.""" + query = ctx.value.lower() + try: + results = await play_utils.PlayUtils.search_local_autocomplete(query, max_results=25) + choices = [] + self.path_mapping = {} # Reset the mapping for each new query + + for path in results: + filename = os.path.splitext(os.path.basename(path))[0] # Remove extension + display = filename # Default display name + + # Try to get metadata for display + try: + if os.path.splitext(path)[1].lower() in ['.mp3', '.wav', '.opus', '.m4a']: + metadata = mutagen.File(path, easy=True) + + # Get metadata fields, falling back to filename if not available + title = metadata.get('title', [filename])[0] + artist = metadata.get('artist', [''])[0] + album_artist = metadata.get('albumartist', [artist])[0] + + # Build display name based on available metadata + if album_artist and title: + display = f"{album_artist} - {title}" + elif artist and title: + display = f"{artist} - {title}" + elif title: + display = title + + except Exception as e: + logger.debug(f"Could not read metadata for {filename}: {str(e)}") + + # Clean up the display name (remove any extra whitespace or brackets) + display = display.strip() + + # Store the mapping (use the cleaned display name as the key) + self.path_mapping[display] = path + + # Truncate display name if too long + if len(display) > 90: # Leave room for potential numbering + display = display[:87] + '...' + + choices.append(OptionChoice(display[:100], display[:100])) + + # Limit to 25 choices (Discord's limit) + if len(choices) >= 25: + break + + return choices + + except Exception as e: + logger.error(f"Error in autocomplete: {str(e)}", exc_info=True) + return [] + + @discord.slash_command(name="grajcuj", description="grajcujmy wszyscy!") + async def play( self, ctx: discord.ApplicationContext, - query: str = discord.Option(str, "co zagrać??", required=True), - source: str = discord.Option(str, "skąd zagrać??", required=True, - choices=["local", "soundcloud", "youtube"]) + file: str = discord.Option( + str, + "Wybierz plik do odtworzenia", + autocomplete=_search_autocomplete, + required=True + ), + source: str = discord.Option( + str, + "Źródło dźwięku", + choices=["local", "soundcloud", "youtube"], + default="local" + ), + compressed: bool = discord.Option( + bool, + "Czy skompresować dźwięk?", + default=False + ) ): - await ctx.defer(ephemeral=True) - if source == "local": - candidates = await play_utils.PlayUtils(self.bot).search_local(query) - if candidates: - embed = discord.Embed(title="Wyniki wyszukiwania", description="Wybierz plik do odtworzenia:", - color=discord.Color.blue()) - view = discord.ui.View(disable_on_timeout=True) - for i, candidate in enumerate(candidates, start=1): - filename = os.path.basename(candidate) - metadata = None - if os.path.splitext(filename)[1].lower() in ['.mp3', '.wav', '.opus']: - metadata = mutagen.File(candidate) - title = metadata.get('title', None) - if metadata and title: - embed.add_field(name=f"{i}. {title}", value="\u200b", inline=False) - else: - embed.add_field(name=f"{i}. {filename}", value="\u200b", inline=False) - button = discord.ui.Button(label=f"Play {i}") - - async def button_callback(interaction: discord.Interaction, candidate=candidate): - await interaction.response.defer(ephemeral=True) + """Odtwarza wybrany plik dźwiękowy z opcjonalną kompresją.""" + await ctx.defer() + + if source != "local": + await ctx.respond("Tylko źródło 'local' jest obecnie obsługiwane w tej komendzie.", ephemeral=True) + return + + # Look up the full path from the display name + full_path = self.path_mapping.get(file, file) + + if not os.path.exists(full_path): + logger.warning(f"File not found: {full_path} (display name: {file})") + await ctx.respond("Wybrany plik nie istnieje.", ephemeral=True) + return + + try: + if compressed: + # Compress the file before playing + compressed_path = await play_utils.PlayUtils(self.bot).compress(ctx, full_path) + if not compressed_path: + await ctx.respond("Nie udało się skompresować pliku.", ephemeral=True) + return + + try: + await play_utils.PlayUtils(self.bot).play(ctx, compressed_path, delete_after_play=True) + except Exception as e: + # Clean up the compressed file if play fails + if os.path.exists(compressed_path): try: - await ctx.interaction.delete_original_response() - # Compress the file before playing - compressed_path = await play_utils.PlayUtils(self.bot).compress(ctx, candidate) - if compressed_path: - try: - await play_utils.PlayUtils(self.bot).play(ctx, compressed_path, delete_after_play=True) - except Exception as e: - # Clean up the compressed file if play fails - if os.path.exists(compressed_path): - try: - os.remove(compressed_path) - logger.info(f"Cleaned up temporary file after play error: {compressed_path}") - except Exception as cleanup_error: - logger.error(f"Error cleaning up file {compressed_path}: {str(cleanup_error)}") - raise e - except Exception as e: - logger.error(f"Error in play_compressed: {str(e)}", exc_info=True) - await ctx.respond(f"Error playing compressed file: {str(e)}", ephemeral=True) - - button.callback = button_callback - view.add_item(button) - await ctx.respond(embed=embed, view=view, ephemeral=True) + os.remove(compressed_path) + logger.info(f"Cleaned up temporary file after play error: {compressed_path}") + except Exception as cleanup_error: + logger.error(f"Error cleaning up file {compressed_path}: {str(cleanup_error)}") + raise e else: - await ctx.respond("Nie znaleziono pliku do odtworzenia", ephemeral=True) - else: - # For YouTube/SoundCloud - try: - url, title = await play_utils.PlayUtils(self.bot).get_url(ctx, query) - if not url: - return - - filename = await play_utils.PlayUtils(self.bot).download(ctx, url) - if not filename: - return - - # Compress the downloaded file - compressed_filename = await play_utils.PlayUtils(self.bot).compress(ctx, filename) - if compressed_filename: - await play_utils.PlayUtils(self.bot).play(ctx, compressed_filename, title) - except Exception as e: - logger.error(f"Error in play_compressed: {str(e)}") - await ctx.respond(f"Error: {str(e)}", ephemeral=True) + await play_utils.PlayUtils(self.bot).play(ctx, full_path) + except Exception as e: + logger.error(f"Error playing file: {str(e)}", exc_info=True) + await ctx.respond(f"Wystąpił błąd podczas odtwarzania: {str(e)}", ephemeral=True) @discord.slash_command(name="grajcuj_plik", description="grajcuj plik") async def play_file(self, ctx: discord.ApplicationContext, attachment: discord.Attachment = discord.Option(discord.Attachment, "co zagrać??", required=True)): 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 import math import os import asyncio +import re from pathlib import Path import discord @@ -33,27 +34,90 @@ class Utils: @staticmethod def combined_score(query, candidate): - """Calculate a combined fuzzy matching score using multiple strategies.""" - # Use several scorers and combine with weights - r_ratio = fuzz.ratio(query, candidate) # overall edit similarity - r_partial = fuzz.partial_ratio(query, candidate) # substring allowance - r_token_set = fuzz.token_set_ratio(query, candidate) # word overlap ignoring extras + """Calculate a combined fuzzy matching score using multiple strategies. - # Calculate prefix score using partial_ratio of the start of the strings - min_len = min(len(query), len(candidate)) - if min_len > 0: - r_prefix = fuzz.ratio(query, candidate[:min_len]) # prefix matching - else: - r_prefix = 0 - - # Weighted sum (tune weights as needed) - score = (0.35 * r_token_set + # Increased weight for token set - 0.30 * r_partial + # Slightly increased weight for partial matches - 0.25 * r_ratio + # Standard ratio - 0.10 * r_prefix) # Reduced weight for prefix matches + 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() + + # First check for direct substring match (fast path) + if query in candidate: + # Base score of 100 for exact match, with bonus for being at word boundary + base_score = 100 + # Check if match is at word boundary + if re.search(rf'\b{re.escape(query)}\b', candidate): + base_score = 110 # Slight bonus for word boundary match + return min(100, base_score) + + # For very short queries (2-3 chars), be more lenient + if len(query) <= 3: + # Check for partial matches at word boundaries + if re.search(rf'\b{re.escape(query)}', candidate): + return 90 # High score for short query matching word start + if query in candidate: + return 80 # Still good for any substring match with short query + + # For longer queries, check word presence + query_words = set(word for word in re.split(r'\W+', query) if len(word) > 1) # Reduced min word length to 1 + if query_words: + candidate_words = set(re.split(r'\W+', candidate)) + # Check if all query words are in candidate (with some leniency) + matching_words = sum(1 for word in query_words if word in candidate_words) + word_ratio = matching_words / len(query_words) + + # If no words match at all, return 0 + if matching_words == 0: + return 0.0 - score *= Utils.length_penalty(query, candidate) # apply length penalty - return min(100, score) # Ensure score doesn't exceed 100 + # Calculate different similarity metrics + r_ratio = fuzz.ratio(query, candidate) # Overall edit similarity + r_partial = fuzz.partial_ratio(query, candidate) # Best substring match + r_token_set = fuzz.token_set_ratio(query, candidate) # Word overlap + + # 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 similarity for word overlap + jaccard_sim = word_intersection / word_union if word_union > 0 else 0 + + # Calculate prefix score (bonus for matching start of words) + 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 # Increased weight + + # Weighted sum with adjusted weights + score = ( + 0.40 * r_token_set + # Increased weight for token set + 0.25 * r_ratio + # Overall similarity + 0.15 * r_partial + # Substring matches + 0.15 * (jaccard_sim * 100) + # Word overlap ratio + 0.05 * min(prefix_bonus, 100) # Prefix bonus (capped at 100) + ) + + # Apply length penalty (less severe) + score *= min(1.0, 0.9 + (0.1 * Utils.length_penalty(query, candidate))) + + # Boost score based on word ratio + score = score * (0.7 + (0.3 * word_ratio)) + + # Ensure score is within bounds + return min(100, max(0, score)) + + return 0.0 # No match found class PlayUtils: def __init__(self, bot): @@ -235,18 +299,24 @@ class PlayUtils: # -- LOCAL PLAY -- - @staticmethod - async def search_local(query, max_results=5): - """Search for audio files that contain the query in their name (case-insensitive). + @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 + 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 + 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: with open('settings.json') as f: data = json.load(f) @@ -254,7 +324,7 @@ class PlayUtils: logger.info(f"Searching in paths: {search_paths}") matches = [] - q = query.lower() + file_count = 0 for path in search_paths: path = Path(path) @@ -273,19 +343,51 @@ class PlayUtils: if file_path.suffix.lower() not in {'.mp4', '.mp3', '.wav', '.opus', '.m4a'}: continue - # Check if query is in filename (case-insensitive) - if q in file_path.name.lower(): - matches.append(str(file_path.absolute())) - if len(matches) >= max_results: - logger.info(f"Found {len(matches)} matches, stopping search") - return matches + file_count += 1 + file_name = file_path.name.lower() + + # For autocomplete, return all files if query is empty + if not query and autocomplete: + matches.append((0, str(file_path.absolute()))) + # For regular search, calculate fuzzy match score + elif query and len(query) >= 2: + # First do a simple substring check for debugging + if query in file_name: + logger.debug(f"Found direct match: {file_name} (contains '{query}')") + + score = Utils.combined_score(query, file_name) + logger.debug(f"File: {file_name} - Score: {score}") + + if score >= 30: # Minimum score threshold + matches.append((score, str(file_path.absolute()))) + + # Early return if we have enough high-scoring matches + if len(matches) >= max_results * 5: # Get more candidates for sorting + break + + if len(matches) >= max_results * 5: # Limit memory usage + break + + logger.info(f"Scanned {file_count} files total") - logger.info(f"Search complete. Found {len(matches)} matches.") - return matches + # Sort by score (descending) and take top results + 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 + return await cls.search_local(query, max_results=max_results, autocomplete=True) -- cgit v1.2.3