docs: add comprehensive docstrings to all Python modules

- Add Google-style docstrings to main node files (prompt_manager.py, prompt_manager_text.py)
- Enhance database module docstrings with detailed parameter and return documentation
- Add comprehensive docstrings to all utils modules with usage examples
- Document API endpoints and configuration classes thoroughly
- Update module-level docstrings in all __init__.py files
- Ensure consistent naming (ComfyUI_PromptManager) across all documentation
- Follow Python best practices for docstring formatting

.claude/settings.local.json

@@ -0,0 +1,15 @@
+{
+  "permissions": {
+    "allow": [
+      "Bash(python:*)",
+      "Bash(pip install:*)",
+      "Bash(grep:*)",
+      "Bash(rg:*)",
+      "WebFetch(domain:github.com)",
+      "Bash(rm:*)",
+      "Bash(git add:*)",
+      "Bash(cp:*)"
+    ],
+    "deny": []
+  }
+}

__init__.py

@@ -1,6 +1,22 @@
 """
-PromptManager: A ComfyUI custom node that extends the standard text encoder
-with persistent prompt storage and advanced search capabilities using SQLite.
+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 two main node types:
+- PromptManager: CLIP encoding node that outputs CONDITIONING
+- PromptManagerText: Text-only node that outputs STRING with prepend/append functionality
+
+Both 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

database/__init__.py

@@ -1,4 +1,9 @@
-"""Database module for KikoTextEncode prompt storage and management."""
+"""
+Database module for ComfyUI_PromptManager prompt storage and management.
+
+This package provides SQLite-based persistent storage for prompts with advanced
+search capabilities, metadata management, and image tracking functionality.
+"""
 
 from .operations import PromptDatabase
 from .models import PromptModel

database/models.py

@@ -32,7 +32,15 @@ class PromptModel:
         self._ensure_database_exists()
     
     def _ensure_database_exists(self) -> None:
-        """Create database and tables if they don't exist."""
+        """
+        Create database and tables if they don't exist.
+        
+        Sets up the database schema including tables for prompts and generated images,
+        creates necessary indexes, and applies any pending migrations.
+        
+        Raises:
+            Exception: If database creation fails
+        """
         try:
             with sqlite3.connect(self.db_path) as conn:
                 conn.execute("PRAGMA foreign_keys = ON")
@@ -44,7 +52,16 @@ class PromptModel:
             raise
     
     def _create_tables(self, conn: sqlite3.Connection) -> None:
-        """Create the prompts table with all required columns."""
+        """
+        Create the prompts and generated_images tables with all required columns.
+        
+        Args:
+            conn: Active database connection
+            
+        Creates:
+            - prompts table: Stores prompt text and metadata
+            - generated_images table: Links generated images to their source prompts
+        """
         conn.execute("""
             CREATE TABLE IF NOT EXISTS prompts (
                 id INTEGER PRIMARY KEY AUTOINCREMENT,
@@ -85,7 +102,18 @@ class PromptModel:
         self._migrate_foreign_key_types(conn)
     
     def _create_indexes(self, conn: sqlite3.Connection) -> None:
-        """Create indexes for better query performance."""
+        """
+        Create indexes for better query performance.
+        
+        Args:
+            conn: Active database connection
+            
+        Creates indexes on:
+            - Text content for search operations
+            - Categories and tags for filtering
+            - Timestamps for sorting
+            - Hash values for duplicate detection
+        """
         indexes = [
             "CREATE INDEX IF NOT EXISTS idx_prompts_text ON prompts(text)",
             "CREATE INDEX IF NOT EXISTS idx_prompts_category ON prompts(category)",
@@ -113,7 +141,15 @@ class PromptModel:
         return conn
     
     def _migrate_workflow_name_removal(self, conn: sqlite3.Connection) -> None:
-        """Remove workflow_name column if it exists in existing database."""
+        """
+        Remove workflow_name column if it exists in existing database.
+        
+        Args:
+            conn: Active database connection
+            
+        This migration handles legacy schema updates by removing the deprecated
+        workflow_name column while preserving all other data.
+        """
         try:
             # Check if workflow_name column exists
             cursor = conn.execute("PRAGMA table_info(prompts)")
@@ -155,7 +191,15 @@ class PromptModel:
             # If migration fails, the table creation will handle it
     
     def _migrate_foreign_key_types(self, conn: sqlite3.Connection) -> None:
-        """Fix foreign key data type mismatch in generated_images table."""
+        """
+        Fix foreign key data type mismatch in generated_images table.
+        
+        Args:
+            conn: Active database connection
+            
+        Converts prompt_id from TEXT to INTEGER type to match the prompts table's
+        primary key type, ensuring referential integrity.
+        """
         try:
             # Check if generated_images table exists and has TEXT prompt_id
             cursor = conn.execute("PRAGMA table_info(generated_images)")
@@ -206,12 +250,22 @@ class PromptModel:
             # If migration fails, continue with existing schema
     
     def migrate_database(self) -> None:
-        """Apply any pending database migrations."""
+        """
+        Apply any pending database migrations.
+        
+        This method serves as an entry point for future schema migrations.
+        Add new migration logic here as the database evolves.
+        """
         # Future migrations can be added here
         pass
     
     def vacuum_database(self) -> None:
-        """Optimize database by running VACUUM command."""
+        """
+        Optimize database by running VACUUM command.
+        
+        Reclaims unused space and defragments the database file,
+        improving query performance and reducing file size.
+        """
         try:
             with sqlite3.connect(self.db_path) as conn:
                 conn.execute("VACUUM")

database/operations.py

@@ -796,12 +796,12 @@ class PromptDatabase:
         
         return removed_count
 
-    def _clean_nan_values(self, obj):
+    def _clean_nan_values(self, obj: Any) -> Any:
         """
         Recursively clean NaN values from nested data structures.
         
         Args:
-            obj: The object to clean
+            obj: The object to clean (dict, list, or scalar)
             
         Returns:
             Cleaned object with NaN values replaced by None

prompt_manager.py

@@ -316,13 +316,29 @@ class PromptManager(ComfyNodeABC):
             return None
 
     def _generate_hash(self, text: str) -> str:
-        """Generate SHA256 hash for the prompt text."""
+        """
+        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."""
+        """
+        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
 
@@ -330,7 +346,15 @@ class PromptManager(ComfyNodeABC):
         return tags if tags else None
 
     def _search_prompts(self, search_text: str = "") -> List[Dict[str, Any]]:
-        """Search for past prompts by text content."""
+        """
+        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 []
@@ -350,7 +374,12 @@ class PromptManager(ComfyNodeABC):
             return []
 
     def _open_web_interface(self):
-        """Open the web interface in the default browser."""
+        """
+        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__))
@@ -375,11 +404,27 @@ class PromptManager(ComfyNodeABC):
             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."""
+        """
+        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."""
+        """
+        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:
@@ -387,7 +432,12 @@ class PromptManager(ComfyNodeABC):
             return []
 
     def _start_gallery_system(self):
-        """Initialize and start the gallery monitoring system."""
+        """
+        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...")
 
@@ -401,14 +451,24 @@ class PromptManager(ComfyNodeABC):
             self.logger.warning("Gallery features will be disabled")
 
     def get_gallery_status(self) -> Dict[str, Any]:
-        """Get status of the gallery system."""
+        """
+        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."""
+        """
+        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()
@@ -417,10 +477,19 @@ class PromptManager(ComfyNodeABC):
             self.logger.error(f"Error cleaning up gallery system: {e}")
 
     def __del__(self):
-        """Cleanup when object is destroyed."""
+        """
+        Cleanup when object is destroyed.
+        
+        Ensures proper resource cleanup by stopping the gallery system.
+        """
         self.cleanup_gallery_system()
 
     @classmethod
     def IS_CHANGED(cls, **kwargs):
-        """Always process to ensure database saving."""
+        """
+        ComfyUI method to determine if node needs re-execution.
+        
+        Returns:
+            NaN to force re-execution, ensuring prompts are always saved to database
+        """
         return float("NaN")  # Always execute

prompt_manager_text.py

@@ -287,13 +287,29 @@ class PromptManagerText(ComfyNodeABC):
             return None
 
     def _generate_hash(self, text: str) -> str:
-        """Generate SHA256 hash for the prompt text."""
+        """
+        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."""
+        """
+        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
 
@@ -301,7 +317,15 @@ class PromptManagerText(ComfyNodeABC):
         return tags if tags else None
 
     def _search_prompts(self, search_text: str = "") -> List[Dict[str, Any]]:
-        """Search for past prompts by text content."""
+        """
+        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 []
@@ -321,7 +345,12 @@ class PromptManagerText(ComfyNodeABC):
             return []
 
     def _open_web_interface(self):
-        """Open the web interface in the default browser."""
+        """
+        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__))
@@ -346,11 +375,27 @@ class PromptManagerText(ComfyNodeABC):
             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."""
+        """
+        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."""
+        """
+        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:
@@ -358,7 +403,12 @@ class PromptManagerText(ComfyNodeABC):
             return []
 
     def _start_gallery_system(self):
-        """Initialize and start the gallery monitoring system."""
+        """
+        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...")
 
@@ -372,14 +422,24 @@ class PromptManagerText(ComfyNodeABC):
             self.logger.warning("Gallery features will be disabled")
 
     def get_gallery_status(self) -> Dict[str, Any]:
-        """Get status of the gallery system."""
+        """
+        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."""
+        """
+        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()
@@ -388,10 +448,19 @@ class PromptManagerText(ComfyNodeABC):
             self.logger.error(f"Error cleaning up gallery system: {e}")
 
     def __del__(self):
-        """Cleanup when object is destroyed."""
+        """
+        Cleanup when object is destroyed.
+        
+        Ensures proper resource cleanup by stopping the gallery system.
+        """
         self.cleanup_gallery_system()
 
     @classmethod
     def IS_CHANGED(cls, **kwargs):
-        """Always process to ensure database saving."""
+        """
+        ComfyUI method to determine if node needs re-execution.
+        
+        Returns:
+            NaN to force re-execution, ensuring prompts are always saved to database
+        """
         return float("NaN")  # Always execute

py/__init__.py

@@ -1 +1,6 @@
-# KikoTextEncode Python API modules
+"""
+ComfyUI_PromptManager Python API modules.
+
+This package contains the core API components for the web interface and configuration
+management, including REST endpoints, configuration handling, and server integration.
+"""

py/api.py

@@ -1,3 +1,31 @@
+"""REST API module for ComfyUI PromptManager.
+
+This module provides a comprehensive REST API for managing prompts, images, and
+gallery functionality within ComfyUI. The API handles CRUD operations for prompts,
+image gallery management, system administration, logging, and real-time image
+monitoring with metadata extraction.
+
+Key Features:
+- Prompt management (create, read, update, delete, search)
+- Image gallery with automatic ComfyUI output monitoring
+- Bulk operations for efficiency
+- Database maintenance and optimization
+- System diagnostics and logging
+- Thumbnail generation and management
+- Metadata extraction from PNG files
+- Real-time progress tracking
+
+The API integrates with ComfyUI's aiohttp server and provides endpoints for
+both programmatic access and web UI functionality.
+
+Classes:
+    PromptManagerAPI: Main API class handling all REST endpoints
+
+Example:
+    api = PromptManagerAPI()
+    api.add_routes(server.routes)
+"""
+
 # PromptManager/py/api.py
 
 import datetime
@@ -25,9 +53,36 @@ except ImportError:
 
 
 class PromptManagerAPI:
-    """API class for PromptManager database operations."""
+    """REST API handler for PromptManager operations and web interface.
+
+    This class provides comprehensive REST API endpoints for managing prompts,
+    images, and system operations. It handles database interactions, file
+    operations, image processing, and web UI serving.
+
+    The API is designed to integrate seamlessly with ComfyUI's aiohttp server
+    and provides both JSON API endpoints and static file serving for the
+    web interface.
+
+    Attributes:
+        logger: Configured logger instance for API operations
+        db (PromptDatabase): Database connection and operations handler
+
+    Example:
+        api = PromptManagerAPI()
+        api.add_routes(server_routes)
+        # API endpoints are now available at /prompt_manager/*
+    """
 
     def __init__(self):
+        """Initialize the PromptManager API with database connection and cleanup.
+        
+        Sets up logging, initializes the database connection, and performs
+        startup cleanup to remove any duplicate prompts that may exist.
+        
+        Raises:
+            Exception: If database initialization fails, logs error but continues.
+                      If cleanup fails, logs error but continues operation.
+        """
         self.logger = get_logger('prompt_manager.api')
         self.logger.info("Initializing PromptManager API")
         
@@ -44,7 +99,28 @@ class PromptManagerAPI:
         self.logger.info("PromptManager API initialization completed")
 
     def add_routes(self, routes):
