commit b868dd4efdaebfdf851545093facda8931131043 Author: BRADSEC <7948876+bradsec@users.noreply.github.com> Date: Sat Jun 27 16:59:01 2026 +1000 Initial commit diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml new file mode 100644 index 0000000..edd08c5 --- /dev/null +++ b/.github/workflows/publish.yml @@ -0,0 +1,26 @@ +name: Publish to Comfy registry +on: + workflow_dispatch: + push: + branches: + - main + - master + paths: + - "pyproject.toml" + +permissions: + issues: write + +jobs: + publish-node: + name: Publish Custom Node to registry + runs-on: ubuntu-latest + if: ${{ github.repository_owner == 'bradsec' }} + steps: + - name: Check out code + uses: actions/checkout@v4 + - name: Publish Custom Node + uses: Comfy-Org/publish-node-action@7578cdb0c1410f2515e6f9b1443c6859cab94f05 # v1.0.1 + with: + ## Add your own personal access token to your Github Repository secrets and reference it here. + personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }} diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..8c49d77 --- /dev/null +++ b/.gitignore @@ -0,0 +1,24 @@ +# Python bytecode and caches +__pycache__/ +*.py[cod] +*.egg-info/ +.pytest_cache/ + +# Environments +.venv/ +venv/ + +# Editor / OS +.vscode/ +.idea/ +.DS_Store + +# AI tooling and planning artifacts (do not ship to the registry/repo) +docs/superpowers/ +.claude/ +CLAUDE.md +*.plan.md + +# Private role prose - never commit (user roles stay local) +roles/*.md +!roles/.gitkeep diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..59b9c67 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 BRADSEC + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. \ No newline at end of file diff --git a/README.md b/README.md new file mode 100644 index 0000000..d87705a --- /dev/null +++ b/README.md @@ -0,0 +1,117 @@ +# LLM Set Role - Custom node for ComfyUI + +A single node that presents a dropdown of LLM **roles** (system prompts) and outputs the selected role's text as a STRING. Wire that STRING into the `system_prompt` input of an LLM node (Ollama, etc.) to give the model a consistent role: a style transform, a summarizer, a translator, a captioner, an editor, anything you can write as a system prompt. + +The node will be located under **Add Node > LLM**. Node name: **Set Role**. + +Designed to pair with LLM nodes such as the [stavsap/comfyui-ollama](https://github.com/stavsap/comfyui-ollama) pack: this node supplies the system prompt, the LLM node does the generation. + +## 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. + +## Nodes included + +- **Set Role** - dropdown of roles from `roles/*.md`, outputs the selected role text, its name, and the list position used. + +## Inputs + +| 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. | +| `index` | INT | `0` | Position used in `increment` mode. Set its `control_after_generate` to **increment** to advance one role per generation. Wraps within `start..end`. | +| `seed` | INT | `0` | Seed used in `random` mode. Set `control_after_generate` to **randomize** for a new pick each generation; the same seed reproduces the same role. | +| `start` | INT | `0` | Lower bound position (0-based) into the alphabetically sorted role list. | +| `end` | INT | `-1` | Upper bound position; `-1` means the last role. Out-of-range values are clamped. | + +## Outputs + +| Output | Type | Description | +|---|---|---| +| `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. | + +## Selection modes + +Roles are sorted alphabetically by display name; `start`/`end` index into that +list (0-based, `end = -1` means the last). The range is inclusive. + +- **fixed**: output the `role` dropdown. `index`/`seed`/`start`/`end` ignored. +- **increment**: output the role at `start + (index mod span)`, wrapping back to + `start` after `end`. Set the `index` widget's `control_after_generate` to + **increment** so each generation advances to the next role. +- **random**: output a role picked from `start..end` using `seed`. The same seed + always yields the same role; set `control_after_generate` to **randomize** for a + fresh pick per run. + +Note: `control_after_generate` is the same per-widget control ComfyUI puts on the +`seed` widget of samplers. Increment/random still work without it by changing +`index`/`seed` yourself or by using ComfyUI's batch count. + +## Adding roles + +1. Drop a new `.md` file into the `roles/` directory. The file content is the system prompt, kept verbatim. See `roles/example.md.example` for the format. +2. The dropdown label defaults to the prettified filename: `summarizer.md` becomes `Summarizer`. +3. To set an explicit label, make the **first line** of the file an HTML comment: + + ```markdown + + ``` + + That line is stripped from the output and is used only as the dropdown label. (An HTML comment is used rather than a `#` heading because role files commonly start with their own `#` heading.) +4. Restart ComfyUI so the new file appears in the dropdown. Editing an existing role does not need a restart: the node re-runs when the file changes. + +Your `roles/*.md` files are gitignored, so your prompts stay private and local. The shipped `roles/example.md.example` is tracked only as a format reference and is not loaded by the node. + +## Usage + +``` +[Set Role].system_prompt -> [Ollama Generate].system_prompt +your prompt text -> [Ollama Generate].prompt +``` + +## Testing + +### Unit tests (no ComfyUI needed) + +```bash +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). + +### In ComfyUI + +1. Restart ComfyUI so the node loads. Confirm it appears once under **Add Node > LLM** as **Set Role**, with no import error in `user/comfyui.log`. +2. Add the node, pick a role and mode, and wire `system_prompt` into your LLM node. +3. Run the workflow. The LLM node receives the selected role text. + +## Installation + +### ComfyUI Manager (recommended) + +In ComfyUI, open **Manager > Custom Nodes Manager**, search for **ComfyUI_LLMSetRole**, click **Install**, then restart ComfyUI. + +### Manual install + +Clone into your ComfyUI `custom_nodes` directory and restart ComfyUI: + +```bash +cd ComfyUI/custom_nodes +git clone https://github.com/bradsec/ComfyUI_LLMSetRole +``` + +No extra dependencies (Python stdlib only). + +## Compatibility + +Dual API: exports the legacy V1 `NODE_CLASS_MAPPINGS` and a V3 `comfy_entrypoint` (`comfy_api`). On builds whose loader reads `NODE_CLASS_MAPPINGS` first, the V1 path is used and `comfy_entrypoint` is skipped, so the node is never registered twice. + +--- + +## Support + +If you find this useful, please consider [starring the repo](https://github.com/bradsec/ComfyUI_LLMSetRole). Stars help other people discover these nodes. diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..943997d --- /dev/null +++ b/__init__.py @@ -0,0 +1,15 @@ +"""ComfyUI_LLMSetRole: dropdown of system-prompt roles for LLM nodes. + +Exports V1 NODE_CLASS_MAPPINGS (authoritative on the 0.25.0 if/elif loader) and +comfy_entrypoint for builds whose loader prefers the V3 schema API. +""" + +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] + +try: + from .nodes import comfy_entrypoint # noqa: F401 + __all__.append("comfy_entrypoint") +except ImportError: + pass diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..aa7bcdb --- /dev/null +++ b/nodes.py @@ -0,0 +1,134 @@ +"""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 by fixed dropdown, by an incrementing index, or randomly by +seed, bounded to a start..end slice of the list. + +Dual-API: the shared core (role_core.select) 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. +""" + +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." + +_ROLE_TIP = ("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.") +_MODE_TIP = ("fixed: use the role dropdown. increment: walk the list by 'index', " + "wrapping within start..end. random: pick within start..end from 'seed'.") +_INDEX_TIP = ("Position used in 'increment' mode. Set control_after_generate to " + "increment to advance one role per generation. Wraps within start..end.") +_SEED_TIP = ("Seed used in 'random' mode. Set control_after_generate to randomize for " + "a new pick per generation; the same seed reproduces the same role.") +_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 ------------------------------------------------------------------- + +class LLMSetRoleNode: + CATEGORY = CATEGORY + FUNCTION = "set_role" + OUTPUT_NODE = False + DESCRIPTION = DESCRIPTION + RETURN_TYPES = ("STRING", "STRING", "INT") + RETURN_NAMES = ("system_prompt", "role_name", "resolved_index") + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "mode": (list(core.MODES), {"default": "fixed", "tooltip": _MODE_TIP}), + "role": (_ROLE_OPTIONS, {"default": _DEFAULT_ROLE, "tooltip": _ROLE_TIP}), + "index": ("INT", {"default": 0, "min": 0, "max": _MAX_INT, + "control_after_generate": True, "tooltip": _INDEX_TIP}), + "seed": ("INT", {"default": 0, "min": 0, "max": _MAX_INT, + "control_after_generate": True, "tooltip": _SEED_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}), + } + } + + def set_role(self, mode, role, index, seed, start, end): + return core.select(mode, role, index, seed, start, end) + + @classmethod + def IS_CHANGED(cls, mode, role, index, seed, start, end): + return core.select_fingerprint(mode, role, index, seed, start, end) + + +NODE_CLASS_MAPPINGS = {NODE_ID: LLMSetRoleNode} +NODE_DISPLAY_NAME_MAPPINGS = {NODE_ID: DISPLAY_NAME} + + +# -- V3 node (import-guarded comfy_api schema) --------------------------------- + +try: + from comfy_api.v0_0_2 import io, ComfyExtension + + def _int_input(name, *, control=False, **kw): + """io.Int.Input, tolerating comfy_api builds without control_after_generate.""" + if control: + try: + return io.Int.Input(name, control_after_generate=True, **kw) + except TypeError: + pass + return io.Int.Input(name, **kw) + + class LLMSetRoleV3(io.ComfyNode): + @classmethod + def define_schema(cls) -> io.Schema: + return io.Schema( + node_id=NODE_ID, + 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), + _int_input("index", control=True, default=0, min=0, max=_MAX_INT, + tooltip=_INDEX_TIP), + _int_input("seed", control=True, default=0, min=0, max=_MAX_INT, + tooltip=_SEED_TIP), + _int_input("start", default=0, min=0, max=_MAX_INT, tooltip=_START_TIP), + _int_input("end", default=-1, min=-1, max=_MAX_INT, tooltip=_END_TIP), + ], + outputs=[ + io.String.Output(id="system_prompt", display_name="system_prompt"), + io.String.Output(id="role_name", display_name="role_name"), + io.Int.Output(id="resolved_index", display_name="resolved_index"), + ], + ) + + @classmethod + def fingerprint_inputs(cls, mode, role, index, seed, start, end): + return core.select_fingerprint(mode, role, index, seed, start, end) + + @classmethod + def execute(cls, mode, role, index, seed, start, end) -> io.NodeOutput: + body, title, pos = core.select(mode, role, index, seed, start, end) + return io.NodeOutput(body, title, pos) + + class LLMSetRoleExtension(ComfyExtension): + async def get_node_list(self): + return [LLMSetRoleV3] + + async def comfy_entrypoint() -> "LLMSetRoleExtension": + return LLMSetRoleExtension() + +except ImportError: + # Older ComfyUI without comfy_api: V1 mappings above are the only path. + pass diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..702f286 --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,16 @@ +[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, increment, or random over a start..end slice, with index/seed control_after_generate. 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.0.0" +license = { file = "LICENSE" } +requires-python = ">=3.10" +dependencies = [] + +[project.urls] +Repository = "https://github.com/bradsec/ComfyUI_LLMSetRole" +# Used by Comfy Registry https://comfyregistry.org + +[tool.comfy] +PublisherId = "bradsec" +DisplayName = "ComfyUI_LLMSetRole" +Icon = "" diff --git a/role_core.py b/role_core.py new file mode 100644 index 0000000..56f28d0 --- /dev/null +++ b/role_core.py @@ -0,0 +1,239 @@ +"""Shared core for ComfyUI_LLMSetRole. + +Pure functions, no ComfyUI import. Scans roles/*.md, resolves dropdown labels, +and loads role text. Both API wrappers (V1 and V3) call into here so the logic +cannot drift. + +Role files are plain Markdown kept verbatim as the LLM system prompt. A file +may declare its dropdown label on its first line as an HTML comment: + + + +When present that line is stripped from the emitted text. Without it the label +is derived from the filename (professional_photo.md -> "Professional Photo"). + +Edit a role and the node re-runs (mtime fingerprint). Add a new role file and +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)" + +# First-line title marker: +_TITLE_RE = re.compile(r"^\s*\s*$", re.IGNORECASE) + + +def prettify_filename(filename): + """professional_photo.md -> 'Professional Photo'.""" + stem = os.path.splitext(filename)[0] + return stem.replace("_", " ").replace("-", " ").strip().title() + + +def parse_title(text, filename): + """Return (title, body). + + If the first line is an HTML title comment, use it as the title and drop + that line from the body. Otherwise the title is the prettified filename and + the body is the text unchanged. + """ + lines = text.split("\n", 1) + m = _TITLE_RE.match(lines[0]) if lines else None + if m: + title = m.group(1).strip() + body = lines[1] if len(lines) > 1 else "" + return title, body + return prettify_filename(filename), text + + +def _validate_name(filename): + """Reject anything but a bare basename, so a label can never escape ROLES_DIR.""" + if os.path.sep in filename or (os.path.altsep and os.path.altsep in filename) \ + or filename in (os.curdir, os.pardir): + raise ValueError(f"Invalid role filename: {filename!r}") + + +@lru_cache(maxsize=None) +def _load_cached(filename, mtime): + """Read and parse one role file, keyed on (filename, mtime). + + Including mtime in the key means editing a file invalidates the cached + entry on the next access, so edits are picked up without a restart. + errors="replace": one mis-encoded role file must not crash the import-time + dropdown build and take the whole node down; it degrades to replacement chars. + """ + with open(os.path.join(ROLES_DIR, filename), encoding="utf-8", errors="replace") as f: + return parse_title(f.read(), filename) + + +def load_role_text(filename): + """Return (title, body) for one role file. + + Cached on the file's mtime: editing a role is picked up on the next run + without restarting ComfyUI; only adding or removing files needs a restart. + """ + _validate_name(filename) + path = os.path.join(ROLES_DIR, filename) + try: + mtime = os.path.getmtime(path) + except OSError: + raise FileNotFoundError( + f"Role file not found: {filename!r} in {ROLES_DIR}. " + f"It may have been removed; restart ComfyUI to refresh the list." + ) + return _load_cached(filename, mtime) + + +def _scan_files(): + """Sorted list of *.md basenames in ROLES_DIR, or [] if none/missing.""" + if not os.path.isdir(ROLES_DIR): + return [] + return sorted( + e.name for e in os.scandir(ROLES_DIR) + if e.is_file() and e.name.lower().endswith(".md") + ) + + +def list_roles(): + """Return [(label, filename)] sorted by label. + + Labels that collide after prettifying are disambiguated by appending the + filename, so every label maps to exactly one file. + """ + loaded = [(fn, load_role_text(fn)[0]) for fn in _scan_files()] + counts = {} + for _, title in loaded: + counts[title] = counts.get(title, 0) + 1 + + pairs = [(f"{title} ({fn})" if counts[title] > 1 else title, fn) + for fn, title in loaded] + pairs.sort(key=lambda p: p[0].lower()) + return pairs + + +def role_labels(): + """Combo option list. Never empty: falls back to a sentinel label.""" + labels = [label for label, _ in list_roles()] + return labels or [NO_ROLES_LABEL] + + +def resolve_filename(label): + """Map a dropdown label back to its role filename. + + Raises ValueError if the label is the no-roles sentinel or is unknown, so a + crafted label can never reference a file outside ROLES_DIR. + """ + if label == NO_ROLES_LABEL: + raise ValueError( + f"No role files found in {ROLES_DIR}. Add a .md file and restart ComfyUI." + ) + for lbl, fn in list_roles(): + if lbl == label: + return fn + raise ValueError(f"Unknown role: {label!r}. Restart ComfyUI to refresh the list.") + + +def apply_role(label): + """Resolve a role label to (system_prompt, role_name).""" + fn = resolve_filename(label) + title, body = load_role_text(fn) + return body, title + + +def role_fingerprint(label): + """Identity for caching: filename + mtime, or the label itself if unresolved.""" + try: + fn = resolve_filename(label) + mtime = os.path.getmtime(os.path.join(ROLES_DIR, fn)) + return f"{fn}:{mtime}" + except (ValueError, FileNotFoundError, OSError): + 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 + + +def label_index(label): + """Position of a dropdown label in the sorted role list, or 0 if absent.""" + for i, (lbl, _) in enumerate(list_roles()): + if lbl == label: + return i + return 0 + + +def pick_index(mode, index, seed, start, end, n): + """Resolve a role position for increment/random modes. + + increment: lo + (index mod span), wrapping at the end back to the start. + random: lo + Random(seed).randrange(span), deterministic per seed. + 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(seed).randrange(span) + # increment (and any unknown mode falls back to deterministic stepping) + return lo + (index % 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, index, seed, start, end): + """Resolve (system_prompt, role_name, resolved_index) for any mode. + + fixed: use the dropdown `role`. increment/random: compute the position from + index/seed within [start, end]. + """ + if mode == "fixed": + pos = label_index(role) + body, title = apply_role(role) + return body, title, pos + pos = pick_index(mode, index, seed, start, end, len(list_roles())) + _, fn = role_at(pos) + title, body = load_role_text(fn) + return body, title, pos + + +def select_fingerprint(mode, role, index, seed, start, end): + """Cache identity covering inputs and the resolved file's mtime.""" + try: + _, _, pos = select(mode, role, index, seed, start, end) + _, fn = role_at(pos) + mtime = os.path.getmtime(os.path.join(ROLES_DIR, fn)) + return f"{mode}:{pos}:{fn}:{mtime}" + except (ValueError, FileNotFoundError, OSError): + return f"{mode}:{role}:{index}:{seed}:{start}:{end}" diff --git a/roles/.gitkeep b/roles/.gitkeep new file mode 100644 index 0000000..e69de29 diff --git a/roles/example.md.example b/roles/example.md.example new file mode 100644 index 0000000..b49a350 --- /dev/null +++ b/roles/example.md.example @@ -0,0 +1,20 @@ + +You are a precise summarizer. Given any input text, return a faithful summary +in 3 to 5 sentences. Preserve key facts, names, numbers, and conclusions. Do +not add opinions, preamble, or closing remarks. Output only the summary. + +# How to use this file + +This is an EXAMPLE role, shown so you can see the format. It is named +`example.md.example` so the node does NOT load it (the node only reads `*.md`). + +To create a real role: +1. Copy this file to `roles/.md` (note: `.md`, not `.md.example`). +2. Replace this content with your system prompt. The whole file becomes the + prompt text sent to the LLM as its system role. +3. Optionally set the dropdown label with an HTML comment on the FIRST line: + ``. That line is stripped from the output. Without + it the label is the prettified filename (`my_role.md` -> "My Role"). +4. Restart ComfyUI so the new role appears in the dropdown. + +Your `roles/*.md` files are gitignored, so your prompts stay private and local. diff --git a/test_set_role.py b/test_set_role.py new file mode 100644 index 0000000..6fc01e7 --- /dev/null +++ b/test_set_role.py @@ -0,0 +1,214 @@ +"""Unit tests for the LLMSetRole shared core. + +Pure logic, no ComfyUI required. +Run: python -m unittest test_set_role -v +""" + +import os +import tempfile +import unittest + +import role_core as core + + +class TestParseTitle(unittest.TestCase): + def test_html_title_extracted_and_stripped(self): + text = "\n# RULE\nbody line\n" + title, body = core.parse_title(text, "whatever.md") + self.assertEqual(title, "Summarizer") + self.assertEqual(body, "# RULE\nbody line\n") + + def test_title_case_insensitive_marker(self): + title, body = core.parse_title("\nx", "f.md") + self.assertEqual(title, "Foo Bar") + self.assertEqual(body, "x") + + def test_no_title_uses_prettified_filename(self): + text = "# CRITICAL OUTPUT RULE\nprose" + title, body = core.parse_title(text, "professional_photo.md") + self.assertEqual(title, "Professional Photo") + self.assertEqual(body, text) + + def test_heading_is_not_treated_as_title(self): + # Files that begin with a '#' heading must not turn it into the title. + title, _ = core.parse_title("# CRITICAL OUTPUT RULE\nx", "selfie_photo.md") + self.assertEqual(title, "Selfie Photo") + + +class TestPrettify(unittest.TestCase): + def test_underscores_and_hyphens(self): + self.assertEqual(core.prettify_filename("professional_photo.md"), "Professional Photo") + self.assertEqual(core.prettify_filename("film-noir.md"), "Film Noir") + + +class TestResolveSecurity(unittest.TestCase): + def test_reject_path_separator(self): + with self.assertRaises(ValueError): + core.load_role_text("../secret.md") + + def test_reject_nested_path(self): + with self.assertRaises(ValueError): + core.load_role_text("sub/dir.md") + + def test_unknown_label_raises(self): + with self.assertRaises(ValueError): + core.resolve_filename("Definitely Not A Real Role") + + def test_no_roles_sentinel_raises(self): + with self.assertRaises(ValueError): + core.resolve_filename(core.NO_ROLES_LABEL) + + +class TestMissingFile(unittest.TestCase): + def test_load_missing_raises(self): + core._load_cached.cache_clear() + with self.assertRaises(FileNotFoundError): + core.load_role_text("does_not_exist_12345.md") + + +class TestCollisionDisambiguation(unittest.TestCase): + def test_same_title_gets_filename_suffix(self): + with tempfile.TemporaryDirectory() as d: + for name in ("a.md", "b.md"): + with open(os.path.join(d, name), "w", encoding="utf-8") as f: + f.write("\nbody") + orig_dir = core.ROLES_DIR + core.ROLES_DIR = d + core._load_cached.cache_clear() + try: + labels = [lbl for lbl, _ in core.list_roles()] + finally: + core.ROLES_DIR = orig_dir + core._load_cached.cache_clear() + self.assertEqual(labels, ["Same (a.md)", "Same (b.md)"]) + + +class TestHotReload(unittest.TestCase): + def test_edit_picked_up_via_mtime(self): + # Regression: a content edit must not return a stale cached body. + with tempfile.TemporaryDirectory() as d: + path = os.path.join(d, "s.md") + orig_dir = core.ROLES_DIR + core.ROLES_DIR = d + core._load_cached.cache_clear() + try: + with open(path, "w", encoding="utf-8") as f: + f.write("\nfirst") + os.utime(path, (1000, 1000)) + _, body1 = core.load_role_text("s.md") + + with open(path, "w", encoding="utf-8") as f: + f.write("\nsecond") + os.utime(path, (2000, 2000)) + _, body2 = core.load_role_text("s.md") + finally: + core.ROLES_DIR = orig_dir + core._load_cached.cache_clear() + self.assertEqual(body1, "first") + self.assertEqual(body2, "second") + + +class TestBadEncoding(unittest.TestCase): + def test_non_utf8_file_does_not_crash(self): + # A mis-encoded file must not crash the dropdown build. + with tempfile.TemporaryDirectory() as d: + with open(os.path.join(d, "bad.md"), "wb") as f: + f.write(b"\n\xff\xfe body") + orig_dir = core.ROLES_DIR + core.ROLES_DIR = d + core._load_cached.cache_clear() + try: + labels = [lbl for lbl, _ in core.list_roles()] + body, _ = core.apply_role("Bad") + finally: + core.ROLES_DIR = orig_dir + core._load_cached.cache_clear() + self.assertEqual(labels, ["Bad"]) + 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, 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, 0, 1, 2, 3) for i in range(4)] + self.assertEqual(got, [1, 2, 1, 2]) + + def test_random_deterministic_per_seed(self): + a = core.pick_index("random", 0, 42, 0, -1, 3) + b = core.pick_index("random", 0, 42, 0, -1, 3) + self.assertEqual(a, b) + self.assertTrue(0 <= a <= 2) + + def test_random_stays_in_subrange(self): + for s in range(50): + p = core.pick_index("random", 0, s, 1, 2, 3) + self.assertIn(p, (1, 2)) + + def test_no_roles_raises(self): + with self.assertRaises(ValueError): + core.pick_index("increment", 0, 0, 0, -1, 0) + + +class TestSelect(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): + 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() + try: + body, name, pos = core.select("fixed", "B", 0, 0, 0, -1) + self.assertEqual((name, pos, body), ("B", 1, "body-b.md")) + + body, name, pos = core.select("increment", "A", 2, 0, 0, -1) + self.assertEqual((name, pos), ("C", 2)) + + body, name, pos = core.select("increment", "A", 3, 0, 0, -1) + self.assertEqual((name, pos), ("A", 0)) # wrapped + + _, name_r, pos_r = core.select("random", "A", 0, 7, 0, -1) + self.assertIn(pos_r, (0, 1, 2)) + self.assertIn(name_r, ("A", "B", "C")) + finally: + core.ROLES_DIR = orig + core._load_cached.cache_clear() + + +class TestShippedRoles(unittest.TestCase): + def test_shipped_roles_present(self): + core._load_cached.cache_clear() + labels = dict((lbl, fn) for lbl, fn in core.list_roles()) + self.assertIn("Professional Photo", labels) + self.assertIn("Selfie Photo", labels) + + def test_apply_returns_body_and_title(self): + core._load_cached.cache_clear() + body, title = core.apply_role("Professional Photo") + self.assertEqual(title, "Professional Photo") + self.assertTrue(body.strip()) + + +if __name__ == "__main__": + unittest.main()