ComfyAssets-ComfyUI_PromptManager/__init__.py

"""
ComfyUI_PromptManager: A comprehensive ComfyUI custom node that extends text encoding
with persistent prompt storage, advanced search capabilities, and an automatic image
gallery system.

This module provides prompt node types:
- PromptManager: CLIP encoding node that outputs CONDITIONING
- PromptManagerFocused: Editor-first layout using the same encoding contract
- PromptManagerSplit: Side-by-side editor and library using the same contract
- PromptManagerText: Text-only node that outputs STRING with prepend/append
  functionality

The nodes share the same database backend for persistent prompt storage and include
an automatic image gallery system that monitors ComfyUI output directories and links
generated images to their source prompts.

Features:
- SQLite-based persistent prompt storage with deduplication
- Advanced search and categorization system
- Real-time image gallery with metadata extraction
- Web-based admin dashboard for prompt management
- Comprehensive logging and diagnostics
"""

import re
from pathlib import Path


def get_version():
    """Parse version from pyproject.toml"""
    try:
        pyproject_path = Path(__file__).parent / "pyproject.toml"
        if pyproject_path.exists():
            content = pyproject_path.read_text()
            match = re.search(r'version\s*=\s*["\']([^"\']+)["\']', content)
            if match:
                return match.group(1)
    except Exception:
        pass
    return "unknown"


from .prompt_manager import (  # noqa: E402
    PromptManager,
    PromptManagerFocused,
    PromptManagerSplit,
)
from .prompt_manager_pair import PromptManagerPair  # noqa: E402
from .prompt_manager_text import PromptManagerText  # noqa: E402
from .prompt_search_list import PromptSearchList  # noqa: E402

NODE_CLASS_MAPPINGS = {
    "PromptManager": PromptManager,
    "PromptManagerFocused": PromptManagerFocused,
    "PromptManagerSplit": PromptManagerSplit,
    "PromptManagerPair": PromptManagerPair,
    "PromptManagerText": PromptManagerText,
    "PromptSearchList": PromptSearchList,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "PromptManager": "Prompt Manager (Classic)",
    "PromptManagerFocused": "Prompt Manager (Focused)",
    "PromptManagerSplit": "Prompt Manager (Split)",
    "PromptManagerPair": "Prompt Manager (+/-)",
    "PromptManagerText": "Prompt Manager Text",
    "PromptSearchList": "Prompt Search List",
}

# Canvas extension only. ComfyUI loads every .js under this directory into the
# canvas page, so the admin/gallery bundles and vendored libraries under web/
# are served by our own routes instead (see py/api).
WEB_DIRECTORY = "web/comfy"

# Add API routes (same pattern as ComfyUI_Assets)
try:
    # Set extension URI
    import os

    from .py import config
    from .py.api import PromptManagerAPI

    extension_uri = os.path.dirname(__file__)
    config.extension_uri = extension_uri

    # Register API routes using the same pattern as ComfyUI_Assets
    routes = config.routes
    api = PromptManagerAPI()
    api.add_routes(routes)
    # Temp files of preview writes a crash interrupted (#317), off the
    # import path: the walk may cover a large thumbnails folder.
    import threading

    threading.Thread(
        target=api.sweep_thumbnail_temps,
        name="PromptManagerPreviewTempSweep",
        daemon=True,
    ).start()
    host_app = getattr(config.server_instance, "app", None)
    if host_app is not None and hasattr(host_app, "on_cleanup"):
        host_app.on_cleanup.append(api.cleanup_history_reads)
        host_app.on_cleanup.append(api.cleanup_annotation_captures)
        host_app.on_cleanup.append(api.cleanup_gallery_listings)
    if host_app is not None and hasattr(host_app, "on_startup"):
        # Build each output root's Images listing once the server's loop runs,
        # so the first open does not wait for the walk (#348).
        host_app.on_startup.append(api.warm_gallery_listings)
    if host_app is not None and hasattr(host_app, "on_shutdown"):
        # Open Needs review change streams end before the server waits on them.
        host_app.on_shutdown.append(api.close_review_events)

except Exception as e:
    # Log to internal logging system without console spam
    try:
        from .utils.logging_config import get_logger

        logger = get_logger("prompt_manager.init")
        logger.error(f"Failed to register API routes: {e}")
    except Exception:
        pass

