DemonGatanjieu-Anomalous_Model_Browser/tests/test_ai_toolbox.py

"""What the user hands the AI to work with (api/ai_toolbox.py): picks checked and kept with the
task, given to the AI once as a message and in begin_task, favourites standing in, and "only"
refusing other custom or paid node types. A stand-in /object_info; temporary folders only."""
import os
import sys
import tempfile
import time
import unittest
from pathlib import Path
from unittest import mock

PLUGIN = Path(__file__).resolve().parents[1]
sys.path[:0] = [str(PLUGIN), str(PLUGIN.parents[1]), str(PLUGIN / "tests")]

from api import ai_panels, ai_tasks, ai_toolbox, mcp_canvas_edit, mcp_task_tools, reference_blocks  # noqa: E402
from api.mcp_tools import ToolError  # noqa: E402
from test_ai_tasks import isolate, ok  # noqa: E402

FAVOURITES = os.path.join(tempfile.mkdtemp(prefix="amb_ai_toolbox_"), "ai_toolbox.json")
INFO = {
    "KSampler": {"python_module": "nodes", "display_name": "KSampler", "input": {"required": {"model": ["MODEL"]}}, "output": ["LATENT"]},
    "ImageBlur": {"python_module": "comfy_extras.nodes_post_processing", "input": {}, "output": ["IMAGE"]},
    "FaceDetailer": {"python_module": "custom_nodes.ComfyUI-Impact-Pack", "display_name": "FaceDetailer", "input": {}, "output": ["IMAGE"]},
    "SAMLoader": {"python_module": "custom_nodes.ComfyUI-Impact-Pack", "input": {}, "output": ["SAM_MODEL"]},
    "KlingTextToVideo": {"python_module": "comfy_api_nodes.nodes_kling", "input": {}, "output": ["VIDEO"]},
}


