From 4f62066c635f617f1163e10c88faa5398621aced Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Sat, 7 Feb 2026 07:02:10 -0800 Subject: [PATCH] fix: harden nodes, validators, and utility modules Phase 0-3 fixes from codebase audit: - Extract shared node logic into prompt_manager_base.py (DRY) - Add input validation guards to validators.py - Fix logging_config.py buffer management - Harden image_monitor.py against race conditions - Fix prompt_tracker.py cleanup edge cases - Fix diagnostics.py import path - Wire config.py server instance correctly --- __init__.py | 4 +- prompt_manager.py | 279 ++-------------------------------------- prompt_manager_base.py | 231 +++++++++++++++++++++++++++++++++ prompt_manager_text.py | 275 ++------------------------------------- py/config.py | 2 +- utils/diagnostics.py | 2 +- utils/image_monitor.py | 21 ++- utils/logging_config.py | 13 +- utils/prompt_tracker.py | 11 +- utils/validators.py | 16 +-- 10 files changed, 293 insertions(+), 561 deletions(-) create mode 100644 prompt_manager_base.py diff --git a/__init__.py b/__init__.py index bc93503..159140a 100644 --- a/__init__.py +++ b/__init__.py @@ -79,7 +79,7 @@ except Exception as e: logger = get_logger("prompt_manager.init") logger.error(f"Failed to register API routes: {e}") - except: + except Exception: pass # Start image monitoring globally at module import time @@ -107,7 +107,7 @@ except Exception as e: 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: + except Exception: print(f"[ComfyUI-PromptManager] Warning: Failed to start image monitoring: {e}") __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"] diff --git a/prompt_manager.py b/prompt_manager.py index 8f0e9bb..c3ad1db 100644 --- a/prompt_manager.py +++ b/prompt_manager.py @@ -3,22 +3,8 @@ PromptManager: Main custom node implementation that extends CLIPTextEncode with persistent prompt storage and search capabilities. """ -import datetime -import hashlib -import json -import os import time -import webbrowser -from typing import Any, Dict, List, Optional, Tuple - -# Import logging system -try: - from .utils.logging_config import get_logger -except ImportError: - import sys - - sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - from utils.logging_config import get_logger +from typing import Any, Tuple try: from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict @@ -35,23 +21,16 @@ except ImportError: InputTypeDict = dict try: - from .database.operations import PromptDatabase - from .utils.comfyui_integration import get_comfyui_integration - from .utils.image_monitor import get_image_monitor - from .utils.prompt_tracker import PromptExecutionContext, get_prompt_tracker + from .prompt_manager_base import PromptManagerBase except ImportError: - # For direct imports when not in a package import os import sys sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - from database.operations import PromptDatabase - from utils.comfyui_integration import get_comfyui_integration - from utils.image_monitor import get_image_monitor - from utils.prompt_tracker import PromptExecutionContext, get_prompt_tracker + from prompt_manager_base import PromptManagerBase -class PromptManager(ComfyNodeABC): +class PromptManager(PromptManagerBase, ComfyNodeABC): """ A ComfyUI custom node that functions like CLIPTextEncode but adds: - Persistent storage of all prompts in SQLite database @@ -61,18 +40,7 @@ class PromptManager(ComfyNodeABC): """ def __init__(self): - self.logger = get_logger("prompt_manager.node") - self.logger.debug("Initializing PromptManager node") - - self.db = PromptDatabase() - # Use singleton getters to ensure only one tracker/monitor exists - self.prompt_tracker = get_prompt_tracker(self.db) - self.image_monitor = get_image_monitor(self.db, self.prompt_tracker) - self.comfyui_integration = get_comfyui_integration() - - # Start image monitoring automatically - self._start_gallery_system() - self.logger.debug("PromptManager node initialization completed") + super().__init__(logger_name="prompt_manager.node") @classmethod def INPUT_TYPES(cls) -> InputTypeDict: @@ -183,9 +151,6 @@ class PromptManager(ComfyNodeABC): # For database storage, save the original main text with metadata about prepend/append storage_text = text - # Search functionality is now handled by the JavaScript UI - # The search parameters are still available for backend processing if needed - # Validate CLIP model if clip is None: error_msg = ( @@ -210,7 +175,7 @@ class PromptManager(ComfyNodeABC): try: prompt_id = self._save_prompt_to_database( - text=storage_text.strip(), # Always strip whitespace + text=storage_text.strip(), category=category.strip() if category else None, tags=extended_tags if extended_tags else None, ) @@ -218,7 +183,7 @@ class PromptManager(ComfyNodeABC): # Set current prompt for image tracking if prompt_id: execution_id = self.prompt_tracker.set_current_prompt( - prompt_text=encoding_text.strip(), # Use final combined text for tracking + prompt_text=encoding_text.strip(), additional_data={ "category": category.strip() if category else None, "tags": extended_tags, @@ -227,7 +192,7 @@ class PromptManager(ComfyNodeABC): prepend_text.strip() if prepend_text else None ), "append_text": append_text.strip() if append_text else None, - "final_text": encoding_text.strip(), # Store final combined text + "final_text": encoding_text.strip(), }, ) self.logger.debug( @@ -237,7 +202,6 @@ class PromptManager(ComfyNodeABC): except Exception as e: # Log error but don't fail the encoding self.logger.warning(f"Failed to save prompt to database: {e}") - # Already logged above, no need for additional print # Perform standard CLIP text encoding using the combined text self.logger.debug( @@ -247,7 +211,7 @@ class PromptManager(ComfyNodeABC): conditioning = clip.encode_from_tokens_scheduled(tokens) # Register with ComfyUI integration for standard metadata compatibility - node_id = f"promptmanager_{int(time.time() * 1000)}" # Unique node ID + node_id = f"promptmanager_{int(time.time() * 1000)}" self.comfyui_integration.register_prompt( node_id, encoding_text.strip(), @@ -263,228 +227,6 @@ class PromptManager(ComfyNodeABC): self.logger.info(f"CLIP encoding completed, text: {repr(encoding_text)[:80]}") return (conditioning, encoding_text) - def _save_prompt_to_database( - self, text: str, category: Optional[str] = None, tags: Optional[list] = None - ) -> Optional[int]: - """ - Save the prompt to the SQLite database. - - Args: - text: The prompt text - category: Optional category - tags: List of tags - - Returns: - The prompt ID if saved successfully, None otherwise - """ - try: - # Generate hash for duplicate detection - prompt_hash = self._generate_hash(text) - self.logger.debug(f"Generated hash for prompt: {prompt_hash[:16]}...") - - # Check if prompt already exists - existing = self.db.get_prompt_by_hash(prompt_hash) - if existing: - self.logger.info( - f"Found existing prompt with ID {existing['id']}, updating metadata" - ) - # Update metadata if this is a duplicate with new info - if any([category, tags]): - self.db.update_prompt_metadata( - prompt_id=existing["id"], category=category, tags=tags - ) - self.logger.debug("Updated metadata for existing prompt") - return existing["id"] - - # Save new prompt - self.logger.debug( - f"Saving new prompt with category: {category}, tags: {tags}" - ) - prompt_id = self.db.save_prompt( - text=text, category=category, tags=tags, prompt_hash=prompt_hash - ) - - if prompt_id: - self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}") - else: - self.logger.warning("Failed to save prompt - no ID returned") - - return prompt_id - - except Exception as e: - self.logger.error(f"Error saving prompt to database: {e}") - # Already logged above, no need for additional print - return None - - def _generate_hash(self, text: str) -> str: - """ - Generate SHA256 hash for the prompt text. - - Args: - text: The prompt text to hash - - Returns: - Hexadecimal string representation of the SHA256 hash - """ - # Normalize text for consistent hashing (strip whitespace, normalize case) - normalized_text = text.strip().lower() - return hashlib.sha256(normalized_text.encode("utf-8")).hexdigest() - - def _parse_tags(self, tags_string: str) -> Optional[list]: - """ - Parse comma-separated tags string into a list. - - Args: - tags_string: Comma-separated string of tags - - Returns: - List of parsed tags, or None if no valid tags found - """ - if not tags_string or not tags_string.strip(): - return None - - tags = [tag.strip() for tag in tags_string.split(",") if tag.strip()] - return tags if tags else None - - def _search_prompts(self, search_text: str = "") -> List[Dict[str, Any]]: - """ - Search for past prompts by text content. - - Args: - search_text: Text to search for in prompt database - - Returns: - List of matching prompt dictionaries with metadata - """ - try: - if not search_text or not search_text.strip(): - return [] - - results = self.db.search_prompts( - text=search_text.strip(), - category=None, - tags=None, - rating_min=None, - limit=50, - ) - - return results - - except Exception as e: - self.logger.error(f"Error searching prompts: {e}") - return [] - - def _open_web_interface(self): - """ - Open the web interface in the default browser. - - Attempts to locate and open the web interface HTML file. - Logs warnings if the interface is not properly configured. - """ - try: - # Look for a web interface directory - current_dir = os.path.dirname(os.path.abspath(__file__)) - web_dir = os.path.join(current_dir, "web_interface") - - if os.path.exists(web_dir): - # If web interface exists, try to start it - index_path = os.path.join(web_dir, "index.html") - if os.path.exists(index_path): - webbrowser.open(f"file://{index_path}") - self.logger.info("Web interface opened in browser") - else: - self.logger.warning( - f"Web interface directory found but no index.html. Please check {web_dir} for setup instructions" - ) - else: - self.logger.info( - "Web interface not yet implemented. This feature will open a web-based prompt management interface when the web_interface directory is created." - ) - - except Exception as e: - self.logger.error(f"Error opening web interface: {e}") - - def search_prompts_api(self, search_text: str = "") -> List[Dict[str, Any]]: - """ - API method for JavaScript UI to search prompts. - - Args: - search_text: Text to search for in prompts - - Returns: - List of matching prompt dictionaries - """ - return self._search_prompts(search_text=search_text) - - def get_recent_prompts_api(self, limit: int = 20) -> List[Dict[str, Any]]: - """ - API method for JavaScript UI to get recent prompts. - - Args: - limit: Maximum number of recent prompts to retrieve - - Returns: - List of recent prompt dictionaries ordered by creation time - """ - try: - return self.db.get_recent_prompts(limit=limit) - except Exception as e: - self.logger.error(f"Error getting recent prompts: {e}") - return [] - - def _start_gallery_system(self): - """ - Initialize and start the gallery monitoring system. - - Starts the image monitor which watches for new generated images - and links them to their source prompts in the database. - """ - try: - self.logger.debug("Starting gallery system...") - - # Start image monitoring - self.image_monitor.start_monitoring() - - self.logger.debug("Gallery system started successfully") - - except Exception as e: - self.logger.error(f"Failed to start gallery system: {e}") - self.logger.warning("Gallery features will be disabled") - - def get_gallery_status(self) -> Dict[str, Any]: - """ - Get status of the gallery system. - - Returns: - Dictionary containing status information for image monitor and prompt tracker - """ - return { - "image_monitor": self.image_monitor.get_status(), - "prompt_tracker": self.prompt_tracker.get_status(), - } - - def cleanup_gallery_system(self): - """ - Clean up gallery system resources. - - Stops image monitoring and releases associated resources. - Called automatically during object destruction. - """ - try: - if hasattr(self, "image_monitor"): - self.image_monitor.stop_monitoring() - self.logger.debug("Gallery system cleaned up") - except Exception as e: - self.logger.error(f"Error cleaning up gallery system: {e}") - - def __del__(self): - """ - Cleanup when object is destroyed. - - Ensures proper resource cleanup by stopping the gallery system. - """ - self.cleanup_gallery_system() - @classmethod def IS_CHANGED(cls, clip, text="", category="", tags="", search_text="", prepend_text="", append_text="", **kwargs): @@ -496,8 +238,5 @@ class PromptManager(ComfyNodeABC): """ import hashlib - # Combine all text inputs that affect the conditioning output - # Note: search_text doesn't affect output, so it's excluded combined = f"{text}|{prepend_text}|{append_text}" - return hashlib.sha256(combined.encode()).hexdigest() diff --git a/prompt_manager_base.py b/prompt_manager_base.py new file mode 100644 index 0000000..db46a66 --- /dev/null +++ b/prompt_manager_base.py @@ -0,0 +1,231 @@ +""" +PromptManagerBase: Shared logic for PromptManager node variants. + +Provides database initialization, prompt saving, hashing, tag parsing, +search, gallery system management, and cleanup — extracted from the +duplicate code in prompt_manager.py and prompt_manager_text.py. +""" + +import hashlib +import os +import webbrowser +from typing import Any, Dict, List, Optional + +try: + from .utils.logging_config import get_logger +except ImportError: + import sys + + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + from utils.logging_config import get_logger + +try: + from .database.operations import PromptDatabase + from .utils.comfyui_integration import get_comfyui_integration + from .utils.image_monitor import get_image_monitor + from .utils.prompt_tracker import get_prompt_tracker +except ImportError: + import sys + + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + from database.operations import PromptDatabase + from utils.comfyui_integration import get_comfyui_integration + from utils.image_monitor import get_image_monitor + from utils.prompt_tracker import get_prompt_tracker + + +class PromptManagerBase: + """Mixin providing shared prompt management logic for ComfyUI nodes. + + Handles database connection, prompt saving with deduplication, + tag parsing, search, gallery system lifecycle, and cleanup. + Subclasses only need to define ComfyUI-specific class attributes + (INPUT_TYPES, RETURN_TYPES, FUNCTION) and their execution method. + """ + + def __init__(self, logger_name: str = "prompt_manager.node"): + self.logger = get_logger(logger_name) + self.logger.debug(f"Initializing {self.__class__.__name__} node") + + self.db = PromptDatabase() + self.prompt_tracker = get_prompt_tracker(self.db) + self.image_monitor = get_image_monitor(self.db, self.prompt_tracker) + self.comfyui_integration = get_comfyui_integration() + + self._start_gallery_system() + self.logger.debug(f"{self.__class__.__name__} node initialization completed") + + def _save_prompt_to_database( + self, text: str, category: Optional[str] = None, tags: Optional[list] = None + ) -> Optional[int]: + """Save the prompt to the SQLite database. + + Args: + text: The prompt text + category: Optional category + tags: List of tags + + Returns: + The prompt ID if saved successfully, None otherwise + """ + try: + prompt_hash = self._generate_hash(text) + self.logger.debug(f"Generated hash for prompt: {prompt_hash[:16]}...") + + existing = self.db.get_prompt_by_hash(prompt_hash) + if existing: + self.logger.info( + f"Found existing prompt with ID {existing['id']}, updating metadata" + ) + if any([category, tags]): + self.db.update_prompt_metadata( + prompt_id=existing["id"], category=category, tags=tags + ) + self.logger.debug("Updated metadata for existing prompt") + return existing["id"] + + self.logger.debug( + f"Saving new prompt with category: {category}, tags: {tags}" + ) + prompt_id = self.db.save_prompt( + text=text, category=category, tags=tags, prompt_hash=prompt_hash + ) + + if prompt_id: + self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}") + else: + self.logger.warning("Failed to save prompt - no ID returned") + + return prompt_id + + except Exception as e: + self.logger.error(f"Error saving prompt to database: {e}") + return None + + def _generate_hash(self, text: str) -> str: + """Generate SHA256 hash for the prompt text. + + Args: + text: The prompt text to hash + + Returns: + Hexadecimal string representation of the SHA256 hash + """ + normalized_text = text.strip().lower() + return hashlib.sha256(normalized_text.encode("utf-8")).hexdigest() + + def _parse_tags(self, tags_string: str) -> Optional[list]: + """Parse comma-separated tags string into a list. + + Args: + tags_string: Comma-separated string of tags + + Returns: + List of parsed tags, or None if no valid tags found + """ + if not tags_string or not tags_string.strip(): + return None + + tags = [tag.strip() for tag in tags_string.split(",") if tag.strip()] + return tags if tags else None + + def _search_prompts(self, search_text: str = "") -> List[Dict[str, Any]]: + """Search for past prompts by text content. + + Args: + search_text: Text to search for in prompt database + + Returns: + List of matching prompt dictionaries with metadata + """ + try: + if not search_text or not search_text.strip(): + return [] + + try: + from .py.config import PromptManagerConfig + max_results = PromptManagerConfig.MAX_SEARCH_RESULTS + except Exception: + max_results = 100 + + results = self.db.search_prompts( + text=search_text.strip(), + category=None, + tags=None, + rating_min=None, + limit=max_results, + ) + + return results + + except Exception as e: + self.logger.error(f"Error searching prompts: {e}") + return [] + + def _open_web_interface(self): + """Open the web interface in the default browser.""" + try: + current_dir = os.path.dirname(os.path.abspath(__file__)) + web_dir = os.path.join(current_dir, "web_interface") + + if os.path.exists(web_dir): + index_path = os.path.join(web_dir, "index.html") + if os.path.exists(index_path): + webbrowser.open(f"file://{index_path}") + self.logger.info("Web interface opened in browser") + else: + self.logger.warning( + f"Web interface directory found but no index.html. " + f"Please check {web_dir} for setup instructions" + ) + else: + self.logger.info( + "Web interface not yet implemented. This feature will open a " + "web-based prompt management interface when the web_interface " + "directory is created." + ) + + except Exception as e: + self.logger.error(f"Error opening web interface: {e}") + + def search_prompts_api(self, search_text: str = "") -> List[Dict[str, Any]]: + """API method for JavaScript UI to search prompts.""" + return self._search_prompts(search_text=search_text) + + def get_recent_prompts_api(self, limit: int = 20) -> List[Dict[str, Any]]: + """API method for JavaScript UI to get recent prompts.""" + try: + return self.db.get_recent_prompts(limit=limit) + except Exception as e: + self.logger.error(f"Error getting recent prompts: {e}") + return [] + + def _start_gallery_system(self): + """Initialize and start the gallery monitoring system.""" + try: + self.logger.debug("Starting gallery system...") + self.image_monitor.start_monitoring() + self.logger.debug("Gallery system started successfully") + except Exception as e: + self.logger.error(f"Failed to start gallery system: {e}") + self.logger.warning("Gallery features will be disabled") + + def get_gallery_status(self) -> Dict[str, Any]: + """Get status of the gallery system.""" + return { + "image_monitor": self.image_monitor.get_status(), + "prompt_tracker": self.prompt_tracker.get_status(), + } + + def cleanup_gallery_system(self): + """Clean up gallery system resources.""" + try: + if hasattr(self, "image_monitor"): + self.image_monitor.stop_monitoring() + self.logger.debug("Gallery system cleaned up") + except Exception as e: + self.logger.error(f"Error cleaning up gallery system: {e}") + + def __del__(self): + """Cleanup when object is destroyed.""" + self.cleanup_gallery_system() diff --git a/prompt_manager_text.py b/prompt_manager_text.py index f4bd6fd..7932bfe 100644 --- a/prompt_manager_text.py +++ b/prompt_manager_text.py @@ -3,22 +3,8 @@ PromptManagerText: A text-only version of PromptManager that outputs STRING without CLIP encoding, while maintaining all database and search features. """ -import datetime -import hashlib -import json -import os import time -import webbrowser -from typing import Any, Dict, List, Optional, Tuple - -# Import logging system -try: - from .utils.logging_config import get_logger -except ImportError: - import sys - - sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - from utils.logging_config import get_logger +from typing import Tuple try: from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict @@ -35,23 +21,16 @@ except ImportError: InputTypeDict = dict try: - from .database.operations import PromptDatabase - from .utils.comfyui_integration import get_comfyui_integration - from .utils.image_monitor import get_image_monitor - from .utils.prompt_tracker import PromptExecutionContext, get_prompt_tracker + from .prompt_manager_base import PromptManagerBase except ImportError: - # For direct imports when not in a package import os import sys sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - from database.operations import PromptDatabase - from utils.comfyui_integration import get_comfyui_integration - from utils.image_monitor import get_image_monitor - from utils.prompt_tracker import PromptExecutionContext, get_prompt_tracker + from prompt_manager_base import PromptManagerBase -class PromptManagerText(ComfyNodeABC): +class PromptManagerText(PromptManagerBase, ComfyNodeABC): """ A ComfyUI custom node that provides all PromptManager features but outputs only a STRING without CLIP encoding. Includes: @@ -63,18 +42,7 @@ class PromptManagerText(ComfyNodeABC): """ def __init__(self): - self.logger = get_logger("prompt_manager_text.node") - self.logger.debug("Initializing PromptManagerText node") - - self.db = PromptDatabase() - # Use singleton getters to ensure only one tracker/monitor exists - self.prompt_tracker = get_prompt_tracker(self.db) - self.image_monitor = get_image_monitor(self.db, self.prompt_tracker) - self.comfyui_integration = get_comfyui_integration() - - # Start image monitoring automatically - self._start_gallery_system() - self.logger.debug("PromptManagerText node initialization completed") + super().__init__(logger_name="prompt_manager_text.node") @classmethod def INPUT_TYPES(cls) -> InputTypeDict: @@ -174,9 +142,6 @@ class PromptManagerText(ComfyNodeABC): # For database storage, save the original main text with metadata about prepend/append storage_text = text - # Search functionality is now handled by the JavaScript UI - # The search parameters are still available for backend processing if needed - # Save prompt to database and set execution context for gallery tracking prompt_id = None if storage_text and storage_text.strip(): @@ -191,7 +156,7 @@ class PromptManagerText(ComfyNodeABC): try: prompt_id = self._save_prompt_to_database( - text=storage_text.strip(), # Always strip whitespace + text=storage_text.strip(), category=category.strip() if category else None, tags=extended_tags if extended_tags else None, ) @@ -199,7 +164,7 @@ class PromptManagerText(ComfyNodeABC): # Set current prompt for image tracking if prompt_id: execution_id = self.prompt_tracker.set_current_prompt( - prompt_text=final_text.strip(), # Use final combined text for tracking + prompt_text=final_text.strip(), additional_data={ "category": category.strip() if category else None, "tags": extended_tags, @@ -219,7 +184,7 @@ class PromptManagerText(ComfyNodeABC): self.logger.warning(f"Failed to save prompt to database: {e}") # Register with ComfyUI integration for standard metadata compatibility - node_id = f"promptmanagertext_{int(time.time() * 1000)}" # Unique node ID + node_id = f"promptmanagertext_{int(time.time() * 1000)}" self.comfyui_integration.register_prompt( node_id, final_text.strip(), @@ -235,227 +200,6 @@ class PromptManagerText(ComfyNodeABC): self.logger.debug(f"Text processing completed: {final_text[:100]}...") return (final_text,) - def _save_prompt_to_database( - self, text: str, category: Optional[str] = None, tags: Optional[list] = None - ) -> Optional[int]: - """ - Save the prompt to the SQLite database. - - Args: - text: The prompt text - category: Optional category - tags: List of tags - - Returns: - The prompt ID if saved successfully, None otherwise - """ - try: - # Generate hash for duplicate detection - prompt_hash = self._generate_hash(text) - self.logger.debug(f"Generated hash for prompt: {prompt_hash[:16]}...") - - # Check if prompt already exists - existing = self.db.get_prompt_by_hash(prompt_hash) - if existing: - self.logger.info( - f"Found existing prompt with ID {existing['id']}, updating metadata" - ) - # Update metadata if this is a duplicate with new info - if any([category, tags]): - self.db.update_prompt_metadata( - prompt_id=existing["id"], category=category, tags=tags - ) - self.logger.debug("Updated metadata for existing prompt") - return existing["id"] - - # Save new prompt - self.logger.debug( - f"Saving new prompt with category: {category}, tags: {tags}" - ) - prompt_id = self.db.save_prompt( - text=text, category=category, tags=tags, prompt_hash=prompt_hash - ) - - if prompt_id: - self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}") - else: - self.logger.warning("Failed to save prompt - no ID returned") - - return prompt_id - - except Exception as e: - self.logger.error(f"Error saving prompt to database: {e}") - return None - - def _generate_hash(self, text: str) -> str: - """ - Generate SHA256 hash for the prompt text. - - Args: - text: The prompt text to hash - - Returns: - Hexadecimal string representation of the SHA256 hash - """ - # Normalize text for consistent hashing (strip whitespace, normalize case) - normalized_text = text.strip().lower() - return hashlib.sha256(normalized_text.encode("utf-8")).hexdigest() - - def _parse_tags(self, tags_string: str) -> Optional[list]: - """ - Parse comma-separated tags string into a list. - - Args: - tags_string: Comma-separated string of tags - - Returns: - List of parsed tags, or None if no valid tags found - """ - if not tags_string or not tags_string.strip(): - return None - - tags = [tag.strip() for tag in tags_string.split(",") if tag.strip()] - return tags if tags else None - - def _search_prompts(self, search_text: str = "") -> List[Dict[str, Any]]: - """ - Search for past prompts by text content. - - Args: - search_text: Text to search for in prompt database - - Returns: - List of matching prompt dictionaries with metadata - """ - try: - if not search_text or not search_text.strip(): - return [] - - results = self.db.search_prompts( - text=search_text.strip(), - category=None, - tags=None, - rating_min=None, - limit=50, - ) - - return results - - except Exception as e: - self.logger.error(f"Error searching prompts: {e}") - return [] - - def _open_web_interface(self): - """ - Open the web interface in the default browser. - - Attempts to locate and open the web interface HTML file. - Logs warnings if the interface is not properly configured. - """ - try: - # Look for a web interface directory - current_dir = os.path.dirname(os.path.abspath(__file__)) - web_dir = os.path.join(current_dir, "web_interface") - - if os.path.exists(web_dir): - # If web interface exists, try to start it - index_path = os.path.join(web_dir, "index.html") - if os.path.exists(index_path): - webbrowser.open(f"file://{index_path}") - self.logger.info("Web interface opened in browser") - else: - self.logger.warning( - f"Web interface directory found but no index.html. Please check {web_dir} for setup instructions" - ) - else: - self.logger.info( - "Web interface not yet implemented. This feature will open a web-based prompt management interface when the web_interface directory is created." - ) - - except Exception as e: - self.logger.error(f"Error opening web interface: {e}") - - def search_prompts_api(self, search_text: str = "") -> List[Dict[str, Any]]: - """ - API method for JavaScript UI to search prompts. - - Args: - search_text: Text to search for in prompts - - Returns: - List of matching prompt dictionaries - """ - return self._search_prompts(search_text=search_text) - - def get_recent_prompts_api(self, limit: int = 20) -> List[Dict[str, Any]]: - """ - API method for JavaScript UI to get recent prompts. - - Args: - limit: Maximum number of recent prompts to retrieve - - Returns: - List of recent prompt dictionaries ordered by creation time - """ - try: - return self.db.get_recent_prompts(limit=limit) - except Exception as e: - self.logger.error(f"Error getting recent prompts: {e}") - return [] - - def _start_gallery_system(self): - """ - Initialize and start the gallery monitoring system. - - Starts the image monitor which watches for new generated images - and links them to their source prompts in the database. - """ - try: - self.logger.debug("Starting gallery system...") - - # Start image monitoring - self.image_monitor.start_monitoring() - - self.logger.debug("Gallery system started successfully") - - except Exception as e: - self.logger.error(f"Failed to start gallery system: {e}") - self.logger.warning("Gallery features will be disabled") - - def get_gallery_status(self) -> Dict[str, Any]: - """ - Get status of the gallery system. - - Returns: - Dictionary containing status information for image monitor and prompt tracker - """ - return { - "image_monitor": self.image_monitor.get_status(), - "prompt_tracker": self.prompt_tracker.get_status(), - } - - def cleanup_gallery_system(self): - """ - Clean up gallery system resources. - - Stops image monitoring and releases associated resources. - Called automatically during object destruction. - """ - try: - if hasattr(self, "image_monitor"): - self.image_monitor.stop_monitoring() - self.logger.debug("Gallery system cleaned up") - except Exception as e: - self.logger.error(f"Error cleaning up gallery system: {e}") - - def __del__(self): - """ - Cleanup when object is destroyed. - - Ensures proper resource cleanup by stopping the gallery system. - """ - self.cleanup_gallery_system() - @classmethod def IS_CHANGED(cls, text="", category="", tags="", search_text="", prepend_text="", append_text="", **kwargs): @@ -467,8 +211,5 @@ class PromptManagerText(ComfyNodeABC): """ import hashlib - # Combine all text inputs that affect the output - # Note: search_text doesn't affect output, so it's excluded combined = f"{text}|{prepend_text}|{append_text}" - return hashlib.sha256(combined.encode()).hexdigest() diff --git a/py/config.py b/py/config.py index 5de3937..10ece56 100644 --- a/py/config.py +++ b/py/config.py @@ -79,7 +79,7 @@ class GalleryConfig: PROCESSING_DELAY = 2.0 # Seconds to wait before processing new files # Prompt tracking settings - PROMPT_TIMEOUT = 120 # Seconds to keep prompt context active + PROMPT_TIMEOUT = 600 # Seconds to keep prompt context active (10 min for long generations) CLEANUP_INTERVAL = 300 # Seconds between cleanup of expired prompts # Database settings diff --git a/utils/diagnostics.py b/utils/diagnostics.py index 88e67cf..d6bccb4 100644 --- a/utils/diagnostics.py +++ b/utils/diagnostics.py @@ -246,7 +246,7 @@ class GalleryDiagnostics: f.write("test") os.remove(test_file) can_write = True - except: + except (IOError, OSError): can_write = False self.logger.info(f" [EDIT] Can write to directory: {can_write}") diff --git a/utils/image_monitor.py b/utils/image_monitor.py index 271330f..32015d8 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -61,8 +61,16 @@ class ImageGenerationHandler(FileSystemEventHandler): self.db_manager = db_manager self.prompt_tracker = prompt_tracker self.metadata_extractor = ComfyUIMetadataExtractor() - self.processing_delay = 2.0 # Wait 2 seconds before processing self.logger = get_logger('prompt_manager.image_monitor') + + # Read from GalleryConfig if available, otherwise use defaults + try: + from ..py.config import GalleryConfig + self.processing_delay = GalleryConfig.PROCESSING_DELAY + self.supported_extensions = tuple(GalleryConfig.SUPPORTED_EXTENSIONS) + except Exception: + self.processing_delay = 2.0 + self.supported_extensions = ('.png', '.jpg', '.jpeg', '.webp', '.gif') def on_created(self, event): """Handle filesystem creation events. @@ -96,7 +104,7 @@ class ImageGenerationHandler(FileSystemEventHandler): # Skip files in thumbnails directory - those are derivatives, not generated images if '/thumbnails/' in filepath or '\\thumbnails\\' in filepath: return False - return filepath.lower().endswith(('.png', '.jpg', '.jpeg', '.webp', '.gif')) + return filepath.lower().endswith(self.supported_extensions) def process_new_image(self, image_path: str): """Process a newly created image file for gallery integration. @@ -288,6 +296,15 @@ class ImageMonitor: self.logger.warning("Image monitoring already running") return + # Check if monitoring is enabled in config + try: + from ..py.config import GalleryConfig + if not GalleryConfig.MONITORING_ENABLED: + self.logger.info("Image monitoring disabled in config") + return + except Exception: + pass + # Check config first, then auto-detect if not configured if not output_directories: try: diff --git a/utils/logging_config.py b/utils/logging_config.py index ef29122..2d4535d 100644 --- a/utils/logging_config.py +++ b/utils/logging_config.py @@ -9,6 +9,7 @@ This module provides centralized logging configuration with support for: - Log viewer API integration """ +import collections import logging import logging.handlers import os @@ -70,7 +71,7 @@ class PromptManagerLogger: } # Memory buffer for recent logs (for web viewer) - self._log_buffer = [] + self._log_buffer = collections.deque(maxlen=self.config['buffer_size']) self._buffer_lock = threading.Lock() # Initialize loggers @@ -176,10 +177,6 @@ class PromptManagerLogger: } self._log_buffer.append(log_entry) - - # Maintain buffer size limit - if len(self._log_buffer) > self.config['buffer_size']: - self._log_buffer = self._log_buffer[-self.config['buffer_size']:] def get_recent_logs(self, limit: int = 100, level: Optional[str] = None) -> List[Dict[str, Any]]: """Get recent log entries from memory buffer. @@ -192,14 +189,14 @@ class PromptManagerLogger: List of log entry dictionaries, most recent first """ with self._buffer_lock: - logs = self._log_buffer[:] - + logs = list(self._log_buffer) + # Filter by level if specified if level: level_num = getattr(logging, level.upper(), None) if level_num: logs = [log for log in logs if getattr(logging, log['level']) >= level_num] - + # Return most recent first return list(reversed(logs[-limit:])) diff --git a/utils/prompt_tracker.py b/utils/prompt_tracker.py index 597933c..3ab1523 100644 --- a/utils/prompt_tracker.py +++ b/utils/prompt_tracker.py @@ -75,8 +75,15 @@ class PromptTracker: self._local = threading.local() self.active_prompts = {} # Global tracking for multiple threads self.lock = threading.Lock() - self.cleanup_interval = 300 # 5 minutes - self.prompt_timeout = 600 # 10 minutes (increased for longer generations) + + # Read from GalleryConfig if available, otherwise use defaults + try: + from ..py.config import GalleryConfig + self.cleanup_interval = GalleryConfig.CLEANUP_INTERVAL + self.prompt_timeout = GalleryConfig.PROMPT_TIMEOUT + except Exception: + self.cleanup_interval = 300 # 5 minutes + self.prompt_timeout = 600 # 10 minutes # Start cleanup thread self.cleanup_thread = threading.Thread(target=self._cleanup_expired_prompts, daemon=True) diff --git a/utils/validators.py b/utils/validators.py index 0d369ea..6796d2a 100644 --- a/utils/validators.py +++ b/utils/validators.py @@ -135,10 +135,10 @@ def validate_tags(tags: Union[str, List[str], None]) -> bool: if len(tag.strip()) > 50: raise ValueError("Individual tags cannot exceed 50 characters") - - # Check for invalid characters (optional - you can adjust this) - if not re.match(r'^[a-zA-Z0-9\s\-_]+$', tag.strip()): - raise ValueError(f"Tag '{tag}' contains invalid characters") + + # Reject control characters and null bytes + if re.search(r'[\x00-\x1f]', tag.strip()): + raise ValueError(f"Tag '{tag}' contains invalid control characters") if len(tags) > 20: # Reasonable limit raise ValueError("Maximum 20 tags allowed") @@ -177,10 +177,10 @@ def validate_category(category: Optional[str]) -> bool: if len(category) > 100: raise ValueError("Category cannot exceed 100 characters") - - # Check for invalid characters (adjust as needed) - if not re.match(r'^[a-zA-Z0-9\s\-_]+$', category): - raise ValueError("Category contains invalid characters") + + # Reject control characters and null bytes + if re.search(r'[\x00-\x1f]', category): + raise ValueError("Category contains invalid control characters") return True