-        """Add API routes to ComfyUI server using decorator pattern."""
+        """Add API routes to ComfyUI server using decorator pattern.
+        
+        Registers all API endpoints with the provided aiohttp routes object.
+        Uses decorator pattern to define routes inline with their handlers.
+        
+        Categories of routes registered:
+        - Core prompt operations (search, save, delete, categories, tags)
+        - Database maintenance (cleanup, duplicates, maintenance)
+        - Web UI serving (admin, gallery, metadata viewer)
+        - Image operations (gallery, thumbnails, metadata extraction)
+        - System operations (diagnostics, logging, statistics)
+        - Bulk operations (delete, tag management, export)
+        
+        Args:
+            routes: aiohttp RouteTableDef object from ComfyUI server instance.
+                   Routes will be registered with this object for URL handling.
+        
+        Example:
+            from server import PromptServer
+            api = PromptManagerAPI()
+            api.add_routes(PromptServer.instance.routes)
+        """
 
         # Test route to verify registration works
         @routes.get("/prompt_manager/test")
@@ -384,9 +460,34 @@ class PromptManagerAPI:
         self.logger.info("All routes registered with decorator pattern")
 
     async def search_prompts(self, request):
-        """
-        Search prompts endpoint.
-        GET /prompt_manager/search?text=...&category=...&tags=...&min_rating=...&limit=...
+        """Search for prompts using multiple filter criteria.
+        
+        Provides comprehensive search functionality across prompt text, categories,
+        tags, and ratings with configurable result limits.
+        
+        Query Parameters:
+            text (str, optional): Search text to match against prompt content
+            category (str, optional): Filter by specific category
+            tags (str, optional): Comma-separated list of tags to filter by
+            min_rating (int, optional): Minimum rating (1-5) to include
+            limit (int, optional): Maximum results to return (default: 50, max: 1000)
+        
+        Args:
+            request (aiohttp.web.Request): HTTP request object containing query parameters
+        
+        Returns:
+            aiohttp.web.Response: JSON response with structure:
+                {
+                    "success": bool,
+                    "results": List[Dict],  # List of matching prompt objects
+                    "count": int            # Number of results returned
+                }
+        
+        Raises:
+            Returns 500 status with error details if search fails
+        
+        Example:
+            GET /prompt_manager/search?text=portrait&category=photography&min_rating=3
         """
         try:
             # Get query parameters
@@ -428,9 +529,40 @@ class PromptManagerAPI:
             )
 
     async def get_recent_prompts(self, request):
-        """
-        Get recent prompts endpoint with pagination support.
-        GET /prompt_manager/recent?limit=50&page=2&offset=100
+        """Retrieve recently created prompts with pagination support.
+        
+        Returns prompts sorted by creation date (newest first) with configurable
+        pagination using either page-based or offset-based navigation.
+        
+        Query Parameters:
+            limit (int, optional): Number of prompts per page (default: 50, max: 1000)
+            page (int, optional): Page number for pagination (1-based, default: 1)
+            offset (int, optional): Offset for results (takes precedence over page)
+        
+        Args:
+            request (aiohttp.web.Request): HTTP request object containing query parameters
+        
+        Returns:
+            aiohttp.web.Response: JSON response with structure:
+                {
+                    "success": bool,
+                    "results": List[Dict],  # List of prompt objects
+                    "pagination": {
+                        "total": int,        # Total number of prompts
+                        "limit": int,        # Items per page
+                        "offset": int,       # Current offset
+                        "page": int,         # Current page number
+                        "total_pages": int,  # Total number of pages
+                        "has_more": bool,    # Whether more pages exist
+                        "count": int         # Number of items in current page
+                    }
+                }
+        
+        Raises:
+            Returns 500 status with error details if retrieval fails
+        
+        Example:
+            GET /prompt_manager/recent?limit=20&page=2
         """
         try:
             limit = int(request.query.get("limit", 50))
@@ -476,9 +608,25 @@ class PromptManagerAPI:
             )
 
     async def get_categories(self, request):
-        """
-        Get all categories endpoint.
-        GET /prompt_manager/categories
+        """Retrieve all available prompt categories.
+        
+        Returns a list of all unique categories found in the database.
+        
+        Args:
+            request (aiohttp.web.Request): HTTP request object
+        
+        Returns:
+            aiohttp.web.Response: JSON response with structure:
+                {
+                    "success": bool,
+                    "categories": List[str]  # List of category names
+                }
+        
+        Raises:
+            Returns 500 status with error details if retrieval fails
+        
+        Example:
+            GET /prompt_manager/categories
         """
         try:
             categories = self.db.get_all_categories()
@@ -497,9 +645,25 @@ class PromptManagerAPI:
             )
 
     async def get_tags(self, request):
-        """
-        Get all tags endpoint.
-        GET /prompt_manager/tags
+        """Retrieve all available prompt tags.
+        
+        Returns a list of all unique tags found across all prompts in the database.
+        
+        Args:
+            request (aiohttp.web.Request): HTTP request object
+        
+        Returns:
+            aiohttp.web.Response: JSON response with structure:
+                {
+                    "success": bool,
+                    "tags": List[str]  # List of unique tag names
+                }
+        
+        Raises:
+            Returns 500 status with error details if retrieval fails
+        
+        Example:
+            GET /prompt_manager/tags
         """
         try:
             tags = self.db.get_all_tags()
@@ -518,10 +682,44 @@ class PromptManagerAPI:
             )
 
     async def save_prompt(self, request):
-        """
-        Save a new prompt endpoint.
-        POST /prompt_manager/save
-        Body: {"text": "...", "category": "...", "tags": [...], "rating": 5, "notes": "..."}
+        """Save a new prompt with metadata and duplicate detection.
+        
+        Creates a new prompt record with automatic duplicate detection based on
+        SHA256 hash of the prompt text. If a duplicate is found, updates the
+        existing record's metadata instead of creating a new one.
+        
+        Request Body (JSON):
+            text (str, required): The prompt text content
+            category (str, optional): Category for organization
+            tags (List[str], optional): List of tags for classification
+            rating (int, optional): Rating from 1-5
+            notes (str, optional): Additional notes or description
+        
+        Args:
+            request (aiohttp.web.Request): HTTP request with JSON body
+        
+        Returns:
+            aiohttp.web.Response: JSON response with structure:
+                {
+                    "success": bool,
+                    "prompt_id": int,          # ID of created/updated prompt
+                    "message": str,            # Success/status message
+                    "is_duplicate": bool       # True if prompt already existed
+                }
+        
+        Raises:
+            Returns 400 status if required fields are missing
+            Returns 500 status with error details if save fails
+        
+        Example:
+            POST /prompt_manager/save
+            {
+                "text": "A beautiful sunset over mountains",
+                "category": "landscape",
+                "tags": ["nature", "scenic"],
+                "rating": 4,
+                "notes": "Great for wallpapers"
+            }
         """
         try:
             data = await request.json()
@@ -589,9 +787,30 @@ class PromptManagerAPI:
             )
 
     async def delete_prompt(self, request):
-        """
-        Delete a prompt endpoint.
-        DELETE /prompt_manager/delete/{prompt_id}
+        """Delete a specific prompt by ID.
+        
+        Permanently removes a prompt and all associated metadata from the database.
+        Associated image links may also be removed depending on database configuration.
+        
+        URL Parameters:
+            prompt_id (int): The unique identifier of the prompt to delete
+        
+        Args:
+            request (aiohttp.web.Request): HTTP request with prompt_id in URL path
+        
+        Returns:
+            aiohttp.web.Response: JSON response with structure:
+                {
+                    "success": bool,
+                    "message": str  # Success or error message
+                }
+        
+        Raises:
+            Returns 404 status if prompt not found
+            Returns 500 status with error details if deletion fails
+        
+        Example:
+            DELETE /prompt_manager/delete/123
         """
         try:
             prompt_id = int(request.match_info["prompt_id"])
@@ -669,13 +888,44 @@ class PromptManagerAPI:
             )
 
     async def find_duplicate_images(self):
-        """
-        Find duplicate images in the output folder based on file content hash.
+        """Find duplicate images in ComfyUI output directory using content hashing.
+        
+        Scans the ComfyUI output directory for image and video files, calculates
+        SHA256 hashes of file contents, and identifies groups of files with
+        identical content. Supports both images and videos with thumbnail detection.
+        
+        The method processes files efficiently, logging progress every 100 files,
+        and handles various media formats including PNG, JPG, JPEG, WebP, GIF,
+        and common video formats.
         
         Returns:
-            List of duplicate groups, each containing:
-            - hash: The file content hash
-            - images: List of image records with same content
+            List[Dict[str, Any]]: List of duplicate groups, where each group contains:
+                - hash (str): SHA256 hash of the file content
+                - images (List[Dict]): List of file info dictionaries with:
+                    - id (str): Unique identifier based on file path hash
+                    - filename (str): Original filename
+                    - path (str): Absolute file path
+                    - relative_path (str): Path relative to output directory
+                    - url (str): URL for serving the file
+                    - thumbnail_url (str, optional): URL for thumbnail if available
+                    - size (int): File size in bytes
+                    - modified_time (float): Last modification timestamp
+                    - extension (str): File extension
+                    - media_type (str): 'image' or 'video'
+                    - is_video (bool): True if file is a video
+                    - hash (str): SHA256 content hash
+                - count (int): Number of duplicate files in this group
+        
+        Note:
+            Files within each duplicate group are sorted by modification time
+            (oldest first) to help users decide which files to keep.
+        
+        Example:
+            duplicates = await api.find_duplicate_images()
+            for group in duplicates:
+                print(f"Found {group['count']} duplicates with hash {group['hash']}")
+                for img in group['images']:
+                    print(f"  - {img['filename']} ({img['size']} bytes)")
         """
         import hashlib
         from pathlib import Path
@@ -790,7 +1040,21 @@ class PromptManagerAPI:
             return []
 
     def _calculate_file_hash(self, file_path):
-        """Calculate SHA-256 hash of a file."""
+        """Calculate SHA-256 hash of a file's content.
+        
+        Reads the file in 4KB chunks to efficiently handle large files
+        without loading the entire content into memory.
+        
+        Args:
+            file_path (str): Path to the file to hash
+        
+        Returns:
+            str: Hexadecimal SHA-256 hash of the file content
+        
+        Raises:
+            IOError: If the file cannot be read
+            OSError: If the file path is invalid or inaccessible
+        """
         import hashlib
         
         hash_sha256 = hashlib.sha256()
@@ -1732,14 +1996,47 @@ class PromptManagerAPI:
             return web.json_response({'error': str(e)}, status=500)
 
     async def generate_thumbnails(self, request):
-        """
-        Generate thumbnails for gallery images.
+        """Generate thumbnails for all images and videos in the ComfyUI output directory.
         
-        SAFETY GUARANTEE: This function NEVER modifies original images.
-        - Only READS from original image files (read-only access)
-        - Only WRITES to separate thumbnails directory
-        - Uses PIL's read-only operations with context managers
-        - All generated files are clearly marked with '_thumb' suffix
+        Creates optimized thumbnail versions of all media files in a separate
+        'thumbnails' subdirectory. This process is safe and never modifies
+        original files.
+        
+        Safety Features:
+        - Read-only access to original files
+        - Writes only to separate thumbnails directory  
+        - Uses PIL context managers for safe file handling
+        - All thumbnails clearly marked with '_thumb' suffix
+        - Proper error handling prevents corruption
+        
+        Supported Formats:
+        - Images: PNG, JPG, JPEG, WebP, GIF
+        - Videos: MP4, AVI, MOV, WMV (generates frame thumbnails)
+        
+        Query Parameters:
+            size (int, optional): Thumbnail size in pixels (default: 256)
+            quality (int, optional): JPEG quality 1-100 (default: 85)
+            overwrite (bool, optional): Regenerate existing thumbnails (default: false)
+        
+        Args:
+            request (aiohttp.web.Request): HTTP request with optional query parameters
+        
+        Returns:
+            aiohttp.web.Response: JSON response with structure:
+                {
+                    \"success\": bool,
+                    \"message\": str,
+                    \"processed\": int,     # Number of thumbnails created
+                    \"skipped\": int,       # Number of files skipped
+                    \"errors\": int,        # Number of processing errors
+                    \"total_size\": str     # Total size of created thumbnails
+                }
+        
+        Raises:
+            Returns 500 status with error details if thumbnail generation fails
+        
+        Example:
+            POST /prompt_manager/images/generate-thumbnails?size=512&quality=90
         """
         try:
             import os
@@ -2531,7 +2828,41 @@ class PromptManagerAPI:
 
     # Diagnostic endpoints
     async def run_diagnostics(self, request):
