diff --git a/TESTING.md b/TESTING.md index ea33428..c46b73a 100644 --- a/TESTING.md +++ b/TESTING.md @@ -565,3 +565,38 @@ Deploying this change needs no Python change and no ComfyUI restart: the diff fo the editor UI touches only HTML, CSS and JavaScript (plus test fixtures). A browser refresh is enough. New editor modules have additive coverage floors in `tools/vendor/frontend-coverage-baseline.json`; none were lowered. + +## Native vocabulary backend (#236) + +Routes under `/prompt_manager/vocabulary/`: `GET tags` and `GET categories` +(bounded pages, signed epoch-scoped cursors), `POST tags` (create or reuse an exact +name), `PATCH tags/{id}` (rename), `POST tags/{id}/delete`, `GET tags/{id}/impact`, +`POST categories/preview` and `POST categories/apply`. Every command carries +`expected_epoch`; tag commands also carry the exact `{id, name}` identity, and +category apply carries the preview's `impact_token`. + +Error contract: command validation failures are 400 with `code: null`; reads +return `{success: false, error}` with no `code`, and impact on an unknown id is 404. A stale epoch is 409 +`storage-changed`; a changed or missing identity, or a changed category projection, +is 409 `vocabulary-conflict` with a generic message. Both 409 bodies include the +current `library_epoch`. Nothing is written on any refusal. + +Count semantics: a count is the number of eligible Library parents carrying the value +directly or through one active saved pair, exactly what `list_library(tags=[name])` or +`list_library(category=value)` returns. Unused registry tags report 0 and are kept by +`get_all_tags`, `suggest_tags`, backup and restore, duplicate cleanup, orphan pruning +and owner deletion (`TestRetainedTags` in `tests/test_database_backup.py`). Restore +rotates the epoch, so a pre-restore `expected_epoch` is refused with `storage-changed`. +Owner rows and their timestamps never change on a tag command; a category apply only +touches rows carrying the exact source. + +Two legacy routes change additively: `GET /tags` now lists unused registry tags, and +`/tags/suggest` without context appends unused names at count 0. Unused tags become +deletable only through the native Delete tag (vocabulary UI PR), not on the legacy +tags page. + +Both token functions (`library_edit_token`, `category_impact_token`) share +`canonical_token` in `database/library_tokens.py`; `TestSharedTokenHashing` pins their +outputs. + +Deploying this backend needs a ComfyUI restart (Python and route changes). diff --git a/database/library_commands.py b/database/library_commands.py index 8f43cf0..caf6c71 100644 --- a/database/library_commands.py +++ b/database/library_commands.py @@ -26,7 +26,8 @@ _PROMPT_COLUMNS = ("category", "rating", "notes") class LibraryCommandError(Exception): """A native command lost a race against storage or another writer. - ``code`` is a wire code (``metadata-conflict`` or ``storage-changed``), + ``code`` is a wire code (``metadata-conflict``, ``storage-changed`` or + ``vocabulary-conflict`` for a tag or category that changed since it was read), ``epoch`` the storage epoch read inside the transaction and ``positive`` the current owner record when it may be disclosed. """ diff --git a/database/library_tokens.py b/database/library_tokens.py index 853408c..efdf5e1 100644 --- a/database/library_tokens.py +++ b/database/library_tokens.py @@ -6,6 +6,18 @@ import json _DOMAIN = "prompt-manager-library-parent:v1\n" +def canonical_token(domain: str, preimage: dict) -> str: + """SHA-256 of ``domain`` plus the canonical JSON of ``preimage``.""" + payload = json.dumps( + preimage, + sort_keys=True, + separators=(",", ":"), + ensure_ascii=False, + allow_nan=False, + ) + return hashlib.sha256((domain + payload).encode("utf-8")).hexdigest() + + def library_edit_token(epoch: str, positive: dict) -> str: """Hash the parent's owner fields under the current storage epoch. @@ -22,11 +34,4 @@ def library_edit_token(epoch: str, positive: dict) -> str: "rating": positive["rating"], "tags": sorted(set(positive["tags"])), } - payload = json.dumps( - preimage, - sort_keys=True, - separators=(",", ":"), - ensure_ascii=False, - allow_nan=False, - ) - return hashlib.sha256((_DOMAIN + payload).encode("utf-8")).hexdigest() + return canonical_token(_DOMAIN, preimage) diff --git a/database/operations.py b/database/operations.py index 36773ae..e210199 100644 --- a/database/operations.py +++ b/database/operations.py @@ -12,6 +12,7 @@ from .models import PromptModel, normalize_image_path, utc_now_iso from .pairs import PairDependencyError, PairOperations from .library import LibraryOperations from .library_commands import LibraryCommandOperations +from .vocabulary import VocabularyOperations from .execution_history import ExecutionHistoryOperations from .execution_contexts import ExecutionContextOperations, LIVE_HINT_CTES from .execution_outputs import ExecutionOutputOperations @@ -182,6 +183,7 @@ class PromptDatabase( PairOperations, LibraryOperations, LibraryCommandOperations, + VocabularyOperations, ExecutionHistoryOperations, ExecutionContextOperations, ExecutionOutputOperations, @@ -743,18 +745,13 @@ class PromptDatabase( def get_all_tags(self) -> List[str]: """ - Get all unique tags linked to a standalone prompt or reusable pair. + Get every tag in the shared vocabulary, including unused registry rows. Returns: Sorted list of tag names """ with self.model.get_connection() as conn: - cursor = conn.execute( - "SELECT t.name AS tag FROM tags t WHERE " - "EXISTS(SELECT 1 FROM prompt_tags pt WHERE pt.tag_id=t.id) OR " - "EXISTS(SELECT 1 FROM prompt_pair_tags pp WHERE pp.tag_id=t.id) " - "ORDER BY t.name" - ) + cursor = conn.execute("SELECT t.name AS tag FROM tags t ORDER BY t.name") return [row["tag"] for row in cursor.fetchall()] def get_tags_with_counts( @@ -830,6 +827,8 @@ class PromptDatabase( counted over the prompts that carry every tag in ``with_tags``; the context tags themselves are left out. Counts include separate standalone and pair references, with owner types kept distinct. Ordered by count/name. + Without context, unused vocabulary rows are included last with count 0; + with context they never appear. """ seen = set() context = [] @@ -861,8 +860,8 @@ class PromptDatabase( "WITH tag_refs(kind,owner_id,tag_id) AS (" "SELECT 'prompt',prompt_id,tag_id FROM prompt_tags UNION ALL " "SELECT 'pair',pair_id,tag_id FROM prompt_pair_tags) " - "SELECT t.name AS name, COUNT(*) AS count" - " FROM tag_refs pt JOIN tags t ON t.id = pt.tag_id" + "SELECT t.name AS name, COUNT(pt.tag_id) AS count" + " FROM tags t LEFT JOIN tag_refs pt ON t.id = pt.tag_id" f" WHERE {' AND '.join(where)}" " GROUP BY t.id" " ORDER BY count DESC, t.name COLLATE NOCASE ASC" @@ -949,6 +948,23 @@ class PromptDatabase( "has_more": (offset + limit) < total, } + @staticmethod + def _move_tag_links( + conn: sqlite3.Connection, source_id: int, target_id: int + ) -> None: + """Repoint component and pair links from source to target. + + Owners that already carry the target keep that one link; the source's + leftover duplicates are removed. The source registry row is left for the + caller to delete. + """ + for junction in ("prompt_tags", "prompt_pair_tags"): + conn.execute( + f"UPDATE OR IGNORE {junction} SET tag_id = ? WHERE tag_id = ?", + (target_id, source_id), + ) + conn.execute(f"DELETE FROM {junction} WHERE tag_id = ?", (source_id,)) + def rename_tag_all_prompts(self, old_name: str, new_name: str) -> Dict[str, Any]: """ Rename a tag across all prompts that use it. @@ -991,15 +1007,7 @@ class PromptDatabase( if existing_new: # Target tag exists — merge: move links, handle conflicts, delete old - new_tag_id = existing_new["id"] - for junction in ("prompt_tags", "prompt_pair_tags"): - conn.execute( - f"UPDATE OR IGNORE {junction} SET tag_id = ? WHERE tag_id = ?", - (new_tag_id, old_tag_id), - ) - conn.execute( - f"DELETE FROM {junction} WHERE tag_id = ?", (old_tag_id,) - ) + self._move_tag_links(conn, old_tag_id, existing_new["id"]) conn.execute("DELETE FROM tags WHERE id = ?", (old_tag_id,)) else: # Simple rename @@ -1097,15 +1105,7 @@ class PromptDatabase( "SELECT COUNT(*) FROM prompt_pair_tags WHERE tag_id = ?", (src_id,) ).fetchone()[0] if src_count or pair_count: - for junction in ("prompt_tags", "prompt_pair_tags"): - conn.execute( - f"UPDATE OR IGNORE {junction} " - "SET tag_id = ? WHERE tag_id = ?", - (target_id, src_id), - ) - conn.execute( - f"DELETE FROM {junction} WHERE tag_id = ?", (src_id,) - ) + self._move_tag_links(conn, src_id, target_id) tags_merged += 1 affected += src_count diff --git a/database/vocabulary.py b/database/vocabulary.py new file mode 100644 index 0000000..a3152d9 --- /dev/null +++ b/database/vocabulary.py @@ -0,0 +1,435 @@ +"""Bounded tag and category reads, and guarded tag commands. + +A count is the number of eligible Library parents that carry the value +directly or through one active saved pair, which is exactly what +``list_library(tags=[name])`` / ``list_library(category=value)`` return. + +Tag commands run in one ``_pair_transaction``: the storage epoch is checked +first, then the exact ``{id, name}`` identity, and the receipt is read on the +same connection before commit. Owner rows and their timestamps never change. + +Category consolidation is previewed in one ``library_snapshot`` and applied in +one ``_pair_transaction`` that recomputes the owner projection and requires the +preview's ``impact_token``; only rows carrying the exact source change. +""" + +import sqlite3 + +from .library import _ELIGIBLE, _limit +from .library_commands import _TOKEN, LibraryCommandError, _validate_epoch +from .library_tokens import canonical_token +from .models import utc_now_iso + +try: + from ..utils.validators import validate_category, validate_tags +except ImportError: + from utils.validators import validate_category, validate_tags + +_CATEGORY_DOMAIN = "prompt-manager-vocabulary-category:v1\n" + + +def _tag_rows(page: str) -> str: + """Count rows for the registry rows selected by ``page`` (id, name).""" + # Eligible parent components carrying a tag in the page, plus the + # positive parent of every active pair carrying it (always eligible). + return ( + "WITH page AS (" + page + ")," + "owners(tag_id,pid) AS (" + "SELECT pt.tag_id,p.id FROM prompts p " + "JOIN prompt_tags pt ON pt.prompt_id=p.id " + "WHERE pt.tag_id IN (SELECT id FROM page) AND " + _ELIGIBLE + " " + "UNION SELECT ppt.tag_id,pp.positive_prompt_id FROM prompt_pairs pp " + "JOIN prompt_pair_tags ppt ON ppt.pair_id=pp.id " + "WHERE pp.archived=0 AND ppt.tag_id IN (SELECT id FROM page)) " + "SELECT page.id,page.name,COUNT(DISTINCT owners.pid) AS prompt_count " + "FROM page LEFT JOIN owners ON owners.tag_id=page.id " + "GROUP BY page.id ORDER BY page.id" + ) + + +_TAG_PAGE = _tag_rows("SELECT id,name FROM tags WHERE id>? ORDER BY id LIMIT ?") +_TAG_ONE = _tag_rows("SELECT id,name FROM tags WHERE id=?") +# Every owner kind a tag can sit on, for the delete confirmation. +_TAG_IMPACT = ( + "SELECT " + "(SELECT COUNT(*) FROM prompts p JOIN prompt_tags pt ON pt.prompt_id=p.id " + "WHERE pt.tag_id=:tag AND " + _ELIGIBLE + ") AS positive_prompts," + "(SELECT COUNT(*) FROM prompts p JOIN prompt_tags pt ON pt.prompt_id=p.id " + "WHERE pt.tag_id=:tag AND NOT " + _ELIGIBLE + ") AS other_prompts," + "(SELECT COUNT(*) FROM prompt_pairs pp JOIN prompt_pair_tags ppt " + "ON ppt.pair_id=pp.id WHERE ppt.tag_id=:tag AND pp.archived=0) " + "AS negative_variants," + "(SELECT COUNT(*) FROM prompt_pairs pp JOIN prompt_pair_tags ppt " + "ON ppt.pair_id=pp.id WHERE ppt.tag_id=:tag AND pp.archived!=0) " + "AS archived_negative_variants" +) +_MAX_ID = 2**63 - 1 + + +def validate_vocabulary_tag_name(value: object) -> str: + """Return the trimmed name, applying tag rules plus the native comma ban.""" + if not isinstance(value, str): + raise ValueError("Tag name must be a string") + validate_tags([value]) + name = value.strip() + if "," in name: + raise ValueError("Tag names cannot contain commas") + return name + + +def _validate_tag_id(tag_id: object) -> None: + if type(tag_id) is not int or not 1 <= tag_id <= _MAX_ID: + raise ValueError("Invalid tag ID") + + +# The /library/options universe: nonblank values on eligible parents or on +# active pairs, grouped exactly (whitespace and case variants stay distinct). +_CATEGORY_GROUPS = ( + "WITH owners(value,pid) AS (" + "SELECT p.category,p.id FROM prompts p WHERE p.category IS NOT NULL " + "AND TRIM(p.category)!='' AND " + _ELIGIBLE + " " + "UNION SELECT pp.category,pp.positive_prompt_id FROM prompt_pairs pp " + "WHERE pp.archived=0 AND pp.category IS NOT NULL AND TRIM(pp.category)!='')," + "grouped AS (SELECT value,MIN(pid) AS first_pid," + "COUNT(DISTINCT pid) AS prompt_count FROM owners GROUP BY value) " +) + + +# Every owner of one exact category value, in id order, with the flag the +# impact counts split on. +_CATEGORY_COMPONENTS = ( + "SELECT p.id," + _ELIGIBLE + " AS eligible FROM prompts p " + "WHERE p.category=? ORDER BY p.id" +) +_CATEGORY_PAIRS = "SELECT id,archived FROM prompt_pairs WHERE category=? ORDER BY id" +_CATEGORY_EXISTS = ( + "SELECT EXISTS(SELECT 1 FROM prompts WHERE category=:v)" + " OR EXISTS(SELECT 1 FROM prompt_pairs WHERE category=:v)" +) + +# Pairs after components, so a failing pair update rolls both back. +_CATEGORY_UPDATES = ( + "UPDATE prompts SET category=?,updated_at=? WHERE category=?", + "UPDATE prompt_pairs SET category=?,updated_at=? WHERE category=?", +) + + +def category_impact_token( + epoch: str, + source: str, + target: str, + target_exists: bool, + components: list, + pairs: list, +) -> str: + """Hash a category projection like ``library_edit_token`` (no secret).""" + preimage = { + "epoch": epoch, + "source": source, + "target": target, + "target_exists": bool(target_exists), + "components": sorted([int(i), bool(flag)] for i, flag in components), + "pairs": sorted([int(i), bool(flag)] for i, flag in pairs), + } + return canonical_token(_CATEGORY_DOMAIN, preimage) + + +def _validate_category_change(source: object, target: object) -> str: + """Check the exact source and return the trimmed, different target.""" + if not isinstance(source, str) or not source.strip(): + raise ValueError("source must be a nonblank category value") + if not isinstance(target, str): + raise ValueError("Category must be a string") + validate_category(target) + target = target.strip() + if not target: + raise ValueError("New category name cannot be blank") + if target == source: + raise ValueError("New category name must differ from the current one") + return target + + +def _validate_impact_token(impact_token: object) -> None: + if not isinstance(impact_token, str) or not _TOKEN.fullmatch(impact_token): + raise ValueError("impact_token must be 64 lowercase hexadecimal characters") + + +def _tag_after(after: object) -> int: + if after is None: + return 0 + if type(after) is not int or after <= 0: + raise ValueError("Invalid vocabulary cursor") + return after + + +def _category_after(after: object) -> tuple[int, str] | None: + if after is None: + return None + if ( + not isinstance(after, list) + or len(after) != 2 + or type(after[0]) is not int + or not isinstance(after[1], str) + ): + raise ValueError("Invalid vocabulary cursor") + return after[0], after[1] + + +class VocabularyOperations: + """Vocabulary reads and tag commands. + + ``_`` helpers take an open ``library_snapshot`` or ``_pair_transaction`` + connection. + """ + + def _vocabulary_tag(self, conn: sqlite3.Connection, tag_id: int) -> dict | None: + row = conn.execute(_TAG_ONE, (tag_id,)).fetchone() + return dict(row) if row else None + + def _vocabulary_identity( + self, conn: sqlite3.Connection, epoch: str, tag_id: int, name: str + ) -> dict: + """The stored tag, or vocabulary-conflict when it is gone or renamed.""" + tag = self._vocabulary_tag(conn, tag_id) + if tag is None or tag["name"] != name: + raise LibraryCommandError("vocabulary-conflict", epoch) + return tag + + @staticmethod + def _tag_owner_counts(conn: sqlite3.Connection, tag_id: int) -> dict: + components = conn.execute( + "SELECT COUNT(DISTINCT prompt_id) FROM prompt_tags WHERE tag_id=?", + (tag_id,), + ).fetchone()[0] + pairs = conn.execute( + "SELECT COUNT(DISTINCT pair_id) FROM prompt_pair_tags WHERE tag_id=?", + (tag_id,), + ).fetchone()[0] + return {"components_changed": components, "pairs_changed": pairs} + + def create_vocabulary_tag(self, name: str, *, expected_epoch: str) -> dict: + """Insert an unused registry row, or reuse the exact existing name.""" + _validate_epoch(expected_epoch) + name = validate_vocabulary_tag_name(name) + with self._pair_transaction() as conn: + epoch = self._library_epoch(conn, expected_epoch) + row = conn.execute("SELECT id FROM tags WHERE name=?", (name,)).fetchone() + created = row is None + if created: + tag_id = conn.execute( + "INSERT INTO tags(name) VALUES(?)", (name,) + ).lastrowid + else: + tag_id = row[0] + return { + "library_epoch": epoch, + "created": created, + "tag": self._vocabulary_tag(conn, tag_id), + } + + def rename_vocabulary_tag( + self, tag_id: int, *, expected_epoch: str, name: str, new_name: str + ) -> dict: + """Rename ``{tag_id, name}``, merging into an existing exact ``new_name``. + + Links move across components and all pairs, archived included; an + owner already carrying the target keeps one link. + """ + _validate_tag_id(tag_id) + _validate_epoch(expected_epoch) + if not isinstance(name, str): + raise ValueError("name must be the tag's current name") + new_name = validate_vocabulary_tag_name(new_name) + if new_name == name: + raise ValueError("New name must differ from the current name") + with self._pair_transaction() as conn: + epoch = self._library_epoch(conn, expected_epoch) + self._vocabulary_identity(conn, epoch, tag_id, name) + counts = self._tag_owner_counts(conn, tag_id) + target = conn.execute( + "SELECT id FROM tags WHERE name=?", (new_name,) + ).fetchone() + if target is None: + conn.execute("UPDATE tags SET name=? WHERE id=?", (new_name, tag_id)) + result_id = tag_id + else: + result_id = target[0] + self._move_tag_links(conn, tag_id, result_id) + conn.execute("DELETE FROM tags WHERE id=?", (tag_id,)) + return { + "library_epoch": epoch, + "merged": target is not None, + "tag": self._vocabulary_tag(conn, result_id), + **counts, + } + + def delete_vocabulary_tag( + self, tag_id: int, *, expected_epoch: str, name: str + ) -> dict: + """Remove the registry row ``{tag_id, name}``; its links cascade.""" + _validate_tag_id(tag_id) + _validate_epoch(expected_epoch) + if not isinstance(name, str): + raise ValueError("name must be the tag's current name") + with self._pair_transaction() as conn: + epoch = self._library_epoch(conn, expected_epoch) + self._vocabulary_identity(conn, epoch, tag_id, name) + counts = self._tag_owner_counts(conn, tag_id) + conn.execute("DELETE FROM tags WHERE id=?", (tag_id,)) + return {"library_epoch": epoch, "deleted": True, **counts} + + def vocabulary_tag_impact(self, tag_id: int) -> dict | None: + """Owner counts by kind for one tag, or None when it does not exist.""" + _validate_tag_id(tag_id) + with self.library_snapshot() as (conn, epoch, _ceiling): + if not conn.execute("SELECT 1 FROM tags WHERE id=?", (tag_id,)).fetchone(): + return None + row = conn.execute(_TAG_IMPACT, {"tag": tag_id}).fetchone() + return {"library_epoch": epoch, "impact": dict(row)} + + @staticmethod + def _category_projection( + conn: sqlite3.Connection, epoch: str, source: str, target: str + ) -> dict: + """Owners of ``source``, impact counts and the token over them.""" + components = [ + [row[0], bool(row[1])] + for row in conn.execute(_CATEGORY_COMPONENTS, (source,)) + ] + pairs = [ + [row[0], bool(row[1])] for row in conn.execute(_CATEGORY_PAIRS, (source,)) + ] + if not components and not pairs: + raise LibraryCommandError("vocabulary-conflict", epoch) + target_exists = bool( + conn.execute(_CATEGORY_EXISTS, {"v": target}).fetchone()[0] + ) + eligible = sum(flag for _id, flag in components) + archived = sum(flag for _id, flag in pairs) + return { + "target_exists": target_exists, + "impact": { + "positive_prompts": eligible, + "other_prompts": len(components) - eligible, + "negative_variants": len(pairs) - archived, + "archived_negative_variants": archived, + }, + "impact_token": category_impact_token( + epoch, source, target, target_exists, components, pairs + ), + } + + def preview_vocabulary_category( + self, *, expected_epoch: str, source: str, target: str + ) -> dict: + """Read-only impact of renaming the exact ``source`` to ``target``.""" + _validate_epoch(expected_epoch) + target = _validate_category_change(source, target) + with self.library_snapshot() as (conn, epoch, _ceiling): + if epoch != expected_epoch: + raise LibraryCommandError("storage-changed", epoch) + projection = self._category_projection(conn, epoch, source, target) + return { + "library_epoch": epoch, + "source": source, + "target": target, + **projection, + } + + def apply_vocabulary_category( + self, *, expected_epoch: str, source: str, target: str, impact_token: str + ) -> dict: + """Set the exact ``source`` to ``target`` on every component and pair. + + The recomputed projection must still match the preview's token; a + mismatch is ``vocabulary-conflict`` with nothing written. + """ + _validate_epoch(expected_epoch) + target = _validate_category_change(source, target) + _validate_impact_token(impact_token) + with self._pair_transaction() as conn: + epoch = self._library_epoch(conn, expected_epoch) + projection = self._category_projection(conn, epoch, source, target) + if projection["impact_token"] != impact_token: + raise LibraryCommandError("vocabulary-conflict", epoch) + stamp = utc_now_iso() + for statement in _CATEGORY_UPDATES: + conn.execute(statement, (target, stamp, source)) + return { + "library_epoch": epoch, + "source": source, + "target": target, + "consolidated": projection["target_exists"], + "impact": projection["impact"], + } + + def _vocabulary_tags( + self, + conn: sqlite3.Connection, + epoch: str, + *, + limit: int = 100, + after: int | None = None, + ) -> dict: + _limit(limit) + start = _tag_after(after) + total = conn.execute("SELECT COUNT(*) FROM tags").fetchone()[0] + rows = [dict(row) for row in conn.execute(_TAG_PAGE, (start, limit + 1))] + has_more = len(rows) > limit + rows = rows[:limit] + return { + "library_epoch": epoch, + "tags": rows, + "total": total, + "has_more": has_more, + "next_key": rows[-1]["id"] if has_more else None, + } + + def _vocabulary_categories( + self, + conn: sqlite3.Connection, + epoch: str, + *, + limit: int = 100, + after: list | None = None, + ) -> dict: + _limit(limit) + key = _category_after(after) + total = conn.execute( + _CATEGORY_GROUPS + "SELECT COUNT(*) FROM grouped" + ).fetchone()[0] + where, params = "", [] + if key is not None: + where = " WHERE first_pid>? OR (first_pid=? AND value>?)" + params = [key[0], key[0], key[1]] + rows = conn.execute( + _CATEGORY_GROUPS + + "SELECT value,first_pid,prompt_count FROM grouped" + + where + + " ORDER BY first_pid,value LIMIT ?", + [*params, limit + 1], + ).fetchall() + has_more = len(rows) > limit + rows = rows[:limit] + return { + "library_epoch": epoch, + "categories": [ + {"value": row["value"], "prompt_count": row["prompt_count"]} + for row in rows + ], + "total": total, + "has_more": has_more, + "next_key": ( + [rows[-1]["first_pid"], rows[-1]["value"]] if has_more else None + ), + } + + def list_vocabulary_tags(self, *, limit: int = 100, after: int | None = None): + with self.library_snapshot() as (conn, epoch, _ceiling): + return self._vocabulary_tags(conn, epoch, limit=limit, after=after) + + def list_vocabulary_categories( + self, *, limit: int = 100, after: list | None = None + ) -> dict: + with self.library_snapshot() as (conn, epoch, _ceiling): + return self._vocabulary_categories(conn, epoch, limit=limit, after=after) diff --git a/py/api/__init__.py b/py/api/__init__.py index 404e0c4..ed24041 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -24,6 +24,7 @@ from PIL import Image from .prompts import PromptRoutesMixin from .pairs import PairRoutesMixin from .library import LibraryRoutesMixin +from .vocabulary import VocabularyRoutesMixin from .history import HistoryRoutesMixin from .images import ImageRoutesMixin from .admin import AdminRoutesMixin @@ -246,6 +247,7 @@ async def _gzip_middleware(request, handler): class PromptManagerAPI( PairRoutesMixin, LibraryRoutesMixin, + VocabularyRoutesMixin, HistoryRoutesMixin, PromptRoutesMixin, ImageRoutesMixin, @@ -406,6 +408,7 @@ class PromptManagerAPI( self._register_pair_routes(routes) self._register_library_routes(routes) + self._register_vocabulary_routes(routes) self._register_history_routes(routes) self._register_prompt_routes(routes) self._register_image_routes(routes) diff --git a/py/api/library.py b/py/api/library.py index 568d1da..3532e49 100644 --- a/py/api/library.py +++ b/py/api/library.py @@ -169,6 +169,9 @@ _ADD_NEGATIVE_FIELDS = frozenset({"expected_epoch", "if_match", "negative_text"} _COMMAND_MESSAGES = { "metadata-conflict": "This prompt changed since it was loaded; review it first", "storage-changed": "Prompt storage was restored; refresh the library", + "vocabulary-conflict": ( + "This tag or category changed since it was loaded; review it first" + ), } _GENERIC_CONFLICT = "Library command conflicted; refresh the library" diff --git a/py/api/vocabulary.py b/py/api/vocabulary.py new file mode 100644 index 0000000..8a579a5 --- /dev/null +++ b/py/api/vocabulary.py @@ -0,0 +1,164 @@ +"""Native vocabulary: bounded, signed-cursor reads and guarded tag and +category commands.""" + +from aiohttp import web + +from .library import ( + _command_body, + _decode, + _encode, + _scope, + library_command_response, +) +from .pairs import pair_response + +_DEFAULT_LIMIT = 100 +_CREATE_TAG_FIELDS = frozenset({"expected_epoch", "name"}) +_RENAME_TAG_FIELDS = frozenset({"expected_epoch", "name", "new_name"}) +_DELETE_TAG_FIELDS = frozenset({"expected_epoch", "name"}) +_PREVIEW_CATEGORY_FIELDS = frozenset({"expected_epoch", "source", "target"}) +_APPLY_CATEGORY_FIELDS = _PREVIEW_CATEGORY_FIELDS | {"impact_token"} + + +def _tag_id(request: web.Request) -> int: + try: + value = int(request.match_info["tag_id"]) + except ValueError: + raise ValueError("Invalid tag ID") from None + if not 1 <= value <= 2**63 - 1: + raise ValueError("Invalid tag ID") + return value + + +def _vocabulary_limit(request: web.Request) -> int: + value = int(request.query.get("limit", str(_DEFAULT_LIMIT))) + if not 1 <= value <= 200: + raise ValueError("limit must be between 1 and 200") + return value + + +class VocabularyRoutesMixin: + def _register_vocabulary_routes(self, routes: web.RouteTableDef) -> None: + routes.get("/prompt_manager/vocabulary/tags")(self.list_vocabulary_tags) + routes.post("/prompt_manager/vocabulary/tags")(self.create_vocabulary_tag) + routes.patch("/prompt_manager/vocabulary/tags/{tag_id}")( + self.rename_vocabulary_tag + ) + routes.post("/prompt_manager/vocabulary/tags/{tag_id}/delete")( + self.delete_vocabulary_tag + ) + routes.get("/prompt_manager/vocabulary/tags/{tag_id}/impact")( + self.get_vocabulary_tag_impact + ) + routes.get("/prompt_manager/vocabulary/categories")( + self.list_vocabulary_categories + ) + routes.post("/prompt_manager/vocabulary/categories/preview")( + self.preview_vocabulary_category + ) + routes.post("/prompt_manager/vocabulary/categories/apply")( + self.apply_vocabulary_category + ) + + async def _vocabulary_page( + self, request: web.Request, kind: str, reader_name: str, name: bool + ) -> web.Response: + limit = _vocabulary_limit(request) + cursor = request.query.get("cursor") + + def read() -> dict: + # The snapshot stays open while the cursor is decoded and encoded, + # so its epoch scope is the epoch the page was read from. + with self.db.library_snapshot() as (conn, epoch, _ceiling): + scope = _scope({"kind": kind, "epoch": epoch}) + result = getattr(self.db, reader_name)( + conn, epoch, limit=limit, after=_decode(cursor, scope) + ) + result["next_cursor"] = _encode( + result.pop("next_key"), scope, name=name + ) + return result + + result = await self._run_in_executor(read) + return web.json_response({"success": True, **result}) + + @pair_response + async def list_vocabulary_tags(self, request: web.Request) -> web.Response: + return await self._vocabulary_page( + request, "vocabulary-tags", "_vocabulary_tags", False + ) + + @pair_response + async def list_vocabulary_categories(self, request: web.Request) -> web.Response: + return await self._vocabulary_page( + request, "vocabulary-categories", "_vocabulary_categories", True + ) + + @library_command_response + async def create_vocabulary_tag(self, request: web.Request) -> web.Response: + body = await _command_body(request, _CREATE_TAG_FIELDS) + result = await self._run_in_executor( + self.db.create_vocabulary_tag, + body["name"], + expected_epoch=body["expected_epoch"], + ) + return web.json_response({"success": True, **result}) + + @library_command_response + async def rename_vocabulary_tag(self, request: web.Request) -> web.Response: + tag_id = _tag_id(request) + body = await _command_body(request, _RENAME_TAG_FIELDS) + result = await self._run_in_executor( + self.db.rename_vocabulary_tag, + tag_id, + expected_epoch=body["expected_epoch"], + name=body["name"], + new_name=body["new_name"], + ) + return web.json_response({"success": True, **result}) + + @library_command_response + async def delete_vocabulary_tag(self, request: web.Request) -> web.Response: + tag_id = _tag_id(request) + body = await _command_body(request, _DELETE_TAG_FIELDS) + result = await self._run_in_executor( + self.db.delete_vocabulary_tag, + tag_id, + expected_epoch=body["expected_epoch"], + name=body["name"], + ) + return web.json_response({"success": True, **result}) + + @pair_response + async def get_vocabulary_tag_impact(self, request: web.Request) -> web.Response: + result = await self._run_in_executor( + self.db.vocabulary_tag_impact, _tag_id(request) + ) + if result is None: + return web.json_response( + {"success": False, "error": "Tag not found"}, status=404 + ) + return web.json_response({"success": True, **result}) + + @library_command_response + async def preview_vocabulary_category(self, request: web.Request) -> web.Response: + body = await _command_body(request, _PREVIEW_CATEGORY_FIELDS) + result = await self._run_in_executor( + self.db.preview_vocabulary_category, + expected_epoch=body["expected_epoch"], + source=body["source"], + target=body["target"], + ) + return web.json_response({"success": True, **result}) + + @library_command_response + async def apply_vocabulary_category(self, request: web.Request) -> web.Response: + body = await _command_body(request, _APPLY_CATEGORY_FIELDS) + result = await self._run_in_executor( + self.db.apply_vocabulary_category, + expected_epoch=body["expected_epoch"], + source=body["source"], + target=body["target"], + impact_token=body["impact_token"], + ) + return web.json_response({"success": True, **result}) diff --git a/tests/test_api_prompts.py b/tests/test_api_prompts.py index 18b1c21..21bb3ca 100644 --- a/tests/test_api_prompts.py +++ b/tests/test_api_prompts.py @@ -255,6 +255,19 @@ class TestListRoutesErrorPaths(RouteCoverageCase): self.assertEqual(c_data["categories"], []) self.assertEqual(t_data["tags"], []) + async def test_tags_lists_unused_row_and_suggest_shape_is_stable(self): + pid = self._save_prompt("P1", tags=["used", "unused"]) + self.api.db.set_prompt_tags(pid, ["used"]) + status, data = await self._json("GET", "/prompt_manager/tags") + self.assertEqual(status, 200) + self.assertEqual(data["tags"], ["unused", "used"]) + s_status, s_data = await self._json("GET", "/prompt_manager/tags/suggest?q=") + self.assertEqual(s_status, 200) + self.assertEqual( + s_data["suggestions"], + [{"name": "used", "count": 1}, {"name": "unused", "count": 0}], + ) + async def test_subfolders_success_and_failure(self): ok_status, ok_data = await self._json( "GET", "/prompt_manager/subfolders?include_ancestors=true" diff --git a/tests/test_api_vocabulary.py b/tests/test_api_vocabulary.py new file mode 100644 index 0000000..a1d4a5a --- /dev/null +++ b/tests/test_api_vocabulary.py @@ -0,0 +1,524 @@ +"""Vocabulary read routes page with signed, epoch-scoped cursors.""" + +import json +import unittest +from unittest.mock import patch + +from database.execution_history_schema import rotate_history_epoch + +try: + from test_api_images import ImageAPITestCase +except ImportError: + from tests.test_api_images import ImageAPITestCase + + +class TestVocabularyAPI(ImageAPITestCase): + async def get(self, path, **params): + response = await self.client.get( + "/prompt_manager/vocabulary/" + path, params=params + ) + return response, await response.json() + + def add_unused(self, *names): + with self.api.db.model.get_connection() as conn: + for name in names: + conn.execute("INSERT INTO tags(name) VALUES(?)", (name,)) + + async def test_tag_and_category_envelopes_and_continuation(self): + self._save_prompt("Cat", category="Animal", tags=["cat"]) + self.api.db.save_prompt_pair("Dog", "Noise", category="Pet", tags=["dog"]) + self.add_unused("unused") + epoch = self.api.db.get_execution_epoch() + response, first = await self.get("tags", limit=2) + self.assertEqual(response.status, 200) + self.assertEqual( + set(first), + {"success", "library_epoch", "tags", "total", "has_more", "next_cursor"}, + ) + self.assertEqual(first["library_epoch"], epoch) + self.assertEqual(first["total"], 3) + self.assertTrue(first["has_more"]) + _, second = await self.get("tags", limit=2, cursor=first["next_cursor"]) + self.assertEqual( + [(row["name"], row["prompt_count"]) for row in first["tags"]] + + [(row["name"], row["prompt_count"]) for row in second["tags"]], + [("cat", 1), ("dog", 1), ("unused", 0)], + ) + self.assertFalse(second["has_more"]) + self.assertIsNone(second["next_cursor"]) + response, categories = await self.get("categories", limit=1) + self.assertEqual(response.status, 200) + self.assertEqual( + set(categories), + { + "success", + "library_epoch", + "categories", + "total", + "has_more", + "next_cursor", + }, + ) + self.assertEqual( + categories["categories"], [{"value": "Animal", "prompt_count": 1}] + ) + _, rest = await self.get( + "categories", limit=1, cursor=categories["next_cursor"] + ) + self.assertEqual(rest["categories"], [{"value": "Pet", "prompt_count": 1}]) + self.assertFalse(rest["has_more"]) + + async def test_vocabulary_cursor_is_signed_epoch_scoped(self): + self.add_unused("a", "b", "c") + self._save_prompt("One", category="One") + self._save_prompt("Two", category="Two") + _, tags = await self.get("tags", limit=1) + _, categories = await self.get("categories", limit=1) + for path, params in ( + ("tags", {"cursor": "tampered"}), + ("tags", {"cursor": tags["next_cursor"][:-2] + "AA"}), + ("tags", {"cursor": categories["next_cursor"]}), + ("categories", {"cursor": tags["next_cursor"]}), + ("categories", {"cursor": "tampered"}), + ("tags", {"limit": "0"}), + ("tags", {"limit": "201"}), + ("tags", {"limit": "x"}), + ("categories", {"limit": "0"}), + ("categories", {"limit": "201"}), + ("categories", {"limit": "x"}), + ): + response, result = await self.get(path, **params) + self.assertEqual(response.status, 400, (path, params)) + self.assertEqual(result["success"], False) + self.assertIsInstance(result["error"], str) + self.assertNotIn(str(self.comfy_dir), json.dumps(result)) + rotate_history_epoch(self.api.db.model.get_connection()) + for path, cursor in ( + ("tags", tags["next_cursor"]), + ("categories", categories["next_cursor"]), + ): + response, result = await self.get(path, limit=1, cursor=cursor) + self.assertEqual(response.status, 400, path) + self.assertIn("refresh", result["error"]) + _, fresh = await self.get("tags", limit=1) + self.assertNotEqual(fresh["library_epoch"], tags["library_epoch"]) + response, _ = await self.get("tags", limit=1, cursor=fresh["next_cursor"]) + self.assertEqual(response.status, 200) + + async def test_read_failures_keep_the_read_envelope(self): + with patch.object( + self.api.db, + "_vocabulary_tags", + create=True, + side_effect=RuntimeError("secret path"), + ): + response, result = await self.get("tags") + self.assertEqual(response.status, 500) + self.assertEqual(result["success"], False) + self.assertNotIn("secret", result["error"]) + + async def test_vocabulary_tags_query_budget(self): + for index in range(150): + self._save_prompt(f"Prompt {index}", tags=[f"tag-{index}", "shared"]) + self.api.db.save_prompt_pair("Pair", "Noise", tags=["pair"]) + statements = [] + original = self.api.db.model.get_connection + + def trace_connection(): + connection = original() + connection.set_trace_callback(statements.append) + return connection + + for path in ("tags", "categories"): + statements.clear() + with patch.object( + self.api.db.model, "get_connection", side_effect=trace_connection + ): + response, result = await self.get(path, limit=200) + self.assertEqual(response.status, 200) + selects = [ + query + for query in statements + if query.lstrip().upper().startswith(("SELECT", "WITH")) + ] + self.assertLessEqual(len(selects), 4, selects) + self.assertEqual( + [query for query in selects if "storage_epoch" in query], + selects[:1], + ) + _, tags = await self.get("tags", limit=200) + self.assertEqual(len(tags["tags"]), 152) + counts = {row["name"]: row["prompt_count"] for row in tags["tags"]} + self.assertEqual(counts["shared"], 150) + self.assertEqual(counts["pair"], 1) + + async def send(self, method, path, body): + response = await self.client.request( + method, "/prompt_manager/vocabulary/" + path, json=body + ) + return response, await response.json() + + def tag_rows(self): + with self.api.db.model.get_connection() as conn: + return [ + conn.execute("SELECT * FROM tags ORDER BY id").fetchall(), + conn.execute("SELECT * FROM prompt_tags ORDER BY 1,2").fetchall(), + conn.execute("SELECT * FROM prompt_pair_tags ORDER BY 1,2").fetchall(), + ] + + async def test_tag_command_routes_return_receipts(self): + epoch = self.api.db.get_execution_epoch() + parent = self._save_prompt("Parent", tags=["moody"]) + pair = self.api.db.save_prompt_pair("Pair", "Noise", tags=["moody"]) + response, created = await self.send( + "POST", "tags", {"expected_epoch": epoch, "name": " sunset "} + ) + self.assertEqual(response.status, 200, created) + self.assertEqual( + created, + { + "success": True, + "library_epoch": epoch, + "created": True, + "tag": { + "id": created["tag"]["id"], + "name": "sunset", + "prompt_count": 0, + }, + }, + ) + _, page = await self.get("tags") + self.assertIn("sunset", [row["name"] for row in page["tags"]]) + moody = next(row for row in page["tags"] if row["name"] == "moody") + response, impact = await self.get(f"tags/{moody['id']}/impact") + self.assertEqual(response.status, 200, impact) + self.assertEqual( + impact, + { + "success": True, + "library_epoch": epoch, + "impact": { + "positive_prompts": 1, + "other_prompts": 0, + "negative_variants": 1, + "archived_negative_variants": 0, + }, + }, + ) + response, renamed = await self.send( + "PATCH", + f"tags/{moody['id']}", + {"expected_epoch": epoch, "name": "moody", "new_name": "sunset"}, + ) + self.assertEqual(response.status, 200, renamed) + self.assertEqual( + renamed, + { + "success": True, + "library_epoch": epoch, + "merged": True, + "tag": { + "id": created["tag"]["id"], + "name": "sunset", + "prompt_count": 2, + }, + "components_changed": 1, + "pairs_changed": 1, + }, + ) + self.assertEqual(self.api.db.get_prompt_by_id(parent)["tags"], ["sunset"]) + self.assertEqual(self.api.db.get_prompt_pair(pair)["tags"], ["sunset"]) + response, deleted = await self.send( + "POST", + f"tags/{created['tag']['id']}/delete", + {"expected_epoch": epoch, "name": "sunset"}, + ) + self.assertEqual(response.status, 200, deleted) + self.assertEqual( + deleted, + { + "success": True, + "library_epoch": epoch, + "deleted": True, + "components_changed": 1, + "pairs_changed": 1, + }, + ) + self.assertEqual(self.api.db.get_all_tags(), []) + response, missing = await self.get(f"tags/{created['tag']['id']}/impact") + self.assertEqual(response.status, 404) + self.assertEqual(missing, {"success": False, "error": "Tag not found"}) + for bad in ("0", "x", str(2**63)): + response, result = await self.get(f"tags/{bad}/impact") + self.assertEqual(response.status, 400, bad) + self.assertNotIn("code", result) + + async def test_create_route_validates_names(self): + epoch = self.api.db.get_execution_epoch() + before = self.tag_rows() + bodies = [ + {"expected_epoch": epoch, "name": name} + for name in ("", " ", "x" * 51, "bad\x01", "a,b", 7, None, ["a"]) + ] + [ + {"expected_epoch": epoch}, + {"name": "ok"}, + {"expected_epoch": epoch, "name": "ok", "extra": 1}, + {"expected_epoch": "not-a-uuid", "name": "ok"}, + ["ok"], + ] + for body in bodies: + response, result = await self.send("POST", "tags", body) + self.assertEqual(response.status, 400, body) + self.assertEqual(result["success"], False) + self.assertIsNone(result["code"]) + self.assertIsInstance(result["error"], str) + self.assertEqual(self.tag_rows(), before) + tag_id = self.api.db.create_vocabulary_tag("ok", expected_epoch=epoch)["tag"][ + "id" + ] + before = self.tag_rows() + for path, method, body in ( + (f"tags/{tag_id}", "PATCH", {"expected_epoch": epoch, "name": "ok"}), + ( + f"tags/{tag_id}", + "PATCH", + {"expected_epoch": epoch, "name": "ok", "new_name": "a,b"}, + ), + ( + f"tags/{tag_id}", + "PATCH", + {"expected_epoch": epoch, "name": "ok", "new_name": "ok"}, + ), + ( + "tags/x", + "PATCH", + {"expected_epoch": epoch, "name": "ok", "new_name": "new"}, + ), + (f"tags/{tag_id}/delete", "POST", {"expected_epoch": epoch}), + ( + f"tags/{tag_id}/delete", + "POST", + {"expected_epoch": epoch, "name": "ok", "new_name": "x"}, + ), + ): + response, result = await self.send(method, path, body) + self.assertEqual(response.status, 400, (path, body)) + self.assertIsNone(result["code"]) + self.assertEqual(self.tag_rows(), before) + + async def test_command_conflicts_use_command_envelope(self): + epoch = self.api.db.get_execution_epoch() + self._save_prompt("Owner", tags=["moody"]) + tag_id = self.tag_rows()[0][0]["id"] + before = self.tag_rows() + conflict = "This tag or category changed since it was loaded; review it first" + for path, method, body in ( + ( + f"tags/{tag_id}", + "PATCH", + {"expected_epoch": epoch, "name": "Moody", "new_name": "x"}, + ), + ( + f"tags/{tag_id + 9}", + "PATCH", + {"expected_epoch": epoch, "name": "moody", "new_name": "x"}, + ), + (f"tags/{tag_id}/delete", "POST", {"expected_epoch": epoch, "name": "x"}), + ( + f"tags/{tag_id + 9}/delete", + "POST", + {"expected_epoch": epoch, "name": "moody"}, + ), + ): + response, result = await self.send(method, path, body) + self.assertEqual(response.status, 409, (path, body)) + self.assertEqual( + result, + { + "success": False, + "code": "vocabulary-conflict", + "error": conflict, + "library_epoch": epoch, + }, + ) + stale = "00000000-0000-4000-8000-000000000000" + for path, method, body in ( + ("tags", "POST", {"expected_epoch": stale, "name": "fresh"}), + ( + f"tags/{tag_id}", + "PATCH", + {"expected_epoch": stale, "name": "moody", "new_name": "x"}, + ), + ( + f"tags/{tag_id}/delete", + "POST", + {"expected_epoch": stale, "name": "moody"}, + ), + ): + response, result = await self.send(method, path, body) + self.assertEqual(response.status, 409, (path, body)) + self.assertEqual(result["code"], "storage-changed") + self.assertEqual(result["library_epoch"], epoch) + self.assertEqual(self.tag_rows(), before) + with patch.object( + self.api.db, + "create_vocabulary_tag", + create=True, + side_effect=RuntimeError("secret path"), + ): + response, result = await self.send( + "POST", "tags", {"expected_epoch": epoch, "name": "fresh"} + ) + self.assertEqual(response.status, 500) + self.assertIsNone(result["code"]) + self.assertNotIn("secret", result["error"]) + + def category_rows(self): + with self.api.db.model.get_connection() as conn: + return [ + conn.execute("SELECT * FROM prompts ORDER BY id").fetchall(), + conn.execute("SELECT * FROM prompt_pairs ORDER BY id").fetchall(), + ] + + async def test_category_routes_preview_then_apply(self): + epoch = self.api.db.get_execution_epoch() + parent = self._save_prompt("Parent", category="Scenery") + pair = self.api.db.save_prompt_pair("Pair", "Noise", category="Scenery") + body = {"expected_epoch": epoch, "source": "Scenery", "target": " Land "} + response, preview = await self.send("POST", "categories/preview", body) + self.assertEqual(response.status, 200, preview) + impact = { + "positive_prompts": 1, + "other_prompts": 0, + "negative_variants": 1, + "archived_negative_variants": 0, + } + self.assertEqual( + preview, + { + "success": True, + "library_epoch": epoch, + "source": "Scenery", + "target": "Land", + "target_exists": False, + "impact": impact, + "impact_token": preview["impact_token"], + }, + ) + response, applied = await self.send( + "POST", + "categories/apply", + {**body, "impact_token": preview["impact_token"]}, + ) + self.assertEqual(response.status, 200, applied) + self.assertEqual( + applied, + { + "success": True, + "library_epoch": epoch, + "source": "Scenery", + "target": "Land", + "consolidated": False, + "impact": impact, + }, + ) + self.assertEqual(self.api.db.get_prompt_by_id(parent)["category"], "Land") + self.assertEqual(self.api.db.get_prompt_pair(pair)["category"], "Land") + response, replay = await self.send( + "POST", + "categories/apply", + {**body, "impact_token": preview["impact_token"]}, + ) + self.assertEqual(response.status, 409) + self.assertEqual(replay["code"], "vocabulary-conflict") + + async def test_preview_validation(self): + epoch = self.api.db.get_execution_epoch() + self._save_prompt("Owner", category="Scenery") + conflict = "This tag or category changed since it was loaded; review it first" + _, preview = await self.send( + "POST", + "categories/preview", + {"expected_epoch": epoch, "source": "Scenery", "target": "Land"}, + ) + token = preview["impact_token"] + before = self.category_rows() + response, result = await self.send( + "POST", + "categories/preview", + {"expected_epoch": epoch, "source": "Missing", "target": "Land"}, + ) + self.assertEqual(response.status, 409) + self.assertEqual( + result, + { + "success": False, + "code": "vocabulary-conflict", + "error": conflict, + "library_epoch": epoch, + }, + ) + bad_previews = [ + {"expected_epoch": epoch, "source": "Scenery", "target": target} + for target in ("", " ", "x" * 101, "bad\x01", "Scenery", 7, None) + ] + [ + {"expected_epoch": epoch, "source": 7, "target": "Land"}, + {"expected_epoch": epoch, "source": "Scenery"}, + {"expected_epoch": epoch, "source": "Scenery", "target": "L", "x": 1}, + {"expected_epoch": "nope", "source": "Scenery", "target": "Land"}, + ["Scenery"], + ] + for body in bad_previews: + response, result = await self.send("POST", "categories/preview", body) + self.assertEqual(response.status, 400, body) + self.assertEqual(result["success"], False) + self.assertIsNone(result["code"]) + apply_body = {"expected_epoch": epoch, "source": "Scenery", "target": "Land"} + for body in ( + apply_body, + {**apply_body, "impact_token": token, "extra": 1}, + {**apply_body, "impact_token": "A" * 64}, + {**apply_body, "impact_token": 7}, + {**apply_body, "target": "Scenery", "impact_token": token}, + ): + response, result = await self.send("POST", "categories/apply", body) + self.assertEqual(response.status, 400, body) + self.assertIsNone(result["code"]) + stale = "00000000-0000-4000-8000-000000000000" + for path, extra in ( + ("categories/preview", {}), + ("categories/apply", {"impact_token": token}), + ): + response, result = await self.send( + "POST", path, {**apply_body, "expected_epoch": stale, **extra} + ) + self.assertEqual(response.status, 409, path) + self.assertEqual(result["code"], "storage-changed") + self.assertEqual(result["library_epoch"], epoch) + self.assertEqual(self.category_rows(), before) + + async def test_blank_category_source_is_rejected_without_writing(self): + epoch = self.api.db.get_execution_epoch() + self._save_prompt("Owner", category="Scenery") + before = self.category_rows() + token = "0" * 64 + for source in ("", " ", "\t\n"): + for path, extra in ( + ("categories/preview", {}), + ("categories/apply", {"impact_token": token}), + ): + body = { + "expected_epoch": epoch, + "source": source, + "target": "Land", + **extra, + } + response, result = await self.send("POST", path, body) + self.assertEqual(response.status, 400, (path, source)) + self.assertFalse(result["success"]) + self.assertIsNone(result["code"]) + self.assertEqual(self.category_rows(), before) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_database.py b/tests/test_database.py index c16f9d5..e7b57a0 100644 --- a/tests/test_database.py +++ b/tests/test_database.py @@ -181,6 +181,46 @@ class TestTagSuggestions(DatabaseTestCase): [("asian", 3), ("portrait", 3), ("dog", 2), ("smile", 1), ("snow_day", 1)], ) + def _add_unused(self, *names): + with self.db.model.get_connection() as conn: + for name in names: + conn.execute("INSERT INTO tags(name) VALUES(?)", (name,)) + + def test_no_context_appends_unused_names_at_zero_after_used(self): + self._add_unused("zz_unused", "asian_unused") + used = [ + ("asian", 3), + ("portrait", 3), + ("dog", 2), + ("smile", 1), + ("snow_day", 1), + ] + self.assertEqual( + self._names("", []), + used + [("asian_unused", 0), ("zz_unused", 0)], + ) + self.assertEqual(self._names("asian", []), [("asian", 3), ("asian_unused", 0)]) + + def test_context_suggestions_never_include_unused_names(self): + self._add_unused("zz_unused", "asian_unused") + self.assertEqual( + self._names("", ["asian"]), [("portrait", 2), ("dog", 1), ("smile", 1)] + ) + + def test_limit_counts_unused_rows_last(self): + self._add_unused("zz_unused", "asian_unused") + self.assertEqual( + self._names("", [], limit=6), + [ + ("asian", 3), + ("portrait", 3), + ("dog", 2), + ("smile", 1), + ("snow_day", 1), + ("asian_unused", 0), + ], + ) + def test_prefix_is_case_insensitive_and_anchored_at_the_start(self): self.assertEqual(self._names("P", []), [("portrait", 3)]) self.assertEqual(self._names("s", []), [("smile", 1), ("snow_day", 1)]) @@ -203,6 +243,24 @@ class TestTagSuggestions(DatabaseTestCase): self.assertEqual(self._names("snow_", []), [("snow_day", 1)]) self.assertEqual(self._names("snow%", []), []) + def test_like_wildcards_never_match_unused_names(self): + self._add_unused("snowXday", "snowfall", "plain") + self.assertEqual(self._names("_", []), []) + self.assertEqual(self._names("%", []), []) + self.assertEqual(self._names("snow_", []), [("snow_day", 1)]) + self.assertEqual(self._names("snow%", []), []) + self.assertEqual( + self._names("snow", []), + [("snow_day", 1), ("snowfall", 0), ("snowXday", 0)], + ) + + def test_context_limit_ignores_unused_rows(self): + self._add_unused("aaa_unused", "portrait_unused") + self.assertEqual( + self._names("", ["asian"], limit=2), [("portrait", 2), ("dog", 1)] + ) + self.assertEqual(self._names("p", ["asian"], limit=1), [("portrait", 2)]) + class TestTagJunctionTables(DatabaseTestCase): """Test normalized tag storage via junction tables.""" @@ -218,6 +276,11 @@ class TestTagJunctionTables(DatabaseTestCase): all_tags = self.db.get_all_tags() self.assertEqual(sorted(all_tags), ["a", "b", "c"]) + def test_get_all_tags_includes_unused_registry_rows(self): + pid = self._save("P1", tags=["a", "b"]) + self.db.set_prompt_tags(pid, ["b"]) + self.assertEqual(self.db.get_all_tags(), ["a", "b"]) + def test_get_tags_with_counts(self): self._save("P1", tags=["common", "rare"]) self._save("P2", tags=["common"]) diff --git a/tests/test_database_backup.py b/tests/test_database_backup.py index 9df8ff9..68c68f5 100644 --- a/tests/test_database_backup.py +++ b/tests/test_database_backup.py @@ -18,8 +18,10 @@ from unittest.mock import patch sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) +from database.library_commands import LibraryCommandError from database.models import PromptModel from database.operations import PromptDatabase +from tests.open_handles import assert_closed from utils.hashing import generate_prompt_hash @@ -442,6 +444,47 @@ if __name__ == "__main__": unittest.main() +class TestRetainedTags(BackupTestCase): + """Unused registry tags are user data: they outlive restore and cleanup.""" + + def _tag_names(self, db): + return sorted(db.get_all_tags()) + + def test_unused_tag_survives_backup_restore_and_restore_rotates_epoch(self): + old_epoch = self.db.get_execution_epoch() + self.db.create_vocabulary_tag("keeper", expected_epoch=old_epoch) + backup = os.path.join(self.tmpdir, "backup.db") + self.assertTrue(self.db.model.backup_database(backup)) + self.db.close_all() + target_path = os.path.join(self.tmpdir, "target.db") + target = PromptDatabase(target_path) + try: + target.model.restore_from_file(backup) + self.assertIn("keeper", self._tag_names(target)) + new_epoch = target.get_execution_epoch() + self.assertNotEqual(new_epoch, old_epoch) + with self.assertRaises(LibraryCommandError) as caught: + target.create_vocabulary_tag("other", expected_epoch=old_epoch) + self.assertEqual(caught.exception.code, "storage-changed") + self.assertNotIn("other", self._tag_names(target)) + finally: + target.close_all() + assert_closed(self, self.tmpdir) + + def test_unused_tag_survives_startup_maintenance(self): + epoch = self.db.get_execution_epoch() + self.db.create_vocabulary_tag("keeper", expected_epoch=epoch) + owner = self.db.save_prompt( + text="owner", prompt_hash=generate_prompt_hash("owner"), tags=["used"] + ) + self.assertTrue(self.db.delete_prompt(owner)) + self.db.cleanup_duplicates() + self.db.prune_orphaned_prompts() + self.assertEqual(self._tag_names(self.db), ["keeper", "used"]) + self.db.close_all() + assert_closed(self, self.tmpdir) + + class TestConnectionsAcrossInstances(BackupTestCase): """Several PromptModel instances share one file (nodes, API, monitor).""" diff --git a/tests/test_pair_tag_maintenance.py b/tests/test_pair_tag_maintenance.py index c56e968..4a65691 100644 --- a/tests/test_pair_tag_maintenance.py +++ b/tests/test_pair_tag_maintenance.py @@ -102,6 +102,40 @@ class TestPairTagMaintenance(unittest.TestCase): self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["source"]) self.assertNotIn("target", self.db.get_all_tags()) + def test_pair_junction_failure_rolls_back_native_rename(self): + pair = self.pair("portrait", ["source", "target"]) + other = self.pair("landscape", ["source"]) + prompt = self.db.save_prompt("standalone", tags=["source"]) + conn = self.db.model.get_connection() + before = [ + conn.execute(f"SELECT * FROM {table} ORDER BY 1,2").fetchall() + for table in ("tags", "prompt_tags", "prompt_pair_tags") + ] + source = conn.execute("SELECT id FROM tags WHERE name='source'").fetchone()[0] + conn.execute( + "CREATE TRIGGER reject_pair_tag_change BEFORE UPDATE " + "ON prompt_pair_tags BEGIN SELECT RAISE(ABORT,'pair tag failure'); END" + ) + conn.commit() + rename = getattr(self.db, "rename_vocabulary_tag", None) + self.assertTrue(callable(rename), "rename_vocabulary_tag is missing") + with self.assertRaises(sqlite3.IntegrityError): + rename( + source, + expected_epoch=self.db.get_execution_epoch(), + name="source", + new_name="target", + ) + self.assertFalse(conn.in_transaction) + after = [ + conn.execute(f"SELECT * FROM {table} ORDER BY 1,2").fetchall() + for table in ("tags", "prompt_tags", "prompt_pair_tags") + ] + self.assertEqual(after, before) + self.assertEqual(self.tags(pair), ["source", "target"]) + self.assertEqual(self.tags(other), ["source"]) + self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["source"]) + if __name__ == "__main__": unittest.main() diff --git a/tests/test_vocabulary.py b/tests/test_vocabulary.py new file mode 100644 index 0000000..8ab623c --- /dev/null +++ b/tests/test_vocabulary.py @@ -0,0 +1,828 @@ +"""Bounded vocabulary reads count exactly what Library Browse shows.""" + +import re +import sqlite3 +import tempfile +import unittest +from pathlib import Path +from unittest.mock import patch + +from database.library_commands import LibraryCommandError +from database.operations import PromptDatabase +from utils.hashing import generate_prompt_hash + + +class TestVocabularyReads(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.db = PromptDatabase(str(Path(self.temp.name) / "vocabulary.db")) + + def tearDown(self): + try: + # Windows refuses to remove a database that still has open handles. + self.db.close_all() + finally: + self.temp.cleanup() + + def save(self, text, **metadata): + return self.db.save_prompt( + text, prompt_hash=generate_prompt_hash(text), **metadata + ) + + def read(self, name, **options): + method = getattr(self.db, name, None) + self.assertTrue(callable(method), f"PromptDatabase.{name} is missing") + return method(**options) + + def add_unused(self, *names): + with self.db.model.get_connection() as conn: + for name in names: + conn.execute("INSERT INTO tags(name) VALUES(?)", (name,)) + + def all_tags(self): + rows, after = [], None + while True: + page = self.read("list_vocabulary_tags", limit=200, after=after) + rows.extend(page["tags"]) + if not page["has_more"]: + return rows + after = page["next_key"] + + def test_tag_rows_are_registry_rows_in_creation_order_with_unused_at_zero(self): + self.save("First", tags=["c"]) + self.save("Second", tags=["a"]) + self.save("Third", tags=["b", "a"]) + self.add_unused("z") + page = self.read("list_vocabulary_tags") + self.assertEqual([row["name"] for row in page["tags"]], ["c", "a", "b", "z"]) + self.assertEqual(page["total"], 4) + self.assertFalse(page["has_more"]) + self.assertIsNone(page["next_key"]) + by_name = {row["name"]: row for row in page["tags"]} + self.assertEqual(by_name["z"]["prompt_count"], 0) + self.assertEqual(by_name["a"]["prompt_count"], 2) + self.assertEqual(set(by_name["c"]), {"id", "name", "prompt_count"}) + + def test_prompt_count_equals_library_browse_total_for_every_tag(self): + self.save("Parent only", tags=["parent-only"]) + self.db.save_prompt_pair("Pair cat", "Pair noise", tags=["pair-only"]) + archived = self.db.save_prompt_pair( + "Archived cat", "Archived noise", tags=["archived-only"] + ) + self.db.archive_prompt_pair(archived) + self.save("Legacy blur", tags=["negative-only"]) + self.db.save_prompt_pair("Negative owner", "Legacy blur") + self.save("Dual dog", tags=["dual"]) + self.db.save_prompt_pair("Dual dog", "Dual noise", tags=["dual"]) + retained = self.save("Retained negative", tags=["retained"]) + self.db.save_prompt_pair("Retained owner", "Retained negative") + with self.db.model.get_connection() as conn: + conn.execute( + "INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)", + (retained, "Kept"), + ) + self.add_unused("unused") + rows = self.all_tags() + self.assertEqual(len(rows), 7) + for row in rows: + self.assertEqual( + row["prompt_count"], + self.db.list_library(tags=[row["name"]])["total"], + row["name"], + ) + counts = {row["name"]: row["prompt_count"] for row in rows} + self.assertEqual( + counts, + { + "parent-only": 1, + "pair-only": 1, + "archived-only": 0, + "negative-only": 0, + "dual": 1, + "retained": 1, + "unused": 0, + }, + ) + + def test_category_rows_match_library_options_and_browse_totals(self): + self.save("Zeta first", category="Zeta") + self.db.save_prompt_pair("Alpha owner", "Alpha noise", category="Alpha") + self.save("Padded", category=" Landscape ") + self.save("Plain", category="Landscape") + self.save("Blank", category=" ") + self.save("Null") + self.save("Zeta again", category="Zeta") + archived = self.db.save_prompt_pair( + "Archived owner", "Archived noise", category="Archived-only" + ) + self.db.archive_prompt_pair(archived) + self.save("Legacy negative", category="Negative-only") + self.db.save_prompt_pair("Negative owner", "Legacy negative") + page = self.read("list_vocabulary_categories") + values = [row["value"] for row in page["categories"]] + self.assertEqual(values, ["Zeta", "Alpha", " Landscape ", "Landscape"]) + self.assertEqual(set(values), set(self.db.get_library_options()["categories"])) + self.assertEqual(page["total"], 4) + self.assertFalse(page["has_more"]) + for row in page["categories"]: + self.assertEqual(set(row), {"value", "prompt_count"}) + self.assertEqual( + row["prompt_count"], + self.db.list_library(category=row["value"])["total"], + row["value"], + ) + self.assertEqual(page["categories"][0]["prompt_count"], 2) + + def test_category_pages_continue_by_first_use_then_value(self): + for index in range(5): + self.save(f"Owner {index}", category=f"C{4 - index}") + first = self.read("list_vocabulary_categories", limit=2) + second = self.read( + "list_vocabulary_categories", limit=2, after=first["next_key"] + ) + third = self.read( + "list_vocabulary_categories", limit=2, after=second["next_key"] + ) + values = [ + row["value"] + for page in (first, second, third) + for row in page["categories"] + ] + self.assertEqual(values, ["C4", "C3", "C2", "C1", "C0"]) + self.assertEqual( + [page["has_more"] for page in (first, second, third)], [True, True, False] + ) + for after in ([1], ["1", "C4"], {"id": 1}, [True, "C4"]): + with self.assertRaises(ValueError): + self.read("list_vocabulary_categories", after=after) + + def test_category_paging_breaks_first_use_ties_by_value(self): + self.save("Shared owner", category="B") + self.db.save_prompt_pair( + "Shared owner", "Pair noise A", category="A", title="A pair" + ) + self.db.save_prompt_pair( + "Shared owner", "Pair noise C", category="C", title="C pair" + ) + pages, after = [], None + for _ in range(3): + page = self.read("list_vocabulary_categories", limit=1, after=after) + pages.append(page) + after = page["next_key"] + values = [row["value"] for page in pages for row in page["categories"]] + self.assertEqual(values, ["A", "B", "C"]) + self.assertEqual([page["has_more"] for page in pages], [True, True, False]) + for page in pages: + row = page["categories"][0] + self.assertEqual( + row["prompt_count"], + self.db.list_library(category=row["value"])["total"], + ) + + def test_pages_are_bounded_and_continue_without_gaps(self): + self.add_unused(*[f"tag-{index:03d}" for index in range(250)]) + pages, after = [], None + for _ in range(3): + page = self.read("list_vocabulary_tags", limit=100, after=after) + pages.append(page) + after = page["next_key"] + self.assertEqual([len(page["tags"]) for page in pages], [100, 100, 50]) + self.assertEqual([page["has_more"] for page in pages], [True, True, False]) + self.assertEqual({page["total"] for page in pages}, {250}) + names = [row["name"] for page in pages for row in page["tags"]] + self.assertEqual(names, [f"tag-{index:03d}" for index in range(250)]) + for options in ({"limit": 0}, {"limit": 201}, {"after": "1"}, {"after": 0}): + with self.assertRaises(ValueError): + self.read("list_vocabulary_tags", **options) + + def test_reads_share_one_snapshot_and_report_epoch(self): + self.save("Cat", category="Animal", tags=["cat"]) + epoch = self.db.get_execution_epoch() + tags = self.read("list_vocabulary_tags") + categories = self.read("list_vocabulary_categories") + self.assertEqual(tags["library_epoch"], epoch) + self.assertEqual(categories["library_epoch"], epoch) + helper = getattr(self.db, "_vocabulary_tags", None) + self.assertTrue(callable(helper), "PromptDatabase._vocabulary_tags is missing") + with self.db.library_snapshot() as (conn, snapshot_epoch, _ceiling): + inner = helper(conn, snapshot_epoch, limit=10, after=None) + self.assertEqual(inner["library_epoch"], epoch) + self.assertTrue(conn.in_transaction) + + +_TAG_TABLES = ('"tags"', '"prompt_tags"', '"prompt_pair_tags"', '"sqlite_sequence"') + + +class TestVocabularyTagCommands(unittest.TestCase): + """Guarded tag create, rename, merge and delete across both junctions.""" + + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.db = PromptDatabase(str(Path(self.temp.name) / "commands.db")) + self.epoch = self.db.get_execution_epoch() + + def tearDown(self): + try: + self.db.close_all() + finally: + self.temp.cleanup() + + def command(self, method_name, /, *args, **options): + method = getattr(self.db, method_name, None) + self.assertTrue(callable(method), f"PromptDatabase.{method_name} is missing") + return method(*args, **options) + + def save(self, text, **metadata): + return self.db.save_prompt( + text, prompt_hash=generate_prompt_hash(text), **metadata + ) + + def tag_id(self, name): + with self.db.model.get_connection() as conn: + row = conn.execute("SELECT id FROM tags WHERE name=?", (name,)).fetchone() + return row[0] if row else None + + def links(self, tag_name): + with self.db.model.get_connection() as conn: + components = conn.execute( + "SELECT pt.prompt_id FROM prompt_tags pt JOIN tags t " + "ON t.id=pt.tag_id WHERE t.name=? ORDER BY 1", + (tag_name,), + ).fetchall() + pairs = conn.execute( + "SELECT ppt.pair_id FROM prompt_pair_tags ppt JOIN tags t " + "ON t.id=ppt.tag_id WHERE t.name=? ORDER BY 1", + (tag_name,), + ).fetchall() + return [row[0] for row in components], [row[0] for row in pairs] + + def junction_rows(self): + with self.db.model.get_connection() as conn: + return ( + conn.execute("SELECT * FROM tags ORDER BY id").fetchall(), + conn.execute("SELECT * FROM prompt_tags ORDER BY 1,2").fetchall(), + conn.execute("SELECT * FROM prompt_pair_tags ORDER BY 1,2").fetchall(), + ) + + def other_rows(self): + """Every row outside the tag registry and junctions, timestamps included.""" + conn = self.db.model.get_connection() + return [ + line + for line in conn.iterdump() + if line.startswith("INSERT INTO") + and not line.startswith(tuple(f"INSERT INTO {t}" for t in _TAG_TABLES)) + ] + + def changes(self): + return self.db.model.get_connection().total_changes + + def owners(self, tag="moody"): + """Parent, active pair, archived pair and negative-only component.""" + parent = self.save("Parent sky", category="Scenery", tags=[tag]) + self.db.link_image_to_prompt(parent, str(Path(self.temp.name) / "a.png")) + active = self.db.save_prompt_pair("Active sky", "Active noise", tags=[tag]) + archived = self.db.save_prompt_pair( + "Archived sky", "Archived noise", tags=[tag] + ) + self.db.archive_prompt_pair(archived) + negative = self.save("Legacy blur", tags=[tag]) + self.db.save_prompt_pair("Negative owner", "Legacy blur") + return parent, negative, active, archived + + def test_create_inserts_unused_tag_and_reuses_exact_existing(self): + self.save("Owner", tags=["other"]) + before = self.junction_rows() + changes = self.changes() + created = self.command( + "create_vocabulary_tag", " sunset ", expected_epoch=self.epoch + ) + self.assertEqual(self.changes() - changes, 1) + self.assertEqual(created["library_epoch"], self.epoch) + self.assertTrue(created["created"]) + self.assertEqual( + created["tag"], + {"id": self.tag_id("sunset"), "name": "sunset", "prompt_count": 0}, + ) + after = self.junction_rows() + self.assertEqual(after[1:], before[1:]) + changes = self.changes() + again = self.command( + "create_vocabulary_tag", "sunset", expected_epoch=self.epoch + ) + self.assertEqual(self.changes(), changes) + self.assertFalse(again["created"]) + self.assertEqual(again["tag"], created["tag"]) + other = self.command( + "create_vocabulary_tag", "Sunset", expected_epoch=self.epoch + ) + self.assertTrue(other["created"]) + self.assertNotEqual(other["tag"]["id"], created["tag"]["id"]) + self.assertIn("sunset", self.db.get_all_tags()) + + def test_create_validates_names(self): + before = self.junction_rows() + for name in ("", " ", "x" * 51, "bad\x01", "a,b", 7, None, ["a"]): + with self.assertRaises(ValueError, msg=repr(name)): + self.command("create_vocabulary_tag", name, expected_epoch=self.epoch) + with self.assertRaises(ValueError): + self.command("create_vocabulary_tag", "ok", expected_epoch="nope") + self.assertEqual(self.junction_rows(), before) + + def test_rename_to_new_name_keeps_id_and_all_links(self): + parent, negative, active, archived = self.owners() + tag_id = self.tag_id("moody") + others = self.other_rows() + result = self.command( + "rename_vocabulary_tag", + tag_id, + expected_epoch=self.epoch, + name="moody", + new_name=" dramatic ", + ) + self.assertEqual( + result, + { + "library_epoch": self.epoch, + "merged": False, + "tag": {"id": tag_id, "name": "dramatic", "prompt_count": 2}, + "components_changed": 2, + "pairs_changed": 2, + }, + ) + self.assertIsNone(self.tag_id("moody")) + self.assertEqual( + self.links("dramatic"), ([parent, negative], [active, archived]) + ) + self.assertEqual(self.db.get_prompt_pair(archived)["tags"], ["dramatic"]) + self.assertEqual(self.other_rows(), others) + + def test_rename_to_existing_merges_both_junctions_without_duplicates(self): + both = self.save("Both", tags=["old", "new"]) + only_old = self.save("Old only", tags=["old"]) + pair_both = self.db.save_prompt_pair("Pair both", "Noise", tags=["old", "new"]) + pair_old = self.db.save_prompt_pair("Pair old", "Noise", tags=["old"]) + self.db.archive_prompt_pair(pair_old) + source, target = self.tag_id("old"), self.tag_id("new") + others = self.other_rows() + result = self.command( + "rename_vocabulary_tag", + source, + expected_epoch=self.epoch, + name="old", + new_name="new", + ) + self.assertTrue(result["merged"]) + self.assertEqual(result["tag"]["id"], target) + self.assertEqual(result["tag"]["name"], "new") + self.assertEqual(result["tag"]["prompt_count"], 3) + self.assertEqual( + (result["components_changed"], result["pairs_changed"]), (2, 2) + ) + self.assertIsNone(self.tag_id("old")) + self.assertEqual(self.links("new"), ([both, only_old], [pair_both, pair_old])) + self.assertEqual(self.db.get_prompt_by_id(both)["tags"], ["new"]) + self.assertEqual(self.db.get_prompt_pair(pair_both)["tags"], ["new"]) + self.assertEqual(self.other_rows(), others) + + def test_case_only_rename_is_a_rename(self): + prompt = self.save("Kitten", tags=["Cat"]) + tag_id = self.tag_id("Cat") + result = self.command( + "rename_vocabulary_tag", + tag_id, + expected_epoch=self.epoch, + name="Cat", + new_name="cat", + ) + self.assertFalse(result["merged"]) + self.assertEqual( + result["tag"], {"id": tag_id, "name": "cat", "prompt_count": 1} + ) + self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["cat"]) + + def test_stale_identity_and_epoch_write_nothing(self): + self.owners() + tag_id = self.tag_id("moody") + before = (self.junction_rows(), self.other_rows()) + changes = self.changes() + for tag, name in ((tag_id, "Moody"), (tag_id + 99, "moody")): + for method, extra in ( + ("rename_vocabulary_tag", {"new_name": "fresh"}), + ("delete_vocabulary_tag", {}), + ): + with self.assertRaises(LibraryCommandError) as caught: + self.command( + method, tag, expected_epoch=self.epoch, name=name, **extra + ) + self.assertEqual(caught.exception.code, "vocabulary-conflict") + self.assertEqual(caught.exception.epoch, self.epoch) + stale = "00000000-0000-4000-8000-000000000000" + with patch.object(type(self.db), "_vocabulary_tag", create=True) as lookup: + for method, args, extra in ( + ("create_vocabulary_tag", ("fresh",), {}), + ( + "rename_vocabulary_tag", + (tag_id,), + {"name": "moody", "new_name": "x"}, + ), + ("delete_vocabulary_tag", (tag_id,), {"name": "moody"}), + ): + with self.assertRaises(LibraryCommandError) as caught: + self.command(method, *args, expected_epoch=stale, **extra) + self.assertEqual(caught.exception.code, "storage-changed") + self.assertEqual(caught.exception.epoch, self.epoch) + lookup.assert_not_called() + for new_name in ("moody", " moody "): + with self.assertRaises(ValueError): + self.command( + "rename_vocabulary_tag", + tag_id, + expected_epoch=self.epoch, + name="moody", + new_name=new_name, + ) + for bad_id in (0, -1, "1", True, 2**63): + with self.assertRaises(ValueError): + self.command( + "delete_vocabulary_tag", + bad_id, + expected_epoch=self.epoch, + name="moody", + ) + self.assertEqual(self.changes(), changes) + self.assertEqual((self.junction_rows(), self.other_rows()), before) + + def test_delete_removes_name_and_links_only(self): + self.owners() + self.save("Bystander", tags=["keep"]) + tag_id = self.tag_id("moody") + impact = self.command("vocabulary_tag_impact", tag_id) + self.assertEqual( + impact, + { + "library_epoch": self.epoch, + "impact": { + "positive_prompts": 1, + "other_prompts": 1, + "negative_variants": 1, + "archived_negative_variants": 1, + }, + }, + ) + self.assertIsNone(self.command("vocabulary_tag_impact", tag_id + 99)) + others = self.other_rows() + result = self.command( + "delete_vocabulary_tag", tag_id, expected_epoch=self.epoch, name="moody" + ) + self.assertEqual( + result, + { + "library_epoch": self.epoch, + "deleted": True, + "components_changed": 2, + "pairs_changed": 2, + }, + ) + self.assertIsNone(self.tag_id("moody")) + self.assertEqual(self.links("moody"), ([], [])) + self.assertEqual(self.other_rows(), others) + self.assertEqual(len(self.links("keep")[0]), 1) + with self.assertRaises(LibraryCommandError) as caught: + self.command( + "delete_vocabulary_tag", tag_id, expected_epoch=self.epoch, name="moody" + ) + self.assertEqual(caught.exception.code, "vocabulary-conflict") + + +_OLD_STAMP = "2000-01-01T00:00:00+00:00" +_DML = re.compile(r"\s*(INSERT|UPDATE|DELETE|REPLACE)\b", re.IGNORECASE) + + +class TestVocabularyCategoryCommands(unittest.TestCase): + """Category preview and token-guarded rename/consolidation.""" + + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.db = PromptDatabase(str(Path(self.temp.name) / "categories.db")) + self.epoch = self.db.get_execution_epoch() + + def tearDown(self): + try: + self.db.close_all() + finally: + self.temp.cleanup() + + def command(self, method_name, /, **options): + method = getattr(self.db, method_name, None) + self.assertTrue(callable(method), f"PromptDatabase.{method_name} is missing") + return method(**options) + + def preview(self, source="Scenery", target="Landscape"): + return self.command( + "preview_vocabulary_category", + expected_epoch=self.epoch, + source=source, + target=target, + ) + + def apply(self, token, source="Scenery", target="Landscape"): + return self.command( + "apply_vocabulary_category", + expected_epoch=self.epoch, + source=source, + target=target, + impact_token=token, + ) + + def save(self, text, **metadata): + return self.db.save_prompt( + text, prompt_hash=generate_prompt_hash(text), **metadata + ) + + def owners(self, source="Scenery"): + """2 eligible, 1 negative-only component; 2 active, 1 archived pair.""" + first = self.save("Parent one", category=source) + second = self.save("Parent two", category=source) + negative = self.save("Legacy blur", category=source) + self.db.save_prompt_pair("Negative owner", "Legacy blur") + active = [ + self.db.save_prompt_pair( + f"Active {n}", + f"Noise {n}", + category=source, + title=f"Title {n}", + tags=[f"tag-{n}"], + rating=n, + notes=f"Notes {n}", + ) + for n in (1, 2) + ] + archived = self.db.save_prompt_pair( + "Archived sky", "Archived noise", category=source, title="Kept" + ) + self.db.archive_prompt_pair(archived) + return [first, second], negative, active, archived + + def age_rows(self): + with self.db.model.get_connection() as conn: + conn.execute("UPDATE prompts SET updated_at=?", (_OLD_STAMP,)) + conn.execute("UPDATE prompt_pairs SET updated_at=?", (_OLD_STAMP,)) + + def rows(self, table): + with self.db.model.get_connection() as conn: + return { + row["id"]: dict(row) + for row in conn.execute(f"SELECT * FROM {table} ORDER BY id") + } + + def snapshot(self): + return self.rows("prompts"), self.rows("prompt_pairs") + + def traced(self, call): + statements = [] + original = self.db.model.get_connection + + def trace_connection(): + connection = original() + connection.set_trace_callback(statements.append) + return connection + + with patch.object( + self.db.model, "get_connection", side_effect=trace_connection + ): + result = call() + conn = original() + conn.set_trace_callback(None) + return result, statements + + def test_preview_counts_every_owner_type(self): + self.owners() + self.save("Elsewhere", category="Other") + before = self.snapshot() + result, statements = self.traced(self.preview) + self.assertEqual( + set(result), + { + "library_epoch", + "source", + "target", + "target_exists", + "impact", + "impact_token", + }, + ) + self.assertEqual(result["library_epoch"], self.epoch) + self.assertEqual((result["source"], result["target"]), ("Scenery", "Landscape")) + self.assertEqual( + result["impact"], + { + "positive_prompts": 2, + "other_prompts": 1, + "negative_variants": 2, + "archived_negative_variants": 1, + }, + ) + self.assertFalse(result["target_exists"]) + self.assertRegex(result["impact_token"], r"\A[0-9a-f]{64}\Z") + self.assertEqual([q for q in statements if _DML.match(q)], []) + self.assertEqual(self.snapshot(), before) + trimmed = self.preview(target=" Other ") + self.assertEqual(trimmed["target"], "Other") + self.assertTrue(trimmed["target_exists"]) + self.assertNotEqual(trimmed["impact_token"], result["impact_token"]) + self.assertEqual(self.preview(), result) + + def test_apply_renames_exact_value_on_components_and_all_pairs(self): + components, negative, active, archived = self.owners() + variant = self.save("Padded", category=" Scenery ") + bystander = self.save("Bystander", category="Other") + self.age_rows() + prompts_before, pairs_before = self.snapshot() + preview = self.preview() + receipt = self.apply(preview["impact_token"]) + self.assertEqual( + receipt, + { + "library_epoch": self.epoch, + "source": "Scenery", + "target": "Landscape", + "consolidated": False, + "impact": preview["impact"], + }, + ) + prompts_after, pairs_after = self.snapshot() + changed_prompts = {*components, negative} + changed_pairs = {*active, archived} + for table_before, table_after, changed in ( + (prompts_before, prompts_after, changed_prompts), + (pairs_before, pairs_after, changed_pairs), + ): + self.assertNotIn("Scenery", [r["category"] for r in table_after.values()]) + self.assertEqual(set(table_after), set(table_before)) + for row_id, row in table_after.items(): + old = table_before[row_id] + if row_id in changed: + self.assertEqual(row["category"], "Landscape") + self.assertNotEqual(row["updated_at"], _OLD_STAMP) + expected = {**old, "category": "Landscape"} + expected["updated_at"] = row["updated_at"] + self.assertEqual(row, expected) + else: + self.assertEqual(row, old) + self.assertEqual(prompts_after[variant]["category"], " Scenery ") + self.assertEqual(prompts_after[bystander]["category"], "Other") + for pair_id in active: + pair = self.db.get_prompt_pair(pair_id) + self.assertEqual(pair["category"], "Landscape") + self.assertEqual(len(pair["tags"]), 1) + self.assertEqual(pairs_after[archived]["archived"], 1) + self.assertEqual(pairs_after[archived]["title"], "Kept") + # Consolidating into an existing value reports it. + merge = self.preview(source=" Scenery ", target="Other") + self.assertTrue(merge["target_exists"]) + receipt = self.apply(merge["impact_token"], source=" Scenery ", target="Other") + self.assertTrue(receipt["consolidated"]) + self.assertEqual( + receipt["impact"], + { + "positive_prompts": 1, + "other_prompts": 0, + "negative_variants": 0, + "archived_negative_variants": 0, + }, + ) + self.assertEqual(self.rows("prompts")[variant]["category"], "Other") + + def test_token_changes_with_membership_eligibility_archive_and_target_presence( + self, + ): + components, negative, active, _archived = self.owners() + + def add_owner(): + self.save("Late owner", category="Scenery") + + def archive_pair(): + self.db.archive_prompt_pair(active[0]) + + def promote_negative(): + with self.db.model.get_connection() as conn: + conn.execute( + "INSERT INTO prompt_library_state(prompt_id,title) VALUES(?,?)", + (negative, "Promoted"), + ) + + def create_target(): + self.save("Target owner", category="Landscape") + + for mutate in (add_owner, archive_pair, promote_negative, create_target): + with self.subTest(mutate.__name__): + token = self.preview()["impact_token"] + mutate() + before = self.snapshot() + changes = self.db.model.get_connection().total_changes + with self.assertRaises(LibraryCommandError) as caught: + self.apply(token) + self.assertEqual(caught.exception.code, "vocabulary-conflict") + self.assertEqual(caught.exception.epoch, self.epoch) + self.assertEqual(self.db.model.get_connection().total_changes, changes) + self.assertEqual(self.snapshot(), before) + token = self.preview()["impact_token"] + with self.db.model.get_connection() as conn: + conn.execute( + "UPDATE prompts SET notes='edited' WHERE id=?", (components[0],) + ) + self.assertTrue(self.apply(token)["consolidated"]) + self.assertEqual(self.rows("prompts")[components[0]]["category"], "Landscape") + + def test_category_command_guards(self): + self.owners() + before = self.snapshot() + stale = "00000000-0000-4000-8000-000000000000" + token = self.preview()["impact_token"] + with self.assertRaises(LibraryCommandError) as caught: + self.command( + "preview_vocabulary_category", + expected_epoch=stale, + source="Scenery", + target="Landscape", + ) + self.assertEqual(caught.exception.code, "storage-changed") + self.assertEqual(caught.exception.epoch, self.epoch) + with self.assertRaises(LibraryCommandError) as caught: + self.command( + "apply_vocabulary_category", + expected_epoch=stale, + source="Scenery", + target="Landscape", + impact_token=token, + ) + self.assertEqual(caught.exception.code, "storage-changed") + for source in ("Missing", "scenery", " Scenery"): + with self.assertRaises(LibraryCommandError) as caught: + self.preview(source=source) + self.assertEqual(caught.exception.code, "vocabulary-conflict") + for source, target in ( + ("Scenery", ""), + ("Scenery", " "), + ("Scenery", "x" * 101), + ("Scenery", "bad\x01"), + ("Scenery", "Scenery"), + ("Scenery", " Scenery "), + ("Scenery", None), + ("Scenery", 7), + (None, "Landscape"), + ("", "Landscape"), + (" ", "Landscape"), + ): + with self.assertRaises(ValueError, msg=(source, target)): + self.preview(source=source, target=target) + for bad in ("", "A" * 64, token[:-1], None, 7): + with self.assertRaises(ValueError): + self.apply(bad) + self.assertEqual(self.snapshot(), before) + + def test_category_apply_failure_rolls_back(self): + self.owners() + token = self.preview()["impact_token"] + with self.db.model.get_connection() as conn: + conn.execute( + "CREATE TRIGGER abort_pairs BEFORE UPDATE ON prompt_pairs " + "BEGIN SELECT RAISE(ABORT, 'boom'); END" + ) + before = self.snapshot() + with self.assertRaises(sqlite3.Error): + self.apply(token) + self.assertEqual(self.snapshot(), before) + self.assertIn("Scenery", [r["category"] for r in before[0].values()]) + self.assertFalse(self.db.model.get_connection().in_transaction) + + +class TestSharedTokenHashing(unittest.TestCase): + def test_both_tokens_keep_their_previous_values(self): + from database.library_tokens import library_edit_token + from database.vocabulary import category_impact_token + + record = { + "id": 3, + "text": "t", + "title": None, + "category": "c", + "notes": "", + "rating": 0, + "tags": ["b", "a", "a"], + } + self.assertEqual( + library_edit_token("e", record), + "32534d4c92e4e9a7d8e31c30ecdbb07563703b5bd7ed216a8d2066c0dc44332b", + ) + self.assertEqual( + category_impact_token( + "e", "S", "T", True, [(1, True), (2, False)], [(5, False)] + ), + "cefe7f6a5544b7f6a203f5527810899c6af568aa7abc1b75adffbfb5ae6251ae", + ) + + +if __name__ == "__main__": + unittest.main()