class ToolboxTests(unittest.TestCase):
    def setUp(self):
        isolate()
        if os.path.exists(FAVOURITES):
            os.remove(FAVOURITES)
        for target, value in ((ai_toolbox, "_path"), (mcp_canvas_edit, "object_info")):
            patcher = mock.patch.object(target, value, (lambda: FAVOURITES) if value == "_path" else (lambda ctx: INFO))
            patcher.start()
            self.addCleanup(patcher.stop)
        ai_tasks.record("Gemini", "describe_canvas", {}, ok({}), time.time(), True)
        self.task = ai_tasks.list_tasks()[0]["id"]
        self.ctx = {"client": "Gemini", "task": self.task, "base_url": "http://127.0.0.1:1"}

    def test_a_pick_is_kept_and_reaches_the_ai_once(self):
        step = ai_toolbox.pick(self.task, {"nodes": ["FaceDetailer", "FaceDetailer"], "combos": ["赛博少女"]}, "prefer")
        self.assertEqual(step["toolbox"]["nodes"], ["FaceDetailer"], "each once")
        self.assertEqual(ai_tasks.get_task(self.task)["toolbox"]["combos"], ["赛博少女"])
        result = ai_panels.attach_pending(ok({}), "Gemini", task_id=self.task)
        message = result["structuredContent"]["user_messages"][0]
        self.assertEqual(message["kind"], "toolbox")
        self.assertIn("FaceDetailer", message["text"])
        self.assertNotIn("user_messages", ai_panels.attach_pending(ok({}), "Gemini", task_id=self.task)["structuredContent"])
        begun = mcp_task_tools.begin_task(self.ctx, "继续")
        self.assertIn("Use these first", begun["toolbox"])

    def test_favourites_stand_in_until_the_task_has_its_own(self):
        self.assertIsNone(ai_toolbox.effective(ai_tasks.get_task(self.task)))
        ai_toolbox.set_favorites({"nodes": ["SAMLoader"]})
        self.assertEqual(ai_toolbox.effective(ai_tasks.get_task(self.task))["from"], "favorites")
        ai_toolbox.pick(self.task, {"nodes": []}, "prefer")
        box = ai_toolbox.effective(ai_tasks.get_task(self.task))
        self.assertEqual((box["from"], box["nodes"]), ("task", []), "an empty pick is the user's choice too")

    def test_only_refuses_other_custom_and_paid_nodes_but_not_built_in_or_page_helpers(self):
        ai_toolbox.pick(self.task, {"nodes": ["FaceDetailer"]}, "only")
        self.assertEqual(ai_toolbox.refused(self.task, ["KSampler", "ImageBlur", "FaceDetailer", "Reroute", "SAMLoader", "KlingTextToVideo"], INFO),
                         ["KlingTextToVideo", "SAMLoader"])
        with mock.patch.object(mcp_canvas_edit, "_page", return_value={}) as page:
            with self.assertRaisesRegex(ToolError, "allowed only these node types.*FaceDetailer.*Not allowed: SAMLoader"):
                mcp_canvas_edit.add_node(self.ctx, "SAMLoader")
            with self.assertRaises(ToolError):
                mcp_canvas_edit.new_workflow(self.ctx, "x", {"1": {"class_type": "KlingTextToVideo", "inputs": {}}})
            page.assert_not_called()
            mcp_canvas_edit.add_node(self.ctx, "KSampler")
            page.assert_called_once()

    def test_a_whole_pack_allows_its_nodes_lists_them_first_and_in_full(self):
        ai_toolbox.pick(self.task, {"packs": ["ComfyUI-Impact-Pack"]}, "only")
        self.assertEqual(ai_toolbox.refused(self.task, ["FaceDetailer", "SAMLoader", "KSampler", "KlingTextToVideo"], INFO),
                         ["KlingTextToVideo"], "every node of the pack, and the built-in ones")
        with mock.patch.object(mcp_canvas_edit, "_page", return_value={}):
            with self.assertRaisesRegex(ToolError, "built-in ones\): every node of the pack ComfyUI-Impact-Pack"):
                mcp_canvas_edit.add_node(self.ctx, "KlingTextToVideo")
        found = mcp_canvas_edit.find_nodes(self.ctx, query="e")
        self.assertEqual({node["type"] for node in found["nodes"][:2]}, {"FaceDetailer", "SAMLoader"})
        self.assertTrue(all(node["picked_by_user"] for node in found["nodes"][:2]))
        box = mcp_task_tools.get_toolbox(self.ctx)
        self.assertEqual(box["packs"], [{"pack": "ComfyUI-Impact-Pack", "count": 2, "node_types": [
            {"type": "FaceDetailer", "display_name": "FaceDetailer", "category": ""},
            {"type": "SAMLoader", "display_name": "SAMLoader", "category": ""}]}])
        self.assertIn("custom node packs", box["note"])

    def test_prefer_refuses_nothing_and_find_nodes_lists_picks_first(self):
        ai_toolbox.pick(self.task, {"nodes": ["SAMLoader"]}, "prefer")
        self.assertEqual(ai_toolbox.refused(self.task, ["KlingTextToVideo"], INFO), [])
        found = mcp_canvas_edit.find_nodes(self.ctx, query="e")
        self.assertEqual(found["nodes"][0]["type"], "SAMLoader")
        self.assertTrue(found["nodes"][0]["picked_by_user"])
        paid = next(node for node in found["nodes"] if node["type"] == "KlingTextToVideo")
        self.assertEqual((paid["source"], paid.get("paid")), ("api", True))
        self.assertEqual(found["nodes"][-1]["type"], "KlingTextToVideo", "paid nodes last")

    def test_get_toolbox_gives_the_picks_in_full(self):
        self.assertFalse(mcp_task_tools.get_toolbox(self.ctx)["picked"])
        ai_toolbox.pick(self.task, {"nodes": ["FaceDetailer", "Fast Groups Muter (rgthree)"], "models": ["loras:0:gone.safetensors"]}, "only")
        with mock.patch("api.mcp_tools.get_model", side_effect=ToolError("gone")):
            box = mcp_task_tools.get_toolbox(self.ctx)
        self.assertEqual((box["mode"], box["nodes"][0]["type"], box["page_only_nodes"]), ("only", "FaceDetailer", ["Fast Groups Muter (rgthree)"]))
        self.assertTrue(box["models"][0]["missing"])

    def test_the_picker_gets_everything_in_one_answer_and_prompts_inside_combos_and_workflows(self):
        from api import ai_toolbox_catalog as catalog

        combos = [{"name": "赛博少女", "model": "a.safetensors", "loras": [], "base": "", "prompt": "1girl, neon", "structure": False,
                   "cover": "/anomalous/image?x", "image": None}]
        flows = [{"name": "高清修复", "filename": "f.json", "tags": ["SDXL"], "prompt": "city", "negative": "blurry"},
                 {"name": "空的", "filename": "e.json", "tags": [], "prompt": "", "negative": ""}]
        with mock.patch.object(catalog, "_combos", return_value=combos), mock.patch.object(catalog, "_workflows", return_value=flows), \
                mock.patch.object(catalog, "_models", return_value=[{"id": "loras:0:x.safetensors"}]), \
                mock.patch("api.mcp_tools.saved_prompts", return_value=iter([{"name": "夜景", "tags": [], "positive": "night", "negative": ""}])):
            data = ai_toolbox.catalog("all")
            ids = [entry["id"] for entry in data["prompts"]]
            self.assertEqual(ids, ["夜景", "combo:赛博少女", "workflow:高清修复"], "a workflow without prompts has none to pick")
            ai_toolbox.pick(self.task, {"prompts": ["workflow:高清修复"]}, "prefer")
            box = mcp_task_tools.get_toolbox(self.ctx)
        self.assertEqual((box["prompts"][0]["positive"], box["prompts"][0]["negative"]), ("city", "blurry"))
        self.assertIn("workflow:<name>", ai_toolbox.words(ai_toolbox.effective(ai_tasks.get_task(self.task))))
        self.assertIs(ai_toolbox.catalog("all"), data, "kept a few seconds")
        catalog._cache.update(at=0.0, data=None)

    def test_a_way_joins_the_picks_as_they_stand_and_wins_its_use(self):
        ai_toolbox.pick(self.task, {"nodes": ["FaceDetailer"]}, "only")
        step = ai_toolbox.pick_way(self.task, "upscale_tile", True)
        box = ai_tasks.get_task(self.task)["toolbox"]
        self.assertEqual((box["ways"], box["nodes"], box["mode"]), (["upscale_tile"], ["FaceDetailer"], "only"))
        self.assertIn("add_nodes(use='upscale', want='tile')", step["text"])
        self.assertEqual(mcp_task_tools.get_toolbox(self.ctx)["ways"], ["add_nodes(use='upscale', want='tile')"])
        self.assertEqual(reference_blocks.pick("upscale", None, None, False, box["ways"]), "upscale_tile", "the AI asked for none")
        self.assertEqual(reference_blocks.pick("upscale", "2x", None, False, box["ways"]), "upscale_tile", "the AI asked for another")
        self.assertEqual(reference_blocks.pick("upscale", "2x", None, False, ["upscale_tile", "upscale_tiled"]), "upscale_tiled")
        self.assertEqual(reference_blocks.pick("upscale", "2x", None, False, ["control_pose"]), "upscale_tiled", "another use's pick")
        ai_toolbox.pick_way(self.task, "upscale_tile", False)
        self.assertEqual(ai_tasks.get_task(self.task)["toolbox"]["ways"], [])
        with self.assertRaises(ai_toolbox.ToolboxError):
            ai_toolbox.pick_way(self.task, "no such way", True)

    def test_favourites_go_with_the_first_way_picked(self):
        ai_toolbox.set_favorites({"nodes": ["FaceDetailer"]})
        ai_toolbox.pick_way(self.task, "upscale", True)
        box = ai_tasks.get_task(self.task)["toolbox"]
        self.assertEqual((box["nodes"], box["ways"]), (["FaceDetailer"], ["upscale"]))

    def test_only_lets_in_what_a_picked_way_builds(self):
        custom = next(iter(sorted(reference_blocks.node_types(["inpaint_lanpaint"]) - {"SetUnionControlNetType"})))
        info = {custom: {"python_module": "custom_nodes.some-pack"}}
        ai_toolbox.pick(self.task, {"nodes": ["FaceDetailer"]}, "only")
        self.assertEqual(ai_toolbox.refused(self.task, [custom], info), [custom])
        ai_toolbox.pick_way(self.task, "inpaint_lanpaint", True)
        self.assertEqual(ai_toolbox.refused(self.task, [custom], info), [])

    def test_what_is_refused(self):
        for picks, mode in [(None, "prefer"), ({"nodes": "FaceDetailer"}, "prefer"), ({"nodes": [3]}, "prefer"),
                            ({"nodes": ["x" * 400]}, "prefer"), ({"nodes": ["a"] * 0 + [str(i) for i in range(61)]}, "prefer"),
                            ({"nodes": []}, "strict")]:
            with self.assertRaises(ai_toolbox.ToolboxError, msg=(picks, mode)):
                ai_toolbox.pick(self.task, picks, mode)
        with self.assertRaises(ai_toolbox.ToolboxError):
            ai_toolbox.pick("no-task", {}, "prefer")
        with self.assertRaises(ai_toolbox.ToolboxError):
            ai_toolbox.catalog("nodes")


if __name__ == "__main__":
    unittest.main()