DemonGatanjieu-Anomalous_Model_Browser/api/mcp_tools.py

"""The MCP endpoint's (mcp_server.py) reading tools: views of the library (mcp_actions.py acts).

`call_tool` runs any tool of a registry, these and mcp_actions.py's. search_models and get_model
are offered through mcp_model_use.py, which adds what loads each model.

Each tool is a plain function over the same stores the browser's pages use (model catalog
and sidecars, output gallery, combos, workflow recipes, saved prompts, generated audio, scan
results, activity log). They never write. Results are small, already-summarised JSON (an AI
reads them), with ids the detail tools take back: model_id "type:path_idx:relative/path" and
image "subfolder/file.png" under ComfyUI's output folder. A tool's mistake (bad id, missing
file) comes back as a tool error the AI can read and correct, never as an exception.
"""

import base64
import datetime
import html
import json
import logging
import os
import re

import folder_paths

from .activity_log import list_entries
from .audio_catalog import _gallery_page as _audio_page
from .gallery_routes import _gallery_listing
from .image_search import read_png_text
from .kept_images import _kept_images
from .material_store import _list_materials, _read_material, get_materials_dir
from .media_routes import _build_card_thumbnail
from .metadata import get_metadata
from .model_catalog import _iter_search_models
from .model_constants import MEDIA_EXTENSIONS, MODEL_EXTENSIONS, PREVIEW_SUFFIXES
from .notebooks import _list_notebooks
from .path_utils import model_file_within, resolve_folder_subdir, resolve_within
from .recipe_store import _list_recipes, get_recipes_dir
from .scan_report import last_scan_path
from .scan_summary import _summarize

MAX_LIMIT = 50
MAX_SEARCH_PICTURES = 6  # search_images(with_pictures): the first few attached
MAX_TEXT = 1500
MAX_IMAGE_BYTES = 1_500_000
IMAGE_TYPES = {".png": "image/png", ".jpg": "image/jpeg", ".jpeg": "image/jpeg", ".webp": "image/webp"}
SOURCES = {"civitai": "civitai", "local": "inferred_from_file", "": "not_scanned"}


class ToolError(Exception):
    """A problem the AI can fix (bad id, nothing there); shown to it as the tool's result."""


# ---------------------------------------------------------------- helpers

def _text(value, limit=MAX_TEXT):
    """Plain text from HTML-ish Civitai descriptions, shortened."""
    value = html.unescape(re.sub(r"<[^>]+>", " ", str(value or "")))
    value = re.sub(r"\s+", " ", value).strip()
    return value if len(value) <= limit else value[:limit] + "…"


def _when(seconds):
    try:
        return datetime.datetime.fromtimestamp(float(seconds)).isoformat(timespec="seconds")
    except (TypeError, ValueError, OverflowError, OSError):
        return None


def _limit(value, default=20):
    try:
        return max(1, min(MAX_LIMIT, int(value)))
    except (TypeError, ValueError):
        return default


def _terms(query):
    return [term for term in str(query or "").lower().split() if term]


def _matches(terms, *texts):
    haystack = " ".join(str(text or "") for text in texts).lower()
    return all(term in haystack for term in terms)


def _model_id(folder_type, path_idx, base_dir, file_path):
    return f"{folder_type}:{path_idx}:{os.path.relpath(file_path, base_dir).replace(os.sep, '/')}"


def _resolve_model(model_id):
    """(type, path_idx, base_dir, file_path) for a model_id from search_models."""
    parts = str(model_id or "").split(":", 2)
    if len(parts) != 3:
        raise ToolError("model_id looks like 'loras:0:folder/file.safetensors'; take it from search_models.")
    folder_type, path_idx, rel = parts
    rel = rel.replace("\\", "/").strip("/")
    subfolder, filename = os.path.split(rel)
    try:
        path_idx = int(path_idx)
        base_dir, target_dir = resolve_folder_subdir(folder_type, path_idx, "/" + subfolder)
        file_path = model_file_within(target_dir, filename)
    except (ValueError, KeyError, TypeError):
        raise ToolError(f"No models folder for {model_id}.")
    if not filename.lower().endswith(MODEL_EXTENSIONS) or not os.path.isfile(file_path):
        raise ToolError(f"No model file at {model_id}; it may have been moved or renamed.")
    return folder_type, path_idx, base_dir, file_path