# Shared database, set up with the monitoring system below. Stays None when that
# fails so the queue hook can report the database as unavailable instead of crashing.
_global_db = None

# Start image monitoring globally at module import time
# This ensures images are linked regardless of which node is used in the workflow
try:
    from .database.operations import PromptDatabase
    from .utils.image_monitor import get_image_monitor
    from .utils.prompt_tracker import get_prompt_tracker
    from .utils.logging_config import get_logger

    _init_logger = get_logger("prompt_manager.init")
    _init_logger.info("Starting global image monitoring system...")

    # Initialize database and get singleton instances
    _global_db = PromptDatabase()
    _global_prompt_tracker = get_prompt_tracker(_global_db)
    _global_image_monitor = get_image_monitor(_global_db, _global_prompt_tracker)

    # The Needs review pass for ambiguous new outputs runs through the API's
    # capture store (the one the scan and the Tag dialog use): a linked output
    # is never queued and its pending items leave the queue. It stops before
    # that store is closed at shutdown.
    _review_api = globals().get("api")
    if _review_api is not None:
        _global_image_monitor.set_review_recorder(_review_api.record_review_file)
        # Every settled output's preview is queued for the API's background
        # worker (#249): the workspace never asks to generate thumbnails.
        _global_image_monitor.set_thumbnail_queue(_review_api.queue_capture_thumbnail)
        _review_host = getattr(config.server_instance, "app", None)
        if _review_host is not None and hasattr(_review_host, "on_cleanup"):

            async def _stop_review_recording(
                app, monitor=_global_image_monitor, api=_review_api
            ):
                monitor.set_review_recorder(None)
                monitor.set_thumbnail_queue(None)
                api.close_auto_thumbnails(final=True)

            _review_host.on_cleanup.insert(0, _stop_review_recording)

    # Start monitoring
    _global_image_monitor.start_monitoring()
    _init_logger.info("Global image monitoring system started successfully")

except Exception as e:
    try:
        from .utils.logging_config import get_logger

        _init_logger = get_logger("prompt_manager.init")
        _init_logger.error(f"Failed to start image monitoring: {e}")
    except Exception:
        print(f"[ComfyUI-PromptManager] Warning: Failed to start image monitoring: {e}")


def _db_factory():
    """The shared database for the queue hook; called lazily per request."""
    if _global_db is None:
        raise RuntimeError("PromptManager database unavailable")
    return _global_db


# Count prompt runs when ComfyUI dequeues workflows (it skips unchanged nodes, so
# node execution alone misses re-runs). Separate from monitoring so either can fail.
try:
    from .py.config import server_instance
    from .utils.usage_tracking import register_queue_hook

    register_queue_hook(server_instance, _db_factory)
except Exception as e:
    try:
        from .utils.logging_config import get_logger

        get_logger("prompt_manager.init").error(f"Failed to register usage hook: {e}")
    except Exception:
        import logging

        logging.getLogger("prompt_manager.init").error(
            f"Failed to register usage hook: {e}"
        )


def _history_db_factory() -> "PromptDatabase":
    """Initialize history on its worker even if monitoring initialization failed."""
    if _global_db is not None:
        return _global_db
    from .database.operations import PromptDatabase

    return PromptDatabase()


_history_registration = None
try:
    from .py.config import server_instance
    from .utils.execution_history import register_history_hook

    _history_registration = register_history_hook(server_instance, _history_db_factory)
except Exception:
    try:
        from .utils.logging_config import get_logger

        get_logger("prompt_manager.init").error("Failed to register execution history")
    except Exception:
        import logging

        logging.getLogger("prompt_manager.init").error(
            "Failed to register execution history"
        )

__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]

# Print startup message with loaded tools
print()
print(f"\033[94m[ComfyUI-PromptManager] Version:\033[0m {get_version()}")
for node_key, display_name in NODE_DISPLAY_NAME_MAPPINGS.items():
    # ASCII only: a Windows console using a legacy code page cannot print emoji,
    # and a UnicodeEncodeError here would abort loading the whole node.
    print(f"  \033[94mLoaded:\033[0m {display_name}")
print(f"\033[94mTotal: {len(NODE_CLASS_MAPPINGS)} tools loaded\033[0m")
print()