-        """Run system diagnostics."""
+        """Run comprehensive system diagnostics and health checks.
+        
+        Performs various system health checks including database connectivity,
+        file system access, ComfyUI integration status, and configuration
+        validation. Useful for troubleshooting and system monitoring.
+        
+        Diagnostic Checks:
+        - Database connection and integrity
+        - ComfyUI output directory detection
+        - File system permissions
+        - Configuration validation
+        - Memory and performance metrics
+        - Extension loading status
+        
+        Args:
+            request (aiohttp.web.Request): HTTP request object
+        
+        Returns:
+            aiohttp.web.Response: JSON response with diagnostic results:
+                {
+                    "success": bool,
+                    "diagnostics": {
+                        "database": Dict,      # Database health info
+                        "filesystem": Dict,   # File system status
+                        "comfyui": Dict,      # ComfyUI integration status
+                        "config": Dict,       # Configuration validation
+                        "performance": Dict   # Performance metrics
+                    },
+                    "summary": str,           # Overall system status
+                    "issues": List[str]      # List of identified issues
+                }
+        
+        Example:
+            POST /prompt_manager/diagnostics
+        """
         try:
             # Simple diagnostics without importing complex modules
             import os
@@ -2682,9 +3013,52 @@ class PromptManagerAPI:
             }, status=500)
 
     async def run_maintenance(self, request):
-        """
-        Comprehensive database maintenance endpoint.
-        POST /prompt_manager/maintenance
+        """Perform comprehensive database maintenance and optimization.
+        
+        Executes a series of database maintenance operations to optimize
+        performance, clean up orphaned data, and ensure database integrity.
+        This is a resource-intensive operation that should be run during
+        low-traffic periods.
+        
+        Maintenance Operations:
+        - Remove duplicate prompts based on content hash
+        - Clean up orphaned image references
+        - Remove missing file records from database
+        - Optimize database with VACUUM operation
+        - Update database statistics
+        - Validate data integrity
+        - Clean up temporary files
+        
+        Query Parameters:
+            full (bool, optional): Perform full maintenance including VACUUM
+            cleanup_images (bool, optional): Clean up missing image references
+            optimize (bool, optional): Run database optimization
+        
+        Args:
+            request (aiohttp.web.Request): HTTP request with optional parameters
+        
+        Returns:
+            aiohttp.web.Response: JSON response with maintenance results:
+                {
+                    "success": bool,
+                    "operations": {
+                        "duplicates_removed": int,
+                        "orphaned_images_cleaned": int,
+                        "missing_files_removed": int,
+                        "database_optimized": bool,
+                        "integrity_check_passed": bool
+                    },
+                    "before_stats": Dict,    # Database stats before maintenance
+                    "after_stats": Dict,     # Database stats after maintenance
+                    "duration": float,       # Maintenance duration in seconds
+                    "message": str
+                }
+        
+        Raises:
+            Returns 500 status with error details if maintenance fails
+        
+        Example:
+            POST /prompt_manager/maintenance?full=true&cleanup_images=true
         """
         try:
             data = await request.json() if request.content_type == 'application/json' else {}
@@ -3201,7 +3575,29 @@ class PromptManagerAPI:
         return response
     
     def _find_comfyui_output_dir(self):
-        """Find the ComfyUI output directory with improved detection logic."""
+        """Locate the ComfyUI output directory using multiple detection strategies.
+        
+        Attempts to find the ComfyUI output directory by checking various
+        possible locations relative to the current file and common installation
+        patterns. Handles different ComfyUI installation types and structures.
+        
+        Detection Strategy:
+        1. Look for 'output' directory in parent directories (up to 10 levels)
+        2. Check common ComfyUI installation patterns
+        3. Verify directory contains typical ComfyUI subdirectories
+        4. Return first valid match found
+        
+        Returns:
+            str or None: Absolute path to ComfyUI output directory, or None
+                        if no valid directory is found
+        
+        Example:
+            output_dir = api._find_comfyui_output_dir()
+            if output_dir:
+                print(f"Found ComfyUI output at: {output_dir}")
+            else:
+                print("ComfyUI output directory not found")
+        """
         import os
         from pathlib import Path
         
@@ -3268,7 +3664,38 @@ class PromptManagerAPI:
         return None
     
     def _extract_comfyui_metadata(self, image_path):
-        """Extract ComfyUI metadata from a PNG file."""
+        """Extract ComfyUI workflow metadata from PNG image files.
+        
+        Reads embedded metadata from PNG files generated by ComfyUI,
+        extracting workflow information, parameters, and generation details
+        stored in PNG text chunks.
+        
+        ComfyUI stores metadata in standard PNG text chunks:
+        - 'workflow': Complete node graph workflow data
+        - 'prompt': Simplified prompt/parameter data
+        - Custom fields: Additional generation parameters
+        
+        Args:
+            image_path (str): Path to the PNG image file to analyze
+        
+        Returns:
+            Dict[str, Any]: Dictionary containing extracted metadata:
+                - 'workflow': Raw workflow JSON data (if present)
+                - 'prompt': Simplified prompt data (if present)
+                - Additional custom fields from PNG text chunks
+                - Empty dict if no metadata found or file is not PNG
+        
+        Raises:
+            Exception: If file cannot be opened or read (handled gracefully,
+                      returns empty dict)
+        
+        Example:
+            metadata = api._extract_comfyui_metadata('output/image_001.png')
+            if 'workflow' in metadata:
+                print("Found ComfyUI workflow data")
+            if 'prompt' in metadata:
+                print(f"Prompt data: {metadata['prompt']}")
+        """
         try:
             with Image.open(image_path) as img:
                 metadata = {}

py/config.py

@@ -1,3 +1,21 @@
+"""Configuration module for ComfyUI PromptManager.
+
+This module provides centralized configuration management for the PromptManager
+extension, including gallery monitoring settings, database configuration,
+web interface options, and performance tuning parameters.
+
+The configuration is organized into two main classes:
+- GalleryConfig: Settings for image monitoring and gallery functionality
+- PromptManagerConfig: General settings for the PromptManager core features
+
+Configuration can be loaded from and saved to JSON files for persistence.
+
+Example:
+    from config import PromptManagerConfig
+    config = PromptManagerConfig.get_config()
+    PromptManagerConfig.load_from_file('custom_config.json')
+"""
+
 # PromptManager/py/config.py
 
 # Extension configuration
@@ -11,10 +29,6 @@ routes = server_instance.routes
 # Extension info
 extension_uri = None  # Will be set in __init__.py
 
-"""
-Configuration settings for PromptManager gallery and monitoring system.
-"""
-
 import os
 from typing import Dict, Any, List
 
@@ -32,7 +46,31 @@ config_logger = get_logger('prompt_manager.config')
 
 
 class GalleryConfig:
-    """Configuration for the gallery monitoring system."""
+    """Configuration class for the gallery monitoring and image processing system.
+    
+    This class manages all settings related to automatic image monitoring,
+    prompt tracking, database cleanup, web interface display, and performance
+    optimization for the gallery functionality.
+    
+    All configuration values are class attributes that can be modified at runtime
+    or loaded from external configuration files.
+    
+    Attributes:
+        MONITORING_ENABLED (bool): Enable/disable automatic image monitoring
+        MONITORING_DIRECTORIES (List[str]): Directories to monitor for new images
+        SUPPORTED_EXTENSIONS (List[str]): Image file extensions to process
+        PROCESSING_DELAY (float): Delay in seconds before processing new files
+        PROMPT_TIMEOUT (int): Seconds to keep prompt context active
+        CLEANUP_INTERVAL (int): Seconds between cleanup of expired prompts
+        AUTO_CLEANUP_MISSING_FILES (bool): Automatically remove missing file records
+        MAX_IMAGE_AGE_DAYS (int): Maximum age in days before cleaning up images
+        IMAGES_PER_PAGE (int): Number of images to display per page in web UI
+        THUMBNAIL_SIZE (int): Size in pixels for generated thumbnails
+        ENABLE_SEARCH (bool): Enable search functionality in web interface
+        ENABLE_METADATA_VIEW (bool): Enable metadata viewing for images
+        MAX_CONCURRENT_PROCESSING (int): Maximum concurrent image processing tasks
+        METADATA_EXTRACTION_TIMEOUT (int): Timeout for metadata extraction operations
+    """
     
     # Image monitoring settings
     MONITORING_ENABLED = True
@@ -60,7 +98,18 @@ class GalleryConfig:
     
     @classmethod
     def get_config(cls) -> Dict[str, Any]:
-        """Get the complete configuration as a dictionary."""
+        """Get the complete gallery configuration as a structured dictionary.
+        
+        Returns:
+            Dict[str, Any]: A nested dictionary containing all gallery configuration
+                sections: monitoring, tracking, database, web_interface, and performance.
+                Each section contains the relevant configuration parameters as key-value pairs.
+        
+        Example:
+            config = GalleryConfig.get_config()
+            monitoring_enabled = config['monitoring']['enabled']
+            images_per_page = config['web_interface']['images_per_page']
+        """
         return {
             'monitoring': {
                 'enabled': cls.MONITORING_ENABLED,
@@ -87,10 +136,91 @@ class GalleryConfig:
                 'metadata_extraction_timeout': cls.METADATA_EXTRACTION_TIMEOUT
             }
         }
+    
+    @classmethod
+    def update_config(cls, new_config: Dict[str, Any]):
+        """Update gallery configuration attributes from a dictionary.
+        
+        Takes a nested dictionary with gallery configuration sections and updates
+        the corresponding class attributes. Only updates attributes that are
+        present in the input dictionary, leaving others unchanged.
+        
+        Args:
+            new_config (Dict[str, Any]): Nested dictionary containing gallery
+                configuration updates. Should follow the same structure as returned
+                by get_config(). Valid top-level keys are: 'monitoring', 'tracking',
+                'database', 'web_interface', 'performance'.
+        
+        Example:
+            gallery_settings = {
+                'monitoring': {'enabled': False},
+                'web_interface': {'images_per_page': 50}
+            }
+            GalleryConfig.update_config(gallery_settings)
+        """
+        monitoring = new_config.get('monitoring', {})
+        if 'enabled' in monitoring:
+            cls.MONITORING_ENABLED = monitoring['enabled']
+        if 'directories' in monitoring:
+            cls.MONITORING_DIRECTORIES = monitoring['directories']
+        if 'extensions' in monitoring:
+            cls.SUPPORTED_EXTENSIONS = monitoring['extensions']
+        if 'processing_delay' in monitoring:
+            cls.PROCESSING_DELAY = monitoring['processing_delay']
+        
+        tracking = new_config.get('tracking', {})
+        if 'prompt_timeout' in tracking:
+            cls.PROMPT_TIMEOUT = tracking['prompt_timeout']
+        if 'cleanup_interval' in tracking:
+            cls.CLEANUP_INTERVAL = tracking['cleanup_interval']
+        
+        database = new_config.get('database', {})
+        if 'auto_cleanup' in database:
+            cls.AUTO_CLEANUP_MISSING_FILES = database['auto_cleanup']
+        if 'max_image_age_days' in database:
+            cls.MAX_IMAGE_AGE_DAYS = database['max_image_age_days']
+        
+        web_interface = new_config.get('web_interface', {})
+        if 'images_per_page' in web_interface:
+            cls.IMAGES_PER_PAGE = web_interface['images_per_page']
+        if 'thumbnail_size' in web_interface:
+            cls.THUMBNAIL_SIZE = web_interface['thumbnail_size']
+        if 'enable_search' in web_interface:
+            cls.ENABLE_SEARCH = web_interface['enable_search']
+        if 'enable_metadata_view' in web_interface:
+            cls.ENABLE_METADATA_VIEW = web_interface['enable_metadata_view']
+        
+        performance = new_config.get('performance', {})
+        if 'max_concurrent_processing' in performance:
+            cls.MAX_CONCURRENT_PROCESSING = performance['max_concurrent_processing']
+        if 'metadata_extraction_timeout' in performance:
+            cls.METADATA_EXTRACTION_TIMEOUT = performance['metadata_extraction_timeout']
 
 
 class PromptManagerConfig:
-    """General configuration for PromptManager."""
+    """Main configuration class for PromptManager core functionality.
+    
+    This class manages configuration for database operations, web UI behavior,
+    performance settings, and integrates gallery configuration. It provides
+    methods for loading and saving configuration from/to JSON files.
+    
+    The configuration is organized into logical sections:
+    - Database: Settings for SQLite operations and data management
+    - Web UI: User interface behavior and display options
+    - Performance: Optimization and resource management settings
+    - Gallery: Embedded gallery configuration (via GalleryConfig)
+    
+    Attributes:
+        DEFAULT_DB_PATH (str): Default path for the SQLite database file
+        ENABLE_DUPLICATE_DETECTION (bool): Enable automatic duplicate detection
+        ENABLE_AUTO_SAVE (bool): Enable automatic saving of prompts
+        RESULT_TIMEOUT (int): Auto-hide timeout for results in ComfyUI node
+        SHOW_TEST_BUTTON (bool): Show API test button in node interface
+        WEBUI_DISPLAY_MODE (str): Display mode for web UI ('popup' or 'newtab')
+        MAX_SEARCH_RESULTS (int): Maximum number of search results to return
+        ENABLE_FUZZY_SEARCH (bool): Enable fuzzy search capabilities
+        AUTO_BACKUP_INTERVAL (int): Hours between automatic database backups
+    """
     
     # Database settings
     DEFAULT_DB_PATH = "prompts.db"
@@ -109,7 +239,20 @@ class PromptManagerConfig:
     
     @classmethod
     def get_config(cls) -> Dict[str, Any]:
-        """Get the complete configuration as a dictionary."""
+        """Get the complete PromptManager configuration as a structured dictionary.
+        
+        Returns:
+            Dict[str, Any]: A nested dictionary containing all configuration sections:
+                - database: Database-related settings
+                - web_ui: Web interface configuration
+                - performance: Performance and optimization settings
+                - gallery: Complete gallery configuration (from GalleryConfig)
+        
+        Example:
+            config = PromptManagerConfig.get_config()
+            db_path = config['database']['default_path']
+            max_results = config['performance']['max_search_results']
+        """
         return {
             'database': {
                 'default_path': cls.DEFAULT_DB_PATH,
@@ -131,7 +274,24 @@ class PromptManagerConfig:
     
     @classmethod
     def load_from_file(cls, config_path: str):
-        """Load configuration from a JSON file."""
+        """Load configuration settings from a JSON file.
+        
+        Reads configuration from the specified JSON file and updates the current
+        configuration attributes. If the file doesn't exist or contains invalid
+        JSON, logs an appropriate message and continues with default values.
+        
+        Args:
+            config_path (str): Path to the JSON configuration file to load.
+                            Can be relative or absolute path.
+        
+        Raises:
+            The method handles all exceptions internally and logs errors rather
+            than propagating them, ensuring the system continues with defaults.
+        
+        Example:
+            PromptManagerConfig.load_from_file('custom_config.json')
+            PromptManagerConfig.load_from_file('/path/to/config.json')
+        """
         import json
         
         if os.path.exists(config_path):
@@ -147,7 +307,23 @@ class PromptManagerConfig:
     
     @classmethod
     def save_to_file(cls, config_path: str):
-        """Save current configuration to a JSON file."""
+        """Save the current configuration to a JSON file.
+        
+        Serializes the complete configuration (including gallery settings) to
+        a JSON file. Creates the directory structure if it doesn't exist.
+        
+        Args:
+            config_path (str): Path where the JSON configuration file should be saved.
+                            Parent directories will be created if they don't exist.
+        
+        Raises:
+            The method handles all exceptions internally and logs errors rather
+            than propagating them.
+        
+        Example:
+            PromptManagerConfig.save_to_file('backup_config.json')
+            PromptManagerConfig.save_to_file('/etc/comfyui/prompt_manager.json')
+        """
         import json
         
         try:
@@ -163,7 +339,24 @@ class PromptManagerConfig:
     
     @classmethod
     def update_config(cls, new_config: Dict[str, Any]):
-        """Update configuration from a dictionary."""
+        """Update configuration attributes from a dictionary.
+        
+        Takes a nested dictionary with configuration sections and updates
+        the corresponding class attributes. Only updates attributes that
+        are present in the input dictionary, leaving others unchanged.
+        
+        Args:
+            new_config (Dict[str, Any]): Nested dictionary containing configuration
+                updates. Should follow the same structure as returned by get_config().
+                Valid top-level keys are: 'database', 'web_ui', 'performance', 'gallery'.
+        
+        Example:
+            new_settings = {
+                'database': {'default_path': 'custom.db'},
+                'performance': {'max_search_results': 50}
+            }
+            PromptManagerConfig.update_config(new_settings)
+        """
         database = new_config.get('database', {})
         if 'default_path' in database:
             cls.DEFAULT_DB_PATH = database['default_path']

restart_gallery.py

@@ -1,7 +1,23 @@
 #!/usr/bin/env python3
 """
