diff --git a/nodes/__init__.py b/nodes/__init__.py index ff17a60..9f78a53 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -8,6 +8,7 @@ Nodes are thin interfaces; all compiler logic lives in the ``compiler`` package. from __future__ import annotations from .category_splitter_node import CategorySplitterNode +from .prompt_builder_node import PromptBuilderNode from .resolver_node import ResolverNode from .scene_analyzer_node import SceneAnalyzerNode from .scene_validator_node import SceneValidatorNode @@ -17,6 +18,7 @@ NODE_CLASS_MAPPINGS: dict[str, type] = { "SceneCompilerValidator": SceneValidatorNode, "SceneCompilerResolver": ResolverNode, "SceneCompilerCategorySplitter": CategorySplitterNode, + "SceneCompilerPromptBuilder": PromptBuilderNode, } NODE_DISPLAY_NAME_MAPPINGS: dict[str, str] = { @@ -24,6 +26,7 @@ NODE_DISPLAY_NAME_MAPPINGS: dict[str, str] = { "SceneCompilerValidator": "Scene Validator", "SceneCompilerResolver": "Resolver", "SceneCompilerCategorySplitter": "Category Splitter", + "SceneCompilerPromptBuilder": "Prompt Builder", } __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/nodes/prompt_builder_node.py b/nodes/prompt_builder_node.py new file mode 100644 index 0000000..5394865 --- /dev/null +++ b/nodes/prompt_builder_node.py @@ -0,0 +1,40 @@ +"""ComfyUI node for the Prompt Builder (MASTER_SPEC §20, §19). + +A thin interface: it adapts the Category Map input to the Prompt Builder module +and exposes one string output per category plus the reserved ``negative`` and +``scene`` outputs (always empty in V1). No compiler logic; it never imports +ComfyUI. +""" + +from __future__ import annotations + +from typing import Any + +from compiler.builder.prompt_builder import RESERVED_OUTPUTS, build_prompts +from compiler.common.categories import CANONICAL_CATEGORIES +from compiler.common.config import Config + +# The stable set of node outputs: one per canonical category, then the reserved +# outputs. Matches the Prompt Outputs produced by build_prompts. +_OUTPUT_NAMES: tuple[str, ...] = (*CANONICAL_CATEGORIES, *RESERVED_OUTPUTS) + + +class PromptBuilderNode: + """Formats a Category Map into one prompt string per category (+ reserved).""" + + CATEGORY = "Scene Compiler" + FUNCTION = "run" + RETURN_TYPES = ("STRING",) * len(_OUTPUT_NAMES) + RETURN_NAMES = _OUTPUT_NAMES + + @classmethod + def INPUT_TYPES(cls) -> dict[str, Any]: + return { + "required": {"category_map": ("CATEGORY_MAP",)}, + "optional": {"config": ("COMPILER_CONFIG",)}, + } + + def run(self, category_map: Any, config: Config | None = None) -> tuple[str, ...]: + result = build_prompts(category_map, config or Config()) + by_name = {output.name: output.value for output in result.data} + return tuple(by_name[name] for name in _OUTPUT_NAMES) diff --git a/tests/test_prompt_builder_node.py b/tests/test_prompt_builder_node.py new file mode 100644 index 0000000..ec755e8 --- /dev/null +++ b/tests/test_prompt_builder_node.py @@ -0,0 +1,55 @@ +"""Tests for the Prompt Builder ComfyUI node (issue #23).""" + +from __future__ import annotations + +from types import MappingProxyType + +import nodes +from compiler.common.categories import CANONICAL_CATEGORIES +from compiler.common.config import Config +from nodes.prompt_builder_node import PromptBuilderNode +from schemas.models import CategoryMap, ResolvedTag + + +def rtag(name: str, category: str) -> ResolvedTag: + return ResolvedTag(tag=name, category=category, source_concept=name, knowledge_base_entry=name) + + +def category_map(**by_category: list[ResolvedTag]) -> CategoryMap: + return CategoryMap(categories=MappingProxyType({k: tuple(v) for k, v in by_category.items()})) + + +def outputs(result: tuple[str, ...]) -> dict[str, str]: + return dict(zip(PromptBuilderNode.RETURN_NAMES, result, strict=True)) + + +def test_emits_one_output_per_category_plus_reserved() -> None: + assert PromptBuilderNode.RETURN_NAMES == (*CANONICAL_CATEGORIES, "negative", "scene") + assert PromptBuilderNode.RETURN_TYPES == ("STRING",) * 21 + + +def test_formats_and_orders_category_strings() -> None: + cm = category_map(hair=[rtag("blonde hair", "hair"), rtag("long hair", "hair")]) + result = PromptBuilderNode().run(cm) + by_name = outputs(result) + assert by_name["hair"] == "blonde hair,long hair" + assert by_name["character"] == "" # empty category -> empty string + + +def test_reserved_outputs_are_empty() -> None: + cm = category_map(character=[rtag("1girl", "character")]) + by_name = outputs(PromptBuilderNode().run(cm)) + assert by_name["negative"] == "" + assert by_name["scene"] == "" + + +def test_custom_separator_from_config() -> None: + cm = category_map(hair=[rtag("a", "hair"), rtag("b", "hair")]) + config = Config.from_json({"prompt_builder": {"separator": " | "}}) + by_name = outputs(PromptBuilderNode().run(cm, config)) + assert by_name["hair"] == "a | b" + + +def test_registered_for_comfyui() -> None: + assert nodes.NODE_CLASS_MAPPINGS["SceneCompilerPromptBuilder"] is PromptBuilderNode + assert nodes.NODE_DISPLAY_NAME_MAPPINGS["SceneCompilerPromptBuilder"] == "Prompt Builder"