feat: reuse positive and negative prompt combinations (#195)
* feat(db): add pair schema and stage restore migrations for #193 * feat(database): persist reusable ordered prompt pairs (#193) * feat(api): expose reusable pairs and atomic metadata enrichment (#193) * feat(metadata): resolve canonical prompt pairs with bounded provenance (#193) * feat(scan): enrich image pairs without guessing or duplicate links (#193) * fix(database): preserve pair tags during global maintenance (#193) * feat(tracking): enrich live image pairs and cancel scans safely (#193) * feat(nodes): encode paired prompts and count cached recipe runs (#193) * fix(database): validate restored pair ownership and provenance (#193) * feat(admin): report pair scan outcomes and cancellation (#193) * feat(ui): add paired prompt editor and explicit recipe library for #193 * fix(autocomplete): include reusable pair metadata vocabulary (#193) * fix(pairs): honor recent sorting and conservative scan outcomes (#193) * fix(ui): protect paired draft saves and clarify negative variants (#193) * fix(metadata): preserve bounded paired execution evidence (#193) * fix(tracking): follow actual pair output roles and prefer execution evidence (#193) * docs(pairs): document live validation and retain renderer screenshots (#193) * test(api): inject disposable databases before API startup (#194) * test(api): canonicalize portable temporary gallery roots (#194) * docs(ui): prototype branching negative variant history (#193) * docs(ui): show selected negative beneath positive in timeline concept (#193) * docs: retain approved timeline for later gallery redesign (#193) * test(db): close restore fixture connections before cleanup (#194) * test(db): reuse existing handle guards in all pair fixtures (#194) * chore: clarify rollback test and intentional exception handling (#193)
@@ -45,6 +45,7 @@ from .prompt_manager import ( # noqa: E402
|
||||
PromptManagerFocused,
|
||||
PromptManagerSplit,
|
||||
)
|
||||
from .prompt_manager_pair import PromptManagerPair # noqa: E402
|
||||
from .prompt_manager_text import PromptManagerText # noqa: E402
|
||||
from .prompt_search_list import PromptSearchList # noqa: E402
|
||||
|
||||
@@ -52,6 +53,7 @@ NODE_CLASS_MAPPINGS = {
|
||||
"PromptManager": PromptManager,
|
||||
"PromptManagerFocused": PromptManagerFocused,
|
||||
"PromptManagerSplit": PromptManagerSplit,
|
||||
"PromptManagerPair": PromptManagerPair,
|
||||
"PromptManagerText": PromptManagerText,
|
||||
"PromptSearchList": PromptSearchList,
|
||||
}
|
||||
@@ -60,6 +62,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"PromptManager": "Prompt Manager (Classic)",
|
||||
"PromptManagerFocused": "Prompt Manager (Focused)",
|
||||
"PromptManagerSplit": "Prompt Manager (Split)",
|
||||
"PromptManagerPair": "Prompt Manager (+/-)",
|
||||
"PromptManagerText": "Prompt Manager Text",
|
||||
"PromptSearchList": "Prompt Search List",
|
||||
}
|
||||
|
||||
@@ -4,14 +4,17 @@ Database schema and models for KikoTextEncode prompt storage.
|
||||
|
||||
import contextlib
|
||||
import datetime
|
||||
import json
|
||||
import math
|
||||
import shutil
|
||||
import sqlite3
|
||||
import os
|
||||
import threading
|
||||
import tempfile
|
||||
import weakref
|
||||
from contextlib import closing
|
||||
from pathlib import Path
|
||||
from typing import Set, Tuple
|
||||
from typing import Optional, Set, Tuple
|
||||
|
||||
# Import logging system
|
||||
try:
|
||||
@@ -29,6 +32,41 @@ except ImportError:
|
||||
ISO_TIMESTAMP_FORMAT = "%Y-%m-%dT%H:%M:%f+00:00"
|
||||
|
||||
|
||||
MAX_PAIR_PROVENANCE_LENGTH = 4_000_000
|
||||
MAX_PAIR_PROVENANCE_DEPTH = 64
|
||||
|
||||
|
||||
def pair_provenance_error(encoded: object) -> Optional[str]:
|
||||
"""Validate the same bounded JSON object for live writes and restores."""
|
||||
if not isinstance(encoded, str):
|
||||
return "Pair provenance must be JSON text"
|
||||
if len(encoded) > MAX_PAIR_PROVENANCE_LENGTH:
|
||||
return "Pair provenance exceeds the size limit"
|
||||
try:
|
||||
value = json.loads(encoded)
|
||||
except (ValueError, TypeError, RecursionError):
|
||||
return "Pair provenance is not readable JSON"
|
||||
if not isinstance(value, dict):
|
||||
return "Pair provenance must be a JSON object"
|
||||
pending = [(iter((value,)), 0)]
|
||||
while pending:
|
||||
children, depth = pending[-1]
|
||||
try:
|
||||
item = next(children)
|
||||
except StopIteration:
|
||||
pending.pop()
|
||||
continue
|
||||
if depth > MAX_PAIR_PROVENANCE_DEPTH:
|
||||
return "Pair provenance exceeds the nesting limit"
|
||||
if isinstance(item, float) and not math.isfinite(item):
|
||||
return "Pair provenance contains nonfinite numbers"
|
||||
if isinstance(item, dict):
|
||||
pending.append((iter(item.values()), depth + 1))
|
||||
elif isinstance(item, list):
|
||||
pending.append((iter(item), depth + 1))
|
||||
return None
|
||||
|
||||
|
||||
def utc_now_iso() -> str:
|
||||
"""Current UTC time as ISO text with microseconds.
|
||||
|
||||
@@ -235,6 +273,224 @@ class PromptModel:
|
||||
|
||||
# Image uniqueness by full path instead of basename (3.2.4)
|
||||
self._migrate_image_uniqueness_by_path(conn)
|
||||
self._migrate_prompt_pairs(conn)
|
||||
|
||||
def _migrate_prompt_pairs(self, conn: sqlite3.Connection) -> None:
|
||||
conn.execute("SAVEPOINT prompt_pair_schema")
|
||||
try:
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS prompt_pairs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
positive_prompt_id INTEGER NOT NULL,
|
||||
negative_prompt_id INTEGER,
|
||||
title TEXT,
|
||||
category TEXT,
|
||||
rating INTEGER CHECK(rating BETWEEN 1 AND 5),
|
||||
notes TEXT,
|
||||
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
updated_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
last_used_at TEXT,
|
||||
run_count INTEGER NOT NULL DEFAULT 0,
|
||||
archived INTEGER NOT NULL DEFAULT 0 CHECK(archived IN (0, 1)),
|
||||
FOREIGN KEY(positive_prompt_id)
|
||||
REFERENCES prompts(id) ON DELETE RESTRICT,
|
||||
FOREIGN KEY(negative_prompt_id)
|
||||
REFERENCES prompts(id) ON DELETE RESTRICT,
|
||||
UNIQUE(positive_prompt_id, negative_prompt_id)
|
||||
)
|
||||
""")
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS prompt_pair_tags (
|
||||
pair_id INTEGER NOT NULL,
|
||||
tag_id INTEGER NOT NULL,
|
||||
PRIMARY KEY(pair_id, tag_id),
|
||||
FOREIGN KEY(pair_id) REFERENCES prompt_pairs(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY(tag_id) REFERENCES tags(id) ON DELETE CASCADE
|
||||
)
|
||||
""")
|
||||
conn.execute("""
|
||||
CREATE TABLE IF NOT EXISTS image_prompt_pairs (
|
||||
id INTEGER PRIMARY KEY AUTOINCREMENT,
|
||||
image_id INTEGER NOT NULL,
|
||||
pair_id INTEGER NOT NULL,
|
||||
provenance TEXT NOT NULL,
|
||||
created_at TEXT NOT NULL DEFAULT CURRENT_TIMESTAMP,
|
||||
FOREIGN KEY(image_id)
|
||||
REFERENCES generated_images(id) ON DELETE CASCADE,
|
||||
FOREIGN KEY(pair_id) REFERENCES prompt_pairs(id) ON DELETE RESTRICT,
|
||||
UNIQUE(pair_id, image_id, provenance)
|
||||
)
|
||||
""")
|
||||
self._create_pair_indexes(conn)
|
||||
error = self._pair_schema_error(conn)
|
||||
if error:
|
||||
raise sqlite3.DatabaseError(error)
|
||||
conn.execute("RELEASE prompt_pair_schema")
|
||||
except BaseException:
|
||||
conn.execute("ROLLBACK TO prompt_pair_schema")
|
||||
conn.execute("RELEASE prompt_pair_schema")
|
||||
raise
|
||||
|
||||
@staticmethod
|
||||
def _create_pair_indexes(conn: sqlite3.Connection) -> None:
|
||||
conn.execute(
|
||||
"CREATE UNIQUE INDEX IF NOT EXISTS idx_prompt_pairs_positive_only "
|
||||
"ON prompt_pairs(positive_prompt_id) WHERE negative_prompt_id IS NULL"
|
||||
)
|
||||
for name, target in (
|
||||
("idx_prompt_pairs_negative", "prompt_pairs(negative_prompt_id)"),
|
||||
("idx_prompt_pair_tags_tag", "prompt_pair_tags(tag_id)"),
|
||||
("idx_image_prompt_pairs_image", "image_prompt_pairs(image_id)"),
|
||||
):
|
||||
conn.execute(f"CREATE INDEX IF NOT EXISTS {name} ON {target}")
|
||||
|
||||
@staticmethod
|
||||
def _pair_schema_error(conn: sqlite3.Connection) -> Optional[str]:
|
||||
required = {
|
||||
"prompt_pairs": {
|
||||
"id",
|
||||
"positive_prompt_id",
|
||||
"negative_prompt_id",
|
||||
"title",
|
||||
"category",
|
||||
"rating",
|
||||
"notes",
|
||||
"created_at",
|
||||
"updated_at",
|
||||
"last_used_at",
|
||||
"run_count",
|
||||
"archived",
|
||||
},
|
||||
"prompt_pair_tags": {"pair_id", "tag_id"},
|
||||
"image_prompt_pairs": {
|
||||
"id",
|
||||
"image_id",
|
||||
"pair_id",
|
||||
"provenance",
|
||||
"created_at",
|
||||
},
|
||||
}
|
||||
present = {
|
||||
row[0]
|
||||
for row in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")
|
||||
}
|
||||
if not present.intersection(required):
|
||||
return None
|
||||
if not set(required).issubset(present):
|
||||
return "Incomplete prompt pair schema"
|
||||
foreign_keys = {
|
||||
"prompt_pairs": {
|
||||
("positive_prompt_id", "prompts", "id", "RESTRICT"),
|
||||
("negative_prompt_id", "prompts", "id", "RESTRICT"),
|
||||
},
|
||||
"prompt_pair_tags": {
|
||||
("pair_id", "prompt_pairs", "id", "CASCADE"),
|
||||
("tag_id", "tags", "id", "CASCADE"),
|
||||
},
|
||||
"image_prompt_pairs": {
|
||||
("pair_id", "prompt_pairs", "id", "RESTRICT"),
|
||||
("image_id", "generated_images", "id", "CASCADE"),
|
||||
},
|
||||
}
|
||||
unique = {
|
||||
"prompt_pairs": ("positive_prompt_id", "negative_prompt_id"),
|
||||
"prompt_pair_tags": ("pair_id", "tag_id"),
|
||||
"image_prompt_pairs": ("pair_id", "image_id", "provenance"),
|
||||
}
|
||||
non_null = {
|
||||
"prompt_pairs": {
|
||||
"positive_prompt_id",
|
||||
"created_at",
|
||||
"updated_at",
|
||||
"run_count",
|
||||
"archived",
|
||||
},
|
||||
"prompt_pair_tags": {"pair_id", "tag_id"},
|
||||
"image_prompt_pairs": {"pair_id", "image_id", "provenance", "created_at"},
|
||||
}
|
||||
integer_columns = {
|
||||
"id",
|
||||
"positive_prompt_id",
|
||||
"negative_prompt_id",
|
||||
"pair_id",
|
||||
"tag_id",
|
||||
"image_id",
|
||||
"rating",
|
||||
"run_count",
|
||||
"archived",
|
||||
}
|
||||
for table, columns in required.items():
|
||||
info = {row[1]: row for row in conn.execute(f"PRAGMA table_info({table})")}
|
||||
if not columns.issubset(info):
|
||||
return f"{table} is missing required pair columns"
|
||||
if any(not info[name][3] for name in non_null[table]):
|
||||
return f"{table} has nullable required pair columns"
|
||||
if any(
|
||||
info[name][2].upper()
|
||||
!= ("INTEGER" if name in integer_columns else "TEXT")
|
||||
for name in columns
|
||||
):
|
||||
return f"{table} has invalid pair column types"
|
||||
primary = tuple(
|
||||
name
|
||||
for name in sorted(info, key=lambda name: info[name][5])
|
||||
if info[name][5]
|
||||
)
|
||||
if primary != (
|
||||
("pair_id", "tag_id") if table == "prompt_pair_tags" else ("id",)
|
||||
):
|
||||
return f"{table} has an invalid pair primary key"
|
||||
actual_fks = {
|
||||
(row[3], row[2], row[4], row[6])
|
||||
for row in conn.execute(f"PRAGMA foreign_key_list({table})")
|
||||
}
|
||||
if actual_fks != foreign_keys[table]:
|
||||
return f"{table} has invalid pair foreign key definitions"
|
||||
if conn.execute(f"PRAGMA foreign_key_check({table})").fetchone():
|
||||
return f"{table} contains foreign key violations"
|
||||
indexes = list(conn.execute(f"PRAGMA index_list({table})"))
|
||||
index_columns = {}
|
||||
for row in indexes:
|
||||
name = row[1].replace('"', '""')
|
||||
index_columns[row[1]] = tuple(
|
||||
item[2] for item in conn.execute(f'PRAGMA index_info("{name}")')
|
||||
)
|
||||
if not any(
|
||||
row[2] and not row[4] and index_columns[row[1]] == unique[table]
|
||||
for row in indexes
|
||||
):
|
||||
return f"{table} is missing pair uniqueness constraints"
|
||||
if table == "prompt_pairs":
|
||||
null_unique = False
|
||||
for row in indexes:
|
||||
if (
|
||||
row[2]
|
||||
and row[4]
|
||||
and index_columns[row[1]] == ("positive_prompt_id",)
|
||||
):
|
||||
sql = conn.execute(
|
||||
"SELECT sql FROM sqlite_master WHERE type='index' AND name"
|
||||
"=?",
|
||||
(row[1],),
|
||||
).fetchone()[0]
|
||||
if " ".join(sql.lower().split()).endswith(
|
||||
"where negative_prompt_id is null"
|
||||
):
|
||||
null_unique = True
|
||||
if not null_unique:
|
||||
return "prompt_pairs is missing empty-negative pair uniqueness"
|
||||
if conn.execute(
|
||||
"SELECT 1 FROM image_prompt_pairs ip "
|
||||
"JOIN prompt_pairs pp ON pp.id=ip.pair_id "
|
||||
"JOIN generated_images gi ON gi.id=ip.image_id "
|
||||
"WHERE pp.positive_prompt_id IS NOT gi.prompt_id LIMIT 1"
|
||||
).fetchone():
|
||||
return "Image pair association has a different positive prompt owner"
|
||||
for row in conn.execute("SELECT provenance FROM image_prompt_pairs"):
|
||||
error = pair_provenance_error(row[0])
|
||||
if error:
|
||||
return error
|
||||
return None
|
||||
|
||||
def _create_indexes(self, conn: sqlite3.Connection) -> None:
|
||||
"""
|
||||
@@ -772,6 +1028,15 @@ class PromptModel:
|
||||
False,
|
||||
"prompts table is missing required columns: " + ", ".join(missing),
|
||||
)
|
||||
forbidden = conn.execute(
|
||||
"SELECT type FROM sqlite_master WHERE type IN ('trigger', 'view') "
|
||||
"LIMIT 1"
|
||||
).fetchone()
|
||||
if forbidden:
|
||||
return False, f"Database contains an unsupported {forbidden[0]}"
|
||||
pair_error = PromptModel._pair_schema_error(conn)
|
||||
if pair_error:
|
||||
return False, pair_error
|
||||
return True, "ok"
|
||||
except sqlite3.Error as e:
|
||||
return False, f"Invalid SQLite database: {e}"
|
||||
@@ -814,13 +1079,11 @@ class PromptModel:
|
||||
"""
|
||||
Replace the live database with the file at src_path.
|
||||
|
||||
The current database is first backed up next to itself as
|
||||
``<db_path>.backup_<YYYYmmdd_HHMMSS>`` (WAL content included), the
|
||||
pages of ``src_path`` are copied into the live database through
|
||||
SQLite's backup API (visible to every open connection) and the schema
|
||||
migrations run on the restored content.
|
||||
|
||||
Callers should run verify_database_file(src_path) first.
|
||||
A disposable SQLite backup of the source is migrated and verified first.
|
||||
Only then is the current database backed up next to itself as
|
||||
``<db_path>.backup_<YYYYmmdd_HHMMSS>`` (WAL content included), and the
|
||||
staged pages replace the live pages through SQLite's backup API.
|
||||
Migration failures never replace live content or modify the source.
|
||||
|
||||
Args:
|
||||
src_path: SQLite file to restore
|
||||
@@ -833,17 +1096,42 @@ class PromptModel:
|
||||
RuntimeError: If the previous database could not be backed up
|
||||
OSError: If the restored file could not be copied into place
|
||||
"""
|
||||
backup_path = ""
|
||||
if os.path.exists(self.db_path):
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
backup_path = f"{self.db_path}.backup_{timestamp}"
|
||||
if not self.backup_database(backup_path):
|
||||
raise RuntimeError("Could not back up the current database")
|
||||
self.logger.info(f"Current database backed up to: {backup_path}")
|
||||
|
||||
self._replace_contents_from(src_path)
|
||||
with PromptModel._init_lock:
|
||||
PromptModel._initialized_paths.discard(self._schema_key())
|
||||
self._ensure_database_exists()
|
||||
self.logger.info(f"Database restored from {src_path}")
|
||||
return backup_path
|
||||
ok, reason = self.verify_database_file(src_path)
|
||||
if not ok:
|
||||
raise ValueError(reason)
|
||||
with tempfile.TemporaryDirectory(prefix="prompt-manager-restore-") as directory:
|
||||
staged_path = str(Path(directory) / "staged.db")
|
||||
source = sqlite3.connect(
|
||||
Path(src_path).resolve().as_uri() + "?mode=ro", uri=True
|
||||
)
|
||||
try:
|
||||
with closing(sqlite3.connect(staged_path)) as destination:
|
||||
source.backup(destination)
|
||||
finally:
|
||||
source.close()
|
||||
staged = None
|
||||
try:
|
||||
staged = PromptModel(staged_path)
|
||||
ok, reason = self.verify_database_file(staged_path)
|
||||
if not ok:
|
||||
raise ValueError(reason)
|
||||
staged.close_all()
|
||||
backup_path = ""
|
||||
if os.path.exists(self.db_path):
|
||||
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
backup_path = f"{self.db_path}.backup_{timestamp}"
|
||||
if not self.backup_database(backup_path):
|
||||
raise RuntimeError("Could not back up the current database")
|
||||
self.logger.info(f"Current database backed up to: {backup_path}")
|
||||
self._replace_contents_from(staged_path)
|
||||
with PromptModel._init_lock:
|
||||
PromptModel._initialized_paths.add(self._schema_key())
|
||||
self.logger.info(f"Database restored from {src_path}")
|
||||
return backup_path
|
||||
finally:
|
||||
if staged is not None:
|
||||
staged.close_all()
|
||||
with PromptModel._init_lock:
|
||||
PromptModel._initialized_paths.discard(
|
||||
os.path.normcase(os.path.realpath(staged_path))
|
||||
)
|
||||
|
||||
@@ -9,6 +9,7 @@ import os
|
||||
from typing import Optional, List, Dict, Any
|
||||
|
||||
from .models import PromptModel, normalize_image_path, utc_now_iso
|
||||
from .pairs import PairDependencyError, PairOperations
|
||||
|
||||
# Import logging system
|
||||
try:
|
||||
@@ -146,7 +147,7 @@ def _coerce_file_info(metadata: Any) -> tuple:
|
||||
)
|
||||
|
||||
|
||||
class PromptDatabase:
|
||||
class PromptDatabase(PairOperations):
|
||||
"""Database operations class for managing prompts."""
|
||||
|
||||
def __init__(self, db_path: Optional[str] = None):
|
||||
@@ -581,7 +582,8 @@ class PromptDatabase:
|
||||
Returns:
|
||||
bool: True if deletion was successful, False otherwise
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
with self._pair_transaction() as conn:
|
||||
self._ensure_prompts_unpaired(conn, [prompt_id])
|
||||
# First delete related images to avoid foreign key constraint
|
||||
conn.execute(
|
||||
"DELETE FROM generated_images WHERE prompt_id = ?", (prompt_id,)
|
||||
@@ -597,7 +599,7 @@ class PromptDatabase:
|
||||
Returns:
|
||||
Number of prompts deleted.
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
with self._pair_transaction() as conn:
|
||||
# Get IDs first for cascade cleanup
|
||||
ids = [
|
||||
r[0]
|
||||
@@ -607,6 +609,7 @@ class PromptDatabase:
|
||||
]
|
||||
if not ids:
|
||||
return 0
|
||||
self._ensure_prompts_unpaired(conn, ids)
|
||||
placeholders = ",".join("?" * len(ids))
|
||||
conn.execute(
|
||||
f"DELETE FROM generated_images WHERE prompt_id IN ({placeholders})",
|
||||
@@ -632,7 +635,9 @@ class PromptDatabase:
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT DISTINCT TRIM(category) as category FROM prompts WHERE "
|
||||
"SELECT TRIM(category) AS category FROM prompts WHERE "
|
||||
"category IS NOT NULL AND TRIM(category) != '' UNION "
|
||||
"SELECT TRIM(category) AS category FROM prompt_pairs WHERE "
|
||||
"category IS NOT NULL AND TRIM(category) != '' ORDER BY category"
|
||||
)
|
||||
return [row["category"] for row in cursor.fetchall()]
|
||||
@@ -699,16 +704,16 @@ class PromptDatabase:
|
||||
|
||||
def get_all_tags(self) -> List[str]:
|
||||
"""
|
||||
Get all unique tags that are in use (linked to at least one prompt).
|
||||
Get all unique tags linked to a standalone prompt or reusable pair.
|
||||
|
||||
Returns:
|
||||
Sorted list of tag names
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute(
|
||||
"SELECT DISTINCT t.name AS tag "
|
||||
"FROM prompt_tags pt "
|
||||
"JOIN tags t ON pt.tag_id = t.id "
|
||||
"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"
|
||||
)
|
||||
return [row["tag"] for row in cursor.fetchall()]
|
||||
@@ -784,7 +789,8 @@ class PromptDatabase:
|
||||
|
||||
Returns tags whose name starts with ``prefix`` (case-insensitive),
|
||||
counted over the prompts that carry every tag in ``with_tags``; the
|
||||
context tags themselves are left out. Ordered by count, then name.
|
||||
context tags themselves are left out. Counts include separate standalone
|
||||
and pair references, with owner types kept distinct. Ordered by count/name.
|
||||
"""
|
||||
seen = set()
|
||||
context = []
|
||||
@@ -799,11 +805,11 @@ class PromptDatabase:
|
||||
if context:
|
||||
placeholders = ",".join("?" * len(context))
|
||||
where.append(
|
||||
"pt.prompt_id IN ("
|
||||
" SELECT pt2.prompt_id FROM prompt_tags pt2"
|
||||
"(pt.kind,pt.owner_id) IN ("
|
||||
" SELECT pt2.kind,pt2.owner_id FROM tag_refs pt2"
|
||||
" JOIN tags t2 ON t2.id = pt2.tag_id"
|
||||
f" WHERE t2.name COLLATE NOCASE IN ({placeholders})"
|
||||
" GROUP BY pt2.prompt_id"
|
||||
" GROUP BY pt2.kind,pt2.owner_id"
|
||||
" HAVING COUNT(DISTINCT lower(t2.name)) = ?)"
|
||||
)
|
||||
params.extend(context)
|
||||
@@ -813,8 +819,11 @@ class PromptDatabase:
|
||||
params.append(max(1, int(limit)))
|
||||
|
||||
sql = (
|
||||
"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 prompt_tags pt JOIN tags t ON t.id = pt.tag_id"
|
||||
" FROM tag_refs pt JOIN tags t ON t.id = pt.tag_id"
|
||||
f" WHERE {' AND '.join(where)}"
|
||||
" GROUP BY t.id"
|
||||
" ORDER BY count DESC, t.name COLLATE NOCASE ASC"
|
||||
@@ -921,6 +930,8 @@ class PromptDatabase:
|
||||
|
||||
old_name = old_name.strip()
|
||||
new_name = new_name.strip()
|
||||
if old_name == new_name:
|
||||
return {"success": True, "affected_count": 0, "skipped_count": 0}
|
||||
|
||||
with self.model.get_connection() as conn:
|
||||
cursor = conn.execute("SELECT id FROM tags WHERE name = ?", (old_name,))
|
||||
@@ -942,12 +953,14 @@ class PromptDatabase:
|
||||
if existing_new:
|
||||
# Target tag exists — merge: move links, handle conflicts, delete old
|
||||
new_tag_id = existing_new["id"]
|
||||
conn.execute(
|
||||
"UPDATE OR IGNORE prompt_tags SET tag_id = ? WHERE tag_id = ?",
|
||||
(new_tag_id, old_tag_id),
|
||||
)
|
||||
# Delete remaining links (conflicts = prompt already had the target tag)
|
||||
conn.execute("DELETE FROM prompt_tags WHERE tag_id = ?", (old_tag_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,)
|
||||
)
|
||||
conn.execute("DELETE FROM tags WHERE id = ?", (old_tag_id,))
|
||||
else:
|
||||
# Simple rename
|
||||
@@ -966,7 +979,7 @@ class PromptDatabase:
|
||||
"""
|
||||
Remove a tag from all prompts that use it.
|
||||
|
||||
With junction tables, deleting from the tags table cascades to prompt_tags.
|
||||
Deleting from tags cascades to standalone and pair tag relationships.
|
||||
|
||||
Args:
|
||||
tag_name: Tag name to remove
|
||||
@@ -1033,20 +1046,27 @@ class PromptDatabase:
|
||||
if not src_row:
|
||||
continue
|
||||
src_id = src_row["id"]
|
||||
if src_id == target_id:
|
||||
continue
|
||||
|
||||
cursor = conn.execute(
|
||||
"SELECT COUNT(*) as c FROM prompt_tags WHERE tag_id = ?", (src_id,)
|
||||
)
|
||||
src_count = cursor.fetchone()["c"]
|
||||
|
||||
if src_count > 0:
|
||||
# Move links from source to target (ignore conflicts)
|
||||
conn.execute(
|
||||
"UPDATE OR IGNORE prompt_tags SET tag_id = ? WHERE tag_id = ?",
|
||||
(target_id, src_id),
|
||||
)
|
||||
# Delete remaining (conflicts = prompt already had target tag)
|
||||
conn.execute("DELETE FROM prompt_tags WHERE tag_id = ?", (src_id,))
|
||||
pair_count = conn.execute(
|
||||
"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,)
|
||||
)
|
||||
tags_merged += 1
|
||||
affected += src_count
|
||||
|
||||
@@ -1290,7 +1310,7 @@ class PromptDatabase:
|
||||
"""
|
||||
self.logger.info("Starting duplicate cleanup process with image preservation")
|
||||
try:
|
||||
with self.model.get_connection() as conn:
|
||||
with self._pair_transaction() as conn:
|
||||
# Find duplicates by text content (case-insensitive)
|
||||
# Note: Removed ORDER BY from GROUP_CONCAT for SQLite compatibility
|
||||
# We'll sort the IDs manually after fetching
|
||||
@@ -1319,6 +1339,11 @@ class PromptDatabase:
|
||||
id_date_pairs.sort(key=lambda x: x[1]) # Sort by date
|
||||
sorted_ids = [int(pair[0]) for pair in id_date_pairs]
|
||||
|
||||
try:
|
||||
self._ensure_prompts_unpaired(conn, sorted_ids)
|
||||
except PairDependencyError:
|
||||
continue
|
||||
|
||||
# Keep the oldest one (first), merge and delete the rest
|
||||
primary_id = sorted_ids[0] # Keep the oldest
|
||||
duplicate_ids = sorted_ids[1:]
|
||||
@@ -1864,7 +1889,8 @@ class PromptDatabase:
|
||||
Raises:
|
||||
sqlite3.IntegrityError: Another prompt already has the new hash.
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
with self._pair_transaction() as conn:
|
||||
self._ensure_prompts_unpaired(conn, [prompt_id])
|
||||
cursor = conn.execute(
|
||||
"UPDATE prompts SET text = ?, hash = ?, updated_at = ? WHERE id = ?",
|
||||
(
|
||||
@@ -1912,7 +1938,8 @@ class PromptDatabase:
|
||||
Number of prompts actually deleted.
|
||||
"""
|
||||
count = 0
|
||||
with self.model.get_connection() as conn:
|
||||
with self._pair_transaction() as conn:
|
||||
self._ensure_prompts_unpaired(conn, prompt_ids)
|
||||
for pid in prompt_ids:
|
||||
conn.execute("DELETE FROM generated_images WHERE prompt_id = ?", (pid,))
|
||||
cursor = conn.execute("DELETE FROM prompts WHERE id = ?", (pid,))
|
||||
@@ -1985,7 +2012,7 @@ class PromptDatabase:
|
||||
cursor = conn.execute(
|
||||
"""SELECT gi.prompt_id, p.text, p.category, p.rating, p.notes,
|
||||
gi.workflow_data, gi.prompt_metadata, gi.generation_time,
|
||||
gi.parameters
|
||||
gi.parameters, gi.id, gi.image_path, gi.file_path
|
||||
FROM generated_images gi
|
||||
JOIN prompts p ON gi.prompt_id = p.id
|
||||
WHERE gi.image_path = ? OR gi.file_path = ?
|
||||
@@ -1997,9 +2024,71 @@ class PromptDatabase:
|
||||
f"%{escape_like(os.path.basename(image_path))}",
|
||||
),
|
||||
)
|
||||
row = cursor.fetchone()
|
||||
if not row:
|
||||
rows = cursor.fetchall()
|
||||
exact = [
|
||||
row
|
||||
for row in rows
|
||||
if row[10] == image_path or row[11] == normalize_image_path(image_path)
|
||||
]
|
||||
matches = exact or rows
|
||||
if len(matches) != 1:
|
||||
return None
|
||||
row = matches[0]
|
||||
associations = self.get_image_prompt_pairs(row[9])
|
||||
pair_ids = {item["pair_id"] for item in associations}
|
||||
pair = (
|
||||
self.get_prompt_pair(next(iter(pair_ids)))
|
||||
if len(pair_ids) == 1
|
||||
else None
|
||||
)
|
||||
candidates = []
|
||||
for item in associations:
|
||||
evidence = item["provenance"]
|
||||
positive_side = evidence.get("positive")
|
||||
negative_side = evidence.get("negative")
|
||||
if not isinstance(positive_side, dict) or not isinstance(
|
||||
negative_side, dict
|
||||
):
|
||||
continue
|
||||
positive, negative = positive_side.get("text"), negative_side.get(
|
||||
"text"
|
||||
)
|
||||
if not isinstance(positive, str) or not isinstance(negative, str):
|
||||
continue
|
||||
sources = []
|
||||
for side in (positive_side, negative_side):
|
||||
provenance = side.get("provenance")
|
||||
sources.append(
|
||||
provenance.get("text_source")
|
||||
if isinstance(provenance, dict)
|
||||
else None
|
||||
)
|
||||
score = sources.count("execution")
|
||||
source = (
|
||||
"execution"
|
||||
if score == 2
|
||||
else (
|
||||
"mixed"
|
||||
if score
|
||||
else (
|
||||
"embedded-parameters"
|
||||
if sources == ["embedded-parameters"] * 2
|
||||
else "graph-inferred"
|
||||
)
|
||||
)
|
||||
)
|
||||
candidates.append((score, positive, negative, source))
|
||||
context = None
|
||||
if candidates and pair:
|
||||
best_score = max(item[0] for item in candidates)
|
||||
contexts = {item[1:] for item in candidates if item[0] == best_score}
|
||||
if len(contexts) == 1:
|
||||
positive, negative, source = next(iter(contexts))
|
||||
context = {
|
||||
"positive_text": positive,
|
||||
"negative_text": negative,
|
||||
"text_source": source,
|
||||
}
|
||||
tags = self._get_prompt_tags(conn, row[0])
|
||||
try:
|
||||
workflow = json.loads(row[5]) if row[5] else None
|
||||
@@ -2010,6 +2099,8 @@ class PromptDatabase:
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
prompt_meta = None
|
||||
return {
|
||||
"pair": pair,
|
||||
"pair_context": context,
|
||||
"prompt_id": row[0],
|
||||
"text": row[1],
|
||||
"category": row[2],
|
||||
@@ -2050,11 +2141,15 @@ class PromptDatabase:
|
||||
|
||||
Returns number of prompts removed.
|
||||
"""
|
||||
with self.model.get_connection() as conn:
|
||||
with self._pair_transaction() as conn:
|
||||
cursor = conn.execute("""
|
||||
SELECT p.id FROM prompts p
|
||||
LEFT JOIN generated_images gi ON p.id = gi.prompt_id
|
||||
WHERE gi.prompt_id IS NULL
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM prompt_pairs pp
|
||||
WHERE pp.positive_prompt_id = p.id OR pp.negative_prompt_id = p.id
|
||||
)
|
||||
AND NOT EXISTS (
|
||||
SELECT 1 FROM prompt_tags pt
|
||||
JOIN tags t ON pt.tag_id = t.id
|
||||
|
||||
@@ -0,0 +1,136 @@
|
||||
"""Atomic metadata enrichment for a single known generation context."""
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Dict
|
||||
|
||||
from .models import normalize_image_path
|
||||
from .operations import _coerce_file_info
|
||||
|
||||
try:
|
||||
from ..utils.hashing import generate_prompt_hash
|
||||
except ImportError:
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
|
||||
def ingest_resolved_pair(
|
||||
db: Any, image_path: str, metadata: Dict[str, Any], resolution: Dict[str, Any]
|
||||
) -> Dict[str, Any]:
|
||||
state = resolution.get("state", "unresolved")
|
||||
if state != "resolved":
|
||||
return {"status": state, "pairs_added": 0, "links_added": 0}
|
||||
positive, negative = resolution["positive"], resolution["negative"]
|
||||
positive_text = positive.get("storage_text", positive["text"])
|
||||
negative_text = negative.get("storage_text", negative["text"])
|
||||
normalized = normalize_image_path(image_path)
|
||||
provenance = {"resolver_version": 1, "positive": positive, "negative": negative}
|
||||
encoded = json.dumps(
|
||||
provenance, sort_keys=True, separators=(",", ":"), allow_nan=False
|
||||
)
|
||||
with db._pair_transaction() as conn:
|
||||
images = conn.execute(
|
||||
"SELECT gi.*,p.hash AS positive_hash FROM generated_images gi "
|
||||
"JOIN prompts p ON p.id=gi.prompt_id "
|
||||
"WHERE gi.file_path=? OR gi.image_path=?",
|
||||
(normalized, image_path),
|
||||
).fetchall()
|
||||
if len(images) > 1 or (
|
||||
images and images[0]["positive_hash"] != generate_prompt_hash(positive_text)
|
||||
):
|
||||
return {"status": "conflict", "pairs_added": 0, "links_added": 0}
|
||||
existing = images[0] if images else None
|
||||
if existing:
|
||||
prior = conn.execute(
|
||||
"SELECT n.hash AS negative_hash FROM image_prompt_pairs ip "
|
||||
"JOIN prompt_pairs pp ON pp.id=ip.pair_id "
|
||||
"LEFT JOIN prompts n ON n.id=pp.negative_prompt_id WHERE ip.image_id=?",
|
||||
(existing["id"],),
|
||||
).fetchall()
|
||||
expected = (
|
||||
generate_prompt_hash(negative_text) if negative_text.strip() else None
|
||||
)
|
||||
if any(row["negative_hash"] != expected for row in prior):
|
||||
return {"status": "conflict", "pairs_added": 0, "links_added": 0}
|
||||
prior_pair = conn.execute(
|
||||
"SELECT pp.id FROM prompt_pairs pp "
|
||||
"JOIN prompts p ON p.id=pp.positive_prompt_id "
|
||||
"LEFT JOIN prompts n ON n.id=pp.negative_prompt_id "
|
||||
"WHERE p.hash=? AND n.hash IS ?",
|
||||
(
|
||||
generate_prompt_hash(positive_text),
|
||||
generate_prompt_hash(negative_text) if negative_text.strip() else None,
|
||||
),
|
||||
).fetchone()
|
||||
pair_id = db._save_pair(conn, positive_text, negative_text)
|
||||
pair = conn.execute(
|
||||
"SELECT * FROM prompt_pairs WHERE id=?", (pair_id,)
|
||||
).fetchone()
|
||||
pairs_added = int(prior_pair is None)
|
||||
file_size, width, height, image_format = _coerce_file_info(metadata)
|
||||
snapshots = {
|
||||
column: json.dumps(metadata[key], sort_keys=True, allow_nan=False)
|
||||
for column, key in (
|
||||
("workflow_data", "workflow"),
|
||||
("prompt_metadata", "prompt"),
|
||||
("parameters", "parameters"),
|
||||
)
|
||||
if key in metadata
|
||||
}
|
||||
info = {
|
||||
"file_size": file_size,
|
||||
"width": width,
|
||||
"height": height,
|
||||
"format": image_format,
|
||||
}
|
||||
snapshots.update(
|
||||
{key: value for key, value in info.items() if value is not None}
|
||||
)
|
||||
changed = False
|
||||
if existing:
|
||||
image_id = existing["id"]
|
||||
updates = {
|
||||
key: value for key, value in snapshots.items() if existing[key] != value
|
||||
}
|
||||
if updates:
|
||||
conn.execute(
|
||||
"UPDATE generated_images SET "
|
||||
+ ",".join(f"{key}=?" for key in updates)
|
||||
+ " WHERE id=?",
|
||||
[*updates.values(), image_id],
|
||||
)
|
||||
changed = True
|
||||
else:
|
||||
columns = ["prompt_id", "image_path", "filename", "file_path", *snapshots]
|
||||
values = [
|
||||
pair["positive_prompt_id"],
|
||||
image_path,
|
||||
os.path.basename(image_path),
|
||||
normalized,
|
||||
*snapshots.values(),
|
||||
]
|
||||
image_id = conn.execute(
|
||||
"INSERT INTO generated_images("
|
||||
+ ",".join(columns)
|
||||
+ ") VALUES("
|
||||
+ ",".join("?" for _ in columns)
|
||||
+ ")",
|
||||
values,
|
||||
).lastrowid
|
||||
linked = conn.execute(
|
||||
"SELECT 1 FROM image_prompt_pairs WHERE pair_id=? "
|
||||
"AND image_id=? AND provenance=?",
|
||||
(pair_id, image_id, encoded),
|
||||
).fetchone()
|
||||
db._link_image_to_pair(conn, pair_id, image_id, provenance)
|
||||
status = (
|
||||
"added"
|
||||
if existing is None
|
||||
else "enriched" if changed or not linked else "unchanged"
|
||||
)
|
||||
return {
|
||||
"status": status,
|
||||
"pairs_added": pairs_added,
|
||||
"links_added": int(not linked),
|
||||
"pair_id": pair_id,
|
||||
"image_id": image_id,
|
||||
}
|
||||
@@ -0,0 +1,387 @@
|
||||
"""Ordered reusable prompt pairs and their generation provenance."""
|
||||
|
||||
import json
|
||||
import sqlite3
|
||||
from contextlib import contextmanager
|
||||
from typing import Any, Dict, Iterator, List, Optional
|
||||
|
||||
from .models import pair_provenance_error, utc_now_iso
|
||||
|
||||
try:
|
||||
from ..utils.hashing import generate_prompt_hash
|
||||
except ImportError:
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
|
||||
class PairDependencyError(ValueError):
|
||||
"""An operation would change or remove content used by a saved pair."""
|
||||
|
||||
|
||||
def _validate_metadata(
|
||||
title: Optional[str],
|
||||
category: Optional[str],
|
||||
tags: Optional[List[str]],
|
||||
rating: Optional[int],
|
||||
notes: Optional[str],
|
||||
) -> Optional[List[str]]:
|
||||
for name, value in (("title", title), ("category", category), ("notes", notes)):
|
||||
if value is not None and not isinstance(value, str):
|
||||
raise ValueError(f"{name} must be a string")
|
||||
if rating is not None and (type(rating) is not int or not 1 <= rating <= 5):
|
||||
raise ValueError("Rating must be between 1 and 5")
|
||||
if tags is None:
|
||||
return None
|
||||
if not isinstance(tags, list) or any(not isinstance(tag, str) for tag in tags):
|
||||
raise ValueError("tags must be a list of strings")
|
||||
return list(dict.fromkeys(tag.strip() for tag in tags if tag.strip()))
|
||||
|
||||
|
||||
class PairOperations:
|
||||
"""Connection-level pair operations mixed into PromptDatabase."""
|
||||
|
||||
@contextmanager
|
||||
def _pair_transaction(self) -> Iterator[sqlite3.Connection]:
|
||||
conn = self.model.get_connection()
|
||||
if conn.in_transaction:
|
||||
raise RuntimeError("Pair operations require their own transaction")
|
||||
with conn:
|
||||
conn.execute("BEGIN IMMEDIATE")
|
||||
yield conn
|
||||
|
||||
def _ensure_prompts_unpaired(
|
||||
self, conn: sqlite3.Connection, prompt_ids: List[int]
|
||||
) -> None:
|
||||
if not prompt_ids:
|
||||
return
|
||||
placeholders = ",".join("?" for _ in prompt_ids)
|
||||
row = conn.execute(
|
||||
f"SELECT id FROM prompt_pairs WHERE positive_prompt_id IN ({placeholders}) "
|
||||
f"OR negative_prompt_id IN ({placeholders}) LIMIT 1",
|
||||
[*prompt_ids, *prompt_ids],
|
||||
).fetchone()
|
||||
if row:
|
||||
raise PairDependencyError(
|
||||
f"Prompt belongs to saved pair {row['id']}; delete the pair first "
|
||||
"or save changed text as a new prompt"
|
||||
)
|
||||
|
||||
def _pair_component(self, conn: sqlite3.Connection, text: str) -> int:
|
||||
prompt_hash = generate_prompt_hash(text)
|
||||
conn.execute(
|
||||
"INSERT INTO prompts(text, hash, created_at, updated_at) VALUES(?,?,?,?) "
|
||||
"ON CONFLICT(hash) DO NOTHING",
|
||||
(text.strip(), prompt_hash, utc_now_iso(), utc_now_iso()),
|
||||
)
|
||||
return conn.execute(
|
||||
"SELECT id FROM prompts WHERE hash = ?", (prompt_hash,)
|
||||
).fetchone()["id"]
|
||||
|
||||
def _sync_pair_tags(
|
||||
self, conn: sqlite3.Connection, pair_id: int, tags: List[str]
|
||||
) -> None:
|
||||
conn.execute("DELETE FROM prompt_pair_tags WHERE pair_id = ?", (pair_id,))
|
||||
for tag_id in self._ensure_tags(conn, tags).values():
|
||||
conn.execute(
|
||||
"INSERT INTO prompt_pair_tags(pair_id,tag_id) VALUES(?,?)",
|
||||
(pair_id, tag_id),
|
||||
)
|
||||
|
||||
def save_prompt_pair(
|
||||
self,
|
||||
positive_text: str,
|
||||
negative_text: str = "",
|
||||
*,
|
||||
title: Optional[str] = None,
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[List[str]] = None,
|
||||
rating: Optional[int] = None,
|
||||
notes: Optional[str] = None,
|
||||
) -> int:
|
||||
"""Save an ordered pair atomically; repeat saves preserve its metadata."""
|
||||
with self._pair_transaction() as conn:
|
||||
return self._save_pair(
|
||||
conn,
|
||||
positive_text,
|
||||
negative_text,
|
||||
title=title,
|
||||
category=category,
|
||||
tags=tags,
|
||||
rating=rating,
|
||||
notes=notes,
|
||||
)
|
||||
|
||||
def _save_pair(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
positive_text: str,
|
||||
negative_text: str = "",
|
||||
*,
|
||||
title: Optional[str] = None,
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[List[str]] = None,
|
||||
rating: Optional[int] = None,
|
||||
notes: Optional[str] = None,
|
||||
) -> int:
|
||||
"""Save an ordered pair atomically; repeat saves preserve its metadata."""
|
||||
if not isinstance(positive_text, str) or not positive_text.strip():
|
||||
raise ValueError("Positive prompt text cannot be empty")
|
||||
if not isinstance(negative_text, str):
|
||||
raise ValueError("Negative prompt text must be a string; use '' for empty")
|
||||
tags = _validate_metadata(title, category, tags, rating, notes)
|
||||
positive = self._pair_component(conn, positive_text)
|
||||
negative = (
|
||||
self._pair_component(conn, negative_text) if negative_text.strip() else None
|
||||
)
|
||||
existing = conn.execute(
|
||||
"SELECT id FROM prompt_pairs WHERE positive_prompt_id = ? "
|
||||
"AND negative_prompt_id IS ?",
|
||||
(positive, negative),
|
||||
).fetchone()
|
||||
if existing:
|
||||
return existing["id"]
|
||||
now = utc_now_iso()
|
||||
pair_id = conn.execute(
|
||||
"INSERT INTO prompt_pairs(positive_prompt_id,negative_prompt_id,title,"
|
||||
"category,rating,notes,created_at,updated_at) VALUES(?,?,?,?,?,?,?,?)",
|
||||
(positive, negative, title, category, rating, notes, now, now),
|
||||
).lastrowid
|
||||
if tags:
|
||||
self._sync_pair_tags(conn, pair_id, tags)
|
||||
return pair_id
|
||||
|
||||
def _read_pair_rows(
|
||||
self, conn: sqlite3.Connection, where: str, params: List[Any], suffix: str = ""
|
||||
) -> List[Dict[str, Any]]:
|
||||
rows = conn.execute(
|
||||
"SELECT pp.*, p.text AS positive_text, "
|
||||
"COALESCE(n.text,'') AS negative_text,"
|
||||
"(SELECT COUNT(DISTINCT image_id) FROM image_prompt_pairs ip "
|
||||
"WHERE ip.pair_id=pp.id) AS image_count FROM prompt_pairs pp "
|
||||
"JOIN prompts p ON p.id=pp.positive_prompt_id "
|
||||
"LEFT JOIN prompts n ON n.id=pp.negative_prompt_id WHERE " + where + suffix,
|
||||
params,
|
||||
).fetchall()
|
||||
pairs = [dict(row) for row in rows]
|
||||
if not pairs:
|
||||
return pairs
|
||||
by_id = {pair["id"]: pair for pair in pairs}
|
||||
for pair in pairs:
|
||||
pair["archived"] = bool(pair["archived"])
|
||||
pair["tags"] = []
|
||||
placeholders = ",".join("?" for _ in by_id)
|
||||
for row in conn.execute(
|
||||
"SELECT pt.pair_id,t.name FROM prompt_pair_tags pt JOIN tags t ON "
|
||||
f"t.id=pt.tag_id WHERE pt.pair_id IN ({placeholders}) ORDER BY t.name",
|
||||
list(by_id),
|
||||
):
|
||||
by_id[row["pair_id"]]["tags"].append(row["name"])
|
||||
return pairs
|
||||
|
||||
def get_prompt_pair(self, pair_id: int) -> Optional[Dict[str, Any]]:
|
||||
"""Return components, pair metadata and image count, including archives."""
|
||||
rows = self._read_pair_rows(self.model.get_connection(), "pp.id = ?", [pair_id])
|
||||
return rows[0] if rows else None
|
||||
|
||||
def list_prompt_pairs(
|
||||
self,
|
||||
*,
|
||||
text: Optional[str] = None,
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[List[str]] = None,
|
||||
include_archived: bool = False,
|
||||
limit: int = 50,
|
||||
offset: int = 0,
|
||||
sort: str = "newest",
|
||||
) -> Dict[str, Any]:
|
||||
"""List pairs with all-tag filtering and stable paginated ordering."""
|
||||
orders = {
|
||||
"newest": "pp.id DESC",
|
||||
"last_used_desc": "julianday(pp.last_used_at) DESC, pp.id DESC",
|
||||
}
|
||||
if sort not in orders:
|
||||
raise ValueError("Unknown pair sort order")
|
||||
if type(limit) is not int or not 1 <= limit <= 500:
|
||||
raise ValueError("limit must be between 1 and 500")
|
||||
if type(offset) is not int or offset < 0:
|
||||
raise ValueError("offset must be a nonnegative integer")
|
||||
tags = _validate_metadata(None, category, tags, None, None)
|
||||
if text is not None and not isinstance(text, str):
|
||||
raise ValueError("text must be a string")
|
||||
clauses, params = ["1=1"], []
|
||||
if not include_archived:
|
||||
clauses.append("pp.archived=0")
|
||||
if category is not None:
|
||||
clauses.append("pp.category = ?")
|
||||
params.append(category)
|
||||
if text:
|
||||
clauses.append("(p.text LIKE ? OR n.text LIKE ? OR pp.title LIKE ?)")
|
||||
params.extend([f"%{text}%"] * 3)
|
||||
for tag in tags or []:
|
||||
clauses.append(
|
||||
"EXISTS(SELECT 1 FROM prompt_pair_tags pt JOIN tags t "
|
||||
"ON t.id=pt.tag_id WHERE pt.pair_id=pp.id AND t.name=?)"
|
||||
)
|
||||
params.append(tag)
|
||||
where = " AND ".join(clauses)
|
||||
conn = self.model.get_connection()
|
||||
total = conn.execute(
|
||||
"SELECT COUNT(*) FROM prompt_pairs pp "
|
||||
"JOIN prompts p ON p.id=pp.positive_prompt_id "
|
||||
"LEFT JOIN prompts n ON n.id=pp.negative_prompt_id WHERE " + where,
|
||||
params,
|
||||
).fetchone()[0]
|
||||
pairs = self._read_pair_rows(
|
||||
conn,
|
||||
where,
|
||||
[*params, limit, offset],
|
||||
f" ORDER BY {orders[sort]} LIMIT ? OFFSET ?",
|
||||
)
|
||||
return {
|
||||
"pairs": pairs,
|
||||
"total": total,
|
||||
"limit": limit,
|
||||
"offset": offset,
|
||||
"has_more": offset + len(pairs) < total,
|
||||
}
|
||||
|
||||
def update_prompt_pair_metadata(
|
||||
self,
|
||||
pair_id: int,
|
||||
*,
|
||||
title: Optional[str] = None,
|
||||
category: Optional[str] = None,
|
||||
tags: Optional[List[str]] = None,
|
||||
rating: Optional[int] = None,
|
||||
notes: Optional[str] = None,
|
||||
) -> bool:
|
||||
"""Update pair metadata explicitly; None preserves and empty values clear."""
|
||||
tags = _validate_metadata(title, category, tags, rating, notes)
|
||||
changes = {
|
||||
key: value
|
||||
for key, value in (
|
||||
("title", title),
|
||||
("category", category),
|
||||
("rating", rating),
|
||||
("notes", notes),
|
||||
)
|
||||
if value is not None
|
||||
}
|
||||
if not changes and tags is None:
|
||||
return False
|
||||
with self._pair_transaction() as conn:
|
||||
if not conn.execute(
|
||||
"SELECT 1 FROM prompt_pairs WHERE id=?", (pair_id,)
|
||||
).fetchone():
|
||||
return False
|
||||
changes["updated_at"] = utc_now_iso()
|
||||
conn.execute(
|
||||
"UPDATE prompt_pairs SET "
|
||||
+ ",".join(f"{key}=?" for key in changes)
|
||||
+ " WHERE id=?",
|
||||
[*changes.values(), pair_id],
|
||||
)
|
||||
if tags is not None:
|
||||
self._sync_pair_tags(conn, pair_id, tags)
|
||||
return True
|
||||
|
||||
def archive_prompt_pair(self, pair_id: int, archived: bool = True) -> bool:
|
||||
"""Hide a reusable pair without changing components or generation history."""
|
||||
if type(archived) is not bool:
|
||||
raise ValueError("archived must be a boolean")
|
||||
with self._pair_transaction() as conn:
|
||||
return (
|
||||
conn.execute(
|
||||
"UPDATE prompt_pairs SET archived=?,updated_at=? WHERE id=?",
|
||||
(int(archived), utc_now_iso(), pair_id),
|
||||
).rowcount
|
||||
> 0
|
||||
)
|
||||
|
||||
def delete_prompt_pair(self, pair_id: int) -> bool:
|
||||
"""Delete unused pair metadata; generated pairs must be archived instead."""
|
||||
with self._pair_transaction() as conn:
|
||||
if conn.execute(
|
||||
"SELECT 1 FROM image_prompt_pairs WHERE pair_id=? LIMIT 1", (pair_id,)
|
||||
).fetchone():
|
||||
raise PairDependencyError(
|
||||
"Pair has generated images; archive it instead"
|
||||
)
|
||||
return (
|
||||
conn.execute("DELETE FROM prompt_pairs WHERE id=?", (pair_id,)).rowcount
|
||||
> 0
|
||||
)
|
||||
|
||||
def record_prompt_pair_use(self, pair_id: int, times: int = 1) -> bool:
|
||||
"""Count pair uses independently of standalone prompt usage."""
|
||||
if type(times) is not int or times < 1:
|
||||
raise ValueError("times must be a positive integer")
|
||||
with self._pair_transaction() as conn:
|
||||
return (
|
||||
conn.execute(
|
||||
"UPDATE prompt_pairs SET run_count=run_count+?,last_used_at=? "
|
||||
"WHERE id=?",
|
||||
(times, utc_now_iso(), pair_id),
|
||||
).rowcount
|
||||
> 0
|
||||
)
|
||||
|
||||
def link_image_to_pair(
|
||||
self, pair_id: int, image_id: int, provenance: Dict[str, Any]
|
||||
) -> int:
|
||||
"""Associate an image with immutable run provenance, idempotently."""
|
||||
with self._pair_transaction() as conn:
|
||||
return self._link_image_to_pair(conn, pair_id, image_id, provenance)
|
||||
|
||||
def _link_image_to_pair(
|
||||
self,
|
||||
conn: sqlite3.Connection,
|
||||
pair_id: int,
|
||||
image_id: int,
|
||||
provenance: Dict[str, Any],
|
||||
) -> int:
|
||||
"""Associate an image with immutable run provenance, idempotently."""
|
||||
if not isinstance(provenance, dict):
|
||||
raise ValueError("provenance must be an object")
|
||||
try:
|
||||
encoded = json.dumps(
|
||||
provenance, sort_keys=True, separators=(",", ":"), allow_nan=False
|
||||
)
|
||||
except (TypeError, ValueError, RecursionError) as error:
|
||||
raise ValueError("provenance must contain JSON values") from error
|
||||
error = pair_provenance_error(encoded)
|
||||
if error:
|
||||
raise ValueError(error)
|
||||
pair = conn.execute(
|
||||
"SELECT positive_prompt_id FROM prompt_pairs WHERE id=?", (pair_id,)
|
||||
).fetchone()
|
||||
image = conn.execute(
|
||||
"SELECT prompt_id FROM generated_images WHERE id=?", (image_id,)
|
||||
).fetchone()
|
||||
if not pair or not image:
|
||||
raise ValueError("Pair and image must exist")
|
||||
if pair["positive_prompt_id"] != image["prompt_id"]:
|
||||
raise ValueError("Image must belong to the pair's positive prompt")
|
||||
conn.execute(
|
||||
"INSERT INTO image_prompt_pairs(pair_id,image_id,provenance,created_at) "
|
||||
"VALUES(?,?,?,?) ON CONFLICT(pair_id,image_id,provenance) DO NOTHING",
|
||||
(pair_id, image_id, encoded, utc_now_iso()),
|
||||
)
|
||||
return conn.execute(
|
||||
"SELECT id FROM image_prompt_pairs "
|
||||
"WHERE pair_id=? AND image_id=? AND provenance=?",
|
||||
(pair_id, image_id, encoded),
|
||||
).fetchone()["id"]
|
||||
|
||||
def get_image_prompt_pairs(self, image_id: int) -> List[Dict[str, Any]]:
|
||||
"""Return every recorded pair association with decoded provenance."""
|
||||
rows = (
|
||||
self.model.get_connection()
|
||||
.execute(
|
||||
"SELECT * FROM image_prompt_pairs WHERE image_id=? ORDER BY id",
|
||||
(image_id,),
|
||||
)
|
||||
.fetchall()
|
||||
)
|
||||
return [
|
||||
{**dict(row), "provenance": json.loads(row["provenance"])} for row in rows
|
||||
]
|
||||
@@ -0,0 +1,111 @@
|
||||
# Reusable positive and negative prompts
|
||||
|
||||
Tracking: [#193](https://github.com/ComfyAssets/ComfyUI_PromptManager/issues/193).
|
||||
This work builds on the merged Split node in PR #192. Keep the feature unmerged
|
||||
until maintainer live testing is complete.
|
||||
|
||||
## Storage contract
|
||||
|
||||
Keep existing prompt IDs, text hashes, metadata, image links and node schemas.
|
||||
A component can serve either role. Ordered combinations refer to shared prompt
|
||||
rows; an explicitly empty negative uses a null reference. Missing or ambiguous
|
||||
metadata does not mean an empty negative. Pair metadata belongs to the
|
||||
combination; saving another combination does not overwrite component metadata.
|
||||
Existing trim/lowercase hash normalization remains unchanged.
|
||||
|
||||
Saved components are immutable while referenced. To change text, create or reuse
|
||||
a component and save a new combination. Archive combinations that have generation
|
||||
history. Pruning and legacy editing/deletion must respect these dependencies.
|
||||
Historical workflow snapshots remain intact, including their original text.
|
||||
|
||||
## Migration and recovery
|
||||
|
||||
Create pair tables after existing legacy rebuild migrations, in an atomic,
|
||||
repeatable migration. Never scan media during schema initialization. Test fresh,
|
||||
legacy and current databases, repeated startup, injected failures and retries.
|
||||
Preserve old backup compatibility. Migrate and validate a disposable backup copy
|
||||
before replacing live database pages. Reject broken pair schemas, foreign keys,
|
||||
triggers and views. Use SQLite backup APIs so committed WAL data is included.
|
||||
|
||||
## Metadata and scanning
|
||||
|
||||
Resolve positive and negative from the same sampler/output context. Track both
|
||||
node ID and output slot; bound graph traversal and text construction. Unknown
|
||||
conditioning operations and conflicting sampler candidates remain unresolved.
|
||||
Never guess that the first encoder or most recent negative belongs to an image.
|
||||
|
||||
Import high-confidence pairs atomically with their image association and source
|
||||
provenance. Repeated scans must not duplicate components, pairs or image links,
|
||||
regenerate thumbnails, or increment generation run counts. Preserve existing
|
||||
conflicting/manual associations for review. Report new, enriched, unchanged,
|
||||
ambiguous, unresolved and failed results separately.
|
||||
|
||||
## Live node
|
||||
|
||||
Use the shared toolkit and Split foundation. Positive and negative editors have
|
||||
equal size. Only positive receives prepend, append and automatic LoRA triggers.
|
||||
Keep upstream connections authoritative. Library selection previews both fields
|
||||
before an explicit load. Saving text changes creates another reusable combination.
|
||||
Verify Classic and Nodes 2.0 startup, switching both directions, reload, caching,
|
||||
linked inputs, popup behavior and lifecycle cleanup before maintainer testing.
|
||||
|
||||
## Delivery gate
|
||||
|
||||
Keep all existing coverage floors and test suites. Record regression failures
|
||||
before implementing each behavior. Tests use disposable state and the managed
|
||||
CPU host; never mutate the maintainer's running ComfyUI or personal database.
|
||||
Python changes require a ComfyUI restart. Live approval is required before merge.
|
||||
|
||||
## Live testing
|
||||
|
||||
1. Download a database backup from Admin before the first restart. The additive
|
||||
schema migration runs on startup; existing workflows and prompt IDs stay valid.
|
||||
2. Restart ComfyUI, reload the browser, and add **Prompt Manager (+/-)** from
|
||||
**ComfyAssets / Prompts**. Connect CLIP and the two conditioning outputs to the
|
||||
sampler's positive and negative inputs.
|
||||
3. Enter both texts, save a pair, then keep positive unchanged and change negative.
|
||||
Recent should show both combinations with distinct negative excerpts. Selecting
|
||||
one previews both; **Use pair** loads both together.
|
||||
4. Test category/tag autocomplete and metadata opt-in, resizing, light/dark,
|
||||
switching Classic / Nodes 2.0 both ways, and workflow reload. Connected fields
|
||||
stay authoritative; run the node to save connected metadata values.
|
||||
5. Queue the same recipe twice and change only negative once. Usage should count
|
||||
dequeued jobs even when ComfyUI caches the node. Prepend, append and LoRA triggers
|
||||
affect positive only. Gallery text is labelled as recorded during execution or
|
||||
recovered from workflow metadata.
|
||||
6. Run an output metadata scan to enrich existing images. Repeat it: pairs and
|
||||
links must remain deduplicated, with unchanged results rather than new links.
|
||||
Conflicting or ambiguous files need review and are preserved.
|
||||
|
||||
Fresh executions embed bounded text evidence in PNG metadata. A cached node may
|
||||
not emit fresh evidence, so its graph-derived text is explicitly marked inferred.
|
||||
Multiple distinct batch records without an image anchor remain unresolved.
|
||||
Rescanning after interruption is safe through idempotent per-file writes; this
|
||||
version does not persist a filesystem scan cursor across server restarts.
|
||||
The API can request cancellation with `POST /prompt_manager/scan/cancel`; the
|
||||
existing dashboard Hide button continues to hide progress without cancelling.
|
||||
|
||||
Pair APIs are under `/prompt_manager/pairs`: POST saves a combination, GET lists
|
||||
and filters, GET/PATCH/DELETE `/{id}` reads, edits metadata or removes an unused
|
||||
pair. Changing text creates another combination. Null metadata updates preserve
|
||||
existing values; empty strings/lists clear supported fields. Archive a pair with
|
||||
`PATCH {"archived":true}`. GET `/export` provides versioned JSON containing both
|
||||
texts and pair metadata; this is an export format, not a JSON import endpoint.
|
||||
|
||||
## Gallery redesign concept retained
|
||||
|
||||
The maintainer approved the interactive
|
||||
[negative variant timeline concept](concepts/prompt-variant-timeline.html) on
|
||||
2026-10-17. Preserve it for the later gallery redesign; it is not production UI
|
||||
and does not expand the current node work.
|
||||
|
||||
Keep one positive parent and one shared gallery, with red branches for reusable
|
||||
negative variants and smaller dots for repeat runs. Show the selected negative
|
||||
in full directly beneath the unchanged positive; selecting another variant or
|
||||
run updates that text and filters the shared gallery. Comparison shows both
|
||||
negatives. Distinguish explicitly empty negatives from an all-variants view and
|
||||
saved variants from executed runs. Production history needs persisted execution
|
||||
records; the prototype's timeline uses illustrative sample data.
|
||||
|
||||
The Admin upgrade/scan notice remains a separate follow-up: explain automatic
|
||||
schema migration and offer an optional metadata enrichment scan or Later.
|
||||
|
After Width: | Height: | Size: 88 KiB |
|
After Width: | Height: | Size: 106 KiB |
|
After Width: | Height: | Size: 89 KiB |
|
After Width: | Height: | Size: 98 KiB |
|
After Width: | Height: | Size: 91 KiB |
|
After Width: | Height: | Size: 105 KiB |
|
After Width: | Height: | Size: 91 KiB |
|
After Width: | Height: | Size: 101 KiB |
@@ -0,0 +1,226 @@
|
||||
"""Paired positive and negative conditioning with reusable recipe storage."""
|
||||
|
||||
import hashlib
|
||||
import json
|
||||
import time
|
||||
from typing import Any, Dict, Optional, Tuple
|
||||
|
||||
try:
|
||||
from .prompt_manager import IO, InputTypeDict, PromptManagerFocused
|
||||
from .utils.hashing import generate_prompt_hash
|
||||
from .utils.pair_execution import record_pair_execution
|
||||
from .utils.prompt_graph import intrinsic_pair_nodes, run_prompt_components
|
||||
from .utils.prompt_pair_resolver import resolve_text_input
|
||||
from .utils.usage_tracking import (
|
||||
PENDING_FIRST_USE,
|
||||
PENDING_PAIR_FIRST_USE,
|
||||
pair_run_key,
|
||||
)
|
||||
except ImportError:
|
||||
from prompt_manager import IO, InputTypeDict, PromptManagerFocused
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.pair_execution import record_pair_execution
|
||||
from utils.prompt_graph import intrinsic_pair_nodes, run_prompt_components
|
||||
from utils.prompt_pair_resolver import resolve_text_input
|
||||
from utils.usage_tracking import (
|
||||
PENDING_FIRST_USE,
|
||||
PENDING_PAIR_FIRST_USE,
|
||||
pair_run_key,
|
||||
)
|
||||
|
||||
|
||||
class PromptManagerPair(PromptManagerFocused):
|
||||
"""Keep raw editors as one recipe while encoding their effective texts."""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls) -> InputTypeDict:
|
||||
inputs = super().INPUT_TYPES()
|
||||
inputs["required"] = {
|
||||
"text": inputs["required"]["text"],
|
||||
"negative_text": (
|
||||
IO.STRING,
|
||||
{"default": "", "multiline": True, "dynamicPrompts": True},
|
||||
),
|
||||
"clip": inputs["required"]["clip"],
|
||||
}
|
||||
inputs["hidden"]["extra_pnginfo"] = "EXTRA_PNGINFO"
|
||||
return inputs
|
||||
|
||||
RETURN_TYPES = (IO.CONDITIONING, IO.CONDITIONING, IO.STRING, IO.STRING)
|
||||
RETURN_NAMES = ("positive", "negative", "positive_text", "negative_text")
|
||||
OUTPUT_TOOLTIPS = (
|
||||
"Positive conditioning with composition and optional LoRA trigger words.",
|
||||
"Negative conditioning from the negative editor, including empty text.",
|
||||
"The effective positive text that was encoded.",
|
||||
"The unchanged negative text that was encoded.",
|
||||
)
|
||||
DESCRIPTION = "Encode and save a reusable positive/negative prompt pair."
|
||||
|
||||
def encode_prompt(
|
||||
self,
|
||||
clip: Any,
|
||||
text: str,
|
||||
negative_text: str = "",
|
||||
category: str = "",
|
||||
tags: str = "",
|
||||
search_text: str = "",
|
||||
prepend_text: str = "",
|
||||
append_text: str = "",
|
||||
prompt: Optional[Dict[str, Any]] = None,
|
||||
unique_id: Optional[str] = None,
|
||||
extra_pnginfo: Optional[Dict[str, Any]] = None,
|
||||
) -> Tuple[Any, Any, str, str]:
|
||||
"""Encode both sides once; database and tracking failures never block CLIP."""
|
||||
if clip is None:
|
||||
raise RuntimeError("ERROR: clip input is invalid: None")
|
||||
if not isinstance(text, str) or not isinstance(negative_text, str):
|
||||
raise ValueError("Positive and negative prompt text must be strings")
|
||||
positive = self._inject_lora_trigger_words(
|
||||
" ".join(
|
||||
part
|
||||
for part in (prepend_text.strip(), text, append_text.strip())
|
||||
if part
|
||||
)
|
||||
)
|
||||
positive_conditioning = clip.encode_from_tokens_scheduled(
|
||||
clip.tokenize(positive)
|
||||
)
|
||||
negative_conditioning = clip.encode_from_tokens_scheduled(
|
||||
clip.tokenize(negative_text)
|
||||
)
|
||||
record_pair_execution(
|
||||
extra_pnginfo,
|
||||
unique_id,
|
||||
text=text,
|
||||
negative_text=negative_text,
|
||||
positive_text=positive,
|
||||
negative_effective_text=negative_text,
|
||||
prepend_text=prepend_text.strip(),
|
||||
append_text=append_text.strip(),
|
||||
)
|
||||
pair_id = prompt_id = None
|
||||
known = isinstance(prompt, dict) and unique_id is not None
|
||||
node_id = str(unique_id)
|
||||
used_fields = (
|
||||
{
|
||||
field
|
||||
for source, field in run_prompt_components(prompt)
|
||||
if source == node_id
|
||||
}
|
||||
if known
|
||||
else {"text"}
|
||||
)
|
||||
is_positive = "text" in used_fields
|
||||
pair_used = node_id in intrinsic_pair_nodes(prompt) if known else True
|
||||
try:
|
||||
pair_id = self.db.save_prompt_pair(
|
||||
text,
|
||||
negative_text,
|
||||
category=category.strip() or None,
|
||||
tags=self._parse_tags(tags),
|
||||
)
|
||||
pair = self.db.get_prompt_pair(pair_id)
|
||||
prompt_id = pair["positive_prompt_id"]
|
||||
resolved_positive = (
|
||||
resolve_text_input(prompt, node_id, "text") if known else None
|
||||
)
|
||||
resolved_negative = (
|
||||
resolve_text_input(prompt, node_id, "negative_text") if known else None
|
||||
)
|
||||
positive_seen = resolved_positive is not None and generate_prompt_hash(
|
||||
resolved_positive
|
||||
) == generate_prompt_hash(text)
|
||||
pair_seen = (
|
||||
positive_seen
|
||||
and resolved_negative is not None
|
||||
and generate_prompt_hash(resolved_negative)
|
||||
== generate_prompt_hash(negative_text)
|
||||
)
|
||||
positive_pending = getattr(self, "pending_first_use", PENDING_FIRST_USE)
|
||||
pair_pending = getattr(
|
||||
self, "pending_pair_first_use", PENDING_PAIR_FIRST_USE
|
||||
)
|
||||
counted = set()
|
||||
for field, component_text, resolved_text, component_id in (
|
||||
("text", text, resolved_positive, prompt_id),
|
||||
(
|
||||
"negative_text",
|
||||
negative_text,
|
||||
resolved_negative,
|
||||
pair["negative_prompt_id"],
|
||||
),
|
||||
):
|
||||
component_hash = generate_prompt_hash(component_text)
|
||||
if (
|
||||
field not in used_fields
|
||||
or not component_id
|
||||
or component_hash in counted
|
||||
):
|
||||
continue
|
||||
counted.add(component_hash)
|
||||
seen_at_queue = (
|
||||
resolved_text is not None
|
||||
and generate_prompt_hash(resolved_text) == component_hash
|
||||
)
|
||||
runs = positive_pending.consume(component_hash) + int(not seen_at_queue)
|
||||
if runs:
|
||||
self.db.record_prompt_use(component_id, times=runs)
|
||||
if pair_used:
|
||||
runs = pair_pending.consume(pair_run_key(text, negative_text)) + int(
|
||||
not pair_seen
|
||||
)
|
||||
if runs:
|
||||
self.db.record_prompt_pair_use(pair_id, times=runs)
|
||||
if is_positive:
|
||||
self.prompt_tracker.set_current_prompt(
|
||||
prompt_text=positive,
|
||||
additional_data={
|
||||
"prompt_id": prompt_id,
|
||||
"pair_id": pair_id if pair_used else None,
|
||||
"positive_text": positive,
|
||||
"negative_text": negative_text,
|
||||
"positive_storage_text": text.strip(),
|
||||
"negative_storage_text": negative_text.strip(),
|
||||
"prepend_text": prepend_text.strip(),
|
||||
"append_text": append_text.strip(),
|
||||
"final_text": positive,
|
||||
"category": category.strip() or None,
|
||||
"tags": self._parse_tags(tags) or [],
|
||||
},
|
||||
push_to_queue=not positive_seen or (pair_used and not pair_seen),
|
||||
)
|
||||
except Exception as error:
|
||||
self.logger.warning(f"Could not save or track paired prompt: {error}")
|
||||
try:
|
||||
self.comfyui_integration.register_prompt(
|
||||
f"promptmanager_pair_{time.time_ns()}",
|
||||
positive,
|
||||
{
|
||||
"prompt_id": prompt_id,
|
||||
"pair_id": pair_id if pair_used else None,
|
||||
"negative_text": negative_text,
|
||||
"category": category.strip() or None,
|
||||
},
|
||||
)
|
||||
except Exception as error:
|
||||
self.logger.warning(f"Could not register paired metadata: {error}")
|
||||
return positive_conditioning, negative_conditioning, positive, negative_text
|
||||
|
||||
@classmethod
|
||||
def IS_CHANGED(
|
||||
cls,
|
||||
clip: Any,
|
||||
text: str = "",
|
||||
negative_text: str = "",
|
||||
category: str = "",
|
||||
tags: str = "",
|
||||
search_text: str = "",
|
||||
prepend_text: str = "",
|
||||
append_text: str = "",
|
||||
**kwargs: Any,
|
||||
) -> str:
|
||||
"""Either editor, composition or pair metadata invalidates the node cache."""
|
||||
serialized = json.dumps(
|
||||
[text, negative_text, prepend_text, append_text, category, tags]
|
||||
)
|
||||
return hashlib.sha256(serialized.encode("utf-8")).hexdigest()
|
||||
@@ -22,6 +22,7 @@ from aiohttp import web
|
||||
from PIL import Image
|
||||
|
||||
from .prompts import PromptRoutesMixin
|
||||
from .pairs import PairRoutesMixin
|
||||
from .images import ImageRoutesMixin
|
||||
from .admin import AdminRoutesMixin
|
||||
from .logging_routes import LoggingRoutesMixin
|
||||
@@ -240,6 +241,7 @@ async def _gzip_middleware(request, handler):
|
||||
|
||||
|
||||
class PromptManagerAPI(
|
||||
PairRoutesMixin,
|
||||
PromptRoutesMixin,
|
||||
ImageRoutesMixin,
|
||||
AdminRoutesMixin,
|
||||
@@ -262,12 +264,12 @@ class PromptManagerAPI(
|
||||
db (PromptDatabase): Database connection and operations handler
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""Initialize the PromptManager API with database connection and cleanup."""
|
||||
def __init__(self, db: PromptDatabase | None = None) -> None:
|
||||
"""Initialize the API, using an injected database before any cleanup."""
|
||||
self.logger = get_logger("prompt_manager.api")
|
||||
self.logger.info("Initializing PromptManager API")
|
||||
|
||||
self.db = PromptDatabase()
|
||||
self.db = db if db is not None else PromptDatabase()
|
||||
self._cached_output_dir = None # Lazy-cached by _find_comfyui_output_dir()
|
||||
self._html_cache = {} # Cached HTML file contents keyed by path
|
||||
self._gallery_cache = {} # dict: path_str -> (files, timestamp)
|
||||
@@ -391,6 +393,7 @@ class PromptManagerAPI(
|
||||
|
||||
# ── Register domain-specific routes from mixins ───────────────
|
||||
|
||||
self._register_pair_routes(routes)
|
||||
self._register_prompt_routes(routes)
|
||||
self._register_image_routes(routes)
|
||||
self._register_admin_routes(routes)
|
||||
|
||||
@@ -26,8 +26,12 @@ except ImportError:
|
||||
|
||||
try:
|
||||
from ...database.operations import PromptDatabase
|
||||
from ...database.pair_ingestion import ingest_resolved_pair
|
||||
from ...utils.prompt_pair_resolver import resolve_prompt_pair
|
||||
except ImportError:
|
||||
from database.operations import PromptDatabase
|
||||
from database.pair_ingestion import ingest_resolved_pair
|
||||
from utils.prompt_pair_resolver import resolve_prompt_pair
|
||||
|
||||
|
||||
try:
|
||||
@@ -40,6 +44,20 @@ VIDEO_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv", ".webm", ".m4v", ".wmv")
|
||||
SCAN_BATCH_SIZE = 50 # Files per executor call during the output scan
|
||||
MAX_SAFETY_BACKUPS = 10 # '<db>.backup_*' files kept next to the live database
|
||||
MAX_SCAN_FILES = 50_000 # Upper bound on media files one scan will look at
|
||||
SCAN_COUNT_KEYS = (
|
||||
"processed",
|
||||
"found",
|
||||
"added",
|
||||
"linked",
|
||||
"enriched",
|
||||
"unchanged",
|
||||
"ambiguous",
|
||||
"unresolved",
|
||||
"conflict",
|
||||
"failed",
|
||||
"pairs_added",
|
||||
"pair_links_added",
|
||||
)
|
||||
MAX_SCAN_DEPTH = 12 # Directory levels below an output root that are scanned
|
||||
# Files the duplicate cleaner may delete: everything the scan reports, plus
|
||||
# the common single-f TIFF spelling. Anything else inside the output
|
||||
@@ -359,6 +377,10 @@ class AdminRoutesMixin:
|
||||
async def scan_images_route(request):
|
||||
return await self.scan_images(request)
|
||||
|
||||
@routes.post("/prompt_manager/scan/cancel")
|
||||
async def cancel_scan_route(request):
|
||||
return await self.cancel_scan(request)
|
||||
|
||||
@routes.get("/prompt_manager/scan/status")
|
||||
async def scan_status_route(request):
|
||||
return await self.scan_status(request)
|
||||
@@ -1280,6 +1302,14 @@ class AdminRoutesMixin:
|
||||
job = self._start_scan_job()
|
||||
return await self._stream_scan_job(request, job)
|
||||
|
||||
async def cancel_scan(self, request):
|
||||
"""Stop between file transactions; an in-flight extraction can finish."""
|
||||
job = self._scan_job
|
||||
requested = bool(job and job.running)
|
||||
if requested:
|
||||
job.cancelled = True
|
||||
return web.json_response({"success": True, "cancellation_requested": requested})
|
||||
|
||||
async def scan_status(self, request):
|
||||
"""Whether an output scan is running, and its latest progress event."""
|
||||
job = self._scan_job
|
||||
@@ -1351,7 +1381,7 @@ class AdminRoutesMixin:
|
||||
media_files = await self._run_in_executor(
|
||||
self._collect_output_media_sync, output_dirs
|
||||
)
|
||||
counts = {"processed": 0, "found": 0, "added": 0, "linked": 0}
|
||||
counts = dict.fromkeys(SCAN_COUNT_KEYS, 0)
|
||||
if media_files:
|
||||
yield self._scan_progress(
|
||||
5, f"Found {len(media_files)} media files to process...", counts
|
||||
@@ -1364,7 +1394,9 @@ class AdminRoutesMixin:
|
||||
f"found={counts['found']}, new_prompts_added={counts['added']}, "
|
||||
f"images_linked_to_existing={counts['linked']}"
|
||||
)
|
||||
yield {"type": "complete", **counts, "directories": dir_names}
|
||||
job = self._scan_job
|
||||
final_type = "cancelled" if job and job.cancelled else "complete"
|
||||
yield {"type": final_type, **counts, "directories": dir_names}
|
||||
|
||||
except Exception:
|
||||
self.logger.exception("Scan error")
|
||||
@@ -1380,6 +1412,8 @@ class AdminRoutesMixin:
|
||||
"""
|
||||
total = len(media_files)
|
||||
for batch_start in range(0, total, SCAN_BATCH_SIZE):
|
||||
if self._scan_job and self._scan_job.cancelled:
|
||||
return
|
||||
batch = media_files[batch_start : batch_start + SCAN_BATCH_SIZE]
|
||||
batch_results = await self._run_in_executor(
|
||||
self._extract_batch_metadata_sync, batch
|
||||
@@ -1405,7 +1439,7 @@ class AdminRoutesMixin:
|
||||
"progress": progress,
|
||||
"status": status,
|
||||
"processed": counts.get("processed", 0),
|
||||
"found": counts.get("found", 0),
|
||||
**{key: counts.get(key, 0) for key in SCAN_COUNT_KEYS},
|
||||
}
|
||||
|
||||
def _collect_output_media_sync(self, output_dirs):
|
||||
@@ -1431,25 +1465,64 @@ class AdminRoutesMixin:
|
||||
|
||||
def _ingest_scan_batch_sync(self, batch_results):
|
||||
"""Store the prompts found in a metadata batch (blocking DB work)."""
|
||||
counts = {"processed": 0, "found": 0, "added": 0, "linked": 0}
|
||||
counts = dict.fromkeys(SCAN_COUNT_KEYS, 0)
|
||||
for media_file, metadata in batch_results:
|
||||
if self._scan_job and self._scan_job.cancelled:
|
||||
break
|
||||
counts["processed"] += 1
|
||||
try:
|
||||
outcome = self._ingest_scanned_file(media_file, metadata)
|
||||
outcome = self._ingest_scanned_file(media_file, metadata, detailed=True)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Error processing {media_file.name}: {e}")
|
||||
counts["failed"] += 1
|
||||
continue
|
||||
if outcome is None:
|
||||
continue
|
||||
counts["found"] += 1
|
||||
if outcome in ("added", "linked"):
|
||||
counts[outcome] += 1
|
||||
status = outcome["status"] if isinstance(outcome, dict) else outcome
|
||||
if status in (
|
||||
"added",
|
||||
"linked",
|
||||
"enriched",
|
||||
"unchanged",
|
||||
"ambiguous",
|
||||
"unresolved",
|
||||
"conflict",
|
||||
"failed",
|
||||
):
|
||||
counts[status] += 1
|
||||
if isinstance(outcome, dict):
|
||||
counts["pairs_added"] += outcome.get("pairs_added", 0)
|
||||
counts["pair_links_added"] += outcome.get("links_added", 0)
|
||||
return counts
|
||||
|
||||
def _ingest_scanned_file(self, media_file, metadata):
|
||||
"""Save or link one scanned file; returns 'added', 'linked', 'found' or None."""
|
||||
def _ingest_scanned_file(self, media_file, metadata, detailed=False):
|
||||
"""Enrich exact pairs; retain legacy positive-only import for absent pairs."""
|
||||
if not metadata:
|
||||
return None
|
||||
canonical = dict(metadata)
|
||||
for key in ("prompt", "workflow"):
|
||||
if isinstance(canonical.get(key), str):
|
||||
try:
|
||||
canonical[key] = json.loads(canonical[key])
|
||||
except (ValueError, TypeError):
|
||||
# Preserve malformed metadata for conservative resolver rejection.
|
||||
pass
|
||||
if not canonical.get("prompt") and not canonical.get("parameters"):
|
||||
return None
|
||||
resolution = resolve_prompt_pair(canonical)
|
||||
if resolution["state"] == "resolved":
|
||||
result = ingest_resolved_pair(
|
||||
self.db, str(media_file), canonical, resolution
|
||||
)
|
||||
return result if detailed else result["status"]
|
||||
if resolution["state"] == "ambiguous":
|
||||
return "ambiguous"
|
||||
if resolution["state"] == "unresolved":
|
||||
if canonical.get("prompt") or not isinstance(
|
||||
canonical.get("parameters"), str
|
||||
):
|
||||
return resolution["state"]
|
||||
parsed = self._parse_comfyui_prompt(metadata)
|
||||
if not (parsed.get("prompt") or parsed.get("parameters")):
|
||||
return None
|
||||
@@ -1465,13 +1538,15 @@ class AdminRoutesMixin:
|
||||
existing = self.db.get_prompt_by_hash(prompt_hash)
|
||||
if existing:
|
||||
try:
|
||||
self.db.link_image_to_prompt(existing["id"], str(media_file))
|
||||
image_id = self.db.link_image_to_prompt(
|
||||
existing["id"], str(media_file), metadata=canonical
|
||||
)
|
||||
except Exception as e:
|
||||
self.logger.error(
|
||||
f"Failed to link {media_file.name} to existing prompt: {e}"
|
||||
)
|
||||
return "found"
|
||||
return "linked"
|
||||
return "linked" if image_id else "unchanged"
|
||||
|
||||
prompt_id = self.db.save_prompt(
|
||||
prompt_text,
|
||||
@@ -1484,7 +1559,7 @@ class AdminRoutesMixin:
|
||||
if not prompt_id:
|
||||
return "found"
|
||||
try:
|
||||
self.db.link_image_to_prompt(prompt_id, str(media_file))
|
||||
self.db.link_image_to_prompt(prompt_id, str(media_file), metadata=canonical)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to link {media_file.name} to new prompt: {e}")
|
||||
return "added"
|
||||
|
||||
@@ -0,0 +1,204 @@
|
||||
"""Reusable ordered prompt combinations, separate from legacy prompt CRUD."""
|
||||
|
||||
from functools import wraps
|
||||
from typing import Callable
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
from .prompts import parse_page_params
|
||||
|
||||
try:
|
||||
from ...database.pairs import PairDependencyError
|
||||
from ...utils.validators import (
|
||||
validate_category,
|
||||
validate_prompt_text,
|
||||
validate_rating,
|
||||
validate_tags,
|
||||
)
|
||||
except ImportError:
|
||||
from database.pairs import PairDependencyError
|
||||
from utils.validators import (
|
||||
validate_category,
|
||||
validate_prompt_text,
|
||||
validate_rating,
|
||||
validate_tags,
|
||||
)
|
||||
|
||||
|
||||
def pair_response(handler: Callable) -> Callable:
|
||||
@wraps(handler)
|
||||
async def wrapped(self, request: web.Request) -> web.Response:
|
||||
try:
|
||||
return await handler(self, request)
|
||||
except PairDependencyError as exc:
|
||||
return web.json_response({"success": False, "error": str(exc)}, status=409)
|
||||
except (ValueError, TypeError) as exc:
|
||||
return web.json_response({"success": False, "error": str(exc)}, status=400)
|
||||
except Exception:
|
||||
self.logger.exception("Prompt pair operation failed")
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Prompt pair operation failed"}, status=500
|
||||
)
|
||||
|
||||
return wrapped
|
||||
|
||||
|
||||
def metadata_fields(data: dict) -> dict:
|
||||
fields = {}
|
||||
for name in ("title", "category", "tags", "rating", "notes"):
|
||||
if name not in data:
|
||||
continue
|
||||
value = data[name]
|
||||
if name == "tags":
|
||||
if not isinstance(value, list):
|
||||
raise ValueError("Tags must be a list of strings")
|
||||
validate_tags(value)
|
||||
elif name == "rating":
|
||||
if isinstance(value, bool):
|
||||
raise ValueError("Rating must be an integer")
|
||||
validate_rating(value)
|
||||
else:
|
||||
if value is not None and not isinstance(value, str):
|
||||
raise ValueError(f"{name} must be a string")
|
||||
if name == "category":
|
||||
validate_category(value)
|
||||
elif value is not None and len(value) > 10000:
|
||||
raise ValueError(f"{name} is too long")
|
||||
fields[name] = value
|
||||
return fields
|
||||
|
||||
|
||||
async def object_body(request: web.Request) -> dict:
|
||||
data = await request.json()
|
||||
if not isinstance(data, dict):
|
||||
raise ValueError("Expected a JSON object")
|
||||
return data
|
||||
|
||||
|
||||
def missing_pair() -> web.Response:
|
||||
return web.json_response({"success": False, "error": "Pair not found"}, status=404)
|
||||
|
||||
|
||||
class PairRoutesMixin:
|
||||
def _register_pair_routes(self, routes: web.RouteTableDef) -> None:
|
||||
routes.get("/prompt_manager/pairs/export")(self.export_prompt_pairs)
|
||||
routes.get("/prompt_manager/pairs")(self.list_prompt_pairs)
|
||||
routes.post("/prompt_manager/pairs")(self.save_prompt_pair)
|
||||
routes.get("/prompt_manager/pairs/{pair_id}")(self.get_prompt_pair)
|
||||
routes.patch("/prompt_manager/pairs/{pair_id}")(self.update_prompt_pair)
|
||||
routes.delete("/prompt_manager/pairs/{pair_id}")(self.delete_prompt_pair)
|
||||
|
||||
@pair_response
|
||||
async def save_prompt_pair(self, request: web.Request) -> web.Response:
|
||||
data = await object_body(request)
|
||||
positive = data.get("positive_text")
|
||||
negative = data.get("negative_text")
|
||||
validate_prompt_text(positive)
|
||||
if not isinstance(negative, str):
|
||||
raise ValueError(
|
||||
"negative_text must be explicit; use an empty string for none"
|
||||
)
|
||||
if negative.strip():
|
||||
validate_prompt_text(negative)
|
||||
fields = metadata_fields(data)
|
||||
pair_id = await self._run_in_executor(
|
||||
self.db.save_prompt_pair, positive, negative, **fields
|
||||
)
|
||||
pair = await self._run_in_executor(self.db.get_prompt_pair, pair_id)
|
||||
return web.json_response({"success": True, "pair": pair})
|
||||
|
||||
@pair_response
|
||||
async def get_prompt_pair(self, request: web.Request) -> web.Response:
|
||||
pair = await self._run_in_executor(
|
||||
self.db.get_prompt_pair, int(request.match_info["pair_id"])
|
||||
)
|
||||
return (
|
||||
web.json_response({"success": True, "pair": pair})
|
||||
if pair
|
||||
else missing_pair()
|
||||
)
|
||||
|
||||
@pair_response
|
||||
async def list_prompt_pairs(self, request: web.Request) -> web.Response:
|
||||
limit, offset = parse_page_params(request.query)
|
||||
archived = request.query.get("include_archived", "false")
|
||||
if archived not in ("true", "false"):
|
||||
raise ValueError("include_archived must be true or false")
|
||||
result = await self._run_in_executor(
|
||||
self.db.list_prompt_pairs,
|
||||
text=request.query.get("text"),
|
||||
category=request.query.get("category"),
|
||||
tags=request.query.getall("tag", []),
|
||||
include_archived=archived == "true",
|
||||
sort=request.query.get("sort", "newest"),
|
||||
limit=limit,
|
||||
offset=offset,
|
||||
)
|
||||
return web.json_response({"success": True, **result})
|
||||
|
||||
@pair_response
|
||||
async def update_prompt_pair(self, request: web.Request) -> web.Response:
|
||||
pair_id = int(request.match_info["pair_id"])
|
||||
data = await object_body(request)
|
||||
allowed = {"title", "category", "tags", "rating", "notes", "archived"}
|
||||
if not data or set(data) - allowed:
|
||||
raise ValueError(
|
||||
"Update metadata or archive; save a new pair to change text"
|
||||
)
|
||||
fields = metadata_fields(data)
|
||||
if "archived" in data:
|
||||
if not isinstance(data["archived"], bool) or fields:
|
||||
raise ValueError("Archive must be a separate boolean update")
|
||||
updated = await self._run_in_executor(
|
||||
self.db.archive_prompt_pair, pair_id, data["archived"]
|
||||
)
|
||||
else:
|
||||
updated = await self._run_in_executor(
|
||||
self.db.update_prompt_pair_metadata, pair_id, **fields
|
||||
)
|
||||
pair = await self._run_in_executor(self.db.get_prompt_pair, pair_id)
|
||||
if not updated and not pair:
|
||||
return missing_pair()
|
||||
return web.json_response({"success": True, "pair": pair})
|
||||
|
||||
@pair_response
|
||||
async def delete_prompt_pair(self, request: web.Request) -> web.Response:
|
||||
deleted = await self._run_in_executor(
|
||||
self.db.delete_prompt_pair, int(request.match_info["pair_id"])
|
||||
)
|
||||
return web.json_response({"success": True}) if deleted else missing_pair()
|
||||
|
||||
@pair_response
|
||||
async def export_prompt_pairs(self, request: web.Request) -> web.Response:
|
||||
pairs = []
|
||||
offset = 0
|
||||
while True:
|
||||
result = await self._run_in_executor(
|
||||
self.db.list_prompt_pairs,
|
||||
include_archived=True,
|
||||
limit=500,
|
||||
offset=offset,
|
||||
)
|
||||
for pair in result["pairs"]:
|
||||
pairs.append(
|
||||
{
|
||||
key: pair[key]
|
||||
for key in (
|
||||
"positive_text",
|
||||
"negative_text",
|
||||
"title",
|
||||
"category",
|
||||
"tags",
|
||||
"rating",
|
||||
"notes",
|
||||
"archived",
|
||||
)
|
||||
}
|
||||
)
|
||||
offset += len(result["pairs"])
|
||||
if not result["has_more"] or not result["pairs"]:
|
||||
break
|
||||
return web.json_response(
|
||||
{"schema_version": 1, "kind": "prompt-manager-pairs", "pairs": pairs},
|
||||
headers={"Content-Disposition": 'attachment; filename="prompt-pairs.json"'},
|
||||
)
|
||||
@@ -6,6 +6,11 @@ import os
|
||||
import sqlite3
|
||||
|
||||
from aiohttp import web
|
||||
|
||||
try:
|
||||
from ...database.pairs import PairDependencyError
|
||||
except ImportError:
|
||||
from database.pairs import PairDependencyError
|
||||
from urllib.parse import quote
|
||||
|
||||
# Upper bounds on page sizes and offsets for every list endpoint. The offset
|
||||
@@ -793,6 +798,8 @@ class PromptRoutesMixin:
|
||||
status=404,
|
||||
)
|
||||
|
||||
except PairDependencyError as exc:
|
||||
return web.json_response({"success": False, "error": str(exc)}, status=409)
|
||||
except ValueError:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Invalid prompt ID"}, status=400
|
||||
@@ -840,6 +847,8 @@ class PromptRoutesMixin:
|
||||
{"success": False, "error": "Prompt not found"}, status=404
|
||||
)
|
||||
|
||||
except PairDependencyError as exc:
|
||||
return web.json_response({"success": False, "error": str(exc)}, status=409)
|
||||
except sqlite3.IntegrityError:
|
||||
return web.json_response(
|
||||
{"success": False, "error": "Another prompt already has this text"},
|
||||
@@ -1081,6 +1090,8 @@ class PromptRoutesMixin:
|
||||
}
|
||||
)
|
||||
|
||||
except PairDependencyError as exc:
|
||||
return web.json_response({"success": False, "error": str(exc)}, status=409)
|
||||
except Exception as e:
|
||||
self.logger.error(f"Bulk delete error: {e}")
|
||||
return web.json_response(
|
||||
|
||||
@@ -16,6 +16,7 @@ class ScanJob:
|
||||
def __init__(self):
|
||||
self.last_event = None
|
||||
self.finished = False
|
||||
self.cancelled = False
|
||||
self.task = None
|
||||
self._subscribers = set()
|
||||
|
||||
|
||||
@@ -0,0 +1,186 @@
|
||||
import {readFileSync} from 'node:fs';
|
||||
import {test,expect} from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
|
||||
import {startCoverage,saveCoverage} from '../../tools/vendor/browser-coverage.mjs';
|
||||
|
||||
const pair={id:4,positive_text:'Library positive',negative_text:'Library negative',category:'Portrait',tags:['studio']};
|
||||
test.beforeEach(async({page,browserName})=>{
|
||||
await startCoverage(page,browserName);
|
||||
page.pairErrors=[];page.on('pageerror',error=>page.pairErrors.push(error.message));
|
||||
await page.route('**/scripts/app.js',route=>route.fulfill({contentType:'text/javascript',body:'export const app=window.fixtureApp;'}));
|
||||
await page.route('**/scripts/api.js',route=>route.fulfill({contentType:'text/javascript',body:'export const api={fileURL:path=>path,queuePrompt:async(...args)=>args};'}));
|
||||
await page.route('**/extensions/promptmanager/prompt_manager.js',route=>route.fulfill({contentType:'text/javascript',body:readFileSync(new URL('../../web/comfy/prompt_manager.js',import.meta.url),'utf8')}));
|
||||
await page.route('**/prompt_manager/pairs?**',route=>route.fulfill({json:{success:true,pairs:[pair],total:1}}));
|
||||
await page.goto('/prompt_manager/test');
|
||||
await page.evaluate(async()=>{
|
||||
document.body.replaceChildren();
|
||||
const settings=new EventTarget();settings.getSettingValue=()=> 'dark';
|
||||
window.fixtureApp={registerExtension(extension){this.extension=extension;},ui:{settings},graph:{_nodes:[],setDirtyCanvas(){}},canvas:{},graphToPrompt:async()=>({})};
|
||||
await import('/extensions/promptmanager/prompt_manager.js');
|
||||
class Node {
|
||||
constructor(){this.id=1;this.type='PromptManagerPair';this.widgets=['text','negative_text','category','tags','search_text'].map(name=>({name,type:'text',value:'',options:{},inputEl:document.createElement('textarea')}));this.inputs=['clip','prepend_text','append_text'].map(name=>({name,link:null}));this.properties={};this.size=[1100,700];}
|
||||
addWidget(type,name,value,callback,options){const widget={type,name,value,callback,options};this.widgets.push(widget);return widget;}
|
||||
addDOMWidget(name,type,element){const host=document.createElement('div');host.style.cssText='height:700px;width:1100px;max-width:100%';host.className='fixture-dom-widget';host.append(element);document.body.append(host);return {element};}
|
||||
setSize(size){this.size=size;}setDirtyCanvas(){}
|
||||
serialize(){return {widgets_values:this.widgets.map(widget=>widget.value),properties:{...this.properties}};}
|
||||
configure(data){this.widgets.forEach((widget,index)=>widget.value=data.widgets_values[index]);this.properties={...data.properties};}
|
||||
}
|
||||
await fixtureApp.extension.beforeRegisterNodeDef(Node,{name:'PromptManagerPair'},fixtureApp);
|
||||
window.node=new Node();node.onNodeCreated?.();fixtureApp.graph._nodes=[node];
|
||||
});
|
||||
});
|
||||
test.afterEach(async({page,browserName})=>{await saveCoverage(page,browserName);expect(page.pairErrors).toEqual([]);});
|
||||
|
||||
async function selectPair(page){await page.getByRole('button',{name:'Recent',exact:true}).click();await page.locator('.pm-editor-result').click();}
|
||||
async function draft(page,positive='Draft positive',negative='Draft negative'){await page.getByRole('textbox',{name:'Positive prompt',exact:true}).fill(positive);await page.getByRole('textbox',{name:'Negative prompt',exact:true}).fill(negative);}
|
||||
|
||||
test('paired editor previews both texts before explicitly replacing both',async({page})=>{
|
||||
await expect(page.locator('.pm-pair')).toBeVisible();
|
||||
await draft(page);
|
||||
await selectPair(page);
|
||||
await expect(page.locator('.pm-editor-selection')).toContainText('Library positive');
|
||||
await expect(page.locator('.pm-editor-selection')).toContainText('Library negative');
|
||||
await expect(page.getByRole('textbox',{name:'Positive prompt',exact:true})).toHaveValue('Draft positive');
|
||||
await expect(page.getByRole('textbox',{name:'Negative prompt',exact:true})).toHaveValue('Draft negative');
|
||||
await page.getByRole('switch',{name:'Load metadata'}).uncheck();
|
||||
await page.getByRole('button',{name:'Use pair',exact:true}).click();
|
||||
await expect(page.getByRole('textbox',{name:'Positive prompt',exact:true})).toHaveValue(pair.positive_text);
|
||||
await expect(page.getByRole('textbox',{name:'Negative prompt',exact:true})).toHaveValue(pair.negative_text);
|
||||
expect(await page.evaluate(()=>Object.fromEntries(node.widgets.map(w=>[w.name,w.value])))).toMatchObject({text:pair.positive_text,negative_text:pair.negative_text,category:'',tags:''});
|
||||
await page.evaluate(()=>{const saved=node.serialize();node.configure(saved);});
|
||||
await expect(page.getByRole('textbox',{name:'Negative prompt',exact:true})).toHaveValue(pair.negative_text);
|
||||
});
|
||||
|
||||
test('paired editor sizes both editors equally and composes positive text only',async({page},testInfo)=>{
|
||||
await draft(page);
|
||||
const positive=page.getByRole('textbox',{name:'Positive prompt',exact:true});const negative=page.getByRole('textbox',{name:'Negative prompt',exact:true});
|
||||
const [p,n]=await Promise.all([positive.boundingBox(),negative.boundingBox()]);expect(Math.abs(p.height-n.height)).toBeLessThan(1);expect(Math.abs(p.width-n.width)).toBeLessThan(1);
|
||||
await page.getByText('Composition · prepend / append',{exact:true}).click();
|
||||
await page.getByRole('textbox',{name:'Prepend',exact:true}).fill('trigger');
|
||||
await page.getByRole('textbox',{name:'Append',exact:true}).fill('detail');
|
||||
await page.getByText('Composed text preview',{exact:true}).click();
|
||||
await expect(page.locator('[data-composed-preview]')).toHaveText('trigger Draft positive detail');
|
||||
await expect(negative).toHaveValue('Draft negative');
|
||||
await page.screenshot({path:testInfo.outputPath('pair-editor.png')});
|
||||
});
|
||||
|
||||
for(const field of ['text','negative_text']) test(`paired loading is atomic when ${field} is connected`,async({page})=>{
|
||||
await draft(page);
|
||||
await selectPair(page);
|
||||
await page.evaluate(field=>{node.inputs.push({name:field,link:8});node.onConnectionsChange();},field);
|
||||
await expect(page.getByRole('button',{name:'Use pair',exact:true})).toBeDisabled();
|
||||
await page.getByRole('button',{name:'Use pair',exact:true}).evaluate(button=>button.dispatchEvent(new Event('click',{bubbles:true})));
|
||||
expect(await page.evaluate(()=>Object.fromEntries(node.widgets.map(w=>[w.name,w.value])))).toMatchObject({text:'Draft positive',negative_text:'Draft negative'});
|
||||
await page.evaluate(()=>{node.inputs=node.inputs.filter(input=>input.link!==8);node.onConnectionsChange();});
|
||||
await expect(page.getByRole('button',{name:'Use pair',exact:true})).toBeEnabled();
|
||||
});
|
||||
|
||||
test('save pair accepts empty negatives, recovers after errors and never clears the draft',async({page})=>{
|
||||
let attempts=0;const bodies=[];
|
||||
await page.route('**/prompt_manager/pairs',route=>{bodies.push(route.request().postDataJSON());return ++attempts===1?route.fulfill({status:503,json:{success:false,error:'Temporary save error'}}):route.fulfill({json:{success:true,pair:{...pair,negative_text:''}}});});
|
||||
await draft(page,'Save this','');
|
||||
await page.getByRole('button',{name:'Save pair',exact:true}).click();
|
||||
await expect(page.getByRole('alert')).toContainText('Temporary save error');
|
||||
await expect(page.getByRole('textbox',{name:'Positive prompt',exact:true})).toHaveValue('Save this');
|
||||
await page.getByRole('button',{name:'Save pair',exact:true}).click();
|
||||
await expect(page.getByRole('status').filter({hasText:'Pair saved'})).toBeVisible();
|
||||
expect(bodies).toEqual([{positive_text:'Save this',negative_text:'',category:'',tags:[]},{positive_text:'Save this',negative_text:'',category:'',tags:[]}]);
|
||||
});
|
||||
|
||||
test('pair filters use repeated tags and removal cleans up',async({page})=>{
|
||||
const urls=[];await page.route('**/prompt_manager/pairs?**',route=>{urls.push(new URL(route.request().url()));return route.fulfill({json:{success:true,pairs:[pair]}});});
|
||||
await page.getByRole('searchbox',{name:'Search saved prompts'}).fill('portrait tag:studio');
|
||||
await page.getByLabel('Filter category',{exact:true}).fill('Portrait');
|
||||
await page.getByLabel('Filter tags',{exact:true}).fill('film, studio, ');
|
||||
await page.getByRole('button',{name:'Search',exact:true}).click();await expect(page.locator('.pm-editor-result')).toBeVisible();
|
||||
expect(urls[0].searchParams.get('category')).toBe('Portrait');expect(urls[0].searchParams.getAll('tag')).toEqual(['studio','film']);
|
||||
await page.evaluate(()=>node.onRemoved());await expect(page.locator('.ca-completions')).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('validating and saving prevents duplicate submissions while preserving pair metadata and drafts',async({page})=>{
|
||||
let release;let requests=0;
|
||||
await page.route('**/prompt_manager/pairs',async route=>{requests++;await new Promise(resolve=>release=resolve);await route.fulfill({json:{success:true,pair:{...pair,category:'Original category',tags:['original-tag']}}});});
|
||||
await page.getByRole('button',{name:'Save pair',exact:true}).click();
|
||||
await expect(page.getByRole('alert')).toContainText('positive');expect(requests).toBe(0);
|
||||
await draft(page);
|
||||
await page.getByRole('button',{name:'Save pair',exact:true}).click();
|
||||
await expect(page.getByRole('button',{name:'Save pair',exact:true})).toBeDisabled();
|
||||
await page.getByRole('button',{name:'Save pair',exact:true}).evaluate(button=>button.dispatchEvent(new Event('click',{bubbles:true})));
|
||||
expect(requests).toBe(1);release();
|
||||
await expect(page.getByRole('status').filter({hasText:'Pair saved'})).toBeVisible();
|
||||
expect(await page.evaluate(()=>Object.fromEntries(node.widgets.map(w=>[w.name,w.value])))).toMatchObject({text:'Draft positive',negative_text:'Draft negative',category:'',tags:''});
|
||||
});
|
||||
|
||||
test('a superseded library response never replaces the latest pair selection',async({page})=>{
|
||||
let release;let started=false;
|
||||
await page.route('**/prompt_manager/pairs?**',async route=>{
|
||||
if(new URL(route.request().url()).searchParams.get('text')==='old'){started=true;await new Promise(resolve=>release=resolve);await route.fulfill({json:{success:true,pairs:[{...pair,positive_text:'Stale positive'}]}}).catch(()=>{});}
|
||||
else await route.fulfill({json:{success:true,pairs:[pair]}});
|
||||
});
|
||||
const search=page.getByRole('searchbox',{name:'Search saved prompts'});
|
||||
await search.fill('old');await page.getByRole('button',{name:'Search',exact:true}).click();await expect.poll(()=>started).toBe(true);
|
||||
await search.fill('new');await page.getByRole('button',{name:'Search',exact:true}).click();
|
||||
await expect(page.locator('.pm-editor-result')).toContainText('Library positive');release();
|
||||
await expect(page.locator('.pm-editor-result')).not.toContainText('Stale positive');
|
||||
});
|
||||
|
||||
test('queue hook synchronizes both editors, preserves links and forwards branch execution arguments',async({page})=>{
|
||||
await draft(page);
|
||||
const args=await page.evaluate(async()=>{
|
||||
await fixtureApp.extension.setup();
|
||||
const {api}=await import('/scripts/api.js');
|
||||
const output={1:{inputs:{text:'stale',negative_text:'stale'}}};
|
||||
const first=await api.queuePrompt(7,{output,workflow:{name:'workflow'}},['branch'],{preview:'latent'});
|
||||
node.inputs.push({name:'negative_text',link:17});
|
||||
const linked={1:{inputs:{text:'old',negative_text:['17',0]}}};
|
||||
const second=await api.queuePrompt(8,{output:linked,workflow:{}},['another']);
|
||||
return {first,second};
|
||||
});
|
||||
expect(args.first).toEqual([7,{output:{1:{inputs:{text:'Draft positive',negative_text:'Draft negative'}}},workflow:{name:'workflow'}},['branch'],{preview:'latent'}]);
|
||||
expect(args.second[1].output[1].inputs).toEqual({text:'Draft positive',negative_text:['17',0]});
|
||||
});
|
||||
|
||||
test('paired editors stack at narrow widths while keeping equal editing areas',async({page})=>{
|
||||
await page.locator('.fixture-dom-widget').evaluate(host=>host.style.width='560px');
|
||||
const [p,n]=await Promise.all(['Positive prompt','Negative prompt'].map(name=>page.getByRole('textbox',{name,exact:true}).boundingBox()));expect(n.y).toBeGreaterThanOrEqual(p.y+p.height);expect(Math.abs(p.width-n.width)).toBeLessThan(1);expect(Math.abs(p.height-n.height)).toBeLessThan(1);
|
||||
});
|
||||
|
||||
test('removal aborts an in-flight pair save and restores host-owned textarea references',async({page})=>{
|
||||
await draft(page);
|
||||
const state=await page.evaluate(async()=>{
|
||||
const previous=window.fetch;let signal;
|
||||
window.fetch=(url,options)=>{if(String(url).endsWith('/pairs')){signal=options.signal;return new Promise((resolve,reject)=>signal.addEventListener('abort',()=>reject(new DOMException('Aborted','AbortError')),{once:true}));}return previous(url,options);};
|
||||
const visible=node.widgets.find(w=>w.name==='negative_text').inputEl;
|
||||
[...document.querySelectorAll('button')].find(button=>button.textContent==='Save pair').click();
|
||||
node.onRemoved();
|
||||
return {aborted:signal.aborted,restored:node.widgets.find(w=>w.name==='negative_text').inputEl!==visible};
|
||||
});
|
||||
expect(state).toEqual({aborted:true,restored:true});
|
||||
});
|
||||
|
||||
test('actual pair saves coexist with resetting the disposable browser fixture',async({page,request})=>{
|
||||
await draft(page,'Actual browser pair','Actual negative');await page.getByRole('button',{name:'Save pair',exact:true}).click();await expect(page.getByRole('status').filter({hasText:'Pair saved'})).toBeVisible();
|
||||
const response=await request.get('/prompt_manager/pairs?text=Actual%20browser%20pair');expect(response.ok()).toBe(true);expect((await response.json()).pairs).toEqual(expect.arrayContaining([expect.objectContaining({positive_text:'Actual browser pair',negative_text:'Actual negative'})]));
|
||||
const reset=await request.post('/__test__/reset');expect(reset.ok()).toBe(true);
|
||||
const after=await request.get('/prompt_manager/pairs?text=Actual%20browser%20pair');expect((await after.json()).pairs).toHaveLength(0);
|
||||
});
|
||||
|
||||
for(const field of ['category','tags'])test(`manual pair save waits for execution when ${field} comes from an upstream link`,async({page})=>{
|
||||
let requests=0;await page.route('**/prompt_manager/pairs',route=>{requests++;return route.fulfill({json:{success:true,pair}});});await draft(page);
|
||||
await page.evaluate(field=>{node.inputs.push({name:field,link:19});node.onConnectionsChange();},field);
|
||||
const save=page.getByRole('button',{name:'Save pair',exact:true});await expect(save).toBeDisabled();await expect(page.getByRole('status').filter({hasText:'Run the node to save connected values.'})).toBeVisible();
|
||||
await save.evaluate(button=>button.dispatchEvent(new Event('click',{bubbles:true})));expect(requests).toBe(0);
|
||||
await page.evaluate(()=>{node.inputs=node.inputs.filter(input=>input.link!==19);node.onConnectionsChange();});await expect(save).toBeEnabled();await save.click();await expect.poll(()=>requests).toBe(1);
|
||||
});
|
||||
|
||||
for(const field of ['positive','negative','category','tags'])test(`save completion identifies an earlier draft after editing ${field}`,async({page})=>{
|
||||
let release;await page.route('**/prompt_manager/pairs',async route=>{await new Promise(resolve=>release=resolve);await route.fulfill({json:{success:true,pair}});});await draft(page);
|
||||
await page.getByText('Metadata · category / tags',{exact:true}).click();await page.getByRole('button',{name:'Save pair',exact:true}).click();await expect.poll(()=>Boolean(release)).toBe(true);
|
||||
const input=field==='positive'?page.getByRole('textbox',{name:'Positive prompt',exact:true}):field==='negative'?page.getByRole('textbox',{name:'Negative prompt',exact:true}):page.getByRole('combobox',{name:field==='category'?'Category':'Tags',exact:true});
|
||||
await input.fill(`Changed ${field}`);release();await expect(page.getByRole('status').filter({hasText:'Earlier draft saved. Current changes are not saved yet.'})).toBeVisible();await expect(input).toHaveValue(`Changed ${field}`);
|
||||
});
|
||||
|
||||
test('pair library rows distinguish negative variants and explicitly show empty negatives as text',async({page})=>{
|
||||
const longNegative='Variant one '+ 'x'.repeat(200);
|
||||
await page.route('**/prompt_manager/pairs?**',route=>route.fulfill({json:{success:true,pairs:[{...pair,negative_text:longNegative},{...pair,id:5,negative_text:'Variant two'},{...pair,id:6,negative_text:''},{...pair,id:7,negative_text:'<img src=x onerror=alert(1)>'}]}}));
|
||||
await page.getByRole('button',{name:'Recent',exact:true}).click();const rows=page.locator('.pm-editor-result');await expect(rows).toHaveCount(4);await expect(rows.nth(0)).toContainText('Negative: Variant one');await expect(rows.nth(1)).toContainText('Negative: Variant two');await expect(rows.nth(2)).toContainText('Negative: (empty)');await expect(rows.nth(3)).toContainText('Negative: <img src=x onerror=alert(1)>');await expect(rows.locator('img')).toHaveCount(0);await expect(rows.nth(0)).not.toContainText('x'.repeat(180));await expect(rows.nth(0).locator('span').nth(1)).toHaveAttribute('title',longNegative);
|
||||
});
|
||||
@@ -80,7 +80,15 @@ def create_app(root: Path) -> web.Application:
|
||||
|
||||
def seed(extra_images: int = 0) -> None:
|
||||
with api.db.model.get_connection() as conn:
|
||||
for table in ("generated_images", "prompt_tags", "tags", "prompts"):
|
||||
for table in (
|
||||
"image_prompt_pairs",
|
||||
"prompt_pair_tags",
|
||||
"prompt_pairs",
|
||||
"generated_images",
|
||||
"prompt_tags",
|
||||
"tags",
|
||||
"prompts",
|
||||
):
|
||||
conn.execute(f"DELETE FROM {table}")
|
||||
for name in ("output", "input"):
|
||||
shutil.rmtree(root / name)
|
||||
|
||||
@@ -0,0 +1,83 @@
|
||||
import {test,expect} from '../../tools/vendor/node_modules/@playwright/test/index.mjs';
|
||||
import {installHostHooks} from './host-fixture.mjs';
|
||||
|
||||
installHostHooks(test,['PromptManagerPair']);
|
||||
async function addPair(page){return page.evaluate(async()=>{const {app}=await import('/scripts/app.js');const node=LiteGraph.createNode('PromptManagerPair');app.graph.add(node);app.canvas.ds.scale=.8;app.canvas.centerOnNode(node);return node.id;});}
|
||||
async function seed(request,positive,negative){const response=await request.post('/prompt_manager/pairs',{data:{positive_text:positive,negative_text:negative,category:'Paired host tests',tags:['pair-host']}});expect(response.ok()).toBe(true);return (await response.json()).pair;}
|
||||
async function values(page,id){return page.evaluate(async id=>{const {app}=await import('/scripts/app.js');return Object.fromEntries(app.graph.getNodeById(id).widgets.map(w=>[w.name,w.value]));},id);}
|
||||
async function select(page,text){const root=page.locator('.pm-pair');await root.getByRole('button',{name:'Recent',exact:true}).click();await root.locator('.pm-editor-result').filter({has:page.getByText(text,{exact:true})}).click();return root;}
|
||||
|
||||
test('paired definition exposes both conditioning/text outputs and positive-only composition',async({request})=>{
|
||||
const definitions=await(await request.get('/object_info')).json();const pair=definitions.PromptManagerPair;
|
||||
expect(pair).toBeDefined();expect(pair.output).toEqual(['CONDITIONING','CONDITIONING','STRING','STRING']);
|
||||
expect(pair.input.required.text[0]).toBe('STRING');expect(pair.input.required.negative_text[0]).toBe('STRING');expect(pair.input.required.clip[0]).toBe('CLIP');
|
||||
for(const name of ['prepend_text','append_text'])expect(pair.input.optional[name]).toMatchObject(['STRING',{default:'',forceInput:true}]);
|
||||
expect(pair.input.optional.negative_prepend_text).toBeUndefined();
|
||||
});
|
||||
|
||||
test('paired library previews both, atomically loads, queues and survives renderer switches and reload',async({page,request},testInfo)=>{
|
||||
const positive=`Pair positive ${testInfo.project.name}`;const negative=`Pair negative ${testInfo.project.name}`;await seed(request,positive,negative);const id=await addPair(page);const root=page.locator('.pm-pair');
|
||||
const pos=root.getByRole('textbox',{name:'Positive prompt',exact:true});const neg=root.getByRole('textbox',{name:'Negative prompt',exact:true});
|
||||
await pos.fill('Local positive');await neg.fill('Local negative');
|
||||
for(const [name,field]of [['text',pos],['negative_text',neg]])expect(await field.evaluate(async(element,{id,name})=>{const {app}=await import('/scripts/app.js');return app.graph.getNodeById(id).widgets.find(w=>w.name===name).inputEl===element;},{id,name})).toBe(true);
|
||||
await select(page,positive);await expect(root.locator('.pm-editor-selection')).toContainText(negative);await expect(pos).toHaveValue('Local positive');await expect(neg).toHaveValue('Local negative');
|
||||
await root.getByRole('button',{name:'Use pair',exact:true}).click();await expect(pos).toHaveValue(positive);await expect(neg).toHaveValue(negative);
|
||||
await root.getByText('Composition · prepend / append',{exact:true}).click();await root.getByRole('textbox',{name:'Prepend',exact:true}).fill('trigger');await root.getByRole('textbox',{name:'Append',exact:true}).fill('detail');
|
||||
for(const renderer of [!testInfo.project.metadata.initialRenderer,testInfo.project.metadata.initialRenderer]){
|
||||
await page.evaluate(async renderer=>{const {app}=await import('/scripts/app.js');await app.ui.settings.setSettingValue('Comfy.VueNodes.Enabled',renderer);app.canvas.centerOnNode(app.graph._nodes[0]);},renderer);
|
||||
await expect.poll(()=>page.evaluate(()=>LiteGraph.vueNodesMode)).toBe(renderer);await expect(pos).toHaveValue(positive);await expect(neg).toHaveValue(negative);
|
||||
}
|
||||
const queued=await page.evaluate(async id=>{const {app}=await import('/scripts/app.js');const inputs=(await app.graphToPrompt()).output[String(id)].inputs;await app.loadGraphData(app.graph.serialize());return inputs;},id);
|
||||
expect(queued).toMatchObject({text:positive,negative_text:negative,prepend_text:'trigger',append_text:'detail'});await expect(pos).toHaveValue(positive);await expect(neg).toHaveValue(negative);
|
||||
await page.screenshot({path:testInfo.outputPath('pair-reloaded.png')});
|
||||
await page.evaluate(async id=>{const {app}=await import('/scripts/app.js');app.graph.remove(app.graph.getNodeById(id));},id);await expect(root).toHaveCount(0);await expect(page.locator('.ca-completions')).toHaveCount(0);
|
||||
});
|
||||
|
||||
test('paired real STRING connections stay linked and restore both local editors after disconnect',async({page,request},testInfo)=>{
|
||||
const positive=`Connected pair ${testInfo.project.name}`;await seed(request,positive,'Connected library negative');const id=await addPair(page);const root=page.locator('.pm-pair');
|
||||
await root.getByRole('textbox',{name:'Positive prompt',exact:true}).fill('Local positive');await root.getByRole('textbox',{name:'Negative prompt',exact:true}).fill('Local negative');
|
||||
await root.getByText('Composition · prepend / append',{exact:true}).click();await root.getByRole('textbox',{name:'Prepend',exact:true}).fill('Local prepend');await root.getByRole('textbox',{name:'Append',exact:true}).fill('Local append');
|
||||
const links=await page.evaluate(async id=>{const {app}=await import('/scripts/app.js');const {convertToInput}=await import('/extensions/core/widgetInputs.js');const node=app.graph.getNodeById(id);for(const name of ['text','negative_text'])convertToInput(node,node.widgets.find(w=>w.name===name));const result={};for(const name of ['text','negative_text','prepend_text','append_text']){const source=LiteGraph.createNode('PrimitiveString');app.graph.add(source);source.widgets.find(w=>w.name==='value').value=`Connected ${name}`;source.connect(0,node,node.inputs.findIndex(i=>i.name===name));result[name]=[String(source.id),0];}app.canvas.centerOnNode(node);return result;},id);
|
||||
await select(page,positive);await expect(root.getByRole('button',{name:'Use pair',exact:true})).toBeDisabled();await expect(root.getByRole('button',{name:'Save pair',exact:true})).toBeDisabled();
|
||||
for(const name of ['Positive prompt','Negative prompt','Prepend','Append'])await expect(root.getByRole('textbox',{name,exact:true})).toBeHidden();
|
||||
const before=await values(page,id);await root.locator('textarea,input').evaluateAll(inputs=>{for(const input of inputs){if(!input.hidden&&!input.closest('[hidden]'))continue;input.value='Overwrite attempt';input.dispatchEvent(new Event('input',{bubbles:true}));}});expect(await values(page,id)).toMatchObject(before);
|
||||
for(const reload of [false,true]){const queued=await page.evaluate(async({id,reload})=>{const {app}=await import('/scripts/app.js');if(reload)await app.loadGraphData(app.graph.serialize());return(await app.graphToPrompt()).output[String(id)].inputs;},{id,reload});expect(queued).toMatchObject(links);}
|
||||
const local=await page.evaluate(async id=>{const {app}=await import('/scripts/app.js');const node=app.graph.getNodeById(id);for(const name of ['text','negative_text','prepend_text','append_text'])node.disconnectInput(node.inputs.findIndex(i=>i.name===name));return(await app.graphToPrompt()).output[String(id)].inputs;},id);
|
||||
expect(local).toMatchObject({text:'Local positive',negative_text:'Local negative',prepend_text:'Local prepend',append_text:'Local append'});
|
||||
for(const [name,value]of [['Positive prompt','Local positive'],['Negative prompt','Local negative'],['Prepend','Local prepend'],['Append','Local append']]){const field=root.getByRole('textbox',{name,exact:true});await expect(field).toHaveCount(1);await expect(field).toBeEditable();await expect(field).toHaveValue(value);}
|
||||
});
|
||||
|
||||
test('paired default layout uses equal editors and follows real host themes',async({page},testInfo)=>{
|
||||
await addPair(page);const root=page.locator('.pm-pair');await root.getByRole('textbox',{name:'Positive prompt',exact:true}).fill('Positive draft');await root.getByRole('textbox',{name:'Negative prompt',exact:true}).fill('Negative draft');
|
||||
for(const palette of ['light','dark']){await page.evaluate(async palette=>{const {app}=await import('/scripts/app.js');await app.ui.settings.setSettingValue('Comfy.ColorPalette',palette);},palette);await expect(root).toHaveAttribute('data-ca-theme',palette);
|
||||
const [p,n]=await Promise.all(['Positive prompt','Negative prompt'].map(name=>root.getByRole('textbox',{name,exact:true}).boundingBox()));expect(Math.abs(p.width-n.width)).toBeLessThan(1);expect(Math.abs(p.height-n.height)).toBeLessThan(1);expect(p.height).toBeGreaterThan(100);
|
||||
await page.screenshot({path:testInfo.outputPath(`pair-${palette}.png`)});
|
||||
}
|
||||
});
|
||||
|
||||
test('paired explicit Save persists an empty negative through actual API and preserves local draft',async({page,request},testInfo)=>{
|
||||
const id=await addPair(page);const root=page.locator('.pm-pair');const text=`Saved pair empty negative ${testInfo.project.name}`;await root.getByRole('textbox',{name:'Positive prompt',exact:true}).fill(text);await root.getByRole('button',{name:'Save pair',exact:true}).click();
|
||||
await expect(root.getByRole('status').filter({hasText:'Pair saved.'})).toBeVisible();const response=await request.get(`/prompt_manager/pairs?text=${encodeURIComponent(text)}`);expect(response.ok()).toBe(true);expect((await response.json()).pairs).toEqual(expect.arrayContaining([expect.objectContaining({positive_text:text,negative_text:''})]));expect(await values(page,id)).toMatchObject({text,negative_text:''});
|
||||
});
|
||||
|
||||
test('paired library filters search both texts, persist on reload and keep editor metadata separate',async({page,request},testInfo)=>{
|
||||
const token=`paired-filter-${testInfo.project.name}`;await seed(request,`Matching ${token}`,`Negative ${token}`);const other=await request.post('/prompt_manager/pairs',{data:{positive_text:`Other ${token}`,negative_text:'different',category:'Other paired category',tags:['other-pair-tag']}});expect(other.ok()).toBe(true);
|
||||
const id=await addPair(page);const root=page.locator('.pm-pair');await root.getByRole('textbox',{name:'Positive prompt',exact:true}).fill('Draft positive');await root.getByRole('textbox',{name:'Negative prompt',exact:true}).fill('Draft negative');
|
||||
await root.getByText('Metadata · category / tags',{exact:true}).click();await root.getByRole('combobox',{name:'Category',exact:true}).fill('Keep category');await root.getByRole('combobox',{name:'Tags',exact:true}).fill('keep-tag');await root.getByText('Metadata · category / tags',{exact:true}).click();
|
||||
await root.getByRole('searchbox',{name:'Search saved prompts'}).fill(token);await root.getByLabel('Filter category',{exact:true}).fill('Paired host tests');await root.getByLabel('Filter tags',{exact:true}).fill('pair-host, ');await root.getByRole('button',{name:'Search',exact:true}).click();
|
||||
await expect(root.locator('.pm-editor-result')).toHaveCount(1);await root.locator('.pm-editor-result').click();await expect(root.locator('.pm-editor-selection')).toContainText(`Negative ${token}`);
|
||||
expect(await values(page,id)).toMatchObject({text:'Draft positive',negative_text:'Draft negative',category:'Keep category',tags:'keep-tag'});await root.getByRole('switch',{name:'Load metadata'}).uncheck();await root.getByRole('button',{name:'Use pair',exact:true}).click();expect(await values(page,id)).toMatchObject({text:`Matching ${token}`,negative_text:`Negative ${token}`,category:'Keep category',tags:'keep-tag'});
|
||||
await page.screenshot({path:testInfo.outputPath('pair-library-preview.png')});
|
||||
await page.evaluate(async()=>{const {app}=await import('/scripts/app.js');await app.loadGraphData(app.graph.serialize());});await expect(root.getByLabel('Filter category',{exact:true})).toHaveValue('Paired host tests');await expect(root.getByLabel('Filter tags',{exact:true})).toHaveValue('pair-host, ');
|
||||
await page.evaluate(async id=>{const {app}=await import('/scripts/app.js');const node=app.graph.getNodeById(id);node.setSize([560,1000]);app.canvas.centerOnNode(node);},id);
|
||||
const [p,n]=await Promise.all(['Positive prompt','Negative prompt'].map(name=>root.getByRole('textbox',{name,exact:true}).boundingBox()));expect(n.y).toBeGreaterThan(p.y+p.height-1);expect(Math.abs(p.width-n.width)).toBeLessThan(1);await page.screenshot({path:testInfo.outputPath('pair-narrow.png')});
|
||||
});
|
||||
|
||||
for(const field of ['category','tags'])test(`manual Save pair guards real connected ${field} and resumes after disconnect`,async({page},testInfo)=>{
|
||||
const id=await addPair(page);const root=page.locator('.pm-pair');await root.getByRole('textbox',{name:'Positive prompt',exact:true}).fill(`Metadata link ${field} ${testInfo.project.name}`);
|
||||
let saves=0;await page.route('**/prompt_manager/pairs',async route=>{if(route.request().method()==='POST')saves++;await route.continue();});
|
||||
const link=await page.evaluate(async({id,field})=>{const {app}=await import('/scripts/app.js');const {convertToInput}=await import('/extensions/core/widgetInputs.js');const target=app.graph.getNodeById(id);convertToInput(target,target.widgets.find(widget=>widget.name===field));const source=LiteGraph.createNode('PrimitiveString');app.graph.add(source);source.widgets.find(widget=>widget.name==='value').value='Connected metadata';source.connect(0,target,target.inputs.findIndex(input=>input.name===field));app.canvas.centerOnNode(target);return [String(source.id),0];},{id,field});
|
||||
const save=root.getByRole('button',{name:'Save pair',exact:true});await expect(save).toBeDisabled();await expect(root.getByRole('status').filter({hasText:'Run the node to save connected values.'})).toBeVisible();await save.evaluate(button=>button.dispatchEvent(new Event('click',{bubbles:true})));expect(saves).toBe(0);
|
||||
const input=await page.evaluate(async({id,field})=>{const {app}=await import('/scripts/app.js');return(await app.graphToPrompt()).output[String(id)].inputs[field];},{id,field});expect(input).toEqual(link);
|
||||
await page.evaluate(async({id,field})=>{const {app}=await import('/scripts/app.js');const target=app.graph.getNodeById(id);target.disconnectInput(target.inputs.findIndex(input=>input.name===field));},{id,field});await expect(save).toBeEnabled();await save.click();await expect(root.getByRole('status').filter({hasText:'Pair saved.'})).toBeVisible();expect(saves).toBe(1);
|
||||
});
|
||||
@@ -0,0 +1,61 @@
|
||||
const test = require('node:test');
|
||||
const assert = require('node:assert/strict');
|
||||
const fs = require('node:fs');
|
||||
const vm = require('node:vm');
|
||||
const path = require('node:path');
|
||||
const ImageHelpers = require('../../web/js/image-helpers.js');
|
||||
const {extractGenerationParams} = require('../../web/js/comfy-metadata.js');
|
||||
|
||||
function graph(positive = 0, negative = 1) {
|
||||
return {
|
||||
1:{class_type:'PromptManagerPair',inputs:{text:'cat',negative_text:'blur',prepend_text:'trigger',append_text:'detail'}},
|
||||
2:{class_type:'KSampler',inputs:{positive:['1',positive],negative:['1',negative],model:['3',0]}},
|
||||
3:{class_type:'CheckpointLoaderSimple',inputs:{ckpt_name:'m.safetensors'}},
|
||||
};
|
||||
}
|
||||
|
||||
test('paired conditioning uses actual output roles and composes only positive', () => {
|
||||
const metadata = extractGenerationParams({prompt:graph()});
|
||||
assert.equal(metadata.positivePrompt,'trigger cat detail');
|
||||
assert.equal(metadata.negativePrompt,'blur');
|
||||
const swapped = extractGenerationParams({prompt:graph(1,0)});
|
||||
assert.equal(swapped.positivePrompt,'blur');
|
||||
assert.equal(swapped.negativePrompt,'trigger cat detail');
|
||||
});
|
||||
|
||||
test('paired strings passed through encoders retain the connected output slot', () => {
|
||||
const prompt=graph();
|
||||
prompt[4]={class_type:'CLIPTextEncode',inputs:{text:['1',2]}};
|
||||
prompt[5]={class_type:'CLIPTextEncode',inputs:{text:['1',3]}};
|
||||
prompt[2].inputs.positive=['4',0];prompt[2].inputs.negative=['5',0];
|
||||
const metadata=extractGenerationParams({prompt});
|
||||
assert.equal(metadata.positivePrompt,'trigger cat detail');
|
||||
assert.equal(metadata.negativePrompt,'blur');
|
||||
});
|
||||
|
||||
test('recorded image pair overrides unrelated PNG fallback, including an empty negative', async () => {
|
||||
for (const negative of ['blur','']) {
|
||||
const context=vm.createContext({document:{addEventListener(){}},window:{},console,ImageHelpers,
|
||||
ApiClient:{fetchUrl:async()=>({ok:true,json:async()=>({success:true,prompt:{text:'raw cat',pair:{negative_text:negative},pair_context:{positive_text:'trigger cat',negative_text:negative}}})})},
|
||||
});
|
||||
const Manager=vm.runInContext(fs.readFileSync(path.join(__dirname,'../../web/js/gallery.js'),'utf8')+'\nGalleryManager;',context);
|
||||
const manager=Object.create(Manager.prototype);
|
||||
manager.extractImageMetadata=async()=>({positivePrompt:'wrong',negativePrompt:'wrong'});
|
||||
manager.displayMetadata=data=>{manager.last=data;};
|
||||
await manager.loadImageMetadata('/prompt_manager/images/serve/example.png',{});
|
||||
assert.equal(manager.last.positivePrompt,'trigger cat');
|
||||
assert.equal(manager.last.negativePrompt,negative);
|
||||
}
|
||||
});
|
||||
|
||||
test('gallery labels graph-inferred text separately from recorded execution text', () => {
|
||||
const context=vm.createContext({document:{addEventListener(){}},window:{},console,escapeHtml:value=>String(value)});
|
||||
const Manager=vm.runInContext(fs.readFileSync(path.join(__dirname,'../../web/js/gallery.js'),'utf8')+'\nGalleryManager;',context);
|
||||
const manager=Object.create(Manager.prototype);
|
||||
manager.formatPromptText=value=>String(value||'');manager.setupMetadataActions=()=>{};
|
||||
for (const [source,label] of [['graph-inferred','Recovered from workflow metadata'],['execution','Recorded during execution'],['embedded-parameters','Read from embedded parameters'],['mixed','Execution and workflow metadata']]) {
|
||||
const panel={};
|
||||
manager.displayMetadata({positivePrompt:'cat',negativePrompt:'blur',textSource:source},panel);
|
||||
assert.ok(panel.innerHTML.includes(label),source);
|
||||
}
|
||||
});
|
||||
@@ -76,3 +76,30 @@ test("progressLabel summarises a running scan for the minimized pill", () => {
|
||||
assert.equal(progressLabel(running), "Scanning 42% · 84 files · 12 prompts");
|
||||
assert.equal(progressLabel(initialState()), "");
|
||||
});
|
||||
|
||||
test('pair totals survive completion and cancellation closes the running state', () => {
|
||||
const running=reduceScanEvent(initialState(),{type:'progress',progress:30,processed:3});
|
||||
const done=reduceScanEvent(running,{type:'complete',pairs_added:2,pair_links_added:3,enriched:1,unchanged:4,ambiguous:1,failed:2});
|
||||
assert.deepEqual([done.pairs_added,done.pair_links_added,done.enriched,done.unchanged,done.ambiguous,done.failed],[2,3,1,4,1,2]);
|
||||
const cancelled=reduceScanEvent(running,{type:'cancelled',processed:3,pairs_added:1});
|
||||
assert.equal(cancelled.status,'done');
|
||||
assert.equal(cancelled.cancelled,true);
|
||||
assert.equal(cancelled.statusText,'Scan cancelled.');
|
||||
assert.equal(cancelled.progress,30);
|
||||
});
|
||||
|
||||
test('completion label reports pair enrichment, unresolved work and failures', () => {
|
||||
const {completionLabel}=require('../../web/js/scan-progress.js');
|
||||
const state=reduceScanEvent(initialState(),{type:'complete',processed:10,found:8,added:1,linked:2,pairs_added:3,pair_links_added:4,enriched:2,unchanged:1,ambiguous:1,unresolved:2,conflict:1,failed:1});
|
||||
assert.match(completionLabel(state),/3 new pairs, 4 pair links/);
|
||||
assert.match(completionLabel(state),/4 need review, 1 failed/);
|
||||
assert.match(completionLabel({...state,cancelled:true}),/^Scan cancelled!/);
|
||||
});
|
||||
|
||||
test('older servers without pair counts retain a complete readable summary', () => {
|
||||
const {completionLabel}=require('../../web/js/scan-progress.js');
|
||||
const state=reduceScanEvent(initialState(),{type:'complete',processed:2,found:1,added:1,linked:0});
|
||||
assert.match(completionLabel(initialState()),/0 new pairs, 0 pair links; 0 need review, 0 failed/);
|
||||
assert.match(completionLabel(state),/0 enriched, 0 unchanged/);
|
||||
assert.equal(fromStatus({running:true,last_event:{type:'unknown'}}).status,'running');
|
||||
});
|
||||
|
||||
@@ -31,8 +31,7 @@ class APITestCase(AioHTTPTestCase):
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(self._temp_db.name)
|
||||
self.api = PromptManagerAPI(db=PromptDatabase(self._temp_db.name))
|
||||
self.api.add_routes(routes)
|
||||
app.router.add_routes(routes)
|
||||
return app
|
||||
|
||||
@@ -70,8 +70,7 @@ class AdminAPITestCase(AioHTTPTestCase):
|
||||
db_path = os.path.join(self.tmpdir, "prompts.db")
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(db_path)
|
||||
self.api = PromptManagerAPI(db=PromptDatabase(db_path))
|
||||
self.api.add_routes(routes)
|
||||
app.router.add_routes(routes)
|
||||
return app
|
||||
@@ -476,10 +475,13 @@ class TestScansRunOffTheEventLoop(AdminAPITestCase):
|
||||
).text()
|
||||
|
||||
self.assertIn('"added": 1', first)
|
||||
self.assertIn('"linked": 1', second)
|
||||
self.assertIn('"unchanged": 1', second)
|
||||
self.assertIn('"added": 0', second)
|
||||
prompts = self.api.db.get_recent_prompts(limit=10)["prompts"]
|
||||
self.assertEqual([p["text"] for p in prompts], ["a red square"])
|
||||
self.assertEqual({p["text"] for p in prompts}, {"a red square", "bad"})
|
||||
pairs = self.api.db.list_prompt_pairs()["pairs"]
|
||||
self.assertEqual(len(pairs), 1)
|
||||
self.assertEqual(pairs[0]["negative_text"], "bad")
|
||||
|
||||
async def test_output_scan_without_directories_reports_error(self):
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.outside_dir)]
|
||||
@@ -1556,7 +1558,7 @@ class TestOutputScanInternals(AdminAPITestCase):
|
||||
"found",
|
||||
)
|
||||
self.assertEqual(self.api._ingest_scanned_file(f, meta), "added")
|
||||
self.assertEqual(self.api._ingest_scanned_file(f, meta), "linked")
|
||||
self.assertEqual(self.api._ingest_scanned_file(f, meta), "unchanged")
|
||||
|
||||
def test_ingest_link_failures_are_tolerated(self):
|
||||
f = Path("gen.png")
|
||||
@@ -1589,7 +1591,23 @@ class TestOutputScanInternals(AdminAPITestCase):
|
||||
counts = self.api._ingest_scan_batch_sync(
|
||||
[(Path("a.png"), meta), (Path("b.png"), {})]
|
||||
)
|
||||
self.assertEqual(counts, {"processed": 2, "found": 0, "added": 0, "linked": 0})
|
||||
self.assertEqual(
|
||||
counts,
|
||||
{
|
||||
"processed": 2,
|
||||
"found": 0,
|
||||
"added": 0,
|
||||
"linked": 0,
|
||||
"enriched": 0,
|
||||
"unchanged": 0,
|
||||
"ambiguous": 0,
|
||||
"unresolved": 0,
|
||||
"conflict": 0,
|
||||
"failed": 1,
|
||||
"pairs_added": 0,
|
||||
"pair_links_added": 0,
|
||||
},
|
||||
)
|
||||
|
||||
async def test_scan_internal_error_streams_error_event(self):
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
|
||||
@@ -1637,7 +1655,10 @@ class TestOutputScanIsABackgroundJob(AdminAPITestCase):
|
||||
await asyncio.wait_for(self.api._scan_job.task, 10)
|
||||
|
||||
prompts = self.api.db.get_recent_prompts(limit=10)["prompts"]
|
||||
self.assertEqual([p["text"] for p in prompts], ["a red square"])
|
||||
self.assertEqual({p["text"] for p in prompts}, {"a red square", "bad"})
|
||||
pairs = self.api.db.list_prompt_pairs()["pairs"]
|
||||
self.assertEqual(len(pairs), 1)
|
||||
self.assertEqual(pairs[0]["negative_text"], "bad")
|
||||
self.assertEqual(self.api._scan_job.last_event["type"], "complete")
|
||||
self.assertFalse(self.api._scan_job.running)
|
||||
|
||||
@@ -1704,7 +1725,7 @@ class TestOutputScanIsABackgroundJob(AdminAPITestCase):
|
||||
body = await (await self.client.request("POST", "/prompt_manager/scan")).text()
|
||||
|
||||
self.assertEqual(len(calls), 2)
|
||||
self.assertIn('"linked": 1', body)
|
||||
self.assertIn('"unchanged": 1', body)
|
||||
|
||||
async def test_failed_scan_is_not_left_running(self):
|
||||
GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)]
|
||||
|
||||
@@ -119,8 +119,7 @@ class AutotagAPITestCase(AioHTTPTestCase):
|
||||
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(self._temp_db.name)
|
||||
self.api = PromptManagerAPI(db=PromptDatabase(self._temp_db.name))
|
||||
self.api._get_all_output_dirs = lambda: [self.output_dir]
|
||||
self.api.add_routes(routes)
|
||||
app.router.add_routes(routes)
|
||||
|
||||
@@ -45,8 +45,7 @@ class BackupRestoreTestCase(AioHTTPTestCase):
|
||||
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(self.db_path)
|
||||
self.api = PromptManagerAPI(db=PromptDatabase(self.db_path))
|
||||
self.api.add_routes(routes)
|
||||
app.router.add_routes(routes)
|
||||
return app
|
||||
|
||||
@@ -322,7 +322,9 @@ class TestLongJobLock(unittest.TestCase):
|
||||
"""_acquire_long_job: one job per family at a time, released on exit."""
|
||||
|
||||
def setUp(self):
|
||||
self.api = PromptManagerAPI()
|
||||
database = MagicMock()
|
||||
database.cleanup_duplicates.return_value = 0
|
||||
self.api = PromptManagerAPI(db=database)
|
||||
|
||||
def _run(self, coro):
|
||||
import asyncio
|
||||
@@ -373,8 +375,9 @@ class CoreRoutesTestCase(AioHTTPTestCase):
|
||||
self.tmpdir = tempfile.mkdtemp()
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(os.path.join(self.tmpdir, "prompts.db"))
|
||||
self.api = PromptManagerAPI(
|
||||
db=PromptDatabase(os.path.join(self.tmpdir, "prompts.db"))
|
||||
)
|
||||
self.api.add_routes(routes)
|
||||
app.router.add_routes(routes)
|
||||
return app
|
||||
@@ -565,8 +568,9 @@ class TestInstanceHelpers(FolderPathsFixture):
|
||||
|
||||
def setUp(self):
|
||||
super().setUp()
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(os.path.join(self.tmpdir, "prompts.db"))
|
||||
self.api = PromptManagerAPI(
|
||||
db=PromptDatabase(os.path.join(self.tmpdir, "prompts.db"))
|
||||
)
|
||||
self.addCleanup(self.api.db.close_all) # runs before the rmtree cleanup
|
||||
orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES)
|
||||
self.addCleanup(setattr, GalleryConfig, "MONITORING_DIRECTORIES", orig_dirs)
|
||||
@@ -668,7 +672,9 @@ class TestInstanceHelpers(FolderPathsFixture):
|
||||
class TestMetadataParsing(unittest.TestCase):
|
||||
|
||||
def setUp(self):
|
||||
self.api = PromptManagerAPI()
|
||||
database = MagicMock()
|
||||
database.cleanup_duplicates.return_value = 0
|
||||
self.api = PromptManagerAPI(db=database)
|
||||
self.tmpdir = tempfile.mkdtemp()
|
||||
self.addCleanup(shutil.rmtree, self.tmpdir, True)
|
||||
|
||||
|
||||
@@ -0,0 +1,71 @@
|
||||
"""API fixtures never open the configured database before installing test state."""
|
||||
|
||||
import unittest.mock
|
||||
|
||||
from py.api import PromptManagerAPI
|
||||
|
||||
|
||||
class TestAPIDatabaseIsolation(unittest.TestCase):
|
||||
def test_injected_database_never_constructs_the_configured_database(self):
|
||||
database = unittest.mock.MagicMock()
|
||||
database.cleanup_duplicates.return_value = 0
|
||||
with (
|
||||
unittest.mock.patch(
|
||||
"py.api.PromptDatabase", side_effect=AssertionError("default DB opened")
|
||||
) as factory,
|
||||
unittest.mock.patch("py.api.get_logger"),
|
||||
):
|
||||
api = PromptManagerAPI(db=database)
|
||||
self.assertIs(api.db, database)
|
||||
factory.assert_not_called()
|
||||
database.cleanup_duplicates.assert_called_once_with()
|
||||
|
||||
def test_falsey_injected_database_is_still_used(self):
|
||||
database = unittest.mock.MagicMock()
|
||||
database.__bool__.return_value = False
|
||||
database.cleanup_duplicates.return_value = 0
|
||||
with (
|
||||
unittest.mock.patch(
|
||||
"py.api.PromptDatabase", side_effect=AssertionError("default DB opened")
|
||||
) as factory,
|
||||
unittest.mock.patch("py.api.get_logger"),
|
||||
):
|
||||
api = PromptManagerAPI(db=database)
|
||||
self.assertIs(api.db, database)
|
||||
factory.assert_not_called()
|
||||
|
||||
def test_default_runtime_construction_still_uses_the_configured_factory(self):
|
||||
database = unittest.mock.MagicMock()
|
||||
database.cleanup_duplicates.return_value = 2
|
||||
with (
|
||||
unittest.mock.patch(
|
||||
"py.api.PromptDatabase", return_value=database
|
||||
) as factory,
|
||||
unittest.mock.patch("py.api.get_logger") as logger,
|
||||
):
|
||||
api = PromptManagerAPI()
|
||||
self.assertIs(api.db, database)
|
||||
factory.assert_called_once_with()
|
||||
logger.return_value.info.assert_any_call(
|
||||
"Startup cleanup: removed 2 duplicate prompts"
|
||||
)
|
||||
|
||||
def test_injected_cleanup_failure_is_logged_without_falling_back_to_default(self):
|
||||
database = unittest.mock.MagicMock()
|
||||
database.cleanup_duplicates.side_effect = RuntimeError("test cleanup")
|
||||
with (
|
||||
unittest.mock.patch(
|
||||
"py.api.PromptDatabase", side_effect=AssertionError("default DB opened")
|
||||
) as factory,
|
||||
unittest.mock.patch("py.api.get_logger") as logger,
|
||||
):
|
||||
api = PromptManagerAPI(db=database)
|
||||
self.assertIs(api.db, database)
|
||||
factory.assert_not_called()
|
||||
logger.return_value.error.assert_called_once_with(
|
||||
"Startup cleanup failed: test cleanup"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -64,8 +64,7 @@ class ImageAPITestCase(AioHTTPTestCase):
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(self._temp_db.name)
|
||||
self.api = PromptManagerAPI(db=PromptDatabase(self._temp_db.name))
|
||||
self.api._get_all_output_dirs = lambda: [self.output_dir]
|
||||
self.api._find_comfyui_output_dir = lambda: str(self.output_dir)
|
||||
self.api.add_routes(routes)
|
||||
|
||||
@@ -104,8 +104,9 @@ class LoggingAPITestCase(AioHTTPTestCase):
|
||||
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(os.path.join(self.tmpdir, "prompts.db"))
|
||||
self.api = PromptManagerAPI(
|
||||
db=PromptDatabase(os.path.join(self.tmpdir, "prompts.db"))
|
||||
)
|
||||
self.api._get_logger_manager = lambda: self.manager
|
||||
self.api.add_routes(routes)
|
||||
app.router.add_routes(routes)
|
||||
|
||||
@@ -76,8 +76,7 @@ class LoraAPITestCase(AioHTTPTestCase):
|
||||
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(self._temp_db.name)
|
||||
self.api = PromptManagerAPI(db=PromptDatabase(self._temp_db.name))
|
||||
self.api.add_routes(routes)
|
||||
app.router.add_routes(routes)
|
||||
return app
|
||||
|
||||
@@ -0,0 +1,154 @@
|
||||
"""Pair routes preserve shared text and expose portable combinations."""
|
||||
|
||||
from tests.test_api_prompts import PromptAPITestCase
|
||||
|
||||
|
||||
class TestPairAPI(PromptAPITestCase):
|
||||
async def _create(self, **overrides):
|
||||
payload = {"positive_text": "cat", "negative_text": "blur"}
|
||||
payload.update(overrides)
|
||||
return await self.client.post("/prompt_manager/pairs", json=payload)
|
||||
|
||||
async def test_create_reuse_read_and_search(self):
|
||||
response = await self._create(tags=["studio"], title="Portrait")
|
||||
self.assertEqual(response.status, 200)
|
||||
pair = (await response.json())["pair"]
|
||||
response = await self._create(negative_text="noise")
|
||||
second = (await response.json())["pair"]
|
||||
self.assertEqual(pair["positive_prompt_id"], second["positive_prompt_id"])
|
||||
self.assertNotEqual(pair["id"], second["id"])
|
||||
response = await self.client.get("/prompt_manager/pairs?text=blur")
|
||||
data = await response.json()
|
||||
self.assertEqual(data["total"], 1)
|
||||
self.assertEqual(data["pairs"][0]["tags"], ["studio"])
|
||||
response = await self.client.get(f'/prompt_manager/pairs/{pair["id"]}')
|
||||
self.assertEqual((await response.json())["pair"]["title"], "Portrait")
|
||||
|
||||
async def test_negative_must_be_explicit_but_can_be_empty(self):
|
||||
response = await self.client.post(
|
||||
"/prompt_manager/pairs", json={"positive_text": "cat"}
|
||||
)
|
||||
self.assertEqual(response.status, 400)
|
||||
response = await self._create(negative_text="")
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertIsNone((await response.json())["pair"]["negative_prompt_id"])
|
||||
|
||||
async def test_invalid_payloads_are_400_without_writes(self):
|
||||
for body in (
|
||||
[],
|
||||
None,
|
||||
{},
|
||||
{"positive_text": 12, "negative_text": ""},
|
||||
{"positive_text": "cat", "negative_text": None},
|
||||
{"positive_text": "cat", "negative_text": "blur", "rating": True},
|
||||
{"positive_text": "cat", "negative_text": "blur", "tags": [12]},
|
||||
):
|
||||
response = await self.client.post("/prompt_manager/pairs", json=body)
|
||||
self.assertEqual(response.status, 400, body)
|
||||
response = await self.client.get("/prompt_manager/pairs")
|
||||
self.assertEqual((await response.json())["total"], 0)
|
||||
|
||||
async def test_repeated_save_does_not_overwrite_metadata(self):
|
||||
response = await self._create(title="Original", tags=["kept"])
|
||||
pair = (await response.json())["pair"]
|
||||
response = await self._create(title="Other", tags=["ignored"])
|
||||
repeated = (await response.json())["pair"]
|
||||
self.assertEqual(repeated["id"], pair["id"])
|
||||
self.assertEqual(repeated["title"], "Original")
|
||||
self.assertEqual(repeated["tags"], ["kept"])
|
||||
|
||||
async def test_metadata_update_and_archive(self):
|
||||
pair = (await (await self._create()).json())["pair"]
|
||||
url = f'/prompt_manager/pairs/{pair["id"]}'
|
||||
response = await self.client.patch(url, json={"title": "New", "tags": ["a"]})
|
||||
self.assertEqual(response.status, 200)
|
||||
response = await self.client.patch(url, json={"positive_text": "dog"})
|
||||
self.assertEqual(response.status, 400)
|
||||
response = await self.client.patch(url, json={"archived": True})
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual(
|
||||
(await (await self.client.get("/prompt_manager/pairs")).json())["total"], 0
|
||||
)
|
||||
response = await self.client.get("/prompt_manager/pairs?include_archived=true")
|
||||
self.assertEqual((await response.json())["total"], 1)
|
||||
response = await self.client.delete(url)
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertIsNotNone(self.api.db.get_prompt_by_id(pair["positive_prompt_id"]))
|
||||
|
||||
async def test_export_is_versioned_and_contains_text_not_only_local_ids(self):
|
||||
await self._create(tags=["a"])
|
||||
response = await self.client.get("/prompt_manager/pairs/export")
|
||||
self.assertEqual(response.status, 200)
|
||||
data = await response.json()
|
||||
self.assertEqual(data["schema_version"], 1)
|
||||
self.assertEqual(data["pairs"][0]["positive_text"], "cat")
|
||||
self.assertEqual(data["pairs"][0]["negative_text"], "blur")
|
||||
self.assertNotIn("positive_prompt_id", data["pairs"][0])
|
||||
|
||||
async def test_missing_invalid_ids_and_pagination(self):
|
||||
for verb in ("get", "delete"):
|
||||
for suffix, status in (("abc", 400), ("9999", 404)):
|
||||
response = await getattr(self.client, verb)(
|
||||
f"/prompt_manager/pairs/{suffix}"
|
||||
)
|
||||
self.assertEqual(response.status, status)
|
||||
for query in ("limit=abc", "offset=abc", "include_archived=maybe"):
|
||||
response = await self.client.get("/prompt_manager/pairs?" + query)
|
||||
self.assertEqual(response.status, 400)
|
||||
response = await self.client.get("/prompt_manager/pairs?limit=99999&offset=-1")
|
||||
data = await response.json()
|
||||
self.assertEqual(data["limit"], 500)
|
||||
self.assertEqual(data["offset"], 0)
|
||||
|
||||
async def test_old_edit_and_delete_routes_report_pair_dependency(self):
|
||||
pair = (await (await self._create()).json())["pair"]
|
||||
prompt_id = pair["positive_prompt_id"]
|
||||
response = await self.client.put(
|
||||
f"/prompt_manager/prompts/{prompt_id}", json={"text": "dog"}
|
||||
)
|
||||
self.assertEqual(response.status, 409)
|
||||
response = await self.client.delete(f"/prompt_manager/delete/{prompt_id}")
|
||||
self.assertEqual(response.status, 409)
|
||||
response = await self.client.post(
|
||||
"/prompt_manager/bulk/delete", json={"prompt_ids": [prompt_id]}
|
||||
)
|
||||
self.assertEqual(response.status, 409)
|
||||
|
||||
async def test_recent_pairs_sort_by_use_and_reject_unknown_sort(self):
|
||||
first = (await (await self._create()).json())["pair"]
|
||||
await self._create(negative_text="noise")
|
||||
self.api.db.record_prompt_pair_use(first["id"])
|
||||
response = await self.client.get("/prompt_manager/pairs?sort=last_used_desc")
|
||||
self.assertEqual((await response.json())["pairs"][0]["id"], first["id"])
|
||||
response = await self.client.get("/prompt_manager/pairs?sort=invalid")
|
||||
self.assertEqual(response.status, 400)
|
||||
|
||||
async def test_null_metadata_patch_is_noop_for_existing_pair(self):
|
||||
pair = (await (await self._create(rating=4)).json())["pair"]
|
||||
response = await self.client.patch(
|
||||
f'/prompt_manager/pairs/{pair["id"]}', json={"rating": None}
|
||||
)
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertEqual((await response.json())["pair"]["rating"], 4)
|
||||
response = await self.client.patch(
|
||||
"/prompt_manager/pairs/99999", json={"rating": None}
|
||||
)
|
||||
self.assertEqual(response.status, 404)
|
||||
|
||||
async def test_database_failures_are_generic_and_bad_json_is_rejected(self):
|
||||
from unittest.mock import patch
|
||||
|
||||
with patch.object(
|
||||
self.api.db,
|
||||
"save_prompt_pair",
|
||||
side_effect=RuntimeError("/private/db/path"),
|
||||
):
|
||||
response = await self._create()
|
||||
self.assertEqual(response.status, 500)
|
||||
self.assertNotIn("/private", await response.text())
|
||||
response = await self.client.post(
|
||||
"/prompt_manager/pairs",
|
||||
data="{",
|
||||
headers={"Content-Type": "application/json"},
|
||||
)
|
||||
self.assertEqual(response.status, 400)
|
||||
@@ -17,6 +17,11 @@ sys.modules.setdefault("server", MagicMock())
|
||||
from aiohttp import web # noqa: E402
|
||||
from aiohttp.test_utils import AioHTTPTestCase # noqa: E402
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.operations import PromptDatabase # noqa: E402
|
||||
from py.api import PromptManagerAPI # noqa: E402
|
||||
import py.api.prompts as prompts_module # noqa: E402
|
||||
@@ -38,8 +43,7 @@ class PromptAPITestCase(AioHTTPTestCase):
|
||||
app = web.Application()
|
||||
routes = web.RouteTableDef()
|
||||
|
||||
self.api = PromptManagerAPI()
|
||||
self.api.db = PromptDatabase(self._temp_db.name)
|
||||
self.api = PromptManagerAPI(db=PromptDatabase(self._temp_db.name))
|
||||
self.api.add_routes(routes)
|
||||
app.router.add_routes(routes)
|
||||
return app
|
||||
@@ -48,6 +52,7 @@ class PromptAPITestCase(AioHTTPTestCase):
|
||||
await super().tearDownAsync() # closes the aiohttp test client
|
||||
# Windows refuses to unlink a database that still has open handles.
|
||||
self.api.db.close_all()
|
||||
assert_closed(self, self._temp_db.name)
|
||||
for path in (
|
||||
self._temp_db.name,
|
||||
self._temp_db.name + "-wal",
|
||||
|
||||
@@ -299,7 +299,9 @@ class TestGetAllOutputDirsFiltersRoots(GalleryRootTestCase):
|
||||
super().setUp()
|
||||
from py.api import PromptManagerAPI
|
||||
|
||||
self.api = PromptManagerAPI()
|
||||
database = MagicMock()
|
||||
database.cleanup_duplicates.return_value = 0
|
||||
self.api = PromptManagerAPI(db=database)
|
||||
self.api._cached_output_dir = None
|
||||
self._orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES)
|
||||
self.addCleanup(self._restore_dirs)
|
||||
|
||||
@@ -138,6 +138,7 @@ class TestPackageWiring(PackageTestCase):
|
||||
[
|
||||
"PromptManager",
|
||||
"PromptManagerFocused",
|
||||
"PromptManagerPair",
|
||||
"PromptManagerSplit",
|
||||
"PromptManagerText",
|
||||
"PromptSearchList",
|
||||
|
||||
@@ -9,7 +9,7 @@ import sys
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||||
|
||||
@@ -26,7 +26,16 @@ class TestFindComfyuiOutputDir(unittest.TestCase):
|
||||
"""Test _find_comfyui_output_dir() respects user-configured directories."""
|
||||
|
||||
def setUp(self):
|
||||
self.api = PromptManagerAPI()
|
||||
roots = patch.object(
|
||||
GalleryConfig,
|
||||
"allowed_gallery_parents",
|
||||
return_value=[os.path.normcase(os.path.realpath(tempfile.gettempdir()))],
|
||||
)
|
||||
roots.start()
|
||||
self.addCleanup(roots.stop)
|
||||
database = MagicMock()
|
||||
database.cleanup_duplicates.return_value = 0
|
||||
self.api = PromptManagerAPI(db=database)
|
||||
self.api._cached_output_dir = None
|
||||
self._orig_dirs = list(GalleryConfig.MONITORING_DIRECTORIES)
|
||||
|
||||
@@ -112,5 +121,22 @@ class TestFindComfyuiOutputDir(unittest.TestCase):
|
||||
self.assertEqual(result, str(Path(tmpdir).resolve()))
|
||||
|
||||
|
||||
class TestOutputDirectoryFixturePortability(unittest.TestCase):
|
||||
def test_aliased_temporary_parent_uses_the_real_gallery_allowlist(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
parent = Path(directory)
|
||||
(parent / "alias").mkdir()
|
||||
(parent / "actual").mkdir()
|
||||
aliased = str(parent / "alias" / ".." / "actual")
|
||||
with patch("tempfile.gettempdir", return_value=aliased):
|
||||
case = TestFindComfyuiOutputDir(
|
||||
"test_configured_directory_takes_priority"
|
||||
)
|
||||
result = unittest.TestResult()
|
||||
case.run(result)
|
||||
self.assertEqual(result.errors, [])
|
||||
self.assertEqual(result.failures, [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
"""Dependency checks exclude competing pair writers until mutation finishes."""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import unittest.mock
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
|
||||
class TestPairDependencyConcurrency(unittest.TestCase):
|
||||
def test_legacy_mutators_lock_out_pair_creation_before_dependency_checks(self):
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = os.path.join(directory, "concurrent.db")
|
||||
db = PromptDatabase(path)
|
||||
try:
|
||||
original = db._ensure_prompts_unpaired
|
||||
checks = []
|
||||
|
||||
def check_with_competing_writer(conn, ids):
|
||||
original(conn, ids)
|
||||
competitor = sqlite3.connect(path, timeout=0)
|
||||
try:
|
||||
with self.assertRaisesRegex(sqlite3.OperationalError, "locked"):
|
||||
competitor.execute("BEGIN IMMEDIATE")
|
||||
finally:
|
||||
competitor.close()
|
||||
checks.append(list(ids))
|
||||
|
||||
actions = (
|
||||
lambda pid: db.update_prompt_text(pid, "changed"),
|
||||
db.delete_prompt,
|
||||
lambda pid: db.bulk_delete_prompts([pid]),
|
||||
lambda _pid: db.delete_prompts_by_category("temporary"),
|
||||
)
|
||||
with unittest.mock.patch.object(
|
||||
db, "_ensure_prompts_unpaired", check_with_competing_writer
|
||||
):
|
||||
for index, action in enumerate(actions):
|
||||
with self.subTest(action=index):
|
||||
text = f"unpaired {index}"
|
||||
pid = db.save_prompt(
|
||||
text,
|
||||
category="temporary",
|
||||
prompt_hash=generate_prompt_hash(text),
|
||||
)
|
||||
self.assertTrue(action(pid))
|
||||
self.assertEqual(len(checks), len(actions))
|
||||
finally:
|
||||
db.close_all()
|
||||
assert_closed(self, directory)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,328 @@
|
||||
"""Embedded actual encoding evidence is used only for its traced paired node."""
|
||||
|
||||
import json
|
||||
import logging
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import Mock
|
||||
|
||||
from PIL import Image
|
||||
from PIL.PngImagePlugin import PngInfo
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from prompt_manager_pair import PromptManagerPair
|
||||
from utils.metadata_extractor import ComfyUIMetadataExtractor
|
||||
from utils.prompt_pair_resolver import resolve_prompt_pair
|
||||
|
||||
KEY = "comfyassets_prompt_pairs"
|
||||
|
||||
|
||||
def graph(positive="cat", negative="blur"):
|
||||
return {
|
||||
"1": {
|
||||
"class_type": "PromptManagerPair",
|
||||
"inputs": {"text": positive, "negative_text": negative},
|
||||
},
|
||||
"5": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"positive": ["1", 0], "negative": ["1", 1]},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class TestPairExecution(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
self.addCleanup(assert_closed, self, self.tmp.name)
|
||||
self.db = PromptDatabase(str(Path(self.tmp.name) / "test.db"))
|
||||
self.addCleanup(self.db.close_all)
|
||||
self.node = PromptManagerPair.__new__(PromptManagerPair)
|
||||
self.node.db = self.db
|
||||
self.node.logger = logging.getLogger("test.execution")
|
||||
self.node.prompt_tracker = Mock()
|
||||
self.node.comfyui_integration = Mock()
|
||||
self.node._inject_lora_trigger_words = lambda text: text + " LORA"
|
||||
self.extra = {"workflow": {"unchanged": True}, "sibling": {"keep": True}}
|
||||
|
||||
def execute(self, text="cat", negative="blur", **kwargs):
|
||||
return self.node.encode_prompt(
|
||||
Mock(),
|
||||
text,
|
||||
negative_text=negative,
|
||||
unique_id="1",
|
||||
extra_pnginfo=self.extra,
|
||||
**kwargs,
|
||||
)
|
||||
|
||||
def resolve(self, source=None, extra=None):
|
||||
return resolve_prompt_pair(
|
||||
{"prompt": source or graph(), **(self.extra if extra is None else extra)}
|
||||
)
|
||||
|
||||
def test_schema_requests_extra_pnginfo_and_actual_positive_lora_is_preserved(self):
|
||||
self.assertEqual(
|
||||
PromptManagerPair.INPUT_TYPES()["hidden"]["extra_pnginfo"], "EXTRA_PNGINFO"
|
||||
)
|
||||
self.execute(prepend_text="trigger")
|
||||
source = graph()
|
||||
source["1"]["inputs"]["prepend_text"] = "trigger"
|
||||
result = self.resolve(source)
|
||||
self.assertEqual(result["state"], "resolved")
|
||||
self.assertEqual(result["positive"]["text"], "trigger cat LORA")
|
||||
self.assertEqual(result["positive"]["storage_text"], "cat")
|
||||
self.assertEqual(result["positive"]["provenance"]["text_source"], "execution")
|
||||
self.assertEqual(result["negative"]["text"], "blur")
|
||||
self.assertEqual(self.extra["workflow"], {"unchanged": True})
|
||||
self.assertEqual(self.extra["sibling"], {"keep": True})
|
||||
|
||||
def test_one_unknown_linked_execution_is_recoverable_without_fifo_guess(self):
|
||||
source = graph(positive=["2", 0])
|
||||
source["2"] = {"class_type": "PromptSearchList", "inputs": {}}
|
||||
self.execute(text="actual batch cat")
|
||||
result = self.resolve(source)
|
||||
self.assertEqual(result["state"], "resolved")
|
||||
self.assertEqual(result["positive"]["storage_text"], "actual batch cat")
|
||||
self.assertEqual(result["positive"]["text"], "actual batch cat LORA")
|
||||
|
||||
def test_distinct_batch_records_are_unresolved_instead_of_using_last_execution(
|
||||
self,
|
||||
):
|
||||
source = graph(positive=["2", 0])
|
||||
source["2"] = {"class_type": "PromptSearchList", "inputs": {}}
|
||||
self.execute(text="first")
|
||||
self.execute(text="second")
|
||||
result = self.resolve(source)
|
||||
self.assertEqual(result["state"], "unresolved")
|
||||
self.assertIn("multiple", result["reason"].lower())
|
||||
|
||||
def test_empty_negative_runtime_record_remains_explicit(self):
|
||||
self.execute(negative="")
|
||||
result = self.resolve(graph(negative=""))
|
||||
self.assertEqual(result["state"], "resolved")
|
||||
self.assertEqual(result["negative"]["text"], "")
|
||||
|
||||
def test_cached_graph_without_execution_evidence_is_explicitly_inferred(self):
|
||||
result = self.resolve(extra={})
|
||||
self.assertEqual(result["state"], "resolved")
|
||||
self.assertEqual(
|
||||
result["positive"]["provenance"]["text_source"], "graph-inferred"
|
||||
)
|
||||
|
||||
def test_stale_or_malformed_evidence_is_not_accepted(self):
|
||||
self.execute(text="stale")
|
||||
self.assertEqual(self.resolve()["state"], "unresolved")
|
||||
for value in (
|
||||
"not json",
|
||||
[],
|
||||
{"version": 900, "nodes": {}},
|
||||
{"version": 1, "nodes": {"1": {"records": [{"bad": True}]}}},
|
||||
):
|
||||
with self.subTest(value=value):
|
||||
self.assertEqual(
|
||||
self.resolve(extra={KEY: value})["state"], "unresolved"
|
||||
)
|
||||
|
||||
def test_unrelated_node_evidence_does_not_override_plain_encoders(self):
|
||||
source = graph()
|
||||
source["1"] = {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "plain positive"},
|
||||
}
|
||||
source["2"] = {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "plain negative"},
|
||||
}
|
||||
source["5"]["inputs"]["negative"] = ["2", 0]
|
||||
result = self.resolve(source, extra={KEY: "broken"})
|
||||
self.assertEqual(result["state"], "resolved")
|
||||
self.assertEqual(result["positive"]["text"], "plain positive")
|
||||
|
||||
def test_png_extractor_preserves_namespace_for_live_and_raw_scan_readers(self):
|
||||
self.execute()
|
||||
path = Path(self.tmp.name) / "image.png"
|
||||
info = PngInfo()
|
||||
info.add_text("prompt", json.dumps(graph()))
|
||||
info.add_text(KEY, json.dumps(self.extra[KEY]))
|
||||
Image.new("RGB", (2, 2)).save(path, pnginfo=info)
|
||||
parsed = ComfyUIMetadataExtractor().extract_metadata(str(path))
|
||||
self.assertEqual(resolve_prompt_pair(parsed)["positive"]["text"], "cat LORA")
|
||||
with Image.open(path) as image:
|
||||
raw = dict(image.text)
|
||||
raw["prompt"] = json.loads(raw["prompt"])
|
||||
self.assertEqual(resolve_prompt_pair(raw)["positive"]["text"], "cat LORA")
|
||||
|
||||
def test_recording_limits_fail_closed_and_repeat_identical_values_are_safe(self):
|
||||
from utils.pair_execution import MAX_EXECUTION_RECORDS
|
||||
|
||||
self.execute()
|
||||
self.execute()
|
||||
self.assertEqual(self.resolve()["state"], "resolved")
|
||||
for index in range(MAX_EXECUTION_RECORDS + 1):
|
||||
self.execute(text=str(index))
|
||||
self.assertEqual(self.resolve()["state"], "unresolved")
|
||||
|
||||
def test_live_metadata_lookup_uses_unique_execution_for_unknown_source(self):
|
||||
from utils.image_monitor import ImageGenerationHandler
|
||||
|
||||
source = graph(positive=["2", 0])
|
||||
source["2"] = {"class_type": "PromptSearchList", "inputs": {}}
|
||||
self.execute(text="actual")
|
||||
handler = ImageGenerationHandler.__new__(ImageGenerationHandler)
|
||||
handler.db_manager = self.db
|
||||
handler.logger = logging.getLogger("test.execution")
|
||||
state, context = handler._find_prompt_from_metadata(
|
||||
{"prompt": source, **self.extra}
|
||||
)
|
||||
self.assertEqual(state, "linked")
|
||||
expected = self.db.list_prompt_pairs()["pairs"][0]["positive_prompt_id"]
|
||||
self.assertEqual(context["id"], expected)
|
||||
|
||||
def test_submission_clears_old_execution_evidence_without_changing_workflow(self):
|
||||
from utils.usage_tracking import handle_queued_prompt, QueuedRuns
|
||||
|
||||
self.execute()
|
||||
request = {"prompt": graph(), "extra_data": {"extra_pnginfo": self.extra}}
|
||||
handle_queued_prompt(request, QueuedRuns(), pair_queued=QueuedRuns())
|
||||
self.assertNotIn(KEY, self.extra)
|
||||
self.assertEqual(self.extra["workflow"], {"unchanged": True})
|
||||
self.assertEqual(
|
||||
self.resolve()["positive"]["provenance"]["text_source"], "graph-inferred"
|
||||
)
|
||||
|
||||
def test_known_invalid_output_slot_or_cycles_cannot_be_overridden(self):
|
||||
self.execute()
|
||||
source = graph(positive=["2", 4])
|
||||
source["2"] = {"class_type": "PrimitiveString", "inputs": {"value": "cat"}}
|
||||
self.assertEqual(self.resolve(source)["state"], "unresolved")
|
||||
source["2"] = {
|
||||
"class_type": "StringConcatenate",
|
||||
"inputs": {"string_a": ["2", 0], "string_b": "cat"},
|
||||
}
|
||||
source["1"]["inputs"]["text"] = ["2", 0]
|
||||
self.assertEqual(self.resolve(source)["state"], "unresolved")
|
||||
|
||||
def test_live_image_association_keeps_actual_encoding_and_repeated_scan_is_stable(
|
||||
self,
|
||||
):
|
||||
from utils.image_monitor import ImageGenerationHandler
|
||||
from database.pair_ingestion import ingest_resolved_pair
|
||||
|
||||
self.execute()
|
||||
metadata = {"prompt": graph(), **self.extra}
|
||||
handler = ImageGenerationHandler.__new__(ImageGenerationHandler)
|
||||
handler.db_manager = self.db
|
||||
handler.logger = logging.getLogger("test.execution")
|
||||
_, context = handler._find_prompt_from_metadata(metadata)
|
||||
path = str(Path(self.tmp.name) / "actual.png")
|
||||
handler.link_image_to_prompt(path, context, metadata)
|
||||
pair_context = self.db.get_image_prompt_info(path)["pair_context"]
|
||||
self.assertEqual(pair_context["positive_text"], "cat LORA")
|
||||
repeated = ingest_resolved_pair(
|
||||
self.db, path, metadata, resolve_prompt_pair(metadata)
|
||||
)
|
||||
self.assertEqual(repeated["status"], "unchanged")
|
||||
|
||||
def test_unknown_batch_with_multiple_executions_is_not_exactly_linked(self):
|
||||
from utils.image_monitor import ImageGenerationHandler
|
||||
|
||||
source = graph(positive=["2", 0])
|
||||
source["2"] = {"class_type": "PromptSearchList", "inputs": {}}
|
||||
self.execute(text="first")
|
||||
self.execute(text="second")
|
||||
handler = ImageGenerationHandler.__new__(ImageGenerationHandler)
|
||||
handler.db_manager = self.db
|
||||
handler.logger = logging.getLogger("test.execution")
|
||||
state, context = handler._find_prompt_from_metadata(
|
||||
{"prompt": source, **self.extra}
|
||||
)
|
||||
self.assertEqual(state, "unknown")
|
||||
self.assertIsNone(context)
|
||||
|
||||
def test_execution_evidence_bounds_and_malformed_structures_fail_closed(self):
|
||||
from utils import pair_execution as evidence
|
||||
from unittest.mock import patch
|
||||
|
||||
kwargs = dict(
|
||||
text="cat",
|
||||
negative_text="blur",
|
||||
positive_text="cat LORA",
|
||||
negative_effective_text="blur",
|
||||
)
|
||||
evidence.record_pair_execution(None, "1", **kwargs)
|
||||
envelope = {}
|
||||
evidence.record_pair_execution(envelope, None, **kwargs)
|
||||
self.assertEqual(envelope, {})
|
||||
evidence.record_pair_execution(envelope, "1", **kwargs)
|
||||
self.assertIsNone(evidence.read_pair_execution(envelope, "untraced"))
|
||||
malformed = (
|
||||
{"version": 1, "nodes": []},
|
||||
{"version": 1, "truncated": True, "nodes": {}},
|
||||
{"version": 1, "nodes": {"1": []}},
|
||||
{"version": 1, "nodes": {"1": {"records": []}}},
|
||||
{"version": 1, "nodes": {"1": {"records": "bad"}}},
|
||||
{
|
||||
"version": 1,
|
||||
"nodes": {
|
||||
"1": {
|
||||
"records": [
|
||||
{
|
||||
**kwargs,
|
||||
"prepend_text": "",
|
||||
"append_text": "",
|
||||
"negative_effective_text": "changed",
|
||||
}
|
||||
]
|
||||
}
|
||||
},
|
||||
},
|
||||
)
|
||||
for value in malformed:
|
||||
with self.subTest(value=value):
|
||||
with self.assertRaises(evidence.ExecutionEvidenceError):
|
||||
evidence.read_pair_execution({KEY: value}, "1")
|
||||
oversized = {}
|
||||
evidence.record_pair_execution(
|
||||
oversized,
|
||||
"1",
|
||||
**{**kwargs, "text": "x" * (evidence.MAX_EXECUTION_TEXT + 1)},
|
||||
)
|
||||
with self.assertRaises(evidence.ExecutionEvidenceError):
|
||||
evidence.read_pair_execution(oversized, "1")
|
||||
with patch.object(evidence, "MAX_EXECUTION_JSON", 10):
|
||||
with self.assertRaises(evidence.ExecutionEvidenceError):
|
||||
evidence.read_pair_execution({KEY: json.dumps(envelope[KEY])}, "1")
|
||||
limited = {}
|
||||
evidence.record_pair_execution(limited, "1", **kwargs)
|
||||
with self.assertRaises(evidence.ExecutionEvidenceError):
|
||||
evidence.read_pair_execution(limited, "1")
|
||||
with patch.object(evidence, "MAX_EXECUTION_NODES", 1):
|
||||
evidence.record_pair_execution(envelope, "2", **kwargs)
|
||||
evidence.record_pair_execution(envelope, "3", **kwargs)
|
||||
with self.assertRaises(evidence.ExecutionEvidenceError):
|
||||
evidence.read_pair_execution(envelope, "1")
|
||||
|
||||
def test_corrupted_existing_namespace_cannot_break_valid_encoding(self):
|
||||
from utils.pair_execution import (
|
||||
record_pair_execution,
|
||||
read_pair_execution,
|
||||
ExecutionEvidenceError,
|
||||
)
|
||||
|
||||
kwargs = dict(
|
||||
text="cat",
|
||||
negative_text="blur",
|
||||
positive_text="cat LORA",
|
||||
negative_effective_text="blur",
|
||||
)
|
||||
for nodes in ([], {"1": {"records": "bad"}}, {"other": {"bad": object()}}):
|
||||
extra = {KEY: {"version": 1, "nodes": nodes}}
|
||||
record_pair_execution(extra, "1", **kwargs)
|
||||
with self.assertRaises(ExecutionEvidenceError):
|
||||
read_pair_execution(extra, "1")
|
||||
@@ -0,0 +1,61 @@
|
||||
"""Live metadata annotation uses the same conservative ingestion as scan."""
|
||||
|
||||
import logging
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.image_monitor import ImageGenerationHandler
|
||||
|
||||
|
||||
def metadata(positive="cat", negative="blur"):
|
||||
return {
|
||||
"prompt": {
|
||||
"1": {"class_type": "PromptManager", "inputs": {"text": positive}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": negative}},
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
class TestPairImageLinking(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.db = PromptDatabase(str(Path(self.temp.name) / "db.sqlite"))
|
||||
self.path = str(Path(self.temp.name) / "image.png")
|
||||
self.pid = self.db.save_prompt("cat", prompt_hash=generate_prompt_hash("cat"))
|
||||
self.handler = ImageGenerationHandler.__new__(ImageGenerationHandler)
|
||||
self.handler.db_manager = self.db
|
||||
self.handler.logger = logging.getLogger("test.pairlink")
|
||||
|
||||
def tearDown(self):
|
||||
self.db.close_all()
|
||||
assert_closed(self, self.temp.name)
|
||||
self.temp.cleanup()
|
||||
|
||||
def test_exact_metadata_creates_pair_link_without_counting_a_run(self):
|
||||
self.handler.link_image_to_prompt(self.path, {"id": self.pid}, metadata())
|
||||
self.handler.link_image_to_prompt(self.path, {"id": self.pid}, metadata())
|
||||
pairs = self.db.list_prompt_pairs()["pairs"]
|
||||
self.assertEqual(len(pairs), 1)
|
||||
self.assertEqual(pairs[0]["image_count"], 1)
|
||||
self.assertEqual(pairs[0]["run_count"], 0)
|
||||
|
||||
def test_stale_context_or_ambiguous_metadata_does_not_create_pair(self):
|
||||
meta = metadata(positive="dog")
|
||||
self.handler.link_image_to_prompt(self.path, {"id": self.pid}, meta)
|
||||
self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
|
||||
meta = metadata()
|
||||
meta["prompt"]["4"] = meta["prompt"]["3"]
|
||||
self.handler.link_image_to_prompt(self.path + "2", {"id": self.pid}, meta)
|
||||
self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
|
||||
@@ -0,0 +1,162 @@
|
||||
"""Atomic, repeatable metadata enrichment without changing usage counts."""
|
||||
|
||||
import tempfile
|
||||
import unittest
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from database.pair_ingestion import ingest_resolved_pair
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
|
||||
def resolved(positive="cat", negative="blur"):
|
||||
return {
|
||||
"state": "resolved",
|
||||
"positive": {"text": positive, "provenance": {"source": "api"}},
|
||||
"negative": {"text": negative, "provenance": {"source": "api"}},
|
||||
}
|
||||
|
||||
|
||||
class TestPairIngestion(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.db = PromptDatabase(str(Path(self.temp.name) / "test.db"))
|
||||
self.path = str(Path(self.temp.name) / "image.png")
|
||||
|
||||
def tearDown(self):
|
||||
self.db.close_all()
|
||||
assert_closed(self, self.temp.name)
|
||||
self.temp.cleanup()
|
||||
|
||||
def test_first_scan_repeat_and_negative_variants_reuse_components(self):
|
||||
first = ingest_resolved_pair(self.db, self.path, {}, resolved())
|
||||
self.assertEqual(first["status"], "added")
|
||||
self.assertEqual(first["pairs_added"], 1)
|
||||
second = ingest_resolved_pair(self.db, self.path, {}, resolved())
|
||||
self.assertEqual(second["status"], "unchanged")
|
||||
self.assertEqual(second["pairs_added"], 0)
|
||||
pairs = self.db.list_prompt_pairs()["pairs"]
|
||||
self.assertEqual(len(pairs), 1)
|
||||
self.assertEqual(pairs[0]["image_count"], 1)
|
||||
self.assertEqual(pairs[0]["run_count"], 0)
|
||||
self.assertEqual(
|
||||
self.db.get_prompt_by_hash(generate_prompt_hash("cat"))["run_count"], 0
|
||||
)
|
||||
ingest_resolved_pair(self.db, self.path + "2", {}, resolved(negative="noise"))
|
||||
self.assertEqual(
|
||||
self.db.model.get_connection()
|
||||
.execute("SELECT count(*) FROM prompts")
|
||||
.fetchone()[0],
|
||||
3,
|
||||
)
|
||||
|
||||
def test_existing_image_enriched_in_place_and_metadata_preserved(self):
|
||||
pid = self.db.save_prompt(
|
||||
"cat", category="Personal", prompt_hash=generate_prompt_hash("cat")
|
||||
)
|
||||
image_id = self.db.link_image_to_prompt(pid, self.path)
|
||||
result = ingest_resolved_pair(
|
||||
self.db, self.path, {"prompt": {"a": 1}}, resolved()
|
||||
)
|
||||
self.assertEqual(result["status"], "enriched")
|
||||
self.assertEqual(result["image_id"], image_id)
|
||||
self.assertEqual(self.db.get_prompt_by_id(pid)["category"], "Personal")
|
||||
row = self.db.get_image_by_id(image_id)
|
||||
self.assertIn("a", row["prompt_metadata"])
|
||||
|
||||
def test_conflicting_image_is_not_reassigned_or_duplicated(self):
|
||||
pid = self.db.save_prompt("dog", prompt_hash=generate_prompt_hash("dog"))
|
||||
self.db.link_image_to_prompt(pid, self.path)
|
||||
result = ingest_resolved_pair(self.db, self.path, {}, resolved())
|
||||
self.assertEqual(result["status"], "conflict")
|
||||
self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
|
||||
self.assertEqual(
|
||||
self.db.model.get_connection()
|
||||
.execute("SELECT count(*) FROM generated_images")
|
||||
.fetchone()[0],
|
||||
1,
|
||||
)
|
||||
|
||||
def test_unknown_and_missing_negative_never_become_empty_pairs(self):
|
||||
for state in ("absent", "ambiguous", "unresolved"):
|
||||
value = resolved()
|
||||
value["state"] = state
|
||||
result = ingest_resolved_pair(self.db, self.path, {}, value)
|
||||
self.assertEqual(result["status"], state)
|
||||
self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
|
||||
result = ingest_resolved_pair(self.db, self.path, {}, resolved(negative=""))
|
||||
self.assertEqual(result["status"], "added")
|
||||
|
||||
def test_failure_after_inserts_rolls_back_all_rows(self):
|
||||
with patch.object(
|
||||
self.db, "_link_image_to_pair", side_effect=RuntimeError("injected")
|
||||
):
|
||||
with self.assertRaises(RuntimeError):
|
||||
ingest_resolved_pair(self.db, self.path, {}, resolved())
|
||||
conn = self.db.model.get_connection()
|
||||
for table in (
|
||||
"prompts",
|
||||
"prompt_pairs",
|
||||
"generated_images",
|
||||
"image_prompt_pairs",
|
||||
):
|
||||
self.assertEqual(
|
||||
conn.execute(f"SELECT count(*) FROM {table}").fetchone()[0], 0
|
||||
)
|
||||
|
||||
def test_distinct_paths_are_not_collapsed_and_effective_text_stays_in_provenance(
|
||||
self,
|
||||
):
|
||||
value = resolved("prefix cat suffix")
|
||||
value["positive"]["storage_text"] = "cat"
|
||||
a = ingest_resolved_pair(self.db, self.path, {}, value)
|
||||
b = ingest_resolved_pair(
|
||||
self.db, self.path.replace("image.png", "other/image.png"), {}, value
|
||||
)
|
||||
self.assertEqual(a["pair_id"], b["pair_id"])
|
||||
pair = self.db.get_prompt_pair(a["pair_id"])
|
||||
self.assertEqual(pair["positive_text"], "cat")
|
||||
self.assertEqual(pair["image_count"], 2)
|
||||
association = self.db.get_image_prompt_pairs(a["image_id"])[0]
|
||||
self.assertEqual(
|
||||
association["provenance"]["positive"]["text"], "prefix cat suffix"
|
||||
)
|
||||
|
||||
def test_image_lookup_exposes_its_recorded_pair_not_latest_negative(self):
|
||||
value = resolved("trigger cat", "blur")
|
||||
value["positive"]["storage_text"] = "cat"
|
||||
first = ingest_resolved_pair(self.db, self.path, {}, value)
|
||||
ingest_resolved_pair(self.db, self.path + "2", {}, resolved("cat", "noise"))
|
||||
info = self.db.get_image_prompt_info(self.path)
|
||||
self.assertEqual(info["pair"]["id"], first["pair_id"])
|
||||
self.assertEqual(info["pair"]["negative_text"], "blur")
|
||||
self.assertEqual(info["pair_context"]["positive_text"], "trigger cat")
|
||||
self.assertEqual(info["pair_context"]["negative_text"], "blur")
|
||||
|
||||
def test_overwritten_file_with_different_negative_reports_conflict(self):
|
||||
ingest_resolved_pair(self.db, self.path, {}, resolved())
|
||||
result = ingest_resolved_pair(
|
||||
self.db, self.path, {}, resolved(negative="noise")
|
||||
)
|
||||
self.assertEqual(result["status"], "conflict")
|
||||
self.assertEqual(self.db.list_prompt_pairs()["total"], 1)
|
||||
|
||||
def test_execution_evidence_wins_over_a_prior_inferred_context(self):
|
||||
inferred = resolved()
|
||||
for side in ("positive", "negative"):
|
||||
inferred[side]["provenance"]["text_source"] = "graph-inferred"
|
||||
ingest_resolved_pair(self.db, self.path, {}, inferred)
|
||||
execution = resolved("cat, lora_trigger", "blur")
|
||||
execution["positive"]["storage_text"] = "cat"
|
||||
for side in ("positive", "negative"):
|
||||
execution[side]["provenance"]["text_source"] = "execution"
|
||||
ingest_resolved_pair(self.db, self.path, {}, execution)
|
||||
info = self.db.get_image_prompt_info(self.path)
|
||||
self.assertEqual(info["pair_context"]["positive_text"], "cat, lora_trigger")
|
||||
self.assertEqual(info["pair_context"]["text_source"], "execution")
|
||||
@@ -0,0 +1,83 @@
|
||||
"""Actual sampler wires determine which component counts as its positive text."""
|
||||
|
||||
import unittest
|
||||
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.usage_tracking import queued_prompt_runs, queued_pair_runs
|
||||
|
||||
|
||||
def paired_graph(positive_slot=0, negative_slot=1):
|
||||
return {
|
||||
"1": {
|
||||
"class_type": "PromptManagerPair",
|
||||
"inputs": {"text": "cat", "negative_text": "blur"},
|
||||
},
|
||||
"2": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"model": ["3", 0],
|
||||
"positive": ["1", positive_slot],
|
||||
"negative": ["1", negative_slot],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class TestPairOutputRoles(unittest.TestCase):
|
||||
def test_swapped_conditioning_counts_actual_positive_component_only(self):
|
||||
graph = paired_graph(1, 0)
|
||||
self.assertEqual(
|
||||
queued_prompt_runs(graph), [(generate_prompt_hash("blur"), "blur")]
|
||||
)
|
||||
self.assertEqual(queued_pair_runs(graph), [])
|
||||
|
||||
def test_negative_string_encoded_as_positive_counts_actual_string_component(self):
|
||||
graph = paired_graph()
|
||||
graph["4"] = {"class_type": "CLIPTextEncode", "inputs": {"text": ["1", 3]}}
|
||||
graph["2"]["inputs"]["positive"] = ["4", 0]
|
||||
self.assertEqual(
|
||||
queued_prompt_runs(graph), [(generate_prompt_hash("blur"), "blur")]
|
||||
)
|
||||
self.assertEqual(queued_pair_runs(graph), [])
|
||||
|
||||
def test_two_sampler_positive_roles_count_each_component_once(self):
|
||||
graph = paired_graph()
|
||||
graph["5"] = {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"model": ["3", 0], "positive": ["1", 1], "negative": ["1", 0]},
|
||||
}
|
||||
self.assertEqual(
|
||||
{text for _, text in queued_prompt_runs(graph)}, {"cat", "blur"}
|
||||
)
|
||||
|
||||
|
||||
class TestPairRoleExecution(unittest.TestCase):
|
||||
def setUp(self):
|
||||
from tests.test_prompt_pair_node import TestPromptPairNode
|
||||
|
||||
TestPromptPairNode.setUp(self)
|
||||
|
||||
def test_new_swapped_component_consumes_its_queued_pending_use(self):
|
||||
self.node.pending_first_use.add(generate_prompt_hash("blur"))
|
||||
self.node.encode_prompt(
|
||||
self.clip, "cat", "blur", prompt=paired_graph(1, 0), unique_id="1"
|
||||
)
|
||||
pair = self.db.list_prompt_pairs()["pairs"][0]
|
||||
self.assertEqual(
|
||||
self.db.get_prompt_by_id(pair["negative_prompt_id"])["run_count"], 1
|
||||
)
|
||||
self.assertEqual(
|
||||
self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 0
|
||||
)
|
||||
self.assertEqual(pair["run_count"], 0)
|
||||
|
||||
def test_unknown_negative_link_used_as_positive_counts_at_execution(self):
|
||||
graph = paired_graph(1, 0)
|
||||
graph["1"]["inputs"]["negative_text"] = ["8", 0]
|
||||
graph["8"] = {"class_type": "UnknownBatch", "inputs": {}}
|
||||
self.node.encode_prompt(self.clip, "cat", "blur", prompt=graph, unique_id="1")
|
||||
pair = self.db.list_prompt_pairs()["pairs"][0]
|
||||
self.assertEqual(
|
||||
self.db.get_prompt_by_id(pair["negative_prompt_id"])["run_count"], 1
|
||||
)
|
||||
self.assertEqual(pair["run_count"], 0)
|
||||
@@ -0,0 +1,119 @@
|
||||
"""Restores preserve pair ownership and bounded, readable JSON provenance."""
|
||||
|
||||
import json
|
||||
import os
|
||||
import tempfile
|
||||
import unittest.mock
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.models import PromptModel
|
||||
from database.operations import PromptDatabase
|
||||
|
||||
|
||||
class TestPairRestoreSemantics(unittest.TestCase):
|
||||
def setUp(self):
|
||||
directory = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(directory.cleanup)
|
||||
self.addCleanup(assert_closed, self, directory.name)
|
||||
self.path = os.path.join(directory.name, "restore.db")
|
||||
self.db = PromptDatabase(self.path)
|
||||
self.addCleanup(self.db.close_all)
|
||||
self.pair_id = self.db.save_prompt_pair("portrait", "blurry")
|
||||
pair = self.db.get_prompt_pair(self.pair_id)
|
||||
self.image_id = self.db.link_image_to_prompt(
|
||||
pair["positive_prompt_id"], os.path.join(directory.name, "image.png")
|
||||
)
|
||||
|
||||
def raw_link(self, provenance, pair_id=None):
|
||||
with self.db.model.get_connection() as conn:
|
||||
conn.execute("DELETE FROM image_prompt_pairs")
|
||||
conn.execute(
|
||||
"INSERT INTO image_prompt_pairs(image_id,pair_id,provenance) "
|
||||
"VALUES(?,?,?)",
|
||||
(self.image_id, pair_id or self.pair_id, provenance),
|
||||
)
|
||||
conn.commit()
|
||||
|
||||
def test_restore_rejects_pair_with_different_image_positive_owner(self):
|
||||
other = self.db.save_prompt_pair("landscape", "blurry")
|
||||
self.raw_link("{}", other)
|
||||
ok, reason = PromptModel.verify_database_file(self.path)
|
||||
self.assertFalse(ok)
|
||||
self.assertIn("positive", reason.lower())
|
||||
|
||||
def test_restore_rejects_unreadable_nonobject_and_nonfinite_provenance(self):
|
||||
for value in (
|
||||
"not JSON",
|
||||
"[]",
|
||||
"null",
|
||||
"false",
|
||||
"42",
|
||||
'"text"',
|
||||
'{"value":NaN}',
|
||||
'{"value":Infinity}',
|
||||
'{"value":-Infinity}',
|
||||
'{"value":1e999}',
|
||||
b'{"blob":true}',
|
||||
):
|
||||
with self.subTest(value=value):
|
||||
self.raw_link(value)
|
||||
ok, reason = PromptModel.verify_database_file(self.path)
|
||||
self.assertFalse(ok)
|
||||
self.assertIn("provenance", reason.lower())
|
||||
|
||||
def test_restore_retains_valid_nested_object_provenance(self):
|
||||
provenance = {"format": "future", "nested": [{"empty": "", "value": 2.5}]}
|
||||
self.raw_link(json.dumps(provenance))
|
||||
self.assertEqual(PromptModel.verify_database_file(self.path), (True, "ok"))
|
||||
self.assertEqual(
|
||||
self.db.get_image_prompt_pairs(self.image_id)[0]["provenance"], provenance
|
||||
)
|
||||
|
||||
def test_restore_rejects_excessive_json_nesting_without_raising(self):
|
||||
self.raw_link('{"nested":' * 1500 + "0" + "}" * 1500)
|
||||
ok, reason = PromptModel.verify_database_file(self.path)
|
||||
self.assertFalse(ok)
|
||||
self.assertIn("provenance", reason.lower())
|
||||
|
||||
def test_restore_and_runtime_share_the_provenance_size_limit(self):
|
||||
provenance = {"large": "x" * (4_000_000 + 1)}
|
||||
with self.assertRaises(ValueError):
|
||||
self.db.link_image_to_pair(self.pair_id, self.image_id, provenance)
|
||||
self.raw_link(json.dumps(provenance))
|
||||
ok, reason = PromptModel.verify_database_file(self.path)
|
||||
self.assertFalse(ok)
|
||||
self.assertIn("provenance", reason.lower())
|
||||
|
||||
def test_runtime_normalizes_deep_json_encoding_error(self):
|
||||
provenance = {}
|
||||
child = provenance
|
||||
for _ in range(1500):
|
||||
child["nested"] = {}
|
||||
child = child["nested"]
|
||||
with self.assertRaises(ValueError):
|
||||
self.db.link_image_to_pair(self.pair_id, self.image_id, provenance)
|
||||
self.assertEqual(self.db.get_image_prompt_pairs(self.image_id), [])
|
||||
|
||||
def test_runtime_converts_encoder_recursion_failure_to_validation_error(self):
|
||||
with unittest.mock.patch(
|
||||
"database.pairs.json.dumps", side_effect=RecursionError("depth")
|
||||
):
|
||||
with self.assertRaisesRegex(ValueError, "JSON"):
|
||||
self.db.link_image_to_pair(self.pair_id, self.image_id, {})
|
||||
|
||||
def test_restore_converts_parser_recursion_failure_to_rejection(self):
|
||||
self.raw_link("{}")
|
||||
with unittest.mock.patch(
|
||||
"database.models.json.loads", side_effect=RecursionError("depth")
|
||||
):
|
||||
ok, reason = PromptModel.verify_database_file(self.path)
|
||||
self.assertFalse(ok)
|
||||
self.assertIn("provenance", reason.lower())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,113 @@
|
||||
"""Metadata scanning enriches known pairs and conservatively leaves conflicts."""
|
||||
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
from tests.test_api_admin import AdminAPITestCase
|
||||
|
||||
|
||||
def graph(positive="cat", negative="blur"):
|
||||
return {
|
||||
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": positive}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": negative}},
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class TestPairScan(AdminAPITestCase):
|
||||
def test_pair_scan_repeated_is_unchanged_and_empty_negative_is_supported(self):
|
||||
path = self.output_dir / "pair.png"
|
||||
meta = {"prompt": json.dumps(graph(negative=""))}
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, meta), "added")
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, meta), "unchanged")
|
||||
pairs = self.api.db.list_prompt_pairs()["pairs"]
|
||||
self.assertEqual(len(pairs), 1)
|
||||
self.assertEqual(pairs[0]["negative_text"], "")
|
||||
self.assertEqual(pairs[0]["run_count"], 0)
|
||||
|
||||
def test_ambiguous_sampler_scan_does_not_save_a_guess(self):
|
||||
g = graph()
|
||||
g["4"] = {**g["3"]}
|
||||
meta = {"prompt": json.dumps(g)}
|
||||
self.assertEqual(
|
||||
self.api._ingest_scanned_file(self.output_dir / "multi.png", meta),
|
||||
"ambiguous",
|
||||
)
|
||||
self.assertEqual(self.api.db.list_prompt_pairs()["total"], 0)
|
||||
self.assertEqual(self.api.db.get_recent_prompts()["prompts"], [])
|
||||
|
||||
def test_batch_exposes_pair_counts_and_actual_unchanged_status(self):
|
||||
meta = {"parameters": "cat\nNegative prompt: blur\nSteps: 20, Sampler: Euler"}
|
||||
batch = [(Path(self.output_dir / "a.png"), meta)]
|
||||
first = self.api._ingest_scan_batch_sync(batch)
|
||||
self.assertEqual(first["pairs_added"], 1)
|
||||
self.assertEqual(first["pair_links_added"], 1)
|
||||
second = self.api._ingest_scan_batch_sync(batch)
|
||||
self.assertEqual(second["unchanged"], 1)
|
||||
self.assertEqual(second["linked"], 0)
|
||||
self.assertEqual(second["pairs_added"], 0)
|
||||
|
||||
def test_legacy_single_prompt_second_pass_not_reported_as_new_link(self):
|
||||
meta = {"prompt": json.dumps({"1": graph()["1"]})}
|
||||
path = self.output_dir / "old.png"
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, meta), "added")
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, meta), "unchanged")
|
||||
|
||||
async def test_cancel_finishes_at_safe_boundary_and_restart_is_idempotent(self):
|
||||
from py.api.scan_job import ScanJob
|
||||
|
||||
job = ScanJob()
|
||||
self.api._scan_job = job
|
||||
response = await self.client.post("/prompt_manager/scan/cancel")
|
||||
self.assertEqual(response.status, 200)
|
||||
self.assertTrue((await response.json())["cancellation_requested"])
|
||||
counts = self.api._ingest_scan_batch_sync(
|
||||
[(self.output_dir / "cancel.png", {"prompt": json.dumps(graph())})]
|
||||
)
|
||||
self.assertEqual(counts["processed"], 0)
|
||||
self.assertEqual(self.api.db.list_prompt_pairs()["total"], 0)
|
||||
job.finish()
|
||||
self.api._scan_job = ScanJob()
|
||||
first = self.api._ingest_scan_batch_sync(
|
||||
[(self.output_dir / "cancel.png", {"prompt": json.dumps(graph())})]
|
||||
)
|
||||
self.assertEqual(first["pairs_added"], 1)
|
||||
self.assertEqual(
|
||||
self.api._ingest_scan_batch_sync(
|
||||
[(self.output_dir / "cancel.png", {"prompt": json.dumps(graph())})]
|
||||
)["unchanged"],
|
||||
1,
|
||||
)
|
||||
|
||||
def test_legacy_found_outcome_is_not_counted_twice(self):
|
||||
result = self.api._ingest_scan_batch_sync(
|
||||
[(self.output_dir / "blank.png", {"parameters": "Negative prompt: blur"})]
|
||||
)
|
||||
self.assertEqual(result["processed"], 1)
|
||||
self.assertEqual(result["found"], 1)
|
||||
|
||||
def test_ambiguous_a1111_does_not_fall_back_to_legacy_guess(self):
|
||||
meta = {
|
||||
"parameters": (
|
||||
"cat\nNegative prompt: blur\nNegative prompt: noise"
|
||||
"\nSteps: 20, Sampler: Euler"
|
||||
)
|
||||
}
|
||||
self.assertEqual(
|
||||
self.api._ingest_scanned_file(self.output_dir / "ambiguous.png", meta),
|
||||
"ambiguous",
|
||||
)
|
||||
self.assertEqual(self.api.db.get_recent_prompts()["prompts"], [])
|
||||
|
||||
def test_raw_png_workflow_is_stored_as_object_and_repeat_is_unchanged(self):
|
||||
path = self.output_dir / "workflow.png"
|
||||
workflow = {"nodes": [{"id": 1, "type": "PromptManagerPair"}]}
|
||||
metadata = {"prompt": json.dumps(graph()), "workflow": json.dumps(workflow)}
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, metadata), "added")
|
||||
self.assertEqual(
|
||||
self.api.db.get_image_prompt_info(str(path))["workflow_data"], workflow
|
||||
)
|
||||
self.assertEqual(self.api._ingest_scanned_file(path, metadata), "unchanged")
|
||||
@@ -0,0 +1,325 @@
|
||||
"""Pair schema installation and restore failures use disposable databases only."""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import unittest
|
||||
from contextlib import closing
|
||||
from pathlib import Path
|
||||
from unittest.mock import patch
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.models import PromptModel
|
||||
|
||||
|
||||
class PairSchemaTests(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.temp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.temp.cleanup)
|
||||
self.addCleanup(assert_closed, self, self.temp.name)
|
||||
self.addCleanup(PromptModel.reset_schema_cache)
|
||||
self.models = []
|
||||
self.addCleanup(self.close_models)
|
||||
self.path = str(Path(self.temp.name) / "live.db")
|
||||
|
||||
def close_models(self):
|
||||
for model in self.models:
|
||||
model.close_all()
|
||||
|
||||
def model(self, path=None):
|
||||
model = PromptModel(path or self.path)
|
||||
self.models.append(model)
|
||||
return model
|
||||
|
||||
def populate(self, model):
|
||||
with model.get_connection() as conn:
|
||||
conn.execute(
|
||||
"INSERT INTO prompts(id,text,hash,category,notes,run_count) VALUES"
|
||||
"(7,'positive','p','keep','notes',12)"
|
||||
)
|
||||
conn.execute("INSERT INTO prompts(id,text,hash) VALUES(9,'negative','n')")
|
||||
conn.execute(
|
||||
"INSERT INTO generated_images(id,prompt_id,image_path,filename,fil"
|
||||
"e_path) VALUES(4,7,'/tmp/image.png','image.png','/tmp/image.png')"
|
||||
)
|
||||
conn.execute("INSERT INTO tags(id,name) VALUES(3,'keep-tag')")
|
||||
conn.execute("INSERT INTO prompt_tags VALUES(7,3)")
|
||||
|
||||
def legacy(self):
|
||||
model = self.model(str(Path(self.temp.name) / "legacy.db"))
|
||||
self.populate(model)
|
||||
with model.get_connection() as conn:
|
||||
for table in ("image_prompt_pairs", "prompt_pair_tags", "prompt_pairs"):
|
||||
conn.execute(f"DROP TABLE IF EXISTS {table}")
|
||||
return model
|
||||
|
||||
def test_restore_scenarios_close_all_sqlite_handles(self):
|
||||
connections = []
|
||||
connect = sqlite3.connect
|
||||
|
||||
def tracked_connect(*args, **kwargs):
|
||||
connection = connect(*args, **kwargs)
|
||||
connections.append(connection)
|
||||
return connection
|
||||
|
||||
scenarios = (
|
||||
"test_pair_ddl_failure_rolls_back_and_retry_succeeds",
|
||||
"test_restore_legacy_wal_migrates_stage_and_preserves_safety_backup",
|
||||
"test_restore_migration_failure_keeps_original_and_source_untouched",
|
||||
"test_verification_rejects_partial_pair_schema_and_orphan_pair",
|
||||
)
|
||||
try:
|
||||
with patch.object(sqlite3, "connect", side_effect=tracked_connect):
|
||||
result = unittest.TestResult()
|
||||
unittest.TestSuite(type(self)(name) for name in scenarios).run(result)
|
||||
self.assertTrue(result.wasSuccessful(), result.errors + result.failures)
|
||||
for index, connection in enumerate(connections):
|
||||
with self.subTest(connection=index):
|
||||
with self.assertRaises(sqlite3.ProgrammingError):
|
||||
connection.execute("SELECT 1")
|
||||
finally:
|
||||
for connection in connections:
|
||||
connection.close()
|
||||
|
||||
def test_fresh_schema_constraints_and_foreign_keys(self):
|
||||
model = self.model()
|
||||
self.populate(model)
|
||||
with model.get_connection() as conn:
|
||||
conn.execute(
|
||||
"INSERT INTO prompt_pairs(id,positive_prompt_id,negative_prompt_id"
|
||||
") VALUES(2,7,9)"
|
||||
)
|
||||
conn.execute("INSERT INTO prompt_pairs(positive_prompt_id) VALUES(7)")
|
||||
conn.execute("INSERT INTO prompt_pair_tags VALUES(2,3)")
|
||||
conn.execute(
|
||||
"INSERT INTO image_prompt_pairs(pair_id,image_id,provenance) VALUE"
|
||||
"S(2,4,'{}')"
|
||||
)
|
||||
for sql in (
|
||||
"INSERT INTO prompt_pairs(positive_prompt_id,negative_prompt_id) V"
|
||||
"ALUES(7,9)",
|
||||
"INSERT INTO prompt_pairs(positive_prompt_id) VALUES(7)",
|
||||
"INSERT INTO prompt_pairs(positive_prompt_id) VALUES(404)",
|
||||
"INSERT INTO prompt_pairs(positive_prompt_id,rating) VALUES(9,6)",
|
||||
"INSERT INTO prompt_pairs(positive_prompt_id,archived) VALUES(9,2)",
|
||||
"DELETE FROM prompts WHERE id=9",
|
||||
"DELETE FROM prompt_pairs WHERE id=2",
|
||||
"INSERT INTO image_prompt_pairs(pair_id,image_id,provenance) VALUE"
|
||||
"S(2,4,'{}')",
|
||||
):
|
||||
with self.subTest(sql=sql), self.assertRaises(sqlite3.IntegrityError):
|
||||
with model.get_connection() as conn:
|
||||
conn.execute(sql)
|
||||
with model.get_connection() as conn:
|
||||
conn.execute(
|
||||
"INSERT INTO prompt_pairs(positive_prompt_id,negative_prompt_id) V"
|
||||
"ALUES(9,7)"
|
||||
)
|
||||
conn.execute("DELETE FROM generated_images WHERE id=4")
|
||||
self.assertEqual(
|
||||
conn.execute("SELECT COUNT(*) FROM image_prompt_pairs").fetchone()[0], 0
|
||||
)
|
||||
self.assertEqual(conn.execute("PRAGMA foreign_key_check").fetchall(), [])
|
||||
|
||||
def test_legacy_upgrade_preserves_ids_metadata_images_counts_and_repeats(self):
|
||||
legacy = self.legacy()
|
||||
path = legacy.db_path
|
||||
legacy.close_all()
|
||||
PromptModel.reset_schema_cache()
|
||||
for _ in range(2):
|
||||
model = self.model(path)
|
||||
with model.get_connection() as conn:
|
||||
self.assertEqual(
|
||||
tuple(
|
||||
conn.execute(
|
||||
"SELECT id,text,category,notes,run_count FROM prompts WHER"
|
||||
"E id=7"
|
||||
).fetchone()
|
||||
),
|
||||
(7, "positive", "keep", "notes", 12),
|
||||
)
|
||||
self.assertEqual(
|
||||
conn.execute("SELECT id FROM generated_images").fetchone()[0], 4
|
||||
)
|
||||
self.assertEqual(
|
||||
tuple(conn.execute("SELECT * FROM prompt_tags").fetchone()), (7, 3)
|
||||
)
|
||||
self.assertEqual(
|
||||
conn.execute("SELECT COUNT(*) FROM prompt_pairs").fetchone()[0], 0
|
||||
)
|
||||
model.close_all()
|
||||
PromptModel.reset_schema_cache()
|
||||
|
||||
def test_pair_ddl_failure_rolls_back_and_retry_succeeds(self):
|
||||
with patch.object(
|
||||
PromptModel,
|
||||
"_create_pair_indexes",
|
||||
side_effect=sqlite3.OperationalError("injected index failure"),
|
||||
):
|
||||
with self.assertRaisesRegex(sqlite3.OperationalError, "injected"):
|
||||
self.model()
|
||||
with closing(sqlite3.connect(self.path)) as conn, conn:
|
||||
self.assertEqual(
|
||||
conn.execute(
|
||||
"SELECT name FROM sqlite_master WHERE name IN ('prompt_pairs','pro"
|
||||
"mpt_pair_tags','image_prompt_pairs')"
|
||||
).fetchall(),
|
||||
[],
|
||||
)
|
||||
model = self.model()
|
||||
self.assertTrue(PromptModel.verify_database_file(model.db_path)[0])
|
||||
|
||||
def test_verification_rejects_partial_pair_schema_and_orphan_pair(self):
|
||||
model = self.model()
|
||||
self.populate(model)
|
||||
with model.get_connection() as conn:
|
||||
conn.execute(
|
||||
"INSERT INTO prompt_pairs(positive_prompt_id,negative_prompt_id) V"
|
||||
"ALUES(7,9)"
|
||||
)
|
||||
model.close_all()
|
||||
with closing(sqlite3.connect(model.db_path)) as conn, conn:
|
||||
conn.execute("PRAGMA foreign_keys=OFF")
|
||||
conn.execute("UPDATE prompt_pairs SET negative_prompt_id=404")
|
||||
ok, reason = PromptModel.verify_database_file(model.db_path)
|
||||
self.assertFalse(ok)
|
||||
self.assertIn("foreign key", reason.lower())
|
||||
with closing(sqlite3.connect(model.db_path)) as conn, conn:
|
||||
conn.execute("DROP TABLE prompt_pair_tags")
|
||||
ok, reason = PromptModel.verify_database_file(model.db_path)
|
||||
self.assertFalse(ok)
|
||||
self.assertIn("pair", reason.lower())
|
||||
|
||||
def test_restore_migration_failure_keeps_original_and_source_untouched(self):
|
||||
live = self.model()
|
||||
self.populate(live)
|
||||
legacy = self.legacy()
|
||||
with live.get_connection() as conn:
|
||||
conn.execute("UPDATE prompts SET text='original live text' WHERE id=7")
|
||||
with patch.object(
|
||||
PromptModel,
|
||||
"_migrate_prompt_pairs",
|
||||
side_effect=sqlite3.OperationalError("migration failed"),
|
||||
):
|
||||
with self.assertRaisesRegex(sqlite3.OperationalError, "migration failed"):
|
||||
live.restore_from_file(legacy.db_path)
|
||||
with live.get_connection() as conn:
|
||||
self.assertEqual(
|
||||
conn.execute("SELECT text FROM prompts WHERE id=7").fetchone()[0],
|
||||
"original live text",
|
||||
)
|
||||
self.assertEqual(
|
||||
conn.execute("SELECT run_count FROM prompts WHERE id=7").fetchone()[0],
|
||||
12,
|
||||
)
|
||||
self.assertEqual(
|
||||
conn.execute("SELECT COUNT(*) FROM generated_images").fetchone()[0], 1
|
||||
)
|
||||
with closing(sqlite3.connect(legacy.db_path)) as conn, conn:
|
||||
self.assertIsNone(
|
||||
conn.execute(
|
||||
"SELECT name FROM sqlite_master WHERE name='prompt_pairs'"
|
||||
).fetchone()
|
||||
)
|
||||
|
||||
def test_restore_legacy_wal_migrates_stage_and_preserves_safety_backup(self):
|
||||
live = self.model()
|
||||
with live.get_connection() as conn:
|
||||
conn.execute("INSERT INTO prompts(text,hash) VALUES('original','original')")
|
||||
legacy = self.legacy()
|
||||
with legacy.get_connection() as conn:
|
||||
conn.execute("INSERT INTO prompts(text,hash) VALUES('wal-only','wal-only')")
|
||||
self.assertGreater(os.path.getsize(legacy.db_path + "-wal"), 0)
|
||||
backup = live.restore_from_file(legacy.db_path)
|
||||
self.assertTrue(Path(backup).is_file())
|
||||
with closing(sqlite3.connect(backup)) as conn, conn:
|
||||
self.assertEqual(
|
||||
conn.execute("SELECT text FROM prompts").fetchall(), [("original",)]
|
||||
)
|
||||
with live.get_connection() as conn:
|
||||
self.assertEqual(
|
||||
conn.execute("SELECT COUNT(*) FROM prompts").fetchone()[0], 3
|
||||
)
|
||||
self.assertEqual(
|
||||
conn.execute("SELECT COUNT(*) FROM prompt_pairs").fetchone()[0], 0
|
||||
)
|
||||
self.assertTrue(PromptModel.verify_database_file(live.db_path)[0])
|
||||
|
||||
def test_restore_staged_verification_failure_does_not_replace_live_pages(self):
|
||||
live = self.model()
|
||||
self.populate(live)
|
||||
source = self.legacy()
|
||||
with patch.object(
|
||||
PromptModel,
|
||||
"verify_database_file",
|
||||
side_effect=[(True, "ok"), (False, "staged validation failed")],
|
||||
):
|
||||
with self.assertRaisesRegex(ValueError, "staged validation failed"):
|
||||
live.restore_from_file(source.db_path)
|
||||
with live.get_connection() as conn:
|
||||
self.assertEqual(
|
||||
conn.execute("SELECT run_count FROM prompts WHERE id=7").fetchone()[0],
|
||||
12,
|
||||
)
|
||||
|
||||
def test_verification_rejects_missing_null_uniqueness(self):
|
||||
model = self.model()
|
||||
with model.get_connection() as conn:
|
||||
conn.execute("DROP INDEX idx_prompt_pairs_positive_only")
|
||||
ok, reason = PromptModel.verify_database_file(model.db_path)
|
||||
self.assertFalse(ok)
|
||||
self.assertIn("uniqueness", reason)
|
||||
|
||||
def test_restore_round_trips_pair_metadata_tags_and_provenance(self):
|
||||
source = self.model(str(Path(self.temp.name) / "paired.db"))
|
||||
self.populate(source)
|
||||
with source.get_connection() as conn:
|
||||
conn.execute(
|
||||
"INSERT INTO prompt_pairs(id,positive_prompt_id,negative_prompt_id"
|
||||
",title,notes,run_count,archived) VALUES(6,7,9,'pair title','pair "
|
||||
"notes',19,1)"
|
||||
)
|
||||
conn.execute("INSERT INTO prompt_pair_tags VALUES(6,3)")
|
||||
conn.execute(
|
||||
"INSERT INTO image_prompt_pairs(id,pair_id,image_id,provenance) VA"
|
||||
'LUES(8,6,4,\'{"execution":"kept"}\')'
|
||||
)
|
||||
live = self.model()
|
||||
live.restore_from_file(source.db_path)
|
||||
with live.get_connection() as conn:
|
||||
self.assertEqual(
|
||||
tuple(
|
||||
conn.execute(
|
||||
"SELECT id,title,notes,run_count,archived FROM prompt_pairs"
|
||||
).fetchone()
|
||||
),
|
||||
(6, "pair title", "pair notes", 19, 1),
|
||||
)
|
||||
self.assertEqual(
|
||||
tuple(conn.execute("SELECT * FROM prompt_pair_tags").fetchone()), (6, 3)
|
||||
)
|
||||
self.assertEqual(
|
||||
conn.execute(
|
||||
"SELECT provenance FROM image_prompt_pairs WHERE id=8"
|
||||
).fetchone()[0],
|
||||
'{"execution":"kept"}',
|
||||
)
|
||||
|
||||
def test_restore_rejects_trigger_before_running_migrations(self):
|
||||
live = self.model()
|
||||
source = self.legacy()
|
||||
with source.get_connection() as conn:
|
||||
conn.execute(
|
||||
"CREATE TRIGGER dangerous AFTER INSERT ON prompts BEGIN DELETE FRO"
|
||||
"M prompts; END"
|
||||
)
|
||||
with self.assertRaisesRegex(ValueError, "trigger"):
|
||||
live.restore_from_file(source.db_path)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,107 @@
|
||||
"""Global tag maintenance preserves pair and standalone tag relationships."""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
|
||||
|
||||
class TestPairTagMaintenance(unittest.TestCase):
|
||||
def setUp(self):
|
||||
directory = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(directory.cleanup)
|
||||
self.addCleanup(assert_closed, self, directory.name)
|
||||
self.db = PromptDatabase(os.path.join(directory.name, "pair-tags.db"))
|
||||
self.addCleanup(self.db.close_all)
|
||||
|
||||
def pair(self, positive: str, tags: list) -> int:
|
||||
return self.db.save_prompt_pair(positive, "blurry", tags=tags)
|
||||
|
||||
def tags(self, pair_id: int) -> list:
|
||||
return self.db.get_prompt_pair(pair_id)["tags"]
|
||||
|
||||
def test_rename_to_existing_tag_preserves_pair_links_and_deduplicates(self):
|
||||
first = self.pair("portrait", ["old", "new", "keep"])
|
||||
second = self.pair("landscape", ["old"])
|
||||
third = self.pair("studio", ["unrelated"])
|
||||
prompt = self.db.save_prompt("standalone", tags=["old", "new"])
|
||||
result = self.db.rename_tag_all_prompts("old", "new")
|
||||
self.assertEqual(result["affected_count"], 1)
|
||||
self.assertEqual(self.tags(first), ["keep", "new"])
|
||||
self.assertEqual(self.tags(second), ["new"])
|
||||
self.assertEqual(self.tags(third), ["unrelated"])
|
||||
self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["new"])
|
||||
self.assertEqual(self.db.list_prompt_pairs(tags=["new"])["total"], 2)
|
||||
self.assertNotIn("old", self.db.get_all_tags())
|
||||
|
||||
def test_rename_pair_only_tag_to_new_name_preserves_links(self):
|
||||
pair = self.pair("portrait", ["old"])
|
||||
self.db.rename_tag_all_prompts("old", "new")
|
||||
self.assertEqual(self.tags(pair), ["new"])
|
||||
|
||||
def test_merge_pair_only_source_preserves_its_links(self):
|
||||
pair = self.pair("portrait", ["source"])
|
||||
other = self.pair("landscape", ["target"])
|
||||
self.db.merge_tags(["source"], "target")
|
||||
self.assertEqual(self.tags(pair), ["target"])
|
||||
self.assertEqual(self.tags(other), ["target"])
|
||||
self.assertEqual(self.db.list_prompt_pairs(tags=["target"])["total"], 2)
|
||||
|
||||
def test_merge_multiple_sources_preserves_shared_and_unrelated_tags(self):
|
||||
pair = self.pair("portrait", ["first", "second", "target", "keep"])
|
||||
other = self.pair("landscape", ["second"])
|
||||
prompt = self.db.save_prompt("standalone", tags=["first", "second"])
|
||||
self.db.merge_tags(["first", "second"], "target")
|
||||
self.assertEqual(self.tags(pair), ["keep", "target"])
|
||||
self.assertEqual(self.tags(other), ["target"])
|
||||
self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["target"])
|
||||
self.assertEqual(self.db.list_prompt_pairs(tags=["target"])["total"], 2)
|
||||
|
||||
def test_self_rename_does_not_remove_tag_or_relationships(self):
|
||||
pair = self.pair("portrait", ["same"])
|
||||
prompt = self.db.save_prompt("standalone", tags=["same"])
|
||||
self.db.rename_tag_all_prompts("same", "same")
|
||||
self.assertEqual(self.tags(pair), ["same"])
|
||||
self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["same"])
|
||||
|
||||
def test_merge_skips_target_even_when_it_is_in_source_list(self):
|
||||
pair = self.pair("portrait", ["target", "source"])
|
||||
prompt = self.db.save_prompt("standalone", tags=["target", "source"])
|
||||
self.db.merge_tags(["target", "source", "target"], "target")
|
||||
self.assertEqual(self.tags(pair), ["target"])
|
||||
self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["target"])
|
||||
self.assertIn("target", self.db.get_all_tags())
|
||||
|
||||
def test_explicit_global_delete_removes_pair_tags_and_keeps_other_tags(self):
|
||||
pair = self.pair("portrait", ["remove", "keep"])
|
||||
prompt = self.db.save_prompt("standalone", tags=["remove", "keep"])
|
||||
self.db.delete_tag_all_prompts("remove")
|
||||
self.assertEqual(self.tags(pair), ["keep"])
|
||||
self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["keep"])
|
||||
self.assertIsNotNone(self.db.get_prompt_pair(pair))
|
||||
|
||||
def test_pair_link_failure_rolls_back_standalone_merge_and_source_delete(self):
|
||||
pair = self.pair("portrait", ["source"])
|
||||
prompt = self.db.save_prompt("standalone", tags=["source"])
|
||||
conn = self.db.model.get_connection()
|
||||
conn.execute(
|
||||
"CREATE TRIGGER reject_pair_tag_change BEFORE UPDATE "
|
||||
"ON prompt_pair_tags BEGIN SELECT RAISE(ABORT,'pair tag failure'); END"
|
||||
)
|
||||
conn.commit()
|
||||
with self.assertRaises(sqlite3.IntegrityError):
|
||||
self.db.merge_tags(["source"], "target")
|
||||
self.assertEqual(self.tags(pair), ["source"])
|
||||
self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["source"])
|
||||
self.assertNotIn("target", self.db.get_all_tags())
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,254 @@
|
||||
"""Paired usage is counted at dequeue, including cached node executions."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import Mock
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from utils import usage_tracking
|
||||
from utils.prompt_graph import run_prompt_nodes
|
||||
from utils.usage_tracking import (
|
||||
PendingFirstUse,
|
||||
QueuedRuns,
|
||||
handle_queued_prompt,
|
||||
on_dequeue,
|
||||
)
|
||||
|
||||
try:
|
||||
from tests.test_prompt_pair_node import pair_graph
|
||||
except ImportError:
|
||||
from test_prompt_pair_node import pair_graph
|
||||
|
||||
|
||||
class TestPairUsageTracking(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
self.addCleanup(assert_closed, self, self.tmp.name)
|
||||
self.db = PromptDatabase(os.path.join(self.tmp.name, "pairs.db"))
|
||||
self.addCleanup(self.db.close_all)
|
||||
self.queued = QueuedRuns()
|
||||
self.pair_queued = QueuedRuns()
|
||||
self.pending = PendingFirstUse()
|
||||
self.pair_pending = PendingFirstUse()
|
||||
|
||||
def submit(self, graph):
|
||||
return handle_queued_prompt(
|
||||
{"prompt": graph}, self.queued, pair_queued=self.pair_queued
|
||||
)
|
||||
|
||||
def dequeue(self, request):
|
||||
on_dequeue(
|
||||
request["prompt_id"],
|
||||
self.queued,
|
||||
lambda: self.db,
|
||||
self.pending,
|
||||
pair_queued=self.pair_queued,
|
||||
pair_pending=self.pair_pending,
|
||||
)
|
||||
|
||||
def node(self):
|
||||
from prompt_manager_pair import PromptManagerPair
|
||||
|
||||
node = PromptManagerPair.__new__(PromptManagerPair)
|
||||
node.db = self.db
|
||||
node.logger = logging.getLogger("test.pair_usage")
|
||||
node.pending_first_use = self.pending
|
||||
node.pending_pair_first_use = self.pair_pending
|
||||
node._inject_lora_trigger_words = lambda text: text
|
||||
node.prompt_tracker = Mock()
|
||||
node.comfyui_integration = Mock()
|
||||
return node
|
||||
|
||||
def execute(self, graph, text="cat", negative="blur"):
|
||||
return self.node().encode_prompt(
|
||||
Mock(), text, negative_text=negative, prompt=graph, unique_id="1"
|
||||
)
|
||||
|
||||
def test_existing_pair_and_positive_count_on_every_dequeue_without_encoding(self):
|
||||
pid = self.db.save_prompt_pair("cat", "blur")
|
||||
for _ in range(3):
|
||||
self.dequeue(self.submit(pair_graph()))
|
||||
pair = self.db.get_prompt_pair(pid)
|
||||
self.assertEqual(pair["run_count"], 3)
|
||||
self.assertEqual(
|
||||
self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 3
|
||||
)
|
||||
self.assertEqual(
|
||||
self.db.get_prompt_by_id(pair["negative_prompt_id"])["run_count"], 0
|
||||
)
|
||||
|
||||
def test_validation_rejection_never_counts_or_saves(self):
|
||||
self.submit(pair_graph())
|
||||
self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
|
||||
self.assertEqual(self.db.search_prompts(), [])
|
||||
|
||||
def test_new_pair_pending_uses_are_consumed_once_when_first_encoding_saves(self):
|
||||
graph = pair_graph()
|
||||
self.dequeue(self.submit(graph))
|
||||
self.dequeue(self.submit(graph))
|
||||
self.execute(graph)
|
||||
pair = self.db.list_prompt_pairs()["pairs"][0]
|
||||
self.assertEqual(pair["run_count"], 2)
|
||||
self.assertEqual(
|
||||
self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 2
|
||||
)
|
||||
self.execute(graph)
|
||||
self.assertEqual(self.db.get_prompt_pair(pair["id"])["run_count"], 2)
|
||||
|
||||
def test_negative_only_changes_create_distinct_pair_counts_and_positive_counts(
|
||||
self,
|
||||
):
|
||||
self.db.save_prompt_pair("cat", "blur")
|
||||
self.db.save_prompt_pair("cat", "watermark")
|
||||
self.dequeue(self.submit(pair_graph()))
|
||||
self.dequeue(self.submit(pair_graph(negative="watermark")))
|
||||
self.assertEqual(
|
||||
[p["run_count"] for p in self.db.list_prompt_pairs()["pairs"]], [1, 1]
|
||||
)
|
||||
positive = self.db.get_prompt_by_hash(
|
||||
usage_tracking.generate_prompt_hash("cat")
|
||||
)
|
||||
self.assertEqual(positive["run_count"], 2)
|
||||
|
||||
def test_negative_output_slots_are_not_positive_prompt_sources(self):
|
||||
for slot in (1, 3):
|
||||
with self.subTest(slot=slot):
|
||||
self.assertEqual(
|
||||
run_prompt_nodes(pair_graph(positive_slot=slot, negative_slot=0)),
|
||||
[],
|
||||
)
|
||||
|
||||
def test_mismatched_negative_does_not_count_unused_intrinsic_pair(self):
|
||||
pid = self.db.save_prompt_pair("cat", "blur")
|
||||
graph = pair_graph()
|
||||
graph["2"] = {
|
||||
"class_type": "CLIPTextEncode",
|
||||
"inputs": {"text": "another negative"},
|
||||
}
|
||||
graph["5"]["inputs"]["negative"] = ["2", 0]
|
||||
self.dequeue(self.submit(graph))
|
||||
self.execute(graph)
|
||||
self.assertEqual(self.db.get_prompt_pair(pid)["run_count"], 0)
|
||||
self.assertEqual(
|
||||
self.db.get_prompt_by_hash(usage_tracking.generate_prompt_hash("cat"))[
|
||||
"run_count"
|
||||
],
|
||||
1,
|
||||
)
|
||||
|
||||
def test_unknown_batch_text_counts_each_execution_once_and_tracks_fifo(self):
|
||||
graph = pair_graph(text=["2", 0])
|
||||
graph["2"] = {"class_type": "PromptSearchList", "inputs": {}}
|
||||
self.submit(graph)
|
||||
node = self.node()
|
||||
node.encode_prompt(
|
||||
Mock(), "batch cat", negative_text="blur", prompt=graph, unique_id="1"
|
||||
)
|
||||
pair = self.db.list_prompt_pairs()["pairs"][0]
|
||||
self.assertEqual(pair["run_count"], 1)
|
||||
self.assertEqual(
|
||||
self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 1
|
||||
)
|
||||
self.assertTrue(
|
||||
node.prompt_tracker.set_current_prompt.call_args.kwargs["push_to_queue"]
|
||||
)
|
||||
|
||||
def test_standalone_pair_counts_cached_dequeues(self):
|
||||
pid = self.db.save_prompt_pair("cat", "")
|
||||
graph = {
|
||||
"1": {
|
||||
"class_type": "PromptManagerPair",
|
||||
"inputs": {"text": "cat", "negative_text": ""},
|
||||
}
|
||||
}
|
||||
self.dequeue(self.submit(graph))
|
||||
self.dequeue(self.submit(graph))
|
||||
self.assertEqual(self.db.get_prompt_pair(pid)["run_count"], 2)
|
||||
|
||||
def test_paired_string_outputs_through_encoders_keep_roles(self):
|
||||
from utils.prompt_graph import intrinsic_pair_nodes
|
||||
|
||||
graph = pair_graph()
|
||||
graph["2"] = {"class_type": "CLIPTextEncode", "inputs": {"text": ["1", 2]}}
|
||||
graph["3"] = {"class_type": "CLIPTextEncode", "inputs": {"text": ["1", 3]}}
|
||||
graph["4"] = {"class_type": "Reroute", "inputs": {"value": ["2", 0]}}
|
||||
graph["5"]["inputs"].update(positive=["4", 0], negative=["3", 0])
|
||||
self.assertEqual(intrinsic_pair_nodes(graph), ["1"])
|
||||
pid = self.db.save_prompt_pair("cat", "blur")
|
||||
self.dequeue(self.submit(graph))
|
||||
self.assertEqual(self.db.get_prompt_pair(pid)["run_count"], 1)
|
||||
graph["3"]["inputs"]["text"] = ["1", 2]
|
||||
self.assertEqual(intrinsic_pair_nodes(graph), [])
|
||||
|
||||
def test_two_samplers_using_same_pair_count_one_run_and_dequeue_is_idempotent(self):
|
||||
graph = pair_graph()
|
||||
graph["6"] = {**graph["5"]}
|
||||
pid = self.db.save_prompt_pair("cat", "blur")
|
||||
request = self.submit(graph)
|
||||
self.dequeue(request)
|
||||
self.dequeue(request)
|
||||
self.assertEqual(self.db.get_prompt_pair(pid)["run_count"], 1)
|
||||
|
||||
def test_new_negative_with_existing_positive_counts_pending_pair_only_once(self):
|
||||
self.db.save_prompt_pair("cat", "old negative")
|
||||
graph = pair_graph(negative="new negative")
|
||||
self.dequeue(self.submit(graph))
|
||||
self.execute(graph, negative="new negative")
|
||||
pair = next(
|
||||
p
|
||||
for p in self.db.list_prompt_pairs()["pairs"]
|
||||
if p["negative_text"] == "new negative"
|
||||
)
|
||||
self.assertEqual(pair["run_count"], 1)
|
||||
self.assertEqual(
|
||||
self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 1
|
||||
)
|
||||
|
||||
def test_unknown_negative_counts_at_execution_without_double_positive_count(self):
|
||||
graph = pair_graph(negative=["2", 0])
|
||||
graph["2"] = {"class_type": "PromptSearchList", "inputs": {}}
|
||||
self.dequeue(self.submit(graph))
|
||||
self.execute(graph, negative="batch negative")
|
||||
pair = self.db.list_prompt_pairs()["pairs"][0]
|
||||
self.assertEqual(pair["run_count"], 1)
|
||||
self.assertEqual(
|
||||
self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 1
|
||||
)
|
||||
|
||||
def test_cyclic_unknown_and_oversized_pair_graphs_are_bounded(self):
|
||||
from utils.prompt_graph import intrinsic_pair_nodes
|
||||
from utils.prompt_pair_resolver import resolve_text_input
|
||||
|
||||
self.assertEqual(intrinsic_pair_nodes(None), [])
|
||||
self.assertEqual(intrinsic_pair_nodes({str(n): {} for n in range(501)}), [])
|
||||
graph = pair_graph()
|
||||
graph["2"] = {"class_type": "Reroute", "inputs": {"value": ["2", 0]}}
|
||||
graph["5"]["inputs"]["positive"] = ["2", 0]
|
||||
self.assertEqual(intrinsic_pair_nodes(graph), [])
|
||||
self.assertIsNone(resolve_text_input(None, "1", "text"))
|
||||
self.assertIsNone(resolve_text_input({1: {}, "1": {}}, "1", "text"))
|
||||
self.assertIsNone(resolve_text_input(graph, "missing", "text"))
|
||||
self.assertIsNone(resolve_text_input(pair_graph(), "1", "text", max_depth=0))
|
||||
graph["1"]["inputs"]["text"] = ["1", 3]
|
||||
self.assertEqual(resolve_text_input(graph, "1", "text"), "blur")
|
||||
self.assertIsNone(resolve_text_input(graph, "1", "text", max_text_length=2))
|
||||
|
||||
def test_replacing_client_prompt_id_clears_obsolete_pair_records(self):
|
||||
pid = self.db.save_prompt_pair("cat", "blur")
|
||||
request = {"prompt_id": "same-id", "prompt": pair_graph()}
|
||||
handle_queued_prompt(request, self.queued, pair_queued=self.pair_queued)
|
||||
replacement = pair_graph()
|
||||
replacement["5"]["inputs"]["positive"] = ["1", 1]
|
||||
replacement["5"]["inputs"]["negative"] = ["1", 0]
|
||||
request = {"prompt_id": "same-id", "prompt": replacement}
|
||||
handle_queued_prompt(request, self.queued, pair_queued=self.pair_queued)
|
||||
self.dequeue(request)
|
||||
self.assertEqual(self.db.get_prompt_pair(pid)["run_count"], 0)
|
||||
@@ -0,0 +1,112 @@
|
||||
"""Autocomplete shares vocabulary without mixing pair and prompt identities."""
|
||||
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
|
||||
class TestPairVocabulary(unittest.TestCase):
|
||||
def setUp(self):
|
||||
directory = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(directory.cleanup)
|
||||
self.addCleanup(assert_closed, self, directory.name)
|
||||
self.db = PromptDatabase(os.path.join(directory.name, "vocabulary.db"))
|
||||
self.addCleanup(self.db.close_all)
|
||||
|
||||
def prompt(self, text: str, **kwargs) -> int:
|
||||
return self.db.save_prompt(
|
||||
text, prompt_hash=generate_prompt_hash(text), **kwargs
|
||||
)
|
||||
|
||||
def suggestions(self, prefix="", context=None, limit=20):
|
||||
return self.db.suggest_tags(prefix, context, limit)
|
||||
|
||||
def test_pair_only_categories_and_tags_are_available_without_component_metadata(
|
||||
self,
|
||||
):
|
||||
pair = self.db.get_prompt_pair(
|
||||
self.db.save_prompt_pair(
|
||||
"portrait", "blurry", category="Pairs", tags=["studio"]
|
||||
)
|
||||
)
|
||||
self.assertEqual(self.db.get_all_categories(), ["Pairs"])
|
||||
self.assertEqual(self.db.get_all_tags(), ["studio"])
|
||||
self.assertEqual(self.suggestions("stu"), [{"name": "studio", "count": 1}])
|
||||
for name in ("positive_prompt_id", "negative_prompt_id"):
|
||||
component = self.db.get_prompt_by_id(pair[name])
|
||||
self.assertIsNone(component["category"])
|
||||
self.assertEqual(component["tags"], [])
|
||||
|
||||
def test_categories_are_trimmed_deduplicated_and_include_archived_pairs(self):
|
||||
self.prompt("standalone", category=" Portrait ")
|
||||
pair = self.db.save_prompt_pair(
|
||||
"portrait", "", category="Portrait", tags=["archived"]
|
||||
)
|
||||
self.db.archive_prompt_pair(pair)
|
||||
self.db.save_prompt_pair("landscape", "", category=" Landscape ")
|
||||
self.db.save_prompt_pair("empty", "", category=" ")
|
||||
self.assertEqual(self.db.get_all_categories(), ["Landscape", "Portrait"])
|
||||
self.assertEqual(self.db.get_all_tags(), ["archived"])
|
||||
self.assertEqual(self.suggestions("arc"), [{"name": "archived", "count": 1}])
|
||||
|
||||
def test_counts_include_distinct_pair_and_prompt_references(self):
|
||||
self.prompt("portrait", tags=["shared", "legacy"])
|
||||
self.db.save_prompt_pair(
|
||||
"portrait", "first negative", tags=["shared", "paired"]
|
||||
)
|
||||
self.db.save_prompt_pair("portrait", "second negative", tags=["shared"])
|
||||
self.assertEqual(self.suggestions("sha"), [{"name": "shared", "count": 3}])
|
||||
self.assertEqual(
|
||||
self.suggestions("sha", ["paired"]), [{"name": "shared", "count": 1}]
|
||||
)
|
||||
self.assertEqual(
|
||||
self.suggestions("sha", ["legacy"]), [{"name": "shared", "count": 1}]
|
||||
)
|
||||
|
||||
def test_context_does_not_join_colliding_pair_and_prompt_ids(self):
|
||||
prompt = self.prompt("portrait", tags=["legacy-context", "legacy-only"])
|
||||
pair = self.db.save_prompt_pair(
|
||||
"portrait", "blurry", tags=["pair-context", "pair-only"]
|
||||
)
|
||||
self.assertEqual(prompt, pair)
|
||||
self.assertEqual(
|
||||
self.suggestions("", ["legacy-context"]),
|
||||
[{"name": "legacy-only", "count": 1}],
|
||||
)
|
||||
self.assertEqual(
|
||||
self.suggestions("", ["pair-context"]), [{"name": "pair-only", "count": 1}]
|
||||
)
|
||||
self.assertEqual(self.suggestions("", ["legacy-context", "pair-context"]), [])
|
||||
|
||||
def test_pair_context_matches_all_tags_case_insensitively(self):
|
||||
self.db.save_prompt_pair("portrait", "blurry", tags=["One", "two", "three"])
|
||||
self.db.save_prompt_pair("other", "blurry", tags=["One", "extra"])
|
||||
self.assertEqual(
|
||||
self.suggestions("TH", ["one", "TWO", "one"]),
|
||||
[{"name": "three", "count": 1}],
|
||||
)
|
||||
self.assertEqual(self.suggestions("", ["missing"]), [])
|
||||
|
||||
def test_pair_prefix_wildcards_are_literal_and_limit_is_preserved(self):
|
||||
self.db.save_prompt_pair(
|
||||
"portrait", "", tags=["tag_one", "tag_two", "tag%literal"]
|
||||
)
|
||||
self.assertEqual(
|
||||
self.suggestions("tag_", limit=1), [{"name": "tag_one", "count": 1}]
|
||||
)
|
||||
self.assertEqual(
|
||||
self.suggestions("tag%"), [{"name": "tag%literal", "count": 1}]
|
||||
)
|
||||
self.assertEqual(self.suggestions("%"), [])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,160 @@
|
||||
"""Paired node encoding preserves recipe storage and output roles."""
|
||||
|
||||
import logging
|
||||
import os
|
||||
import tempfile
|
||||
import unittest
|
||||
from unittest.mock import Mock
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from prompt_manager_pair import PromptManagerPair
|
||||
from utils.usage_tracking import PendingFirstUse
|
||||
|
||||
|
||||
def pair_graph(text="cat", negative="blur", positive_slot=0, negative_slot=1):
|
||||
return {
|
||||
"1": {
|
||||
"class_type": "PromptManagerPair",
|
||||
"inputs": {"text": text, "negative_text": negative},
|
||||
},
|
||||
"5": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {
|
||||
"model": ["9", 0],
|
||||
"positive": ["1", positive_slot],
|
||||
"negative": ["1", negative_slot],
|
||||
},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class TestPromptPairNode(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
self.addCleanup(assert_closed, self, self.tmp.name)
|
||||
self.db = PromptDatabase(os.path.join(self.tmp.name, "pairs.db"))
|
||||
self.addCleanup(self.db.close_all)
|
||||
self.node = PromptManagerPair.__new__(PromptManagerPair)
|
||||
self.node.db = self.db
|
||||
self.node.logger = logging.getLogger("test.pair_node")
|
||||
self.node.prompt_tracker = Mock()
|
||||
self.node.comfyui_integration = Mock()
|
||||
self.node._inject_lora_trigger_words = Mock(
|
||||
side_effect=lambda text: text + " LORA"
|
||||
)
|
||||
self.node.pending_first_use = PendingFirstUse()
|
||||
self.node.pending_pair_first_use = PendingFirstUse()
|
||||
self.clip = Mock()
|
||||
self.clip.encode_from_tokens_scheduled.side_effect = lambda tokens: (
|
||||
"conditioning",
|
||||
tokens,
|
||||
)
|
||||
self.clip.tokenize.side_effect = lambda text: text
|
||||
|
||||
def test_schema_has_positive_composition_and_four_ordered_outputs(self):
|
||||
schema = PromptManagerPair.INPUT_TYPES()
|
||||
self.assertEqual(list(schema["required"]), ["text", "negative_text", "clip"])
|
||||
self.assertTrue(schema["required"]["negative_text"][1]["multiline"])
|
||||
self.assertEqual(schema["required"]["negative_text"][1]["default"], "")
|
||||
for name in ("prepend_text", "append_text"):
|
||||
self.assertTrue(schema["optional"][name][1]["forceInput"])
|
||||
self.assertEqual(
|
||||
PromptManagerPair.RETURN_TYPES,
|
||||
("CONDITIONING", "CONDITIONING", "STRING", "STRING"),
|
||||
)
|
||||
self.assertNotEqual(
|
||||
PromptManagerPair.IS_CHANGED(None, "cat", negative_text="blur"),
|
||||
PromptManagerPair.IS_CHANGED(None, "cat", negative_text="watermark"),
|
||||
)
|
||||
|
||||
def test_encodes_each_side_once_positive_composition_and_lora_only(self):
|
||||
graph = pair_graph()
|
||||
result = self.node.encode_prompt(
|
||||
self.clip,
|
||||
"cat",
|
||||
negative_text=" blur ",
|
||||
prepend_text=" trigger ",
|
||||
append_text=" detail ",
|
||||
category="Animals",
|
||||
tags="studio",
|
||||
prompt=graph,
|
||||
unique_id="1",
|
||||
)
|
||||
self.assertEqual(
|
||||
result,
|
||||
(
|
||||
("conditioning", "trigger cat detail LORA"),
|
||||
("conditioning", " blur "),
|
||||
"trigger cat detail LORA",
|
||||
" blur ",
|
||||
),
|
||||
)
|
||||
self.node._inject_lora_trigger_words.assert_called_once_with(
|
||||
"trigger cat detail"
|
||||
)
|
||||
self.assertEqual(self.clip.tokenize.call_count, 2)
|
||||
pair = self.db.list_prompt_pairs()["pairs"][0]
|
||||
self.assertEqual(
|
||||
(
|
||||
pair["positive_text"],
|
||||
pair["negative_text"],
|
||||
pair["category"],
|
||||
pair["tags"],
|
||||
),
|
||||
("cat", "blur", "Animals", ["studio"]),
|
||||
)
|
||||
self.assertIsNone(
|
||||
self.db.get_prompt_by_id(pair["positive_prompt_id"])["category"]
|
||||
)
|
||||
tracked = self.node.prompt_tracker.set_current_prompt.call_args.kwargs
|
||||
self.assertEqual(tracked["additional_data"]["pair_id"], pair["id"])
|
||||
self.assertEqual(tracked["additional_data"]["negative_text"], " blur ")
|
||||
self.assertEqual(tracked["additional_data"]["positive_storage_text"], "cat")
|
||||
|
||||
def test_empty_negative_encodes_empty_text_and_stores_null_component(self):
|
||||
result = self.node.encode_prompt(self.clip, "cat", negative_text="")
|
||||
self.assertEqual(result[3], "")
|
||||
self.assertIsNone(self.db.list_prompt_pairs()["pairs"][0]["negative_prompt_id"])
|
||||
|
||||
def test_invalid_clip_cannot_save_or_count(self):
|
||||
with self.assertRaises(RuntimeError):
|
||||
self.node.encode_prompt(None, "cat", negative_text="blur")
|
||||
self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
|
||||
self.node.prompt_tracker.set_current_prompt.assert_not_called()
|
||||
|
||||
def test_persistence_failure_does_not_block_conditioning(self):
|
||||
self.node.db = Mock()
|
||||
self.node.db.save_prompt_pair.side_effect = RuntimeError("unavailable")
|
||||
result = self.node.encode_prompt(self.clip, "cat", negative_text="blur")
|
||||
self.assertEqual(result[2:], ("cat LORA", "blur"))
|
||||
self.node.prompt_tracker.set_current_prompt.assert_not_called()
|
||||
|
||||
def test_negative_output_used_as_positive_does_not_claim_positive_usage(self):
|
||||
self.node.encode_prompt(
|
||||
self.clip,
|
||||
"cat",
|
||||
negative_text="blur",
|
||||
prompt=pair_graph(positive_slot=1, negative_slot=0),
|
||||
unique_id="1",
|
||||
)
|
||||
pair = self.db.list_prompt_pairs()["pairs"][0]
|
||||
self.assertEqual(pair["run_count"], 0)
|
||||
self.assertEqual(
|
||||
self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 0
|
||||
)
|
||||
self.node.prompt_tracker.set_current_prompt.assert_not_called()
|
||||
|
||||
def test_invalid_text_does_not_write_and_integration_failures_do_not_block(self):
|
||||
with self.assertRaises(ValueError):
|
||||
self.node.encode_prompt(self.clip, [], negative_text="blur")
|
||||
self.node.comfyui_integration.register_prompt.side_effect = RuntimeError(
|
||||
"metadata unavailable"
|
||||
)
|
||||
result = self.node.encode_prompt(self.clip, "cat", negative_text="blur")
|
||||
self.assertEqual(result[2:], ("cat LORA", "blur"))
|
||||
@@ -0,0 +1,427 @@
|
||||
"""Conservative prompt-pair resolution contracts, independent of ComfyUI."""
|
||||
|
||||
import unittest
|
||||
|
||||
from utils.prompt_pair_resolver import OutputAdapter, resolve_prompt_pair
|
||||
|
||||
|
||||
def graph(positive: str = "portrait", negative: str = "blurry") -> dict:
|
||||
return {
|
||||
"1": {"class_type": "CLIPTextEncode", "inputs": {"text": positive}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": negative}},
|
||||
"3": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"positive": ["1", 0], "negative": ["2", 0]},
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
class TestPromptPairResolver(unittest.TestCase):
|
||||
def test_resolves_same_sampler_and_preserves_output_provenance(self):
|
||||
result = resolve_prompt_pair({"prompt": graph()})
|
||||
self.assertEqual(result["state"], "resolved")
|
||||
self.assertEqual(result["positive"]["text"], "portrait")
|
||||
self.assertEqual(result["negative"]["text"], "blurry")
|
||||
self.assertEqual(result["positive"]["provenance"]["node_id"], "1")
|
||||
self.assertEqual(result["positive"]["provenance"]["output_slot"], 0)
|
||||
self.assertEqual(result["sampler_id"], "3")
|
||||
|
||||
def test_explicit_empty_negative_is_resolved_but_missing_is_absent(self):
|
||||
self.assertEqual(resolve_prompt_pair(graph(negative=""))["state"], "resolved")
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(graph(negative=""))["negative"]["text"], ""
|
||||
)
|
||||
source = graph()
|
||||
del source["3"]["inputs"]["negative"]
|
||||
result = resolve_prompt_pair(source)
|
||||
self.assertEqual(result["state"], "absent")
|
||||
self.assertEqual(result["positive"]["text"], "portrait")
|
||||
self.assertIsNone(result["negative"])
|
||||
|
||||
def test_multiple_samplers_are_ambiguous_even_when_texts_match(self):
|
||||
source = graph()
|
||||
source["4"] = {
|
||||
"class_type": "KSamplerAdvanced",
|
||||
"inputs": source["3"]["inputs"],
|
||||
}
|
||||
result = resolve_prompt_pair(source)
|
||||
self.assertEqual(result["state"], "ambiguous")
|
||||
self.assertEqual(len(result["candidates"]), 2)
|
||||
self.assertIsNone(result["positive"])
|
||||
|
||||
def test_saver_anchor_selects_only_its_sampler(self):
|
||||
source = graph()
|
||||
source.update(
|
||||
{
|
||||
"4": {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {"positive": ["2", 0], "negative": ["1", 0]},
|
||||
},
|
||||
"5": {"class_type": "VAEDecode", "inputs": {"samples": ["4", 0]}},
|
||||
"6": {"class_type": "SaveImage", "inputs": {"images": ["5", 0]}},
|
||||
}
|
||||
)
|
||||
result = resolve_prompt_pair(source, output_node_id="6")
|
||||
self.assertEqual(result["state"], "resolved")
|
||||
self.assertEqual(result["sampler_id"], "4")
|
||||
self.assertEqual(result["positive"]["text"], "blurry")
|
||||
|
||||
def test_unknown_anchor_transform_and_multistage_are_unresolved(self):
|
||||
source = graph()
|
||||
source["5"] = {
|
||||
"class_type": "MysteryImageTransform",
|
||||
"inputs": {"images": ["3", 0]},
|
||||
}
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source, output_node_id="5")["state"], "unresolved"
|
||||
)
|
||||
source["4"] = {
|
||||
"class_type": "KSampler",
|
||||
"inputs": {**source["3"]["inputs"], "latent_image": ["3", 0]},
|
||||
}
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source, output_node_id="4")["state"], "unresolved"
|
||||
)
|
||||
|
||||
def test_combines_unknown_encoders_and_sdxl_channels_are_not_guessed(self):
|
||||
for name in ("ConditioningCombine", "CLIPTextEncodeSDXL", "UnknownEncoder"):
|
||||
with self.subTest(name=name):
|
||||
source = graph()
|
||||
source["1"] = {
|
||||
"class_type": name,
|
||||
"inputs": {
|
||||
"text": "guess",
|
||||
"text_g": "global",
|
||||
"text_l": "local",
|
||||
"conditioning_1": ["2", 0],
|
||||
},
|
||||
}
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_prompt_manager_composition_and_string_slot_are_resolved(self):
|
||||
source = graph()
|
||||
source["8"] = {
|
||||
"class_type": "PromptManagerSplit",
|
||||
"inputs": {
|
||||
"text": "portrait",
|
||||
"prepend_text": " trigger ",
|
||||
"append_text": " detail ",
|
||||
},
|
||||
}
|
||||
source["1"]["inputs"]["text"] = ["8", 1]
|
||||
result = resolve_prompt_pair(source)
|
||||
self.assertEqual(result["positive"]["text"], "trigger portrait detail")
|
||||
self.assertIn(
|
||||
{"node_id": "8", "output_slot": 1}, result["positive"]["provenance"]["path"]
|
||||
)
|
||||
source["3"]["inputs"]["positive"] = ["8", 0]
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source)["positive"]["text"], "trigger portrait detail"
|
||||
)
|
||||
source["1"]["inputs"]["text"] = ["8", 0]
|
||||
source["3"]["inputs"]["positive"] = ["1", 0]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_concat_and_reroute_resolve_safe_linked_strings(self):
|
||||
source = graph()
|
||||
source["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "portrait"}}
|
||||
source["9"] = {
|
||||
"class_type": "StringConcatenate",
|
||||
"inputs": {"string_a": ["8", 0], "string_b": ["8", 0], "delimiter": ", "},
|
||||
}
|
||||
source["10"] = {"class_type": "Reroute", "inputs": {"value": ["9", 0]}}
|
||||
source["1"]["inputs"]["text"] = ["10", 0]
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source)["positive"]["text"], "portrait, portrait"
|
||||
)
|
||||
|
||||
def test_slot_adapter_does_not_conflate_two_outputs_of_one_node(self):
|
||||
source = graph()
|
||||
source["8"] = {
|
||||
"class_type": "FuturePair",
|
||||
"inputs": {"first": "positive", "second": "negative"},
|
||||
}
|
||||
source["3"]["inputs"] = {"positive": ["8", 0], "negative": ["8", 1]}
|
||||
adapters = {
|
||||
("FuturePair", 0): OutputAdapter("first", "conditioning"),
|
||||
("FuturePair", 1): OutputAdapter("second", "conditioning"),
|
||||
}
|
||||
result = resolve_prompt_pair(source, adapters=adapters)
|
||||
self.assertEqual(result["positive"]["text"], "positive")
|
||||
self.assertEqual(result["negative"]["text"], "negative")
|
||||
self.assertNotEqual(
|
||||
result["positive"]["provenance"]["output_slot"],
|
||||
result["negative"]["provenance"]["output_slot"],
|
||||
)
|
||||
|
||||
def test_prompt_manager_text_has_only_its_real_string_output_zero(self):
|
||||
source = graph()
|
||||
source["8"] = {
|
||||
"class_type": "PromptManagerText",
|
||||
"inputs": {"text": "text node", "prepend_text": "before"},
|
||||
}
|
||||
source["1"]["inputs"]["text"] = ["8", 0]
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source)["positive"]["text"], "before text node"
|
||||
)
|
||||
source["1"]["inputs"]["text"] = ["8", 1]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
source["3"]["inputs"]["positive"] = ["8", 0]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_malformed_class_names_and_unknown_samplers_do_not_claim_absence(self):
|
||||
for name in ([], {}, None, "UnknownSampler"):
|
||||
source = graph()
|
||||
source["3"]["class_type"] = name
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_pair_adapter_composes_only_positive_and_retains_recipe_text(self):
|
||||
source = graph()
|
||||
source["8"] = {
|
||||
"class_type": "PromptManagerPair",
|
||||
"inputs": {
|
||||
"text": "positive",
|
||||
"negative_text": "negative",
|
||||
"prepend_text": "trigger",
|
||||
"append_text": "detail",
|
||||
},
|
||||
}
|
||||
source["3"]["inputs"] = {"positive": ["8", 0], "negative": ["8", 1]}
|
||||
result = resolve_prompt_pair(source)
|
||||
self.assertEqual(result["positive"]["text"], "trigger positive detail")
|
||||
self.assertEqual(result["positive"]["storage_text"], "positive")
|
||||
self.assertEqual(result["negative"]["text"], "negative")
|
||||
self.assertEqual(result["negative"]["storage_text"], "negative")
|
||||
source["1"]["inputs"]["text"] = ["8", 2]
|
||||
source["2"]["inputs"]["text"] = ["8", 3]
|
||||
source["3"]["inputs"] = {"positive": ["1", 0], "negative": ["2", 0]}
|
||||
result = resolve_prompt_pair(source)
|
||||
self.assertEqual(result["positive"]["storage_text"], "positive")
|
||||
self.assertEqual(result["negative"]["text"], "negative")
|
||||
source["3"]["inputs"]["negative"] = ["8", 3]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_linked_recipe_and_ordinary_clip_storage_values_are_explicit(self):
|
||||
source = graph()
|
||||
result = resolve_prompt_pair(source)
|
||||
self.assertEqual(result["positive"]["storage_text"], "portrait")
|
||||
source["8"] = {
|
||||
"class_type": "PrimitiveString",
|
||||
"inputs": {"value": "linked main"},
|
||||
}
|
||||
source["1"] = {
|
||||
"class_type": "PromptManager",
|
||||
"inputs": {"text": ["8", 0], "prepend_text": "before"},
|
||||
}
|
||||
result = resolve_prompt_pair(source)
|
||||
self.assertEqual(result["positive"]["text"], "before linked main")
|
||||
self.assertEqual(result["positive"]["storage_text"], "linked main")
|
||||
a1111 = resolve_prompt_pair(
|
||||
{
|
||||
"parameters": (
|
||||
"portrait\nNegative prompt: blurry\n" "Steps: 20, Sampler: Euler"
|
||||
)
|
||||
}
|
||||
)
|
||||
self.assertEqual(a1111["positive"]["storage_text"], "portrait")
|
||||
self.assertEqual(a1111["negative"]["storage_text"], "blurry")
|
||||
|
||||
def test_cycles_depth_visits_and_text_growth_are_bounded(self):
|
||||
source = graph()
|
||||
source["8"] = {"class_type": "Reroute", "inputs": {"value": ["8", 0]}}
|
||||
source["1"]["inputs"]["text"] = ["8", 0]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(graph(), max_visits=1)["state"], "unresolved"
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(graph(), max_depth=0)["state"], "unresolved"
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(graph(), max_text_length=3)["state"], "unresolved"
|
||||
)
|
||||
|
||||
def test_malformed_metadata_and_links_never_raise(self):
|
||||
for value in (
|
||||
None,
|
||||
[],
|
||||
"raw text",
|
||||
{},
|
||||
{"prompt": []},
|
||||
{"prompt": {"1": None}},
|
||||
):
|
||||
with self.subTest(value=value):
|
||||
self.assertIn(
|
||||
resolve_prompt_pair(value)["state"], ("absent", "unresolved")
|
||||
)
|
||||
for link in (["missing", 0], ["1", -1], ["1", True], ["1", 17], ["1"], 5):
|
||||
source = graph()
|
||||
source["3"]["inputs"]["positive"] = link
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_a1111_explicit_parameters_and_empty_negative(self):
|
||||
for negative in ("bad anatomy\nblurry", ""):
|
||||
result = resolve_prompt_pair(
|
||||
{
|
||||
"parameters": (
|
||||
f"portrait\nNegative prompt: {negative}\n"
|
||||
"Steps: 20, Sampler: Euler, Seed: 1"
|
||||
)
|
||||
}
|
||||
)
|
||||
self.assertEqual(result["state"], "resolved")
|
||||
self.assertEqual(result["negative"]["text"], negative)
|
||||
self.assertEqual(result["positive"]["provenance"]["format"], "a1111")
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair({"parameters": "portrait\nSteps: 20, Sampler: Euler"})[
|
||||
"state"
|
||||
],
|
||||
"absent",
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair({"parameters": "unstructured arbitrary text"})["state"],
|
||||
"unresolved",
|
||||
)
|
||||
|
||||
def test_was_concat_preserves_configured_whitespace_and_rejects_bad_delimiters(
|
||||
self,
|
||||
):
|
||||
source = graph()
|
||||
source["8"] = {
|
||||
"class_type": "Text Concatenate",
|
||||
"inputs": {
|
||||
"text_a": " portrait ",
|
||||
"text_b": "",
|
||||
"text_c": "studio",
|
||||
"delimiter": ", ",
|
||||
},
|
||||
}
|
||||
source["1"]["inputs"]["text"] = ["8", 0]
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source)["positive"]["text"], "portrait, studio"
|
||||
)
|
||||
source["8"]["inputs"]["clean_whitespace"] = "false"
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source)["positive"]["text"], " portrait , studio"
|
||||
)
|
||||
source["8"]["inputs"]["delimiter"] = []
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
source["8"]["inputs"] = {1: "bad key"}
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_valid_empty_latent_and_reroute_but_not_unknown_or_cyclic_latents(self):
|
||||
source = graph()
|
||||
source["8"] = {"class_type": "EmptyLatentImage", "inputs": {}}
|
||||
source["9"] = {"class_type": "Reroute", "inputs": {"value": ["8", 0]}}
|
||||
source["3"]["inputs"]["latent_image"] = ["9", 0]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "resolved")
|
||||
source["9"]["inputs"]["value"] = ["9", 0]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
source["3"]["inputs"]["latent_image"] = ["8", 0]
|
||||
source["8"]["class_type"] = "UnknownLatentTransform"
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
source["3"]["inputs"]["latent_image"] = None
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_cyclic_or_invalid_output_anchors_never_guess(self):
|
||||
source = graph()
|
||||
source["8"] = {"class_type": "Reroute", "inputs": {"value": ["8", 0]}}
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source, output_node_id="8")["state"], "unresolved"
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source, output_node_id="missing")["state"], "unresolved"
|
||||
)
|
||||
source["8"] = {"class_type": "SaveImage", "inputs": {"images": ["3", 1]}}
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source, output_node_id="8")["state"], "unresolved"
|
||||
)
|
||||
source["8"] = {"class_type": "SaveImage", "inputs": {"images": ["3", 0]}}
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(source, output_node_id="8", max_depth=1)["state"],
|
||||
"unresolved",
|
||||
)
|
||||
|
||||
def test_malformed_reroute_empty_positive_and_duplicate_string_ids(self):
|
||||
source = graph(positive="")
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
source = graph()
|
||||
source["8"] = {
|
||||
"class_type": "Reroute",
|
||||
"inputs": {"value": ["1", 0], "input": ["2", 0]},
|
||||
}
|
||||
source["3"]["inputs"]["positive"] = ["8", 0]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
source[3] = source["3"]
|
||||
self.assertEqual(resolve_prompt_pair(source)["state"], "unresolved")
|
||||
|
||||
def test_a1111_malformed_ambiguous_and_bounded_parameters(self):
|
||||
self.assertEqual(resolve_prompt_pair({"parameters": None})["state"], "absent")
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair({"parameters": "x" * 20}, max_text_length=10)["state"],
|
||||
"unresolved",
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(
|
||||
{"parameters": "\nNegative prompt: bad\nSteps: 20, Sampler: Euler"}
|
||||
)["state"],
|
||||
"unresolved",
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(
|
||||
{
|
||||
"parameters": (
|
||||
"portrait\nNegative prompt: bad\nNegative prompt: worse\n"
|
||||
"Steps: 20, Sampler: Euler"
|
||||
)
|
||||
}
|
||||
)["state"],
|
||||
"ambiguous",
|
||||
)
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(
|
||||
{
|
||||
"parameters": (
|
||||
"portrait\nNegative prompt: bad\n"
|
||||
"Steps: 20, Sampler: Euler\nextra text"
|
||||
)
|
||||
}
|
||||
)["state"],
|
||||
"unresolved",
|
||||
)
|
||||
|
||||
def test_exponential_concat_text_is_bounded_before_building_a_large_string(self):
|
||||
source = graph()
|
||||
source["8"] = {"class_type": "PrimitiveString", "inputs": {"value": "xx"}}
|
||||
for index in range(9, 19):
|
||||
source[str(index)] = {
|
||||
"class_type": "StringConcatenate",
|
||||
"inputs": {
|
||||
"string_a": [str(index - 1), 0],
|
||||
"string_b": [str(index - 1), 0],
|
||||
},
|
||||
}
|
||||
source["1"]["inputs"]["text"] = ["18", 0]
|
||||
result = resolve_prompt_pair(source, max_text_length=100)
|
||||
self.assertEqual(result["state"], "unresolved")
|
||||
self.assertIn("length limit", result["reason"])
|
||||
self.assertEqual(
|
||||
resolve_prompt_pair(graph(), max_depth=1)["state"], "unresolved"
|
||||
)
|
||||
|
||||
def test_api_graph_is_authoritative_over_unrelated_parameters(self):
|
||||
source = graph()
|
||||
source["1"]["class_type"] = "UnknownEncoder"
|
||||
result = resolve_prompt_pair(
|
||||
{
|
||||
"prompt": source,
|
||||
"parameters": (
|
||||
"fallback\nNegative prompt: wrong\n" "Steps: 20, Sampler: Euler"
|
||||
),
|
||||
}
|
||||
)
|
||||
self.assertEqual(result["state"], "unresolved")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,301 @@
|
||||
"""Reusable pairs own their metadata without changing shared prompt content."""
|
||||
|
||||
import os
|
||||
import sqlite3
|
||||
import tempfile
|
||||
import unittest
|
||||
from concurrent.futures import ThreadPoolExecutor
|
||||
from unittest.mock import patch
|
||||
|
||||
try:
|
||||
from tests.open_handles import assert_closed
|
||||
except ImportError: # discovered with tests/ as the top-level directory
|
||||
from open_handles import assert_closed
|
||||
|
||||
from database.operations import PromptDatabase
|
||||
from utils.hashing import generate_prompt_hash
|
||||
|
||||
|
||||
class TestPromptPairs(unittest.TestCase):
|
||||
def setUp(self):
|
||||
self.tmp = tempfile.TemporaryDirectory()
|
||||
self.addCleanup(self.tmp.cleanup)
|
||||
self.addCleanup(assert_closed, self, self.tmp.name)
|
||||
self.db = PromptDatabase(os.path.join(self.tmp.name, "pairs.db"))
|
||||
self.addCleanup(self.db.close_all)
|
||||
|
||||
def save(self, positive="a cat", negative="blurry", **kwargs):
|
||||
return self.db.save_prompt_pair(positive, negative, **kwargs)
|
||||
|
||||
def test_pair_dedup_is_ordered_and_empty_negative_is_explicit(self):
|
||||
first = self.save()
|
||||
self.assertEqual(first, self.save(" A CAT ", " BLURRY "))
|
||||
self.assertNotEqual(first, self.save("blurry", "a cat"))
|
||||
self.assertNotEqual(first, self.save("a cat", "watermark"))
|
||||
empty = self.save("a cat", "")
|
||||
self.assertEqual(empty, self.save("a cat", " "))
|
||||
self.assertIsNone(self.db.get_prompt_pair(empty)["negative_prompt_id"])
|
||||
self.assertEqual(self.db.get_prompt_pair(empty)["negative_text"], "")
|
||||
self.assertNotEqual(empty, first)
|
||||
|
||||
def test_same_text_can_serve_both_roles_and_existing_metadata_survives(self):
|
||||
prompt = self.db.save_prompt(
|
||||
"a cat",
|
||||
category="original",
|
||||
tags=["old"],
|
||||
prompt_hash=generate_prompt_hash("a cat"),
|
||||
)
|
||||
pair = self.db.get_prompt_pair(
|
||||
self.save("a cat", "a cat", category="pair", tags=["new"])
|
||||
)
|
||||
self.assertEqual(pair["positive_prompt_id"], prompt)
|
||||
self.assertEqual(pair["negative_prompt_id"], prompt)
|
||||
self.assertEqual(pair["tags"], ["new"])
|
||||
self.assertEqual(self.db.get_prompt_by_id(prompt)["category"], "original")
|
||||
self.assertEqual(self.db.get_prompt_by_id(prompt)["tags"], ["old"])
|
||||
|
||||
def test_repeated_save_preserves_pair_metadata_and_explicit_update_changes_it(self):
|
||||
pid = self.save(
|
||||
title="Title",
|
||||
category="Portrait",
|
||||
tags=["studio", "studio"],
|
||||
rating=4,
|
||||
notes="keep",
|
||||
)
|
||||
self.assertEqual(pid, self.save(title="overwrite", tags=["bad"], rating=1))
|
||||
pair = self.db.get_prompt_pair(pid)
|
||||
self.assertEqual(
|
||||
(
|
||||
pair["title"],
|
||||
pair["category"],
|
||||
pair["tags"],
|
||||
pair["rating"],
|
||||
pair["notes"],
|
||||
),
|
||||
("Title", "Portrait", ["studio"], 4, "keep"),
|
||||
)
|
||||
self.assertTrue(
|
||||
self.db.update_prompt_pair_metadata(
|
||||
pid, title="", category="", tags=[], notes=""
|
||||
)
|
||||
)
|
||||
pair = self.db.get_prompt_pair(pid)
|
||||
self.assertEqual(
|
||||
(pair["title"], pair["category"], pair["tags"], pair["notes"]),
|
||||
("", "", [], ""),
|
||||
)
|
||||
self.assertEqual(pair["rating"], 4)
|
||||
self.assertFalse(self.db.update_prompt_pair_metadata(999999, title="missing"))
|
||||
self.assertFalse(self.db.update_prompt_pair_metadata(pid))
|
||||
|
||||
def test_list_filters_pagination_archive_and_pair_usage_are_independent(self):
|
||||
first = self.save(category="Animals", tags=["studio", "cat"])
|
||||
second = self.save("a dog", "blur", category="Animals", tags=["studio"])
|
||||
other = self.save("sky", "", category="Outdoor")
|
||||
result = self.db.list_prompt_pairs(category="Animals", tags=["studio"], limit=1)
|
||||
self.assertEqual(result["total"], 2)
|
||||
self.assertTrue(result["has_more"])
|
||||
self.assertEqual(len(result["pairs"]), 1)
|
||||
self.assertEqual(
|
||||
self.db.list_prompt_pairs(text="dog")["pairs"][0]["id"], second
|
||||
)
|
||||
self.assertEqual(self.db.list_prompt_pairs(tags=["studio", "cat"])["total"], 1)
|
||||
self.assertTrue(self.db.archive_prompt_pair(first))
|
||||
self.assertEqual(self.db.list_prompt_pairs()["total"], 2)
|
||||
self.assertEqual(self.db.list_prompt_pairs(include_archived=True)["total"], 3)
|
||||
self.assertTrue(self.db.get_prompt_pair(first)["archived"])
|
||||
self.assertTrue(self.db.archive_prompt_pair(first, False))
|
||||
self.assertTrue(self.db.record_prompt_pair_use(other, 2))
|
||||
pair = self.db.get_prompt_pair(other)
|
||||
self.assertEqual(pair["run_count"], 2)
|
||||
self.assertIsNotNone(pair["last_used_at"])
|
||||
self.assertEqual(
|
||||
self.db.get_prompt_by_id(pair["positive_prompt_id"])["run_count"], 0
|
||||
)
|
||||
|
||||
def test_failure_during_tag_save_rolls_back_pair_and_components(self):
|
||||
with patch.object(
|
||||
self.db, "_ensure_tags", side_effect=sqlite3.OperationalError("injected")
|
||||
):
|
||||
with self.assertRaises(sqlite3.OperationalError):
|
||||
self.save(tags=["studio"])
|
||||
self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
|
||||
self.assertEqual(self.db.search_prompts(), [])
|
||||
|
||||
def test_concurrent_same_pair_returns_one_id_and_preserves_first_metadata(self):
|
||||
first = self.save(title="original")
|
||||
|
||||
def save(index):
|
||||
try:
|
||||
return self.save(title=f"writer {index}")
|
||||
finally:
|
||||
self.db.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=6) as executor:
|
||||
ids = list(executor.map(save, range(24)))
|
||||
self.assertEqual(set(ids), {first})
|
||||
self.assertEqual(self.db.get_prompt_pair(first)["title"], "original")
|
||||
self.assertEqual(self.db.list_prompt_pairs()["total"], 1)
|
||||
|
||||
def test_concurrent_empty_negative_pair_creation_is_atomic(self):
|
||||
def save(_):
|
||||
try:
|
||||
return self.save(negative="")
|
||||
finally:
|
||||
self.db.close()
|
||||
|
||||
with ThreadPoolExecutor(max_workers=6) as executor:
|
||||
ids = list(executor.map(save, range(12)))
|
||||
self.assertEqual(len(set(ids)), 1)
|
||||
self.assertEqual(len(self.db.search_prompts()), 1)
|
||||
|
||||
def test_component_edit_and_all_delete_paths_reject_dependencies_atomically(self):
|
||||
from database.operations import PairDependencyError
|
||||
|
||||
pair = self.db.get_prompt_pair(self.save())
|
||||
positive = pair["positive_prompt_id"]
|
||||
image = self.db.link_image_to_prompt(
|
||||
positive, os.path.join(self.tmp.name, "keep.png")
|
||||
)
|
||||
unused = self.db.save_prompt(
|
||||
"unused", category="remove", prompt_hash=generate_prompt_hash("unused")
|
||||
)
|
||||
self.db.update_prompt_metadata(positive, category="remove")
|
||||
for action in (
|
||||
lambda: self.db.update_prompt_text(positive, "replacement"),
|
||||
lambda: self.db.delete_prompt(positive),
|
||||
lambda: self.db.bulk_delete_prompts([unused, positive]),
|
||||
lambda: self.db.delete_prompts_by_category("remove"),
|
||||
):
|
||||
with self.subTest(action=action):
|
||||
with self.assertRaises(PairDependencyError):
|
||||
action()
|
||||
self.assertIsNotNone(self.db.get_prompt_by_id(unused))
|
||||
self.assertEqual(self.db.get_prompt_by_id(positive)["text"], "a cat")
|
||||
self.assertIsNotNone(self.db.get_image_by_id(image))
|
||||
self.assertTrue(self.db.delete_prompt_pair(pair["id"]))
|
||||
self.assertTrue(self.db.update_prompt_text(positive, "replacement"))
|
||||
|
||||
def test_prune_skips_members_even_in_archived_pairs_and_preserves_other_behavior(
|
||||
self,
|
||||
):
|
||||
pair = self.db.get_prompt_pair(self.save())
|
||||
self.db.archive_prompt_pair(pair["id"])
|
||||
unused = self.db.save_prompt(
|
||||
"unused", prompt_hash=generate_prompt_hash("unused")
|
||||
)
|
||||
self.assertEqual(self.db.prune_orphaned_prompts(), 1)
|
||||
self.assertIsNone(self.db.get_prompt_by_id(unused))
|
||||
self.assertIsNotNone(self.db.get_prompt_by_id(pair["negative_prompt_id"]))
|
||||
|
||||
def test_duplicate_cleanup_skips_referenced_group_without_losing_pairs(self):
|
||||
pair = self.db.get_prompt_pair(self.save())
|
||||
with self.db.model.get_connection() as conn:
|
||||
conn.execute(
|
||||
"INSERT INTO prompts(text,category) VALUES(' A CAT ', 'duplicate')"
|
||||
)
|
||||
self.assertEqual(self.db.cleanup_duplicates(), 0)
|
||||
self.assertEqual(self.db.get_prompt_pair(pair["id"])["positive_text"], "a cat")
|
||||
self.assertEqual(len(self.db.search_prompts()), 3)
|
||||
|
||||
def test_image_association_preserves_provenance_and_pair_deletion_requires_archive(
|
||||
self,
|
||||
):
|
||||
from database.operations import PairDependencyError
|
||||
|
||||
pid = self.save()
|
||||
pair = self.db.get_prompt_pair(pid)
|
||||
image = self.db.link_image_to_prompt(
|
||||
pair["positive_prompt_id"], os.path.join(self.tmp.name, "render.png")
|
||||
)
|
||||
provenance = {
|
||||
"run_id": "abc",
|
||||
"negative_text": "actual expanded negative",
|
||||
"positive_text": "trigger a cat",
|
||||
}
|
||||
association = self.db.link_image_to_pair(pid, image, provenance)
|
||||
self.assertEqual(
|
||||
association,
|
||||
self.db.link_image_to_pair(
|
||||
pid, image, dict(reversed(list(provenance.items())))
|
||||
),
|
||||
)
|
||||
self.assertEqual(
|
||||
self.db.get_image_prompt_pairs(image)[0]["provenance"], provenance
|
||||
)
|
||||
self.assertEqual(self.db.get_prompt_pair(pid)["image_count"], 1)
|
||||
with self.assertRaises(PairDependencyError):
|
||||
self.db.delete_prompt_pair(pid)
|
||||
self.assertTrue(self.db.archive_prompt_pair(pid))
|
||||
with self.assertRaises(ValueError):
|
||||
self.db.link_image_to_pair(99999, image, provenance)
|
||||
self.assertTrue(self.db.delete_image(image))
|
||||
self.assertEqual(self.db.get_image_prompt_pairs(image), [])
|
||||
self.assertTrue(self.db.delete_prompt_pair(pid))
|
||||
self.assertIsNotNone(self.db.get_prompt_by_id(pair["positive_prompt_id"]))
|
||||
|
||||
def test_pair_validates_values_before_writing(self):
|
||||
for positive, negative, kwargs in (
|
||||
("", "blur", {}),
|
||||
("cat", None, {}),
|
||||
("cat", "blur", {"rating": 6}),
|
||||
("cat", "blur", {"tags": "one"}),
|
||||
):
|
||||
with self.subTest(positive=positive, negative=negative, kwargs=kwargs):
|
||||
with self.assertRaises((TypeError, ValueError)):
|
||||
self.save(positive, negative, **kwargs)
|
||||
self.assertEqual(self.db.search_prompts(), [])
|
||||
for kwargs in ({"limit": 0}, {"offset": -1}, {"limit": True}):
|
||||
with self.assertRaises(ValueError):
|
||||
self.db.list_prompt_pairs(**kwargs)
|
||||
self.assertIsNone(self.db.get_prompt_pair(99999))
|
||||
self.assertFalse(self.db.delete_prompt_pair(99999))
|
||||
|
||||
def test_connection_helpers_leave_commit_and_rollback_to_the_caller(self):
|
||||
conn = self.db.model.get_connection()
|
||||
|
||||
def rollback_transaction():
|
||||
with conn:
|
||||
conn.execute("BEGIN IMMEDIATE")
|
||||
self.db._save_pair(conn, "transient", "temporary", tags=["tag"])
|
||||
self.assertRaisesRegex(RuntimeError, "own transaction", self.save)
|
||||
raise RuntimeError("rollback")
|
||||
|
||||
self.assertRaisesRegex(RuntimeError, "rollback", rollback_transaction)
|
||||
self.assertEqual(self.db.list_prompt_pairs()["total"], 0)
|
||||
self.assertEqual(self.db.search_prompts(), [])
|
||||
|
||||
def test_metadata_update_failure_rolls_back_prior_fields_and_tags(self):
|
||||
pid = self.save(title="keep", tags=["old"])
|
||||
with patch.object(
|
||||
self.db, "_ensure_tags", side_effect=sqlite3.OperationalError("injected")
|
||||
):
|
||||
with self.assertRaises(sqlite3.OperationalError):
|
||||
self.db.update_prompt_pair_metadata(pid, title="replace", tags=["new"])
|
||||
pair = self.db.get_prompt_pair(pid)
|
||||
self.assertEqual(pair["title"], "keep")
|
||||
self.assertEqual(pair["tags"], ["old"])
|
||||
|
||||
def test_invalid_associations_metadata_and_missing_mutations_are_safe(self):
|
||||
pid = self.save()
|
||||
other = self.db.save_prompt("other", prompt_hash=generate_prompt_hash("other"))
|
||||
image = self.db.link_image_to_prompt(
|
||||
other, os.path.join(self.tmp.name, "other.png")
|
||||
)
|
||||
for provenance in ([], {"number": float("nan")}, {"object": object()}):
|
||||
with self.assertRaises(ValueError):
|
||||
self.db.link_image_to_pair(pid, image, provenance)
|
||||
with self.assertRaisesRegex(ValueError, "positive prompt"):
|
||||
self.db.link_image_to_pair(pid, image, {})
|
||||
self.assertEqual(self.db.get_image_prompt_pairs(image), [])
|
||||
self.assertFalse(self.db.archive_prompt_pair(999999))
|
||||
self.assertFalse(self.db.record_prompt_pair_use(999999))
|
||||
for times in (0, True):
|
||||
with self.assertRaises(ValueError):
|
||||
self.db.record_prompt_pair_use(pid, times)
|
||||
with self.assertRaises(ValueError):
|
||||
self.db.archive_prompt_pair(pid, "yes")
|
||||
with self.assertRaises(ValueError):
|
||||
self.db.update_prompt_pair_metadata(pid, title=[])
|
||||
with self.assertRaises(ValueError):
|
||||
self.db.list_prompt_pairs(text=[])
|
||||
@@ -0,0 +1,46 @@
|
||||
"""Serve only the timeline concept and its pinned toolkit assets."""
|
||||
|
||||
import argparse
|
||||
from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
|
||||
from pathlib import Path
|
||||
from urllib.parse import urlsplit
|
||||
|
||||
ROOT = Path(__file__).resolve().parents[1]
|
||||
FILES = {
|
||||
"/": ("docs/concepts/prompt-variant-timeline.html", "text/html; charset=utf-8"),
|
||||
"/toolkit/core.js": ("web/js/shared/toolkit/core.js", "text/javascript"),
|
||||
"/toolkit/controls.css": ("web/js/shared/toolkit/controls.css", "text/css"),
|
||||
}
|
||||
|
||||
|
||||
class PreviewHandler(BaseHTTPRequestHandler):
|
||||
def do_GET(self) -> None:
|
||||
target = FILES.get(urlsplit(self.path).path)
|
||||
if target is None:
|
||||
self.send_error(404)
|
||||
return
|
||||
filename, mime = target
|
||||
body = (ROOT / filename).read_bytes()
|
||||
self.send_response(200)
|
||||
self.send_header("Content-Type", mime)
|
||||
self.send_header("Content-Length", str(len(body)))
|
||||
self.send_header("Cache-Control", "no-store")
|
||||
self.send_header("X-Content-Type-Options", "nosniff")
|
||||
self.end_headers()
|
||||
self.wfile.write(body)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--bind", default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8766)
|
||||
args = parser.parse_args()
|
||||
server = ThreadingHTTPServer((args.bind, args.port), PreviewHandler)
|
||||
print(f"Timeline prototype: http://{args.bind}:{args.port}/", flush=True)
|
||||
try:
|
||||
server.serve_forever()
|
||||
except KeyboardInterrupt:
|
||||
# Ctrl+C exits normally; the finally block closes the listening socket.
|
||||
pass
|
||||
finally:
|
||||
server.server_close()
|
||||
@@ -35,6 +35,12 @@ from .hashing import generate_prompt_hash
|
||||
from .metadata_extractor import ComfyUIMetadataExtractor
|
||||
from .logging_config import get_logger
|
||||
from .prompt_graph import resolve_text, run_prompt_nodes
|
||||
from .prompt_pair_resolver import resolve_prompt_pair
|
||||
|
||||
try:
|
||||
from ..database.pair_ingestion import ingest_resolved_pair
|
||||
except ImportError:
|
||||
from database.pair_ingestion import ingest_resolved_pair
|
||||
|
||||
# Upper bound on images waiting to be processed (a burst beyond this is skipped)
|
||||
MAX_PENDING_IMAGES = 1000
|
||||
@@ -380,6 +386,29 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
return "unknown", None
|
||||
|
||||
prompt_text = None
|
||||
resolution = resolve_prompt_pair(metadata)
|
||||
graph = metadata.get("prompt")
|
||||
if resolution["state"] == "resolved" and isinstance(graph, dict):
|
||||
positive = resolution["positive"]
|
||||
path = positive["provenance"].get("path", [])
|
||||
paired_source = any(
|
||||
isinstance(graph.get(item.get("node_id")), dict)
|
||||
and graph[item["node_id"]].get("class_type") == "PromptManagerPair"
|
||||
for item in path
|
||||
)
|
||||
if paired_source:
|
||||
try:
|
||||
existing = self.db_manager.get_prompt_by_hash(
|
||||
generate_prompt_hash(positive["storage_text"])
|
||||
)
|
||||
if existing:
|
||||
return "linked", {
|
||||
"id": existing["id"],
|
||||
"text": existing["text"],
|
||||
"from_metadata": True,
|
||||
}
|
||||
except Exception as error:
|
||||
self.logger.warning(f"Paired metadata lookup failed: {error}")
|
||||
|
||||
# The executed graph: use the PromptManager node feeding the positive input
|
||||
prompt_data = metadata.get("prompt")
|
||||
@@ -529,6 +558,17 @@ class ImageGenerationHandler(FileSystemEventHandler):
|
||||
etc.)
|
||||
"""
|
||||
try:
|
||||
resolution = resolve_prompt_pair(metadata)
|
||||
if resolution["state"] == "resolved":
|
||||
positive = resolution["positive"]
|
||||
expected = self.db_manager.get_prompt_by_hash(
|
||||
generate_prompt_hash(positive.get("storage_text", positive["text"]))
|
||||
)
|
||||
if expected and expected["id"] == prompt_context["id"]:
|
||||
ingest_resolved_pair(
|
||||
self.db_manager, image_path, metadata, resolution
|
||||
)
|
||||
return
|
||||
image_id = self.db_manager.link_image_to_prompt(
|
||||
prompt_id=prompt_context["id"], image_path=image_path, metadata=metadata
|
||||
)
|
||||
|
||||
@@ -116,6 +116,7 @@ class ComfyUIMetadataExtractor:
|
||||
|
||||
# Extract other common metadata fields
|
||||
metadata_fields = [
|
||||
"comfyassets_prompt_pairs",
|
||||
"parameters",
|
||||
"model",
|
||||
"sampler",
|
||||
|
||||
@@ -0,0 +1,145 @@
|
||||
"""Bounded per-node execution evidence embedded in ComfyUI PNG metadata."""
|
||||
|
||||
import json
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
EXECUTION_KEY = "comfyassets_prompt_pairs"
|
||||
MAX_EXECUTION_RECORDS = 16
|
||||
MAX_EXECUTION_NODES = 64
|
||||
MAX_EXECUTION_TEXT = 100_000
|
||||
MAX_EXECUTION_JSON = 500_000
|
||||
_FIELDS = (
|
||||
"text",
|
||||
"negative_text",
|
||||
"positive_text",
|
||||
"negative_effective_text",
|
||||
"prepend_text",
|
||||
"append_text",
|
||||
)
|
||||
|
||||
|
||||
class ExecutionEvidenceError(ValueError):
|
||||
"""Execution metadata is malformed, stale, truncated or ambiguous."""
|
||||
|
||||
|
||||
def record_pair_execution(
|
||||
extra_pnginfo: Any,
|
||||
node_id: Any,
|
||||
*,
|
||||
text: str,
|
||||
negative_text: str,
|
||||
positive_text: str,
|
||||
negative_effective_text: str,
|
||||
prepend_text: str = "",
|
||||
append_text: str = "",
|
||||
) -> None:
|
||||
"""Append evidence without changing workflow or sibling extension metadata."""
|
||||
if not isinstance(extra_pnginfo, dict) or node_id is None:
|
||||
return
|
||||
envelope = extra_pnginfo.get(EXECUTION_KEY)
|
||||
if not isinstance(envelope, dict) or envelope.get("version") != 1:
|
||||
envelope = {"version": 1, "nodes": {}}
|
||||
extra_pnginfo[EXECUTION_KEY] = envelope
|
||||
if envelope.get("truncated"):
|
||||
return
|
||||
nodes = envelope.get("nodes")
|
||||
if not isinstance(nodes, dict):
|
||||
envelope["truncated"] = True
|
||||
return
|
||||
if str(node_id) not in nodes and len(nodes) >= MAX_EXECUTION_NODES:
|
||||
envelope["truncated"] = True
|
||||
return
|
||||
node = nodes.setdefault(str(node_id), {"records": []})
|
||||
if not isinstance(node, dict) or not isinstance(node.get("records"), list):
|
||||
envelope["truncated"] = True
|
||||
return
|
||||
record = {
|
||||
"text": text,
|
||||
"negative_text": negative_text,
|
||||
"positive_text": positive_text,
|
||||
"negative_effective_text": negative_effective_text,
|
||||
"prepend_text": prepend_text,
|
||||
"append_text": append_text,
|
||||
}
|
||||
if any(
|
||||
not isinstance(value, str) or len(value) > MAX_EXECUTION_TEXT
|
||||
for value in record.values()
|
||||
):
|
||||
node["truncated"] = True
|
||||
return
|
||||
if record not in node["records"]:
|
||||
if len(node["records"]) >= MAX_EXECUTION_RECORDS:
|
||||
node["truncated"] = True
|
||||
return
|
||||
node["records"].append(record)
|
||||
try:
|
||||
if len(json.dumps(envelope, ensure_ascii=False)) > MAX_EXECUTION_JSON:
|
||||
extra_pnginfo[EXECUTION_KEY] = {
|
||||
"version": 1,
|
||||
"truncated": True,
|
||||
"nodes": {},
|
||||
}
|
||||
except (TypeError, ValueError, RecursionError):
|
||||
extra_pnginfo[EXECUTION_KEY] = {"version": 1, "truncated": True, "nodes": {}}
|
||||
|
||||
|
||||
def read_pair_execution(metadata: Any, node_id: str) -> Optional[Dict[str, str]]:
|
||||
"""Return one distinct record for a traced node; missing evidence stays absent."""
|
||||
if not isinstance(metadata, dict) or EXECUTION_KEY not in metadata:
|
||||
return None
|
||||
envelope = metadata[EXECUTION_KEY]
|
||||
if isinstance(envelope, str):
|
||||
if len(envelope) > MAX_EXECUTION_JSON:
|
||||
raise ExecutionEvidenceError("Execution evidence exceeds the size limit")
|
||||
try:
|
||||
envelope = json.loads(envelope)
|
||||
except (ValueError, TypeError, RecursionError) as error:
|
||||
raise ExecutionEvidenceError("Malformed execution evidence JSON") from error
|
||||
if not isinstance(envelope, dict) or envelope.get("version") != 1:
|
||||
raise ExecutionEvidenceError("Unsupported execution evidence format")
|
||||
nodes = envelope.get("nodes")
|
||||
if (
|
||||
envelope.get("truncated")
|
||||
or not isinstance(nodes, dict)
|
||||
or len(nodes) > MAX_EXECUTION_NODES
|
||||
):
|
||||
raise ExecutionEvidenceError("Truncated or malformed execution evidence")
|
||||
if node_id not in nodes:
|
||||
return None
|
||||
node = nodes[node_id]
|
||||
if not isinstance(node, dict) or node.get("truncated"):
|
||||
raise ExecutionEvidenceError("Truncated or malformed node execution evidence")
|
||||
records = node.get("records")
|
||||
if (
|
||||
not isinstance(records, list)
|
||||
or not records
|
||||
or len(records) > MAX_EXECUTION_RECORDS
|
||||
):
|
||||
raise ExecutionEvidenceError("Malformed node execution records")
|
||||
unique = {}
|
||||
for record in records:
|
||||
if not isinstance(record, dict) or any(
|
||||
not isinstance(record.get(key), str)
|
||||
or len(record[key]) > MAX_EXECUTION_TEXT
|
||||
for key in _FIELDS
|
||||
):
|
||||
raise ExecutionEvidenceError("Malformed execution text values")
|
||||
clean = {key: record[key] for key in _FIELDS}
|
||||
if clean["negative_effective_text"] != clean["negative_text"]:
|
||||
raise ExecutionEvidenceError("Negative execution text must be unchanged")
|
||||
unique[json.dumps(clean, sort_keys=True)] = clean
|
||||
if len(unique) != 1:
|
||||
raise ExecutionEvidenceError(
|
||||
"Multiple distinct executions have no image anchor"
|
||||
)
|
||||
return next(iter(unique.values()))
|
||||
|
||||
|
||||
def clear_pair_execution(request: Any) -> None:
|
||||
"""Discard client-carried execution records before a new validated run."""
|
||||
if not isinstance(request, dict):
|
||||
return
|
||||
extra = request.get("extra_data")
|
||||
png = extra.get("extra_pnginfo") if isinstance(extra, dict) else None
|
||||
if isinstance(png, dict):
|
||||
png.pop(EXECUTION_KEY, None)
|
||||
@@ -9,10 +9,16 @@ link-tracing rules of web/js/comfy-metadata.js.
|
||||
"""
|
||||
|
||||
import re
|
||||
from typing import Any, Dict, List, Optional
|
||||
from typing import Any, Dict, List, Optional, Tuple
|
||||
|
||||
PROMPT_MANAGER_TYPES = frozenset(
|
||||
{"PromptManager", "PromptManagerText", "PromptManagerFocused", "PromptManagerSplit"}
|
||||
{
|
||||
"PromptManager",
|
||||
"PromptManagerText",
|
||||
"PromptManagerFocused",
|
||||
"PromptManagerSplit",
|
||||
"PromptManagerPair",
|
||||
}
|
||||
)
|
||||
|
||||
# Upper bound on nodes expanded per graph; queued and embedded graphs are untrusted
|
||||
@@ -93,32 +99,65 @@ def has_sampler(graph: Any) -> bool:
|
||||
return isinstance(graph, dict) and bool(_positive_roots(graph))
|
||||
|
||||
|
||||
def positive_prompt_nodes(graph: Any) -> List[str]:
|
||||
"""Ids of PromptManager family nodes that feed a positive input.
|
||||
|
||||
Ordered by discovery, each id once. Malformed graphs yield an empty list.
|
||||
"""
|
||||
def positive_prompt_components(graph: Any) -> List[Tuple[str, str]]:
|
||||
"""Node/input identities reaching actual sampler positive wires."""
|
||||
if not isinstance(graph, dict):
|
||||
return []
|
||||
|
||||
found: List[str] = []
|
||||
found: List[Tuple[str, str]] = []
|
||||
seen = set()
|
||||
stack = list(reversed(_positive_roots(graph)))
|
||||
while stack and len(seen) < MAX_VISITS:
|
||||
node_id = str(stack.pop()[0])
|
||||
if node_id in seen:
|
||||
link = stack.pop()
|
||||
node_id, slot = str(link[0]), link[1]
|
||||
if (node_id, slot) in seen:
|
||||
continue
|
||||
seen.add(node_id)
|
||||
seen.add((node_id, slot))
|
||||
node = graph.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
if node.get("class_type") in PROMPT_MANAGER_TYPES:
|
||||
found.append(node_id)
|
||||
field = "text"
|
||||
if node.get("class_type") == "PromptManagerPair":
|
||||
field = {
|
||||
0: "text",
|
||||
1: "negative_text",
|
||||
2: "text",
|
||||
3: "negative_text",
|
||||
}.get(slot)
|
||||
if field is None:
|
||||
continue
|
||||
identity = (node_id, field)
|
||||
if identity not in found:
|
||||
found.append(identity)
|
||||
continue
|
||||
stack.extend(reversed(_upstream_links(_inputs(node))))
|
||||
return found
|
||||
|
||||
|
||||
def positive_prompt_nodes(graph: Any) -> List[str]:
|
||||
"""Nodes whose main editor reaches a sampler's positive input."""
|
||||
return list(
|
||||
dict.fromkeys(
|
||||
node_id
|
||||
for node_id, field in positive_prompt_components(graph)
|
||||
if field == "text"
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def run_prompt_components(graph: Any) -> List[Tuple[str, str]]:
|
||||
"""Actual positive components, retaining standalone main-editor fallback."""
|
||||
components = positive_prompt_components(graph)
|
||||
if components or not isinstance(graph, dict) or has_sampler(graph):
|
||||
return components
|
||||
return [
|
||||
(str(node_id), "text")
|
||||
for node_id, node in graph.items()
|
||||
if isinstance(node, dict) and node.get("class_type") in PROMPT_MANAGER_TYPES
|
||||
]
|
||||
|
||||
|
||||
def run_prompt_nodes(graph: Any) -> List[str]:
|
||||
"""Prompt nodes that represent a run: the positive ones, if any can be found.
|
||||
|
||||
@@ -228,3 +267,67 @@ def resolve_text(graph: Any, node_id: str) -> Optional[str]:
|
||||
|
||||
text = resolve(_inputs(graph.get(str(node_id))).get("text"))
|
||||
return text.strip() if isinstance(text, str) and text.strip() else None
|
||||
|
||||
|
||||
def intrinsic_pair_nodes(graph: Any) -> List[str]:
|
||||
"""Paired nodes whose matching outputs feed the same sampler, or standalone."""
|
||||
if not isinstance(graph, dict) or len(graph) > MAX_VISITS:
|
||||
return []
|
||||
if not has_sampler(graph):
|
||||
return [
|
||||
str(key)
|
||||
for key, value in graph.items()
|
||||
if isinstance(value, dict)
|
||||
and value.get("class_type") == "PromptManagerPair"
|
||||
]
|
||||
|
||||
def origin(link: Any, role: str, kind: str = "conditioning") -> Optional[str]:
|
||||
seen = set()
|
||||
while _is_link(link) and len(seen) < MAX_DEPTH:
|
||||
node_id, slot = str(link[0]), link[1]
|
||||
if (node_id, slot, kind) in seen:
|
||||
return None
|
||||
seen.add((node_id, slot, kind))
|
||||
node = graph.get(node_id)
|
||||
if not isinstance(node, dict):
|
||||
return None
|
||||
name, inputs = node.get("class_type"), _inputs(node)
|
||||
if name == "PromptManagerPair":
|
||||
expected = {
|
||||
("positive", "conditioning"): 0,
|
||||
("negative", "conditioning"): 1,
|
||||
("positive", "string"): 2,
|
||||
("negative", "string"): 3,
|
||||
}
|
||||
return node_id if slot == expected[(role, kind)] else None
|
||||
if name == "CLIPTextEncode" and slot == 0 and kind == "conditioning":
|
||||
link, kind = inputs.get("text"), "string"
|
||||
elif name == "Reroute" and slot == 0:
|
||||
keys = [key for key in ("value", "input", "in") if key in inputs]
|
||||
if len(keys) != 1:
|
||||
return None
|
||||
link = inputs[keys[0]]
|
||||
else:
|
||||
return None
|
||||
return None
|
||||
|
||||
found = []
|
||||
for node in graph.values():
|
||||
if not isinstance(node, dict):
|
||||
continue
|
||||
inputs = _inputs(node)
|
||||
if not _is_link(inputs.get("positive")) or not _is_link(inputs.get("negative")):
|
||||
continue
|
||||
if (
|
||||
not _is_link(inputs.get("model"))
|
||||
and "sampler" not in str(node.get("class_type", "")).lower()
|
||||
):
|
||||
continue
|
||||
positive = origin(inputs["positive"], "positive")
|
||||
if (
|
||||
positive
|
||||
and positive == origin(inputs["negative"], "negative")
|
||||
and positive not in found
|
||||
):
|
||||
found.append(positive)
|
||||
return found
|
||||
|
||||
@@ -0,0 +1,490 @@
|
||||
"""Conservatively recover a prompt pair from one sampler or A1111 parameters.
|
||||
|
||||
``resolve_prompt_pair`` returns a dictionary with ``state``, ``positive``,
|
||||
``negative``, ``sampler_id``, ``reason`` and ``candidates``. Each recovered side
|
||||
contains effective ``text``, deduplicated-recipe ``storage_text`` and
|
||||
``provenance``; ComfyUI provenance records the exact output
|
||||
slot and complete followed link path. States distinguish a resolved pair (an
|
||||
explicit empty negative is valid), absent data, unsupported/unresolved data and
|
||||
ambiguous sampler selection. Unknown conditioning semantics are never flattened.
|
||||
|
||||
Future nodes can register declarative ``OutputAdapter`` entries by
|
||||
``(class_type, output_slot)``. This does not prescribe a future paired-node schema.
|
||||
Effective text uses a unique bounded execution record for traced paired nodes
|
||||
when available; cached graph-only results remain explicitly inferred.
|
||||
The resolver has no
|
||||
ComfyUI, database, configuration or filesystem dependencies.
|
||||
"""
|
||||
|
||||
import re
|
||||
from dataclasses import dataclass
|
||||
from typing import Any, Dict, Mapping, Optional, Tuple
|
||||
|
||||
from .pair_execution import ExecutionEvidenceError, read_pair_execution
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class OutputAdapter:
|
||||
"""A known output's text input, output kind, and optional PM composition.
|
||||
|
||||
``kind`` is ``conditioning`` or ``string``. Composition is opt-in and follows
|
||||
the existing PromptManager prepend/main/append rule. Adapters describe pure
|
||||
text provenance, not arbitrary executable transforms. ``storage_input`` keeps
|
||||
the resolved input as the recipe even without composition (negative editor).
|
||||
"""
|
||||
|
||||
input_name: str
|
||||
kind: str
|
||||
compose: bool = False
|
||||
storage_input: bool = False
|
||||
|
||||
|
||||
_SAMPLERS = frozenset({"KSampler", "KSamplerAdvanced", "SamplerCustom"})
|
||||
_PROMPT_NODES = frozenset(
|
||||
{"PromptManager", "PromptManagerFocused", "PromptManagerSplit"}
|
||||
)
|
||||
_DEFAULT_ADAPTERS = {
|
||||
(name, slot): OutputAdapter("text", kind, True)
|
||||
for name in _PROMPT_NODES
|
||||
for slot, kind in ((0, "conditioning"), (1, "string"))
|
||||
}
|
||||
for _slot, _kind in ((0, "conditioning"), (2, "string")):
|
||||
_DEFAULT_ADAPTERS[("PromptManagerPair", _slot)] = OutputAdapter("text", _kind, True)
|
||||
for _slot, _kind in ((1, "conditioning"), (3, "string")):
|
||||
_DEFAULT_ADAPTERS[("PromptManagerPair", _slot)] = OutputAdapter(
|
||||
"negative_text", _kind, storage_input=True
|
||||
)
|
||||
_DEFAULT_ADAPTERS[("PromptManagerText", 0)] = OutputAdapter("text", "string", True)
|
||||
_DEFAULT_ADAPTERS[("CLIPTextEncode", 0)] = OutputAdapter("text", "conditioning")
|
||||
for _name in ("PrimitiveString", "PrimitiveStringMultiline"):
|
||||
_DEFAULT_ADAPTERS[(_name, 0)] = OutputAdapter("value", "string")
|
||||
for _name in ("Text Multiline", "String Literal", "ShowText|pysssss"):
|
||||
_DEFAULT_ADAPTERS[(_name, 0)] = OutputAdapter("text", "string")
|
||||
|
||||
|
||||
class _Unresolved(Exception):
|
||||
pass
|
||||
|
||||
|
||||
def _link(value: Any) -> bool:
|
||||
return (
|
||||
isinstance(value, (list, tuple))
|
||||
and len(value) == 2
|
||||
and isinstance(value[0], (str, int))
|
||||
and not isinstance(value[0], bool)
|
||||
and isinstance(value[1], int)
|
||||
and not isinstance(value[1], bool)
|
||||
and value[1] >= 0
|
||||
)
|
||||
|
||||
|
||||
def _result(
|
||||
state: str,
|
||||
positive: Optional[dict] = None,
|
||||
negative: Optional[dict] = None,
|
||||
sampler_id: Optional[str] = None,
|
||||
reason: str = "",
|
||||
candidates: Optional[list] = None,
|
||||
) -> Dict[str, Any]:
|
||||
return {
|
||||
"state": state,
|
||||
"positive": positive,
|
||||
"negative": negative,
|
||||
"sampler_id": sampler_id,
|
||||
"reason": reason,
|
||||
"candidates": candidates or [],
|
||||
}
|
||||
|
||||
|
||||
class _GraphResolver:
|
||||
def __init__(
|
||||
self,
|
||||
graph: dict,
|
||||
adapters: Mapping[Tuple[str, int], OutputAdapter],
|
||||
max_visits: int,
|
||||
max_depth: int,
|
||||
max_text_length: int,
|
||||
metadata: Optional[dict] = None,
|
||||
) -> None:
|
||||
self.graph = graph
|
||||
self.adapters = {**_DEFAULT_ADAPTERS, **adapters}
|
||||
self.budget = max_visits
|
||||
self.max_depth = max_depth
|
||||
self.max_text_length = max_text_length
|
||||
self.memo: dict = {}
|
||||
self.active: set = set()
|
||||
self.metadata = metadata or {}
|
||||
self.text_sources: dict = {}
|
||||
|
||||
def node(self, node_id: str) -> Tuple[str, dict]:
|
||||
node = self.graph.get(node_id)
|
||||
if not isinstance(node, dict) or not isinstance(node.get("inputs"), dict):
|
||||
raise _Unresolved("Missing or malformed graph node")
|
||||
if any(not isinstance(key, str) for key in node["inputs"]):
|
||||
raise _Unresolved("Malformed node input names")
|
||||
return str(node.get("class_type", "")), node["inputs"]
|
||||
|
||||
def bounded(self, text: str) -> str:
|
||||
if len(text) > self.max_text_length:
|
||||
raise _Unresolved("Resolved text exceeds the length limit")
|
||||
return text
|
||||
|
||||
def text(self, value: Any, kind: str, depth: int = 0) -> Tuple[str, list, str]:
|
||||
if depth >= self.max_depth:
|
||||
raise _Unresolved("Resolution depth limit reached")
|
||||
if isinstance(value, str) and kind == "string":
|
||||
return self.bounded(value), [], value
|
||||
if not _link(value):
|
||||
raise _Unresolved("Expected a typed graph link or text input")
|
||||
node_id, slot = str(value[0]), value[1]
|
||||
key = (node_id, slot, kind)
|
||||
if key in self.active:
|
||||
raise _Unresolved("Cyclic prompt graph")
|
||||
if key in self.memo:
|
||||
return self.memo[key]
|
||||
self.budget -= 1
|
||||
if self.budget < 0:
|
||||
raise _Unresolved("Resolution visit limit reached")
|
||||
name, inputs = self.node(node_id)
|
||||
self.active.add(key)
|
||||
try:
|
||||
if name == "Reroute" and slot == 0:
|
||||
keys = [key for key in ("value", "input", "in") if key in inputs]
|
||||
if len(keys) != 1:
|
||||
raise _Unresolved("Reroute input is not uniquely identifiable")
|
||||
text, path, storage_text = self.text(inputs[keys[0]], kind, depth + 1)
|
||||
elif (name, slot) in self.adapters:
|
||||
adapter = self.adapters[(name, slot)]
|
||||
if not isinstance(adapter, OutputAdapter) or adapter.kind != kind:
|
||||
raise _Unresolved("Output slot has an incompatible type")
|
||||
evidence = (
|
||||
self.execution_record(node_id, inputs)
|
||||
if name == "PromptManagerPair"
|
||||
else None
|
||||
)
|
||||
if evidence is not None:
|
||||
text = evidence[
|
||||
(
|
||||
"positive_text"
|
||||
if adapter.input_name == "text"
|
||||
else "negative_effective_text"
|
||||
)
|
||||
]
|
||||
storage_text = evidence[adapter.input_name]
|
||||
path = []
|
||||
self.text_sources[(node_id, slot)] = "execution"
|
||||
else:
|
||||
text, path, storage_text = self.text(
|
||||
inputs.get(adapter.input_name), "string", depth + 1
|
||||
)
|
||||
if adapter.compose or adapter.storage_input:
|
||||
storage_text = text
|
||||
if adapter.compose:
|
||||
before, before_path, _ = self.text(
|
||||
inputs.get("prepend_text", ""), "string", depth + 1
|
||||
)
|
||||
after, after_path, _ = self.text(
|
||||
inputs.get("append_text", ""), "string", depth + 1
|
||||
)
|
||||
text = self.bounded(
|
||||
" ".join(
|
||||
part
|
||||
for part in (before.strip(), text, after.strip())
|
||||
if part
|
||||
)
|
||||
)
|
||||
path = before_path + path + after_path
|
||||
elif (
|
||||
name in ("StringConcatenate", "Text Concatenate")
|
||||
and slot == 0
|
||||
and kind == "string"
|
||||
):
|
||||
if name == "StringConcatenate":
|
||||
keys = ["string_a", "string_b"]
|
||||
delimiter = inputs.get("delimiter", "")
|
||||
clean = False
|
||||
else:
|
||||
keys = sorted(
|
||||
key for key in inputs if re.fullmatch(r"text_[a-z]", key)
|
||||
)
|
||||
delimiter = inputs.get("delimiter", " ")
|
||||
clean = (
|
||||
str(inputs.get("clean_whitespace", "true")).lower() == "true"
|
||||
)
|
||||
if not keys or not isinstance(delimiter, str):
|
||||
raise _Unresolved("Malformed concatenation")
|
||||
parts, path, size = [], [], 0
|
||||
for input_name in keys:
|
||||
part, followed, _ = self.text(
|
||||
inputs.get(input_name, ""), "string", depth + 1
|
||||
)
|
||||
if clean:
|
||||
part = part.strip()
|
||||
path.extend(followed)
|
||||
if part or name == "StringConcatenate":
|
||||
size += len(part) + (len(delimiter) if parts else 0)
|
||||
if size > self.max_text_length:
|
||||
raise _Unresolved("Resolved text exceeds the length limit")
|
||||
parts.append(part)
|
||||
text = delimiter.join(parts)
|
||||
storage_text = text
|
||||
else:
|
||||
raise _Unresolved("Unsupported node semantics or output slot")
|
||||
# Memoized paths stay bounded even for diamond-shaped string graphs.
|
||||
path = [{"node_id": node_id, "output_slot": slot}] + path
|
||||
if len(path) > self.max_depth * max(1, len(self.graph)):
|
||||
raise _Unresolved("Provenance path limit reached")
|
||||
self.memo[key] = (text, path, storage_text)
|
||||
return text, path, storage_text
|
||||
finally:
|
||||
self.active.discard(key)
|
||||
|
||||
def execution_record(self, node_id: str, inputs: dict) -> Optional[dict]:
|
||||
try:
|
||||
record = read_pair_execution(self.metadata, node_id)
|
||||
except ExecutionEvidenceError as error:
|
||||
raise _Unresolved(str(error)) from error
|
||||
if record is None:
|
||||
return None
|
||||
for name in ("text", "negative_text", "prepend_text", "append_text"):
|
||||
try:
|
||||
expected = self.text(inputs.get(name, ""), "string")[0]
|
||||
except _Unresolved as error:
|
||||
if str(error) == "Unsupported node semantics or output slot":
|
||||
value = inputs.get(name)
|
||||
if _link(value):
|
||||
source_name, _ = self.node(str(value[0]))
|
||||
known_slots = {
|
||||
slot
|
||||
for (kind, slot) in self.adapters
|
||||
if kind == source_name
|
||||
}
|
||||
if known_slots and value[1] not in known_slots:
|
||||
raise
|
||||
if (
|
||||
source_name
|
||||
in ("StringConcatenate", "Text Concatenate", "Reroute")
|
||||
and value[1] != 0
|
||||
):
|
||||
raise
|
||||
continue
|
||||
raise
|
||||
if name in ("prepend_text", "append_text"):
|
||||
expected = expected.strip()
|
||||
if expected != record[name]:
|
||||
raise _Unresolved(
|
||||
"Execution evidence does not match recorded node inputs"
|
||||
)
|
||||
return record
|
||||
|
||||
def side(self, value: Any, sampler_id: str) -> dict:
|
||||
text, path, storage_text = self.text(value, "conditioning")
|
||||
origin = path[0]
|
||||
return {
|
||||
"text": text,
|
||||
"storage_text": storage_text,
|
||||
"provenance": {
|
||||
"format": "comfyui",
|
||||
"sampler_id": sampler_id,
|
||||
**origin,
|
||||
"path": path,
|
||||
"text_source": self.text_sources.get(
|
||||
(origin["node_id"], origin["output_slot"]), "graph-inferred"
|
||||
),
|
||||
},
|
||||
}
|
||||
|
||||
def pair(self, sampler_id: str) -> dict:
|
||||
positive = negative = None
|
||||
try:
|
||||
_, inputs = self.node(sampler_id)
|
||||
# A second sampler in the latent ancestry denotes a multistage run;
|
||||
# associating the final image with one prompt pair would be a guess.
|
||||
if "latent_image" in inputs:
|
||||
self.latent_source(inputs["latent_image"], set(), 0)
|
||||
positive = self.side(inputs.get("positive"), sampler_id)
|
||||
if not positive["text"].strip():
|
||||
raise _Unresolved("Positive prompt is empty")
|
||||
if "negative" not in inputs:
|
||||
return _result(
|
||||
"absent",
|
||||
positive,
|
||||
sampler_id=sampler_id,
|
||||
reason="Negative input is absent",
|
||||
)
|
||||
negative = self.side(inputs["negative"], sampler_id)
|
||||
return _result("resolved", positive, negative, sampler_id)
|
||||
except _Unresolved as error:
|
||||
return _result("unresolved", positive, negative, sampler_id, str(error))
|
||||
|
||||
def latent_source(self, value: Any, seen: set, depth: int) -> None:
|
||||
if not _link(value) or depth >= self.max_depth:
|
||||
raise _Unresolved("Unknown latent ancestry")
|
||||
node_id, slot = str(value[0]), value[1]
|
||||
if node_id in seen or len(seen) >= max(0, self.budget):
|
||||
raise _Unresolved("Cyclic or excessive latent ancestry")
|
||||
seen.add(node_id)
|
||||
name, inputs = self.node(node_id)
|
||||
if name in _SAMPLERS:
|
||||
raise _Unresolved("Multistage sampler ancestry is ambiguous")
|
||||
if name in ("EmptyLatentImage", "EmptySD3LatentImage") and slot == 0:
|
||||
return
|
||||
if name == "Reroute" and slot == 0 and "value" in inputs:
|
||||
self.latent_source(inputs["value"], seen, depth + 1)
|
||||
return
|
||||
raise _Unresolved("Unsupported latent ancestry")
|
||||
|
||||
def anchor(self, node_id: str) -> str:
|
||||
seen = set()
|
||||
slot = 0
|
||||
while len(seen) < self.max_depth and len(seen) < self.budget:
|
||||
if node_id in seen:
|
||||
raise _Unresolved("Cyclic output ancestry")
|
||||
seen.add(node_id)
|
||||
name, inputs = self.node(node_id)
|
||||
if name in _SAMPLERS and slot == 0:
|
||||
return node_id
|
||||
key = {
|
||||
"SaveImage": "images",
|
||||
"PreviewImage": "images",
|
||||
"VAEDecode": "samples",
|
||||
"VAEDecodeTiled": "samples",
|
||||
"Reroute": "value",
|
||||
}.get(name)
|
||||
value = inputs.get(key) if key and slot == 0 else None
|
||||
if not _link(value):
|
||||
raise _Unresolved("Unsupported output ancestry")
|
||||
node_id, slot = str(value[0]), value[1]
|
||||
raise _Unresolved("Output ancestry limit reached")
|
||||
|
||||
|
||||
def _parameters(value: Any, max_text_length: int) -> dict:
|
||||
if not isinstance(value, str) or not value:
|
||||
return _result("absent", reason="No A1111 parameters")
|
||||
if len(value) > max_text_length:
|
||||
return _result("unresolved", reason="Parameters exceed the length limit")
|
||||
footers = list(re.finditer(r"(?m)^Steps:\s*\d+\s*,[^\n]*$", value))
|
||||
if len(footers) != 1 or value[footers[0].end() :].strip():
|
||||
return _result("unresolved", reason="A1111 settings footer is not unique")
|
||||
body = value[: footers[0].start()].rstrip("\r\n")
|
||||
delimiter = "\nNegative prompt:"
|
||||
if body.count(delimiter) > 1:
|
||||
return _result("ambiguous", reason="Multiple negative prompt delimiters")
|
||||
positive, found, negative = body.partition(delimiter)
|
||||
positive = positive.strip()
|
||||
if not positive:
|
||||
return _result("unresolved", reason="Positive prompt is empty")
|
||||
provenance = {
|
||||
"format": "a1111",
|
||||
"field": "parameters",
|
||||
"text_source": "embedded-parameters",
|
||||
}
|
||||
positive_side = {
|
||||
"text": positive,
|
||||
"storage_text": positive,
|
||||
"provenance": {**provenance, "role": "positive"},
|
||||
}
|
||||
if not found:
|
||||
return _result(
|
||||
"absent", positive_side, reason="Negative prompt delimiter is absent"
|
||||
)
|
||||
negative_side = {
|
||||
"text": negative.strip(),
|
||||
"storage_text": negative.strip(),
|
||||
"provenance": {**provenance, "role": "negative"},
|
||||
}
|
||||
return _result("resolved", positive_side, negative_side)
|
||||
|
||||
|
||||
def resolve_prompt_pair(
|
||||
metadata: Any,
|
||||
*,
|
||||
output_node_id: Optional[str] = None,
|
||||
adapters: Optional[Mapping[Tuple[str, int], OutputAdapter]] = None,
|
||||
max_visits: int = 500,
|
||||
max_depth: int = 64,
|
||||
max_text_length: int = 100_000,
|
||||
) -> Dict[str, Any]:
|
||||
"""Resolve parsed metadata; a present API graph is authoritative over text.
|
||||
|
||||
Accepts ``{"prompt": api_graph, "parameters": optional_a1111_text}`` or a
|
||||
direct API graph. An optional saver/sampler node ID anchors branch selection.
|
||||
Unanchored multiple samplers are ambiguous, even when their texts are equal.
|
||||
"""
|
||||
if not isinstance(metadata, dict) or not metadata:
|
||||
return _result("absent", reason="No parsed prompt metadata")
|
||||
if "prompt" in metadata:
|
||||
graph = metadata["prompt"]
|
||||
elif "parameters" in metadata:
|
||||
return _parameters(metadata["parameters"], max_text_length)
|
||||
else:
|
||||
graph = metadata
|
||||
if not isinstance(graph, dict) or not graph:
|
||||
return _result("unresolved", reason="API graph is missing or malformed")
|
||||
if len(graph) > max_visits or max_depth <= 0 or max_text_length <= 0:
|
||||
return _result("unresolved", reason="Resolution limit reached")
|
||||
normalized = {str(node_id): node for node_id, node in graph.items()}
|
||||
if len(normalized) != len(graph):
|
||||
return _result("unresolved", reason="Graph node IDs are not unique")
|
||||
resolver = _GraphResolver(
|
||||
normalized,
|
||||
adapters or {},
|
||||
max_visits,
|
||||
max_depth,
|
||||
max_text_length,
|
||||
metadata if "prompt" in metadata else None,
|
||||
)
|
||||
if output_node_id is not None:
|
||||
try:
|
||||
sampler_ids = [resolver.anchor(str(output_node_id))]
|
||||
except _Unresolved as error:
|
||||
return _result("unresolved", reason=str(error))
|
||||
else:
|
||||
sampler_ids = [
|
||||
node_id
|
||||
for node_id, node in normalized.items()
|
||||
if isinstance(node, dict)
|
||||
and isinstance(node.get("class_type"), str)
|
||||
and node["class_type"] in _SAMPLERS
|
||||
]
|
||||
if not sampler_ids:
|
||||
has_pair_inputs = any(
|
||||
isinstance(node, dict)
|
||||
and isinstance(node.get("inputs"), dict)
|
||||
and ("positive" in node["inputs"] or "negative" in node["inputs"])
|
||||
for node in normalized.values()
|
||||
)
|
||||
return _result(
|
||||
"unresolved" if has_pair_inputs else "absent",
|
||||
reason="No supported sampler pair found",
|
||||
)
|
||||
if len(sampler_ids) > 1:
|
||||
return _result(
|
||||
"ambiguous",
|
||||
reason="Multiple sampler branches",
|
||||
candidates=[resolver.pair(node_id) for node_id in sampler_ids],
|
||||
)
|
||||
return resolver.pair(sampler_ids[0])
|
||||
|
||||
|
||||
def resolve_text_input(
|
||||
graph: Any,
|
||||
node_id: str,
|
||||
input_name: str,
|
||||
*,
|
||||
max_visits: int = 500,
|
||||
max_depth: int = 64,
|
||||
max_text_length: int = 100_000,
|
||||
) -> Optional[str]:
|
||||
"""Resolve one raw STRING input; empty text is distinct from unknown input."""
|
||||
if not isinstance(graph, dict) or len(graph) > max_visits:
|
||||
return None
|
||||
normalized = {str(key): value for key, value in graph.items()}
|
||||
if len(normalized) != len(graph) or max_depth <= 0 or max_text_length <= 0:
|
||||
return None
|
||||
resolver = _GraphResolver(normalized, {}, max_visits, max_depth, max_text_length)
|
||||
try:
|
||||
_, inputs = resolver.node(str(node_id))
|
||||
return resolver.text(inputs.get(input_name), "string")[0]
|
||||
except _Unresolved:
|
||||
return None
|
||||
@@ -29,11 +29,19 @@ from typing import Any, Callable, List, Optional, Tuple
|
||||
try:
|
||||
from .hashing import generate_prompt_hash
|
||||
from .logging_config import get_logger
|
||||
from .prompt_graph import resolve_text, run_prompt_nodes
|
||||
from .prompt_graph import intrinsic_pair_nodes, resolve_text, run_prompt_components
|
||||
from .prompt_pair_resolver import resolve_text_input
|
||||
from .pair_execution import clear_pair_execution
|
||||
except ImportError:
|
||||
from utils.hashing import generate_prompt_hash
|
||||
from utils.logging_config import get_logger
|
||||
from utils.prompt_graph import resolve_text, run_prompt_nodes
|
||||
from utils.prompt_graph import (
|
||||
intrinsic_pair_nodes,
|
||||
resolve_text,
|
||||
run_prompt_components,
|
||||
)
|
||||
from utils.prompt_pair_resolver import resolve_text_input
|
||||
from utils.pair_execution import clear_pair_execution
|
||||
|
||||
logger = get_logger("prompt_manager.usage_tracking")
|
||||
|
||||
@@ -88,10 +96,10 @@ class QueuedRuns:
|
||||
self._lock = threading.Lock()
|
||||
|
||||
def put(self, prompt_id: str, runs: List[Run]) -> None:
|
||||
if not runs:
|
||||
return
|
||||
with self._lock:
|
||||
self._entries.pop(prompt_id, None)
|
||||
if not runs:
|
||||
return
|
||||
self._entries[prompt_id] = list(runs)
|
||||
while len(self._entries) > self._max:
|
||||
self._entries.popitem(last=False)
|
||||
@@ -108,6 +116,8 @@ class QueuedRuns:
|
||||
# Shared between the queue hook and the nodes
|
||||
PENDING_FIRST_USE = PendingFirstUse()
|
||||
QUEUED_RUNS = QueuedRuns()
|
||||
PENDING_PAIR_FIRST_USE = PendingFirstUse()
|
||||
QUEUED_PAIR_RUNS = QueuedRuns()
|
||||
|
||||
|
||||
def queued_prompt_runs(graph: Any) -> List[Run]:
|
||||
@@ -117,8 +127,13 @@ def queued_prompt_runs(graph: Any) -> List[Run]:
|
||||
"""
|
||||
runs: List[Run] = []
|
||||
seen = set()
|
||||
for node_id in run_prompt_nodes(graph):
|
||||
text = resolve_text(graph, node_id)
|
||||
for node_id, field in run_prompt_components(graph):
|
||||
node = graph.get(node_id, {}) if isinstance(graph, dict) else {}
|
||||
text = (
|
||||
resolve_text_input(graph, node_id, field)
|
||||
if node.get("class_type") == "PromptManagerPair"
|
||||
else resolve_text(graph, node_id)
|
||||
)
|
||||
if not text:
|
||||
continue
|
||||
prompt_hash = generate_prompt_hash(text)
|
||||
@@ -129,6 +144,62 @@ def queued_prompt_runs(graph: Any) -> List[Run]:
|
||||
return runs
|
||||
|
||||
|
||||
def pair_run_key(positive: str, negative: str) -> str:
|
||||
"""Stable ordered identity; an empty negative remains an explicit component."""
|
||||
return generate_prompt_hash(positive) + ":" + generate_prompt_hash(negative)
|
||||
|
||||
|
||||
def queued_pair_runs(graph: Any) -> List[Tuple[str, str, str]]:
|
||||
"""Resolve raw paired recipes only when both outputs form the used pair."""
|
||||
found, seen = [], set()
|
||||
for node_id in intrinsic_pair_nodes(graph):
|
||||
positive = resolve_text_input(graph, node_id, "text")
|
||||
negative = resolve_text_input(graph, node_id, "negative_text")
|
||||
if positive is None or not positive.strip() or negative is None:
|
||||
continue
|
||||
key = pair_run_key(positive, negative)
|
||||
if key not in seen:
|
||||
seen.add(key)
|
||||
found.append((key, positive, negative))
|
||||
return found
|
||||
|
||||
|
||||
def find_pair_id(db: Any, positive: str, negative: str) -> Optional[int]:
|
||||
"""Find a saved pair without creating or overwriting its metadata."""
|
||||
positive_row = db.get_prompt_by_hash(generate_prompt_hash(positive))
|
||||
negative_row = (
|
||||
db.get_prompt_by_hash(generate_prompt_hash(negative))
|
||||
if negative.strip()
|
||||
else None
|
||||
)
|
||||
if not positive_row or (negative.strip() and not negative_row):
|
||||
return None
|
||||
row = (
|
||||
db.model.get_connection()
|
||||
.execute(
|
||||
"SELECT id FROM prompt_pairs "
|
||||
"WHERE positive_prompt_id=? AND negative_prompt_id IS ?",
|
||||
(positive_row["id"], negative_row["id"] if negative_row else None),
|
||||
)
|
||||
.fetchone()
|
||||
)
|
||||
return row["id"] if row else None
|
||||
|
||||
|
||||
def count_pair_runs(
|
||||
runs: List[Tuple[str, str, str]],
|
||||
db: Any,
|
||||
pending: PendingFirstUse,
|
||||
) -> None:
|
||||
"""Count existing recipes at dequeue and defer new recipes until first save."""
|
||||
for key, positive, negative in runs:
|
||||
pair_id = find_pair_id(db, positive, negative)
|
||||
if pair_id:
|
||||
db.record_prompt_pair_use(pair_id)
|
||||
else:
|
||||
pending.add(key)
|
||||
|
||||
|
||||
def count_runs(runs: List[Run], db: Any, pending: PendingFirstUse) -> None:
|
||||
"""Record one use per prompt in the database; remember the rest for first save."""
|
||||
for prompt_hash, _text in runs:
|
||||
@@ -139,7 +210,12 @@ def count_runs(runs: List[Run], db: Any, pending: PendingFirstUse) -> None:
|
||||
pending.add(prompt_hash)
|
||||
|
||||
|
||||
def handle_queued_prompt(json_data: Any, queued: QueuedRuns = QUEUED_RUNS) -> Any:
|
||||
def handle_queued_prompt(
|
||||
json_data: Any,
|
||||
queued: QueuedRuns = QUEUED_RUNS,
|
||||
*,
|
||||
pair_queued: QueuedRuns = QUEUED_PAIR_RUNS,
|
||||
) -> Any:
|
||||
"""ComfyUI on-prompt handler body: stash the runs, always return the request.
|
||||
|
||||
Adds a canonical UUID4 ``prompt_id`` when the client sent none, so the stash
|
||||
@@ -147,15 +223,21 @@ def handle_queued_prompt(json_data: Any, queued: QueuedRuns = QUEUED_RUNS) -> An
|
||||
the request has not been validated yet.
|
||||
"""
|
||||
try:
|
||||
clear_pair_execution(json_data)
|
||||
graph = json_data.get("prompt") if isinstance(json_data, dict) else None
|
||||
runs = queued_prompt_runs(graph)
|
||||
if not runs:
|
||||
pair_runs = queued_pair_runs(graph)
|
||||
if not runs and not pair_runs:
|
||||
if isinstance(json_data, dict) and json_data.get("prompt_id") is not None:
|
||||
queued.put(str(json_data["prompt_id"]), [])
|
||||
pair_queued.put(str(json_data["prompt_id"]), [])
|
||||
return json_data
|
||||
prompt_id = json_data.get("prompt_id")
|
||||
if prompt_id is None:
|
||||
prompt_id = str(uuid.uuid4())
|
||||
json_data["prompt_id"] = prompt_id
|
||||
queued.put(str(prompt_id), runs)
|
||||
pair_queued.put(str(prompt_id), pair_runs)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not resolve queued prompt runs: {e}")
|
||||
return json_data
|
||||
@@ -166,12 +248,18 @@ def on_dequeue(
|
||||
queued: QueuedRuns,
|
||||
db_factory: Callable[[], Any],
|
||||
pending: PendingFirstUse,
|
||||
*,
|
||||
pair_queued: QueuedRuns = QUEUED_PAIR_RUNS,
|
||||
pair_pending: PendingFirstUse = PENDING_PAIR_FIRST_USE,
|
||||
) -> None:
|
||||
"""Count the runs stashed for a prompt ComfyUI is about to execute."""
|
||||
try:
|
||||
runs = queued.pop(str(prompt_id))
|
||||
if runs:
|
||||
count_runs(runs, db_factory(), pending)
|
||||
pair_runs = pair_queued.pop(str(prompt_id))
|
||||
if runs or pair_runs:
|
||||
db = db_factory()
|
||||
count_runs(runs, db, pending)
|
||||
count_pair_runs(pair_runs, db, pair_pending)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not count runs of dequeued prompt {prompt_id}: {e}")
|
||||
|
||||
@@ -181,10 +269,14 @@ def count_queued_prompt(
|
||||
) -> Any:
|
||||
"""Legacy on-prompt handler body: count at submit time (no prompt_queue)."""
|
||||
try:
|
||||
clear_pair_execution(json_data)
|
||||
graph = json_data.get("prompt") if isinstance(json_data, dict) else None
|
||||
runs = queued_prompt_runs(graph)
|
||||
if runs:
|
||||
count_runs(runs, db_factory(), pending)
|
||||
pair_runs = queued_pair_runs(graph)
|
||||
if runs or pair_runs:
|
||||
db = db_factory()
|
||||
count_runs(runs, db, pending)
|
||||
count_pair_runs(pair_runs, db, PENDING_PAIR_FIRST_USE)
|
||||
except Exception as e:
|
||||
logger.warning(f"Could not count queued prompt runs: {e}")
|
||||
return json_data
|
||||
|
||||
@@ -3,8 +3,8 @@
|
||||
import { app } from "../../scripts/app.js";
|
||||
import { api } from "../../scripts/api.js";
|
||||
const { installClassicControls } = await import(api.fileURL('/prompt_manager/js/classic-node.js'));
|
||||
const { installFocusedNode, installSplitNode } = await import(api.fileURL('/prompt_manager/js/focused-node.js'));
|
||||
const promptManagerTypes = new Set(['PromptManager', 'PromptManagerText', 'PromptManagerFocused', 'PromptManagerSplit']);
|
||||
const { installFocusedNode, installSplitNode, installPairNode } = await import(api.fileURL('/prompt_manager/js/focused-node.js'));
|
||||
const promptManagerTypes = new Set(['PromptManager', 'PromptManagerText', 'PromptManagerFocused', 'PromptManagerSplit', 'PromptManagerPair']);
|
||||
|
||||
// ============================================================================
|
||||
// Top Bar Button Integration (similar to rgthree-comfy approach)
|
||||
@@ -262,20 +262,12 @@ app.registerExtension({
|
||||
if (promptManagerTypes.has(node.type) && output[node.id]) {
|
||||
const outputNode = output[node.id];
|
||||
|
||||
// Skip if text input is connected to another node (value is a link reference array)
|
||||
if (outputNode.inputs && Array.isArray(outputNode.inputs.text)) {
|
||||
continue;
|
||||
for (const name of node.type === 'PromptManagerPair' ? ['text','negative_text'] : ['text']) {
|
||||
if (!outputNode.inputs || Array.isArray(outputNode.inputs[name]) || node.inputs?.some(input => input.name === name && input.link != null)) continue;
|
||||
const widget = node.widgets?.find(widget => widget.name === name);
|
||||
if (widget?.inputEl) outputNode.inputs[name] = widget.inputEl.value;
|
||||
}
|
||||
|
||||
// Only sync widget DOM value for non-connected (widget-based) text inputs
|
||||
const textWidget = node.widgets?.find(w => w.name === "text");
|
||||
if (textWidget && textWidget.inputEl) {
|
||||
const currentValue = textWidget.inputEl.value;
|
||||
|
||||
if (outputNode.inputs && outputNode.inputs.text !== currentValue) {
|
||||
outputNode.inputs.text = currentValue;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -292,8 +284,8 @@ app.registerExtension({
|
||||
},
|
||||
|
||||
async beforeRegisterNodeDef(nodeType, nodeData, app) {
|
||||
if (nodeData.name === 'PromptManagerFocused' || nodeData.name === 'PromptManagerSplit') {
|
||||
const install = nodeData.name === 'PromptManagerSplit' ? installSplitNode : installFocusedNode;
|
||||
if (nodeData.name === 'PromptManagerFocused' || nodeData.name === 'PromptManagerSplit' || nodeData.name === 'PromptManagerPair') {
|
||||
const install = nodeData.name === 'PromptManagerPair' ? installPairNode : nodeData.name === 'PromptManagerSplit' ? installSplitNode : installFocusedNode;
|
||||
install(nodeType, app, api, parseSearchQuery);
|
||||
return;
|
||||
}
|
||||
|
||||
@@ -2353,7 +2353,7 @@
|
||||
if (this.scanState.status === ScanProgress.RUNNING) {
|
||||
// A scan we were following ended while we were disconnected.
|
||||
const last = status.last_event;
|
||||
const terminal = last && (last.type === "complete" || last.type === "error");
|
||||
const terminal = last && (last.type === "complete" || last.type === "cancelled" || last.type === "error");
|
||||
this.applyScanEvent(
|
||||
terminal ? last : { type: "error", message: "the server no longer reports a running scan" }
|
||||
);
|
||||
@@ -2418,13 +2418,8 @@
|
||||
}
|
||||
|
||||
completeScan(state) {
|
||||
const { processed, found, added, linked } = state;
|
||||
document.getElementById("scanPill").classList.add("hidden");
|
||||
const linkedText = linked > 0 ? `, linked ${linked} images to existing prompts` : '';
|
||||
this.showNotification(
|
||||
`Scan completed! Processed ${processed} files, found ${found} prompts, added ${added} new prompts to database${linkedText}.`,
|
||||
"success"
|
||||
);
|
||||
this.showNotification(ScanProgress.completionLabel(state), "success");
|
||||
|
||||
// Auto-close after a short delay so the completion message is visible
|
||||
setTimeout(() => {
|
||||
|
||||
@@ -63,7 +63,7 @@
|
||||
const widgets = new Map();
|
||||
for (const n of workflowNodes(workflow)) {
|
||||
const text = widgetsOf(n)[0];
|
||||
if (/promptmanager/i.test(n.type || "") && typeof text === "string") {
|
||||
if (n.type !== "PromptManagerPair" && /promptmanager/i.test(n.type || "") && typeof text === "string") {
|
||||
widgets.set(String(n.id), text);
|
||||
}
|
||||
}
|
||||
@@ -94,8 +94,9 @@
|
||||
if (!isLink(link) || seen.size >= MAX_VISITS) return [];
|
||||
const id = String(link[0]);
|
||||
const node = ctx.graph[id];
|
||||
if (!node || !node.inputs || seen.has(id)) return [];
|
||||
seen.add(id);
|
||||
const key = `${id}:${link[1]}`;
|
||||
if (!node || !node.inputs || seen.has(key)) return [];
|
||||
seen.add(key);
|
||||
return [id, node];
|
||||
}
|
||||
|
||||
@@ -132,6 +133,10 @@
|
||||
if (typeof value === "string") return value;
|
||||
const [id, node] = enter(ctx, value, seen);
|
||||
if (!node) return "";
|
||||
if (node.class_type === "PromptManagerPair") {
|
||||
if (value[1] === 3) return resolveString(ctx, node.inputs.negative_text, seen);
|
||||
return value[1] === 2 ? composeNodeText(ctx, id, node, seen) : "";
|
||||
}
|
||||
if ("text" in node.inputs) return composeNodeText(ctx, id, node, seen);
|
||||
const delimiter = typeof node.inputs.delimiter === "string" ? node.inputs.delimiter : " ";
|
||||
return Object.entries(node.inputs)
|
||||
@@ -165,6 +170,10 @@
|
||||
function traceConditioning(ctx, link, role, seen = new Set()) {
|
||||
const [id, node] = enter(ctx, link, seen);
|
||||
if (!node) return "";
|
||||
if (node.class_type === "PromptManagerPair") {
|
||||
if (link[1] === 1) return resolveString(ctx, node.inputs.negative_text, seen);
|
||||
return link[1] === 0 ? composeNodeText(ctx, id, node, seen) : "";
|
||||
}
|
||||
if ("text" in node.inputs) return composeNodeText(ctx, id, node, seen);
|
||||
// Pass-through nodes (ControlNet apply, conditioning combine/set area, ...)
|
||||
const { inputs } = node;
|
||||
@@ -237,6 +246,7 @@
|
||||
PromptManagerText: PROMPT_MANAGER_WIDGETS,
|
||||
PromptManagerFocused: PROMPT_MANAGER_WIDGETS,
|
||||
PromptManagerSplit: PROMPT_MANAGER_WIDGETS,
|
||||
PromptManagerPair: ["text", "negative_text", "category", "tags", "search_text", "prepend_text", "append_text"],
|
||||
CheckpointLoader: ["config_name", "ckpt_name"],
|
||||
UNETLoader: ["unet_name", "weight_dtype"],
|
||||
};
|
||||
|
||||
@@ -11,13 +11,18 @@ function loadStyles() {
|
||||
const link = document.createElement('link'); Object.assign(link, {rel:'stylesheet', href}); document.head.append(link);
|
||||
}
|
||||
|
||||
export function installPairNode(nodeType, app, api, parseQuery) {
|
||||
installFocusedNode(nodeType, app, api, parseQuery, {layout:'pair'});
|
||||
}
|
||||
|
||||
export function installSplitNode(nodeType, app, api, parseQuery) {
|
||||
installFocusedNode(nodeType, app, api, parseQuery, {layout:'split'});
|
||||
}
|
||||
|
||||
export function installFocusedNode(nodeType, app, api, parseQuery, {layout = 'focused'} = {}) {
|
||||
const split = layout === 'split';
|
||||
const uiKey = split ? 'splitUI' : 'focusedUI';
|
||||
const pair = layout === 'pair';
|
||||
const split = layout !== 'focused';
|
||||
const uiKey = pair ? 'pairUI' : split ? 'splitUI' : 'focusedUI';
|
||||
const prototype = nodeType.prototype;
|
||||
const created = prototype.onNodeCreated;
|
||||
const removed = prototype.onRemoved;
|
||||
@@ -30,14 +35,14 @@ export function installFocusedNode(nodeType, app, api, parseQuery, {layout = 'fo
|
||||
const editorId = `pm-focused-${++nextEditorId}`;
|
||||
this.properties ||= {};
|
||||
const widgets = new Map();
|
||||
for (const name of names) {
|
||||
for (const name of pair ? [...names,'negative_text'] : names) {
|
||||
let widget = this.widgets?.find(item => item.name === name);
|
||||
if (!widget) widget = this.addWidget('text', name, '', () => {}, {});
|
||||
widget.hidden = true; widget.type = 'hidden'; widget.computeSize = () => [0, -4];
|
||||
widgets.set(name, widget);
|
||||
}
|
||||
const element = document.createElement('div'); element.className = `ca-tools pm-focused${split ? ' pm-split' : ''}`;
|
||||
element.setAttribute('aria-label',split ? 'Split prompt editor' : 'Focused prompt editor');
|
||||
const element = document.createElement('div'); element.className = `ca-tools pm-focused${split ? ' pm-split' : ''}${pair ? ' pm-pair' : ''}`;
|
||||
element.setAttribute('aria-label',pair ? 'Paired prompt editor' : split ? 'Split prompt editor' : 'Focused prompt editor');
|
||||
loadStyles(); mountTheme(element, app.ui?.settings, scope);
|
||||
const toolbar = document.createElement('div'); toolbar.className = 'pm-editor-row pm-editor-toolbar';
|
||||
const title = document.createElement('strong'); title.textContent = 'Prompt Manager';
|
||||
@@ -45,6 +50,9 @@ export function installFocusedNode(nodeType, app, api, parseQuery, {layout = 'fo
|
||||
const web = createButton('Web UI ↗'); toolbar.append(title);
|
||||
if (!split) toolbar.append(browse);
|
||||
toolbar.append(web);
|
||||
const save = pair ? createButton('Save pair','ca-primary') : null;
|
||||
if (save) toolbar.append(save);
|
||||
let saving = false;
|
||||
const workspace = document.createElement('div'); workspace.className = 'pm-editor-workspace';
|
||||
const editor = document.createElement('div'); editor.className = 'pm-editor-compose';
|
||||
workspace.append(editor);
|
||||
@@ -55,24 +63,40 @@ export function installFocusedNode(nodeType, app, api, parseQuery, {layout = 'fo
|
||||
const mode = data?.settings?.webui_display_mode;
|
||||
if (!scope.signal.aborted && data?.success && ['popup','newtab'].includes(mode)) webDisplayMode = mode;
|
||||
}).catch(() => {});
|
||||
const textWidget = widgets.get('text');
|
||||
const nativeInput = textWidget.inputEl;
|
||||
const inputDescriptor = Object.getOwnPropertyDescriptor(textWidget,'inputEl');
|
||||
const main = document.createElement('textarea'); main.rows = 6; main.spellcheck = nativeInput?.spellcheck ?? false;
|
||||
// ComfyUI owns widget.element; inputEl exposes our visible editor to extensions.
|
||||
Object.defineProperty(textWidget,'inputEl',{configurable:true,writable:true,value:main});
|
||||
scope.cleanup(() => {
|
||||
if (inputDescriptor) Object.defineProperty(textWidget,'inputEl',inputDescriptor);
|
||||
else delete textWidget.inputEl;
|
||||
});
|
||||
const mainField = createField('Main prompt', main, `${editorId}-text`);
|
||||
mainField.classList.add('pm-editor-main');
|
||||
const connected = document.createElement('p'); connected.className = 'pm-editor-status'; connected.textContent = 'Text is supplied by an upstream connection.';
|
||||
const controls = new Map([['text', main]]);
|
||||
function textControl(name, label) {
|
||||
const widget = widgets.get(name);
|
||||
const nativeInput = widget.inputEl;
|
||||
const descriptor = Object.getOwnPropertyDescriptor(widget,'inputEl');
|
||||
const input = document.createElement('textarea'); input.rows = 6;
|
||||
input.spellcheck = nativeInput?.spellcheck ?? false;
|
||||
// The host retains widget.element; extensions receive the visible editor.
|
||||
Object.defineProperty(widget,'inputEl',{configurable:true,writable:true,value:input});
|
||||
scope.cleanup(() => {
|
||||
if (descriptor) Object.defineProperty(widget,'inputEl',descriptor);
|
||||
else delete widget.inputEl;
|
||||
});
|
||||
const field = createField(label,input,`${editorId}-${name}`);
|
||||
field.classList.add('pm-editor-main');
|
||||
const note = document.createElement('p'); note.className = 'pm-editor-status';
|
||||
note.textContent = `${pair ? label : 'Text'} is supplied by an upstream connection.`;
|
||||
return {input,field,note};
|
||||
}
|
||||
const positive = textControl('text',pair ? 'Positive prompt' : 'Main prompt');
|
||||
const main = positive.input, mainField = positive.field, connected = positive.note;
|
||||
const negative = pair ? textControl('negative_text','Negative prompt') : null;
|
||||
const controls = new Map([['text',main]]);
|
||||
if (negative) controls.set('negative_text',negative.input);
|
||||
const details = new Map();
|
||||
const completions = new Map();
|
||||
element.append(toolbar, workspace);
|
||||
editor.append(mainField, connected);
|
||||
if (pair) {
|
||||
const fields = document.createElement('div'); fields.className = 'pm-pair-texts';
|
||||
for (const control of [positive,negative]) {
|
||||
const slot = document.createElement('div'); slot.className = 'pm-pair-text-slot';
|
||||
slot.append(control.field,control.note); fields.append(slot);
|
||||
}
|
||||
editor.append(fields);
|
||||
} else editor.append(mainField, connected);
|
||||
function disclosure(key, label) {
|
||||
const section = document.createElement('details'); section.className = 'pm-editor-details';
|
||||
const summary = document.createElement('summary'); summary.textContent = label; section.append(summary); editor.append(section); details.set(key,section);
|
||||
@@ -100,7 +124,14 @@ export function installFocusedNode(nodeType, app, api, parseQuery, {layout = 'fo
|
||||
}
|
||||
const composed = disclosure('preview','Composed text preview');
|
||||
const preview = document.createElement('p'); preview.className = 'pm-editor-preview'; preview.dataset.composedPreview = ''; composed.append(preview);
|
||||
const footer = document.createElement('p'); footer.className = 'pm-editor-status'; footer.setAttribute('role','status'); footer.textContent = 'Main prompt is saved when this node runs.';
|
||||
const footer = document.createElement('p'); footer.className = 'pm-editor-status'; footer.setAttribute('role','status'); footer.textContent = pair ? 'The prompt pair is saved when this node runs, or use Save pair.' : 'Main prompt is saved when this node runs.';
|
||||
const saveConnectionNotice = pair ? document.createElement('p') : null;
|
||||
if (saveConnectionNotice) {
|
||||
saveConnectionNotice.className = 'pm-editor-status'; saveConnectionNotice.setAttribute('role','status');
|
||||
saveConnectionNotice.textContent = 'Run the node to save connected values.'; element.append(saveConnectionNotice);
|
||||
}
|
||||
function hasConnectedPairValues() { return ['text','negative_text','category','tags'].some(isConnected); }
|
||||
function pairPayload() { return {positive_text:value('text'),negative_text:value('negative_text'),category:value('category').trim(),tags:value('tags').split(',').map(tag => tag.trim()).filter(Boolean)}; }
|
||||
function isConnected(name) { return thisNode.inputs?.some(input => input.name === name && input.link != null); }
|
||||
function value(name) { return String(widgets.get(name)?.value ?? ''); }
|
||||
function write(name, text) {
|
||||
@@ -114,9 +145,11 @@ export function installFocusedNode(nodeType, app, api, parseQuery, {layout = 'fo
|
||||
}
|
||||
function render() {
|
||||
mainField.hidden = Boolean(isConnected('text')); connected.hidden = !isConnected('text');
|
||||
if (negative) { negative.field.hidden = Boolean(isConnected('negative_text')); negative.note.hidden = !isConnected('negative_text'); }
|
||||
if (save) { save.disabled = saving || hasConnectedPairValues(); saveConnectionNotice.hidden = !hasConnectedPairValues(); }
|
||||
controls.forEach((input,name) => {
|
||||
input.value = value(name);
|
||||
if (name !== 'text') {
|
||||
if (name !== 'text' && name !== 'negative_text') {
|
||||
if (isConnected(name)) completions.get(name)?.close();
|
||||
input.hidden = Boolean(isConnected(name));
|
||||
input._connectionNote.hidden = !isConnected(name);
|
||||
@@ -125,14 +158,15 @@ export function installFocusedNode(nodeType, app, api, parseQuery, {layout = 'fo
|
||||
preview.textContent = ['prepend_text','text','append_text'].map(name => isConnected(name) ? `[${name} supplied by connection]` : name === 'text' ? value(name) : value(name).trim()).filter(Boolean).join(' ');
|
||||
library.refreshAvailability();
|
||||
}
|
||||
const library = createPromptLibrary({scope,api,settings:app.ui?.settings,parseQuery,isConnected:() => isConnected('text'), onQuery:text => write('search_text',text), onFilters:() => {this.properties[uiKey] = snapshot();}, onUse:(prompt,includeMetadata) => {
|
||||
if (isConnected('text')) return;
|
||||
write('text',prompt.text);
|
||||
const library = createPromptLibrary({scope,api,mode:pair ? 'pairs' : 'prompts',settings:app.ui?.settings,parseQuery,isConnected:() => isConnected('text') || (pair && isConnected('negative_text')), onQuery:text => write('search_text',text), onFilters:() => {this.properties[uiKey] = snapshot();}, onUse:(prompt,includeMetadata) => {
|
||||
if (isConnected('text') || (pair && isConnected('negative_text'))) return;
|
||||
write('text',pair ? prompt.positive_text : prompt.text);
|
||||
if (pair) write('negative_text',prompt.negative_text || '');
|
||||
if (includeMetadata) {
|
||||
write('category',prompt.category || '');
|
||||
write('tags',Array.isArray(prompt.tags) ? prompt.tags.join(', ') : prompt.tags || '');
|
||||
}
|
||||
render(); footer.setAttribute('role','status'); footer.textContent = 'Loaded prompt. Saved when this node runs.'; main.focus();
|
||||
render(); footer.setAttribute('role','status'); footer.textContent = pair ? 'Loaded pair. Saved when this node runs.' : 'Loaded prompt. Saved when this node runs.'; main.focus();
|
||||
}});
|
||||
workspace.append(library.element); element.append(footer);
|
||||
function snapshot() {
|
||||
@@ -144,13 +178,36 @@ export function installFocusedNode(nodeType, app, api, parseQuery, {layout = 'fo
|
||||
library.element.hidden = !split && !state.libraryOpen; browse.textContent = state.libraryOpen ? 'Hide library' : 'Browse library'; browse.setAttribute('aria-expanded', String(Boolean(state.libraryOpen)));
|
||||
library.setQuery(value('search_text')); library.restoreFilters(state.libraryFilters); render();
|
||||
}
|
||||
controls.forEach((input,name) => scope.listen(input,'input',() => {write(name,input.value);render();footer.setAttribute('role','status');footer.textContent='Main prompt is saved when this node runs.';}));
|
||||
controls.forEach((input,name) => scope.listen(input,'input',() => {write(name,input.value);render();footer.setAttribute('role','status');footer.textContent=pair?'The pair is saved when this node runs.':'Main prompt is saved when this node runs.';}));
|
||||
scope.listen(browse,'click',() => {
|
||||
library.element.hidden = !library.element.hidden;
|
||||
library.closeCompletions();
|
||||
browse.textContent = library.element.hidden ? 'Browse library' : 'Hide library'; browse.setAttribute('aria-expanded',String(!library.element.hidden));
|
||||
this.properties[uiKey] = snapshot(); if (!library.element.hidden) library.focus();
|
||||
});
|
||||
if (save) scope.listen(save,'click',async () => {
|
||||
if (saving || hasConnectedPairValues()) return;
|
||||
if (!value('text').trim()) {
|
||||
footer.setAttribute('role','alert'); footer.textContent = 'Enter a positive prompt before saving the pair.';
|
||||
main.focus(); return;
|
||||
}
|
||||
const savedDraft = JSON.stringify(pairPayload());
|
||||
saving = true; render(); footer.setAttribute('role','status'); footer.textContent = 'Saving pair…';
|
||||
try {
|
||||
const response = await fetch(api.fileURL('/prompt_manager/pairs'), {
|
||||
method:'POST',headers:{'Content-Type':'application/json'},signal:scope.signal,
|
||||
body:savedDraft,
|
||||
});
|
||||
const data = await response.json();
|
||||
if (scope.signal.aborted) return;
|
||||
if (!response.ok || !data.success) throw new Error(data.error || `HTTP ${response.status}`);
|
||||
footer.setAttribute('role','status'); footer.textContent = !hasConnectedPairValues() && JSON.stringify(pairPayload()) === savedDraft ? 'Pair saved.' : 'Earlier draft saved. Current changes are not saved yet.';
|
||||
} catch (error) {
|
||||
if (!scope.signal.aborted) { footer.setAttribute('role','alert'); footer.textContent = `Could not save pair: ${error.message}`; }
|
||||
} finally {
|
||||
saving = false; if (!scope.signal.aborted) render();
|
||||
}
|
||||
});
|
||||
scope.listen(web,'click',() => {
|
||||
const url = api.fileURL('/prompt_manager/web');
|
||||
const opened = webDisplayMode === 'newtab' ? window.open(url,'_blank') : window.open(url,'_blank','width=1200,height=800,scrollbars=yes,resizable=yes');
|
||||
@@ -159,7 +216,7 @@ export function installFocusedNode(nodeType, app, api, parseQuery, {layout = 'fo
|
||||
});
|
||||
this._focusedView = {element,render,restore,snapshot};
|
||||
this.addDOMWidget('prompt_manager_focused_ui','div',element,{serialize:false, hideOnZoom:false, getMinHeight:() => 360});
|
||||
this.size = [split ? 900 : 550,650]; this.setSize(this.size); restore();
|
||||
this.size = [pair ? 1100 : split ? 900 : 550,pair ? 700 : 650]; this.setSize(this.size); restore();
|
||||
};
|
||||
prototype.configure = function (...args) { const result = configure?.apply(this,args); this._focusedView?.restore(); return result; };
|
||||
prototype.serialize = function (...args) {
|
||||
|
||||
@@ -1221,8 +1221,9 @@
|
||||
if (databasePrompt) {
|
||||
// Create metadata object from database prompt
|
||||
metadata = {
|
||||
positivePrompt: databasePrompt.text,
|
||||
negativePrompt: 'No negative prompt found', // Database doesn't separate positive/negative
|
||||
textSource: databasePrompt.pair_context?.text_source,
|
||||
positivePrompt: databasePrompt.pair_context?.positive_text ?? databasePrompt.text,
|
||||
negativePrompt: databasePrompt.pair_context?.negative_text ?? databasePrompt.pair?.negative_text ?? 'No negative prompt found',
|
||||
checkpoint: 'Unknown',
|
||||
steps: 'Unknown',
|
||||
cfgScale: 'Unknown',
|
||||
@@ -1248,7 +1249,7 @@
|
||||
metadata.sampler = pngMetadata.sampler;
|
||||
metadata.seed = pngMetadata.seed;
|
||||
// Use negative prompt from PNG if available
|
||||
if (pngMetadata.negativePrompt && pngMetadata.negativePrompt !== 'No negative prompt found') {
|
||||
if (!databasePrompt.pair && pngMetadata.negativePrompt && pngMetadata.negativePrompt !== 'No negative prompt found') {
|
||||
metadata.negativePrompt = pngMetadata.negativePrompt;
|
||||
}
|
||||
}
|
||||
@@ -1289,6 +1290,12 @@
|
||||
return;
|
||||
}
|
||||
|
||||
const textSourceLabel = new Map([
|
||||
['execution','Recorded during execution'],
|
||||
['graph-inferred','Recovered from workflow metadata'],
|
||||
['embedded-parameters','Read from embedded parameters'],
|
||||
['mixed','Execution and workflow metadata'],
|
||||
]).get(metadata.textSource);
|
||||
const html = `
|
||||
<!-- File Path -->
|
||||
<div class="space-y-2">
|
||||
@@ -1309,6 +1316,8 @@
|
||||
</div>
|
||||
</div>
|
||||
|
||||
${textSourceLabel ? `<p class="text-xs text-pm-secondary">${escapeHtml(textSourceLabel)}</p>` : ''}
|
||||
|
||||
<!-- Prompt Section -->
|
||||
<div class="space-y-3">
|
||||
<div class="flex items-center justify-between">
|
||||
|
||||
@@ -38,3 +38,11 @@
|
||||
.pm-split .pm-editor-library { border-block-start: 0; border-inline-start: 1px solid var(--ca-line); padding-block-start: 0; padding-inline-start: 12px; }
|
||||
.pm-split .pm-editor-fields { grid-template-columns: minmax(0,1fr); }
|
||||
}
|
||||
|
||||
.pm-pair-texts { display: grid; grid-template-columns: repeat(2,minmax(0,1fr)); gap: 12px; flex: 1; min-block-size: 180px; min-inline-size: 0; }
|
||||
.pm-pair-text-slot { display: flex; flex-direction: column; min-inline-size: 0; min-block-size: 0; }
|
||||
.pm-pair-text-slot > .pm-editor-main { display: flex; flex-direction: column; align-items: stretch; gap: 5px; flex: 1; min-block-size: 180px; }
|
||||
.pm-pair .pm-editor-toolbar { flex-shrink: 0; }
|
||||
@container pm-editor (inline-size <= 720px) {
|
||||
.pm-pair-texts { grid-template-columns: minmax(0,1fr); grid-template-rows: repeat(2,minmax(180px,1fr)); flex: none; }
|
||||
}
|
||||
|
||||
@@ -2,7 +2,9 @@ import {createScope, createButton, createSwitch, createField} from './shared/too
|
||||
import {attachMetadataAutocomplete} from './metadata-autocomplete.js';
|
||||
let nextLibraryId = 0;
|
||||
|
||||
export function createPromptLibrary({scope, api, settings, parseQuery, onUse, isConnected, onQuery, onFilters}) {
|
||||
export function createPromptLibrary({scope, api, settings, parseQuery, onUse, isConnected, onQuery, onFilters, mode = 'prompts'}) {
|
||||
const pairs = mode === 'pairs';
|
||||
const noun = pairs ? 'pair' : 'prompt';
|
||||
const element = document.createElement('section');
|
||||
element.className = 'pm-editor-library';
|
||||
element.setAttribute('aria-label', 'Prompt library');
|
||||
@@ -18,18 +20,19 @@ export function createPromptLibrary({scope, api, settings, parseQuery, onUse, is
|
||||
const category = document.createElement('input'); category.type = 'text';
|
||||
const tags = document.createElement('input'); tags.type = 'text';
|
||||
filters.append(createField('Filter category',category,`${libraryId}-category`),createField('Filter tags',tags,`${libraryId}-tags`));
|
||||
const hint = document.createElement('p'); hint.className = 'pm-editor-status'; hint.textContent = 'Search combines text, category and tags. Recent shows all recent prompts.';
|
||||
const hint = document.createElement('p'); hint.className = 'pm-editor-status'; hint.textContent = `Search combines text, category and tags. Recent shows all recent ${noun}s.`;
|
||||
const completions = [category,tags].map((input,index) => attachMetadataAutocomplete(input, {
|
||||
scope,api,settings,multiple:index === 1,available:() => !element.hidden && input.getClientRects().length > 0,
|
||||
}));
|
||||
const status = document.createElement('p'); status.className = 'pm-editor-status'; status.setAttribute('role', 'status'); status.textContent = 'Search or browse recent prompts.';
|
||||
const status = document.createElement('p'); status.className = 'pm-editor-status'; status.setAttribute('role', 'status'); status.textContent = `Search or browse recent ${noun}s.`;
|
||||
const results = document.createElement('div'); results.className = 'pm-editor-results';
|
||||
const preview = document.createElement('section'); preview.className = 'pm-editor-selection'; preview.hidden = true;
|
||||
const previewText = document.createElement('p'); previewText.className = 'pm-editor-preview';
|
||||
const previewNegative = document.createElement('p'); previewNegative.className = 'pm-editor-preview';
|
||||
const actions = document.createElement('div'); actions.className = 'pm-editor-row';
|
||||
const includeMetadata = createSwitch('Load metadata', '', `${libraryId}-load-metadata`);
|
||||
const use = createButton('Use prompt', 'ca-primary'); use.disabled = true;
|
||||
actions.append(includeMetadata, use); preview.append(previewText, actions);
|
||||
const use = createButton(pairs ? 'Use pair' : 'Use prompt', 'ca-primary'); use.disabled = true;
|
||||
actions.append(includeMetadata, use); preview.append(previewText); if (pairs) preview.append(previewNegative); preview.append(actions);
|
||||
element.append(heading, queryRow, filters, hint, status, results, preview);
|
||||
let selected = null;
|
||||
let request = null;
|
||||
@@ -56,26 +59,33 @@ export function createPromptLibrary({scope, api, settings, parseQuery, onUse, is
|
||||
const params = new URLSearchParams(recentOnly ? {limit:'20', sort:'last_used_desc'} : {limit:'50'});
|
||||
if (!recentOnly) {
|
||||
if (text) params.set('text', text);
|
||||
if (tagFilters.length) params.set('tags', tagFilters.join(','));
|
||||
if (tagFilters.length) { if (pairs) tagFilters.forEach(tag => params.append('tag',tag)); else params.set('tags', tagFilters.join(',')); }
|
||||
if (categoryFilter) params.set('category', categoryFilter);
|
||||
}
|
||||
status.textContent = 'Loading prompts…';
|
||||
status.textContent = `Loading ${noun}s…`;
|
||||
try {
|
||||
const response = await fetch(api.fileURL(`/prompt_manager/${recentOnly ? 'recent' : 'search'}?${params}`), {signal:request.signal});
|
||||
const response = await fetch(api.fileURL(`/prompt_manager/${pairs ? 'pairs' : recentOnly ? 'recent' : 'search'}?${params}`), {signal:request.signal});
|
||||
if (!response.ok) throw new Error(`HTTP ${response.status}: ${response.statusText}`);
|
||||
const data = await response.json();
|
||||
if (scope.signal.aborted || current !== sequence) return;
|
||||
if (!data.success) throw new Error(data.error || 'Could not load prompts');
|
||||
const found = data.results || [];
|
||||
status.textContent = found.length ? `${found.length} ${found.length === 1 ? 'prompt' : 'prompts'}` : 'No matching prompts.';
|
||||
if (!data.success) throw new Error(data.error || `Could not load ${noun}s`);
|
||||
const found = (pairs ? data.pairs : data.results) || [];
|
||||
status.textContent = found.length ? `${found.length} ${noun}${found.length === 1 ? '' : 's'}` : `No matching ${noun}s.`;
|
||||
found.forEach(item => {
|
||||
const button = createButton('', 'pm-editor-result'); button.setAttribute('aria-pressed','false');
|
||||
const text = document.createElement('span'); text.textContent = item.text;
|
||||
const text = document.createElement('span'); text.textContent = pairs ? item.title || item.positive_text : item.text;
|
||||
const metadata = document.createElement('span'); metadata.className = 'pm-editor-result-meta';
|
||||
metadata.textContent = [item.category, Array.isArray(item.tags) ? item.tags.join(', ') : item.tags, item.rating ? `${item.rating}/5` : ''].filter(Boolean).join(' · ');
|
||||
button.append(text, metadata);
|
||||
button.append(text);
|
||||
if (pairs) {
|
||||
const negative = document.createElement('span'); negative.className = 'pm-editor-result-meta';
|
||||
const summary = item.negative_text || '(empty)';
|
||||
negative.textContent = `Negative: ${summary.length > 160 ? `${summary.slice(0,160)}…` : summary}`;
|
||||
negative.title = summary; button.append(negative);
|
||||
}
|
||||
button.append(metadata);
|
||||
rows.listen(button, 'click', () => {
|
||||
selected = item; previewText.textContent = item.text; preview.hidden = false;
|
||||
selected = item; previewText.textContent = pairs ? `Positive: ${item.positive_text}` : item.text; if (pairs) previewNegative.textContent = `Negative: ${item.negative_text || '(empty)'}`; preview.hidden = false;
|
||||
results.querySelectorAll('button').forEach(row => row.setAttribute('aria-pressed', String(row === button)));
|
||||
status.textContent = 'Previewing selection. Editor unchanged.'; refreshAvailability();
|
||||
});
|
||||
@@ -83,7 +93,7 @@ export function createPromptLibrary({scope, api, settings, parseQuery, onUse, is
|
||||
});
|
||||
} catch (error) {
|
||||
if (scope.signal.aborted || current !== sequence || request.signal.aborted) return;
|
||||
status.classList.add('pm-editor-error'); status.setAttribute('role','alert'); status.textContent = `Could not load prompts: ${error.message}`;
|
||||
status.classList.add('pm-editor-error'); status.setAttribute('role','alert'); status.textContent = `Could not load ${noun}s: ${error.message}`;
|
||||
}
|
||||
}
|
||||
scope.listen(search, 'click', () => load());
|
||||
|
||||
@@ -55,12 +55,18 @@
|
||||
found: num(event.found, state.found),
|
||||
};
|
||||
}
|
||||
if (event.type === "complete") {
|
||||
if (event.type === "complete" || event.type === "cancelled") {
|
||||
const counts = Object.fromEntries(
|
||||
['pairs_added','pair_links_added','enriched','unchanged','ambiguous','unresolved','conflict','failed']
|
||||
.map(key => [key, num(event[key], state[key] || 0)])
|
||||
);
|
||||
return {
|
||||
...state,
|
||||
...counts,
|
||||
status: DONE,
|
||||
progress: 100,
|
||||
statusText: "Scan completed!",
|
||||
cancelled: event.type === "cancelled",
|
||||
progress: event.type === "cancelled" ? state.progress : 100,
|
||||
statusText: event.type === "cancelled" ? "Scan cancelled." : "Scan completed!",
|
||||
processed: num(event.processed, state.processed),
|
||||
found: num(event.found, state.found),
|
||||
added: num(event.added, state.added),
|
||||
@@ -93,8 +99,14 @@
|
||||
return `Scanning ${pct}% · ${state.processed} files · ${state.found} prompts`;
|
||||
}
|
||||
|
||||
function completionLabel(state) {
|
||||
const review = (state.ambiguous || 0) + (state.unresolved || 0) + (state.conflict || 0);
|
||||
return `${state.cancelled ? 'Scan cancelled' : 'Scan completed'}! Processed ${state.processed} files, found ${state.found} prompts; ${state.added} new image records, ${state.linked} links, ${state.enriched || 0} enriched, ${state.unchanged || 0} unchanged; ${state.pairs_added || 0} new pairs, ${state.pair_links_added || 0} pair links; ${review} need review, ${state.failed || 0} failed.`;
|
||||
}
|
||||
|
||||
const api = {
|
||||
initialState,
|
||||
completionLabel,
|
||||
reduceScanEvent,
|
||||
fromStatus,
|
||||
dismissAction,
|
||||
|
||||