diff --git a/database/models.py b/database/models.py index cee627a..e38f7ba 100644 --- a/database/models.py +++ b/database/models.py @@ -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. diff --git a/database/operations.py b/database/operations.py index bceb9ad..52c40cd 100644 --- a/database/operations.py +++ b/database/operations.py @@ -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() diff --git a/prompt_manager_base.py b/prompt_manager_base.py index 43a6326..38e9dd0 100644 --- a/prompt_manager_base.py +++ b/prompt_manager_base.py @@ -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. diff --git a/tests/test_prompt_usage.py b/tests/test_prompt_usage.py new file mode 100644 index 0000000..1d9e6ed --- /dev/null +++ b/tests/test_prompt_usage.py @@ -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()