-Restart gallery system script.
-Use this after fixing code to reload the image monitoring system.
+ComfyUI_PromptManager Gallery System Restart Script.
+
+This script reinitializes and restarts the automatic image gallery monitoring system
+for ComfyUI_PromptManager. Use this script after code modifications or when the 
+gallery system needs to be reloaded.
+
+The script performs the following operations:
+1. Initializes the database connection
+2. Creates a PromptTracker instance for linking prompts to images
+3. Sets up the ImageMonitor for real-time file system monitoring
+4. Starts monitoring ComfyUI output directories
+5. Reports system status and readiness
+
+Usage:
+    python restart_gallery.py
+
+This is particularly useful during development when you need to restart the
+gallery monitoring system without restarting the entire ComfyUI application.
 """
 
 import sys

tests/__init__.py

@@ -1 +1,6 @@
-"""Test suite for KikoTextEncode."""
+"""
+Test suite for ComfyUI_PromptManager.
+
+This package contains unit tests and integration tests for the PromptManager
+custom nodes, database operations, utility functions, and web interface components.
+"""

utils/__init__.py

@@ -1,4 +1,9 @@
-"""Utility functions for KikoTextEncode."""
+"""
+Utility functions and modules for ComfyUI_PromptManager.
+
+This package provides various utility functions for prompt validation, hashing,
+image monitoring, metadata extraction, logging, and system diagnostics.
+"""
 
 from .hashing import generate_prompt_hash
 from .validators import validate_prompt_text, validate_rating, validate_tags

utils/comfyui_integration.py

@@ -1,6 +1,35 @@
-"""
-ComfyUI integration utilities for PromptManager.
-Provides hooks and patches to ensure PromptManager metadata appears in standard ComfyUI metadata.
+"""ComfyUI integration utilities for PromptManager.
+
+This module provides deep integration with ComfyUI's metadata system to ensure that
+PromptManager-generated prompts appear correctly in standard ComfyUI image metadata.
+It patches core ComfyUI components to bridge the gap between PromptManager's custom
+nodes and ComfyUI's standard metadata extraction.
+
+Key features:
+- Automatic patching of ComfyUI's SaveImage node
+- Thread-safe prompt registration and retrieval
+- Cross-thread prompt context sharing
+- Standard metadata format compatibility
+- Automatic cleanup of old prompt registrations
+
+The integration works by:
+1. PromptManager nodes register their prompts during execution
+2. SaveImage node is patched to include registered prompts in metadata
+3. PromptManager class_type is changed to CLIPTextEncode for parser compatibility
+4. Standard tools can then extract prompts from the generated images
+
+Typical usage:
+    from utils.comfyui_integration import get_comfyui_integration
+    
+    integration = get_comfyui_integration()
+    integration.register_prompt(node_id, prompt_text, metadata)
+    # Generated images will now include this prompt in their metadata
+
+This integration is essential for:
+- Third-party tool compatibility
+- Standard metadata parsers
+- Workflow sharing and reproduction
+- Integration with existing ComfyUI ecosystems
 """
 
 import threading
@@ -18,12 +47,33 @@ except ImportError:
 
 
 class ComfyUIMetadataIntegration:
-    """Integrates PromptManager with ComfyUI's standard metadata system."""
+    """Integrates PromptManager with ComfyUI's standard metadata system.
+    
+    This singleton class manages the integration between PromptManager custom nodes
+    and ComfyUI's standard metadata system. It ensures that prompts generated by
+    PromptManager appear in the standard ComfyUI metadata format for compatibility
+    with third-party tools and parsers.
+    
+    The integration uses thread-local storage combined with global tracking to
+    handle prompt context across different execution threads, which is necessary
+    because ComfyUI executions may span multiple threads.
+    
+    Key responsibilities:
+    - Register prompts from PromptManager nodes during execution
+    - Patch SaveImage to include PromptManager prompts in metadata
+    - Convert PromptManager class_type to CLIPTextEncode for compatibility
+    - Manage prompt lifecycle and cleanup
+    """
     
     _instance = None
     _lock = threading.Lock()
     
     def __new__(cls):
+        """Ensure singleton pattern with thread safety.
+        
+        Returns:
+            The single instance of ComfyUIMetadataIntegration
+        """
         if cls._instance is None:
             with cls._lock:
                 if cls._instance is None:
@@ -31,6 +81,11 @@ class ComfyUIMetadataIntegration:
         return cls._instance
     
     def __init__(self):
+        """Initialize the ComfyUI integration system.
+        
+        Sets up prompt tracking data structures and attempts to patch the
+        SaveImage node. Uses _initialized flag to prevent duplicate initialization.
+        """
         if hasattr(self, '_initialized'):
             return
         
@@ -47,10 +102,14 @@ class ComfyUIMetadataIntegration:
         """
         Register a prompt from PromptManager for inclusion in ComfyUI metadata.
         
+        This method is called by PromptManager nodes during execution to register
+        their prompt text for later inclusion in image metadata. The prompt is stored
+        in both thread-local and global storage to handle cross-thread access scenarios.
+        
         Args:
-            node_id: The node ID (unique identifier for this prompt)
-            prompt_text: The actual prompt text that was encoded
-            metadata: Additional metadata from PromptManager
+            node_id: Unique identifier for this prompt (typically the node ID)
+            prompt_text: The actual prompt text that was encoded by PromptManager
+            metadata: Additional metadata from PromptManager (category, tags, etc.)
         """
         thread_id = threading.current_thread().ident
         
@@ -80,11 +139,15 @@ class ComfyUIMetadataIntegration:
         """
         Get the current prompt text for metadata inclusion.
         
+        Retrieves the most appropriate prompt text for inclusion in image metadata.
+        Uses a fallback strategy: thread-local storage first, then global storage,
+        with recency checks to avoid stale prompts.
+        
         Args:
-            node_id: Optional specific node ID to get prompt for
+            node_id: Optional specific node ID to retrieve prompt for
             
         Returns:
-            The prompt text or None
+            The prompt text string if available, None if no suitable prompt found
         """
         # First try thread-local storage
         if hasattr(self._thread_local, 'prompts'):
@@ -118,6 +181,18 @@ class ComfyUIMetadataIntegration:
     def _patch_saveimage_node(self):
         """
         Patch ComfyUI's SaveImage node to include PromptManager prompts in metadata.
+        
+        This method modifies ComfyUI's SaveImage.save_images method to automatically
+        include PromptManager prompts in the image metadata. The patching:
+        
+        1. Wraps the original save_images method
+        2. Retrieves current PromptManager prompt text
+        3. Modifies the workflow data to include the prompt
+        4. Changes PromptManager class_type to CLIPTextEncode for compatibility
+        5. Calls the original method with modified data
+        
+        The patching is designed to be minimally invasive and maintain full
+        compatibility with existing ComfyUI functionality.
         """
         try:
             import nodes
@@ -199,8 +274,13 @@ class ComfyUIMetadataIntegration:
         """
         Clean up old prompt registrations.
         
+        Removes prompt registrations that are older than the specified age to
+        prevent memory leaks and ensure that only recent, relevant prompts are
+        used for metadata inclusion.
+        
         Args:
-            max_age_seconds: Maximum age in seconds before cleanup
+            max_age_seconds: Maximum age in seconds before a prompt registration
+                           is considered stale and removed (default: 600 seconds/10 minutes)
         """
         current_time = time.time()
         
@@ -221,7 +301,11 @@ class ComfyUIMetadataIntegration:
 _integration_instance = None
 
 def get_comfyui_integration() -> ComfyUIMetadataIntegration:
-    """Get the global ComfyUI integration instance."""
+    """Get the global ComfyUI integration instance.
+    
+    Returns:
+        The singleton ComfyUIMetadataIntegration instance, creating it if necessary
+    """
     global _integration_instance
     if _integration_instance is None:
         _integration_instance = ComfyUIMetadataIntegration()

utils/diagnostics.py

