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:
+35
@@ -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).
|
||||
|
||||
@@ -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.
|
||||
"""
|
||||
|
||||
@@ -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
@@ -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
|
||||
|
||||
|
||||
@@ -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)
|
||||
@@ -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)
|
||||
|
||||
@@ -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"
|
||||
|
||||
|
||||
@@ -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})
|
||||
@@ -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"
|
||||
|
||||
@@ -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()
|
||||
@@ -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"])
|
||||
|
||||
@@ -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)."""
|
||||
|
||||
|
||||
@@ -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()
|
||||
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user