moosecrap/MooseCP
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Tag search fixes, WIP on Danbooru wiki

Browse Source
This commit is contained in:
moosecrap
2026-07-28 16:07:18 -07:00
parent b107aff150
commit 3e523e9725
3 changed files with 116 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files
config.py
+5 -2
View File
@@ -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
+1 -1
View File
@@ -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
+108 -38
View File
@@ -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 by count descending
matches.sort(key=lambda x: x["count"], reverse=True)
# Sort substring matches by count descending
substring_matches.sort(key=lambda x: x["tag"]["count"], reverse=True)
# 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"])
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
# 2. Find Similarity Matches (Low Priority)
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 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
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)
# 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)}]