From c448e422c7e2a63d3883bc5f230610871c4e111f Mon Sep 17 00:00:00 2001 From: Marcel Hofer <49859266+goodguy1963@users.noreply.github.com> Date: Sat, 8 Aug 2026 11:04:07 +0200 Subject: [PATCH] fix: harden system prompt preset node --- __init__.py | 6 +++- nodes/system_prompt_preset.py | 11 ++++--- tests/test_new_features.py | 58 +++++++++++++++++++++++++++++++---- web/js/appearance.js | 8 ++++- 4 files changed, 71 insertions(+), 12 deletions(-) diff --git a/__init__.py b/__init__.py index b4634e9..2af8470 100644 --- a/__init__.py +++ b/__init__.py @@ -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", diff --git a/nodes/system_prompt_preset.py b/nodes/system_prompt_preset.py index feb4a5d..ce5f114 100644 --- a/nodes/system_prompt_preset.py +++ b/nodes/system_prompt_preset.py @@ -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 = { diff --git a/tests/test_new_features.py b/tests/test_new_features.py index 168b839..b521dbc 100644 --- a/tests/test_new_features.py +++ b/tests/test_new_features.py @@ -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", ] diff --git a/web/js/appearance.js b/web/js/appearance.js index 4626f45..6dba4b3 100644 --- a/web/js/appearance.js +++ b/web/js/appearance.js @@ -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;