""" mcp_server.py ───────────── MCP 服务层 挂载方式（在 ui_nicegui.py 中）： from mcp_server import mcp app.mount('/mcp', mcp.streamable_http_app()) 接入地址： https://sakizuki-danboorusearch.hf.space/mcp/mcp 支持的工具： search_tags 自然语言搜索标签 get_related_tags 基于共现表查关联推荐 """ import json import asyncio import logging from anyio import BrokenResourceError, ClosedResourceError from mcp.server.fastmcp import FastMCP from mcp.server.transport_security import TransportSecuritySettings from core.engine import DanbooruTagger from core.models import SearchRequest import core.counter as counter import re # ── 过滤客户端断连产生的无害报错噪音 ────────────────────────────────── class _SuppressClientDisconnect(logging.Filter): def filter(self, record: logging.LogRecord) -> bool: exc = record.exc_info[1] if record.exc_info else None if isinstance(exc, (BrokenResourceError, ClosedResourceError)): return False # 丢弃该日志记录 return True _disconnect_filter = _SuppressClientDisconnect() logging.getLogger("mcp.server.streamable_http").addFilter(_disconnect_filter) logging.getLogger("uvicorn.error").addFilter(_disconnect_filter) mcp = FastMCP( name="danbooru-searcher", transport_security=TransportSecuritySettings(enable_dns_rebinding_protection=False), ) @mcp.tool() async def search_tags( query: str, search_mode: str = "full_scene", category: str = "all", show_nsfw: bool = True, include_wiki: bool = False, ) -> str: """ Search Danbooru tags using natural language and return a ready-to-use prompt. Only supported for general, copyright, and character tag searches; **artists and meta tags are not supported.** ## Args - query: Natural language description (Chinese recommended). - search_mode: Preset strategy. Pick the one that matches your intent. "full_scene" — Full scene → prompt (e.g. "一个穿着白色水手服的少女在雨中奔跑") "concept_explore" — Vague concept exploration, broad recall (e.g. "赛博朋克服装", "兔耳朵", "中国风汉服") "subject_describe" — Describe a subject to find matching tags (e.g. "EVA中蓝发的驾驶员", "两侧有开口，前方有拉绳的运动短裤") "precise_lookup" — Precise lookup / spell fix (e.g. "selafuku", "thighhigh") - category: Filter to a specific tag category. Default "all". "all" — All (通用 + 版权 + 人物) "general" — Visual attributes, clothing, pose, background, etc. "character" — Named characters from any series "copyright" — Specific anime/game/franchise titles - show_nsfw: Include NSFW tags. Default True. - include_wiki: Append wiki description to each result. Default False. Set True when tags are unfamiliar and need disambiguation. ## Query writing guide Use **spaces, newlines, Chinese commas (，), or Chinese dunhao (、)** to manually separate concepts. Each delimiter-bounded segment ≤7 characters stays atomic — the engine respects your intent. | Query style | Example | |---|---| | Concept list (spaces) | `运动社团 校队 比赛 运动会` | | Concept list (dun hao) | `反乌托邦、赛博朋克、蒸汽朋克` | | Natural sentence | `一个穿着白色水手服的少女在雨中奔跑` | | Mixed | `运动社团 一个穿水手服的少女` | ## Workflow After search_tags, pass selected tags to get_related_tags to discover complementary tags via co-occurrence. Chain freely: search_tags → get_related_tags → get_related_tags → search_tags for multi-hop exploration. ## Returns JSON with: prompt (comma-separated tags), keywords, results. Each result: tag, cn_name, category, final_score, count[, wiki if include_wiki=True]. """ _SEARCH_MODE_PRESETS: dict[str, dict] = { "precise_lookup": {"top_k": 10, "limit": 10, "popularity_weight": 0.15, "use_segmentation": False, "group_mode": "off", "max_per_group": 2}, "concept_explore": {"top_k": 80, "limit": 80, "popularity_weight": 0.15, "use_segmentation": True, "group_mode": "expand", "max_per_group": 2}, "subject_describe": {"top_k": 20, "limit": 20, "popularity_weight": 0.15, "use_segmentation": False, "group_mode": "off", "max_per_group": 2}, "full_scene": {"top_k": 5, "limit": 80, "popularity_weight": 0.15, "use_segmentation": True, "group_mode": "diverse", "max_per_group": 2}, } preset = _SEARCH_MODE_PRESETS.get(search_mode, _SEARCH_MODE_PRESETS["full_scene"]) _CATEGORY_MAP: dict[str, list[str]] = { "all": ["General", "Character", "Copyright", "Artist", "Meta"], "general": ["General"], "character": ["Character"], "copyright": ["Copyright"], } target_categories = _CATEGORY_MAP.get( category, _CATEGORY_MAP["all"], ) tagger = await DanbooruTagger.get_instance() request = SearchRequest( query=query, top_k=preset["top_k"], limit=preset["limit"], popularity_weight=preset["popularity_weight"], show_nsfw=show_nsfw, use_segmentation=preset["use_segmentation"], target_categories=target_categories, group_mode=preset["group_mode"], max_per_group=preset["max_per_group"], ) response = await tagger.search_async(request) # 计数：每次 MCP 搜索调用均计入搜索、成功、复制；访问不变 await counter.increment() await counter.increment_success() await counter.increment_copy() await counter.increment_mcp() results = [] for r in response.results: if r.nsfw == '1' and not show_nsfw: continue item = { "tag": r.tag, "cn_name": r.cn_name, "category": r.category, "final_score": r.final_score, "count": r.count, } if include_wiki: item["wiki"] = r.wiki results.append(item) payload = { "prompt": response.tags_sfw if not show_nsfw else response.tags_all, "keywords": response.keywords, "results": results, } han_chars = re.findall(r'[\u4e00-\u9fff]', query) if len(query) > 0 and len(han_chars) / len(query) < 0.5: payload["hint"] = ( "检测到英文查询，该搜索引擎对中文查询优化更好，如果搜索结果不合预期，推荐用中文重试" ) return json.dumps(payload, ensure_ascii=False, indent=2) @mcp.tool() async def get_related_tags( tags: list[str], limit: int = 50, show_nsfw: bool = True, include_wiki: bool = False, ) -> str: """ Return co-occurrence-based tag recommendations for a given tag list (NPMI scoring). Only supported for general, copyright, and character tag searches; **artists and meta tags are not supported.** This tool surfaces tags that frequently appear alongside the seeds in Danbooru, mixing categories (General / Character / Copyright) by design. ## Typical use cases - Attribute → characters who have it e.g. ["fingerless_gloves"] → tifa_lockhart, cammy_white, bridget_(guilty_gear), ... - Work → characters in it e.g. ["overlord_(maruyama)"] → shalltear_bloodfallen, ainz_ooal_gown, albedo_(overlord), ... - Character → their visual attributes e.g. ["amiya_(arknights)"] → outfits, expressions, accessories - Theme exploration e.g. ["fighter_jet"] → aircraft types, actions, backgrounds - Multi-tag intersection e.g. ["maid", "twintails"] → tags specific to the combination, scored by summed NPMI For within-category exploration (e.g. "more clothing tags like X"), use search_tags with the `category` parameter instead. ## Workflow Chain freely: search_tags → get_related_tags → get_related_tags → search_tags. Each hop along the co-occurrence graph reveals tags unreachable by semantic search alone. ## Args - tags: List of canonical Danbooru tag names (underscores, no spaces). e.g. ["white_serafuku", "sailor_collar"] - limit: Max recommendations returned. Default 50. - show_nsfw: Include NSFW tags. Default True. - include_wiki: Append wiki description to each result. Default False. Set True when result tags are unfamiliar and need disambiguation. ## Returns JSON array sorted by aggregated NPMI score (descending). Each result: - tag, cn_name, category, count (post_count), cooc_score (normalized to [0,1]) - sources: seed tags that contributed to this score - wiki: only if include_wiki=True """ tagger = await DanbooruTagger.get_instance() # ── 检查标签是否存在，不存在则尝试 search_tags 纠错 ────────────────── valid_tags = [] invalid_tags = [] for t in tags: if t in tagger._name_to_idx: valid_tags.append(t) else: invalid_tags.append(t) corrections = {} if invalid_tags: for bad_tag in invalid_tags: try: req = SearchRequest( query=bad_tag, top_k=5, limit=5, popularity_weight=0.15, use_segmentation=False, target_layers=['英文'] ) resp = await tagger.search_async(req) if resp.results: corrections[bad_tag] = resp.results[0].tag except Exception: pass if not valid_tags and not corrections: return json.dumps({ "error": "所有传入的标签均不存在于标签表中", "invalid_tags": invalid_tags, }, ensure_ascii=False, indent=2) # 用纠错后的标签替换无效标签 corrected_tags = [] for t in tags: if t in valid_tags: corrected_tags.append(t) elif t in corrections: corrected_tags.append(corrections[t]) results = await asyncio.to_thread( tagger.get_related, corrected_tags, set(corrected_tags), limit, show_nsfw, ) # 计数：每次 MCP related 调用均计入搜索、成功、复制；访问不变 await counter.increment() await counter.increment_success() await counter.increment_copy() await counter.increment_mcp() output = [] for r in results: item = { "tag": r.tag, "cn_name": r.cn_name, "category": r.category, "count": r.post_count, "cooc_score": r.cooc_score, "sources": r.sources, } if include_wiki: item["wiki"] = r.wiki output.append(item) payload = {"results": output} if corrections: correction_notes = [ f"{bad} → {good}" for bad, good in corrections.items() ] payload = { "correction_note": "标签拼写错误，已经纠错: " + ", ".join(correction_notes), "corrections": corrections, "results": output, } return json.dumps(payload, ensure_ascii=False, indent=2)