def _resolve_image(image):
    """The output image at `image` ('subfolder/file.png'), checked to stay inside the output folder."""
    rel = str(image or "").replace("\\", "/").strip("/")
    try:
        path = resolve_within(folder_paths.get_output_directory(), rel)
    except ValueError:
        raise ToolError("image is a path under ComfyUI's output folder, as search_images returns it.")
    if os.path.splitext(path)[1].lower() not in IMAGE_TYPES or not os.path.isfile(path):
        raise ToolError(f"No output image at {image}: it is not (or no longer) in ComfyUI's output folder; "
                        "the user may have deleted it. search_images lists the pictures there.")
    return rel, path


def _image_content(path):
    """An MCP image item: the 512 px thumbnail the gallery uses, or None when too big or unreadable."""
    try:
        thumb = _build_card_thumbnail(path)
        with open(thumb, "rb") as source:
            data = source.read(MAX_IMAGE_BYTES + 1)
    except OSError:
        return None
    if len(data) > MAX_IMAGE_BYTES:
        return None
    mime = IMAGE_TYPES.get(os.path.splitext(thumb)[1].lower(), "image/webp")
    return {"type": "image", "data": base64.b64encode(data).decode("ascii"), "mimeType": mime}


def _kept_lookup():
    return {(item["subfolder"], item["filename"]): {kind: (value or {}).get("name") for kind, value in item.items()
                                                    if kind in ("recipe", "combo", "prompt") and value}
            for item in _kept_images()}


def _model_brief(folder_type, path_idx, base_dir, file_path, meta):
    rel_dir = os.path.relpath(os.path.dirname(file_path), base_dir).replace(os.sep, "/")
    try:
        size_mb = round(os.path.getsize(file_path) / 1048576, 1)
    except OSError:
        size_mb = None
    return {
        "model_id": _model_id(folder_type, path_idx, base_dir, file_path),
        "name": meta.get("custom_name") or meta.get("name") or os.path.basename(file_path),
        "file": os.path.basename(file_path),
        "type": folder_type,
        "folder": "" if rel_dir == "." else rel_dir,
        "base_model": meta.get("baseModel") or "",
        "source": SOURCES.get(meta.get("info_source", ""), "not_scanned"),
        "trigger_words": list(meta.get("trainedWords") or [])[:8],
        "size_mb": size_mb,
    }


def _each_model():
    """Every model file once (folders registered twice, like unet and diffusion_models, count once).
    ComfyUI also registers custom_nodes, whose node packs carry files of their own: not the library."""
    seen = set()
    for folder_type, path_idx, base_dir, file_path in _iter_search_models():
        if folder_type == "custom_nodes" or not file_path.lower().endswith(MODEL_EXTENSIONS):
            continue
        real = os.path.normcase(os.path.realpath(file_path))
        if real in seen:
            continue
        seen.add(real)
        yield folder_type, path_idx, base_dir, file_path


# ---------------------------------------------------------------- models and scans

def search_models(query="", type="", base_model="", source="", limit=20):
    terms, wanted_base, found = _terms(query), str(base_model or "").lower(), []
    total = 0
    for folder_type, path_idx, base_dir, file_path in _each_model():
        if type and folder_type != type:
            continue
        meta = get_metadata(file_path)
        if wanted_base and wanted_base not in str(meta.get("baseModel") or "").lower():
            continue
        if source and SOURCES.get(meta.get("info_source", ""), "not_scanned") != source:
            continue
        if not _matches(terms, os.path.relpath(file_path, base_dir), meta.get("name"), meta.get("custom_name"),
                        meta.get("baseModel"), " ".join(meta.get("trainedWords") or [])):
            continue
        total += 1
        if len(found) < _limit(limit):
            found.append(_model_brief(folder_type, path_idx, base_dir, file_path, meta))
    return {"total_matches": total, "models": found}