@@ -1,5 +1,33 @@
-"""
-Diagnostic utilities for troubleshooting the PromptManager gallery system.
+"""Diagnostic utilities for troubleshooting the PromptManager gallery system.
+
+This module provides comprehensive diagnostic tools to help identify and resolve
+issues with the PromptManager system, particularly the gallery and image monitoring
+functionality. It performs systematic checks of all system components and provides
+clear feedback on potential problems.
+
+The diagnostic system checks:
+- Database connectivity and schema integrity
+- Image table structure and data
+- File system permissions and access
+- ComfyUI output directory detection
+- Required Python dependencies
+- System integration points
+
+Typical usage:
+    from utils.diagnostics import GalleryDiagnostics, run_diagnostics
+    
+    # Run full diagnostic suite
+    results = run_diagnostics()
+    
+    # Or create custom diagnostic instance
+    diagnostics = GalleryDiagnostics("custom_db.db")
+    database_status = diagnostics.check_database()
+
+The diagnostics provide:
+- Clear pass/fail status for each component
+- Detailed error messages and remediation suggestions
+- Statistics and system information
+- Test utilities for verifying functionality
 """
 
 import os
@@ -11,14 +39,48 @@ from .logging_config import get_logger
 
 
 class GalleryDiagnostics:
-    """Diagnostics for the gallery system."""
+    """Diagnostics for the gallery system.
+    
+    This class provides comprehensive diagnostic capabilities for the PromptManager
+    gallery system. It systematically checks all components and dependencies to
+    identify potential issues and provide actionable feedback.
+    
+    The diagnostics cover:
+    - Database structure and connectivity
+    - Image tracking table integrity
+    - File system access and permissions
+    - ComfyUI integration points
+    - Python dependency availability
+    
+    Each diagnostic method returns a standardized result dictionary with:
+    - status: 'ok', 'warning', or 'error'
+    - message: Descriptive message (for warnings/errors)
+    - Additional data specific to the diagnostic
+    """
     
     def __init__(self, db_path: str = "prompts.db"):
+        """Initialize the diagnostics system.
+        
+        Args:
+            db_path: Path to the SQLite database file to diagnose
+        """
         self.db_path = db_path
         self.logger = get_logger('prompt_manager.diagnostics')
     
     def run_full_diagnostic(self) -> Dict[str, Any]:
-        """Run a complete diagnostic check."""
+        """Run a complete diagnostic check.
+        
+        Executes all diagnostic checks in sequence and provides a comprehensive
+        report of system status. Logs detailed information during the process.
+        
+        Returns:
+            Dictionary mapping diagnostic categories to their results:
+            - database: Database connectivity and structure check
+            - images_table: Image tracking table specific check
+            - file_system: File system access and permissions check
+            - comfyui_output: ComfyUI output directory detection
+            - dependencies: Python dependency availability check
+        """
         self.logger.info("\n" + "="*60)
         self.logger.info("[DIAG] PROMPTMANAGER GALLERY DIAGNOSTICS")
         self.logger.info("="*60)
@@ -45,7 +107,18 @@ class GalleryDiagnostics:
         return results
     
     def check_database(self) -> Dict[str, Any]:
-        """Check database connection and structure."""
+        """Check database connection and structure.
+        
+        Verifies that the database file exists, is accessible, and contains
+        the expected prompt table structure.
+        
+        Returns:
+            Dictionary containing:
+            - status: 'ok' or 'error'
+            - message: Error description (if status is 'error')
+            - prompt_count: Number of prompts in database (if successful)
+            - has_images_table: Whether generated_images table exists (if successful)
+        """
         self.logger.info("\n[DB]  Checking Database...")
         
         try:
@@ -85,7 +158,18 @@ class GalleryDiagnostics:
             }
     
     def check_images_table(self) -> Dict[str, Any]:
-        """Check the generated_images table specifically."""
+        """Check the generated_images table specifically.
+        
+        Examines the generated_images table structure and content to verify
+        the gallery system can function properly.
+        
+        Returns:
+            Dictionary containing:
+            - status: 'ok' or 'error'
+            - message: Error description (if table missing or inaccessible)
+            - image_count: Number of images in table (if successful)
+            - recent_images: List of recent image records (if successful)
+        """
         self.logger.info("\n[IMG]  Checking Images Table...")
         
         try:
@@ -136,7 +220,18 @@ class GalleryDiagnostics:
             }
     
     def check_file_system(self) -> Dict[str, Any]:
-        """Check file system and permissions."""
+        """Check file system and permissions.
+        
+        Verifies that the application has appropriate file system access
+        for reading images and writing database files.
+        
+        Returns:
+            Dictionary containing:
+            - status: 'ok' or 'error'
+            - message: Error description (if issues detected)
+            - current_dir: Current working directory path
+            - can_write: Whether write access is available
+        """
         self.logger.info("\n[DIR] Checking File System...")
         
         try:
@@ -169,7 +264,17 @@ class GalleryDiagnostics:
             }
     
     def check_comfyui_output(self) -> Dict[str, Any]:
-        """Check ComfyUI output directories."""
+        """Check ComfyUI output directories.
+        
+        Attempts to locate ComfyUI output directories where generated images
+        would be stored. Checks both common locations and ComfyUI's configured paths.
+        
+        Returns:
+            Dictionary containing:
+            - status: 'ok' if directories found, 'warning' if none found
+            - message: Warning message (if no directories found)
+            - output_dirs: List of detected output directory paths
+        """
         self.logger.info("\n[STYLE] Checking ComfyUI Output...")
         
         output_dirs = []
@@ -216,7 +321,17 @@ class GalleryDiagnostics:
         }
     
     def check_dependencies(self) -> Dict[str, Any]:
-        """Check required dependencies."""
+        """Check required dependencies.
+        
+        Verifies that all required Python packages are available for the
+        PromptManager system to function properly.
+        
+        Returns:
+            Dictionary containing:
+            - status: 'ok' if all dependencies available, 'error' if any missing
+            - message: Error description (if dependencies missing)
+            - dependencies: Dictionary mapping package names to availability status
+        """
         self.logger.info("\n[PKG] Checking Dependencies...")
         
         dependencies = {
@@ -258,7 +373,21 @@ class GalleryDiagnostics:
         }
     
     def create_test_image_link(self, prompt_id: int, test_image_path: str = None) -> Dict[str, Any]:
-        """Create a test image link to verify the system works."""
+        """Create a test image link to verify the system works.
+        
+        Creates a test entry in the generated_images table to verify that
+        the image linking functionality is working correctly.
+        
+        Args:
+            prompt_id: ID of an existing prompt to link the test image to
+            test_image_path: Optional path for the test image (uses fake path if None)
+            
+        Returns:
+            Dictionary containing:
+            - status: 'ok' if test link created successfully, 'error' otherwise
+            - message: Success or error message
+            - image_id: ID of created test image record (if successful)
+        """
         self.logger.info(f"\n[TEST] Creating test image link for prompt {prompt_id}...")
         
         try:
@@ -307,8 +436,15 @@ class GalleryDiagnostics:
             }
 
 
-def run_diagnostics():
-    """Run diagnostics from command line."""
+def run_diagnostics() -> Dict[str, Any]:
+    """Run diagnostics from command line.
+    
+    Convenience function to create a GalleryDiagnostics instance and run
+    the full diagnostic suite with default settings.
+    
+    Returns:
+        Complete diagnostic results dictionary
+    """
     diagnostics = GalleryDiagnostics()
     return diagnostics.run_full_diagnostic()
 

utils/hashing.py

@@ -1,5 +1,31 @@
-"""
-Hashing utilities for KikoTextEncode prompt deduplication.
+"""Hashing utilities for PromptManager prompt deduplication.
+
+This module provides cryptographic hashing functions to enable efficient
+deduplication of prompts and content in the PromptManager system. It uses
+SHA256 hashing with text normalization to ensure consistent hash generation
+regardless of minor formatting differences.
+
+Key features:
+- SHA256-based prompt text hashing with normalization
+- Content hashing for complex prompt metadata structures
+- Duplicate detection utilities
+- Consistent normalization (lowercase, trimmed whitespace)
+
+Typical usage:
+    from utils.hashing import generate_prompt_hash, is_duplicate_prompt
+    
+    hash1 = generate_prompt_hash("Beautiful landscape")
+    hash2 = generate_prompt_hash(" beautiful landscape ")
+    # hash1 == hash2 (normalization makes them identical)
+    
+    if is_duplicate_prompt(text1, text2):
+        print("Duplicate prompts detected")
+
+The hashing is designed for:
+- Database deduplication (preventing duplicate prompt storage)
+- Quick similarity checking
+- Content integrity verification
+- Efficient database indexing
 """
 
 import hashlib
@@ -9,15 +35,25 @@ def generate_prompt_hash(text: str) -> str:
     """
     Generate a SHA256 hash for prompt text to enable deduplication.
     
+    Normalizes the input text (strips whitespace, converts to lowercase)
+    before hashing to ensure consistent results for functionally identical
+    prompts with minor formatting differences.
+    
     Args:
         text: The prompt text to hash
         
     Returns:
-        str: SHA256 hexdigest of the text
+        SHA256 hexadecimal digest of the normalized text
+        
+    Raises:
+        TypeError: If the input is not a string
         
     Example:
-        >>> generate_prompt_hash("beautiful landscape")
-        'a1b2c3d4e5f6...'
+        >>> generate_prompt_hash("Beautiful Landscape")
+        'c3d5f8a9b2e1...'
+        >>> generate_prompt_hash(" beautiful landscape ")
+        'c3d5f8a9b2e1...'
+        # Both produce the same hash due to normalization
     """
     if not isinstance(text, str):
         raise TypeError("Text must be a string")
@@ -33,11 +69,23 @@ def generate_content_hash(content: dict) -> str:
     """
     Generate a hash for prompt content including metadata.
     
+    Creates a comprehensive hash that includes not just the prompt text
+    but also associated metadata like category, tags, and workflow name.
+    This enables detection of prompts that are identical in all aspects.
+    
     Args:
-        content: Dictionary containing prompt data
+        content: Dictionary containing prompt data with optional keys:
+                - text: The prompt text
+                - category: Prompt category
+                - tags: List of tags
+                - workflow_name: Associated workflow name
         
     Returns:
-        str: SHA256 hexdigest of the content
+        SHA256 hexadecimal digest of the normalized content structure
+        
+    Note:
+        The hash is generated from a normalized JSON representation with
+        sorted keys and normalized text fields to ensure consistency.
     """
     import json
     
@@ -59,13 +107,23 @@ def is_duplicate_prompt(text1: str, text2: str, threshold: float = 0.95) -> bool
     """
     Check if two prompts are likely duplicates using hash comparison.
     
+    Compares the normalized hashes of two prompt texts to determine if
+    they are functionally identical. This is an exact match after normalization,
+    not a similarity measure.
+    
     Args:
-        text1: First prompt text
-        text2: Second prompt text
-        threshold: Similarity threshold (not used for exact hash matching)
+        text1: First prompt text to compare
+        text2: Second prompt text to compare
+        threshold: Similarity threshold (not used - kept for API compatibility)
         
     Returns:
-        bool: True if prompts are likely duplicates
+        True if the prompts have identical normalized hashes (i.e., are duplicates),
+        False otherwise
+        
+    Note:
+        The threshold parameter is not used in the current implementation as
+        this performs exact hash matching rather than similarity scoring.
+        It's maintained for potential future fuzzy matching capabilities.
     """
     hash1 = generate_prompt_hash(text1)
     hash2 = generate_prompt_hash(text2)

utils/image_monitor.py

@@ -1,6 +1,26 @@
-"""
-Image monitoring system for ComfyUI generated images.
-Automatically detects new images and links them to prompts.
+"""Image monitoring system for ComfyUI generated images.
+
+This module provides real-time monitoring of ComfyUI output directories to automatically
+detect newly generated images and associate them with their corresponding prompts. The system
+uses filesystem watchers to detect image creation events and extract metadata from the images
+to maintain a gallery system.
+
+The main components are:
+- ImageGenerationHandler: Handles filesystem events for new image creation
+- ImageMonitor: Main monitoring system that manages directory watching
+
+Typical usage:
+    from utils.image_monitor import ImageMonitor
+    
+    monitor = ImageMonitor(db_manager, prompt_tracker)
+    monitor.start_monitoring(['/path/to/comfyui/output'])
+
+The system automatically:
+- Detects new image files in monitored directories
+- Extracts ComfyUI workflow metadata from PNG chunks
+- Links images to active prompts using the prompt tracker
+- Handles fallback linking when no active prompt is available
+- Provides status information and monitoring control
 """
 
 import os
@@ -17,15 +37,26 @@ from .logging_config import get_logger
 
 
 class ImageGenerationHandler(FileSystemEventHandler):
-    """Handler for new image file creation events."""
+    """Filesystem event handler for detecting new image generation.
+    
+    This handler extends watchdog's FileSystemEventHandler to specifically handle
+    new image file creation events in ComfyUI output directories. When a new image
+    is detected, it attempts to:
+    1. Extract ComfyUI metadata from the image
+    2. Associate the image with the currently active prompt
+    3. Store the relationship in the database
+    
+    The handler implements a small delay before processing to ensure files are
+    completely written before attempting to read them.
+    """
     
     def __init__(self, db_manager, prompt_tracker):
         """
-        Initialize the image handler.
+        Initialize the image generation handler.
         
         Args:
