aboutsummaryrefslogtreecommitdiffstats
path: root/cogs/utils
diff options
context:
space:
mode:
Diffstat (limited to 'cogs/utils')
-rw-r--r--cogs/utils/file_search.py500
1 files changed, 500 insertions, 0 deletions
diff --git a/cogs/utils/file_search.py b/cogs/utils/file_search.py
new file mode 100644
index 0000000..4a2bdc4
--- /dev/null
+++ b/cogs/utils/file_search.py
@@ -0,0 +1,500 @@
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"""
7Universal File Searcher
8
9This module provides a universal file searcher that can search for different file types:
10- Audio files (mp3, wav, mp4, opus, m4a, flac, aac, ogg)
11- Video files (mp4, avi, mkv, mov, wmv, flv, webm, m4v, 3gp)
12- Image files (jpg, jpeg, png, gif, bmp, tiff, webp, svg, ico)
13
14Usage Examples:
15 # Search all file types
16 results = await UniversalSearcher.search_local("query")
17
18 # Search only audio files
19 results = await UniversalSearcher.search_local("query", filetype="Audio")
20
21 # Search only video files
22 results = await UniversalSearcher.search_local("query", filetype="Video")
23
24 # Search only image files
25 results = await UniversalSearcher.search_local("query", filetype="Image")
26
27 # Autocomplete search
28 results = await UniversalSearcher.search_local_autocomplete("partial_query", filetype="Video")
29
30Settings.json structure:
31{
32 "file_paths": [
33 "/path/to/your/media/files",
34 "/another/path/with/mixed/content",
35 "relative/path/to/files"
36 ]
37}
38
39Note: file_paths should contain directories that may contain any mix of audio, video,
40and image files. The searcher will filter by file type automatically based on the
41filetype parameter.
42"""
43
44import json
45import logging
46import re
47import urllib.parse
48from pathlib import Path
49import aiofiles
50from rapidfuzz import fuzz
51import aiohttp
52import asyncio
53import os
54from dotenv import load_dotenv
55
56load_dotenv()
57
58logger = logging.getLogger(__name__)
59
60class SearchUtils:
61 """Utilities for searching and scoring audio files."""
62
63 @staticmethod
64 def normalize_name(p: Path, keep_ext=False):
65 """Preprocess filenames: lower, replace separators, strip extensions optionally."""
66 name = str(p.name if keep_ext else p.stem)
67 return name.lower().replace('_', ' ').replace('-', ' ')
68
69 @staticmethod
70 def length_penalty(query, candidate):
71 """Penalize tiny filenames when query is long; returns multiplier [0.0,1.0]."""
72 ql = len(query)
73 cl = len(candidate)
74 if ql == 0:
75 return 1.0
76 ratio = cl / ql
77 # Prefer candidate lengths near query length; gentle falloff otherwise
78 return math.exp(-0.5 * ((math.log(max(ratio, 1e-6)))**2))
79
80 @staticmethod
81 def combined_score(query, candidate):
82 """Calculate a combined fuzzy matching score using multiple strategies.
83
84 Args:
85 query: The search query
86 candidate: The candidate string to score against
87
88 Returns:
89 float: Score between 0-100, where higher means better match
90 """
91 if not query or not candidate:
92 return 0.0
93
94 query = query.lower().strip()
95 candidate = candidate.lower()
96
97 # Check for direct matches first
98 direct_match_score = SearchUtils._check_direct_matches(query, candidate)
99 if direct_match_score is not None:
100 return direct_match_score
101
102 # Handle short queries specially
103 short_query_score = SearchUtils._handle_short_queries(query, candidate)
104 if short_query_score is not None:
105 return short_query_score
106
107 # For longer queries, use complex scoring
108 return SearchUtils._calculate_complex_score(query, candidate)
109
110 @staticmethod
111 def _check_direct_matches(query: str, candidate: str) -> float:
112 """Check for direct substring matches."""
113 if query in candidate:
114 base_score = 100
115 if re.search(rf'\b{re.escape(query)}\b', candidate):
116 base_score = 110
117 return min(100, base_score)
118 return 0
119
120 @staticmethod
121 def _handle_short_queries(query: str, candidate: str) -> float:
122 """Handle scoring for short queries (2-3 chars)."""
123 if len(query) <= 3:
124 if re.search(rf'\b{re.escape(query)}', candidate):
125 return 90
126 if query in candidate:
127 return 80
128 return 0
129
130 @staticmethod
131 def _calculate_complex_score(query: str, candidate: str) -> float:
132 """Calculate complex score for longer queries."""
133 query_words = {word for word in re.split(r'\W+', query) if len(word) > 1}
134 if not query_words:
135 return 0.0
136
137 candidate_words = set(re.split(r'\W+', candidate))
138 word_ratio = SearchUtils._calculate_word_ratio(query_words, candidate_words)
139
140 if word_ratio == 0:
141 return 0.0
142
143 # Calculate various similarity metrics
144 similarity_scores = SearchUtils._get_similarity_scores(query, candidate)
145 word_scores = SearchUtils._get_word_scores(query, candidate, query_words, candidate_words)
146
147 # Calculate final weighted score
148 final_score = SearchUtils._calculate_weighted_score(similarity_scores, word_scores, word_ratio)
149
150 # Apply penalties and boosts
151 final_score = SearchUtils._apply_adjustments(final_score, query, candidate, word_ratio)
152
153 return min(100, max(0, int(final_score)))
154
155 @staticmethod
156 def _calculate_word_ratio(query_words: set, candidate_words: set) -> float:
157 """Calculate the ratio of matching words."""
158 matching_words = sum(1 for word in query_words if word in candidate_words)
159 return matching_words / len(query_words)
160
161 @staticmethod
162 def _get_similarity_scores(query: str, candidate: str) -> dict:
163 """Get various fuzzy similarity scores."""
164 return {
165 'ratio': fuzz.ratio(query, candidate),
166 'partial': fuzz.partial_ratio(query, candidate),
167 'token_set': fuzz.token_set_ratio(query, candidate)
168 }
169
170 @staticmethod
171 def _get_word_scores(query: str, candidate: str, query_words: set, candidate_words: set) -> dict:
172 """Calculate word-based scores."""
173 query_tokens = set(query.split())
174 candidate_tokens = set(candidate.split())
175
176 word_intersection = len(query_tokens & candidate_tokens)
177 word_union = len(query_tokens | candidate_tokens)
178 jaccard_sim = word_intersection / word_union if word_union > 0 else 0
179
180 prefix_bonus = SearchUtils._calculate_prefix_bonus(query_words, candidate_words)
181
182 return {
183 'jaccard': jaccard_sim,
184 'prefix': min(prefix_bonus, 100)
185 }
186
187 @staticmethod
188 def _calculate_prefix_bonus(query_words: set, candidate_words: set) -> float:
189 """Calculate prefix matching bonus."""
190 prefix_bonus = 0
191 for q_word in query_words:
192 for c_word in candidate_words:
193 if c_word.startswith(q_word):
194 prefix_bonus += (len(q_word) / max(len(c_word), 1)) * 20
195 return prefix_bonus
196
197 @staticmethod
198 def _calculate_weighted_score(similarity_scores: dict, word_scores: dict, word_ratio: float) -> float:
199 """Calculate the final weighted score."""
200 return (
201 0.40 * similarity_scores['token_set'] +
202 0.25 * similarity_scores['ratio'] +
203 0.15 * similarity_scores['partial'] +
204 0.15 * (word_scores['jaccard'] * 100) +
205 0.05 * word_scores['prefix']
206 )
207
208 @staticmethod
209 def _apply_adjustments(score: float, query: str, candidate: str, word_ratio: float) -> float:
210 """Apply length penalty and word ratio boost."""
211 import math
212
213 # Apply length penalty
214 score *= min(1.0, 0.9 + (0.1 * SearchUtils.length_penalty(query, candidate)))
215
216 # Apply word ratio boost
217 score = score * (0.7 + (0.3 * word_ratio))
218
219 return score
220
221class UniversalSearcher:
222 """Handles searching for files in configured paths with support for multiple file types."""
223
224 FILE_TYPE_EXTENSIONS = {
225 'Audio': {'.mp4', '.mp3', '.wav', '.opus', '.m4a', '.flac', '.aac', '.ogg'},
226 'Video': {'.mp4', '.avi', '.mkv', '.mov', '.wmv', '.flv', '.webm', '.m4v', '.3gp', '.gif'},
227 'Image': {'.jpg', '.jpeg', '.png', '.gif', '.bmp', '.tiff', '.webp', '.svg', '.ico'}
228 }
229
230 @staticmethod
231 def _is_valid_file(file_path: Path, filetype: str = None) -> bool:
232 """Check if file is valid for the specified filetype."""
233 if not file_path.is_file():
234 return False
235
236 if filetype is None:
237 # Search all file types
238 all_extensions = set()
239 for extensions in UniversalSearcher.FILE_TYPE_EXTENSIONS.values():
240 all_extensions.update(extensions)
241 return file_path.suffix.lower() in all_extensions
242
243 filetype = filetype.capitalize()
244 if filetype not in UniversalSearcher.FILE_TYPE_EXTENSIONS:
245 return False
246
247 valid_extensions = UniversalSearcher.FILE_TYPE_EXTENSIONS[filetype]
248 return file_path.suffix.lower() in valid_extensions
249
250 @staticmethod
251 def _calculate_match_score(query: str, file_name: str, autocomplete: bool) -> tuple[int, str]:
252 """Calculate match score for a file."""
253 if not query and autocomplete:
254 return 0, file_name
255
256 if query and len(query) >= 2:
257 if query in file_name:
258 logger.debug(f"Found direct match: {file_name} (contains '{query}')")
259
260 score = SearchUtils.combined_score(query, file_name)
261 logger.debug(f"File: {file_name} - Score: {score}")
262
263 if score >= 30:
264 return score, file_name
265
266 return -1, file_name
267
268 @staticmethod
269 async def _search_single_path(path: Path, query: str, max_results: int, autocomplete: bool, filetype: str = None) -> list[tuple[int, str]]:
270 """Search for files in a single path."""
271 matches = []
272
273 logger.debug(f"Searching in directory: {path} (filetype: {filetype or 'All'})")
274
275 for file_path in path.rglob('*'):
276 if not UniversalSearcher._is_valid_file(file_path, filetype):
277 continue
278
279 file_name = file_path.name.lower()
280 score, _ = UniversalSearcher._calculate_match_score(query, file_name, autocomplete)
281
282 if score >= 0:
283 matches.append((score, str(file_path.absolute())))
284
285 if len(matches) >= max_results * 5:
286 break
287
288 return matches
289
290 @staticmethod
291 def _get_search_paths(data: dict, filetype: str = None) -> list:
292 """Get the search paths from settings."""
293 return data.get('file_paths', [])
294
295 @classmethod
296 async def _search_paths(cls, search_paths: list, query: str, max_results: int, autocomplete: bool, filetype: str = None) -> tuple[list, int]:
297 """Search through the provided paths and return matches and file count."""
298 matches = []
299 file_count = 0
300
301 for path_str in search_paths:
302 path = Path(path_str)
303 if not path.exists():
304 logger.warning(f"Path does not exist: {path}")
305 continue
306
307 path_matches = await cls._search_single_path(path, query, max_results, autocomplete, filetype)
308 matches.extend(path_matches)
309 file_count += len(path_matches)
310
311 if len(matches) >= max_results * 5:
312 break
313
314 return matches, file_count
315
316 @classmethod
317 async def search_local(cls, query: str = None, max_results: int = 25, autocomplete: bool = False, filetype: str = None):
318 """Search for files with fuzzy matching and optional autocomplete support.
319
320 Args:
321 query: Search query string (can be None for autocomplete)
322 max_results: Maximum number of results to return
323 autocomplete: If True, returns results even with empty query
324 filetype: Type of files to search for ('Audio', 'Video', 'Image', or None for all)
325
326 Returns:
327 List of matching file paths with scores
328 """
329 if not query and not autocomplete:
330 return []
331
332 query = (query or '').lower().strip()
333 logger.info(f"Starting local search for query: '{query}' (filetype: {filetype or 'All'})")
334
335 try:
336 async with aiofiles.open('settings.json') as f:
337 content = await f.read()
338 data = json.loads(content)
339
340 search_paths = cls._get_search_paths(data, filetype)
341 logger.info(f"Searching in paths: {search_paths}")
342
343 matches, file_count = await cls._search_paths(search_paths, query, max_results, autocomplete, filetype)
344
345 logger.info(f"Scanned {file_count} files total")
346
347 matches.sort(reverse=True, key=lambda x: x[0])
348 result = [path for score, path in matches[:max_results]]
349
350 logger.info(f"Search complete. Found {len(result)} matches out of {file_count} files.")
351 if not result and query:
352 logger.warning(f"No matches found for query: '{query}'. Try a different search term.")
353 return result
354
355 except Exception as e:
356 logger.error(f"Error during local search: {str(e)}", exc_info=True)
357 return []
358
359 @classmethod
360 async def search_local_autocomplete(cls, query: str = None, max_results: int = 10, filetype: str = None):
361 """Wrapper for autocomplete searches that returns more results and allows empty query.
362
363 Args:
364 query: Search query string (can be None for autocomplete)
365 max_results: Maximum number of results to return
366 filetype: Type of files to search for ('Audio', 'Video', 'Image', or None for all)
367 """
368 # For autocomplete, we want to show more options and allow empty queries
369 if not query:
370 return []
371
372 return await cls.search_local(query, max_results=max_results, autocomplete=True, filetype=filetype)
373
374 @staticmethod
375 async def _load_baseurl():
376 """Load baseurl from settings.json"""
377 try:
378 async with aiofiles.open('settings.json') as f:
379 content = await f.read()
380 data = json.loads(content)
381 return data.get('baseurl', 'http://localhost:8080').rstrip('/')
382 except Exception as e:
383 logger.error(f"Error loading baseurl: {e}")
384 return 'http://localhost:8080'
385
386 @staticmethod
387 async def search_copyparty(query: str, max_results: int = 25):
388 """Search for files using copyparty API (jPOST method).
389
390 Args:
391 query: Search query string
392 max_results: Maximum number of results to return
393
394 Returns:
395 List of matching file URLs
396 """
397
398
399 if not query:
400 return []
401
402 query = query.strip()
403 logger.info(f"Starting copyparty search for query: '{query}'")
404 query = f'name like "*{query}*"'
405
406 try:
407 baseurl = await UniversalSearcher._load_baseurl()
408 password = os.getenv("COPYPARTY_R_PASSWORD")
409
410 async with aiohttp.ClientSession() as session:
411 headers = {"Content-Type": "application/json"}
412 if password:
413 headers["PW"] = password
414 async with session.post(f"{baseurl}/?srch", headers=headers, data=json.dumps({"n": max_results, "q": query})) as response:
415 if response.status == 200:
416 data = await response.json()
417 logger.info(f"Copyparty response: {data}")
418 files = UniversalSearcher._extract_files_from_response(data)
419 results = UniversalSearcher._convert_files_to_urls(files, baseurl, max_results)
420
421 logger.info(f"Copyparty search complete. Found {len(results)} matches.")
422 return results
423 else:
424 logger.error(f"Copyparty API returned status {response.status}")
425 return []
426
427 except aiohttp.ClientError as e:
428 logger.error(f"HTTP error during copyparty search: {e}")
429 return []
430 except Exception as e:
431 logger.error(f"Error during copyparty search: {str(e)}", exc_info=True)
432 return []
433
434 @staticmethod
435 def _extract_files_from_response(data):
436 """Extract files list from copyparty API response."""
437 if isinstance(data, dict):
438 if 'hits' in data:
439 return data['hits']
440 elif 'files' in data:
441 return data['files']
442 elif 'results' in data:
443 return data['results']
444 else:
445 return data if isinstance(data, list) else []
446 elif isinstance(data, list):
447 return data
448 else:
449 return []
450
451 @staticmethod
452 def _convert_files_to_urls(files, baseurl, max_results):
453 """Convert file info objects to URLs."""
454 results = []
455
456 for file_info in files[:max_results]:
457 if isinstance(file_info, dict):
458 file_url = UniversalSearcher._extract_url_from_dict(file_info, baseurl)
459 if file_url:
460 results.append(file_url)
461 elif isinstance(file_info, str):
462 # If it's just a string path/URL
463 if file_info.startswith('http'):
464 results.append(file_info)
465 else:
466 results.append(f"{baseurl}/{file_info.lstrip('/')}")
467
468 return results
469
470 @staticmethod
471 def _extract_url_from_dict(file_info, baseurl):
472 """Extract URL from file info dictionary."""
473 if 'url' in file_info:
474 return file_info['url']
475 elif 'rp' in file_info:
476 return f"{baseurl}/{file_info['rp']}"
477 elif 'path' in file_info:
478 return f"{baseurl}/{file_info['path'].lstrip('/')}"
479 elif 'name' in file_info:
480 return f"{baseurl}/{file_info['name'].lstrip('/')}"
481 else:
482 return None
483
484 @staticmethod
485 async def search_copyparty_autocomplete(query: str, max_results: int = 10):
486 """Wrapper for autocomplete searches using copyparty.
487
488 Args:
489 query: Search query string
490 max_results: Maximum number of results to return
491
492 Returns:
493 List of matching file URLs
494 """
495 if not query:
496 return []
497
498 return await UniversalSearcher.search_copyparty(query, max_results)
499
500