ComfyAssets-ComfyUI_PromptManager/prompt_manager.py

"""
PromptManager: Main custom node implementation that extends CLIPTextEncode
with persistent prompt storage and search capabilities.
"""

import time
from typing import Any, Dict, Optional, Tuple

if __package__:
    from .utils.execution_context_capture import record_runtime_context
else:
    from utils.execution_context_capture import record_runtime_context

try:
    from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
except ImportError:
    # Fallback for older ComfyUI versions
    class ComfyNodeABC:
        pass

    class IO:
        STRING = "STRING"
        CLIP = "CLIP"
        CONDITIONING = "CONDITIONING"

    InputTypeDict = dict

try:
    from .prompt_manager_base import PromptManagerBase
except ImportError:
    import os
    import sys

    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
    from prompt_manager_base import PromptManagerBase


class PromptManager(PromptManagerBase, ComfyNodeABC):
    """
    A ComfyUI custom node that functions like CLIPTextEncode but adds:
    - Persistent storage of all prompts in SQLite database
    - Search and retrieval capabilities
    - Metadata management (categories, tags, ratings, notes)
    - Duplicate detection via SHA256 hashing
    """

    def __init__(self):
        super().__init__(logger_name="prompt_manager.node")

    @classmethod
    def INPUT_TYPES(cls) -> InputTypeDict:
        return {
            "required": {
                "text": (
                    IO.STRING,
                    {
                        "multiline": True,
                        "dynamicPrompts": True,
                        "tooltip": (
                            "The text prompt to be encoded and saved to database."
                        ),
                    },
                ),
                "clip": (
                    IO.CLIP,
                    {"tooltip": "The CLIP model used for encoding the text."},
                ),
            },
            "optional": {
                "category": (
                    IO.STRING,
                    {
                        "default": "",
                        "tooltip": (
                            "Optional category for organizing prompts "
                            "(e.g., 'landscapes', 'portraits')"
                        ),
                    },
                ),
                "tags": (
                    IO.STRING,
                    {
                        "default": "",
                        "tooltip": (
                            "Comma-separated tags for the prompt "
                            "(e.g., 'anime, detailed, sunset')"
                        ),
                    },
                ),
                "search_text": (
                    IO.STRING,
                    {
                        "default": "",
                        "tooltip": (
                            "Search past prompts by text; add tag:name "
                            '(or tag:"two words") to filter by tag'
                        ),
                    },
                ),
                "prepend_text": (
                    IO.STRING,
                    {
                        "tooltip": (
                            "Text to prepend to the main prompt (connected STRING "
                            "nodes will be added before the main text)"
                        )
                    },
                ),
                "append_text": (
                    IO.STRING,
                    {
                        "tooltip": (
                            "Text to append to the main prompt (connected STRING "
                            "nodes will be added after the main text)"
                        )
                    },
                ),
            },
            "hidden": {
                # The queued graph and this node's id: used to tell the positive
                # prompt of a run from negative-prompt PromptManager nodes
                "prompt": "PROMPT",
                "unique_id": "UNIQUE_ID",
            },
        }

    RETURN_TYPES = (IO.CONDITIONING, IO.STRING)
    OUTPUT_TOOLTIPS = (
        "A conditioning containing the embedded text used to guide the diffusion "
        "model.",
        "The final combined text string (with prepend/append applied) that was "
        "encoded.",
    )
    FUNCTION = "encode_prompt"
    OUTPUT_NODE = True
    CATEGORY = "🫶 ComfyAssets/🧠 Prompts"
    DESCRIPTION = (
        "Encodes a text prompt using a CLIP model into an embedding that can be used "
        "to guide the diffusion model towards generating specific images. "
        "Additionally saves all prompts "
        "to a local SQLite database with optional metadata for search and retrieval."
    )

    def encode_prompt(
        self,
        clip,
        text: str,
        category: str = "",
        tags: str = "",
        search_text: str = "",
        prepend_text: str = "",
        append_text: str = "",
        prompt: Optional[Dict[str, Any]] = None,
        unique_id: Optional[str] = None,
    ) -> Tuple[Any]:
        """
        Encode the text prompt and save it to the database.

        Args:
            clip: The CLIP model for encoding
            text: The text prompt to encode
            category: Optional category for organization
            tags: Comma-separated tags
            search_text: Text to search for in past prompts
            prepend_text: Text to prepend to the main prompt
            append_text: Text to append to the main prompt

        Returns:
            Tuple containing the conditioning for the diffusion model and the final
            text string

        Raises:
            RuntimeError: If clip input is invalid
        """
        # Combine prepend, main text, and append text
        parts = []
        if prepend_text and prepend_text.strip():
            parts.append(prepend_text.strip())
        if text:
            parts.append(text)
        if append_text and append_text.strip():
            parts.append(append_text.strip())
        final_text = " ".join(parts)

        # Inject LoRA trigger words if integration is enabled
        final_text = self._inject_lora_trigger_words(final_text)

        # Use the combined text for encoding
        encoding_text = final_text

        # For database storage, save the original main text with metadata
        # about prepend/append
        storage_text = text

        # Validate CLIP model
        if clip is None:
            error_msg = (
                "ERROR: clip input is invalid: None\n\n"
                "If the clip is from a checkpoint loader node your checkpoint does not "
                "contain a valid clip or text encoder model."
            )
            self.logger.error("CLIP validation failed: clip input is None")
            raise RuntimeError(error_msg)

        # Save prompt to database and set execution context for gallery tracking
        prompt_id = None
        extended_tags = []
        if storage_text and storage_text.strip():
            self.logger.debug(f"Processing prompt text: {storage_text[:100]}...")

            # Add prepend/append info to tags if they exist
            extended_tags = self._parse_tags(tags) or []
            if prepend_text and prepend_text.strip():
                extended_tags.append(f"prepend:{prepend_text.strip()[:50]}")
            if append_text and append_text.strip():
                extended_tags.append(f"append:{append_text.strip()[:50]}")

            try:
                prompt_id = self._track_prompt_execution(
                    text=storage_text.strip(),
                    encoding_text=encoding_text.strip(),
                    category=category.strip() if category else None,
                    tags=extended_tags if extended_tags else None,
                    additional_data={
                        "category": category.strip() if category else None,
                        "tags": extended_tags,
                        "prepend_text": (
                            prepend_text.strip() if prepend_text else None
                        ),
                        "append_text": append_text.strip() if append_text else None,
                        "final_text": encoding_text.strip(),
                    },
                    prompt_graph=prompt,
                    unique_id=unique_id,
                )
                self.logger.debug(f"Tracked prompt ID: {prompt_id}")

            except Exception as e:
                # Log error but don't fail the encoding
                self.logger.warning(f"Failed to save prompt to database: {e}")
                self._mark_untracked(prompt, unique_id)
        else:
            self._skip_empty_prompt(prompt, unique_id)

        # Perform standard CLIP text encoding using the combined text
        self.logger.debug(
            f"Performing CLIP text encoding on combined text: {encoding_text[:100]}..."
        )
        tokens = clip.tokenize(encoding_text)
        conditioning = clip.encode_from_tokens_scheduled(tokens)

        # Register with ComfyUI integration for standard metadata compatibility
        node_id = f"promptmanager_{int(time.time() * 1000)}"
        self.comfyui_integration.register_prompt(
            node_id,
            encoding_text.strip(),
            {
                "category": category.strip() if category else None,
                "tags": extended_tags,
                "prompt_id": prompt_id,
                "prepend_text": prepend_text.strip() if prepend_text else None,
                "append_text": append_text.strip() if append_text else None,
            },
        )

        self.logger.info(f"CLIP encoding completed, text: {repr(encoding_text)[:80]}")
        result = (conditioning, encoding_text)
        record_runtime_context(
            type(self).__name__, text, prepend_text, append_text, encoding_text
        )
        return result

    @classmethod
    def IS_CHANGED(
        cls,
        clip,
        text="",
        category="",
        tags="",
        search_text="",
        prepend_text="",
        append_text="",
        **kwargs,
    ):
        """
        ComfyUI method to determine if node needs re-execution.

        Returns a hash of input values that affect the conditioning output or
        the saved prompt's metadata. This enables proper branch execution -
        only re-execute when inputs change.
        """
        import hashlib

        combined = f"{text}|{prepend_text}|{append_text}|{category}|{tags}"
        return hashlib.sha256(combined.encode()).hexdigest()


class PromptManagerFocused(PromptManager):
    """The editor-first UI uses the existing encoding and storage contract."""

    @classmethod
    def INPUT_TYPES(cls) -> InputTypeDict:
        inputs = super().INPUT_TYPES()
        for name in ("prepend_text", "append_text"):
            kind, options = inputs["optional"][name]
            inputs["optional"][name] = (
                kind,
                {**options, "default": "", "forceInput": True},
            )
        return inputs


class PromptManagerSplit(PromptManagerFocused):
    """The persistent library layout shares Focused encoding and storage."""