fix: harden system prompt preset node
This commit is contained in:
+5
-1
@@ -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",
|
||||
|
||||
@@ -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 = {
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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;
|
||||
|
||||
Reference in New Issue
Block a user