Tag search

tools/__init__.py

@@ -25,6 +25,7 @@ from .get_text_context import handle as get_text_context_handler
 from .wikipedia import handle as wikipedia_handler
 from .get_model_info import handle as get_model_info_handler
 from .generate_image import handle as generate_image_handler
+from .search_tags import handle as search_tags_handler
 
 # Central registry of all available tools
 TOOL_REGISTRY = [
@@ -154,4 +155,16 @@ TOOL_REGISTRY = [
         },
         handler=generate_image_handler
     ),
+    Tool(
+        name="search_tags",
+        description="Searches the Danbooru tag database for tags matching a query. Returns the most popular tags including alias matches. Useful for finding the correct booru-style tags for anime models.",
+        schema={
+            "type": "object",
+            "properties": {
+                "query": {"type": "string", "description": "The tag or partial tag to search for (e.g., 'white hair' or 'lesbian')."}
+            },
+            "required": ["query"],
+        },
+        handler=search_tags_handler
+    ),
 ]

tools/search_tags.py

@@ -0,0 +1,68 @@
+import csv
+import os
+from typing import List, Dict, Any
+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.")
+
+    tag_file_path = config.TAG_DATABASE_PATH
+    if not os.path.exists(tag_file_path):
+        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)
+            for row in reader:
+                if not row or len(row) < 3:
+                    continue
+                
+                name = row[0].lower()
+                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)
+                    })
+
+    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)
+
+    # Format top 20 results
+    results = []
+    for m in matches[:20]:
+        count_fmt = format_count(str(m["count"]))
+        type_sfx = get_type_suffix(m["type"])
+        
+        if m["matched_via"] == "alias":
+            line = f"{m['alias_match']} → {m['name']} ({count_fmt}){type_sfx}"
+        else:
+            line = f"{m['name']} ({count_fmt}){type_sfx}"
+        
+        results.append(line)
+
+    if not results:
+        return {"text": f"No tags found matching '{query}'."}
+
+    return {"text": "Top matches:\n" + "\n".join(results)}

tools/utils.py

@@ -117,3 +117,26 @@ def parse_indices(indices_str: str) -> List[int]:
             except ValueError:
                 raise ToolError(f"Invalid index format: {part}")
     return indices
+
+def format_count(count_str: str) -> str:
+    """Formats large numbers into a human-readable string (e.g., 1.2M, 218k)."""
+    try:
+        count = int(count_str)
+        if count >= 1_000_000:
+            return f"{count / 1_000_000:.1f}M".replace(".0", "")
+        if count >= 1_000:
+            return f"{count / 1_000:.0f}k"
+        return str(count)
+    except (ValueError, TypeError):
+        return "0"
+
+def get_type_suffix(type_val: str) -> str:
+    """Returns the human-readable suffix for the tag type."""
+    mapping = {
+        "0": "", # General
+        "1": " [Artist]",
+        "3": " [Copyright]",
+        "4": " [Character]",
+        "5": " [Meta]"
+    }
+    return mapping.get(str(type_val), f" [Type {type_val}]")