feat(vocabulary): backend for native tags and categories (#241)

* feat(tags): keep unused tags in the shared vocabulary

* feat(vocabulary): add bounded tag and category reads

* feat(vocabulary): guard tag create, rename and merge

* feat(vocabulary): preview and apply category consolidation

* refactor(vocabulary): share canonical token hashing

* test(vocabulary): prove retained tags survive restore

* test(vocabulary): cover category paging tie-break and clarify contract docs

* test(vocabulary): assert rejected tag creations write nothing
This commit is contained in:
Vito
2026-10-19 05:47:36 -07:00
committed by GitHub
parent 39dcc40114
commit ab333d3838
14 changed files with 2187 additions and 36 deletions
+35
View File
@@ -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).
+2 -1
View File
@@ -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.
"""
+13 -8
View File
@@ -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)
+27 -27
View File
@@ -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
+435
View File
@@ -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)
+3
View File
@@ -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)
+3
View File
@@ -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"
+164
View File
@@ -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})
+13
View File
@@ -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"
+524
View File
@@ -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()
+63
View File
@@ -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"])
+43
View File
@@ -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)."""
+34
View File
@@ -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()
+828
View File
@@ -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()