nodes/ComfyAssets-ComfyUI_PromptManager
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: track prompt usage (run_count, last_used_at) with server-side sorting

Browse Source 
- Migration adds last_used_at and run_count; existing prompts are estimated
  from linked images (count, newest image time, else created_at). The
  backfill only touches NULL rows, so it self-heals rows written after a
  downgrade or an interrupted upgrade.
- Every node run records a use, including re-runs of an existing prompt;
  a tracking failure is logged and never blocks generation.
- search_prompts/get_recent_prompts accept a whitelisted sort key
  (last_used, created, run_count, rating, text); default stays newest first.

Refs #90
This commit is contained in:
Vito Sansevero
2026-10-07 11:35:55 -07:00
parent 1cffa22ae2
commit 6099d5eb23
4 changed files with 312 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
database/models.py
+56 -1
View File
@@ -17,6 +17,10 @@ except ImportError:
sys.path.insert(0, current_dir)
from utils.logging_config import get_logger
# Single UTC format for last_used_at so ORDER BY can compare it as text
ISO_TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%f+00:00"
SQL_NOW_ISO = f"strftime('{ISO_TIMESTAMP_FORMAT}', 'now')"
class PromptModel:
"""Database model for prompt storage and schema management."""
@@ -78,7 +82,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 +137,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 +159,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 +451,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.
database/operations.py
+52 -7
View File
@@ -8,7 +8,7 @@ import datetime
import os
from typing import Optional, List, Dict, Any, Union
from .models import PromptModel
from .models import PromptModel, SQL_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.
@@ -120,10 +139,11 @@ class PromptDatabase:
with self.model.get_connection() as conn:
cursor = conn.execute(
"""
f"""
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, ?, ?, ?, ?, ?, {SQL_NOW_ISO})
""",
(
text.strip(),
@@ -142,6 +162,25 @@ class PromptDatabase:
self.logger.debug(f"Successfully saved prompt with ID: {prompt_id}")
return prompt_id
def record_prompt_use(self, prompt_id: int) -> bool:
"""
Count one run of a prompt and mark it as most recently used.
Args:
prompt_id: The prompt ID
Returns:
True if the prompt exists and was updated
"""
with self.model.get_connection() as conn:
cursor = conn.execute(
f"UPDATE prompts SET run_count = run_count + 1, last_used_at = {SQL_NOW_ISO}"
" WHERE id = ?",
(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 +230,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 +247,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 +316,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 +332,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 +354,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()
prompt_manager_base.py
+9
View File
@@ -82,6 +82,7 @@ class PromptManagerBase:
prompt_id=existing["id"], category=category, tags=tags
)
self.logger.debug("Updated metadata for existing prompt")
self._record_use(existing["id"])
return existing["id"]
self.logger.debug(
@@ -93,6 +94,7 @@ class PromptManagerBase:
if prompt_id:
self.logger.debug(f"Successfully saved new prompt with ID: {prompt_id}")
self._record_use(prompt_id)
else:
self.logger.warning("Failed to save prompt - no ID returned")
@@ -102,6 +104,13 @@ class PromptManagerBase:
self.logger.error(f"Error saving prompt to database: {e}")
return None
def _record_use(self, prompt_id: int) -> None:
"""Count this execution; a tracking failure must never block generation."""
try:
self.db.record_prompt_use(prompt_id)
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.
tests/test_prompt_usage.py
+195
View File
@@ -0,0 +1,195 @@
"""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_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):
"""Every node execution counts, including re-runs of an existing prompt."""
def _node(self):
node = PromptManagerBase.__new__(PromptManagerBase)
node.db = self.db
node.logger = __import__("logging").getLogger("test.prompt_usage")
return node
def test_new_and_repeated_runs_are_counted(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"], 2)
if __name__ == "__main__":
unittest.main()