- ui_ai_now.js: over the message box, a run's progress with stop, the AI's step still going, a gate to allow or refuse there, and the AI's open question with its options as buttons (own words or skip too; several picks and pictures lead to the stage). progressWords is shared with the run card; the panel's and gate's answers are exported for it. - ui_ai_drop.js: pictures dropped on the box or pasted are kept in output/ai_inbox (POST /anomalous/ai_pictures/upload, ai_pictures.save_upload: read as a picture first, 25 MB cap, never over a file) and attached; ComfyUI output pictures dragged in attach as they are. - Long AI notes are folded to their first lines with "Show all". Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
180 lines
10 KiB
Python
180 lines
10 KiB
Python
"""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()
|