Tag search fixes, WIP on Danbooru wiki

config.py

@@ -10,7 +10,7 @@ ROOT_DIR = Path(__file__).parent.resolve()
 HOST = "127.0.0.1"
 PORT = 8000
 request_host = ContextVar("request_host", default=f"{HOST}:{PORT}")
-LOG_LEVEL = "WARNING" # Options: "DEBUG", "INFO", "WARNING", "ERROR"
+LOG_LEVEL = "DEBUG" # Options: "DEBUG", "INFO", "WARNING", "ERROR"
 LOG_FILE = str(ROOT_DIR / "debug.log")
 USER_AGENT = "MooseCP/1.0 Local MCP Server (https://long-cat.net/)"
 
@@ -18,8 +18,11 @@ USER_AGENT = "MooseCP/1.0 Local MCP Server (https://long-cat.net/)"
 SD_URL = "http://127.0.0.1:7860"
 MODEL_PRESETS_PATH = str(ROOT_DIR / "model_presets.toml")
 RES_PRESETS_PATH = str(ROOT_DIR / "resolution_presets.toml")
-TAG_DATABASE_PATH = "/home/matt/stable-diffusion-webui/extensions/a1111-sd-webui-tagcomplete/tags/danbooru.csv"
+TAG_DATABASE_PATH = str(ROOT_DIR / "danbooru.csv")
 TAG_SEARCH_LIMIT = 20
+ENABLE_TAG_WIKI = False
+DANBOORU_LOGIN = ""  # Your Danbooru username (Optional)
+DANBOORU_API_KEY = "" # Your Danbooru API key (Optional)
 
 # --- Model Specific Token Tuning (Tuned for Gemma 4) ---
 # Patch size is typically (clip.vision.patch_size * n_merge)

tools/__init__.py

@@ -157,7 +157,7 @@ TOOL_REGISTRY = [
     ),
     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.",
