From 470f55dec4e66ae7a12010e2dbbe0069b256db70 Mon Sep 17 00:00:00 2001 From: Isulion Date: Tue, 20 Oct 2026 20:38:19 +0200 Subject: [PATCH] feat(megaprompt-v4): pure roll() with locked/pinned/disabled reels and frozen prompt Co-Authored-By: Claude Opus 5.5 --- Core_Nodes/mega_prompt_v4.py | 67 +++++++++++++++++- tests/test_mega_prompt_v4.py | 131 +++++++++++++++++++++++++++++++++++ 2 files changed, 197 insertions(+), 1 deletion(-) diff --git a/Core_Nodes/mega_prompt_v4.py b/Core_Nodes/mega_prompt_v4.py index f81e576..669edd8 100644 --- a/Core_Nodes/mega_prompt_v4.py +++ b/Core_Nodes/mega_prompt_v4.py @@ -7,10 +7,12 @@ reels with the V3 theme handlers and assembles the prompt. The same ``roll`` serves workflow execution and the preview routes (mega_prompt_v4_routes.py). """ +import copy import json import logging +import random -from .mega_prompt_V3 import get_theme_registry +from .mega_prompt_V3 import GENERATION_LOCK, get_theme_registry logger = logging.getLogger(__name__) @@ -157,3 +159,66 @@ def list_themes(registry=None): registry = registry or get_theme_registry() entries = [theme_entry(name, registry) for name in registry.handlers] return sorted(entries, key=lambda e: e["display"].split(" ", 1)[-1].strip()) + + +# reel name -> key of the dict returned by BaseThemeHandler.generate +COMPONENT_KEYS = {"subject": "subject", "location": "environment", "style": "style", "effects": "effects"} + + +def new_seed(): + return random.SystemRandom().randrange(0, 2**32) + + +def _result(state, prompt, segments, rolled, seed, registry): + return {"state": state, "prompt": prompt, "segments": segments, "rolled": rolled, + "seed": seed, "theme": theme_entry(state["theme"]["name"], registry)} + + +def roll(state, reels, seed, registry=None): + """Re-draw the requested, unlocked reels and assemble the prompt (RollResult). + + Locked reels keep their exact text; a pinned custom subject/location is + passed to the handler, which rebuilds the wording around it. + """ + registry = registry or get_theme_registry() + state = parse_state(state) + reels = set(reels) & set(REEL_NAMES) + + if state["frozen_prompt"].strip(): + prompt, segments = assemble_prompt(state) + return _result(state, prompt, segments, [], seed, registry) + + rolled = [] + theme = state["theme"] + rng = random.Random(seed) + if registry.handlers and (("theme" in reels and not theme["locked"]) or theme["name"] not in registry.handlers): + theme["name"] = rng.choice(sorted(registry.handlers)) + rolled.append("theme") + + r = state["reels"] + to_roll = [name for name in COMPONENT_KEYS + if name in reels and not r[name]["locked"] and r[name].get("enabled", True)] + if to_roll: + before = copy.deepcopy(state) + try: + handler = registry.handlers[theme["name"]] + with GENERATION_LOCK: + registry.config_manager.set_seed(seed) + components = handler.generate( + custom_subject=r["subject"]["custom"], + custom_location=r["location"]["custom"], + include_environment=r["location"]["enabled"], + include_style=r["style"]["enabled"], + include_effects=r["effects"]["enabled"], + ) + except Exception as e: + logger.error("Mega Prompt V4: theme '%s' failed: %s: %s", theme["name"], type(e).__name__, e) + message = f"Error: {type(e).__name__}: {e}" + return _result(before, message, [{"reel": "error", "text": message}], [], seed, registry) + if not isinstance(components, dict): + components = {} + for name in to_roll: + r[name]["text"] = str(components.get(COMPONENT_KEYS[name]) or "") + + prompt, segments = assemble_prompt(state) + return _result(state, prompt, segments, rolled + to_roll, seed, registry) diff --git a/tests/test_mega_prompt_v4.py b/tests/test_mega_prompt_v4.py index 660db22..9100d19 100644 --- a/tests/test_mega_prompt_v4.py +++ b/tests/test_mega_prompt_v4.py @@ -101,3 +101,134 @@ def test_assemble_frozen_prompt_wins(): mp4 = _mp4() prompt, segs = mp4.assemble_prompt(_state(frozen_prompt="exactly this", lora="x")) assert prompt == "exactly this" and segs == [{"reel": "frozen", "text": "exactly this"}] + + +import random +import types + + +class _Handler: + DISPLAY_NAME = "🧪 Alpha" + + def __init__(self, fail=False): + self.calls, self.fail = [], fail + + def generate(self, **kw): + self.calls.append(kw) + if self.fail: + raise KeyError("anime.characters") + n = len(self.calls) + return {"subject": f"subject {kw['custom_subject'] or n}", "environment": f"env {kw['custom_location'] or n}", + "style": f"style {n}", "effects": f"fx {n}"} + + +def _registry(fail=False): + h = _Handler(fail) + cm = types.SimpleNamespace(random=random.Random(), seeds=[]) + cm.set_seed = cm.seeds.append + return types.SimpleNamespace(handlers={"alpha": h, "beta": _Handler()}, config_manager=cm), h + + +def test_roll_same_seed_same_result_with_real_registry(): + mp4 = _mp4() + s = mp4.default_state() + s["theme"] = {"name": "anime", "locked": True} + a, b = mp4.roll(s, mp4.REEL_NAMES, 42), mp4.roll(s, mp4.REEL_NAMES, 42) + assert a == b and a["prompt"] and a["seed"] == 42 and a["theme"]["name"] == "anime" + assert a["state"]["theme"]["name"] == "anime" and "theme" not in a["rolled"] + assert mp4.roll(s, mp4.REEL_NAMES, 43)["prompt"] != a["prompt"] + + +def test_roll_never_touches_locked_reels_and_only_requested_ones(): + mp4 = _mp4() + reg, h = _registry() + s = mp4.default_state() + s["theme"] = {"name": "alpha", "locked": True} + s["reels"]["subject"].update(text="KEEP ME", locked=True) + s["reels"]["style"]["text"] = "old style" + out = mp4.roll(s, ["subject", "location"], 1, registry=reg) + r = out["state"]["reels"] + assert r["subject"]["text"] == "KEEP ME" and r["style"]["text"] == "old style" + assert r["location"]["text"] == "env 1" and out["rolled"] == ["location"] + assert reg.config_manager.seeds == [1] and h.calls[0]["include_style"] is True + assert s["reels"]["location"]["text"] == "" # input not mutated + + +def test_roll_unlocked_theme_is_random_and_recorded(): + mp4 = _mp4() + reg, _ = _registry() + s = mp4.default_state() + out = mp4.roll(s, mp4.REEL_NAMES, 7, registry=reg) + assert out["state"]["theme"]["name"] in ("alpha", "beta") and "theme" in out["rolled"] + assert out["theme"]["icon"] == "🧪" + + +def test_roll_unknown_theme_falls_back_to_random(): + mp4 = _mp4() + reg, _ = _registry() + s = mp4.default_state() + s["theme"] = {"name": "gone", "locked": True} + out = mp4.roll(s, ["style"], 3, registry=reg) + assert out["state"]["theme"]["name"] in ("alpha", "beta") and out["prompt"] + + +def test_roll_pinned_custom_text_is_passed_to_handler(): + mp4 = _mp4() + reg, h = _registry() + s = mp4.default_state() + s["theme"] = {"name": "alpha", "locked": True} + s["reels"]["subject"]["custom"] = "my elf queen" + out = mp4.roll(s, mp4.REEL_NAMES, 1, registry=reg) + assert h.calls[0]["custom_subject"] == "my elf queen" + assert out["state"]["reels"]["subject"]["text"] == "subject my elf queen" + + +def test_roll_disabled_reel_keeps_text_then_fills_when_enabled(): + mp4 = _mp4() + reg, h = _registry() + s = mp4.default_state() + s["theme"] = {"name": "alpha", "locked": True} + s["reels"]["location"]["enabled"] = False + out = mp4.roll(s, mp4.REEL_NAMES, 1, registry=reg) + assert out["state"]["reels"]["location"]["text"] == "" and "location" not in out["rolled"] + assert h.calls[0]["include_environment"] is False and "env" not in out["prompt"] + s2 = out["state"] + s2["reels"]["location"]["enabled"] = True + assert mp4.roll(s2, ["location"], 2, registry=reg)["state"]["reels"]["location"]["text"] == "env 2" + + +def test_roll_frozen_prompt_skips_handlers(): + mp4 = _mp4() + reg, h = _registry() + s = mp4.default_state() + s["frozen_prompt"] = "fixed" + out = mp4.roll(s, mp4.REEL_NAMES, 1, registry=reg) + assert out["prompt"] == "fixed" and out["rolled"] == [] and h.calls == [] + + +def test_roll_handler_exception_returns_error_prompt(): + mp4 = _mp4() + reg, _ = _registry(fail=True) + s = mp4.default_state() + s["theme"] = {"name": "alpha", "locked": True} + s["reels"]["style"]["text"] = "old" + out = mp4.roll(s, mp4.REEL_NAMES, 1, registry=reg) + assert out["prompt"].startswith("Error: KeyError") and out["segments"][0]["reel"] == "error" + assert out["state"]["reels"]["style"]["text"] == "old" and out["rolled"] == [] + + +def test_roll_uses_the_shared_generation_lock(): + mp4 = _mp4() + + class RecordingLock: + entries = 0 + def __enter__(self): RecordingLock.entries += 1 + def __exit__(self, *a): return False + + reg, _ = _registry() + original, mp4.GENERATION_LOCK = mp4.GENERATION_LOCK, RecordingLock() + try: + mp4.roll(mp4.default_state(), mp4.REEL_NAMES, 1, registry=reg) + assert RecordingLock.entries == 1 + finally: + mp4.GENERATION_LOCK = original