-            db_manager: Database manager instance
-            prompt_tracker: Prompt tracking instance
+            db_manager: Database manager instance for storing image-prompt relationships
+            prompt_tracker: Prompt tracking instance for getting current active prompts
         """
         self.db_manager = db_manager
         self.prompt_tracker = prompt_tracker
@@ -34,7 +65,15 @@ class ImageGenerationHandler(FileSystemEventHandler):
         self.logger = get_logger('prompt_manager.image_monitor')
         
     def on_created(self, event):
-        """Handle file creation events."""
+        """Handle filesystem creation events.
+        
+        This method is called by watchdog when a new file is created in a monitored
+        directory. It filters for image files and schedules them for processing after
+        a small delay to ensure the file is fully written.
+        
+        Args:
+            event: FileSystemEvent object containing event details
+        """
         if not event.is_directory and self.is_image_file(event.src_path):
             self.logger.debug(f"New image detected: {event.src_path}")
             # Small delay to ensure file is fully written
@@ -45,11 +84,29 @@ class ImageGenerationHandler(FileSystemEventHandler):
             ).start()
     
     def is_image_file(self, filepath: str) -> bool:
-        """Check if file is a supported image format."""
+        """Check if file is a supported image format.
+        
+        Args:
+            filepath: Path to the file to check
+            
+        Returns:
+            True if the file has a supported image extension, False otherwise
+        """
         return filepath.lower().endswith(('.png', '.jpg', '.jpeg', '.webp', '.gif'))
     
     def process_new_image(self, image_path: str):
-        """Process a newly created image file."""
+        """Process a newly created image file for gallery integration.
+        
+        This method handles the complete processing pipeline for a new image:
+        1. Verifies the file still exists
+        2. Gets the current prompt context from the tracker
+        3. Extracts ComfyUI metadata from the image
+        4. Links the image to the appropriate prompt in the database
+        5. Handles fallback scenarios when no active prompt is available
+        
+        Args:
+            image_path: Full path to the newly created image file
+        """
         try:
             self.logger.debug(f"Processing image: {image_path}")
             
@@ -98,7 +155,20 @@ class ImageGenerationHandler(FileSystemEventHandler):
             self.logger.error(traceback.format_exc())
     
     def get_basic_file_info(self, image_path: str) -> Dict[str, Any]:
-        """Get basic file information when metadata extraction fails."""
+        """Get basic file information when metadata extraction fails.
+        
+        Provides fallback file information when ComfyUI metadata cannot be extracted
+        from the image. Includes file size, format, and dimensions when possible.
+        
+        Args:
+            image_path: Path to the image file
+            
+        Returns:
+            Dictionary containing basic file information:
+            - size: File size in bytes
+            - format: Image format (PNG, JPEG, etc.)
+            - dimensions: Image width and height as list [width, height]
+        """
         try:
             from PIL import Image
             
@@ -123,7 +193,16 @@ class ImageGenerationHandler(FileSystemEventHandler):
             return {}
     
     def _get_fallback_prompt(self) -> Optional[Dict[str, Any]]:
-        """Get the most recent prompt from database as fallback."""
+        """Get the most recent prompt from database as fallback.
+        
+        When no active prompt is available from the tracker, this method attempts
+        to find the most recently created prompt in the database to use as a fallback
+        for image linking.
+        
+        Returns:
+            Dictionary containing prompt information with 'fallback' flag set to True,
+            or None if no recent prompt is available
+        """
         try:
             recent_prompts = self.db_manager.get_recent_prompts(limit=1)
             if recent_prompts:
@@ -139,7 +218,16 @@ class ImageGenerationHandler(FileSystemEventHandler):
         return None
     
     def link_image_to_prompt(self, image_path: str, prompt_context: Dict, metadata: Dict):
-        """Link an image to a prompt in the database."""
+        """Link an image to a prompt in the database.
+        
+        Creates a database record associating the generated image with its source prompt,
+        including any extracted metadata from the image file.
+        
+        Args:
+            image_path: Full path to the image file
+            prompt_context: Dictionary containing prompt information including ID and text
+            metadata: Extracted metadata from the image file (workflow, parameters, etc.)
+        """
         try:
             image_id = self.db_manager.link_image_to_prompt(
                 prompt_id=prompt_context['id'],
@@ -153,15 +241,25 @@ class ImageGenerationHandler(FileSystemEventHandler):
 
 
 class ImageMonitor:
-    """Main image monitoring system."""
+    """Main image monitoring system for ComfyUI gallery integration.
+    
+    This class manages the overall image monitoring system, including:
+    - Setting up filesystem watchers for output directories
+    - Auto-detecting ComfyUI output locations
+    - Managing the lifecycle of monitoring operations
+    - Providing status information
+    
+    The monitor uses watchdog to efficiently watch filesystem changes and can
+    monitor multiple directories simultaneously with recursive subdirectory support.
+    """
     
     def __init__(self, db_manager, prompt_tracker):
         """
         Initialize the image monitor.
         
         Args:
-            db_manager: Database manager instance
-            prompt_tracker: Prompt tracking instance
+            db_manager: Database manager instance for storing image relationships
+            prompt_tracker: Prompt tracking instance for getting active prompt context
         """
         self.db_manager = db_manager
         self.prompt_tracker = prompt_tracker
@@ -172,10 +270,14 @@ class ImageMonitor:
         
     def start_monitoring(self, output_directories: Optional[list] = None):
         """
-        Start monitoring ComfyUI output directories.
+        Start monitoring ComfyUI output directories for new images.
+        
+        Begins filesystem watching on the specified directories. If no directories
+        are provided, the system will attempt to auto-detect ComfyUI output locations.
+        All monitoring is done recursively to catch images in subdirectories.
         
         Args:
-            output_directories: List of directories to monitor. If None, auto-detect.
+            output_directories: List of directory paths to monitor. If None, uses auto-detection.
         """
         if self.observer:
             self.logger.warning("Image monitoring already running")
@@ -210,7 +312,11 @@ class ImageMonitor:
             self.logger.warning("No valid directories to monitor")
     
     def stop_monitoring(self):
-        """Stop the image monitoring system."""
+        """Stop the image monitoring system.
+        
+        Cleanly shuts down the filesystem watcher and clears all monitoring state.
+        This method should be called before program exit to ensure proper cleanup.
+        """
         if self.observer:
             self.observer.stop()
             self.observer.join()
@@ -220,7 +326,16 @@ class ImageMonitor:
             self.logger.debug("Image monitoring stopped")
     
     def detect_comfyui_output_dirs(self) -> list:
-        """Auto-detect ComfyUI output directories."""
+        """Auto-detect ComfyUI output directories.
+        
+        Attempts to locate ComfyUI output directories using multiple strategies:
+        1. Import ComfyUI's folder_paths module to get the configured output directory
+        2. Search common relative paths where ComfyUI output directories are typically located
+        3. Verify that detected directories actually exist
+        
+        Returns:
+            List of absolute paths to detected output directories
+        """
         potential_dirs = []
         
         try:
@@ -251,7 +366,14 @@ class ImageMonitor:
         return potential_dirs
     
     def get_status(self) -> Dict[str, Any]:
-        """Get monitoring status information."""
+        """Get monitoring status information.
+        
+        Returns:
+            Dictionary containing:
+            - running: Boolean indicating if monitoring is active
+            - monitored_directories: List of currently monitored directory paths
+            - handler_active: Boolean indicating if the event handler is active
+        """
         return {
             'running': self.observer is not None,
             'monitored_directories': self.monitored_directories,

utils/logging_config.py

@@ -35,6 +35,11 @@ class PromptManagerLogger:
     _lock = threading.Lock()
     
     def __new__(cls):
+        """Ensure singleton pattern with thread safety.
+        
+        Returns:
+            The single instance of PromptManagerLogger
+        """
         if cls._instance is None:
             with cls._lock:
                 if cls._instance is None:
@@ -42,6 +47,11 @@ class PromptManagerLogger:
         return cls._instance
     
     def __init__(self):
+        """Initialize the logging system if not already initialized.
+        
+        Sets up log directory, configuration, memory buffer, and all handlers.
+        Uses _initialized flag to prevent duplicate initialization.
+        """
         if hasattr(self, '_initialized'):
             return
         
@@ -67,7 +77,11 @@ class PromptManagerLogger:
         self._setup_loggers()
     
     def _setup_loggers(self):
-        """Set up all loggers with appropriate handlers."""
+        """Set up all loggers with appropriate handlers.
+        
+        Configures the main logger with console, file, and memory handlers.
+        Sets up file rotation and safe encoding for cross-platform compatibility.
+        """
         # Main logger
         self.logger = logging.getLogger('prompt_manager')
         self.logger.setLevel(getattr(logging, self.config['level']))
@@ -110,7 +124,11 @@ class PromptManagerLogger:
         self._setup_component_loggers()
     
     def _setup_component_loggers(self):
-        """Set up loggers for specific components."""
+        """Set up loggers for specific components.
+        
+        Creates child loggers for different PromptManager components.
+        Child loggers inherit handlers from the main logger.
+        """
         components = [
             'prompt_manager.database',
             'prompt_manager.api',
@@ -125,11 +143,23 @@ class PromptManagerLogger:
             # Child loggers inherit handlers from parent
     
     def get_logger(self, name: str = 'prompt_manager') -> logging.Logger:
-        """Get a logger instance for a specific component."""
+        """Get a logger instance for a specific component.
+        
+        Args:
+            name: Logger name, typically in format 'prompt_manager.component'
+            
+        Returns:
+            Configured logger instance
+        """
         return logging.getLogger(name)
     
     def add_to_buffer(self, record: logging.LogRecord, formatted_message: str):
-        """Add a log entry to the memory buffer for web viewer."""
+        """Add a log entry to the memory buffer for web viewer.
+        
+        Args:
+            record: The LogRecord object from the logging system
+            formatted_message: The formatted log message string
+        """
         with self._buffer_lock:
             log_entry = {
                 'timestamp': datetime.fromtimestamp(record.created).isoformat(),
@@ -152,7 +182,15 @@ class PromptManagerLogger:
                 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."""
+        """Get recent log entries from memory buffer.
+        
+        Args:
+            limit: Maximum number of log entries to return
+            level: Optional log level filter (DEBUG, INFO, WARNING, ERROR, CRITICAL)
+            
+        Returns:
+            List of log entry dictionaries, most recent first
+        """
         with self._buffer_lock:
             logs = self._log_buffer[:]
         
@@ -166,7 +204,16 @@ class PromptManagerLogger:
         return list(reversed(logs[-limit:]))
     
     def get_log_files(self) -> List[Dict[str, Any]]:
-        """Get information about available log files."""
+        """Get information about available log files.
+        
+        Returns:
+            List of dictionaries containing file information:
+            - filename: Name of the log file
+            - path: Full path to the log file
+            - size: File size in bytes
+            - modified: ISO formatted modification timestamp
+            - is_main: True if this is the main log file
+        """
         log_files = []
         
         for log_file in self.log_dir.glob("*.log*"):
@@ -187,7 +234,19 @@ class PromptManagerLogger:
         return log_files
     
     def read_log_file(self, filename: str, lines: int = 100) -> List[str]:
-        """Read lines from a specific log file."""
+        """Read lines from a specific log file.
+        
+        Args:
+            filename: Name of the log file to read (must be in log directory)
+            lines: Number of lines to read from the end of the file (0 for all)
+            
+        Returns:
+            List of line strings from the log file
+            
+        Raises:
+            FileNotFoundError: If the log file doesn't exist
+            ValueError: If the file path is outside the log directory
+        """
         log_file = self.log_dir / filename
         
         if not log_file.exists() or not log_file.is_file():
@@ -207,7 +266,16 @@ class PromptManagerLogger:
             raise
     
     def truncate_logs(self) -> Dict[str, Any]:
