fix: harden system prompt preset node

This commit is contained in:
Marcel Hofer
2026-08-08 11:04:07 +02:00
parent 850c96460b
commit c448e422c7
4 changed files with 71 additions and 12 deletions
+5 -1
View File
@@ -27,7 +27,11 @@ COMMERCIAL_MODULES = {
"AILab_QwenVL_PromptEnhancer",
"AILab_StreamDisplay",
}
COMMERCIAL_NODE_MODULES = {"story_split_node", "vram_cleanup"}
COMMERCIAL_NODE_MODULES = {
"story_split_node",
"system_prompt_preset",
"vram_cleanup",
}
LEGACY_NODE_REPLACEMENTS = {
"AILab_QwenVL": "ThinkingLLM_QwenVL",
+7 -4
View File
@@ -26,15 +26,19 @@ def _load_system_prompt_presets() -> tuple[list[str], dict[str, str]]:
for name, prompt in system_prompts.items()
if isinstance(prompt, str)
}
system_prompts.setdefault(NO_PRESET_PROMPT, "")
raw_presets = data.get("_preset_prompts") or []
preset_prompts = [str(name) for name in raw_presets if isinstance(name, str)]
preset_prompts = [
name
for name in raw_presets
if isinstance(name, str) and name in system_prompts
]
if not preset_prompts:
preset_prompts = list(system_prompts)
if NO_PRESET_PROMPT not in preset_prompts:
preset_prompts.insert(0, NO_PRESET_PROMPT)
system_prompts.setdefault(NO_PRESET_PROMPT, "")
return preset_prompts, system_prompts
@@ -65,8 +69,7 @@ class ThinkingLLMSystemPromptPreset:
def get_system_prompt(self, preset_prompt: str):
_, system_prompts = _load_system_prompt_presets()
# Match the existing ThinkingLLM node's compatibility fallback.
return (system_prompts.get(preset_prompt, preset_prompt),)
return (system_prompts.get(preset_prompt, ""),)
NODE_CLASS_MAPPINGS = {
+52 -6
View File
@@ -203,9 +203,17 @@ def build_loader_test_stubs() -> dict[str, types.ModuleType]:
}
def load_thinkingllm_loader_subset(module_filenames: list[str]):
def load_thinkingllm_loader_subset(
module_filenames: list[str],
node_filenames: list[str] | None = None,
):
package_name = "thinkingllm_loader_smoke"
tracked_names = [package_name, *[Path(filename).stem for filename in module_filenames]]
node_filenames = node_filenames or []
tracked_names = [
package_name,
*[Path(filename).stem for filename in module_filenames],
*[Path(filename).stem for filename in node_filenames],
]
previous_modules = {
name: sys.modules[name]
for name in tracked_names
@@ -220,7 +228,7 @@ def load_thinkingllm_loader_subset(module_filenames: list[str]):
if resolved == pkg_path:
return ["__init__.py", *module_filenames]
if resolved == nodes_path:
return []
return node_filenames
return real_listdir(path)
try:
@@ -529,7 +537,16 @@ class TestModelRecommendations(unittest.TestCase):
self.assertIn('os.environ.get("THINKINGLLM_COMMERCIAL_RELEASE") == "1"', init_source)
self.assertIn('WEB_DIRECTORY = None if COMMERCIAL_RELEASE else "./web"', init_source)
self.assertNotIn('"AILab_QwenVL_GGUF"', re.search(r"COMMERCIAL_MODULES = \{.*?\}", init_source, re.DOTALL).group(0))
self.assertIn('COMMERCIAL_NODE_MODULES = {"story_split_node", "vram_cleanup"}', init_source)
commercial_node_modules = re.search(
r"COMMERCIAL_NODE_MODULES\s*=\s*(\{.*?\})",
init_source,
re.DOTALL,
)
self.assertIsNotNone(commercial_node_modules)
self.assertEqual(
ast.literal_eval(commercial_node_modules.group(1)),
{"story_split_node", "system_prompt_preset", "vram_cleanup"},
)
self.assertIn("only permits preinstalled, hash-locked models", hf_source)
self.assertIn("if custom.exists() and not COMMERCIAL_RELEASE", hf_source)
self.assertIn("if not COMMERCIAL_RELEASE:\n _scan_local_hf_models()", hf_source)
@@ -843,7 +860,10 @@ class TestModelRecommendations(unittest.TestCase):
node_cls = module.NODE_CLASS_MAPPINGS[node_name]
combo = node_cls.INPUT_TYPES()["required"]["preset_prompt"][0]
self.assertEqual(combo, data["_preset_prompts"])
expected_presets = [
name for name in data["_preset_prompts"] if name in data["qwenvl"]
]
self.assertEqual(combo, expected_presets)
self.assertEqual(node_cls.RETURN_TYPES, ("STRING",))
self.assertEqual(node_cls.RETURN_NAMES, ("system_prompt",))
self.assertEqual(
@@ -857,15 +877,40 @@ class TestModelRecommendations(unittest.TestCase):
legacy_preset = "🍿 Wan 2.2 NSFW T2V Timeline (3s)"
self.assertNotIn(legacy_preset, data["qwenvl"])
self.assertNotIn(legacy_preset, combo)
self.assertEqual(
node_cls().get_system_prompt(legacy_preset),
(legacy_preset,),
("",),
)
def test_standalone_system_prompt_preset_node_is_registered_by_all_loaders(self):
node_name = "ThinkingLLM_SystemPromptPreset"
node_files = ["system_prompt_preset.py"]
community_package = load_thinkingllm_loader_subset([], node_files)
self.assertIn(node_name, community_package.NODE_CLASS_MAPPINGS)
with mock.patch.dict(os.environ, {"THINKINGLLM_COMMERCIAL_RELEASE": "1"}):
commercial_package = load_thinkingllm_loader_subset([], node_files)
self.assertIn(node_name, commercial_package.NODE_CLASS_MAPPINGS)
def test_appearance_js_hooks_preset_dropdown_tooltips(self):
source = read_source("web/js/appearance.js")
self.assertIn("preset_tooltips.json", source)
self.assertIn("hookPresetTooltipPreviews(node)", source)
self.assertIn("PRESET_TOOLTIP_NODE_NAMES.has(node.comfyClass)", source)
inference_nodes = re.search(
r"const THINKINGLLM_NODE_NAMES = new Set\(\[(.*?)\]\);",
source,
re.DOTALL,
).group(1)
tooltip_nodes = re.search(
r"const PRESET_TOOLTIP_NODE_NAMES = new Set\(\[(.*?)\]\);",
source,
re.DOTALL,
).group(1)
self.assertNotIn("ThinkingLLM_SystemPromptPreset", inference_nodes)
self.assertIn("ThinkingLLM_SystemPromptPreset", tooltip_nodes)
for widget_name in ["preset_prompt", "enhancement_style", "preset_system_prompt"]:
with self.subTest(widget_name=widget_name):
self.assertIn(f'"{widget_name}"', source)
@@ -913,6 +958,7 @@ class TestModelRecommendations(unittest.TestCase):
"ThinkingLLM_Gemma4_Audio_GGUF",
"ThinkingLLM_Whisper_ASR",
"ThinkingLLM_QwenVL_GGUF_PromptEnhancer",
"ThinkingLLM_SystemPromptPreset",
"VRAMCleanup",
"StorySplitNode",
]
+7 -1
View File
@@ -23,6 +23,7 @@ const NODE_COLORS = {
"ThinkingLLM_Gemma4_Audio_GGUF": "QwenVLGGUF",
"ThinkingLLM_Whisper_ASR": "QwenVLGGUF",
"ThinkingLLM_QwenVL_GGUF_PromptEnhancer": "Enhancer",
"ThinkingLLM_SystemPromptPreset": "Tools",
"AILab_QwenVL_PromptLibrary": "Tools",
"VRAMCleanup": "Tools",
"StorySplitNode": "Tools",
@@ -45,6 +46,11 @@ const THINKINGLLM_NODE_NAMES = new Set([
"ThinkingLLM_QwenVL_GGUF_PromptEnhancer",
]);
const PRESET_TOOLTIP_NODE_NAMES = new Set([
...THINKINGLLM_NODE_NAMES,
"ThinkingLLM_SystemPromptPreset",
]);
const AUDIO_CAPABLE_NODE_NAMES = new Set([
"ThinkingLLM_Gemma4_Audio_GGUF",
"ThinkingLLM_Whisper_ASR",
@@ -406,7 +412,7 @@ function updatePresetTooltip(widget, promptMap) {
}
function hookPresetTooltipPreviews(node) {
if (!THINKINGLLM_NODE_NAMES.has(node.comfyClass) || node._thinkingllmPresetTooltipsHooked) {
if (!PRESET_TOOLTIP_NODE_NAMES.has(node.comfyClass) || node._thinkingllmPresetTooltipsHooked) {
return;
}
node._thinkingllmPresetTooltipsHooked = true;