Vito
3be0ada6a1
* fix(ingestion): link an image only to a prompt that provably produced it An LLM fed a PromptManagerPair's prepend_text while its text editor was empty. The node refused to save (an empty positive raised and logged a warning), and the image monitor then linked the image to a stale tracker entry or to the newest prompt in the database. - Nodes: an empty text editor saves and tracks nothing (prepend and append text are never the stored prompt), logs at info, and marks the node's run key empty so no earlier entry can be used for its images. - Tracker entries carry a run key: a digest of the prompt node and its whole upstream graph, the identity ComfyUI's cache uses. The monitor computes the same key from the image's embedded graph and uses a queued or snapshot entry only when the keys match (a snapshot also needs a known run with one entry, so a batch's last item is never applied to every image). - The recency fallback is gone; anything uncertain stays unlinked and is logged at info. - Recorded execution history attributes a file when exactly one run reported it, the file was written during that run, and the run's positive contexts name one stored prompt. The monitor and the output scan use it as their last certain source. - The scan and pair ingestion treat a resolved pair with an empty stored positive as unresolved instead of raising. - tools/review_image_links.py reviews selected links by these rules; dry run by default, --apply changes only named ids after a backup. Closes #342 * test(ingestion): real-host attribution case and the #342 testing contract The managed host disables the image monitor, so the owned fixture pack gains a route that starts and stops the plugin's monitor singleton for the attribution case only. The case runs in both renderers: a typed Pair image links, a Pair image whose positive is only linked prepend text stays unlinked with nothing stored, an unknown node feeding text links through its run key, and its cached re-run links again. * fix(ingestion): take a queued entry only from the run the image appeared in Dropping a key's older queue entries whenever the node ran again let a later queued run, executing during the monitor's processing delay, take the earlier image's entry. The snapshot taken when the image appears now names the key's latest run, and the queue is popped by key and run: same-key entries of earlier runs ahead of it are stale, entries of later runs wait for their own images. Also keeps the attribution case out of the stock real-host suite. * test(ingestion): host lookups degrade to no proof * fix(ingestion): address the #344 review (history re-check, keys, repair) - Images without an embedded API graph: the scan now asks recorded history before skipping a file with no prompt metadata, and the monitor gives a keyless unlinked image one history-only re-check when its run's outputs are recorded (a seal notification from the outputs write; at most 1000 waiting images, each for an hour). - Run keys: an image is unkeyable when any positive node has no key, so a two-pass graph is never linked through the other pass; links follow ComfyUI's is_link exactly. - A failed save marks the node's run key like an empty editor, so an older entry of that key can't stand in for this run. - A placeholder is the latest entry: the image stays unlinked without falling through to history, the same rule as a popped placeholder. - An index on execution_run_outputs(filename, subfolder, root_key, path_style) for the history lookup, added by migration and accepted by the output schema check (any other index is still refused). - Repair tool: unlink removes only the link record and only while the image is still linked to the prompt it was judged on; every run prints what --apply does. - Real-host case proves the re-run's PromptManager node was cached. * test(ingestion): real-host bare-output re-check and the review-round contract A metadata-disabled save, held open behind the fixture gate after SaveImage, waits unlinked while its run is open and links once the run's outputs are recorded. The fixture's monitor route lists the waiting files. TESTING.md covers the history re-check, the read index and what the repair tool's --apply does. * fix(ingestion): batch history lookups instead of adding an index An index on execution_run_outputs would make the library unopenable by older builds, whose schema check refuses any extra index there. The migration and the schema-check allowance are reverted, so the history schema is exactly main's. stored_run_prompts instead looks up a batch of files with one parameterized filename IN (...) query per chunk of at most 500 names and matches in memory; the scan computes it once per scan batch. * fix(ingestion): address the #344 re-review (seal race, linked roots) - A sealed history re-check that can't be scheduled yet (the image's own pass still pending, or the pending set full) is kept due and scheduled when a pending image finishes, instead of being dropped. - History matches a file through ComfyUI's spelling of its output folder first, then with both sides resolved, so a symlinked or junction output folder still attributes the scan's resolved paths; waiting files are keyed by their resolved path. - The scan makes its batch history lookup only when a file needs it. - set_current_prompt pushes to the queue last, so marking a failed execution never adds a second queue slot. - A keyless image outside the output folder doesn't take a waiting slot.
692 lines
26 KiB
Python
692 lines
26 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
|
|
from collections import OrderedDict
|
|
from typing import Any, Dict, Optional, Tuple
|
|
|
|
# Import logging system
|
|
try:
|
|
from .hashing import generate_prompt_hash
|
|
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.hashing import generate_prompt_hash
|
|
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,
|
|
run_key: Optional[str] = None,
|
|
run_id: Optional[str] = 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
|
|
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.
|
|
run_key: The node's run key (``utils.run_keys.prompt_node_key``). The
|
|
image monitor links an image to a tracked prompt only when the
|
|
image's own graph gives the same key (#342).
|
|
run_id: The ComfyUI prompt id executing the node, when the host says.
|
|
|
|
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:
|
|
# Same hash as the nodes, so the lookup finds their saved prompt
|
|
prompt_hash = generate_prompt_hash(prompt_text)
|
|
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 {},
|
|
"run_key": run_key,
|
|
"run_id": run_id,
|
|
}
|
|
|
|
# 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
|
|
|
|
self._remember_keyed(execution_context)
|
|
|
|
self.logger.debug(
|
|
f"Set current prompt: {execution_id} -> {prompt_text[:50]}... "
|
|
f"(thread: {threading.current_thread().ident})"
|
|
)
|
|
self.logger.debug(f"Active prompts count: {len(self.active_prompts)}")
|
|
|
|
# Push to batch queue for ordered image linking. Last, with nothing
|
|
# after it that can raise: a caller that marks a failed execution
|
|
# (#342) must never add a second queue slot for the same item.
|
|
if push_to_queue:
|
|
with self._queue_lock:
|
|
self._prompt_queue.append(execution_context)
|
|
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(
|
|
"Using recent prompt from global tracking: "
|
|
f"{recent_prompt['execution_id']}"
|
|
)
|
|
return recent_prompt
|
|
|
|
self.logger.debug(
|
|
f"No prompt context found (thread: {threading.current_thread().ident})"
|
|
)
|
|
return None
|
|
|
|
def mark_empty(
|
|
self,
|
|
run_key: Optional[str],
|
|
run_id: Optional[str] = None,
|
|
push_to_queue: bool = False,
|
|
) -> None:
|
|
"""Record that a positive node ran with nothing to save (#342).
|
|
|
|
The run key's latest entry becomes a placeholder, so no earlier prompt of
|
|
that key can be linked to this run's images; with ``push_to_queue`` the
|
|
placeholder also holds this item's place in the batch queue, so a later
|
|
item's image never takes the wrong queue entry.
|
|
"""
|
|
if not run_key:
|
|
return
|
|
placeholder = {
|
|
"id": None,
|
|
"placeholder": True,
|
|
"execution_id": self.generate_execution_id(),
|
|
"text": "",
|
|
"timestamp": time.time(),
|
|
"additional_data": {},
|
|
"run_key": run_key,
|
|
"run_id": run_id,
|
|
}
|
|
self._remember_keyed(placeholder)
|
|
if push_to_queue:
|
|
with self._queue_lock:
|
|
self._prompt_queue.append(placeholder)
|
|
|
|
def _keyed_runs(self) -> "KeyedRuns":
|
|
keyed = getattr(self, "_keyed", None)
|
|
if keyed is None:
|
|
with _KEYED_INIT_LOCK:
|
|
keyed = getattr(self, "_keyed", None)
|
|
if keyed is None:
|
|
keyed = self._keyed = KeyedRuns()
|
|
return keyed
|
|
|
|
def _remember_keyed(self, context: Dict[str, Any]) -> None:
|
|
if context.get("run_key"):
|
|
self._keyed_runs().record(context)
|
|
|
|
def keyed_snapshot(self) -> Dict[str, Tuple[Dict[str, Any], Optional[str], int]]:
|
|
"""``{run key: (latest entry, its run id, entries of that run)}`` now.
|
|
|
|
The image monitor captures it when an image appears, before later runs
|
|
can move on.
|
|
"""
|
|
return self._keyed_runs().snapshot(self.prompt_timeout)
|
|
|
|
def pop_for_key(
|
|
self, run_key: Optional[str], run_id: Optional[str] = None
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Pop the oldest queued entry of ``run_key`` (FIFO), from ``run_id`` if set.
|
|
|
|
``run_id`` is the run the key's latest entry came from when the image
|
|
appeared (``keyed_snapshot``): entries of that key queued by earlier
|
|
runs ahead of it are stale and dropped, while entries of later runs
|
|
(queued after the image appeared) stay for their own images. Entries
|
|
of other keys stay queued; expired entries are discarded.
|
|
"""
|
|
if not run_key:
|
|
return None
|
|
now = time.time()
|
|
with self._queue_lock:
|
|
live = [
|
|
ctx
|
|
for ctx in self._prompt_queue
|
|
if now - ctx.get("timestamp", 0) < self.prompt_timeout
|
|
]
|
|
index = next(
|
|
(
|
|
i
|
|
for i, ctx in enumerate(live)
|
|
if ctx.get("run_key") == run_key
|
|
and (run_id is None or ctx.get("run_id") == run_id)
|
|
),
|
|
None,
|
|
)
|
|
if index is None:
|
|
self._prompt_queue = live
|
|
return None
|
|
found = live[index]
|
|
# Same-key entries ahead of it came from earlier runs: stale
|
|
stale_or_found = {
|
|
i for i in range(index + 1) if live[i].get("run_key") == run_key
|
|
}
|
|
self._prompt_queue = [
|
|
ctx for i, ctx in enumerate(live) if i not in stale_or_found
|
|
]
|
|
return found
|
|
|
|
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, "
|
|
f"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("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} "
|
|
f"(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(
|
|
"No current prompt to clear "
|
|
f"(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,
|
|
}
|
|
|
|
|
|
_KEYED_INIT_LOCK = threading.Lock()
|
|
# Run keys remembered for the image monitor (oldest forgotten first)
|
|
MAX_KEYED_RUNS = 256
|
|
|
|
|
|
class KeyedRuns:
|
|
"""The latest tracked entry per run key, and how many its run produced.
|
|
|
|
More than one entry for a key in one run (a batch fed by a list) means the
|
|
latest one alone cannot say which item an image belongs to.
|
|
"""
|
|
|
|
def __init__(self, max_keys: int = MAX_KEYED_RUNS):
|
|
self._max = max_keys
|
|
self._latest: "OrderedDict[str, Tuple[Dict[str, Any], Optional[str], int]]"
|
|
self._latest = OrderedDict()
|
|
self._lock = threading.Lock()
|
|
|
|
def record(self, context: Dict[str, Any]) -> None:
|
|
key, run_id = context["run_key"], context.get("run_id")
|
|
with self._lock:
|
|
previous = self._latest.pop(key, None)
|
|
same_run = (
|
|
previous is not None and run_id is not None and previous[1] == run_id
|
|
)
|
|
count = previous[2] + 1 if same_run else 1
|
|
self._latest[key] = (context, run_id, count)
|
|
while len(self._latest) > self._max:
|
|
self._latest.popitem(last=False)
|
|
|
|
def snapshot(self, timeout: float) -> Dict[str, Tuple[Dict[str, Any], Any, int]]:
|
|
now = time.time()
|
|
with self._lock:
|
|
return {
|
|
key: entry
|
|
for key, entry in self._latest.items()
|
|
if now - entry[0].get("timestamp", 0) < timeout
|
|
}
|
|
|
|
|
|
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
|