diff --git a/README.md b/README.md index 25e9e43..3eefeba 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/nodes.py b/nodes.py index 1e24d97..df6b8b1 100644 --- a/nodes.py +++ b/nodes.py @@ -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): diff --git a/pyproject.toml b/pyproject.toml index e22a308..9456507 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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 = [] diff --git a/role_core.py b/role_core.py index 6a3e597..f491e58 100644 --- a/role_core.py +++ b/role_core.py @@ -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) diff --git a/test_set_role.py b/test_set_role.py index acbf7ae..3ee82a6 100644 --- a/test_set_role.py +++ b/test_set_role.py @@ -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"\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):