Compare commits
7
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9affe5743c | ||
|
|
2439bb136b | ||
|
|
0e8108ae9b | ||
|
|
94a9260f89 | ||
|
|
23b0162a06 | ||
|
|
6099d5eb23 | ||
|
|
1cffa22ae2 |
+19
@@ -111,6 +111,25 @@ except Exception as e:
|
||||
except Exception:
|
||||
print(f"[ComfyUI-PromptManager] Warning: Failed to start image monitoring: {e}")
|
||||
|
||||
# Count prompt runs when workflows are queued (ComfyUI skips unchanged nodes, so
|
||||
# node execution alone misses re-runs). Separate from monitoring so either can fail.
|
||||
try:
|
||||
from .py.config import server_instance
|
||||
from .utils.usage_tracking import register_queue_hook
|
||||
|
||||
register_queue_hook(server_instance, lambda: _global_db)
|
||||
except Exception as e:
|
||||
try:
|
||||
from .utils.logging_config import get_logger
|
||||
|
||||
get_logger("prompt_manager.init").error(f"Failed to register usage hook: {e}")
|
||||
except Exception:
|
||||
import logging
|
||||
|
||||
logging.getLogger("prompt_manager.init").error(
|
||||
f"Failed to register usage hook: {e}"
|
||||
)
|
||||
|
||||
__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
||||
|
||||
# Print startup message with loaded tools
|
||||
|
||||
+65
-1
@@ -2,6 +2,7 @@
|
||||
Database schema and models for KikoTextEncode prompt storage.
|
||||
"""
|
||||
|
||||
import datetime
|
||||
import sqlite3
|
||||
import os
|
||||
import threading
|
||||
@@ -17,6 +18,18 @@ except ImportError:
|
||||
sys.path.insert(0, current_dir)
|
||||
from utils.logging_config import get_logger
|
||||
|
||||
# last_used_at is UTC ISO text so ORDER BY can compare it as text. The backfill
|
||||
# writes milliseconds (SQLite's %f); live writes use microseconds so runs in the
|
||||
# same millisecond still order correctly. Both forms sort correctly together.
|
||||
ISO_TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%f+00:00"
|
||||
|
||||
|
||||
def utc_now_iso() -> str:
|
||||
"""Current UTC time as ISO text with microseconds, e.g. 2026-10-07T18:39:06.758123+00:00."""
|
||||
return datetime.datetime.now(datetime.timezone.utc).isoformat(
|
||||
timespec="microseconds"
|
||||
)
|
||||
|
||||
|
||||
class PromptModel:
|
||||
"""Database model for prompt storage and schema management."""
|
||||
@@ -78,7 +91,9 @@ class PromptModel:
|
||||
tags TEXT,
|
||||
rating INTEGER CHECK(rating >= 1 AND rating <= 5),
|
||||
notes TEXT,
|
||||
hash TEXT UNIQUE
|
||||
hash TEXT UNIQUE,
|
||||
last_used_at TIMESTAMP,
|
||||
run_count INTEGER NOT NULL DEFAULT 0
|
||||
)
|
||||
""")
|
||||
|
||||
@@ -131,6 +146,9 @@ class PromptModel:
|
||||
# Migrate JSON tags to normalized junction tables
|
||||
self._migrate_json_tags_to_junction(conn)
|
||||
|
||||
# Usage tracking columns (last, after any migration that rebuilds prompts)
|
||||
self._migrate_add_usage_columns(conn)
|
||||
|
||||
def _create_indexes(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
Create indexes for better query performance.
|
||||
@@ -150,6 +168,8 @@ class PromptModel:
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_created_at ON prompts(created_at)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_hash ON prompts(hash)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_rating ON prompts(rating)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_last_used ON prompts(last_used_at)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompts_run_count ON prompts(run_count)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_prompt_images ON generated_images(prompt_id)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_image_path ON generated_images(image_path)",
|
||||
"CREATE INDEX IF NOT EXISTS idx_generation_time ON generated_images(generation_time)",
|
||||
@@ -440,6 +460,50 @@ class PromptModel:
|
||||
except Exception as e:
|
||||
self.logger.error(f"Tag junction migration error: {e}")
|
||||
|
||||
def _migrate_add_usage_columns(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
Add last_used_at and run_count (3.2.4) and fill in any prompt without usage data.
|
||||
|
||||
The backfill runs on every start but only touches rows whose last_used_at
|
||||
is NULL, so it is a no-op once healed. That also repairs rows written by an
|
||||
older version after a downgrade, or left unfilled by an interrupted upgrade.
|
||||
|
||||
The estimate comes from linked images: run_count is the image count (at
|
||||
least 1) and last_used_at the newest image time, falling back to created_at.
|
||||
Timestamps are compared with julianday() because prompts use ISO 'T'
|
||||
timestamps while generated_images uses SQLite's CURRENT_TIMESTAMP format.
|
||||
"""
|
||||
try:
|
||||
columns = {row[1] for row in conn.execute("PRAGMA table_info(prompts)")}
|
||||
if "last_used_at" not in columns:
|
||||
self.logger.info("Adding prompt usage column last_used_at")
|
||||
conn.execute("ALTER TABLE prompts ADD COLUMN last_used_at TIMESTAMP")
|
||||
if "run_count" not in columns:
|
||||
self.logger.info("Adding prompt usage column run_count")
|
||||
conn.execute(
|
||||
"ALTER TABLE prompts ADD COLUMN run_count INTEGER NOT NULL DEFAULT 0"
|
||||
)
|
||||
|
||||
cursor = conn.execute(f"""
|
||||
UPDATE prompts SET
|
||||
run_count = MAX(run_count, 1, (
|
||||
SELECT COUNT(*) FROM generated_images gi
|
||||
WHERE gi.prompt_id = prompts.id
|
||||
)),
|
||||
last_used_at = strftime('{ISO_TIMESTAMP_FORMAT}', MAX(
|
||||
COALESCE(julianday(created_at), julianday('now')),
|
||||
COALESCE((
|
||||
SELECT MAX(julianday(gi.generation_time))
|
||||
FROM generated_images gi WHERE gi.prompt_id = prompts.id
|
||||
), 0)
|
||||
))
|
||||
WHERE last_used_at IS NULL
|
||||
""")
|
||||
if cursor.rowcount:
|
||||
self.logger.info(f"Estimated usage for {cursor.rowcount} prompts")
|
||||
except Exception as e:
|
||||
self.logger.error(f"Usage tracking migration error: {e}")
|
||||
|
||||
def migrate_database(self) -> None:
|
||||
"""
|
||||
Apply any pending database migrations.
|
||||
|
||||
+53
-6
@@ -8,7 +8,7 @@ import datetime
|
||||
import os
|
||||
from typing import Optional, List, Dict, Any, Union
|
||||
|
||||
from .models import PromptModel
|
||||
from .models import PromptModel, utc_now_iso
|
||||
|
||||
# Import logging system
|
||||
try:
|
||||
@@ -28,6 +28,25 @@ TAG_SUBQUERY = (
|
||||
"WHERE pt.prompt_id = prompts.id) AS _tag_list"
|
||||
)
|
||||
|
||||
# Server-side sort orders, keyed by the value the UI sends. Only these strings ever
|
||||
# reach ORDER BY; unknown keys fall back to DEFAULT_SORT.
|
||||
SORT_ORDERS = {
|
||||
"last_used_desc": "last_used_at DESC, id DESC",
|
||||
"created_desc": "created_at DESC, id DESC",
|
||||
"created_asc": "created_at ASC, id ASC",
|
||||
"run_count_desc": "run_count DESC, last_used_at DESC, id DESC",
|
||||
"rating_desc": "rating IS NULL, rating DESC, id DESC",
|
||||
"rating_asc": "rating IS NULL, rating ASC, id DESC",
|
||||
"text_asc": "text COLLATE NOCASE ASC, id ASC",
|
||||
"text_desc": "text COLLATE NOCASE DESC, id DESC",
|
||||
}
|
||||
DEFAULT_SORT = "created_desc"
|
||||
|
||||
|
||||
def order_by_clause(sort: Optional[str]) -> str:
|
||||
"""Return a whitelisted ORDER BY expression for a UI sort key."""
|
||||
return SORT_ORDERS.get(sort or DEFAULT_SORT, SORT_ORDERS[DEFAULT_SORT])
|
||||
|
||||
|
||||
def _resolve_db_path(db_path: Optional[str] = None) -> str:
|
||||
"""Resolve the database path from config, falling back to defaults.
|
||||
@@ -122,8 +141,9 @@ class PromptDatabase:
|
||||
cursor = conn.execute(
|
||||
"""
|
||||
INSERT INTO prompts (
|
||||
text, category, tags, rating, notes, hash, created_at, updated_at
|
||||
) VALUES (?, ?, NULL, ?, ?, ?, ?, ?)
|
||||
text, category, tags, rating, notes, hash, created_at, updated_at,
|
||||
last_used_at
|
||||
) VALUES (?, ?, NULL, ?, ?, ?, ?, ?, ?)
|
||||
""",
|
||||
(
|
||||
text.strip(),
|
||||
@@ -133,6 +153,7 @@ class PromptDatabase:
|
||||
prompt_hash,
|
||||
datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
||||
datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
||||
utc_now_iso(),
|
||||
),
|
||||
)
|
||||
prompt_id = cursor.lastrowid
|
||||
@@ -142,6 +163,26 @@ class PromptDatabase:
|
||||
self.logger.debug(f"Successfully saved prompt with ID: {prompt_id}")
|
||||
return prompt_id
|
||||
|
||||
def record_prompt_use(self, prompt_id: int, times: int = 1) -> bool:
|
||||
"""
|
||||
Count runs of a prompt and mark it as most recently used.
|
||||
|
||||
Args:
|
||||
prompt_id: The prompt ID
|
||||
times: Number of runs to add (queued runs counted on first save)
|
||||
|
||||
Returns:
|
||||
True if the prompt exists and was updated
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"UPDATE prompts SET run_count = run_count + ?, last_used_at = ?"
|
||||
" WHERE id = ?",
|
||||
(max(1, int(times)), utc_now_iso(), prompt_id),
|
||||
)
|
||||
conn.commit()
|
||||
return cursor.rowcount > 0
|
||||
|
||||
def get_prompt_by_id(self, prompt_id: int) -> Optional[Dict[str, Any]]:
|
||||
"""
|
||||
Get a prompt by its ID.
|
||||
@@ -191,6 +232,7 @@ class PromptDatabase:
|
||||
limit: int = 100,
|
||||
offset: int = 0,
|
||||
tag_partial: bool = False,
|
||||
sort: Optional[str] = None,
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""
|
||||
Search prompts with various filters.
|
||||
@@ -207,6 +249,7 @@ class PromptDatabase:
|
||||
limit: Maximum number of results
|
||||
offset: Number of results to skip
|
||||
tag_partial: Use LIKE matching for tags instead of exact match
|
||||
sort: Sort key from SORT_ORDERS (default: newest first)
|
||||
|
||||
Returns:
|
||||
List of dictionaries containing prompt data
|
||||
@@ -275,7 +318,7 @@ class PromptDatabase:
|
||||
query_parts.append("AND created_at <= ?")
|
||||
params.append(date_to)
|
||||
|
||||
query_parts.append("ORDER BY created_at DESC LIMIT ? OFFSET ?")
|
||||
query_parts.append(f"ORDER BY {order_by_clause(sort)} LIMIT ? OFFSET ?")
|
||||
params.extend([limit, offset])
|
||||
|
||||
query = " ".join(query_parts)
|
||||
@@ -291,13 +334,16 @@ class PromptDatabase:
|
||||
|
||||
return prompts
|
||||
|
||||
def get_recent_prompts(self, limit: int = 10, offset: int = 0) -> Dict[str, Any]:
|
||||
def get_recent_prompts(
|
||||
self, limit: int = 10, offset: int = 0, sort: Optional[str] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""
|
||||
Get the most recent prompts with pagination support.
|
||||
|
||||
Args:
|
||||
limit: Maximum number of prompts to return
|
||||
offset: Number of prompts to skip (for pagination)
|
||||
sort: Sort key from SORT_ORDERS (default: newest first)
|
||||
|
||||
Returns:
|
||||
Dictionary containing prompt data and pagination info
|
||||
@@ -310,7 +356,8 @@ class PromptDatabase:
|
||||
|
||||
# Get paginated results
|
||||
cursor = conn.execute(
|
||||
f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts ORDER BY created_at DESC LIMIT ? OFFSET ?",
|
||||
f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts"
|
||||
f" ORDER BY {order_by_clause(sort)} LIMIT ? OFFSET ?",
|
||||
(limit, offset),
|
||||
)
|
||||
rows = cursor.fetchall()
|
||||
|
||||
+23
-21
@@ -4,7 +4,7 @@ with persistent prompt storage and search capabilities.
|
||||
"""
|
||||
|
||||
import time
|
||||
from typing import Any, Tuple
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
try:
|
||||
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
|
||||
@@ -94,6 +94,12 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
|
||||
},
|
||||
),
|
||||
},
|
||||
"hidden": {
|
||||
# The queued graph and this node's id: used to tell the positive
|
||||
# prompt of a run from negative-prompt PromptManager nodes
|
||||
"prompt": "PROMPT",
|
||||
"unique_id": "UNIQUE_ID",
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (IO.CONDITIONING, IO.STRING)
|
||||
@@ -119,6 +125,8 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
|
||||
search_text: str = "",
|
||||
prepend_text: str = "",
|
||||
append_text: str = "",
|
||||
prompt: Optional[Dict[str, Any]] = None,
|
||||
unique_id: Optional[str] = None,
|
||||
) -> Tuple[Any]:
|
||||
"""
|
||||
Encode the text prompt and save it to the database.
|
||||
@@ -181,30 +189,24 @@ class PromptManager(PromptManagerBase, ComfyNodeABC):
|
||||
extended_tags.append(f"append:{append_text.strip()[:50]}")
|
||||
|
||||
try:
|
||||
prompt_id = self._save_prompt_to_database(
|
||||
prompt_id = self._track_prompt_execution(
|
||||
text=storage_text.strip(),
|
||||
encoding_text=encoding_text.strip(),
|
||||
category=category.strip() if category else None,
|
||||
tags=extended_tags if extended_tags else None,
|
||||
additional_data={
|
||||
"category": category.strip() if category else None,
|
||||
"tags": extended_tags,
|
||||
"prepend_text": (
|
||||
prepend_text.strip() if prepend_text else None
|
||||
),
|
||||
"append_text": append_text.strip() if append_text else None,
|
||||
"final_text": encoding_text.strip(),
|
||||
},
|
||||
prompt_graph=prompt,
|
||||
unique_id=unique_id,
|
||||
)
|
||||
|
||||
# Set current prompt for image tracking
|
||||
if prompt_id:
|
||||
execution_id = self.prompt_tracker.set_current_prompt(
|
||||
prompt_text=encoding_text.strip(),
|
||||
additional_data={
|
||||
"category": category.strip() if category else None,
|
||||
"tags": extended_tags,
|
||||
"prompt_id": prompt_id,
|
||||
"prepend_text": (
|
||||
prepend_text.strip() if prepend_text else None
|
||||
),
|
||||
"append_text": append_text.strip() if append_text else None,
|
||||
"final_text": encoding_text.strip(),
|
||||
},
|
||||
)
|
||||
self.logger.debug(
|
||||
f"Set execution context: {execution_id} for prompt ID: {prompt_id}"
|
||||
)
|
||||
self.logger.debug(f"Tracked prompt ID: {prompt_id}")
|
||||
|
||||
except Exception as e:
|
||||
# Log error but don't fail the encoding
|
||||
|
||||
+84
-1
@@ -24,6 +24,9 @@ try:
|
||||
from .utils.comfyui_integration import get_comfyui_integration
|
||||
from .utils.image_monitor import get_image_monitor
|
||||
from .utils.prompt_tracker import get_prompt_tracker
|
||||
from .utils.hashing import generate_prompt_hash
|
||||
from .utils.prompt_graph import resolve_text, run_prompt_nodes
|
||||
from .utils.usage_tracking import PENDING_FIRST_USE
|
||||
except ImportError:
|
||||
import sys
|
||||
|
||||
@@ -32,6 +35,9 @@ except ImportError:
|
||||
from utils.comfyui_integration import get_comfyui_integration
|
||||
from utils.image_monitor import get_image_monitor
|
||||
from utils.prompt_tracker import get_prompt_tracker
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.prompt_graph import resolve_text, run_prompt_nodes
|
||||
from utils.usage_tracking import PENDING_FIRST_USE
|
||||
|
||||
|
||||
class PromptManagerBase:
|
||||
@@ -56,14 +62,23 @@ class PromptManagerBase:
|
||||
self.logger.debug(f"{self.__class__.__name__} node initialization completed")
|
||||
|
||||
def _save_prompt_to_database(
|
||||
self, text: str, category: Optional[str] = None, tags: Optional[list] = None
|
||||
self,
|
||||
text: str,
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[list] = None,
|
||||
count_run: bool = False,
|
||||
) -> Optional[int]:
|
||||
"""Save the prompt to the SQLite database.
|
||||
|
||||
Runs are normally counted by the queue hook (utils/usage_tracking.py), so a
|
||||
save only counts one when asked to (batch items whose text the hook could not
|
||||
see) or when the hook queued this prompt before it existed.
|
||||
|
||||
Args:
|
||||
text: The prompt text
|
||||
category: Optional category
|
||||
tags: List of tags
|
||||
count_run: Count this execution as a run of the prompt
|
||||
|
||||
Returns:
|
||||
The prompt ID if saved successfully, None otherwise
|
||||
@@ -82,6 +97,8 @@ class PromptManagerBase:
|
||||
prompt_id=existing["id"], category=category, tags=tags
|
||||
)
|
||||
self.logger.debug("Updated metadata for existing prompt")
|
||||
if count_run:
|
||||
self._record_use(existing["id"])
|
||||
return existing["id"]
|
||||
|
||||
self.logger.debug(
|
||||
@@ -93,6 +110,10 @@ class PromptManagerBase:
|
||||
|
||||
if prompt_id:
|
||||
self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}")
|
||||
pending = getattr(self, "pending_first_use", PENDING_FIRST_USE)
|
||||
runs = pending.consume(prompt_hash) + (1 if count_run else 0)
|
||||
if runs:
|
||||
self._record_use(prompt_id, times=runs)
|
||||
else:
|
||||
self.logger.warning("Failed to save prompt - no ID returned")
|
||||
|
||||
@@ -102,6 +123,68 @@ class PromptManagerBase:
|
||||
self.logger.error(f"Error saving prompt to database: {e}")
|
||||
return None
|
||||
|
||||
def _prompt_role(
|
||||
self, prompt_graph: Any, unique_id: Any, text: str
|
||||
) -> Dict[str, bool]:
|
||||
"""How this node's execution relates to the run, from the queued graph.
|
||||
|
||||
Returns flags:
|
||||
is_run_prompt: feeds a sampler's positive input (images belong to it)
|
||||
seen_at_queue: the queue hook resolved exactly this text, so it already
|
||||
counted the run and the image's metadata can identify the prompt
|
||||
known: the graph and node id were available
|
||||
"""
|
||||
if not isinstance(prompt_graph, dict) or unique_id is None:
|
||||
return {"is_run_prompt": True, "seen_at_queue": False, "known": False}
|
||||
node_id = str(unique_id)
|
||||
resolved = resolve_text(prompt_graph, node_id)
|
||||
return {
|
||||
"is_run_prompt": node_id in run_prompt_nodes(prompt_graph),
|
||||
"seen_at_queue": resolved is not None
|
||||
and generate_prompt_hash(resolved) == generate_prompt_hash(text),
|
||||
"known": True,
|
||||
}
|
||||
|
||||
def _track_prompt_execution(
|
||||
self,
|
||||
text: str,
|
||||
encoding_text: str,
|
||||
category: Optional[str],
|
||||
tags: Optional[list],
|
||||
additional_data: Dict[str, Any],
|
||||
prompt_graph: Any = None,
|
||||
unique_id: Any = None,
|
||||
) -> Optional[int]:
|
||||
"""Save the prompt and, for the run's positive prompt, register it for linking.
|
||||
|
||||
Negative-prompt nodes are saved but never become the prompt their images are
|
||||
linked to. Text the queue hook could resolve is counted there and linked from
|
||||
the image's own metadata; anything else (PromptSearchList batches, unknown
|
||||
text nodes, or a resolution that doesn't match) is counted here and pushed
|
||||
onto the FIFO queue. Without graph context the pre-3.2.4 behaviour is kept.
|
||||
"""
|
||||
role = self._prompt_role(prompt_graph, unique_id, text)
|
||||
batch_item = (
|
||||
role["known"] and role["is_run_prompt"] and not role["seen_at_queue"]
|
||||
)
|
||||
prompt_id = self._save_prompt_to_database(
|
||||
text=text, category=category, tags=tags, count_run=batch_item
|
||||
)
|
||||
if prompt_id and role["is_run_prompt"]:
|
||||
self.prompt_tracker.set_current_prompt(
|
||||
prompt_text=encoding_text,
|
||||
additional_data={**additional_data, "prompt_id": prompt_id},
|
||||
push_to_queue=batch_item or not role["known"],
|
||||
)
|
||||
return prompt_id
|
||||
|
||||
def _record_use(self, prompt_id: int, times: int = 1) -> None:
|
||||
"""Count runs; a tracking failure must never block generation."""
|
||||
try:
|
||||
self.db.record_prompt_use(prompt_id, times=times)
|
||||
except Exception as e:
|
||||
self.logger.warning(f"Could not record use of prompt {prompt_id}: {e}")
|
||||
|
||||
def _inject_lora_trigger_words(self, text: str) -> str:
|
||||
"""Append LoRA trigger words if integration is enabled.
|
||||
|
||||
|
||||
+22
-20
@@ -4,7 +4,7 @@ without CLIP encoding, while maintaining all database and search features.
|
||||
"""
|
||||
|
||||
import time
|
||||
from typing import Tuple
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
try:
|
||||
from comfy.comfy_types import IO, ComfyNodeABC, InputTypeDict
|
||||
@@ -94,6 +94,12 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
|
||||
},
|
||||
),
|
||||
},
|
||||
"hidden": {
|
||||
# The queued graph and this node's id: used to tell the positive
|
||||
# prompt of a run from negative-prompt PromptManager nodes
|
||||
"prompt": "PROMPT",
|
||||
"unique_id": "UNIQUE_ID",
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (IO.STRING,)
|
||||
@@ -117,6 +123,8 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
|
||||
search_text: str = "",
|
||||
prepend_text: str = "",
|
||||
append_text: str = "",
|
||||
prompt: Optional[Dict[str, Any]] = None,
|
||||
unique_id: Optional[str] = None,
|
||||
) -> Tuple[str]:
|
||||
"""
|
||||
Process the text prompt and save it to the database.
|
||||
@@ -162,29 +170,23 @@ class PromptManagerText(PromptManagerBase, ComfyNodeABC):
|
||||
extended_tags.append(f"append:{append_text.strip()[:50]}")
|
||||
|
||||
try:
|
||||
prompt_id = self._save_prompt_to_database(
|
||||
prompt_id = self._track_prompt_execution(
|
||||
text=storage_text.strip(),
|
||||
encoding_text=final_text.strip(),
|
||||
category=category.strip() if category else None,
|
||||
tags=extended_tags if extended_tags else None,
|
||||
additional_data={
|
||||
"category": category.strip() if category else None,
|
||||
"tags": extended_tags,
|
||||
"prepend_text": (
|
||||
prepend_text.strip() if prepend_text else None
|
||||
),
|
||||
"append_text": append_text.strip() if append_text else None,
|
||||
},
|
||||
prompt_graph=prompt,
|
||||
unique_id=unique_id,
|
||||
)
|
||||
|
||||
# Set current prompt for image tracking
|
||||
if prompt_id:
|
||||
execution_id = self.prompt_tracker.set_current_prompt(
|
||||
prompt_text=final_text.strip(),
|
||||
additional_data={
|
||||
"category": category.strip() if category else None,
|
||||
"tags": extended_tags,
|
||||
"prompt_id": prompt_id,
|
||||
"prepend_text": (
|
||||
prepend_text.strip() if prepend_text else None
|
||||
),
|
||||
"append_text": append_text.strip() if append_text else None,
|
||||
},
|
||||
)
|
||||
self.logger.debug(
|
||||
f"Set execution context: {execution_id} for prompt ID: {prompt_id}"
|
||||
)
|
||||
self.logger.debug(f"Tracked prompt ID: {prompt_id}")
|
||||
|
||||
except Exception as e:
|
||||
# Log error but don't fail the processing
|
||||
|
||||
+6
-2
@@ -132,6 +132,8 @@ class PromptRoutesMixin:
|
||||
limit = int(request.query.get("limit", 50))
|
||||
|
||||
folder = request.query.get("folder", "").strip() or None
|
||||
# Validated against a whitelist in the database layer
|
||||
sort = request.query.get("sort") or None
|
||||
|
||||
tags = None
|
||||
if tags_str:
|
||||
@@ -150,6 +152,7 @@ class PromptRoutesMixin:
|
||||
rating_min=min_rating,
|
||||
limit=limit,
|
||||
folder=folder,
|
||||
sort=sort,
|
||||
)
|
||||
self._enrich_prompt_images(results)
|
||||
|
||||
@@ -182,7 +185,7 @@ class PromptRoutesMixin:
|
||||
return web.json_response({"success": False, "error": str(e)}, status=500)
|
||||
|
||||
async def get_recent_prompts(self, request):
|
||||
"""Retrieve recently created prompts with pagination support."""
|
||||
"""Retrieve prompts with pagination and an optional sort (default newest first)."""
|
||||
try:
|
||||
limit = int(request.query.get("limit", 50))
|
||||
page = int(request.query.get("page", 1))
|
||||
@@ -196,8 +199,9 @@ class PromptRoutesMixin:
|
||||
elif limit < 1:
|
||||
limit = 1
|
||||
|
||||
sort = request.query.get("sort") or None
|
||||
results = await self._run_in_executor(
|
||||
self.db.get_recent_prompts, limit=limit, offset=offset
|
||||
self.db.get_recent_prompts, limit=limit, offset=offset, sort=sort
|
||||
)
|
||||
self._enrich_prompt_images(results["prompts"])
|
||||
|
||||
|
||||
+1
-1
@@ -1,7 +1,7 @@
|
||||
[project]
|
||||
name = "promptmanager"
|
||||
description = "A powerful ComfyUI custom node that extends the standard text encoder with persistent prompt storage, advanced search capabilities, and an automatic image gallery system using SQLite."
|
||||
version = "3.2.3"
|
||||
version = "3.2.4"
|
||||
license = {file = "LICENSE"}
|
||||
dependencies = ["# Core dependencies for PromptManager", "# Note: Most dependencies are already included with ComfyUI", "# Already included with Python standard library:", "# - sqlite3", "# - hashlib", "# - json", "# - datetime", "# - os", "# - typing", "# - threading", "# - uuid", "# Required for gallery functionality:", "watchdog>=2.1.0 # For file system monitoring", "Pillow>=8.0.0 # For image metadata extraction (usually included with ComfyUI)", "# Optional dependencies for enhanced search functionality:", "# fuzzywuzzy[speedup]>=0.18.0 # For fuzzy string matching (optional)", "# sqlalchemy>=1.4.0 # For advanced ORM features (optional)", "# Development dependencies (optional):", "# pytest>=6.0.0 # For running tests", "# black>=22.0.0 # For code formatting", "# flake8>=4.0.0 # For linting", "# mypy>=0.910 # For type checking"]
|
||||
|
||||
|
||||
@@ -0,0 +1,66 @@
|
||||
// Run with: node --test tests/js/
|
||||
const test = require("node:test");
|
||||
const assert = require("node:assert/strict");
|
||||
const sort = require("../../web/js/prompt-list-sort.js");
|
||||
|
||||
test("default sort is recently used and is a listed option", () => {
|
||||
assert.equal(sort.DEFAULT_SORT, "last_used_desc");
|
||||
assert.ok(sort.SORT_OPTIONS.some((o) => o.value === sort.DEFAULT_SORT));
|
||||
});
|
||||
|
||||
test("options include Recently used and Most used, each with a label", () => {
|
||||
const values = sort.SORT_OPTIONS.map((o) => o.value);
|
||||
assert.ok(values.includes("last_used_desc"));
|
||||
assert.ok(values.includes("run_count_desc"));
|
||||
for (const option of sort.SORT_OPTIONS) assert.ok(option.label.trim(), option.value);
|
||||
assert.equal(new Set(values).size, values.length, "values are unique");
|
||||
});
|
||||
|
||||
test("options cannot be mutated by callers", () => {
|
||||
assert.throws(() => sort.SORT_OPTIONS.push({ value: "x", label: "x" }));
|
||||
// Non-strict code fails silently on frozen objects, so check the value survived
|
||||
const first = sort.SORT_OPTIONS[0].value;
|
||||
sort.SORT_OPTIONS[0].value = "hacked";
|
||||
assert.equal(sort.SORT_OPTIONS[0].value, first);
|
||||
assert.ok(sort.SORT_OPTIONS.every(Object.isFrozen));
|
||||
});
|
||||
|
||||
test("normalizeSort keeps known values and falls back for anything else", () => {
|
||||
assert.equal(sort.normalizeSort("run_count_desc"), "run_count_desc");
|
||||
for (const bad of [undefined, null, "", "created_at; DROP TABLE prompts", 42, "LAST_USED_DESC"]) {
|
||||
assert.equal(sort.normalizeSort(bad), sort.DEFAULT_SORT, String(bad));
|
||||
}
|
||||
});
|
||||
|
||||
test("buildRecentUrl computes offset from page and always sends a valid sort", () => {
|
||||
assert.equal(
|
||||
sort.buildRecentUrl({ page: 3, limit: 25, sort: "rating_desc" }),
|
||||
"/prompt_manager/recent?limit=25&offset=50&page=3&sort=rating_desc",
|
||||
);
|
||||
assert.equal(
|
||||
sort.buildRecentUrl({ sort: "bogus" }),
|
||||
"/prompt_manager/recent?limit=50&offset=0&page=1&sort=last_used_desc",
|
||||
);
|
||||
});
|
||||
|
||||
test("withSort returns a copy with the sort set and leaves the input untouched", () => {
|
||||
const params = new URLSearchParams({ text: "castle", sort: "text_asc" });
|
||||
const result = sort.withSort(params, "run_count_desc");
|
||||
assert.equal(result.get("sort"), "run_count_desc");
|
||||
assert.equal(result.get("text"), "castle");
|
||||
assert.equal(params.get("sort"), "text_asc");
|
||||
assert.equal(sort.withSort(new URLSearchParams(), "nope").get("sort"), sort.DEFAULT_SORT);
|
||||
});
|
||||
|
||||
test("formatRunCount labels zero, one and many runs", () => {
|
||||
assert.equal(sort.formatRunCount(0), "Never run");
|
||||
assert.equal(sort.formatRunCount(1), "1 run");
|
||||
assert.equal(sort.formatRunCount(642), "642 runs");
|
||||
assert.equal(sort.formatRunCount(1234), "1,234 runs");
|
||||
});
|
||||
|
||||
test("formatRunCount hides the badge for missing or invalid counts", () => {
|
||||
for (const bad of [undefined, null, -1, 1.5, "3", NaN]) {
|
||||
assert.equal(sort.formatRunCount(bad), "", String(bad));
|
||||
}
|
||||
});
|
||||
@@ -238,6 +238,65 @@ class TestUpdatePrompt(APITestCase):
|
||||
self.assertTrue(data["success"])
|
||||
|
||||
|
||||
class TestUsageSorting(APITestCase):
|
||||
"""Server-side sort for /recent and /search (3.2.4, #90.1)."""
|
||||
|
||||
def _use(self, prompt_id, times=1):
|
||||
for _ in range(times):
|
||||
self.api.db.record_prompt_use(prompt_id)
|
||||
|
||||
async def _ids(self, path):
|
||||
resp = await self.client.request("GET", path)
|
||||
self.assertEqual(resp.status, 200)
|
||||
data = await resp.json()
|
||||
self.assertTrue(data["success"])
|
||||
return [p["id"] for p in data["results"]]
|
||||
|
||||
async def test_recent_sorted_by_last_used_bubbles_rerun_prompt(self):
|
||||
old = self._save_prompt("Old prompt")
|
||||
self._save_prompt("New prompt")
|
||||
self._use(old)
|
||||
ids = await self._ids("/prompt_manager/recent?limit=1&sort=last_used_desc")
|
||||
self.assertEqual(ids, [old])
|
||||
|
||||
async def test_recent_sorted_by_run_count(self):
|
||||
# Most-run prompt is the oldest, so newest-first would give the opposite order
|
||||
most = self._save_prompt("Run thrice")
|
||||
least = self._save_prompt("Run once")
|
||||
self._use(most, times=3)
|
||||
self._use(least)
|
||||
ids = await self._ids("/prompt_manager/recent?sort=run_count_desc")
|
||||
self.assertEqual(ids, [most, least])
|
||||
|
||||
async def test_recent_default_is_newest_first(self):
|
||||
first = self._save_prompt("First")
|
||||
second = self._save_prompt("Second")
|
||||
self._use(first)
|
||||
self.assertEqual(await self._ids("/prompt_manager/recent"), [second, first])
|
||||
|
||||
async def test_unknown_sort_falls_back_instead_of_erroring(self):
|
||||
self._save_prompt("Only")
|
||||
ids = await self._ids(
|
||||
"/prompt_manager/recent?sort=created_at;DROP%20TABLE%20prompts"
|
||||
)
|
||||
self.assertEqual(len(ids), 1)
|
||||
|
||||
async def test_search_respects_sort(self):
|
||||
a = self._save_prompt("castle at dawn")
|
||||
self._save_prompt("castle at dusk")
|
||||
self._use(a)
|
||||
ids = await self._ids("/prompt_manager/search?text=castle&sort=last_used_desc")
|
||||
self.assertEqual(ids[0], a)
|
||||
|
||||
async def test_results_include_usage_fields(self):
|
||||
pid = self._save_prompt("Tracked")
|
||||
self._use(pid, times=2)
|
||||
resp = await self.client.request("GET", "/prompt_manager/recent")
|
||||
prompt = (await resp.json())["results"][0]
|
||||
self.assertEqual(prompt["run_count"], 2)
|
||||
self.assertTrue(prompt["last_used_at"])
|
||||
|
||||
|
||||
class TestTags(APITestCase):
|
||||
|
||||
async def test_get_all_tags(self):
|
||||
|
||||
@@ -0,0 +1,141 @@
|
||||
"""Image linking picks the positive prompt, preferring the image's own metadata."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
import unittest.mock as mock
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.image_monitor import ImageGenerationHandler
|
||||
from utils.prompt_tracker import PromptTracker
|
||||
|
||||
LOADER = {
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"inputs": {"ckpt_name": "m.safetensors"},
|
||||
}
|
||||
|
||||
|
||||
def graph(positive_text="a cat", negative_text="blurry"):
|
||||
# Negative listed first: dict order must not decide which prompt an image gets
|
||||
return {
|
||||
"176": {"class_type": "PromptManager", "inputs": {"text": negative_text}},
|
||||
"134": {"class_type": "PromptManager", "inputs": {"text": positive_text}},
|
||||
"5": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"positive": ["134", 0],
|
||||
"negative": ["176", 0],
|
||||
"model": ["20", 0],
|
||||
},
|
||||
},
|
||||
"20": LOADER,
|
||||
}
|
||||
|
||||
|
||||
class LinkingTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
fd, self.db_path = tempfile.mkstemp(suffix=".db")
|
||||
os.close(fd)
|
||||
fd, self.image = tempfile.mkstemp(suffix=".png")
|
||||
os.close(fd)
|
||||
self.db = PromptDatabase(self.db_path)
|
||||
self.tracker = PromptTracker(self.db)
|
||||
self.handler = ImageGenerationHandler(self.db, self.tracker)
|
||||
self.linked = []
|
||||
self.handler.link_image_to_prompt = (
|
||||
lambda path, prompt, meta: self.linked.append(prompt["id"])
|
||||
)
|
||||
|
||||
def tearDown(self):
|
||||
for path in (
|
||||
self.image,
|
||||
self.db_path,
|
||||
self.db_path + "-wal",
|
||||
self.db_path + "-shm",
|
||||
):
|
||||
if os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
def _save(self, text):
|
||||
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
||||
|
||||
def _process(self, metadata):
|
||||
with mock.patch.object(
|
||||
self.handler.metadata_extractor, "extract_metadata", return_value=metadata
|
||||
):
|
||||
self.handler.process_new_image(self.image)
|
||||
return self.linked[-1] if self.linked else None
|
||||
|
||||
|
||||
class TestMetadataLinking(LinkingTestCase):
|
||||
def test_links_the_positive_prompt_not_the_first_prompt_manager_node(self):
|
||||
positive = self._save("a cat")
|
||||
self._save("blurry")
|
||||
self.assertEqual(self._process({"prompt": graph()}), positive)
|
||||
|
||||
def test_image_metadata_wins_over_a_stale_queue_entry(self):
|
||||
positive = self._save("a cat")
|
||||
stale = self._save("previous run")
|
||||
self.tracker.set_current_prompt("previous run", {"prompt_id": stale})
|
||||
|
||||
self.assertEqual(self._process({"prompt": graph()}), positive)
|
||||
self.assertEqual(
|
||||
self.tracker.pop_next_prompt()["id"], stale, "queue left intact"
|
||||
)
|
||||
|
||||
def test_text_from_a_known_string_node_links_from_metadata(self):
|
||||
positive = self._save("a cat")
|
||||
linked_graph = graph()
|
||||
linked_graph["134"]["inputs"]["text"] = ["8", 0]
|
||||
linked_graph["8"] = {
|
||||
"class_type": "PrimitiveString",
|
||||
"inputs": {"value": "a cat"},
|
||||
}
|
||||
self.assertEqual(self._process({"prompt": linked_graph}), positive)
|
||||
|
||||
def test_batch_items_with_linked_text_use_the_queue(self):
|
||||
first, second = self._save("item one"), self._save("item two")
|
||||
batch_graph = graph()
|
||||
batch_graph["134"]["inputs"]["text"] = ["7", 0]
|
||||
batch_graph["7"] = {"class_type": "PromptSearchList", "inputs": {}}
|
||||
self.tracker.set_current_prompt("item one", {"prompt_id": first})
|
||||
self.tracker.set_current_prompt("item two", {"prompt_id": second})
|
||||
|
||||
self.assertEqual(self._process({"prompt": batch_graph}), first)
|
||||
self.assertEqual(self._process({"prompt": batch_graph}), second)
|
||||
|
||||
def test_api_graph_present_skips_the_role_unaware_workflow_fallback(self):
|
||||
negative = self._save("blurry")
|
||||
batch_graph = graph()
|
||||
batch_graph["134"]["inputs"]["text"] = ["7", 0]
|
||||
metadata = {
|
||||
"prompt": batch_graph,
|
||||
"text_encoder_nodes": [
|
||||
{"type": "PromptManager", "inputs": [], "widgets_values": ["blurry"]}
|
||||
],
|
||||
}
|
||||
self.assertIsNone(self.handler._find_prompt_from_metadata(metadata))
|
||||
self.assertNotEqual(self._process(metadata), negative)
|
||||
|
||||
def test_workflow_fallback_still_used_without_an_api_graph(self):
|
||||
positive = self._save("from workflow only")
|
||||
metadata = {
|
||||
"text_encoder_nodes": [
|
||||
{
|
||||
"type": "PromptManager",
|
||||
"inputs": [],
|
||||
"widgets_values": ["from workflow only"],
|
||||
}
|
||||
]
|
||||
}
|
||||
self.assertEqual(
|
||||
self.handler._find_prompt_from_metadata(metadata)["id"], positive
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,83 @@
|
||||
"""Nodes pass ComfyUI's queued graph to role-aware tracking; the hook registers once."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import threading
|
||||
import unittest
|
||||
import unittest.mock as mock
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from prompt_manager import PromptManager
|
||||
from prompt_manager_text import PromptManagerText
|
||||
from utils import usage_tracking
|
||||
|
||||
GRAPH = {"134": {"class_type": "PromptManager", "inputs": {"text": "a cat"}}}
|
||||
|
||||
|
||||
def bare(node_cls):
|
||||
"""A node without its database, tracker or gallery side effects."""
|
||||
node = node_cls.__new__(node_cls)
|
||||
node.logger = __import__("logging").getLogger("test.wiring")
|
||||
node.comfyui_integration = mock.Mock()
|
||||
node._inject_lora_trigger_words = lambda text: text
|
||||
node._track_prompt_execution = mock.Mock(return_value=42)
|
||||
return node
|
||||
|
||||
|
||||
class TestHiddenInputs(unittest.TestCase):
|
||||
def test_both_nodes_request_the_queued_graph_and_their_id(self):
|
||||
for node_cls in (PromptManager, PromptManagerText):
|
||||
hidden = node_cls.INPUT_TYPES().get("hidden", {})
|
||||
self.assertEqual(hidden.get("prompt"), "PROMPT", node_cls.__name__)
|
||||
self.assertEqual(hidden.get("unique_id"), "UNIQUE_ID", node_cls.__name__)
|
||||
|
||||
|
||||
class TestNodesForwardGraph(unittest.TestCase):
|
||||
def test_prompt_manager_forwards_graph_and_texts(self):
|
||||
node = bare(PromptManager)
|
||||
clip = mock.Mock()
|
||||
node.encode_prompt(
|
||||
clip, "a cat", prepend_text="best,", prompt=GRAPH, unique_id="134"
|
||||
)
|
||||
|
||||
kwargs = node._track_prompt_execution.call_args.kwargs
|
||||
self.assertEqual(kwargs["prompt_graph"], GRAPH)
|
||||
self.assertEqual(kwargs["unique_id"], "134")
|
||||
self.assertEqual(kwargs["text"], "a cat")
|
||||
self.assertEqual(kwargs["encoding_text"], "best, a cat")
|
||||
self.assertEqual(kwargs["additional_data"]["final_text"], "best, a cat")
|
||||
|
||||
def test_prompt_manager_text_forwards_graph(self):
|
||||
node = bare(PromptManagerText)
|
||||
node.process_text("a castle", prompt=GRAPH, unique_id="7")
|
||||
|
||||
kwargs = node._track_prompt_execution.call_args.kwargs
|
||||
self.assertEqual((kwargs["prompt_graph"], kwargs["unique_id"]), (GRAPH, "7"))
|
||||
self.assertEqual(kwargs["text"], "a castle")
|
||||
|
||||
def test_nodes_still_work_without_graph_context(self):
|
||||
node = bare(PromptManager)
|
||||
node.encode_prompt(mock.Mock(), "a cat")
|
||||
kwargs = node._track_prompt_execution.call_args.kwargs
|
||||
self.assertIsNone(kwargs["prompt_graph"])
|
||||
self.assertIsNone(kwargs["unique_id"])
|
||||
|
||||
|
||||
class TestQueueHookRegistration(unittest.TestCase):
|
||||
def test_registers_once_and_handler_passes_requests_through(self):
|
||||
server = mock.Mock()
|
||||
db = mock.Mock()
|
||||
db.get_prompt_by_hash.return_value = None
|
||||
with mock.patch.object(usage_tracking, "_hook_registered", threading.Event()):
|
||||
self.assertTrue(usage_tracking.register_queue_hook(server, lambda: db))
|
||||
self.assertFalse(usage_tracking.register_queue_hook(server, lambda: db))
|
||||
|
||||
server.add_on_prompt_handler.assert_called_once()
|
||||
handler = server.add_on_prompt_handler.call_args.args[0]
|
||||
request = {"prompt": GRAPH}
|
||||
self.assertIs(handler(request), request)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,244 @@
|
||||
"""Tests for utils.prompt_graph: which PromptManager nodes feed a positive input."""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from utils.prompt_graph import (
|
||||
positive_prompt_nodes,
|
||||
resolve_text,
|
||||
run_prompt_nodes,
|
||||
)
|
||||
|
||||
|
||||
def pm(text="a cat", **inputs):
|
||||
return {"class_type": "PromptManager", "inputs": {"text": text, **inputs}}
|
||||
|
||||
|
||||
def ksampler(positive, negative, model="20"):
|
||||
return {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"positive": [positive, 0],
|
||||
"negative": [negative, 0],
|
||||
"model": [model, 0],
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
LOADER = {
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"inputs": {"ckpt_name": "m.safetensors"},
|
||||
}
|
||||
|
||||
|
||||
class TestPositivePromptNodes(unittest.TestCase):
|
||||
def test_positive_and_negative_prompt_manager_nodes(self):
|
||||
# Negative listed first, as in the #75 / pmtest workflow
|
||||
graph = {
|
||||
"176": pm("blurry"),
|
||||
"134": pm("a cat"),
|
||||
"5": ksampler("134", "176"),
|
||||
"20": LOADER,
|
||||
}
|
||||
self.assertEqual(positive_prompt_nodes(graph), ["134"])
|
||||
|
||||
def test_negative_only_prompt_manager_is_never_positive(self):
|
||||
graph = {
|
||||
"176": pm("blurry"),
|
||||
"6": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat"}},
|
||||
"5": ksampler("6", "176"),
|
||||
"20": LOADER,
|
||||
}
|
||||
self.assertEqual(positive_prompt_nodes(graph), [])
|
||||
|
||||
def test_through_controlnet_and_combine_pass_throughs(self):
|
||||
graph = {
|
||||
"1": pm("subject"),
|
||||
"2": pm("style"),
|
||||
"3": pm("blurry"),
|
||||
"4": {
|
||||
"class_type": "ConditioningCombine",
|
||||
"inputs": {"conditioning_1": ["1", 0], "conditioning_2": ["2", 0]},
|
||||
},
|
||||
"5": {
|
||||
"class_type": "ControlNetApplyAdvanced",
|
||||
"inputs": {"positive": ["4", 0], "negative": ["3", 0], "strength": 1.0},
|
||||
},
|
||||
"6": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"positive": ["5", 0],
|
||||
"negative": ["5", 1],
|
||||
"model": ["20", 0],
|
||||
},
|
||||
},
|
||||
"20": LOADER,
|
||||
}
|
||||
self.assertEqual(sorted(positive_prompt_nodes(graph)), ["1", "2"])
|
||||
|
||||
def test_cfg_guider_and_basic_guider(self):
|
||||
graph = {
|
||||
"1": pm("fox"),
|
||||
"2": pm("watermark"),
|
||||
"3": {
|
||||
"class_type": "CFGGuider",
|
||||
"inputs": {
|
||||
"positive": ["1", 0],
|
||||
"negative": ["2", 0],
|
||||
"model": ["20", 0],
|
||||
"cfg": 5,
|
||||
},
|
||||
},
|
||||
"4": pm("owl"),
|
||||
"5": {
|
||||
"class_type": "BasicGuider",
|
||||
"inputs": {"conditioning": ["4", 0], "model": ["20", 0]},
|
||||
},
|
||||
"20": LOADER,
|
||||
}
|
||||
self.assertEqual(sorted(positive_prompt_nodes(graph)), ["1", "4"])
|
||||
|
||||
def test_prompt_manager_text_wired_into_an_encoder(self):
|
||||
graph = {
|
||||
"1": {"class_type": "PromptManagerText", "inputs": {"text": "a castle"}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": ["1", 0]}},
|
||||
"3": {"class_type": "PromptManagerText", "inputs": {"text": "lowres"}},
|
||||
"4": {"class_type": "CLIPTextEncode", "inputs": {"text": ["3", 0]}},
|
||||
"5": ksampler("2", "4"),
|
||||
"20": LOADER,
|
||||
}
|
||||
self.assertEqual(positive_prompt_nodes(graph), ["1"])
|
||||
|
||||
def test_node_used_as_both_positive_and_negative_counts_as_positive(self):
|
||||
graph = {"1": pm("same"), "5": ksampler("1", "1"), "20": LOADER}
|
||||
self.assertEqual(positive_prompt_nodes(graph), ["1"])
|
||||
|
||||
def test_two_samplers_sharing_a_prompt_list_it_once(self):
|
||||
graph = {
|
||||
"1": pm("hero"),
|
||||
"2": pm("blurry"),
|
||||
"5": ksampler("1", "2"),
|
||||
"14": ksampler("1", "2"),
|
||||
"20": LOADER,
|
||||
}
|
||||
self.assertEqual(positive_prompt_nodes(graph), ["1"])
|
||||
|
||||
def test_graph_without_sampler_or_malformed_input(self):
|
||||
self.assertEqual(positive_prompt_nodes({"1": pm()}), [])
|
||||
for bad in (None, [], "x", {"1": None}, {"1": {"inputs": None}}):
|
||||
self.assertEqual(positive_prompt_nodes(bad), [], repr(bad))
|
||||
|
||||
def test_cycles_terminate(self):
|
||||
loop = {
|
||||
"class_type": "Weird",
|
||||
"inputs": {
|
||||
"conditioning": ["9", 0],
|
||||
"conditioning_1": ["9", 0],
|
||||
"positive": ["9", 0],
|
||||
},
|
||||
}
|
||||
graph = {"9": loop, "5": ksampler("9", "9"), "20": LOADER}
|
||||
self.assertEqual(positive_prompt_nodes(graph), [])
|
||||
|
||||
|
||||
class TestRunPromptNodes(unittest.TestCase):
|
||||
"""The prompt nodes that represent a run, with a safe fallback."""
|
||||
|
||||
def test_uses_positive_nodes_when_found(self):
|
||||
graph = {
|
||||
"176": pm("blurry"),
|
||||
"134": pm("a cat"),
|
||||
"5": ksampler("134", "176"),
|
||||
"20": LOADER,
|
||||
}
|
||||
self.assertEqual(run_prompt_nodes(graph), ["134"])
|
||||
|
||||
def test_unrecognised_sampler_falls_back_to_every_prompt_manager_node(self):
|
||||
graph = {
|
||||
"1": pm("a cat"),
|
||||
"9": {"class_type": "ExoticSampler", "inputs": {"cond": ["1", 0]}},
|
||||
"2": {"class_type": "PromptManagerText", "inputs": {"text": "x"}},
|
||||
}
|
||||
self.assertEqual(sorted(run_prompt_nodes(graph)), ["1", "2"])
|
||||
|
||||
def test_malformed_graph(self):
|
||||
self.assertEqual(run_prompt_nodes(None), [])
|
||||
|
||||
|
||||
class TestResolveText(unittest.TestCase):
|
||||
"""Text known at queue time: typed, or from known pure string nodes."""
|
||||
|
||||
def _linked(self, source):
|
||||
return {"1": pm(["2", 0]), "2": source}
|
||||
|
||||
def test_typed_text(self):
|
||||
self.assertEqual(resolve_text({"1": pm(" a cat ")}, "1"), "a cat")
|
||||
|
||||
def test_core_primitive_string(self):
|
||||
for cls in ("PrimitiveString", "PrimitiveStringMultiline"):
|
||||
graph = self._linked({"class_type": cls, "inputs": {"value": "a cat"}})
|
||||
self.assertEqual(resolve_text(graph, "1"), "a cat", cls)
|
||||
|
||||
def test_core_string_concatenate(self):
|
||||
graph = self._linked(
|
||||
{
|
||||
"class_type": "StringConcatenate",
|
||||
"inputs": {
|
||||
"string_a": "a cat",
|
||||
"string_b": "at dusk",
|
||||
"delimiter": ", ",
|
||||
},
|
||||
}
|
||||
)
|
||||
self.assertEqual(resolve_text(graph, "1"), "a cat, at dusk")
|
||||
|
||||
def test_was_text_concatenate_skips_empty_and_cleans_whitespace(self):
|
||||
graph = self._linked(
|
||||
{
|
||||
"class_type": "Text Concatenate",
|
||||
"inputs": {
|
||||
"delimiter": " ",
|
||||
"clean_whitespace": "true",
|
||||
"text_a": " a cat ",
|
||||
"text_b": "",
|
||||
"text_c": ["3", 0],
|
||||
},
|
||||
}
|
||||
)
|
||||
graph["3"] = {"class_type": "PrimitiveString", "inputs": {"value": "at dusk"}}
|
||||
self.assertEqual(resolve_text(graph, "1"), "a cat at dusk")
|
||||
|
||||
def test_pass_through_nodes(self):
|
||||
for cls in ("ShowText|pysssss", "Text Multiline"):
|
||||
graph = self._linked({"class_type": cls, "inputs": {"text": "a cat"}})
|
||||
self.assertEqual(resolve_text(graph, "1"), "a cat", cls)
|
||||
|
||||
def test_unknown_and_batch_sources_are_never_guessed(self):
|
||||
for source in (
|
||||
{"class_type": "PromptSearchList", "inputs": {"search_text": "cat"}},
|
||||
{
|
||||
"class_type": "SomeTextReplace",
|
||||
"inputs": {"text": "a cat", "find": "cat", "replace": "dog"},
|
||||
},
|
||||
):
|
||||
self.assertIsNone(
|
||||
resolve_text(self._linked(source), "1"), source["class_type"]
|
||||
)
|
||||
|
||||
def test_cycles_and_missing_links(self):
|
||||
graph = {
|
||||
"1": pm(["2", 0]),
|
||||
"2": {
|
||||
"class_type": "StringConcatenate",
|
||||
"inputs": {"string_a": ["2", 0], "string_b": "x", "delimiter": ""},
|
||||
},
|
||||
}
|
||||
self.assertIsNone(resolve_text(graph, "1"))
|
||||
self.assertIsNone(resolve_text({"1": pm(["9", 0])}, "1"))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,230 @@
|
||||
"""Tests for prompt usage tracking: run_count, last_used_at and server-side sorting."""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from prompt_manager_base import PromptManagerBase
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
|
||||
class UsageTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
fd, self.path = tempfile.mkstemp(suffix=".db")
|
||||
os.close(fd)
|
||||
self.db = PromptDatabase(self.path)
|
||||
|
||||
def tearDown(self):
|
||||
for suffix in ("", "-wal", "-shm"):
|
||||
if os.path.exists(self.path + suffix):
|
||||
os.unlink(self.path + suffix)
|
||||
|
||||
def _save(self, text, **kwargs):
|
||||
return self.db.save_prompt(
|
||||
text=text, prompt_hash=generate_prompt_hash(text), **kwargs
|
||||
)
|
||||
|
||||
def _set(self, prompt_id, **columns):
|
||||
with sqlite3.connect(self.path) as conn:
|
||||
for column, value in columns.items():
|
||||
conn.execute(
|
||||
f"UPDATE prompts SET {column} = ? WHERE id = ?", (value, prompt_id)
|
||||
)
|
||||
|
||||
def _ids(self, prompts):
|
||||
return [p["id"] for p in prompts]
|
||||
|
||||
|
||||
class TestRecordPromptUse(UsageTestCase):
|
||||
def test_new_prompt_starts_unused_but_timestamped(self):
|
||||
prompt = self.db.get_prompt_by_id(self._save("fresh"))
|
||||
self.assertEqual(prompt["run_count"], 0)
|
||||
self.assertTrue(prompt["last_used_at"])
|
||||
|
||||
def test_record_use_increments_and_touches_last_used(self):
|
||||
pid = self._save("used")
|
||||
self._set(pid, last_used_at="2020-01-01T00:00:00.000+00:00")
|
||||
|
||||
self.assertTrue(self.db.record_prompt_use(pid))
|
||||
self.assertTrue(self.db.record_prompt_use(pid))
|
||||
|
||||
prompt = self.db.get_prompt_by_id(pid)
|
||||
self.assertEqual(prompt["run_count"], 2)
|
||||
self.assertGreater(prompt["last_used_at"], "2020-01-01T00:00:00.000+00:00")
|
||||
|
||||
def test_rerun_right_after_another_save_still_sorts_first(self):
|
||||
# No hand-set timestamps: creation and re-run happen within the same millisecond
|
||||
old = self._save("old")
|
||||
self._save("new")
|
||||
self.db.record_prompt_use(old)
|
||||
result = self.db.get_recent_prompts(limit=1, sort="last_used_desc")
|
||||
self.assertEqual(self._ids(result["prompts"]), [old])
|
||||
|
||||
def test_microsecond_timestamps_order_correctly_against_backfilled_ones(self):
|
||||
backfilled = self._save("backfilled")
|
||||
self._set(backfilled, last_used_at="2026-01-01T00:00:00.000+00:00")
|
||||
recent = self._save("recent")
|
||||
self.db.record_prompt_use(recent)
|
||||
self.assertRegex(
|
||||
self.db.get_prompt_by_id(recent)["last_used_at"],
|
||||
r"^\d{4}-\d\d-\d\dT\d\d:\d\d:\d\d\.\d{6}\+00:00$",
|
||||
)
|
||||
result = self.db.get_recent_prompts(limit=2, sort="last_used_desc")
|
||||
self.assertEqual(self._ids(result["prompts"]), [recent, backfilled])
|
||||
|
||||
def test_record_use_can_add_several_runs_at_once(self):
|
||||
pid = self._save("several")
|
||||
self.db.record_prompt_use(pid, times=3)
|
||||
self.assertEqual(self.db.get_prompt_by_id(pid)["run_count"], 3)
|
||||
|
||||
def test_record_use_on_missing_prompt_returns_false(self):
|
||||
self.assertFalse(self.db.record_prompt_use(999999))
|
||||
|
||||
|
||||
class TestServerSideSort(UsageTestCase):
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.old = self._save("old prompt", rating=5)
|
||||
self.mid = self._save("middle prompt", rating=1)
|
||||
self.new = self._save("new prompt")
|
||||
self._set(self.old, last_used_at="2026-01-01T00:00:00.000+00:00", run_count=9)
|
||||
self._set(self.mid, last_used_at="2026-01-02T00:00:00.000+00:00", run_count=1)
|
||||
self._set(self.new, last_used_at="2026-01-03T00:00:00.000+00:00", run_count=3)
|
||||
|
||||
def test_rerun_bubbles_old_prompt_to_top_across_pages(self):
|
||||
self.db.record_prompt_use(self.old)
|
||||
page1 = self.db.get_recent_prompts(limit=1, offset=0, sort="last_used_desc")
|
||||
self.assertEqual(self._ids(page1["prompts"]), [self.old])
|
||||
|
||||
def test_most_used(self):
|
||||
result = self.db.get_recent_prompts(limit=10, sort="run_count_desc")
|
||||
self.assertEqual(self._ids(result["prompts"]), [self.old, self.new, self.mid])
|
||||
|
||||
def test_rating_sort_puts_unrated_last_and_spans_pages(self):
|
||||
page1 = self.db.get_recent_prompts(limit=1, offset=0, sort="rating_desc")
|
||||
self.assertEqual(self._ids(page1["prompts"]), [self.old])
|
||||
result = self.db.get_recent_prompts(limit=10, sort="rating_asc")
|
||||
self.assertEqual(self._ids(result["prompts"]), [self.mid, self.old, self.new])
|
||||
|
||||
def test_text_sort(self):
|
||||
result = self.db.get_recent_prompts(limit=10, sort="text_asc")
|
||||
self.assertEqual(self._ids(result["prompts"]), [self.mid, self.new, self.old])
|
||||
|
||||
def test_default_and_unknown_sort_keep_newest_first(self):
|
||||
expected = [self.new, self.mid, self.old]
|
||||
self.assertEqual(
|
||||
self._ids(self.db.get_recent_prompts(limit=10)["prompts"]), expected
|
||||
)
|
||||
injected = self.db.get_recent_prompts(limit=10, sort="id; DROP TABLE prompts")
|
||||
self.assertEqual(self._ids(injected["prompts"]), expected)
|
||||
|
||||
def test_search_accepts_sort(self):
|
||||
self.db.record_prompt_use(self.mid)
|
||||
results = self.db.search_prompts(text="prompt", sort="last_used_desc")
|
||||
self.assertEqual(self._ids(results)[0], self.mid)
|
||||
|
||||
|
||||
@unittest.skipIf(
|
||||
sqlite3.sqlite_version_info < (3, 35, 0),
|
||||
"simulating the old schema needs ALTER TABLE DROP COLUMN (SQLite 3.35+)",
|
||||
)
|
||||
class TestUsageMigration(unittest.TestCase):
|
||||
"""Databases from before 3.2.4 get the columns and an estimated backfill."""
|
||||
|
||||
def setUp(self):
|
||||
fd, self.path = tempfile.mkstemp(suffix=".db")
|
||||
os.close(fd)
|
||||
db = PromptDatabase(self.path)
|
||||
self.with_images = db.save_prompt(text="has images", prompt_hash="h1")
|
||||
self.no_images = db.save_prompt(text="no images", prompt_hash="h2")
|
||||
with sqlite3.connect(self.path) as conn:
|
||||
conn.execute(
|
||||
"UPDATE prompts SET created_at = '2026-01-01T10:00:00.000000+00:00'"
|
||||
)
|
||||
# generated_images uses SQLite's space-separated CURRENT_TIMESTAMP format
|
||||
for i, ts in enumerate(("2026-03-05 08:00:00", "2026-02-01 09:00:00")):
|
||||
conn.execute(
|
||||
"INSERT INTO generated_images (prompt_id, image_path, filename, generation_time)"
|
||||
" VALUES (?, ?, ?, ?)",
|
||||
(self.with_images, f"/out/{i}.png", f"{i}.png", ts),
|
||||
)
|
||||
# Simulate a pre-3.2.4 schema
|
||||
for index in ("idx_prompts_last_used", "idx_prompts_run_count"):
|
||||
conn.execute(f"DROP INDEX IF EXISTS {index}")
|
||||
conn.execute("ALTER TABLE prompts DROP COLUMN last_used_at")
|
||||
conn.execute("ALTER TABLE prompts DROP COLUMN run_count")
|
||||
|
||||
def tearDown(self):
|
||||
for suffix in ("", "-wal", "-shm"):
|
||||
if os.path.exists(self.path + suffix):
|
||||
os.unlink(self.path + suffix)
|
||||
|
||||
def test_backfill_estimates_usage_from_linked_images(self):
|
||||
db = PromptDatabase(self.path)
|
||||
used = db.get_prompt_by_id(self.with_images)
|
||||
unused = db.get_prompt_by_id(self.no_images)
|
||||
|
||||
self.assertEqual(used["run_count"], 2)
|
||||
self.assertTrue(used["last_used_at"].startswith("2026-03-05T08:00:00"))
|
||||
self.assertEqual(unused["run_count"], 1)
|
||||
self.assertTrue(unused["last_used_at"].startswith("2026-01-01T10:00:00"))
|
||||
|
||||
def test_rows_written_while_downgraded_are_healed_on_next_start(self):
|
||||
db = PromptDatabase(self.path)
|
||||
# 3.2.3 does not know the columns: its inserts leave last_used_at NULL
|
||||
with sqlite3.connect(self.path) as conn:
|
||||
conn.execute(
|
||||
"INSERT INTO prompts (text, hash, created_at)"
|
||||
" VALUES ('from 3.2.3', 'h3', '2026-04-01T00:00:00.000000+00:00')"
|
||||
)
|
||||
legacy_id = conn.execute(
|
||||
"SELECT id FROM prompts WHERE hash = 'h3'"
|
||||
).fetchone()[0]
|
||||
self.assertIsNone(db.get_prompt_by_id(legacy_id)["last_used_at"])
|
||||
|
||||
healed = PromptDatabase(self.path).get_prompt_by_id(legacy_id)
|
||||
self.assertTrue(healed["last_used_at"].startswith("2026-04-01T00:00:00"))
|
||||
self.assertEqual(healed["run_count"], 1)
|
||||
|
||||
def test_migration_runs_once(self):
|
||||
db = PromptDatabase(self.path)
|
||||
db.record_prompt_use(self.no_images)
|
||||
reopened = PromptDatabase(self.path)
|
||||
self.assertEqual(reopened.get_prompt_by_id(self.no_images)["run_count"], 2)
|
||||
|
||||
|
||||
class TestNodeRecordsUse(UsageTestCase):
|
||||
"""Node saves count a run only when asked; re-runs are counted by the queue hook.
|
||||
|
||||
ComfyUI skips unchanged nodes, so counting at execution misses re-runs; see
|
||||
tests/test_usage_tracking.py for the hook that counts them.
|
||||
"""
|
||||
|
||||
def _node(self):
|
||||
node = PromptManagerBase.__new__(PromptManagerBase)
|
||||
node.db = self.db
|
||||
node.logger = __import__("logging").getLogger("test.prompt_usage")
|
||||
return node
|
||||
|
||||
def test_save_without_count_run_leaves_counting_to_the_queue_hook(self):
|
||||
node = self._node()
|
||||
first = node._save_prompt_to_database("a lighthouse at dusk")
|
||||
second = node._save_prompt_to_database("a lighthouse at dusk")
|
||||
|
||||
self.assertEqual(first, second)
|
||||
self.assertEqual(self.db.get_prompt_by_id(first)["run_count"], 0)
|
||||
|
||||
def test_count_run_counts_new_and_existing_prompts(self):
|
||||
node = self._node()
|
||||
pid = node._save_prompt_to_database("batch item", count_run=True)
|
||||
node._save_prompt_to_database("batch item", count_run=True)
|
||||
self.assertEqual(self.db.get_prompt_by_id(pid)["run_count"], 2)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,39 @@
|
||||
"""The UI's sort options must match the server's whitelist exactly."""
|
||||
|
||||
import os
|
||||
import re
|
||||
import sys
|
||||
import unittest
|
||||
|
||||
ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
sys.path.insert(0, ROOT)
|
||||
|
||||
from database.operations import SORT_ORDERS
|
||||
|
||||
SORT_MODULE = os.path.join(ROOT, "web", "js", "prompt-list-sort.js")
|
||||
NODE_EXTENSION = os.path.join(ROOT, "web", "prompt_manager.js")
|
||||
|
||||
|
||||
class TestSortContract(unittest.TestCase):
|
||||
def test_ui_sort_options_match_server_whitelist(self):
|
||||
with open(SORT_MODULE, encoding="utf-8") as fh:
|
||||
source = fh.read()
|
||||
ui_values = set(re.findall(r'value:\s*"([a-z_]+)"', source))
|
||||
|
||||
self.assertTrue(ui_values, "no sort options found in prompt-list-sort.js")
|
||||
self.assertEqual(ui_values, set(SORT_ORDERS))
|
||||
|
||||
def test_node_recent_button_requests_recently_used(self):
|
||||
# The node extension can't import the shared module (ComfyUI load order is
|
||||
# not guaranteed), so it sends the key directly; it must be a server key.
|
||||
with open(NODE_EXTENSION, encoding="utf-8") as fh:
|
||||
source = fh.read()
|
||||
match = re.search(r"/prompt_manager/recent\?[^`'\"]*sort=([a-z_]+)", source)
|
||||
|
||||
self.assertIsNotNone(match, "node Recent request does not send a sort")
|
||||
self.assertEqual(match.group(1), "last_used_desc")
|
||||
self.assertIn(match.group(1), SORT_ORDERS)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,283 @@
|
||||
"""Run counting and node roles: only positive PromptManager nodes count and link.
|
||||
|
||||
Runs are counted when a prompt is queued (ComfyUI's on-prompt hook), because
|
||||
ComfyUI skips unchanged nodes, so node execution can't be relied on to count.
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
import unittest.mock as mock
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from prompt_manager_base import PromptManagerBase
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.prompt_tracker import PromptTracker
|
||||
from utils.usage_tracking import PendingFirstUse, handle_queued_prompt
|
||||
|
||||
LOADER = {
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"inputs": {"ckpt_name": "m.safetensors"},
|
||||
}
|
||||
|
||||
|
||||
def workflow(positive_text="a cat", negative_text="blurry"):
|
||||
"""API graph with a positive (#134) and a negative (#176) PromptManager node."""
|
||||
return {
|
||||
"176": {"class_type": "PromptManager", "inputs": {"text": negative_text}},
|
||||
"134": {"class_type": "PromptManager", "inputs": {"text": positive_text}},
|
||||
"5": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"positive": ["134", 0],
|
||||
"negative": ["176", 0],
|
||||
"model": ["20", 0],
|
||||
},
|
||||
},
|
||||
"20": LOADER,
|
||||
}
|
||||
|
||||
|
||||
class DbTestCase(unittest.TestCase):
|
||||
def setUp(self):
|
||||
fd, self.path = tempfile.mkstemp(suffix=".db")
|
||||
os.close(fd)
|
||||
self.db = PromptDatabase(self.path)
|
||||
self.pending = PendingFirstUse()
|
||||
|
||||
def tearDown(self):
|
||||
for suffix in ("", "-wal", "-shm"):
|
||||
if os.path.exists(self.path + suffix):
|
||||
os.unlink(self.path + suffix)
|
||||
|
||||
def _save(self, text):
|
||||
return self.db.save_prompt(text=text, prompt_hash=generate_prompt_hash(text))
|
||||
|
||||
def _runs(self, prompt_id):
|
||||
return self.db.get_prompt_by_id(prompt_id)["run_count"]
|
||||
|
||||
def _queue(self, graph):
|
||||
return handle_queued_prompt(
|
||||
{"prompt": graph, "client_id": "x"}, self.db, self.pending
|
||||
)
|
||||
|
||||
|
||||
class TestQueueHook(DbTestCase):
|
||||
def test_counts_the_positive_prompt_only(self):
|
||||
positive, negative = self._save("a cat"), self._save("blurry")
|
||||
self._queue(workflow())
|
||||
self.assertEqual(self._runs(positive), 1)
|
||||
self.assertEqual(self._runs(negative), 0)
|
||||
|
||||
def test_requeue_counts_again_even_when_comfyui_skips_the_cached_node(self):
|
||||
positive = self._save("a cat")
|
||||
self._queue(workflow())
|
||||
self._queue(workflow()) # no node execution in between
|
||||
self.assertEqual(self._runs(positive), 2)
|
||||
|
||||
def test_new_prompt_is_counted_when_the_node_first_saves_it(self):
|
||||
self._queue(workflow(positive_text="a brand new prompt"))
|
||||
self.assertTrue(
|
||||
self.pending.consume(generate_prompt_hash("a brand new prompt"))
|
||||
)
|
||||
self.assertFalse(
|
||||
self.pending.consume(generate_prompt_hash("a brand new prompt"))
|
||||
)
|
||||
|
||||
def test_same_text_in_two_positive_nodes_counts_once(self):
|
||||
positive = self._save("a cat")
|
||||
graph = workflow()
|
||||
graph["135"] = {"class_type": "PromptManager", "inputs": {"text": "a cat"}}
|
||||
graph["4"] = {
|
||||
"class_type": "ConditioningCombine",
|
||||
"inputs": {"conditioning_1": ["134", 0], "conditioning_2": ["135", 0]},
|
||||
}
|
||||
graph["5"]["inputs"]["positive"] = ["4", 0]
|
||||
self._queue(graph)
|
||||
self.assertEqual(self._runs(positive), 1)
|
||||
|
||||
def test_text_from_a_known_string_node_is_counted_at_queue_time(self):
|
||||
positive = self._save("a cat")
|
||||
graph = workflow()
|
||||
graph["134"]["inputs"]["text"] = ["8", 0]
|
||||
graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}}
|
||||
self._queue(graph)
|
||||
self._queue(graph) # cached re-run: the node won't execute
|
||||
self.assertEqual(self._runs(positive), 2)
|
||||
|
||||
def test_new_prompt_queued_twice_before_first_run_counts_twice(self):
|
||||
self._queue(workflow(positive_text="queued twice"))
|
||||
self._queue(workflow(positive_text="queued twice"))
|
||||
self.assertEqual(self.pending.consume(generate_prompt_hash("queued twice")), 2)
|
||||
|
||||
def test_linked_text_is_left_to_node_execution(self):
|
||||
positive = self._save("from a batch")
|
||||
graph = workflow()
|
||||
graph["134"]["inputs"]["text"] = ["7", 0]
|
||||
graph["7"] = {"class_type": "PromptSearchList", "inputs": {}}
|
||||
self._queue(graph)
|
||||
self.assertEqual(self._runs(positive), 0)
|
||||
|
||||
def test_returns_the_request_unchanged_and_never_raises(self):
|
||||
request = {"prompt": workflow(), "client_id": "x", "extra_data": {}}
|
||||
self.assertIs(handle_queued_prompt(request, self.db, self.pending), request)
|
||||
for bad in (None, {}, {"prompt": "x"}, {"prompt": None}, []):
|
||||
self.assertIs(handle_queued_prompt(bad, self.db, self.pending), bad)
|
||||
broken_db = mock.Mock()
|
||||
broken_db.get_prompt_by_hash.side_effect = RuntimeError("db locked")
|
||||
self.assertIs(handle_queued_prompt(request, broken_db, self.pending), request)
|
||||
|
||||
|
||||
class TestPendingFirstUse(unittest.TestCase):
|
||||
def test_entries_expire(self):
|
||||
pending = PendingFirstUse(ttl_seconds=10)
|
||||
with mock.patch("utils.usage_tracking.time.monotonic", return_value=100.0):
|
||||
pending.add("h")
|
||||
with mock.patch("utils.usage_tracking.time.monotonic", return_value=111.0):
|
||||
self.assertFalse(pending.consume("h"))
|
||||
|
||||
def test_consume_returns_how_many_runs_were_pending(self):
|
||||
pending = PendingFirstUse()
|
||||
pending.add("h")
|
||||
pending.add("h")
|
||||
pending.add("h")
|
||||
self.assertEqual(pending.consume("h"), 3)
|
||||
self.assertEqual(pending.consume("h"), 0)
|
||||
|
||||
def test_size_is_capped(self):
|
||||
pending = PendingFirstUse(max_entries=3)
|
||||
for h in ("a", "b", "c", "d"):
|
||||
pending.add(h)
|
||||
self.assertFalse(pending.consume("a")) # oldest evicted
|
||||
self.assertTrue(pending.consume("d"))
|
||||
|
||||
|
||||
class FakeTracker:
|
||||
def __init__(self):
|
||||
self.calls = []
|
||||
|
||||
def set_current_prompt(self, prompt_text, additional_data=None, push_to_queue=True):
|
||||
self.calls.append({"text": prompt_text, "push_to_queue": push_to_queue})
|
||||
return "exec-1"
|
||||
|
||||
|
||||
class TestNodeRoles(DbTestCase):
|
||||
def _node(self):
|
||||
node = PromptManagerBase.__new__(PromptManagerBase)
|
||||
node.db = self.db
|
||||
node.prompt_tracker = FakeTracker()
|
||||
node.logger = __import__("logging").getLogger("test.usage_tracking")
|
||||
node.pending_first_use = self.pending
|
||||
return node
|
||||
|
||||
def _run(self, node, graph, unique_id, text):
|
||||
return node._track_prompt_execution(
|
||||
text=text,
|
||||
encoding_text=text,
|
||||
category=None,
|
||||
tags=None,
|
||||
additional_data={},
|
||||
prompt_graph=graph,
|
||||
unique_id=unique_id,
|
||||
)
|
||||
|
||||
def test_negative_node_saves_but_never_becomes_the_current_prompt(self):
|
||||
node = self._node()
|
||||
prompt_id = self._run(node, workflow(), "176", "blurry")
|
||||
self.assertIsNotNone(self.db.get_prompt_by_id(prompt_id))
|
||||
self.assertEqual(node.prompt_tracker.calls, [])
|
||||
self.assertEqual(self._runs(prompt_id), 0)
|
||||
|
||||
def test_positive_typed_text_is_current_without_queue_push(self):
|
||||
node = self._node()
|
||||
self._run(node, workflow(), "134", "a cat")
|
||||
self.assertEqual(
|
||||
node.prompt_tracker.calls, [{"text": "a cat", "push_to_queue": False}]
|
||||
)
|
||||
|
||||
def test_first_save_after_queue_hook_counts_one_run(self):
|
||||
self._queue(workflow(positive_text="never seen before"))
|
||||
prompt_id = self._run(
|
||||
self._node(),
|
||||
workflow(positive_text="never seen before"),
|
||||
"134",
|
||||
"never seen before",
|
||||
)
|
||||
self.assertEqual(self._runs(prompt_id), 1)
|
||||
|
||||
def test_new_prompt_queued_twice_counts_both_runs_on_first_save(self):
|
||||
graph = workflow(positive_text="double queued")
|
||||
self._queue(graph)
|
||||
self._queue(graph)
|
||||
prompt_id = self._run(self._node(), graph, "134", "double queued")
|
||||
self.assertEqual(self._runs(prompt_id), 2)
|
||||
|
||||
def test_text_resolved_at_queue_time_is_not_counted_again_or_queued(self):
|
||||
graph = workflow()
|
||||
graph["134"]["inputs"]["text"] = ["8", 0]
|
||||
graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}}
|
||||
positive = self._save("a cat")
|
||||
self._queue(graph)
|
||||
node = self._node()
|
||||
self._run(node, graph, "134", "a cat")
|
||||
self.assertEqual(self._runs(positive), 1)
|
||||
self.assertEqual(
|
||||
node.prompt_tracker.calls, [{"text": "a cat", "push_to_queue": False}]
|
||||
)
|
||||
|
||||
def test_resolution_mismatch_falls_back_to_counting_and_queueing(self):
|
||||
# Hook guessed "a cat" but the node received different text: trust the node
|
||||
graph = workflow()
|
||||
graph["134"]["inputs"]["text"] = ["8", 0]
|
||||
graph["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "a cat"}}
|
||||
node = self._node()
|
||||
prompt_id = self._run(node, graph, "134", "actually a dog")
|
||||
self.assertEqual(self._runs(prompt_id), 1)
|
||||
self.assertEqual(
|
||||
node.prompt_tracker.calls,
|
||||
[{"text": "actually a dog", "push_to_queue": True}],
|
||||
)
|
||||
|
||||
def test_existing_prompt_is_not_counted_twice_by_hook_and_node(self):
|
||||
positive = self._save("a cat")
|
||||
self._queue(workflow())
|
||||
self._run(self._node(), workflow(), "134", "a cat")
|
||||
self.assertEqual(self._runs(positive), 1)
|
||||
|
||||
def test_positive_linked_text_pushes_queue_and_counts_each_item(self):
|
||||
graph = workflow()
|
||||
graph["134"]["inputs"]["text"] = ["7", 0]
|
||||
graph["7"] = {"class_type": "PromptSearchList", "inputs": {}}
|
||||
node = self._node()
|
||||
first = self._run(node, graph, "134", "batch item one")
|
||||
self._run(node, graph, "134", "batch item one")
|
||||
self.assertEqual(self._runs(first), 2)
|
||||
self.assertTrue(all(c["push_to_queue"] for c in node.prompt_tracker.calls))
|
||||
|
||||
def test_without_graph_context_keeps_legacy_behaviour(self):
|
||||
node = self._node()
|
||||
self._run(node, None, None, "standalone")
|
||||
self.assertEqual(
|
||||
node.prompt_tracker.calls, [{"text": "standalone", "push_to_queue": True}]
|
||||
)
|
||||
|
||||
|
||||
class TestTrackerQueueFlag(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tracker = PromptTracker(mock.Mock())
|
||||
|
||||
def test_push_to_queue_false_keeps_the_batch_queue_empty(self):
|
||||
self.tracker.set_current_prompt("a cat", {"prompt_id": 1}, push_to_queue=False)
|
||||
self.assertIsNone(self.tracker.pop_next_prompt())
|
||||
self.assertEqual(self.tracker.get_current_prompt()["id"], 1)
|
||||
|
||||
def test_default_still_pushes(self):
|
||||
self.tracker.set_current_prompt("a cat", {"prompt_id": 1})
|
||||
self.assertEqual(self.tracker.pop_next_prompt()["id"], 1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
+24
-20
@@ -34,6 +34,7 @@ from watchdog.events import FileSystemEventHandler
|
||||
|
||||
from .metadata_extractor import ComfyUIMetadataExtractor
|
||||
from .logging_config import get_logger
|
||||
from .prompt_graph import resolve_text, run_prompt_nodes
|
||||
|
||||
|
||||
class ImageGenerationHandler(FileSystemEventHandler):
|
||||
@@ -146,21 +147,26 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
except Exception as meta_error:
|
||||
self.logger.warning(f"Metadata extraction failed: {meta_error}")
|
||||
|
||||
# Strategy 1: Pop from batch queue (most reliable for batch workflows).
|
||||
# Prompts are queued during CLIP encoding in order; images save in
|
||||
# the same order, so FIFO pop gives the correct prompt per image.
|
||||
current_prompt = self.prompt_tracker.pop_next_prompt()
|
||||
# Strategy 1: The image's own metadata. It records the positive prompt
|
||||
# that produced this image, independent of node caching and of any
|
||||
# negative-prompt PromptManager nodes in the workflow.
|
||||
current_prompt = self._find_prompt_from_metadata(metadata)
|
||||
if current_prompt:
|
||||
self.logger.info(
|
||||
f"Queue match: prompt {current_prompt['id']} for "
|
||||
f"Metadata match: prompt {current_prompt['id']} for "
|
||||
f"{os.path.basename(image_path)}"
|
||||
)
|
||||
|
||||
# Strategy 2: Find prompt from image metadata
|
||||
# Strategy 2: Pop from batch queue. Batch items (text supplied by
|
||||
# PromptSearchList) share one graph, so metadata can't tell them apart;
|
||||
# they are queued in encoding order and images save in the same order.
|
||||
if not current_prompt:
|
||||
current_prompt = self._find_prompt_from_metadata(metadata)
|
||||
current_prompt = self.prompt_tracker.pop_next_prompt()
|
||||
if current_prompt:
|
||||
self.logger.info(f"Metadata match: prompt {current_prompt['id']}")
|
||||
self.logger.info(
|
||||
f"Queue match: prompt {current_prompt['id']} for "
|
||||
f"{os.path.basename(image_path)}"
|
||||
)
|
||||
|
||||
# Strategy 3: Use snapshot captured at file-creation time
|
||||
if not current_prompt and prompt_snapshot:
|
||||
@@ -231,21 +237,19 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
|
||||
prompt_text = None
|
||||
|
||||
# Try to find PromptManager node in the prompt execution data
|
||||
# The executed graph: use the PromptManager node feeding the positive input
|
||||
prompt_data = metadata.get("prompt")
|
||||
if isinstance(prompt_data, dict):
|
||||
for node_id, node_info in prompt_data.items():
|
||||
if not isinstance(node_info, dict):
|
||||
continue
|
||||
class_type = node_info.get("class_type", "")
|
||||
if class_type in ("PromptManager", "PromptManagerText"):
|
||||
inputs = node_info.get("inputs", {})
|
||||
text = inputs.get("text", "")
|
||||
if text and isinstance(text, str) and text.strip():
|
||||
prompt_text = text.strip()
|
||||
break
|
||||
for node_id in run_prompt_nodes(prompt_data):
|
||||
prompt_text = resolve_text(prompt_data, node_id)
|
||||
if prompt_text:
|
||||
break
|
||||
if not prompt_text:
|
||||
# Positive text can't be known from the graph (batch item or unknown
|
||||
# node): leave it to the queue rather than guess from workflow widgets
|
||||
return None
|
||||
|
||||
# Fallback: check text_encoder_nodes from workflow, but only if
|
||||
# Fallback (no executed graph): check text_encoder_nodes from workflow, but only if
|
||||
# the text input is NOT connected (connected inputs override widget values,
|
||||
# so the widget value would be stale in batch workflows).
|
||||
if not prompt_text:
|
||||
|
||||
@@ -0,0 +1,178 @@
|
||||
"""Find which PromptManager nodes feed a sampler's positive input in a ComfyUI graph.
|
||||
|
||||
Works on the API-format prompt graph ({node_id: {"class_type", "inputs"}}) that
|
||||
ComfyUI queues and embeds in saved images. Usage counting and image linking use it
|
||||
so that negative-prompt PromptManager nodes never count as "the prompt" of a run.
|
||||
|
||||
Pure functions with no ComfyUI imports, so they are unit-testable. Mirrors the
|
||||
link-tracing rules of web/js/comfy-metadata.js.
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
PROMPT_MANAGER_TYPES = frozenset({"PromptManager", "PromptManagerText"})
|
||||
|
||||
# Upper bound on nodes expanded per graph; queued and embedded graphs are untrusted
|
||||
MAX_VISITS = 500
|
||||
|
||||
# Inputs that carry positive conditioning or prompt text towards a sampler
|
||||
_CONDITIONING_KEY = re.compile(r"^conditioning(_\w+)?$")
|
||||
_TEXT_KEY = re.compile(r"^(text|string|prompt)(_?[a-z0-9]+)?$", re.IGNORECASE)
|
||||
|
||||
|
||||
def _is_link(value: Any) -> bool:
|
||||
return (
|
||||
isinstance(value, (list, tuple))
|
||||
and len(value) == 2
|
||||
and isinstance(value[0], (str, int))
|
||||
and isinstance(value[1], int)
|
||||
)
|
||||
|
||||
|
||||
def _inputs(node: Any) -> Dict[str, Any]:
|
||||
inputs = node.get("inputs") if isinstance(node, dict) else None
|
||||
return inputs if isinstance(inputs, dict) else {}
|
||||
|
||||
|
||||
def _positive_roots(graph: Dict[str, Any]) -> List[Any]:
|
||||
"""Links into samplers' positive inputs (KSampler, CFGGuider, BasicGuider, ...)."""
|
||||
roots = []
|
||||
for node in graph.values():
|
||||
inputs = _inputs(node)
|
||||
if not _is_link(inputs.get("model")):
|
||||
continue
|
||||
if _is_link(inputs.get("positive")) and _is_link(inputs.get("negative")):
|
||||
roots.append(inputs["positive"])
|
||||
elif (
|
||||
_is_link(inputs.get("conditioning"))
|
||||
and "guider" in str(node.get("class_type", "")).lower()
|
||||
):
|
||||
roots.append(inputs["conditioning"])
|
||||
return roots
|
||||
|
||||
|
||||
def _upstream_links(inputs: Dict[str, Any]) -> List[Any]:
|
||||
"""Inputs to follow from a pass-through or encoder node, never `negative`."""
|
||||
links = []
|
||||
for key, value in inputs.items():
|
||||
if not _is_link(value):
|
||||
continue
|
||||
if key == "positive" or _CONDITIONING_KEY.match(key) or _TEXT_KEY.match(key):
|
||||
links.append(value)
|
||||
return links
|
||||
|
||||
|
||||
def positive_prompt_nodes(graph: Any) -> List[str]:
|
||||
"""Ids of PromptManager/PromptManagerText nodes that feed a positive input.
|
||||
|
||||
Ordered by discovery, each id once. Malformed graphs yield an empty list.
|
||||
"""
|
||||
if not isinstance(graph, dict):
|
||||
return []
|
||||
|
||||
found: List[str] = []
|
||||
seen = set()
|
||||
stack = list(reversed(_positive_roots(graph)))
|
||||
while stack and len(seen) < MAX_VISITS:
|
||||
node_id = str(stack.pop()[0])
|
||||
if node_id in seen:
|
||||
continue
|
||||
seen.add(node_id)
|
||||
node = graph.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
if node.get("class_type") in PROMPT_MANAGER_TYPES:
|
||||
found.append(node_id)
|
||||
continue
|
||||
stack.extend(reversed(_upstream_links(_inputs(node))))
|
||||
return found
|
||||
|
||||
|
||||
def run_prompt_nodes(graph: Any) -> List[str]:
|
||||
"""Prompt nodes that represent a run: the positive ones, if any can be found.
|
||||
|
||||
When no PromptManager node can be traced to a positive input (an unrecognised
|
||||
custom sampler, say), every PromptManager node counts, as before 3.2.4, so an
|
||||
exotic workflow never silently loses usage counting or image linking.
|
||||
"""
|
||||
positive = positive_prompt_nodes(graph)
|
||||
if positive or not isinstance(graph, dict):
|
||||
return positive
|
||||
return [
|
||||
str(node_id)
|
||||
for node_id, node in graph.items()
|
||||
if isinstance(node, dict) and node.get("class_type") in PROMPT_MANAGER_TYPES
|
||||
]
|
||||
|
||||
|
||||
def _join_text_inputs(
|
||||
inputs: Dict[str, Any], keys: List[str], resolve
|
||||
) -> Optional[str]:
|
||||
"""WAS-style join: resolve each key, optionally strip, skip empty, join by delimiter."""
|
||||
delimiter = inputs.get("delimiter", " ")
|
||||
if not isinstance(delimiter, str):
|
||||
return None
|
||||
clean = str(inputs.get("clean_whitespace", "true")).lower() == "true"
|
||||
parts = []
|
||||
for key in keys:
|
||||
value = resolve(inputs.get(key, ""))
|
||||
if value is None:
|
||||
return None
|
||||
value = value.strip() if clean else value
|
||||
if value:
|
||||
parts.append(value)
|
||||
return delimiter.join(parts)
|
||||
|
||||
|
||||
def _resolve_string_node(
|
||||
class_type: str, inputs: Dict[str, Any], resolve
|
||||
) -> Optional[str]:
|
||||
"""Output of a known pure string node, or None when the type isn't known."""
|
||||
if class_type in ("PrimitiveString", "PrimitiveStringMultiline"):
|
||||
return resolve(inputs.get("value"))
|
||||
if class_type in ("ShowText|pysssss", "Text Multiline", "String Literal"):
|
||||
return resolve(inputs.get("text", inputs.get("string")))
|
||||
if class_type == "StringConcatenate":
|
||||
a, b = resolve(inputs.get("string_a", "")), resolve(inputs.get("string_b", ""))
|
||||
delimiter = inputs.get("delimiter", "")
|
||||
if a is None or b is None or not isinstance(delimiter, str):
|
||||
return None
|
||||
return f"{a}{delimiter}{b}"
|
||||
if class_type == "Text Concatenate":
|
||||
keys = sorted(k for k in inputs if re.fullmatch(r"text_[a-z]", k))
|
||||
return _join_text_inputs(inputs, keys, resolve)
|
||||
return None
|
||||
|
||||
|
||||
def resolve_text(graph: Any, node_id: str) -> Optional[str]:
|
||||
"""A prompt node's `text` as it will be at run time, if knowable from the graph.
|
||||
|
||||
Typed text, or text produced by known pure string nodes (core primitives and
|
||||
StringConcatenate, WAS Text Concatenate/Multiline, pysssss ShowText). Anything
|
||||
else, including PromptSearchList batches, returns None rather than a guess.
|
||||
Callers compare the result with the text the node actually received.
|
||||
"""
|
||||
if not isinstance(graph, dict):
|
||||
return None
|
||||
seen = set()
|
||||
|
||||
def resolve(value: Any) -> Optional[str]:
|
||||
if isinstance(value, str):
|
||||
return value
|
||||
if not _is_link(value) or len(seen) >= MAX_VISITS:
|
||||
return None
|
||||
source_id = str(value[0])
|
||||
if source_id in seen:
|
||||
return None
|
||||
seen.add(source_id)
|
||||
source = graph.get(source_id)
|
||||
if not isinstance(source, dict):
|
||||
return None
|
||||
return _resolve_string_node(
|
||||
str(source.get("class_type", "")), _inputs(source), resolve
|
||||
)
|
||||
|
||||
seen.add(str(node_id))
|
||||
text = resolve(_inputs(graph.get(str(node_id))).get("text"))
|
||||
return text.strip() if isinstance(text, str) and text.strip() else None
|
||||
+13
-6
@@ -99,7 +99,10 @@ class PromptTracker:
|
||||
self.logger.debug("PromptTracker initialization completed")
|
||||
|
||||
def set_current_prompt(
|
||||
self, prompt_text: str, additional_data: Optional[Dict[str, Any]] = None
|
||||
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.
|
||||
@@ -113,6 +116,9 @@ class PromptTracker:
|
||||
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
|
||||
@@ -164,11 +170,12 @@ class PromptTracker:
|
||||
self.active_prompts[execution_id] = execution_context
|
||||
|
||||
# Push to batch queue for ordered image linking
|
||||
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}"
|
||||
)
|
||||
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})"
|
||||
|
||||
@@ -0,0 +1,107 @@
|
||||
"""Count prompt runs when a workflow is queued, not when nodes execute.
|
||||
|
||||
ComfyUI skips nodes whose inputs are unchanged, so a PromptManager node does not
|
||||
execute when the same prompt is re-run. The on-prompt hook sees every queued
|
||||
workflow instead: it counts a run for each PromptManager node feeding a sampler's
|
||||
positive input. Prompts not yet in the database are remembered as pending and
|
||||
counted by the node when it first saves them.
|
||||
|
||||
Text from known pure string nodes is resolved at queue time too. Anything else
|
||||
(PromptSearchList batches, unknown nodes) is counted by the node as it executes.
|
||||
"""
|
||||
|
||||
import threading
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
from typing import Any, Tuple
|
||||
|
||||
try:
|
||||
from .hashing import generate_prompt_hash
|
||||
from .logging_config import get_logger
|
||||
from .prompt_graph import resolve_text, run_prompt_nodes
|
||||
except ImportError:
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.logging_config import get_logger
|
||||
from utils.prompt_graph import resolve_text, run_prompt_nodes
|
||||
|
||||
logger = get_logger("prompt_manager.usage_tracking")
|
||||
|
||||
PENDING_TTL_SECONDS = 3600 # a queued prompt that never executes is forgotten
|
||||
PENDING_MAX_ENTRIES = 1000
|
||||
|
||||
|
||||
class PendingFirstUse:
|
||||
"""Queued runs of positive prompts that were not in the database yet.
|
||||
|
||||
Counts per hash: the same new prompt can be queued several times before its
|
||||
first job runs, and later jobs may reuse ComfyUI's cached node.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ttl_seconds: float = PENDING_TTL_SECONDS,
|
||||
max_entries: int = PENDING_MAX_ENTRIES,
|
||||
):
|
||||
self._ttl = ttl_seconds
|
||||
self._max = max_entries
|
||||
self._entries: OrderedDict[str, Tuple[int, float]] = OrderedDict()
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def add(self, prompt_hash: str) -> None:
|
||||
with self._lock:
|
||||
count, _ = self._entries.pop(prompt_hash, (0, 0.0))
|
||||
self._entries[prompt_hash] = (count + 1, time.monotonic())
|
||||
while len(self._entries) > self._max:
|
||||
self._entries.popitem(last=False)
|
||||
|
||||
def consume(self, prompt_hash: str) -> int:
|
||||
"""Number of recent queued runs not yet counted (0 if none); clears them."""
|
||||
with self._lock:
|
||||
count, added = self._entries.pop(prompt_hash, (0, 0.0))
|
||||
return count if count and time.monotonic() - added <= self._ttl else 0
|
||||
|
||||
|
||||
# Shared between the queue hook and the nodes
|
||||
PENDING_FIRST_USE = PendingFirstUse()
|
||||
|
||||
|
||||
def handle_queued_prompt(json_data: Any, db: Any, pending: PendingFirstUse) -> Any:
|
||||
"""ComfyUI on-prompt handler body: count runs, always return the request as-is."""
|
||||
try:
|
||||
graph = json_data.get("prompt") if isinstance(json_data, dict) else None
|
||||
counted = set()
|
||||
for node_id in run_prompt_nodes(graph):
|
||||
text = resolve_text(graph, node_id)
|
||||
if not text:
|
||||
continue
|
||||
prompt_hash = generate_prompt_hash(text)
|
||||
if prompt_hash in counted:
|
||||
continue
|
||||
counted.add(prompt_hash)
|
||||
existing = db.get_prompt_by_hash(prompt_hash)
|
||||
if existing:
|
||||
db.record_prompt_use(existing["id"])
|
||||
else:
|
||||
pending.add(prompt_hash)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not count queued prompt runs: {e}")
|
||||
return json_data
|
||||
|
||||
|
||||
_hook_registered = threading.Event()
|
||||
_register_lock = threading.Lock()
|
||||
|
||||
|
||||
def register_queue_hook(
|
||||
server: Any, db_factory: Any, pending: PendingFirstUse = PENDING_FIRST_USE
|
||||
) -> bool:
|
||||
"""Register the on-prompt handler once; db_factory is called lazily per request."""
|
||||
with _register_lock:
|
||||
if _hook_registered.is_set():
|
||||
return False
|
||||
server.add_on_prompt_handler(
|
||||
lambda json_data: handle_queued_prompt(json_data, db_factory(), pending)
|
||||
)
|
||||
_hook_registered.set()
|
||||
logger.info("Registered prompt usage hook")
|
||||
return True
|
||||
+3
-6
@@ -107,12 +107,8 @@
|
||||
<label class="block text-xs font-medium text-pm-secondary mb-1 uppercase tracking-wide">Sort By</label>
|
||||
<select id="sortBy"
|
||||
class="w-full px-2.5 py-1.5 bg-pm-input border border-pm rounded-pm-sm text-pm text-[13px] focus:outline-none">
|
||||
<option value="created_desc">Newest First</option>
|
||||
<option value="created_asc">Oldest First</option>
|
||||
<option value="rating_desc">Highest Rated</option>
|
||||
<option value="rating_asc">Lowest Rated</option>
|
||||
<option value="text_asc">A-Z</option>
|
||||
<option value="text_desc">Z-A</option>
|
||||
<!-- Filled from PromptListSort.SORT_OPTIONS in admin.js -->
|
||||
<option value="last_used_desc">Recently Used</option>
|
||||
</select>
|
||||
</div>
|
||||
<div class="flex items-end">
|
||||
@@ -1391,6 +1387,7 @@ No sentences, no commentary, no captions. Only tags. Keep length short but descr
|
||||
</div>
|
||||
|
||||
<script src="/prompt_manager/js/comfy-metadata.js"></script>
|
||||
<script src="/prompt_manager/js/prompt-list-sort.js"></script>
|
||||
<script src="/prompt_manager/js/admin.js"></script>
|
||||
<script src="/prompt_manager/js/tags-page.js"></script>
|
||||
</body>
|
||||
|
||||
+35
-40
@@ -1,6 +1,7 @@
|
||||
class PromptAdmin {
|
||||
constructor() {
|
||||
this.prompts = [];
|
||||
this.listMode = "recent"; // "recent" or "search"; a sort change reloads it
|
||||
this.selectedPrompts = new Set();
|
||||
this.settings = {
|
||||
resultTimeout: 5,
|
||||
@@ -114,8 +115,13 @@
|
||||
document.getElementById(id).addEventListener("change", () => this.search());
|
||||
});
|
||||
|
||||
// Sort dropdown change
|
||||
document.getElementById("sortBy").addEventListener("change", () => this.renderPrompts());
|
||||
// Sort dropdown: options come from PromptListSort; sorting runs on the server
|
||||
const sortSelect = document.getElementById("sortBy");
|
||||
sortSelect.innerHTML = PromptListSort.SORT_OPTIONS
|
||||
.map((o) => `<option value="${o.value}">${this.escapeHtml(o.label)}</option>`)
|
||||
.join("");
|
||||
sortSelect.value = PromptListSort.DEFAULT_SORT;
|
||||
sortSelect.addEventListener("change", () => this.reloadPrompts());
|
||||
}
|
||||
|
||||
bindModalEvents() {
|
||||
@@ -355,9 +361,13 @@
|
||||
|
||||
async loadRecentPrompts(page = 1) {
|
||||
try {
|
||||
this.listMode = "recent";
|
||||
this.pagination.currentPage = page;
|
||||
const offset = (page - 1) * this.pagination.limit;
|
||||
const response = await fetch(`/prompt_manager/recent?limit=${this.pagination.limit}&offset=${offset}&page=${page}`);
|
||||
const response = await fetch(PromptListSort.buildRecentUrl({
|
||||
page,
|
||||
limit: this.pagination.limit,
|
||||
sort: this.currentSort(),
|
||||
}));
|
||||
if (response.ok) {
|
||||
const data = await response.json();
|
||||
if (data.success) {
|
||||
@@ -424,7 +434,8 @@
|
||||
if (folder) params.append("folder", folder);
|
||||
params.append("limit", "100");
|
||||
|
||||
const response = await fetch(`/prompt_manager/search?${params}`);
|
||||
this.listMode = "search";
|
||||
const response = await fetch(`/prompt_manager/search?${PromptListSort.withSort(params, this.currentSort())}`);
|
||||
if (response.ok) {
|
||||
const data = await response.json();
|
||||
if (data.success) {
|
||||
@@ -442,36 +453,13 @@
|
||||
}
|
||||
}
|
||||
|
||||
sortPrompts(prompts) {
|
||||
const sortBy = document.getElementById("sortBy")?.value || "created_desc";
|
||||
const [field, direction] = sortBy.split("_");
|
||||
|
||||
return [...prompts].sort((a, b) => {
|
||||
let aVal, bVal;
|
||||
|
||||
switch (field) {
|
||||
case "rating":
|
||||
aVal = a.rating || 0;
|
||||
bVal = b.rating || 0;
|
||||
break;
|
||||
case "created":
|
||||
aVal = new Date(a.created_at);
|
||||
bVal = new Date(b.created_at);
|
||||
break;
|
||||
case "text":
|
||||
aVal = a.text.toLowerCase();
|
||||
bVal = b.text.toLowerCase();
|
||||
break;
|
||||
default:
|
||||
return 0;
|
||||
}
|
||||
|
||||
if (direction === "asc") {
|
||||
return aVal > bVal ? 1 : aVal < bVal ? -1 : 0;
|
||||
} else {
|
||||
return aVal < bVal ? 1 : aVal > bVal ? -1 : 0;
|
||||
}
|
||||
});
|
||||
currentSort() {
|
||||
return PromptListSort.normalizeSort(document.getElementById("sortBy")?.value);
|
||||
}
|
||||
|
||||
/** Re-fetch the list currently shown, e.g. after the sort changes. */
|
||||
reloadPrompts() {
|
||||
return this.listMode === "search" ? this.search() : this.loadRecentPrompts(1);
|
||||
}
|
||||
|
||||
renderPrompts() {
|
||||
@@ -493,14 +481,13 @@
|
||||
return;
|
||||
}
|
||||
|
||||
// Sort prompts before rendering
|
||||
const sortedPrompts = this.sortPrompts(this.prompts);
|
||||
container.innerHTML = sortedPrompts.map((prompt) => this.renderPromptItem(prompt)).join("");
|
||||
// Already in the selected order: the server sorts across all pages
|
||||
container.innerHTML = this.prompts.map((prompt) => this.renderPromptItem(prompt)).join("");
|
||||
this.selectedPrompts.clear();
|
||||
this.updateBulkActionButtons();
|
||||
|
||||
// Add hover behavior to all star ratings
|
||||
sortedPrompts.forEach(prompt => {
|
||||
this.prompts.forEach(prompt => {
|
||||
this.addStarHoverBehavior(prompt.id);
|
||||
});
|
||||
|
||||
@@ -514,7 +501,7 @@
|
||||
});
|
||||
|
||||
// Load film strips for each prompt (async, non-blocking)
|
||||
this.loadAllFilmStrips(sortedPrompts);
|
||||
this.loadAllFilmStrips(this.prompts);
|
||||
}
|
||||
|
||||
async loadAllFilmStrips(prompts) {
|
||||
@@ -550,6 +537,10 @@
|
||||
const category = prompt.category || "No category";
|
||||
const rating = prompt.rating || 0;
|
||||
const created = new Date(prompt.created_at).toLocaleDateString();
|
||||
const runLabel = PromptListSort.formatRunCount(prompt.run_count);
|
||||
const lastUsedTitle = prompt.last_used_at
|
||||
? `Last used ${new Date(prompt.last_used_at).toLocaleString()}`
|
||||
: "";
|
||||
|
||||
return `
|
||||
<div class="bg-pm-surface rounded-pm-md border border-pm hover:border-pm transition-all duration-200" data-id="${prompt.id}">
|
||||
@@ -574,6 +565,10 @@
|
||||
<span>📅</span>
|
||||
<span>${created}</span>
|
||||
</div>
|
||||
${runLabel ? `<div class="flex items-center space-x-1" title="${this.escapeHtml(lastUsedTitle)}">
|
||||
<span>🔁</span>
|
||||
<span>${this.escapeHtml(runLabel)}</span>
|
||||
</div>` : ""}
|
||||
<div class="flex items-center space-x-1">
|
||||
<div class="rating flex space-x-1" data-id="${prompt.id}" data-rating="${rating}">
|
||||
${this.renderStars(rating, prompt.id)}
|
||||
|
||||
@@ -0,0 +1,64 @@
|
||||
/**
|
||||
* Prompt list sorting for the admin dashboard (3.2.4, #90.1).
|
||||
*
|
||||
* Sorting happens on the server so it spans every page; this module only owns the
|
||||
* option list, validation and request URLs. The option values must match
|
||||
* SORT_ORDERS in database/operations.py (enforced by tests/test_sort_contract.py).
|
||||
*
|
||||
* Loadable in the browser (window.PromptListSort) and in Node for tests.
|
||||
*/
|
||||
(function (root) {
|
||||
"use strict";
|
||||
|
||||
const SORT_OPTIONS = Object.freeze(
|
||||
[
|
||||
{ value: "last_used_desc", label: "Recently Used" },
|
||||
{ value: "created_desc", label: "Newest First" },
|
||||
{ value: "created_asc", label: "Oldest First" },
|
||||
{ value: "run_count_desc", label: "Most Used" },
|
||||
{ value: "rating_desc", label: "Highest Rated" },
|
||||
{ value: "rating_asc", label: "Lowest Rated" },
|
||||
{ value: "text_asc", label: "A-Z" },
|
||||
{ value: "text_desc", label: "Z-A" },
|
||||
].map((option) => Object.freeze(option)),
|
||||
);
|
||||
|
||||
const DEFAULT_SORT = "last_used_desc";
|
||||
const DEFAULT_LIMIT = 50;
|
||||
const KNOWN = new Set(SORT_OPTIONS.map((o) => o.value));
|
||||
|
||||
/** A known sort key, or DEFAULT_SORT for anything else. */
|
||||
function normalizeSort(value) {
|
||||
return typeof value === "string" && KNOWN.has(value) ? value : DEFAULT_SORT;
|
||||
}
|
||||
|
||||
/** URL for one page of /prompt_manager/recent. */
|
||||
function buildRecentUrl({ page = 1, limit = DEFAULT_LIMIT, sort } = {}) {
|
||||
const params = new URLSearchParams({
|
||||
limit: String(limit),
|
||||
offset: String((page - 1) * limit),
|
||||
page: String(page),
|
||||
sort: normalizeSort(sort),
|
||||
});
|
||||
return `/prompt_manager/recent?${params}`;
|
||||
}
|
||||
|
||||
/** Copy of `params` with a validated sort key; the input is not modified. */
|
||||
function withSort(params, sort) {
|
||||
const copy = new URLSearchParams(params);
|
||||
copy.set("sort", normalizeSort(sort));
|
||||
return copy;
|
||||
}
|
||||
|
||||
/** Badge text for a prompt's run count; empty when the count is unknown. */
|
||||
function formatRunCount(count) {
|
||||
if (!Number.isInteger(count) || count < 0) return "";
|
||||
if (count === 0) return "Never run";
|
||||
if (count === 1) return "1 run";
|
||||
return `${count.toLocaleString("en-US")} runs`;
|
||||
}
|
||||
|
||||
const api = { SORT_OPTIONS, DEFAULT_SORT, normalizeSort, buildRecentUrl, withSort, formatRunCount };
|
||||
if (typeof module !== "undefined" && module.exports) module.exports = api;
|
||||
else root.PromptListSort = api;
|
||||
})(typeof window !== "undefined" ? window : globalThis);
|
||||
@@ -978,7 +978,7 @@ app.registerExtension({
|
||||
url = `/prompt_manager/search?${queryParams.toString()}`;
|
||||
} else if (method === "get_recent_prompts") {
|
||||
const limit = params.limit || 20;
|
||||
url = `/prompt_manager/recent?limit=${limit}`;
|
||||
url = `/prompt_manager/recent?limit=${limit}&sort=last_used_desc`;
|
||||
} else {
|
||||
throw new Error(`Unknown method: ${method}`);
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user