mcp_server.py

"""
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,
    use_segmentation: bool = True,
    top_k: int = 5,
    limit: int = 80,
    popularity_weight: float = 0.15,
    show_nsfw: bool = True,
    include_wiki: bool = False,
    category: str = "all",
    group_mode: str = "off",
    max_per_group: int = 2,
) -> str:
    """
Search Danbooru tags using natural language and return a ready-to-use prompt.

## Args
- query: Natural language description (Chinese recommended).
- use_segmentation: Split multi-concept input into segments for separate retrieval. True for scene descriptions, False for single-concept queries.
- top_k: Candidates recalled per segment. Semantics change with use_segmentation — see guide below.
- limit: Max tags returned.
- popularity_weight: Influence of tag post count on ranking (0.0–1.0). Default 0.15.
- show_nsfw: Include NSFW tags. Default True.
- include_wiki: Append wiki description to each result. Default False.
- category: Filter results to a specific tag category. Default "all".
    "all"       —  All categories (通用 + 版权 + 人物 )
    "general"   —  General: visual attributes, clothing, pose, background, etc.
    "copyright" —  Copyright: specific anime/game/franchise titles
    "character" —  Character: named characters from any series
    Use this when you know what kind of tag you need — e.g. looking for a
    character name vs. describing a scene visually.
- group_mode: Tag group processing mode. Default "off".
    "off"     — No group processing (backward compatible)
    "expand"  — Boost same-group tags for concept exploration
    "diverse" — Limit tags per group for scene diversity
- max_per_group: Max tags per group in diverse mode. Default 2.

## Query writing guide

The `query` parameter supports explicit delimiter control for precise segmentation.

### Explicit delimiters

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 | When to use |
|---|---|---|
| Concept list (spaces) | `运动社团 校队 比赛 运动会` | You know the exact concepts to search |
| Concept list (dun hao) | `反乌托邦、赛博朋克、蒸汽朋克` | Same, with Chinese list punctuation |
| Natural sentence | `一个穿着白色水手服的少女在雨中奔跑` | Scene description, let the engine auto-split |
| Mixed | `运动社团 一个穿水手服的少女` | Mix concepts with descriptive phrases |

Segments >7 characters are still auto-split by jieba, but the raw segment is kept as an additional query
to preserve clause-level semantics (multi-granularity retrieval).

### Recommendations

1. **Concept lists → use explicit delimiters:** Group independent concepts with spaces or dunhao.
   `运动社团 校队 比赛 体育祭 田径部` is better than `运动社团校队比赛体育祭田径部`.

2. **Scene descriptions → write naturally:** Natural Chinese with Chinese commas works well for full scenes.
   `一个穿着白色水手服，蓝色短裙的少女在雨中奔跑` — commas here are grammatical, not delimiters.

3. **Precise lookup → turn off segmentation:** For finding a specific character or copyright title,
   set `use_segmentation=False` and combine with `category` filter.
   e.g. `query="EVA中蓝发的零号机驾驶员"` with `category="character"` and `use_segmentation=False`.

4. **Category filtering:** Use `category` to narrow results. Looking for a character?
   `category="character"`. Building a scene prompt? `category="general"`.

## Parameter guide

### Step 1 — Decide use_segmentation + top_k together

top_k means "candidates per segment"; its effect depends on whether segmentation is on.

Multi-concept input (scene description) → use_segmentation=True

| Sub-scenario          | top_k | Reason                                                   |
|-----------------------|-------|----------------------------------------------------------|
| Full scene → prompt   | 5     | Many segments; low top_k distributes result slots fairly |
| Vague concept explore | 80    | Few segments; high top_k needed for broad recall         |

Single-concept input → use_segmentation=False

top_k acts as total candidate pool size. Use 20 for all single-concept cases.

| Sub-scenario                  | top_k |
|-------------------------------|-------|
| Describe subject / find tag   | 20    |
| Precise lookup / spell fix    | 20    |

### Step 2 — Decide limit independently

| Goal                         | limit |
|------------------------------|-------|
| Full prompt for image gen    | 80    |
| Concept exploration          | 20–80 |
| Precise lookup / role search | 10–20 |

### Auxiliary params

popularity_weight (default 0.15, rarely needs changing):
- Higher (0.3+): favor common, well-established tags
- Lower (0.0): surface niche/rare tags

include_wiki (default False):
- True: The meaning of the tag is important — disambiguation, explaining tags to users, exploring unfamiliar domains, or when you are unsure of the tag's meaning
- False: Prompt generation (Wiki is irrelevant to the downstream task), tags are known

### Quick reference

| Scenario                      | use_segmentation | top_k | limit |
|-------------------------------|------------------|-------|-------|
| Full scene → prompt (default) | True             | 5     | 80    |
| Vague concept exploration     | True             | 80    | 80    |
| Describe subject / find tag   | False            | 20    | 20    |
| Precise lookup / spell fix    | False            | 20    | 10    |

### Workflow

After search_tags, pass selected tags to get_related_tags to discover complementary tags via co-occurrence (accessories, character features, scene atmosphere).
Supports chained exploration / iterative loops – take the interesting tags from the returned results as input to call get_related_tags again,
and use the results from get_related to feed back into a new round of search, enabling multi-hop deep traversal along the co-occurrence graph.

## Examples

Precise lookup / spell fix — e.g. "selafuku", "thighhigh", "twintail"
→ use_segmentation=False, top_k=20, limit=10

Vague concept exploration — e.g. "兔耳朵", "赛博朋克服装", "假肢"
→ use_segmentation=True, top_k=80, limit=80

Describe subject / find tag — e.g. "EVA中蓝发的零号机驾驶员", "命运石之门中的助手"
→ use_segmentation=False, top_k=20, limit=20

Full scene → prompt — e.g. "一个穿着白色水手服，蓝色短裙的少女在雨中的城市里奔跑"
→ use_segmentation=True, top_k=5, limit=80

## Returns
JSON with: prompt (comma-separated tags), keywords, results.
Each result: tag, cn_name, category, final_score, count[, wiki if include_wiki=True].
    """
    _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"],  # unrecognized value → fall back to all
    )

    tagger = await DanbooruTagger.get_instance()
    request = SearchRequest(
        query=query,
        top_k=top_k,
        limit=limit,
        popularity_weight=popularity_weight,
        show_nsfw=show_nsfw,
        use_segmentation=use_segmentation,
        target_categories=target_categories,
        group_mode=group_mode,
        max_per_group=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).

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 = 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)