nodes/bradsec-ComfyUI_LLMSetRole
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Step roles via combo control_after_generate; bump 1.3.2

Browse Source 
Replace the custom mode/value/start/end/counter machinery with ComfyUI's
native per-widget control_after_generate on the role combo. Set it to
'increment wrap' to advance one role per generation (wrapping), or
'randomize' for a random role; 'fixed' holds the selection. The frontend
steps the combo, so the role visibly updates on the node, and no internal
index, seed, modulo, or NaN re-run logic is needed. Outputs unchanged:
system_prompt, role_name, resolved_index. Dual-API preserved; V3 passes
control_after_generate inside try/except TypeError for older comfy_api.
This commit is contained in:
BRADSEC
2026-06-28 12:27:19 +10:00
parent 734c8cc160
commit 33bcce9d3e
5 changed files with 67 additions and 232 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+14 -19
View File
@@ -8,7 +8,7 @@ Designed to pair with LLM nodes such as the [stavsap/comfyui-ollama](https://git
## The problem it solves
Multi-stage prompt pipelines reuse the same long system prompts. Pasting that prose into a text box on every workflow is error-prone and hard to keep consistent. This node keeps each role as a Markdown file on disk and exposes them as a dropdown, so a workflow picks a role by name and the exact, version-controlled text flows into the LLM. Modes let you also step or randomize through your roles across a batch.
Multi-stage prompt pipelines reuse the same long system prompts. Pasting that prose into a text box on every workflow is error-prone and hard to keep consistent. This node keeps each role as a Markdown file on disk and exposes them as a dropdown, so a workflow picks a role by name and the exact, version-controlled text flows into the LLM. The role widget's `control_after_generate` lets you step or randomize through your roles across a batch.
## Nodes included
@@ -18,10 +18,7 @@ Multi-stage prompt pipelines reuse the same long system prompts. Pasting that pr
| Input | Type | Default | Description |
|---|---|---|---|
| `mode` | combo | `fixed` | How the role is chosen: `fixed`, `increment`, or `random`. See Selection modes below. |
| `role` | combo | first role | Role used in `fixed` mode. The list is built from `roles/*.md` at startup. Add a `.md` file and restart ComfyUI to add a role. |
| `start` | INT | `0` | Lower bound position (0-based) into the alphabetically sorted role list. Bounds `increment` and `random`. |
| `end` | INT | `-1` | Upper bound position; `-1` means the last role. Out-of-range values are clamped. |
| `role` | combo | first role | The role whose text is output. The list is built from `roles/*.md` at startup (add a `.md` file and restart ComfyUI to add a role). Use the widget's `control_after_generate` to step or randomize across a batch. See Stepping below. |
## Outputs
@@ -29,23 +26,21 @@ Multi-stage prompt pipelines reuse the same long system prompts. Pasting that pr
|---|---|---|
| `system_prompt` | STRING | The full text of the selected role file, ready to wire into an LLM node's `system_prompt` input. |
| `role_name` | STRING | The selected role's display name, e.g. for logging or filename suffixes. |
| `resolved_index` | INT | The list position actually used. Useful for logging or labelling batched runs. |
| `resolved_index` | INT | The selected role's position in the alphabetically sorted list. Useful for logging or labelling batched runs. |
## Selection modes
## Stepping across a batch
Roles are sorted alphabetically by display name; `start`/`end` index into that
list (0-based, `end = -1` means the last). The range is inclusive.
The `role` widget carries ComfyUI's per-widget `control_after_generate` control,
the same one samplers put on `seed`. Use it to choose how the role advances
between generations:
- **fixed**: output the `role` dropdown. `start`/`end` ignored.
- **increment**: output the next role alphabetically, advancing one step per
generation and wrapping from `end` back to `start`. The node steps on its own;
there is no extra widget to set. The counter is per-node and process-local, so
it restarts from `start` when ComfyUI restarts.
- **random**: output a role picked at random within `start..end`, fresh each
generation.
- **fixed**: keep the selected role.
- **increment wrap**: advance to the next role each generation, wrapping from the
last role back to the first. Drive a batch through every role in order.
- **randomize**: pick a random role each generation.
`increment` and `random` re-run the node on every generation (so a downstream LLM
node re-runs too); `fixed` only re-runs when the role or its file text changes.
The frontend steps the combo itself, so the role you see selected on the node is
the role being used. Roles are sorted alphabetically by display name.
## Adding roles
@@ -78,7 +73,7 @@ cd ComfyUI_LLMSetRole
python -m unittest test_set_role -v
```
Expect 25 passing tests (title parsing, filename prettify, label-collision disambiguation, path-traversal rejection, missing-file error, edit hot-reload, bad-encoding tolerance, range bounds/clamp/swap, increment wrap, random determinism, mode selection, shipped-role discovery).
Expect 17 passing tests (title parsing, filename prettify, label-collision disambiguation, path-traversal rejection, missing-file error, edit hot-reload, bad-encoding tolerance, role resolution, shipped-role discovery).
### In ComfyUI
nodes.py
+33 -36
View File
@@ -1,11 +1,12 @@
"""ComfyUI_LLMSetRole nodes.
A role picker: a dropdown of system-prompt roles loaded from roles/*.md, output
as a STRING to wire into an LLM node's system_prompt input. A mode control
selects the role: a fixed dropdown pick, an automatic alphabetical step (one role
per generation), or a fresh random pick, bounded to a start..end slice of the list.
as a STRING to wire into an LLM node's system_prompt input. To step or randomize
roles across a batch, use the role widget's own control_after_generate control
("fixed", "increment wrap", or "randomize"); the frontend advances the combo, so
the node needs no mode, counter, or index logic of its own.
Dual-API: the shared core (role_core.select) is wrapped by a V1 class
Dual-API: the shared core (role_core.resolve) is wrapped by a V1 class
(NODE_CLASS_MAPPINGS, authoritative on the 0.25.0 if/elif loader) and an
import-guarded V3 class (comfy_entrypoint) for builds whose loader prefers the
comfy_api schema. node_id "LLMSetRole" is fixed in both.
@@ -16,21 +17,17 @@ from . import role_core as core
NODE_ID = "LLMSetRole"
DISPLAY_NAME = "Set Role"
CATEGORY = "LLM"
DESCRIPTION = "Pick an LLM role system prompt from roles/*.md (fixed, increment, or random) and output it as a STRING for an LLM node's system_prompt."
DESCRIPTION = "Pick an LLM role system prompt from roles/*.md and output it as a STRING for an LLM node's system_prompt. Use the role widget's control_after_generate to step or randomize roles across a batch."
_ROLE_TIP = ("Role used in 'fixed' mode. The list is built from roles/*.md at "
_ROLE_TIP = ("The role whose Markdown becomes the system prompt. Set this widget's "
"control_after_generate to step roles across a batch: 'fixed' holds it, "
"'increment wrap' advances one role per generation and wraps around, "
"'randomize' picks at random. The list is built from roles/*.md at "
"startup; add a .md file and restart ComfyUI to add a role.")
_MODE_TIP = ("fixed: use the role dropdown. increment: step to the next role "
"alphabetically each generation, wrapping within start..end. random: a "
"fresh random role within start..end each generation. increment and random "
"advance on their own; no extra widget to set.")
_START_TIP = "Lower bound position (0-based) into the alphabetically sorted role list."
_END_TIP = "Upper bound position; -1 means the last role. Out-of-range values are clamped."
# Combo options built once at import. role_core guarantees a non-empty list.
_ROLE_OPTIONS = core.role_labels()
_DEFAULT_ROLE = _ROLE_OPTIONS[0]
_MAX_INT = 0xFFFFFFFFFFFFFFFF
# -- V1 node -------------------------------------------------------------------
@@ -47,20 +44,18 @@ class LLMSetRoleNode:
def INPUT_TYPES(cls):
return {
"required": {
"mode": (list(core.MODES), {"default": "fixed", "tooltip": _MODE_TIP}),
"role": (_ROLE_OPTIONS, {"default": _DEFAULT_ROLE, "tooltip": _ROLE_TIP}),
"start": ("INT", {"default": 0, "min": 0, "max": _MAX_INT, "tooltip": _START_TIP}),
"end": ("INT", {"default": -1, "min": -1, "max": _MAX_INT, "tooltip": _END_TIP}),
},
"hidden": {"unique_id": "UNIQUE_ID"},
"role": (_ROLE_OPTIONS, {"default": _DEFAULT_ROLE,
"control_after_generate": True,
"tooltip": _ROLE_TIP}),
}
}
def set_role(self, mode, role, start, end, unique_id=None):
return core.select(mode, role, start, end, unique_id)
def set_role(self, role):
return core.resolve(role)
@classmethod
def IS_CHANGED(cls, mode, role, start, end, unique_id=None):
return core.select_fingerprint(mode, role, start, end)
def IS_CHANGED(cls, role):
return core.role_fingerprint(role)
NODE_CLASS_MAPPINGS = {NODE_ID: LLMSetRoleNode}
@@ -72,6 +67,16 @@ NODE_DISPLAY_NAME_MAPPINGS = {NODE_ID: DISPLAY_NAME}
try:
from comfy_api.v0_0_2 import io, ComfyExtension
def _role_input():
"""io.Combo.Input for the role, tolerating comfy_api builds whose Combo
does not accept control_after_generate (older schema versions)."""
try:
return io.Combo.Input("role", options=_ROLE_OPTIONS, default=_DEFAULT_ROLE,
control_after_generate=True, tooltip=_ROLE_TIP)
except TypeError:
return io.Combo.Input("role", options=_ROLE_OPTIONS, default=_DEFAULT_ROLE,
tooltip=_ROLE_TIP)
class LLMSetRoleV3(io.ComfyNode):
@classmethod
def define_schema(cls) -> io.Schema:
@@ -80,15 +85,7 @@ try:
display_name=DISPLAY_NAME,
category=CATEGORY,
description=DESCRIPTION,
inputs=[
io.Combo.Input("mode", options=list(core.MODES),
default="fixed", tooltip=_MODE_TIP),
io.Combo.Input("role", options=_ROLE_OPTIONS,
default=_DEFAULT_ROLE, tooltip=_ROLE_TIP),
io.Int.Input("start", default=0, min=0, max=_MAX_INT, tooltip=_START_TIP),
io.Int.Input("end", default=-1, min=-1, max=_MAX_INT, tooltip=_END_TIP),
],
hidden=[io.Hidden.unique_id],
inputs=[_role_input()],
outputs=[
io.String.Output(id="system_prompt", display_name="system_prompt"),
io.String.Output(id="role_name", display_name="role_name"),
@@ -97,12 +94,12 @@ try:
)
@classmethod
def fingerprint_inputs(cls, mode, role, start, end):
return core.select_fingerprint(mode, role, start, end)
def fingerprint_inputs(cls, role):
return core.role_fingerprint(role)
@classmethod
def execute(cls, mode, role, start, end) -> io.NodeOutput:
body, title, pos = core.select(mode, role, start, end, cls.hidden.unique_id)
def execute(cls, role) -> io.NodeOutput:
body, title, pos = core.resolve(role)
return io.NodeOutput(body, title, pos)
class LLMSetRoleExtension(ComfyExtension):
pyproject.toml
+2 -2
View File
@@ -1,7 +1,7 @@
[project]
name = "comfyui_llmsetrole"
description = "Pick an LLM role (system prompt) from a dropdown and output it as a STRING for an LLM node's system_prompt. Modes: fixed (use the dropdown), increment (step to the next role alphabetically each generation), or random (a fresh pick each generation), over a start..end slice. Roles are plain Markdown files in roles/; add a .md file and restart to add a role. Node name: Set Role. Located under Add Node > LLM."
version = "1.3.1"
description = "Pick an LLM role (system prompt) from a dropdown and output it as a STRING for an LLM node's system_prompt. Use the role widget's control_after_generate to step (increment wrap) or randomize roles across a batch. Roles are plain Markdown files in roles/; add a .md file and restart to add a role. Node name: Set Role. Located under Add Node > LLM."
version = "1.3.2"
license = { file = "LICENSE" }
requires-python = ">=3.10"
dependencies = []
role_core.py
+9 -115
View File
@@ -17,12 +17,9 @@ restart ComfyUI, since the dropdown options are built once at import.
"""
import os
import random
import re
from functools import lru_cache
MODES = ("fixed", "increment", "random")
ROLES_DIR = os.path.join(os.path.dirname(__file__), "roles")
NO_ROLES_LABEL = "(no roles found)"
@@ -157,24 +154,11 @@ def role_fingerprint(label):
return str(label)
# -- Position-based selection (fixed / increment / random) ---------------------
def _clamp(v, lo, hi):
return lo if v < lo else hi if v > hi else v
def _bounds(start, end, n):
"""Return (lo, hi, span) inclusive position bounds within a list of size n.
end < 0 means "last". Out-of-range values are clamped; reversed bounds are
swapped. Assumes n >= 1 (callers handle the empty case first).
"""
lo = _clamp(start, 0, n - 1)
hi = _clamp(end if end >= 0 else n - 1, 0, n - 1)
if lo > hi:
lo, hi = hi, lo
return lo, hi, hi - lo + 1
# -- Selection -----------------------------------------------------------------
# Stepping (one-by-one, random) is handled by the frontend's per-widget
# control_after_generate on the `role` combo ("increment wrap" / "randomize"),
# so the core only resolves the currently-selected label. No mode, counter, or
# index math lives here.
def label_index(label):
"""Position of a dropdown label in the sorted role list, or 0 if absent."""
@@ -184,97 +168,7 @@ def label_index(label):
return 0
def pick_index(mode, value, start, end, n):
"""Resolve a role position for increment/random modes from an integer `value`.
Pure helper used by select() and the tests:
increment: lo + (value mod span), wrapping at the end back to the start.
random: lo + Random(value).randrange(span), deterministic per value.
Raises ValueError if there are no roles.
"""
if n < 1:
raise ValueError(
f"No role files found in {ROLES_DIR}. Add a .md file and restart ComfyUI."
)
lo, _, span = _bounds(start, end, n)
if mode == "random":
return lo + random.Random(value).randrange(span)
# increment (and any unknown mode falls back to deterministic stepping)
return lo + (value % span)
# -- Increment state -----------------------------------------------------------
# Per-node counters for increment mode, keyed by the node's unique_id. Process-
# local: stepping restarts from the first role when ComfyUI restarts. select()
# advances the counter, and the node forces a re-run each generation
# (select_fingerprint returns NaN) so the walk progresses instead of caching.
_increment_counters = {}
def reset_counters():
"""Clear all increment state. Used by tests; harmless at runtime."""
_increment_counters.clear()
def _next_increment_pos(node_key, start, end, n):
"""Advance this node's counter and return the next position within bounds."""
c = _increment_counters.get(node_key, 0)
_increment_counters[node_key] = c + 1
return pick_index("increment", c, start, end, n)
def _random_pos(start, end, n):
"""A fresh random position within bounds, non-reproducible by design."""
lo, _, span = _bounds(start, end, n)
return lo + random.randrange(span)
def role_at(pos):
"""(label, filename) at a position. Raises ValueError if no roles."""
roles = list_roles()
if not roles:
raise ValueError(
f"No role files found in {ROLES_DIR}. Add a .md file and restart ComfyUI."
)
return roles[pos]
def select(mode, role, start, end, node_key=None):
"""Resolve (system_prompt, role_name, resolved_index) for any mode.
fixed: use the dropdown `role`. increment: the next role alphabetically,
advancing this node's counter and wrapping within [start, end]. random: a
fresh random role within [start, end], picked anew each call.
"""
if mode == "fixed":
pos = label_index(role)
body, title = apply_role(role)
return body, title, pos
n = len(list_roles())
if n < 1:
raise ValueError(
f"No role files found in {ROLES_DIR}. Add a .md file and restart ComfyUI."
)
pos = _random_pos(start, end, n) if mode == "random" \
else _next_increment_pos(node_key, start, end, n)
_, fn = role_at(pos)
title, body = load_role_text(fn)
return body, title, pos
def select_fingerprint(mode, role, start, end):
"""Cache identity.
fixed caches on the resolved file's mtime so the node only re-runs when the
role or its text changes. increment/random return NaN, which never equals
itself, so ComfyUI re-runs the node every generation and the walk advances.
"""
if mode != "fixed":
return float("nan")
try:
pos = label_index(role)
_, fn = role_at(pos)
mtime = os.path.getmtime(os.path.join(ROLES_DIR, fn))
return f"fixed:{pos}:{fn}:{mtime}"
except (ValueError, FileNotFoundError, OSError):
return f"fixed:{role}"
def resolve(role):
"""Resolve (system_prompt, role_name, resolved_index) for the selected role."""
body, title = apply_role(role)
return body, title, label_index(role)
test_set_role.py
+9 -60
View File
@@ -127,80 +127,29 @@ class TestBadEncoding(unittest.TestCase):
self.assertIn("body", body)
class TestBounds(unittest.TestCase):
def test_full_range_with_end_minus_one(self):
self.assertEqual(core._bounds(0, -1, 3), (0, 2, 3))
def test_clamp_out_of_range(self):
self.assertEqual(core._bounds(-5, 99, 3), (0, 2, 3))
def test_subrange(self):
self.assertEqual(core._bounds(1, 2, 4), (1, 2, 2))
def test_swapped_bounds(self):
self.assertEqual(core._bounds(2, 0, 3), (0, 2, 3))
class TestPickIndex(unittest.TestCase):
def test_increment_wraps(self):
got = [core.pick_index("increment", i, 0, -1, 3) for i in range(5)]
self.assertEqual(got, [0, 1, 2, 0, 1])
def test_increment_within_subrange(self):
got = [core.pick_index("increment", i, 1, 2, 3) for i in range(4)]
self.assertEqual(got, [1, 2, 1, 2])
def test_random_deterministic_per_value(self):
a = core.pick_index("random", 42, 0, -1, 3)
b = core.pick_index("random", 42, 0, -1, 3)
self.assertEqual(a, b)
self.assertTrue(0 <= a <= 2)
def test_random_stays_in_subrange(self):
for v in range(50):
p = core.pick_index("random", v, 1, 2, 3)
self.assertIn(p, (1, 2))
def test_no_roles_raises(self):
with self.assertRaises(ValueError):
core.pick_index("increment", 0, 0, -1, 0)
class TestSelect(unittest.TestCase):
class TestResolve(unittest.TestCase):
def _make(self, d, names):
for n in names:
with open(os.path.join(d, n), "w", encoding="utf-8") as f:
f.write(f"<!-- title: {n[:-3].upper()} -->\nbody-{n}")
def test_select_modes(self):
def test_resolve_returns_body_title_index(self):
with tempfile.TemporaryDirectory() as d:
self._make(d, ("a.md", "b.md", "c.md")) # titles A, B, C -> positions 0,1,2
orig = core.ROLES_DIR
core.ROLES_DIR = d
core._load_cached.cache_clear()
core.reset_counters()
try:
body, name, pos = core.select("fixed", "B", 0, -1)
self.assertEqual((name, pos, body), ("B", 1, "body-b.md"))
# increment walks alphabetically, advancing once per call, wrapping
core.reset_counters()
walk = [core.select("increment", "A", 0, -1, "node1")[1] for _ in range(4)]
self.assertEqual(walk, ["A", "B", "C", "A"])
# distinct node ids keep independent counters
core.reset_counters()
self.assertEqual(core.select("increment", "A", 0, -1, "x")[1], "A")
self.assertEqual(core.select("increment", "A", 0, -1, "y")[1], "A")
self.assertEqual(core.select("increment", "A", 0, -1, "x")[1], "B")
_, name_r, pos_r = core.select("random", "A", 0, -1)
self.assertIn(pos_r, (0, 1, 2))
self.assertIn(name_r, ("A", "B", "C"))
self.assertEqual(core.resolve("A"), ("body-a.md", "A", 0))
self.assertEqual(core.resolve("B"), ("body-b.md", "B", 1))
self.assertEqual(core.resolve("C"), ("body-c.md", "C", 2))
finally:
core.ROLES_DIR = orig
core._load_cached.cache_clear()
core.reset_counters()
def test_resolve_unknown_role_raises(self):
with self.assertRaises(ValueError):
core.resolve("Definitely Not A Real Role")
class TestShippedRoles(unittest.TestCase):