Vito
6c7e4d5b35
* chore: start 3.2.4 * feat: track prompt usage (run_count, last_used_at) with server-side sorting - Migration adds last_used_at and run_count; existing prompts are estimated from linked images (count, newest image time, else created_at). The backfill only touches NULL rows, so it self-heals rows written after a downgrade or an interrupted upgrade. - Every node run records a use, including re-runs of an existing prompt; a tracking failure is logged and never blocks generation. - search_prompts/get_recent_prompts accept a whitelisted sort key (last_used, created, run_count, rating, text); default stays newest first. Refs #90 * feat: sort param on /recent and /search; microsecond last_used_at - /prompt_manager/recent and /search accept ?sort= (whitelisted in the DB layer; unknown values fall back to newest first). - last_used_at is written with microseconds from Python so a re-run in the same millisecond as another save still sorts first; ordering stays correct against millisecond values from the backfill. Refs #90 * feat: server-side sort in admin UI with Recently Used / Most Used and run badges - New web/js/prompt-list-sort.js owns the sort options, validation and request URLs; admin.js fills the dropdown from it and re-fetches from the server on change, so every sort spans all pages (the old client-side sort only reordered the current page). - Default sort is Recently Used; prompt cards show a run-count badge with the last-used time on hover. - The node's Recent button now lists recently used prompts. - Contract test keeps the UI sort keys and the node's key in sync with the server whitelist. Refs #90 * fix: link images and count runs by the positive prompt, not node execution With a negative-prompt PromptManager node in the workflow, every executed PromptManager node pushed onto the image-linking queue, so images could be linked to the negative prompt or to a previous run. ComfyUI also skips unchanged nodes, so re-runs of the same prompt were never counted. - utils/prompt_graph.py finds the PromptManager nodes feeding a sampler's positive input (KSampler, CFG/Basic guiders, ControlNet and combine pass-throughs, PromptManagerText into encoders), falling back to all PromptManager nodes when none can be traced. - Runs are counted by a ComfyUI on-prompt hook when a workflow is queued; prompts not yet saved are counted when the node first saves them. - Nodes receive the queued graph (hidden PROMPT/UNIQUE_ID). Only the run's positive prompt becomes the current prompt; only batch items with linked text are pushed onto the FIFO queue. - The image monitor links from the image's own metadata first, using the positive node, and falls back to the queue for batch items. Refs #90 * fix: address review — count cached runs of linked text, keep repeated pending runs - Text from known pure string nodes (core PrimitiveString/StringConcatenate, WAS Text Concatenate/Multiline, pysssss ShowText) is resolved at queue time, so cached re-runs of such prompts are counted and their images link from metadata. Unknown nodes are never guessed; the node compares its real text with the queue-time resolution and falls back to counting and queueing itself on any mismatch. - A new prompt queued several times before its first run keeps every pending run; the first save records all of them. - Registration flag is a threading.Event; consistent unittest imports; logging instead of print in the init fallback. * refactor: unquote PendingFirstUse annotation so Tuple usage is visible
545 lines
21 KiB
Python
545 lines
21 KiB
Python
"""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
|
|
import time
|
|
import uuid
|
|
import hashlib
|
|
from typing import Optional, Dict, Any
|
|
from datetime import datetime, timezone
|
|
|
|
# Import logging system
|
|
try:
|
|
from .logging_config import get_logger
|
|
except ImportError:
|
|
import sys
|
|
import os
|
|
|
|
current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, current_dir)
|
|
from utils.logging_config import get_logger
|
|
|
|
|
|
class PromptTracker:
|
|
"""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 for prompt operations
|
|
"""
|
|
self.logger = get_logger("prompt_manager.prompt_tracker")
|
|
self.logger.debug("Initializing PromptTracker")
|
|
|
|
self.db_manager = db_manager
|
|
self._local = threading.local()
|
|
self.active_prompts = {} # Global tracking for multiple threads
|
|
self.lock = threading.Lock()
|
|
# FIFO queue for batch workflows: prompt_ids are pushed during encoding
|
|
# and popped during image linking, preserving correct ordering.
|
|
self._prompt_queue = []
|
|
self._queue_lock = threading.Lock()
|
|
|
|
# Read from GalleryConfig if available, otherwise use defaults
|
|
try:
|
|
from ..py.config import GalleryConfig
|
|
|
|
self.cleanup_interval = GalleryConfig.CLEANUP_INTERVAL
|
|
self.prompt_timeout = GalleryConfig.PROMPT_TIMEOUT
|
|
except Exception:
|
|
self.cleanup_interval = 300 # 5 minutes
|
|
self.prompt_timeout = 600 # 10 minutes
|
|
|
|
# Start cleanup thread
|
|
self.cleanup_thread = threading.Thread(
|
|
target=self._cleanup_expired_prompts, daemon=True
|
|
)
|
|
self.cleanup_thread.start()
|
|
self.logger.debug("PromptTracker initialization completed")
|
|
|
|
def set_current_prompt(
|
|
self,
|
|
prompt_text: str,
|
|
additional_data: Optional[Dict[str, Any]] = None,
|
|
push_to_queue: bool = True,
|
|
) -> 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
|
|
push_to_queue: Also push onto the FIFO batch queue. Only batch items whose
|
|
text comes from another node need it; typed prompts are linked from
|
|
the image's own metadata instead.
|
|
|
|
Returns:
|
|
Unique execution ID for this prompt execution
|
|
"""
|
|
execution_id = self.generate_execution_id()
|
|
|
|
# Get prompt_id from additional_data (passed from PromptManager)
|
|
# PromptTracker should NOT save prompts to avoid duplicates
|
|
prompt_id = additional_data.get("prompt_id") if additional_data else None
|
|
|
|
if not prompt_id:
|
|
# Fallback: try to find existing prompt using consistent hash calculation
|
|
try:
|
|
# Use consistent hash calculation (same as PromptManager)
|
|
import hashlib
|
|
|
|
normalized_text = prompt_text.strip().lower()
|
|
prompt_hash = hashlib.sha256(
|
|
normalized_text.encode("utf-8")
|
|
).hexdigest()
|
|
existing_prompt = self.db_manager.get_prompt_by_hash(prompt_hash)
|
|
|
|
if existing_prompt:
|
|
prompt_id = existing_prompt["id"]
|
|
self.logger.debug(f"Found existing prompt ID: {prompt_id}")
|
|
else:
|
|
# Generate a temporary ID for tracking (prompt should be saved by PromptManager)
|
|
prompt_id = f"temp_{int(time.time())}"
|
|
self.logger.debug(f"Using temporary ID for tracking: {prompt_id}")
|
|
except Exception as e:
|
|
self.logger.error(f"Error finding prompt: {e}")
|
|
prompt_id = f"temp_{int(time.time())}"
|
|
|
|
# Create execution context
|
|
execution_context = {
|
|
"id": prompt_id,
|
|
"execution_id": execution_id,
|
|
"text": prompt_text,
|
|
"timestamp": time.time(),
|
|
"thread_id": threading.current_thread().ident,
|
|
"additional_data": additional_data or {},
|
|
}
|
|
|
|
# Store in thread-local storage
|
|
self._local.current_prompt = execution_context
|
|
|
|
# Store in global tracking with thread safety
|
|
with self.lock:
|
|
self.active_prompts[execution_id] = execution_context
|
|
|
|
# Push to batch queue for ordered image linking
|
|
if push_to_queue:
|
|
with self._queue_lock:
|
|
self._prompt_queue.append(execution_context)
|
|
self.logger.debug(
|
|
f"Queue size: {len(self._prompt_queue)} after push for prompt {prompt_id}"
|
|
)
|
|
|
|
self.logger.debug(
|
|
f"Set current prompt: {execution_id} -> {prompt_text[:50]}... (thread: {threading.current_thread().ident})"
|
|
)
|
|
self.logger.debug(f"Active prompts count: {len(self.active_prompts)}")
|
|
return execution_id
|
|
|
|
def get_current_prompt(self) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
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:
|
|
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)
|
|
|
|
if current:
|
|
# Check if prompt hasn't expired
|
|
if time.time() - current["timestamp"] < self.prompt_timeout:
|
|
return current
|
|
else:
|
|
self.logger.debug(f"Current prompt expired: {current['execution_id']}")
|
|
self.clear_current_prompt()
|
|
|
|
# Fallback: try to find recent prompt from global tracking
|
|
# This is crucial for image monitoring which runs in different threads
|
|
recent_prompt = self._find_recent_prompt()
|
|
if recent_prompt:
|
|
self.logger.debug(
|
|
f"Using recent prompt from global tracking: {recent_prompt['execution_id']}"
|
|
)
|
|
return recent_prompt
|
|
|
|
self.logger.debug(
|
|
f"No prompt context found (thread: {threading.current_thread().ident})"
|
|
)
|
|
return None
|
|
|
|
def pop_next_prompt(self) -> Optional[Dict[str, Any]]:
|
|
"""Pop the next prompt from the batch queue (FIFO).
|
|
|
|
In batch workflows, prompts are encoded in order and images are saved
|
|
in the same order. This method pops the oldest queued prompt, ensuring
|
|
each image links to the correct prompt by position.
|
|
|
|
Skips and discards expired entries to prevent stale prompts from
|
|
accumulating across sessions.
|
|
|
|
Returns:
|
|
Prompt context dict, or None if queue is empty or all entries expired
|
|
"""
|
|
now = time.time()
|
|
with self._queue_lock:
|
|
while self._prompt_queue:
|
|
ctx = self._prompt_queue.pop(0)
|
|
if now - ctx.get("timestamp", 0) < self.prompt_timeout:
|
|
self.logger.debug(
|
|
f"Queue pop: prompt {ctx.get('id')} "
|
|
f"({len(self._prompt_queue)} remaining)"
|
|
)
|
|
return ctx
|
|
self.logger.debug(f"Queue discard expired: prompt {ctx.get('id')}")
|
|
return None
|
|
|
|
def _find_recent_prompt(self) -> Optional[Dict[str, Any]]:
|
|
"""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"
|
|
)
|
|
|
|
recent_prompts = []
|
|
for exec_id, prompt in self.active_prompts.items():
|
|
age_seconds = current_time - prompt["timestamp"]
|
|
self.logger.debug(
|
|
f"Prompt {exec_id}: age={age_seconds:.1f}s, timeout={self.prompt_timeout}s"
|
|
)
|
|
|
|
if age_seconds < self.prompt_timeout:
|
|
recent_prompts.append(prompt)
|
|
else:
|
|
self.logger.debug(f"Prompt {exec_id} is expired")
|
|
|
|
if recent_prompts:
|
|
# Return the most recent one
|
|
most_recent = max(recent_prompts, key=lambda p: p["timestamp"])
|
|
self.logger.debug(f"Found recent prompt: {most_recent['execution_id']}")
|
|
return most_recent
|
|
else:
|
|
self.logger.debug(f"No recent prompts found")
|
|
|
|
return None
|
|
|
|
def clear_current_prompt(self):
|
|
"""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"]
|
|
self.logger.debug(
|
|
f"Clearing current prompt: {execution_id} (thread: {threading.current_thread().ident})"
|
|
)
|
|
|
|
# Clear from thread-local storage
|
|
self._local.current_prompt = None
|
|
|
|
# Remove from global tracking
|
|
with self.lock:
|
|
removed = self.active_prompts.pop(execution_id, None)
|
|
if removed:
|
|
self.logger.debug(
|
|
f"Removed prompt from global tracking: {execution_id}"
|
|
)
|
|
self.logger.debug(
|
|
f"Remaining active prompts: {len(self.active_prompts)}"
|
|
)
|
|
else:
|
|
self.logger.debug(
|
|
f"Prompt {execution_id} was not in global tracking"
|
|
)
|
|
else:
|
|
self.logger.debug(
|
|
f"No current prompt to clear (thread: {threading.current_thread().ident})"
|
|
)
|
|
|
|
def extend_prompt_timeout(self, execution_id: str, additional_seconds: int = 60):
|
|
"""
|
|
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 (currently unused,
|
|
method just resets the timestamp to current time)
|
|
"""
|
|
with self.lock:
|
|
if execution_id in self.active_prompts:
|
|
self.active_prompts[execution_id]["timestamp"] = time.time()
|
|
self.logger.debug(f"Extended timeout for prompt: {execution_id}")
|
|
|
|
def generate_execution_id(self) -> str:
|
|
"""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.
|
|
|
|
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()
|
|
expired_ids = []
|
|
|
|
with self.lock:
|
|
for exec_id, prompt_data in self.active_prompts.items():
|
|
if (
|
|
current_time - prompt_data["timestamp"]
|
|
> self.prompt_timeout
|
|
):
|
|
expired_ids.append(exec_id)
|
|
|
|
for exec_id in expired_ids:
|
|
self.active_prompts.pop(exec_id, None)
|
|
|
|
if expired_ids:
|
|
self.logger.debug(f"Cleaned up {len(expired_ids)} expired prompts")
|
|
|
|
# Also prune expired entries from the batch queue
|
|
with self._queue_lock:
|
|
before = len(self._prompt_queue)
|
|
self._prompt_queue = [
|
|
ctx
|
|
for ctx in self._prompt_queue
|
|
if current_time - ctx.get("timestamp", 0) < self.prompt_timeout
|
|
]
|
|
pruned = before - len(self._prompt_queue)
|
|
if pruned:
|
|
self.logger.debug(f"Pruned {pruned} expired queue entries")
|
|
|
|
time.sleep(self.cleanup_interval)
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error in cleanup thread: {e}")
|
|
time.sleep(60) # Wait a minute before retrying
|
|
|
|
def get_active_prompts(self) -> Dict[str, Dict[str, Any]]:
|
|
"""
|
|
Get all currently active prompts.
|
|
|
|
Returns:
|
|
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.
|
|
|
|
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 from global tracking
|
|
"""
|
|
with self.lock:
|
|
cleared_count = len(self.active_prompts)
|
|
self.active_prompts.clear()
|
|
|
|
# Clear thread-local storage as well
|
|
self._local.current_prompt = None
|
|
|
|
self.logger.debug(f"Cleared {cleared_count} active prompts")
|
|
return cleared_count
|
|
|
|
def get_status(self) -> Dict[str, Any]:
|
|
"""
|
|
Get tracker status information.
|
|
|
|
Returns:
|
|
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)
|
|
|
|
current_prompt = self.get_current_prompt()
|
|
|
|
return {
|
|
"active_prompts_count": active_count,
|
|
"current_prompt_id": current_prompt["id"] if current_prompt else None,
|
|
"current_execution_id": (
|
|
current_prompt["execution_id"] if current_prompt else None
|
|
),
|
|
"thread_id": threading.current_thread().ident,
|
|
"prompt_timeout": self.prompt_timeout,
|
|
"cleanup_interval": self.cleanup_interval,
|
|
}
|
|
|
|
|
|
class PromptExecutionContext:
|
|
"""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 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
|
|
self.additional_data = kwargs
|
|
self.execution_id = None
|
|
|
|
def __enter__(self):
|
|
"""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
|
|
)
|
|
return self.execution_id
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
"""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
|
|
|
|
|
|
# Singleton instance management
|
|
_tracker_instance: Optional[PromptTracker] = None
|
|
_tracker_lock = threading.Lock()
|
|
|
|
|
|
def get_prompt_tracker(db_manager) -> PromptTracker:
|
|
"""Get or create the singleton PromptTracker instance.
|
|
|
|
This ensures only one PromptTracker exists across all PromptManager nodes,
|
|
preventing duplicate image linking when multiple nodes are used.
|
|
|
|
Args:
|
|
db_manager: Database manager instance for prompt operations
|
|
|
|
Returns:
|
|
The singleton PromptTracker instance
|
|
"""
|
|
global _tracker_instance
|
|
if _tracker_instance is None:
|
|
with _tracker_lock:
|
|
if _tracker_instance is None:
|
|
_tracker_instance = PromptTracker(db_manager)
|
|
return _tracker_instance
|