def get_model(model_id, include_cover=False):
    folder_type, path_idx, base_dir, file_path = _resolve_model(model_id)
    meta = get_metadata(file_path)
    detail = _model_brief(folder_type, path_idx, base_dir, file_path, meta)
    detail.update({
        "trigger_words": list(meta.get("trainedWords") or []),
        "civitai_name": meta.get("name") or "",
        "civitai_url": meta.get("civitai_url") or "",
        "download_link": meta.get("source_url") or "",
        "user_notes": meta.get("custom_notes") or "",
        "fields_set_by_user": list(meta.get("user_fields") or []),
        "description": _text(meta.get("description")),
        "sha256": meta.get("hash") or "",
    })
    if meta.get("info_source") == "local":
        detail["why_not_on_civitai"] = meta.get("unmatched_reason") or "unknown"
    extra = []
    if include_cover:
        base = os.path.splitext(file_path)[0]
        cover = next((base + suffix for suffix in PREVIEW_SUFFIXES + MEDIA_EXTENSIONS
                      if os.path.splitext(suffix)[1].lower() in IMAGE_TYPES and os.path.isfile(base + suffix)), None)
        item = _image_content(cover) if cover else None
        detail["cover"] = "attached" if item else "none"
        extra = [item] if item else []
    return detail, extra


def library_overview():
    by_type, by_base = {}, {}
    for folder_type, _idx, _base, file_path in _each_model():
        by_type[folder_type] = by_type.get(folder_type, 0) + 1
        base = get_metadata(file_path).get("baseModel") or "unknown"
        by_base[base] = by_base.get(base, 0) + 1
    summary = _summarize()
    return {
        "models_by_type": dict(sorted(by_type.items(), key=lambda item: -item[1])),
        "models_by_base_model": dict(sorted(by_base.items(), key=lambda item: -item[1])[:20]),
        "scan": {key: summary[key] for key in ("total", "matched", "unmatched", "new", "pending")},
        "combos": len(_list_notebooks()),
        "workflow_recipes": len(_list_recipes(get_recipes_dir())),
        "saved_prompts": sum(1 for item in _list_materials(get_materials_dir()) if item.get("kind") == "prompt_plan"),
    }


def scan_report():
    summary = _summarize()
    try:
        with open(last_scan_path(), encoding="utf-8") as source:
            last = json.load(source)
    except (OSError, ValueError):
        last = None
    report = {
        "counts": {key: summary[key] for key in ("total", "matched", "unmatched", "new", "pending")},
        "meaning": "matched: found on Civitai; unmatched: base model inferred from the file (reason says why); "
                   "new: never scanned; pending: unmatched only because Civitai was not asked (offline or no connection).",
        "unmatched_models": summary["unmatched_models"][:MAX_LIMIT],
        "not_scanned_models": summary["new_models"][:MAX_LIMIT],
    }
    if isinstance(last, dict) and last.get("counts"):
        report["last_scan"] = {"finished": _when(last.get("finished")), "kind": last.get("kind"),
                               "counts": last.get("counts"), "civitai_unreachable": bool(last.get("civitai_down")),
                               "models": (last.get("files") or [])[:30]}
    return report


# ---------------------------------------------------------------- outputs

