Tag not found update

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
+# Global cache to avoid reading the CSV from disk on every request
-async def handle(args: Dict[str, Any]) -> Dict[str, Any]:
+_TAG_CACHE: Optional[List[Dict[str, Any]]] = None
-    """
-    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.")
 
+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,27 +24,53 @@ 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 []
 
+                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 name
+        is_direct = query in tag["name_lower"]
-                is_alias = any(query in alias.strip() for alias in aliases)
+        is_alias = any(query in alias.strip() for alias in tag["aliases"])
 
         if is_direct or is_alias:
             matches.append({
-                        "name": row[0],
+                "name": tag["name"],
-                        "type": tag_type,
+                "type": tag["type"],
-                        "count": count,
+                "count": tag["count"],
-                        "matched_via": name if is_direct else "alias",
+                "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)
+                "alias_match": "" if is_direct else next((a.strip() for a in tag["aliases"] if query in a.strip()), query)
             })
 
-    except Exception as e:
-        raise ToolError(f"Error reading tag database: {str(e)}")
-
     # 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:
+        # 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)}