+        description="Searches the Danbooru tag database for tags matching a query. Returns the Danbooru wiki page on an exact match. On any query, returns a list of similar tags as well as aliases. Useful for finding the correct booru-style tags for anime models or getting more information.",
         schema={
             "type": "object",
             "properties": {

tools/search_tags.py

@@ -2,6 +2,8 @@ import csv
 import os
 import config
 import difflib
+import httpx
+import re
 from typing import List, Dict, Any, Optional
 from tools.utils import ToolError, format_count, get_type_suffix
 
@@ -41,14 +43,47 @@ def _load_tags() -> List[Dict[str, Any]]:
     
     return tags
 
+    return tags
+
+async def _fetch_wiki_info(tag_name: str) -> Optional[str]:
+    """Fetches the wiki description for a tag from Danbooru."""
+    if not config.ENABLE_TAG_WIKI:
+        return None
+    
+    url = f"https://danbooru.donmai.us/wiki_pages/{tag_name}.json"
+    try:
+        # Use Basic Auth if credentials are provided
+        auth = None
+        if config.DANBOORU_LOGIN and config.DANBOORU_API_KEY:
+            auth = (config.DANBOORU_LOGIN, config.DANBOORU_API_KEY)
+
+        async with httpx.AsyncClient(timeout=5.0) as client:
+            # Using a browser-like UA to avoid potential blocks
+            headers = {"User-Agent": config.USER_AGENT}
+            resp = await client.get(url, headers=headers, auth=auth)
+            
+            if resp.status_code == 403:
+                return " (Wiki access blocked by Danbooru/Cloudflare)"
+            
+            if resp.status_code == 200:
+                data = resp.json()
+                wiki_body = data.get("wiki_page", {}).get("body")
+                if wiki_body:
+                    # Strip HTML tags for the LLM
+                    return re.sub(r'<[^>]*>', '', wiki_body).strip()
+    except Exception as e:
+        return f" (Error fetching wiki: {str(e)})"
+    return None
+
 async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     """
     Searches the Danbooru tag database for tags matching a query. 
-    Returns the most popular tags including alias matches.
+    Returns a unified list prioritized by substring matches then similarity.
     """
     global _TAG_CACHE
     
-    query = args.get("query", "").lower()
+    # Normalize query: treat spaces and underscores as identical
+    query = args.get("query", "").lower().replace(" ", "_")
     if not query:
         raise ToolError("The 'query' argument is required.")
 
@@ -56,55 +91,90 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
     if _TAG_CACHE is None:
         _TAG_CACHE = _load_tags()
 
-    matches = []
+    # 1. Find Substring/Alias Matches (High Priority)
+    substring_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({
+            # Find which alias actually matched for reporting
-                "name": tag["name"],
+            matched_alias = ""
-                "type": tag["type"],
+            if not is_direct:
-                "count": tag["count"],
+                matched_alias = next((a.strip() for a in tag["aliases"] if query in a.strip()), query)
+            
+            substring_matches.append({
+                "tag": tag,
                 "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)
+                "alias_match": matched_alias
             })
     
-    # Sort by count descending
+    # Sort substring matches by count descending
-    matches.sort(key=lambda x: x["count"], reverse=True)
+    substring_matches.sort(key=lambda x: x["tag"]["count"], reverse=True)
 
-    # Format top results
+    # 2. Find Similarity Matches (Low Priority)
-    results = []
-    for m in matches[:config.TAG_SEARCH_LIMIT]:
-        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:
-        # 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=config.TAG_SEARCH_LIMIT, cutoff=0.5)
+    similar_names_lower = difflib.get_close_matches(query, all_names_lower, n=config.TAG_SEARCH_LIMIT, cutoff=0.5)
-        
+    similar_matches = []
-        if not suggestions_lower:
+    for s_lower in similar_names_lower:
-            return [{"type": "text", "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:
+            similar_matches.append(tag)
+
+    # Build Unified List
+    final_results = []
+    
+    # Add substring matches first
+    for m in substring_matches:
+        final_results.append(m)
+        if len(final_results) >= config.TAG_SEARCH_LIMIT:
+            break
+            
+    # Fill remaining slots with similar matches
+    if len(final_results) < config.TAG_SEARCH_LIMIT:
+        mentioned_names = {m["tag"]["name_lower"] for m in final_results}
+        for tag in similar_matches:
+            if tag["name_lower"] not in mentioned_names:
+                final_results.append({"tag": tag, "matched_via": "similarity", "alias_match": ""})
+                if len(final_results) >= config.TAG_SEARCH_LIMIT:
+                    break
+
+    if not final_results:
+        return [{"type": "text", "text": f"No tags found matching '{query}'."}]
+
+    # Format output lines
+    output_lines = []
+    
+    # Special Case: Exact Match Wiki Header
+    # Check if the very first result is an exact match
+    first_res = final_results[0]
+    if first_res["tag"]["name_lower"] == query:
+        exact_tag = first_res["tag"]
+        count_fmt = format_count(str(exact_tag["count"]))
+        type_sfx = get_type_suffix(exact_tag["type"])
+        output_lines.append(f"Exact Match: {exact_tag['name']} ({count_fmt}){type_sfx}")
+        
+        wiki_info = await _fetch_wiki_info(exact_tag["name"])
+        if wiki_info:
+            output_lines.append(f"Wiki: {wiki_info}\n")
+        else:
+            output_lines.append("") # spacer
+            
+    # List the tags
+    for res in final_results:
+        tag = res["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 [{"type": "text", "text": f"No exact matches for '{query}'. Did you mean:\n" + "\n".join(suggestions)}]
+        # If this is the exact match we already listed in the header, skip it
+        if first_res["tag"]["name_lower"] == query and tag["name_lower"] == query:
+            continue
 
-    return [{"type": "text", "text": "Top matches:\n" + "\n".join(results)}]
+        if res["matched_via"] == "alias":
+            line = f"- {res['alias_match']} → {tag['name']} ({count_fmt}){type_sfx}"
+        else:
+            line = f"- {tag['name']} ({count_fmt}){type_sfx}"
+        
+        output_lines.append(line)
+
+    return [{"type": "text", "text": "\n".join(output_lines)}]