feat(megaprompt-v4): pure roll() with locked/pinned/disabled reels and frozen prompt
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Isulion
co-authored by
Claude Opus 5.5
Claude Opus 5.5
parent
4bc54702c3
commit
470f55dec4
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user