diff --git a/tools/search_tags.py b/tools/search_tags.py index 5c0b70a..c94f535 100644 --- a/tools/search_tags.py +++ b/tools/search_tags.py @@ -1,23 +1,22 @@ import csv import os -from typing import List, Dict, Any +import config +import difflib +from typing import List, Dict, Any, Optional from tools.utils import ToolError, format_count, get_type_suffix -# Tag database path is now managed in config.py -async def handle(args: Dict[str, Any]) -> Dict[str, Any]: - """ - Searches the Danbooru tag database for tags matching a query. - Returns the most popular tags including alias matches. - """ - query = args.get("query", "").lower() - if not query: - raise ToolError("The 'query' argument is required.") +# Global cache to avoid reading the CSV from disk on every request +_TAG_CACHE: Optional[List[Dict[str, Any]]] = None +def _load_tags() -> List[Dict[str, Any]]: + """Reads the Danbooru tag CSV and caches it in memory.""" tag_file_path = config.TAG_DATABASE_PATH + tags = [] + if not os.path.exists(tag_file_path): + # We raise ToolError here because it's a fatal configuration issue raise ToolError(f"Tag database not found at {tag_file_path}. Please ensure the tag-autocomplete extension is installed.") - matches = [] try: with open(tag_file_path, mode='r', encoding='utf-8') as f: reader = csv.reader(f) @@ -25,26 +24,52 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]: if not row or len(row) < 3: continue - name = row[0].lower() + name = row[0] tag_type = row[1] count = int(row[2]) if row[2].isdigit() else 0 aliases = row[3].lower().split(',') if len(row) > 3 else [] - # Check for match in name or aliases - is_direct = query in name - is_alias = any(query in alias.strip() for alias in aliases) - - if is_direct or is_alias: - matches.append({ - "name": row[0], - "type": tag_type, - "count": count, - "matched_via": name if is_direct else "alias", - "alias_match": "" if is_direct else next((a.strip() for a in aliases if query in a.strip()), query) - }) - + tags.append({ + "name": name, + "name_lower": name.lower(), + "type": tag_type, + "count": count, + "aliases": aliases + }) except Exception as e: raise ToolError(f"Error reading tag database: {str(e)}") + + return tags + +async def handle(args: Dict[str, Any]) -> Dict[str, Any]: + """ + Searches the Danbooru tag database for tags matching a query. + Returns the most popular tags including alias matches. + """ + global _TAG_CACHE + + query = args.get("query", "").lower() + if not query: + raise ToolError("The 'query' argument is required.") + + # Lazy-load the tags into memory + if _TAG_CACHE is None: + _TAG_CACHE = _load_tags() + + matches = [] + for tag in _TAG_CACHE: + # Check for match in name or aliases + is_direct = query in tag["name_lower"] + is_alias = any(query in alias.strip() for alias in tag["aliases"]) + + if is_direct or is_alias: + matches.append({ + "name": tag["name"], + "type": tag["type"], + "count": tag["count"], + "matched_via": "name" if is_direct else "alias", + "alias_match": "" if is_direct else next((a.strip() for a in tag["aliases"] if query in a.strip()), query) + }) # Sort by count descending matches.sort(key=lambda x: x["count"], reverse=True) @@ -63,6 +88,23 @@ async def handle(args: Dict[str, Any]) -> Dict[str, Any]: results.append(line) if not results: - return {"text": f"No tags found matching '{query}'."} + # Attempt to find similar tags using difflib + all_names_lower = [t["name_lower"] for t in _TAG_CACHE] + suggestions_lower = difflib.get_close_matches(query, all_names_lower, n=10, cutoff=0.5) + + if not suggestions_lower: + return {"text": f"No tags found matching '{query}'."} + + # Map lowercased suggestions back to original tag objects + suggestions = [] + for s_lower in suggestions_lower: + # Find the first tag that matches this lowercased name + tag = next((t for t in _TAG_CACHE if t["name_lower"] == s_lower), None) + if tag: + count_fmt = format_count(str(tag["count"])) + type_sfx = get_type_suffix(tag["type"]) + suggestions.append(f"{tag['name']} ({count_fmt}){type_sfx}") + + return {"text": f"No exact matches for '{query}'. Did you mean:\n" + "\n".join(suggestions)} return {"text": "Top matches:\n" + "\n".join(results)}