"""A run's shape and what changed (ai_run_map.py), the task's own canvas tab (ai_task_canvas.py), pictures handed to the AI (ai_pictures.py) and show_run_map.""" import io import json import os import time import unittest from unittest import mock from PIL import Image from PIL.PngImagePlugin import PngInfo from test_ai_tasks import OUTPUT, isolate, ok from api import ai_panels, ai_pictures, ai_run_map, ai_task_canvas, ai_tasks, mcp_task_tools def lora_workflow(steps=20, strength=0.65): return { "4": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "sdxl/Kirazuri_v4.safetensors"}, "_meta": {"title": "Kirazuri"}}, "10": {"class_type": "LoraLoader", "inputs": {"model": ["4", 0], "clip": ["4", 1], "lora_name": "agefx.safetensors", "strength_model": strength, "strength_clip": 1}}, "11": {"class_type": "Reroute", "inputs": {"": ["10", 0]}}, "6": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["10", 1], "text": "irena, grey hair, witch hat"}}, "7": {"class_type": "CLIPTextEncode", "inputs": {"clip": ["10", 1], "text": "lowres"}}, "5": {"class_type": "EmptyLatentImage", "inputs": {"width": 1024, "height": 1536, "batch_size": 2}}, "3": {"class_type": "KSampler", "inputs": {"model": ["11", 0], "positive": ["6", 0], "negative": ["7", 0], "latent_image": ["5", 0], "seed": 123, "steps": steps, "cfg": 5, "sampler_name": "euler", "scheduler": "normal", "denoise": 1}}, "8": {"class_type": "VAEDecode", "inputs": {"samples": ["3", 0], "vae": ["4", 2]}}, "9": {"class_type": "SaveImage", "inputs": {"images": ["8", 0], "filename_prefix": "irena"}}, } class MapTests(unittest.TestCase): def test_boxes_roles_facts_and_links_through_plumbing(self): shape = ai_run_map.map_of(lora_workflow()) boxes = {box["id"]: box for box in shape["boxes"]} self.assertNotIn("11", boxes, "a reroute is wired through, not shown") self.assertEqual(boxes["4"]["role"], "model") self.assertEqual(boxes["4"]["facts"]["file"], "Kirazuri_v4") self.assertEqual((boxes["10"]["role"], boxes["10"]["facts"]["strength"]), ("lora", 0.65)) self.assertEqual(boxes["6"]["facts"], {"text": "irena, grey hair, witch hat", "polarity": "positive"}) self.assertEqual(boxes["7"]["facts"]["polarity"], "negative") self.assertEqual(boxes["5"]["facts"], {"width": 1024, "height": 1536, "batch": 2}) self.assertEqual(boxes["3"]["facts"], {"steps": 20, "cfg": 5, "sampler_name": "euler", "scheduler": "normal", "denoise": 1}) self.assertIn({"from": "10", "to": "3", "type": "MODEL", "input": "model"}, [ {**link, "type": "MODEL"} for link in shape["links"] if link["from"] == "10" and link["to"] == "3"]) self.assertNotIn("3.seed", shape["params"], "the seed changes every run; it is no change to report") def test_changes_since_the_run_before(self): before = ai_run_map.map_of(lora_workflow()) after_workflow = lora_workflow(steps=30, strength=0.8) after_workflow["12"] = {"class_type": "LoraLoader", "inputs": {"model": ["10", 0], "lora_name": "light.safetensors"}} del after_workflow["9"] changes = ai_run_map.changes(before, ai_run_map.map_of(after_workflow)) self.assertIn({"node": "12", "title": "LoraLoader", "change": "added"}, changes) self.assertIn({"node": "9", "title": "SaveImage", "change": "removed"}, changes) self.assertIn({"node": "3", "title": "KSampler", "key": "steps", "from": "20", "to": "30"}, changes) self.assertIn({"node": "10", "title": "LoraLoader", "key": "strength_model", "from": "0.65", "to": "0.8"}, changes) self.assertEqual(ai_run_map.changes(None, before), []) def test_a_huge_workflow_keeps_what_makes_the_picture(self): workflow = lora_workflow() for index in range(40): workflow[f"x{index}"] = {"class_type": "SomeHelper", "inputs": {"value": index}} shape = ai_run_map.map_of(workflow) self.assertEqual(len(shape["boxes"]), ai_run_map.MAX_BOXES) self.assertTrue({"3", "4", "10", "6", "5"} <= {box["id"] for box in shape["boxes"]}) class TaskCanvasTests(unittest.TestCase): def setUp(self): isolate() self.task = ai_tasks.begin_task("A", "Irena portrait")["id"] def test_the_task_gets_its_tab_and_the_ai_is_told_once(self): self.assertEqual(ai_task_canvas.default_mode(), "new") self.assertIsNone(ai_task_canvas.spec("set_values"), "outside a tool call there is no task") holder = ai_task_canvas.begin_call(self.task) self.assertEqual(ai_task_canvas.spec("set_values"), {"mode": "new", "tab": None, "title": "Irena portrait"}) self.assertIsNone(ai_task_canvas.spec("open_model"), "not a canvas action") ai_task_canvas.remember({"path": "workflows/Irena portrait 101500.json", "name": "Irena portrait 101500", "made": True}) self.assertEqual(ai_tasks.get_task(self.task)["canvas"]["tab"]["path"], "workflows/Irena portrait 101500.json") told = ai_task_canvas.tell(ok({"done": True}), holder) self.assertIn("Irena portrait 101500", told["structuredContent"]["canvas_tab"]["note"]) again = ai_task_canvas.begin_call(self.task) self.assertEqual(ai_task_canvas.spec("add_node")["tab"]["name"], "Irena portrait 101500") self.assertNotIn("canvas_tab", ai_task_canvas.tell(ok({}), again)["structuredContent"]) def test_the_user_chooses_per_task_and_by_default(self): ai_task_canvas.set_mode(self.task, "current") ai_task_canvas.begin_call(self.task) self.assertEqual(ai_task_canvas.spec("set_values")["mode"], "current") other = ai_tasks.begin_task("B", "Other", new_task=True)["id"] ai_task_canvas.set_default("current") self.assertEqual(ai_task_canvas.mode_of(ai_tasks.get_task(other)), "current") with self.assertRaises(ValueError): ai_task_canvas.set_mode(self.task, "sideways") with self.assertRaises(ValueError): ai_task_canvas.set_default("sideways") ai_task_canvas.set_default("new") def picture(name, workflow=None): info = PngInfo() if workflow: info.add_text("prompt", json.dumps(workflow)) Image.new("RGB", (64, 96), (40, 80, 120)).save(os.path.join(OUTPUT, name), pnginfo=info) return name class PictureTests(unittest.TestCase): def setUp(self): isolate() self.task = ai_tasks.begin_task("A", "Look at these")["id"] self.patch = mock.patch.object(ai_tasks.folder_paths, "get_output_directory", return_value=OUTPUT) self.patch.start() self.addCleanup(self.patch.stop) def test_pictures_reach_the_ai_themselves_with_what_made_them(self): ai_pictures.hand(self.task, [picture("one.png", lora_workflow())]) ai_pictures.hand(self.task, [picture("two.png")]) steps = ai_tasks.get_task(self.task)["steps"] self.assertEqual(len(steps), 1, "pictures handed before the AI looked join one step") self.assertEqual([item["filename"] for item in steps[0]["results"]], ["one.png", "two.png"]) result = ai_panels.attach_pending(ok({}), "A", task_id=self.task) message = result["structuredContent"]["user_messages"][0] self.assertEqual((message["kind"], message["images"]), ("pictures", ["one.png", "two.png"])) made = message["made_with"][0] self.assertEqual((made["models"], made["loras"], made["positive"], made["size"]), (["Kirazuri_v4"], ["agefx (0.65)"], "irena, grey hair, witch hat", "1024x1536")) self.assertEqual(message["made_with"][1], {"image": "two.png"}) self.assertEqual(sum(1 for item in result["content"] if item.get("type") == "image"), 2) self.assertNotIn("_pictures", message) def test_only_output_pictures_and_open_tasks(self): for images in ([], ["../secret.png"], ["missing.png"], [picture("x.png")] * 9): with self.assertRaises(ai_pictures.PicturesError): ai_pictures.hand(self.task, images) with self.assertRaises(ai_pictures.PicturesError): ai_pictures.hand("nope", [picture("y.png")]) def test_a_dropped_picture_is_kept_in_the_inbox_never_over_a_file(self): data = io.BytesIO() Image.new("RGB", (8, 8), "red").save(data, "JPEG") first = ai_pictures.save_upload(data.getvalue(), "../../草稿 参考.png", now=1000) second = ai_pictures.save_upload(data.getvalue(), "../../草稿 参考.png", now=1000) self.addCleanup(lambda: [os.remove(os.path.join(OUTPUT, rel)) for rel in (first, second)]) self.assertTrue(first.startswith("ai_inbox/草稿 参考_") and first.endswith(".jpg"), first) # its stem, its real type self.assertNotEqual(first, second) ai_pictures.hand(self.task, [first]) # an output picture like any other for bad in (b"", b"not a picture", b"x" * (ai_pictures.MAX_UPLOAD + 1)): with self.assertRaises(ai_pictures.PicturesError): ai_pictures.save_upload(bad, "x.png") class ShowRunMapTests(unittest.TestCase): def setUp(self): isolate() self.task = ai_tasks.begin_task("A", "Run")["id"] self.ctx = {"client": "A", "task": self.task} def test_the_stage_shows_the_latest_run_that_has_its_shape(self): from api.mcp_tools import ToolError with self.assertRaises(ToolError): mcp_task_tools.show_run_map(self.ctx, "结构") step = ai_tasks.record("A", "run_workflow", {}, ok({}), time.time()) with ai_tasks._lock: task = ai_tasks._load()[self.task] task["steps"][-1].update(map=ai_run_map.map_of(lora_workflow()), changes=[]) ai_tasks._save(task) answer = mcp_task_tools.show_run_map(self.ctx, "结构") self.assertEqual(answer["run"], step["id"]) panel = ai_tasks.get_task(self.task)["steps"][-1]["panel"] self.assertEqual(panel["blocks"][0]["type"], "run_map") self.assertEqual(panel["blocks"][0]["step"], step["id"]) if __name__ == "__main__": unittest.main()