From 3e523e9725afb4c1a5b329744eee250081da2d4c Mon Sep 17 00:00:00 2001 From: moosecrap Date: Tue, 28 Jul 2026 16:07:18 -0700 Subject: [PATCH] Tag search fixes, WIP on Danbooru wiki --- config.py | 7 +- tools/__init__.py | 2 +- tools/search_tags.py | 150 +++++++++++++++++++++++++++++++------------ 3 files changed, 116 insertions(+), 43 deletions(-) diff --git a/config.py b/config.py index 0f70080..d0d02ac 100644 --- a/config.py +++ b/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) diff --git a/tools/__init__.py b/tools/__init__.py index a68a560..5b07e9d 100644 --- a/tools/__init__.py +++ b/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": { diff --git a/tools/search_tags.py b/tools/search_tags.py index 2d177b4..96d75fe 100644 --- a/tools/search_tags.py +++ b/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({ - "name": tag["name"], - "type": tag["type"], - "count": tag["count"], + # Find which alias actually matched for reporting + matched_alias = "" + if not is_direct: + 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 substring matches by count descending + substring_matches.sort(key=lambda x: x["tag"]["count"], reverse=True) - # Sort by count descending - matches.sort(key=lambda x: x["count"], reverse=True) + # 2. Find Similarity Matches (Low Priority) + all_names_lower = [t["name_lower"] for t in _TAG_CACHE] + similar_names_lower = difflib.get_close_matches(query, all_names_lower, n=config.TAG_SEARCH_LIMIT, cutoff=0.5) + similar_matches = [] + for s_lower in similar_names_lower: + tag = next((t for t in _TAG_CACHE if t["name_lower"] == s_lower), None) + if tag: + similar_matches.append(tag) - # Format top results - results = [] - for m in matches[:config.TAG_SEARCH_LIMIT]: - count_fmt = format_count(str(m["count"])) - type_sfx = get_type_suffix(m["type"]) + # 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}") - if m["matched_via"] == "alias": - line = f"{m['alias_match']} → {m['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: - line = f"{m['name']} ({count_fmt}){type_sfx}" + 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"]) - results.append(line) + # 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 - 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) + 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}" - if not suggestions_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: - 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)}] + output_lines.append(line) - return [{"type": "text", "text": "Top matches:\n" + "\n".join(results)}] + return [{"type": "text", "text": "\n".join(output_lines)}]