def search_images(query="", limit=20, page=1, with_pictures=False):
    images = _gallery_listing(folder_paths.get_output_directory(), False, str(query or "").strip())
    limit = _limit(limit)
    try:
        page = max(1, int(page))
    except (TypeError, ValueError):
        page = 1
    kept = _kept_lookup()
    chosen = images[(page - 1) * limit:page * limit]
    found = {
        "total_matches": len(images), "page": page,
        "images": [{"image": "/".join(part for part in (item["subfolder"], item["filename"]) if part),
                    "created": _when(item.get("mtime")),
                    **({"kept_as": kept[(item["subfolder"], item["filename"])]}
                       if (item["subfolder"], item["filename"]) in kept else {})}
                   for item in chosen],
    }
    if not with_pictures:
        return found
    pictures = []  # the first few themselves, in the order listed, so an AI that sees pictures need not ask for each
    for entry in found["images"][:MAX_SEARCH_PICTURES]:
        item = _image_content(_resolve_image(entry["image"])[1])
        entry["attached"] = bool(item)
        if item:
            pictures.append(item)
    return found, pictures


def get_image_info(image, include_image=False):
    rel, path = _resolve_image(image)
    info = {"image": rel, "created": _when(os.path.getmtime(path))}
    chunks = read_png_text(path) if path.lower().endswith(".png") else {}
    try:
        prompt = json.loads(chunks.get("prompt") or "null")
    except ValueError:
        prompt = None
    if isinstance(prompt, dict):
        nodes, models = [], set()
        for node_id, node in list(prompt.items())[:80]:
            if not isinstance(node, dict):
                continue
            inputs = {}
            for key, value in (node.get("inputs") or {}).items():
                if isinstance(value, str):
                    inputs[key] = _text(value, 2000) if len(value) > 2000 else value
                    if value.lower().endswith(MODEL_EXTENSIONS):
                        models.add(value.replace("\\", "/"))
                elif isinstance(value, (int, float, bool)):
                    inputs[key] = value
            nodes.append({"id": node_id, "type": node.get("class_type"), "inputs": inputs})
        info.update(models_used=sorted(models), nodes=nodes)
    elif chunks.get("parameters"):
        info["parameters"] = _text(chunks["parameters"], 4000)
    else:
        info["note"] = "No generation settings are stored in this image."
    kept = _kept_lookup().get(tuple(rel.rsplit("/", 1)) if "/" in rel else ("", rel))
    if kept:
        info["kept_as"] = kept
    extra = []
    if include_image:
        item = _image_content(path)
        info["image_attached"] = bool(item)
        extra = [item] if item else []
    return info, extra


def search_audio(query="", limit=20):
    items, total, _more = _audio_page(folder_paths.get_output_directory(), 1, _limit(limit), str(query or "").lower().strip())
    return {"total_matches": total, "audio": [{
        "file": "/".join(part for part in (item["subfolder"], item["filename"]) if part),
        "created": _when(item.get("mtime")),
        "speech": (item.get("generation") or {}).get("speech") or "",
        "voice_sample": (item.get("generation") or {}).get("sample") or "",
        "seed": (item.get("generation") or {}).get("seed"),
    } for item in items]}


# ---------------------------------------------------------------- what the user kept

def _model_name(model):
    if isinstance(model, dict):
        return model.get("filename") or model.get("name") or ""
    return str(model or "")


def combo_view(note):
    """A combo (one of _list_notebooks) as the AI and the materials picker see it, with its LoRAs and texts."""
    data = note.get("data") or {}
    if data.get("kind") == "nodes":  # its own node structure: model and text slots by their names
        structure = data.get("structure") or {}
        values = data.get("values") or {}
        slots = structure.get("slots") or []
        models = [{"slot": s.get("label") or s.get("widget"), "folder": s.get("folder") or "",
                   "file": os.path.basename(str(values.get(s["id"]) or "").replace("\\", "/"))}
                  for s in slots if s.get("kind") == "model"]
        texts = [{"slot": s.get("label") or s.get("widget"), "text": _text(values.get(s["id"]), 400)}
                 for s in slots if s.get("kind") == "text"]
        combo = {"name": note.get("name"), "kind": "node structure", "models": models, "texts": texts,
                 "nodes": [n.get("name") or n.get("type") for n in structure.get("nodes") or []],
                 "main_model": models[0]["file"] if models else "", "prompt": texts[0]["text"] if texts else ""}
        return combo, [model["file"] for model in models[1:]] + [text["text"] for text in texts[1:]]
    loras = [_model_name(lora) for lora in data.get("loras") or []]
    return {"name": note.get("name"), "base_model": data.get("baseModel") or "", "main_model": _model_name(data.get("mainModel")),
            "loras": loras, "prompt": _text(data.get("promptEn"), 800)}, loras


