Tag search

README.md

@@ -43,6 +43,11 @@ Provides the "manual" for available generation models.
 *   **Best for**: Learning the prompting style, recommended settings, and available resolution presets for a specific model.
 *   **Workflow**: Call without arguments to see the catalog; call with `model_name` for the detailed guide.
 
+### 🏷️ `search_tags`
+Searches the Danbooru tag database for recognized tags and aliases.
+*   **Best for**: Finding the correct booru-style tags, checking tag popularity, and resolving aliases (e.g., 'lesbian' → 'yuri').
+*   **Workflow**: Provide a query string to get a list of the most popular matching tags.
+
 ### 🌐 `browse_wikipedia`
 Allows the model to browse Wikipedia using its API.
 *   **Best for**: Quickly retrieving summaries, structural maps (ToC), or specific section content from Wikipedia without dumping the entire page.

config.py

@@ -16,6 +16,7 @@ USER_AGENT = "Mozilla/5.0 (X11; Linux x86_64; rv:151.0) Gecko/20100101 Firefox/1
 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"
 
 # --- Model Specific Token Tuning (Tuned for Gemma 4) ---
 # Patch size is typically (clip.vision.patch_size * n_merge)

model_presets.toml

@@ -6,7 +6,7 @@
 ["noobaiXLNAIXL_vPred10Version"]
 description = "Anime-style model, very stylistic and not aesthetic-tuned. Can do any NSFW."
 guide = """
-This model is not aesthetic tuned, it must be given explicit tags for everything. Omitted parts of the prompt will not default to something 'good'.
+This model is not aesthetic tuned, it must be given explicit tags for everything. Omitted parts of the prompt will not default to something 'good'. This model is very fickle, you will almost always need to iterate on the prompt or resubmit to roll the best picture.
 Has very excellent understanding of characters and artists down to extremely niche. Unless prompting an original character, their name is enough to decribe their appearance completely except for clothing.
 Accepts a list of comma-separated booru-style tags. Use spaces, not underscores, for tags.
 Prompts MUST follow this format: <1girl/1boy/1other/solo/couple/(can use multiple)>, <character(s)>, <series>, <artist>, <tags>

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/generate_image.py

@@ -141,7 +141,7 @@ async def handle(args: Dict[str, Any]) -> List[Dict[str, Any]]:
         return [
             {
                 "type": "text", 
-                "text": f"Image successfully generated using model '{model_name}'.\n\nTo show the image to the user, you can include this markdown link in your response: ![Generated Image]({proxy_url})"
+                "text": f"Image successfully generated using model '{model_name}'.\n\nLocal path: {raw_path}\nTo show the image to the user, you can include this markdown link in your response: ![Generated Image]({proxy_url})"
             },
             {
                 "type": "image",

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}]")

tools/wikipedia.py

@@ -50,7 +50,7 @@ async def _search(title: str, limit: int = 5, fallback: bool = False) -> str:
         return f"No Wikipedia results found for '{title}'."
     
     if fallback:
-        header = f"Article not found. Please select one of the following similar articles to proceed with '{title}':"
+        header = f"Article not found. You MUST select one of the following articles and submit another query:"
     else:
         header = f"Search results for '{title}':"
         
@@ -62,7 +62,7 @@ async def _search(title: str, limit: int = 5, fallback: bool = False) -> str:
         lines.append(f"{i}. [{title_res}] - Snippet: {snippet}")
     
     if fallback:
-        lines.append("\n[Tip: Use the title in brackets [ ] to explore the selected page.]")
+        lines.append("\nUse the title in brackets [ ] to explore the selected page.")
         
     return "\n".join(lines)