Vito
6c7e4d5b35
* chore: start 3.2.4 * feat: track prompt usage (run_count, last_used_at) with server-side sorting - Migration adds last_used_at and run_count; existing prompts are estimated from linked images (count, newest image time, else created_at). The backfill only touches NULL rows, so it self-heals rows written after a downgrade or an interrupted upgrade. - Every node run records a use, including re-runs of an existing prompt; a tracking failure is logged and never blocks generation. - search_prompts/get_recent_prompts accept a whitelisted sort key (last_used, created, run_count, rating, text); default stays newest first. Refs #90 * feat: sort param on /recent and /search; microsecond last_used_at - /prompt_manager/recent and /search accept ?sort= (whitelisted in the DB layer; unknown values fall back to newest first). - last_used_at is written with microseconds from Python so a re-run in the same millisecond as another save still sorts first; ordering stays correct against millisecond values from the backfill. Refs #90 * feat: server-side sort in admin UI with Recently Used / Most Used and run badges - New web/js/prompt-list-sort.js owns the sort options, validation and request URLs; admin.js fills the dropdown from it and re-fetches from the server on change, so every sort spans all pages (the old client-side sort only reordered the current page). - Default sort is Recently Used; prompt cards show a run-count badge with the last-used time on hover. - The node's Recent button now lists recently used prompts. - Contract test keeps the UI sort keys and the node's key in sync with the server whitelist. Refs #90 * fix: link images and count runs by the positive prompt, not node execution With a negative-prompt PromptManager node in the workflow, every executed PromptManager node pushed onto the image-linking queue, so images could be linked to the negative prompt or to a previous run. ComfyUI also skips unchanged nodes, so re-runs of the same prompt were never counted. - utils/prompt_graph.py finds the PromptManager nodes feeding a sampler's positive input (KSampler, CFG/Basic guiders, ControlNet and combine pass-throughs, PromptManagerText into encoders), falling back to all PromptManager nodes when none can be traced. - Runs are counted by a ComfyUI on-prompt hook when a workflow is queued; prompts not yet saved are counted when the node first saves them. - Nodes receive the queued graph (hidden PROMPT/UNIQUE_ID). Only the run's positive prompt becomes the current prompt; only batch items with linked text are pushed onto the FIFO queue. - The image monitor links from the image's own metadata first, using the positive node, and falls back to the queue for batch items. Refs #90 * fix: address review — count cached runs of linked text, keep repeated pending runs - Text from known pure string nodes (core PrimitiveString/StringConcatenate, WAS Text Concatenate/Multiline, pysssss ShowText) is resolved at queue time, so cached re-runs of such prompts are counted and their images link from metadata. Unknown nodes are never guessed; the node compares its real text with the queue-time resolution and falls back to counting and queueing itself on any mismatch. - A new prompt queued several times before its first run keeps every pending run; the first save records all of them. - Registration flag is a threading.Event; consistent unittest imports; logging instead of print in the init fallback. * refactor: unquote PendingFirstUse annotation so Tuple usage is visible
231 lines
9.6 KiB
Python
231 lines
9.6 KiB
Python
"""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()
|