def list_combos(query="", limit=20, *, exact_name=False):
    terms, found = _terms(query), []
    for note in _list_notebooks():
        combo, loras = combo_view(note)
        if exact_name and combo["name"] != query:
            continue
        if not exact_name and not _matches(terms, combo["name"], combo["main_model"], combo["prompt"], " ".join(loras)):
            continue
        found.append(combo)
    return {"total_matches": len(found), "combos": found[:_limit(limit)]}


def list_workflows(query="", limit=20, *, exact_name=False):
    terms, found = _terms(query), []
    for recipe in _list_recipes(get_recipes_dir()):
        data = recipe.get("data") or {}
        params = json.dumps(data.get("params") or {}, ensure_ascii=False)
        if exact_name and (data.get("name") or recipe["filename"]) != query:
            continue
        if not exact_name and not _matches(terms, data.get("name"), " ".join(data.get("tags") or []), data.get("notes"), params):
            continue
        found.append({"name": data.get("name") or recipe["filename"], "tags": data.get("tags") or [],
                      "notes": _text(data.get("notes"), 500), "saved": _when((data.get("timestamp") or 0) / 1000),
                      "settings": params if len(params) <= MAX_TEXT else params[:MAX_TEXT] + "…"})
    return {"total_matches": len(found), "workflows": found[:_limit(limit)]}


def saved_prompts():
    """Every saved prompt plan: {name, tags, positive, negative}."""
    directory = get_materials_dir()
    for summary in _list_materials(directory):
        if summary.get("kind") != "prompt_plan" or summary.get("moved_to_recipe"):
            continue
        try:
            plan = _read_material(os.path.join(directory, summary["filename"])).get("plan") or {}
        except (OSError, ValueError, KeyError):
            continue
        yield {"name": summary.get("name"), "tags": summary.get("tags") or [],
               "positive": _text(plan.get("positive"), 1200), "negative": _text(plan.get("negative"), 600)}


def list_saved_prompts(query="", limit=20):
    terms = _terms(query)
    found = [prompt for prompt in saved_prompts()
             if _matches(terms, prompt["name"], " ".join(prompt["tags"]), prompt["positive"], prompt["negative"])]
    return {"total_matches": len(found), "prompts": found[:_limit(limit)]}


def recent_activity(limit=20, source=""):
    entries = list_entries(_limit(limit), source if source in ("canvas", "file") else None)
    return {"entries": [{"time": _when(entry.get("time")), "source": entry.get("source"), "action": entry.get("action"),
                         "target": entry.get("target") or "", "detail": entry.get("detail") or {}} for entry in entries]}


# ---------------------------------------------------------------- registry

def _schema(properties=None, required=()):
    schema = {"type": "object", "properties": properties or {}, "additionalProperties": False}
    if required:
        schema["required"] = list(required)
    return schema


_QUERY = {"type": "string", "description": "Words that must all appear (case-insensitive); empty lists everything."}
_LIMIT = {"type": "integer", "minimum": 1, "maximum": MAX_LIMIT, "default": 20}


def _tool(name, title, description, run, properties=None, required=()):
    return name, {"run": run, "spec": {
        "name": name, "title": title, "description": description,
        "inputSchema": _schema(properties, required),
        "annotations": {"title": title, "readOnlyHint": True, "destructiveHint": False,
                        "idempotentHint": True, "openWorldHint": False},
    }}