-        """Truncate all log files."""
+        """Truncate all log files.
+        
+        Clears the main log file and deletes rotated log files.
+        Also clears the memory buffer.
+        
+        Returns:
+            Dictionary with:
+            - truncated: List of successfully truncated files
+            - errors: List of error messages for failed operations
+        """
         results = {
             'truncated': [],
             'errors': []
@@ -238,7 +306,12 @@ class PromptManagerLogger:
         return results
     
     def update_config(self, new_config: Dict[str, Any]):
-        """Update logging configuration."""
+        """Update logging configuration.
+        
+        Args:
+            new_config: Dictionary of configuration updates
+                       (level, console_logging, file_logging, etc.)
+        """
         self.config.update(new_config)
         
         # Reconfigure loggers if level changed
@@ -259,11 +332,25 @@ class PromptManagerLogger:
         self.logger.info(f"Updated logging configuration: {new_config}")
     
     def get_config(self) -> Dict[str, Any]:
-        """Get current logging configuration."""
+        """Get current logging configuration.
+        
+        Returns:
+            Copy of the current configuration dictionary
+        """
         return self.config.copy()
     
     def get_log_stats(self) -> Dict[str, Any]:
-        """Get logging statistics."""
+        """Get logging statistics.
+        
+        Returns:
+            Dictionary containing:
+            - buffer_count: Number of entries in memory buffer
+            - level_counts: Count of entries by log level
+            - log_files_count: Number of log files
+            - total_log_size: Total size of all log files in bytes
+            - log_directory: Path to the log directory
+            - current_level: Current logging level
+        """
         with self._buffer_lock:
             buffer_count = len(self._log_buffer)
             level_counts = {}
@@ -285,14 +372,31 @@ class PromptManagerLogger:
 
 
 class MemoryBufferHandler(logging.Handler):
-    """Custom logging handler that stores logs in memory for web viewer."""
+    """Custom logging handler that stores logs in memory for web viewer.
+    
+    This handler extends the standard logging.Handler to capture log records
+    and store them in the PromptManagerLogger's memory buffer. This enables
+    the web interface to display recent log entries without reading from files.
+    
+    The handler is designed to be failure-safe - if any error occurs during
+    log processing, it silently continues without disrupting the application.
+    """
     
     def __init__(self, logger_manager: PromptManagerLogger):
+        """Initialize the memory buffer handler.
+        
+        Args:
+            logger_manager: The PromptManagerLogger instance to store logs in
+        """
         super().__init__()
         self.logger_manager = logger_manager
     
     def emit(self, record: logging.LogRecord):
-        """Handle a log record by adding it to the memory buffer."""
+        """Handle a log record by adding it to the memory buffer.
+        
+        Args:
+            record: The LogRecord to process and store
+        """
         try:
             formatted_message = self.format(record)
             self.logger_manager.add_to_buffer(record, formatted_message)
@@ -305,14 +409,25 @@ class MemoryBufferHandler(logging.Handler):
 _logger_manager = None
 
 def get_logger_manager() -> PromptManagerLogger:
-    """Get the global logger manager instance."""
+    """Get the global logger manager instance.
+    
+    Returns:
+        The singleton PromptManagerLogger instance, creating it if necessary
+    """
     global _logger_manager
     if _logger_manager is None:
         _logger_manager = PromptManagerLogger()
     return _logger_manager
 
 def get_logger(name: str = 'prompt_manager') -> logging.Logger:
-    """Convenience function to get a logger."""
+    """Convenience function to get a logger.
+    
+    Args:
+        name: Logger name, defaults to 'prompt_manager'
+        
+    Returns:
+        Configured logger instance for the specified name
+    """
     return get_logger_manager().get_logger(name)
 
 # Initialize logging on import

utils/metadata_extractor.py

@@ -1,6 +1,33 @@
-"""
-ComfyUI metadata extraction utilities.
-Extracts workflow and prompt information from generated images.
+"""ComfyUI metadata extraction utilities.
+
+This module provides comprehensive metadata extraction capabilities for ComfyUI-generated
+images. It can extract workflow information, prompt data, and generation parameters from
+PNG images that contain embedded ComfyUI metadata in their text chunks.
+
+The extractor supports:
+- Complete workflow data extraction from PNG text chunks
+- Text encoder node identification and analysis
+- Generation parameter extraction (steps, cfg_scale, sampler, etc.)
+- Basic file information fallback when metadata is unavailable
+- Flexible node type detection for various ComfyUI extensions
+
+Typical usage:
+    from utils.metadata_extractor import ComfyUIMetadataExtractor
+    
+    extractor = ComfyUIMetadataExtractor()
+    metadata = extractor.extract_metadata('/path/to/image.png')
+    
+    if metadata:
+        workflow = metadata.get('workflow', {})
+        text_nodes = metadata.get('text_encoder_nodes', [])
+        params = extractor.get_generation_parameters(metadata)
+
+The extracted metadata includes:
+- file_info: Basic file stats (size, dimensions, format, timestamps)
+- workflow: Complete ComfyUI workflow data structure
+- text_encoder_nodes: List of identified text encoding nodes
+- prompt: ComfyUI prompt execution data
+- Generation parameters: steps, cfg_scale, sampler, seed, etc.
 """
 
 import os
@@ -13,21 +40,46 @@ from .logging_config import get_logger
 
 
 class ComfyUIMetadataExtractor:
-    """Extracts ComfyUI metadata from generated images."""
+    """Extracts ComfyUI metadata from generated images.
+    
+    This class handles extraction of ComfyUI workflow and generation metadata
+    from PNG images. It parses the text chunks embedded by ComfyUI and extracts
+    structured information about workflows, prompts, and generation parameters.
+    
+    The extractor is designed to be robust and handle various ComfyUI workflow
+    formats, including custom nodes and different metadata structures.
+    """
     
     def __init__(self):
-        """Initialize the metadata extractor."""
+        """Initialize the metadata extractor.
+        
+        Sets up logging and prepares the extractor for metadata parsing operations.
+        """
         self.logger = get_logger('prompt_manager.metadata_extractor')
     
     def extract_metadata(self, image_path: str) -> Optional[Dict[str, Any]]:
         """
         Extract ComfyUI workflow and prompt metadata from an image.
         
+        This method opens a PNG image and extracts all available ComfyUI metadata
+        from the embedded text chunks. It handles JSON parsing, workflow analysis,
+        and parameter extraction.
+        
         Args:
-            image_path: Path to the image file
+            image_path: Path to the PNG image file to analyze
             
         Returns:
-            Dictionary containing extracted metadata or None if extraction fails
+            Dictionary containing extracted metadata with keys:
+            - file_info: Basic file information (always present)
+            - workflow: Parsed workflow data (if available)
+            - text_encoder_nodes: List of text encoding nodes (if found)
+            - prompt: ComfyUI prompt data (if available)
+            - Additional fields: parameters, model, sampler, steps, etc. (if present)
+            Returns None if no ComfyUI metadata is found or extraction fails
+            
+        Raises:
+            Exception: Re-raises any exception that occurs during extraction,
+                      with appropriate error logging
         """
         try:
             with Image.open(image_path) as image:
@@ -85,12 +137,21 @@ class ComfyUIMetadataExtractor:
         """
         Get basic file information.
         
+        Extracts fundamental file properties including size, dimensions, format,
+        and filesystem timestamps.
+        
         Args:
             image_path: Path to the image file
-            image: PIL Image object
+            image: Opened PIL Image object
             
         Returns:
-            Dictionary containing file information
+            Dictionary containing:
+            - size: File size in bytes
+            - dimensions: Image dimensions as [width, height]
+            - format: Image format (PNG, JPEG, etc.)
+            - mode: Color mode (RGB, RGBA, etc.)
+            - created_time: File creation timestamp
+            - modified_time: File modification timestamp
         """
         try:
             stat = os.stat(image_path)
@@ -110,11 +171,17 @@ class ComfyUIMetadataExtractor:
         """
         Find text encoder nodes in the workflow data.
         
+        Searches through the workflow structure to identify nodes that perform
+        text encoding operations. This includes standard ComfyUI nodes like
+        CLIPTextEncode as well as custom nodes and extensions.
+        
         Args:
-            workflow_data: ComfyUI workflow data
+            workflow_data: ComfyUI workflow data structure (can be in various formats)
             
         Returns:
-            List of text encoder node data
+            List of dictionaries representing text encoder nodes, each containing
+            the node's configuration, inputs, and metadata. For dictionary-based
+            workflows, adds 'node_id' field to each node.
         """
         text_encoder_nodes = []
         
@@ -155,11 +222,15 @@ class ComfyUIMetadataExtractor:
         """
         Check if a node is a text encoder node.
         
+        Determines whether a given node performs text encoding by examining
+        its type, class, and title for known text encoding patterns. Supports
+        various ComfyUI node types and custom extensions.
+        
         Args:
-            node_data: Node data to check
+            node_data: Node configuration dictionary to analyze
             
         Returns:
-            True if the node is a text encoder
+            True if the node is identified as a text encoder, False otherwise
         """
         if not isinstance(node_data, dict):
             return False
@@ -195,11 +266,15 @@ class ComfyUIMetadataExtractor:
         """
         Extract the actual prompt text from workflow data.
         
+        Searches through text encoder nodes to find and extract the actual
+        prompt text that was used for generation. Handles various input field
+        names and data structures.
+        
         Args:
-            workflow_data: ComfyUI workflow data
+            workflow_data: ComfyUI workflow data structure
             
         Returns:
-            Extracted prompt text or None
+            The extracted prompt text string, or None if no prompt text is found
         """
         text_encoder_nodes = self.find_text_encoder_nodes(workflow_data)
         
@@ -223,11 +298,23 @@ class ComfyUIMetadataExtractor:
         """
         Extract generation parameters from metadata.
         
+        Parses the metadata to extract key generation parameters used for
+        image creation, including sampling settings, model information,
+        and other configuration values.
+        
         Args:
-            metadata: Full metadata dictionary
+            metadata: Full metadata dictionary from extract_metadata()
             
         Returns:
-            Dictionary of generation parameters
+            Dictionary containing generation parameters such as:
+            - steps: Number of sampling steps
+            - cfg_scale: Classifier-free guidance scale
+            - sampler: Sampling method used
+            - scheduler: Noise scheduler
+            - seed: Random seed value
+            - model: Model name/path
+            - width/height: Image dimensions
+            - batch_size: Number of images generated
         """
         parameters = {}
         
@@ -252,11 +339,16 @@ class ComfyUIMetadataExtractor:
         """
         Extract generation parameters from workflow data.
         
+        Analyzes the workflow structure to identify and extract generation
+        parameters from various node types. This method can be extended
+        to support specific workflow patterns and custom nodes.
+        
         Args:
-            workflow_data: ComfyUI workflow data
+            workflow_data: ComfyUI workflow data structure
             
         Returns:
-            Dictionary of extracted parameters
+            Dictionary of extracted parameters. Currently returns empty dict
+            but can be extended based on specific workflow analysis needs.
         """
         parameters = {}
         

utils/prompt_tracker.py

@@ -1,6 +1,28 @@
-"""
-Prompt tracking system for linking generated images to prompts.
-Tracks active prompt executions to associate them with generated images.
+"""Prompt tracking system for linking generated images to prompts.
+
+This module provides thread-safe tracking of active prompt executions to enable automatic
+association of generated images with their source prompts. The system maintains prompt
+context across different execution threads and provides fallback mechanisms for reliable
+image-prompt linking.
+
+Key features:
+- Thread-safe prompt tracking using threading.local and locks
+- Automatic cleanup of expired prompt contexts
+- Fallback mechanisms for cross-thread prompt access
+- Support for execution timeouts and extensions
+- Context manager for automatic prompt lifecycle management
+
+Typical usage:
+    from utils.prompt_tracker import PromptTracker
+    
+    tracker = PromptTracker(db_manager)
+    execution_id = tracker.set_current_prompt("beautiful landscape")
+    # Images generated after this point will be linked to this prompt
+    
+Or using the context manager:
+    with PromptExecutionContext(tracker, "beautiful landscape"):
+        # Generate images here
+        pass
 """
 
 import threading
@@ -22,14 +44,29 @@ except ImportError:
 
 
 class PromptTracker:
-    """Thread-safe tracking of current prompt executions."""
+    """Thread-safe tracking of current prompt executions.
+    
+    This class manages active prompt contexts across multiple threads, enabling
+    the image monitoring system to correctly associate generated images with their
+    source prompts. It uses both thread-local storage and global tracking to handle
+    various execution scenarios.
+    
+    The tracker automatically cleans up expired prompts and provides fallback
+    mechanisms when prompts are accessed from different threads (e.g., image
+    monitoring running in a separate thread from prompt execution).
+    
+    Attributes:
+        active_prompts: Global dictionary of active prompt contexts
+        cleanup_interval: Seconds between cleanup operations (default: 300)
+        prompt_timeout: Seconds before a prompt expires (default: 600)
+    """
     
     def __init__(self, db_manager):
         """
         Initialize the prompt tracker.
         
         Args:
-            db_manager: Database manager instance
+            db_manager: Database manager instance for prompt operations
         """
         self.logger = get_logger('prompt_manager.prompt_tracker')
         self.logger.debug("Initializing PromptTracker")
@@ -49,14 +86,19 @@ class PromptTracker:
     def set_current_prompt(self, prompt_text: str, additional_data: Optional[Dict[str, Any]] = None) -> str:
         """
         Set the current prompt for this thread for image tracking.
+        
+        This method establishes the prompt context that will be used to link any
+        subsequently generated images. The prompt is stored in both thread-local
+        storage and global tracking to handle cross-thread access scenarios.
+        
         Note: Prompt saving is handled by PromptManager to avoid duplicates.
         
         Args:
             prompt_text: The prompt text being executed
-            additional_data: Additional prompt metadata (should include prompt_id from PromptManager)
+            additional_data: Additional prompt metadata, should include prompt_id from PromptManager
             
         Returns:
-            Execution ID for this prompt
+            Unique execution ID for this prompt execution
         """
         execution_id = self.generate_execution_id()
         
@@ -109,8 +151,19 @@ class PromptTracker:
         """
         Get the current prompt context for this thread.
         
+        Attempts to retrieve the current prompt context, first from thread-local
+        storage, then from global tracking as a fallback. This enables image
+        monitoring (which may run in a different thread) to access prompt context.
+        
         Returns:
-            Current prompt context or None
+            Dictionary containing prompt context with keys:
+            - id: Prompt ID in database
+            - execution_id: Unique execution identifier
+            - text: Prompt text
+            - timestamp: When the prompt was set
+            - thread_id: Thread identifier
+            - additional_data: Any additional metadata
+            Returns None if no valid prompt context is found
         """
         current = getattr(self._local, 'current_prompt', None)
         
@@ -133,7 +186,15 @@ class PromptTracker:
         return None
     
     def _find_recent_prompt(self) -> Optional[Dict[str, Any]]:
-        """Find the most recent prompt that's still valid."""
+        """Find the most recent prompt that's still valid.
+        
+        Searches the global prompt tracking for the most recently set prompt
+        that hasn't expired. This is used as a fallback when thread-local
+        storage doesn't contain a current prompt.
+        
+        Returns:
+            Most recent valid prompt context, or None if no valid prompts found
+        """
         with self.lock:
             current_time = time.time()
             self.logger.debug(f"Searching for recent prompt among {len(self.active_prompts)} active prompts")
@@ -159,7 +220,12 @@ class PromptTracker:
         return None
     
     def clear_current_prompt(self):
-        """Clear the current prompt context for this thread."""
+        """Clear the current prompt context for this thread.
+        
+        Removes the prompt context from both thread-local storage and global
+        tracking. This should be called when prompt execution is complete,
+        though the system also handles automatic cleanup via timeouts.
+        """
         current = getattr(self._local, 'current_prompt', None)
         if current:
             execution_id = current['execution_id']
@@ -183,9 +249,14 @@ class PromptTracker:
         """
         Extend the timeout for a specific prompt execution.
         
+        This is useful for long-running image generation processes where images
+        may continue to be generated well after the initial prompt execution.
+        Updates the timestamp to prevent the prompt from being cleaned up.
+        
         Args:
             execution_id: The execution ID to extend
-            additional_seconds: Additional seconds to add to timeout
+            additional_seconds: Additional seconds to add to timeout (currently unused,
+                              method just resets the timestamp to current time)
         """
         with self.lock:
             if execution_id in self.active_prompts:
@@ -193,11 +264,22 @@ class PromptTracker:
                 self.logger.debug(f"Extended timeout for prompt: {execution_id}")
     
     def generate_execution_id(self) -> str:
-        """Generate a unique execution ID."""
+        """Generate a unique execution ID.
+        
+        Creates a unique identifier for this prompt execution using UUID and timestamp.
+        
+        Returns:
+            Unique execution ID in format 'exec_{uuid8}_{timestamp}'
+        """
         return f"exec_{uuid.uuid4().hex[:8]}_{int(time.time())}"
     
     def _cleanup_expired_prompts(self):
-        """Background thread to clean up expired prompts."""
+        """Background thread to clean up expired prompts.
+        
+        Runs continuously as a daemon thread, periodically removing expired
+        prompt contexts from the global tracking dictionary. This prevents
+        memory leaks from accumulating old prompt data.
+        """
         while True:
             try:
                 current_time = time.time()
@@ -225,17 +307,21 @@ class PromptTracker:
         Get all currently active prompts.
         
         Returns:
-            Dictionary of active prompts
+            Dictionary mapping execution IDs to prompt context dictionaries.
+            This is a copy of the internal tracking dictionary.
         """
         with self.lock:
             return self.active_prompts.copy()
     
     def clear_all_active_prompts(self) -> int:
         """
-        Clear all active prompts (useful before scanning to avoid conflicts).
+        Clear all active prompts.
+        
+        Removes all prompt contexts from both global and thread-local storage.
+        This is useful before batch operations or when resetting the tracking state.
         
         Returns:
-            Number of prompts that were cleared
+            Number of prompts that were cleared from global tracking
         """
         with self.lock:
             cleared_count = len(self.active_prompts)
@@ -252,7 +338,13 @@ class PromptTracker:
         Get tracker status information.
         
         Returns:
-            Status information dictionary
+            Dictionary containing:
+            - active_prompts_count: Number of currently active prompts
+            - current_prompt_id: ID of current prompt in this thread (if any)
+            - current_execution_id: Execution ID of current prompt in this thread (if any)
+            - thread_id: Current thread identifier
+            - prompt_timeout: Configured prompt timeout in seconds
+            - cleanup_interval: Configured cleanup interval in seconds
         """
         with self.lock:
             active_count = len(self.active_prompts)
@@ -270,16 +362,30 @@ class PromptTracker:
 
 
 class PromptExecutionContext:
-    """Context manager for prompt executions."""
+    """Context manager for prompt executions.
+    
+    This context manager provides automatic prompt lifecycle management,
+    ensuring prompt context is properly set and cleaned up. While the current
+    implementation doesn't automatically clear prompts on exit (to allow for
+    images generated after prompt execution), it provides a clean interface
+    for prompt management.
+    
+    Example:
+        with PromptExecutionContext(tracker, "beautiful landscape") as exec_id:
+            # Generate images here
+            # Images will be linked to this prompt context
+            pass
+        # Prompt context remains active for timeout period
+    """
     
     def __init__(self, prompt_tracker: PromptTracker, prompt_text: str, **kwargs):
         """
         Initialize execution context.
         
         Args:
-            prompt_tracker: PromptTracker instance
-            prompt_text: The prompt text
-            **kwargs: Additional prompt metadata
+            prompt_tracker: PromptTracker instance to use for tracking
+            prompt_text: The prompt text to track
+            **kwargs: Additional prompt metadata to include in the context
         """
         self.prompt_tracker = prompt_tracker
         self.prompt_text = prompt_text
@@ -287,7 +393,13 @@ class PromptExecutionContext:
         self.execution_id = None
     
     def __enter__(self):
-        """Enter the execution context."""
+        """Enter the execution context.
+        
+        Sets the current prompt in the tracker and returns the execution ID.
+        
+        Returns:
+            Unique execution ID for this prompt
+        """
         self.execution_id = self.prompt_tracker.set_current_prompt(
             self.prompt_text, 
             self.additional_data
@@ -295,7 +407,17 @@ class PromptExecutionContext:
         return self.execution_id
     
     def __exit__(self, exc_type, exc_val, exc_tb):
-        """Exit the execution context."""
+        """Exit the execution context.
+        
+        Currently does not clear the prompt context to allow images to be
+        generated after prompt execution completes. The prompt will be cleaned
+        up automatically after the timeout period.
+        
+        Args:
+            exc_type: Exception type (if an exception occurred)
+            exc_val: Exception value (if an exception occurred) 
+            exc_tb: Exception traceback (if an exception occurred)
+        """
         # Don't clear immediately - let the timeout handle it
         # This allows images to be generated after the prompt execution completes
         pass

utils/validators.py

@@ -1,5 +1,33 @@
-"""
-Input validation utilities for KikoTextEncode.
+"""Input validation utilities for PromptManager.
+
+This module provides comprehensive input validation and sanitization functions
+for the PromptManager system. It ensures data integrity and security by validating
+user inputs before they are processed or stored in the database.
+
+Validation functions include:
+- Prompt text validation (length limits, type checking)
+- Rating validation (1-5 scale with None support)
+- Tag validation and parsing (comma-separated strings or lists)
+- Category validation (optional string fields)
+- Workflow name validation
+- Input sanitization and cleaning utilities
+
+All validation functions follow a consistent pattern:
+- Type checking with descriptive error messages
+- Reasonable limits to prevent abuse
+- Support for None/optional values where appropriate
+- Raise ValueError with clear messages on validation failures
+
+Typical usage:
+    from utils.validators import validate_prompt_text, sanitize_input
+    
+    try:
+        validate_prompt_text(user_input)
+        clean_text = sanitize_input(user_input)
+        # Process the validated and cleaned input
+    except ValueError as e:
+        # Handle validation error with user-friendly message
+        print(f"Invalid input: {e}")
 """
 
 import re
@@ -10,14 +38,20 @@ def validate_prompt_text(text: str) -> bool:
     """
     Validate prompt text input.
     
+    Ensures the prompt text is a valid string with reasonable length limits.
+    Empty or whitespace-only strings are rejected.
+    
     Args:
         text: The prompt text to validate
         
     Returns:
-        bool: True if valid, False otherwise
+        True if the text passes validation
         
     Raises:
-        ValueError: If text is invalid with descriptive message
+        ValueError: If text is invalid with descriptive message including:
+                   - Not a string type
+                   - Empty or whitespace-only
+                   - Exceeds maximum length (10,000 characters)
     """
     if not isinstance(text, str):
         raise ValueError("Prompt text must be a string")
@@ -35,14 +69,18 @@ def validate_rating(rating: Optional[int]) -> bool:
     """
     Validate rating input.
     
+    Validates rating values on a 1-5 scale, with None allowed for unrated prompts.
+    
     Args:
-        rating: The rating to validate (1-5 or None)
+        rating: The rating to validate (1-5 scale or None for no rating)
         
     Returns:
-        bool: True if valid, False otherwise
+        True if the rating is valid
         
     Raises:
-        ValueError: If rating is invalid
+        ValueError: If rating is invalid:
+                   - Not an integer (when not None)
+                   - Outside the 1-5 range
     """
     if rating is None:
         return True
@@ -60,14 +98,22 @@ def validate_tags(tags: Union[str, List[str], None]) -> bool:
     """
     Validate tags input.
     
+    Accepts tags as comma-separated string, list of strings, or None.
+    Validates each tag for length and character restrictions.
+    
     Args:
-        tags: Tags as string, list, or None
+        tags: Tags as comma-separated string, list of strings, or None
         
     Returns:
-        bool: True if valid, False otherwise
+        True if all tags are valid
         
     Raises:
-        ValueError: If tags are invalid
+        ValueError: If tags are invalid:
+                   - Wrong input type (not string, list, or None)
+                   - Individual tag is empty or only whitespace
+                   - Individual tag exceeds 50 characters
+                   - Tag contains invalid characters (non-alphanumeric, spaces, hyphens, underscores)
+                   - More than 20 tags provided
     """
     if tags is None:
         return True
@@ -104,14 +150,20 @@ def validate_category(category: Optional[str]) -> bool:
     """
     Validate category input.
     
+    Validates optional category strings with reasonable length limits
+    and character restrictions.
+    
     Args:
-        category: The category to validate
+        category: The category string to validate (None allowed for no category)
         
     Returns:
-        bool: True if valid, False otherwise
+        True if the category is valid or None
         
     Raises:
-        ValueError: If category is invalid
+        ValueError: If category is invalid:
+                   - Not a string type (when not None)
+                   - Exceeds 100 characters
+                   - Contains invalid characters (non-alphanumeric, spaces, hyphens, underscores)
     """
     if category is None:
         return True
@@ -137,14 +189,19 @@ def validate_workflow_name(workflow_name: Optional[str]) -> bool:
     """
     Validate workflow name input.
     
+    Validates optional workflow name strings with generous length limits
+    to accommodate descriptive workflow names.
+    
     Args:
-        workflow_name: The workflow name to validate
+        workflow_name: The workflow name to validate (None allowed for no workflow)
         
     Returns:
-        bool: True if valid, False otherwise
+        True if the workflow name is valid or None
         
     Raises:
-        ValueError: If workflow name is invalid
+        ValueError: If workflow name is invalid:
+                   - Not a string type (when not None)
+                   - Exceeds 200 characters
     """
     if workflow_name is None:
         return True
@@ -166,11 +223,20 @@ def sanitize_input(text: str) -> str:
     """
     Sanitize text input by removing potentially harmful content.
     
+    Cleans text input by removing control characters, normalizing whitespace,
+    and limiting excessive empty lines. Preserves the semantic content while
+    ensuring safe storage and display.
+    
     Args:
-        text: The text to sanitize
+        text: The text string to sanitize
         
     Returns:
-        str: Sanitized text
+        Sanitized text string with:
+        - Null bytes and control characters removed
+        - Normalized line endings (\r\n and \r converted to \n)
+        - Trimmed whitespace on each line
+        - Limited consecutive empty lines (maximum 2)
+        - Overall trimmed result
     """
     if not isinstance(text, str):
         return ""
@@ -202,11 +268,15 @@ def parse_tags_string(tags_string: str) -> List[str]:
     """
     Parse a comma-separated tags string into a clean list.
     
+    Converts comma-separated tag strings into a clean, deduplicated list
+    of tags. Each tag is sanitized and trimmed.
+    
     Args:
-        tags_string: Comma-separated tags string
+        tags_string: Comma-separated tags string (e.g., "tag1, tag2, tag3")
         
     Returns:
-        List[str]: Cleaned list of unique tags
+        List of unique, cleaned tag strings. Empty input returns empty list.
+        Limited to maximum 20 tags to prevent abuse.
     """
     if not tags_string or not isinstance(tags_string, str):
         return []