TOOLS = dict([
    _tool("library_overview", "Library overview",
          "How many models of each type and base model the user has, how many are scanned, and how many "
          "combos, workflow recipes and saved prompts they kept. A good first call.", library_overview),
    _tool("scan_report", "Scan status",
          "Which models are matched on Civitai, which are not and why, which were never scanned, and what the last scan did.",
          scan_report),
    _tool("search_images", "Search output images",
          "Find images ComfyUI generated, newest first, by words in their prompt, model or settings (e.g. a LoRA "
          "name or 'seed:123'). Returns image paths for get_image_info. with_pictures attaches the first 6 "
          "themselves (512 px), in the listed order, so you can look at them at once.",
          search_images, {"query": _QUERY, "limit": _LIMIT, "page": {"type": "integer", "minimum": 1, "default": 1},
                          "with_pictures": {"type": "boolean", "default": False}}),
    _tool("get_image_info", "Image generation settings",
          "The nodes and settings that produced an output image (prompts, models, seed, sampler…) and whether the "
          "user kept it as a recipe, combo or prompt. include_image attaches a 512 px copy.",
          get_image_info, {"image": {"type": "string", "description": "From search_images, e.g. 'ComfyUI_00012_.png'."},
                           "include_image": {"type": "boolean", "default": False}}, ["image"]),
    _tool("list_combos", "Combos",
          "The user's combos (搭配): a main model, LoRAs and a prompt kept together, or (kind \"node "
          "structure\") model nodes and text nodes saved from the canvas with their links, listed by slot.", list_combos,
          {"query": _QUERY, "limit": _LIMIT}),
    _tool("list_workflows", "Workflow recipes",
          "The user's saved workflow recipes with tags, notes and their main settings.", list_workflows,
          {"query": _QUERY, "limit": _LIMIT}),
    _tool("list_saved_prompts", "Saved prompts",
          "Prompts the user saved in the Prompt Studio, with positive and negative text.", list_saved_prompts,
          {"query": _QUERY, "limit": _LIMIT}),
    _tool("search_audio", "Generated audio",
          "Speech the user generated (Anomalous TTS), newest first, with its text, voice and seed.", search_audio,
          {"query": _QUERY, "limit": _LIMIT}),
    _tool("recent_activity", "Activity log",
          "What Anomalous changed recently, newest first: canvas edits it made and files it saved, scanned or deleted.",
          recent_activity, {"limit": _LIMIT, "source": {"type": "string", "enum": ["canvas", "file"]}}),
])


def call_tool(tools, name, arguments, ctx=None):
    """An MCP tools/call result for tool `name` of `tools` (run in a worker thread). Tools marked
    `context` get `ctx` ({base_url} of this server) first."""
    tool = tools[name]
    spec = tool["spec"]["inputSchema"]
    unknown = set(arguments) - set(spec["properties"])
    missing = [key for key in spec.get("required", []) if key not in arguments]
    if unknown or missing:
        return _failure(f"Unknown arguments {sorted(unknown)}" if unknown else f"Missing arguments {missing}")
    try:
        value = tool["run"](ctx, **arguments) if tool.get("context") else tool["run"](**arguments)
    except ToolError as error:
        return _failure(str(error))
    except Exception as error:  # noqa: BLE001 - an AI app gets a readable failure, never a dropped request
        logging.getLogger(__name__).warning("Anomalous MCP tool %s failed", name, exc_info=True)
        return _failure(f"{name} failed: {error}")
    data, extra = value if isinstance(value, tuple) else (value, [])
    return {"content": [{"type": "text", "text": json.dumps(data, ensure_ascii=False)}, *extra],
            "structuredContent": data, "isError": False}


def _failure(message):
    return {"content": [{"type": "text", "text": message}], "isError": True}