diff --git a/__init__.py b/__init__.py index f05e9736..c1e1dd52 100644 --- a/__init__.py +++ b/__init__.py @@ -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", } diff --git a/database/models.py b/database/models.py index 84e7392e..37a6dcd6 100644 --- a/database/models.py +++ b/database/models.py @@ -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 - ``.backup_`` (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 + ``.backup_`` (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)) + ) diff --git a/database/operations.py b/database/operations.py index c501d5bd..732d49bb 100644 --- a/database/operations.py +++ b/database/operations.py @@ -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 diff --git a/database/pair_ingestion.py b/database/pair_ingestion.py new file mode 100644 index 00000000..cad3641a --- /dev/null +++ b/database/pair_ingestion.py @@ -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, + } diff --git a/database/pairs.py b/database/pairs.py new file mode 100644 index 00000000..c250b026 --- /dev/null +++ b/database/pairs.py @@ -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 + ] diff --git a/docs/concepts/prompt-variant-timeline.html b/docs/concepts/prompt-variant-timeline.html new file mode 100644 index 00000000..e177be08 --- /dev/null +++ b/docs/concepts/prompt-variant-timeline.html @@ -0,0 +1,116 @@ + + + + + + +Prompt history · ComfyAssets prototype + + + + + +
ComfyAssets/ Prompt Manager / AdminInteractive concept
+
+

Prompt history

One positive prompt. Every negative experiment.

Sample data and illustrative previews.
Everything here stays in this browser tab.

+
+
Positive prompt · unchanged

A remote alpine observatory at blue hour, warm light through glass, dramatic mountain peaks, cinematic composition, atmospheric mist, fine detail

Landscapes
+

Negative variants

Click a red dot to inspect a negative. Small dots are repeat runs. Select a variant for its full run history.
+
+ Positive− Negative variantRunSaved, not run
+

Shared gallery

1 positive · 1 gallery
+
+
Select a negative or a run to explore its images.
+
Concept only · No ComfyUI or database connectionTime flows left to right · 3 images per sample run
+
+

Add a negative variant

The positive stays unchanged. A saved variant starts with a hollow dot until its first run.

Leave empty to explicitly use no negative prompt. Existing text is reused.

 +

Generated image

 + + + diff --git a/docs/reusable-prompt-pairs.md b/docs/reusable-prompt-pairs.md new file mode 100644 index 00000000..98003b9a --- /dev/null +++ b/docs/reusable-prompt-pairs.md @@ -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. diff --git a/docs/screenshots/paired-toolkit/classic-pair-dark.png b/docs/screenshots/paired-toolkit/classic-pair-dark.png new file mode 100644 index 00000000..00ffcf96 Binary files /dev/null and b/docs/screenshots/paired-toolkit/classic-pair-dark.png differ diff --git a/docs/screenshots/paired-toolkit/classic-pair-library-preview.png b/docs/screenshots/paired-toolkit/classic-pair-library-preview.png new file mode 100644 index 00000000..49fd994b Binary files /dev/null and b/docs/screenshots/paired-toolkit/classic-pair-library-preview.png differ diff --git a/docs/screenshots/paired-toolkit/classic-pair-light.png b/docs/screenshots/paired-toolkit/classic-pair-light.png new file mode 100644 index 00000000..ef75be6b Binary files /dev/null and b/docs/screenshots/paired-toolkit/classic-pair-light.png differ diff --git a/docs/screenshots/paired-toolkit/classic-pair-narrow.png b/docs/screenshots/paired-toolkit/classic-pair-narrow.png new file mode 100644 index 00000000..1c9fbc39 Binary files /dev/null and b/docs/screenshots/paired-toolkit/classic-pair-narrow.png differ diff --git a/docs/screenshots/paired-toolkit/nodes2-pair-dark.png b/docs/screenshots/paired-toolkit/nodes2-pair-dark.png new file mode 100644 index 00000000..e8aa68ed Binary files /dev/null and b/docs/screenshots/paired-toolkit/nodes2-pair-dark.png differ diff --git a/docs/screenshots/paired-toolkit/nodes2-pair-library-preview.png b/docs/screenshots/paired-toolkit/nodes2-pair-library-preview.png new file mode 100644 index 00000000..c5bf1044 Binary files /dev/null and b/docs/screenshots/paired-toolkit/nodes2-pair-library-preview.png differ diff --git a/docs/screenshots/paired-toolkit/nodes2-pair-light.png b/docs/screenshots/paired-toolkit/nodes2-pair-light.png new file mode 100644 index 00000000..03b08c2d Binary files /dev/null and b/docs/screenshots/paired-toolkit/nodes2-pair-light.png differ diff --git a/docs/screenshots/paired-toolkit/nodes2-pair-narrow.png b/docs/screenshots/paired-toolkit/nodes2-pair-narrow.png new file mode 100644 index 00000000..417af7c4 Binary files /dev/null and b/docs/screenshots/paired-toolkit/nodes2-pair-narrow.png differ diff --git a/prompt_manager_pair.py b/prompt_manager_pair.py new file mode 100644 index 00000000..75ae89b7 --- /dev/null +++ b/prompt_manager_pair.py @@ -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() diff --git a/py/api/__init__.py b/py/api/__init__.py index cf267936..f243b6b0 100644 --- a/py/api/__init__.py +++ b/py/api/__init__.py @@ -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) diff --git a/py/api/admin.py b/py/api/admin.py index d13c0341..206792e5 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -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 # '.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" diff --git a/py/api/pairs.py b/py/api/pairs.py new file mode 100644 index 00000000..775f0e8a --- /dev/null +++ b/py/api/pairs.py @@ -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"'}, + ) diff --git a/py/api/prompts.py b/py/api/prompts.py index 0f5df1fb..122d0735 100644 --- a/py/api/prompts.py +++ b/py/api/prompts.py @@ -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( diff --git a/py/api/scan_job.py b/py/api/scan_job.py index d3eb9ce6..e6304707 100644 --- a/py/api/scan_job.py +++ b/py/api/scan_job.py @@ -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() diff --git a/tests/browser/pair-node.spec.mjs b/tests/browser/pair-node.spec.mjs new file mode 100644 index 00000000..d1245741 --- /dev/null +++ b/tests/browser/pair-node.spec.mjs @@ -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:''}]}})); + 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: ');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); +}); diff --git a/tests/browser/server.py b/tests/browser/server.py index f5f011bd..2eee3505 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -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) diff --git a/tests/comfyui/pair-node.spec.mjs b/tests/comfyui/pair-node.spec.mjs new file mode 100644 index 00000000..7cb97559 --- /dev/null +++ b/tests/comfyui/pair-node.spec.mjs @@ -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); +}); diff --git a/tests/js/pair-metadata.test.js b/tests/js/pair-metadata.test.js new file mode 100644 index 00000000..0d0b91c3 --- /dev/null +++ b/tests/js/pair-metadata.test.js @@ -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); + } +}); diff --git a/tests/js/scan-progress.test.js b/tests/js/scan-progress.test.js index 6f0cfb92..837bc7a3 100644 --- a/tests/js/scan-progress.test.js +++ b/tests/js/scan-progress.test.js @@ -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'); +}); diff --git a/tests/test_api.py b/tests/test_api.py index f093393b..051881a8 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -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 diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index c4f242eb..1f910933 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -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)] diff --git a/tests/test_api_autotag.py b/tests/test_api_autotag.py index 8db42d90..832221dd 100644 --- a/tests/test_api_autotag.py +++ b/tests/test_api_autotag.py @@ -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) diff --git a/tests/test_api_backup_restore.py b/tests/test_api_backup_restore.py index fb473d98..eb55b47a 100644 --- a/tests/test_api_backup_restore.py +++ b/tests/test_api_backup_restore.py @@ -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 diff --git a/tests/test_api_core.py b/tests/test_api_core.py index 16d4a55d..91c4bce0 100644 --- a/tests/test_api_core.py +++ b/tests/test_api_core.py @@ -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) diff --git a/tests/test_api_database_isolation.py b/tests/test_api_database_isolation.py new file mode 100644 index 00000000..75593f49 --- /dev/null +++ b/tests/test_api_database_isolation.py @@ -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() diff --git a/tests/test_api_images.py b/tests/test_api_images.py index 2a820ade..c43ceddd 100644 --- a/tests/test_api_images.py +++ b/tests/test_api_images.py @@ -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) diff --git a/tests/test_api_logging.py b/tests/test_api_logging.py index cf15f96d..97c18c4c 100644 --- a/tests/test_api_logging.py +++ b/tests/test_api_logging.py @@ -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) diff --git a/tests/test_api_lora.py b/tests/test_api_lora.py index 1c340b12..36b99553 100644 --- a/tests/test_api_lora.py +++ b/tests/test_api_lora.py @@ -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 diff --git a/tests/test_api_pairs.py b/tests/test_api_pairs.py new file mode 100644 index 00000000..846e9964 --- /dev/null +++ b/tests/test_api_pairs.py @@ -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) diff --git a/tests/test_api_prompts.py b/tests/test_api_prompts.py index e2a6c834..18b1c21b 100644 --- a/tests/test_api_prompts.py +++ b/tests/test_api_prompts.py @@ -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", diff --git a/tests/test_gallery_roots.py b/tests/test_gallery_roots.py index 9d1dfef8..5ae871b4 100644 --- a/tests/test_gallery_roots.py +++ b/tests/test_gallery_roots.py @@ -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) diff --git a/tests/test_init_package.py b/tests/test_init_package.py index 53ec2bbf..39fc9bf3 100644 --- a/tests/test_init_package.py +++ b/tests/test_init_package.py @@ -138,6 +138,7 @@ class TestPackageWiring(PackageTestCase): [ "PromptManager", "PromptManagerFocused", + "PromptManagerPair", "PromptManagerSplit", "PromptManagerText", "PromptSearchList", diff --git a/tests/test_output_dir.py b/tests/test_output_dir.py index dd518b03..3d643323 100644 --- a/tests/test_output_dir.py +++ b/tests/test_output_dir.py @@ -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() diff --git a/tests/test_pair_dependency_concurrency.py b/tests/test_pair_dependency_concurrency.py new file mode 100644 index 00000000..7a95c851 --- /dev/null +++ b/tests/test_pair_dependency_concurrency.py @@ -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() diff --git a/tests/test_pair_execution.py b/tests/test_pair_execution.py new file mode 100644 index 00000000..5f6bea26 --- /dev/null +++ b/tests/test_pair_execution.py @@ -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") diff --git a/tests/test_pair_image_linking.py b/tests/test_pair_image_linking.py new file mode 100644 index 00000000..6c873098 --- /dev/null +++ b/tests/test_pair_image_linking.py @@ -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) diff --git a/tests/test_pair_ingestion.py b/tests/test_pair_ingestion.py new file mode 100644 index 00000000..4cfdfa9d --- /dev/null +++ b/tests/test_pair_ingestion.py @@ -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") diff --git a/tests/test_pair_output_roles.py b/tests/test_pair_output_roles.py new file mode 100644 index 00000000..1fa364dd --- /dev/null +++ b/tests/test_pair_output_roles.py @@ -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) diff --git a/tests/test_pair_restore_semantics.py b/tests/test_pair_restore_semantics.py new file mode 100644 index 00000000..ad2272c8 --- /dev/null +++ b/tests/test_pair_restore_semantics.py @@ -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() diff --git a/tests/test_pair_scan.py b/tests/test_pair_scan.py new file mode 100644 index 00000000..c8d055d6 --- /dev/null +++ b/tests/test_pair_scan.py @@ -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") diff --git a/tests/test_pair_schema_restore.py b/tests/test_pair_schema_restore.py new file mode 100644 index 00000000..16e919da --- /dev/null +++ b/tests/test_pair_schema_restore.py @@ -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() diff --git a/tests/test_pair_tag_maintenance.py b/tests/test_pair_tag_maintenance.py new file mode 100644 index 00000000..c56e968d --- /dev/null +++ b/tests/test_pair_tag_maintenance.py @@ -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() diff --git a/tests/test_pair_usage_tracking.py b/tests/test_pair_usage_tracking.py new file mode 100644 index 00000000..9039f29d --- /dev/null +++ b/tests/test_pair_usage_tracking.py @@ -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) diff --git a/tests/test_pair_vocabulary.py b/tests/test_pair_vocabulary.py new file mode 100644 index 00000000..147274ac --- /dev/null +++ b/tests/test_pair_vocabulary.py @@ -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() diff --git a/tests/test_prompt_pair_node.py b/tests/test_prompt_pair_node.py new file mode 100644 index 00000000..693a6c13 --- /dev/null +++ b/tests/test_prompt_pair_node.py @@ -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")) diff --git a/tests/test_prompt_pair_resolver.py b/tests/test_prompt_pair_resolver.py new file mode 100644 index 00000000..bb4c8cda --- /dev/null +++ b/tests/test_prompt_pair_resolver.py @@ -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() diff --git a/tests/test_prompt_pairs.py b/tests/test_prompt_pairs.py new file mode 100644 index 00000000..5fb880a6 --- /dev/null +++ b/tests/test_prompt_pairs.py @@ -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=[]) diff --git a/tools/preview-prompt-timeline.py b/tools/preview-prompt-timeline.py new file mode 100644 index 00000000..bb3b1ed7 --- /dev/null +++ b/tools/preview-prompt-timeline.py @@ -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() diff --git a/utils/image_monitor.py b/utils/image_monitor.py index 5a8167b8..a42504e9 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -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 ) diff --git a/utils/metadata_extractor.py b/utils/metadata_extractor.py index c6e85416..c1c044a8 100644 --- a/utils/metadata_extractor.py +++ b/utils/metadata_extractor.py @@ -116,6 +116,7 @@ class ComfyUIMetadataExtractor: # Extract other common metadata fields metadata_fields = [ + "comfyassets_prompt_pairs", "parameters", "model", "sampler", diff --git a/utils/pair_execution.py b/utils/pair_execution.py new file mode 100644 index 00000000..1dc04e3b --- /dev/null +++ b/utils/pair_execution.py @@ -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) diff --git a/utils/prompt_graph.py b/utils/prompt_graph.py index acc3b2e1..fc1ee362 100644 --- a/utils/prompt_graph.py +++ b/utils/prompt_graph.py @@ -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 diff --git a/utils/prompt_pair_resolver.py b/utils/prompt_pair_resolver.py new file mode 100644 index 00000000..4db5507f --- /dev/null +++ b/utils/prompt_pair_resolver.py @@ -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 diff --git a/utils/usage_tracking.py b/utils/usage_tracking.py index 7dcaadf2..2acbdf2d 100644 --- a/utils/usage_tracking.py +++ b/utils/usage_tracking.py @@ -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 diff --git a/web/comfy/prompt_manager.js b/web/comfy/prompt_manager.js index 59045c6f..e1f10efb 100644 --- a/web/comfy/prompt_manager.js +++ b/web/comfy/prompt_manager.js @@ -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; } diff --git a/web/js/admin.js b/web/js/admin.js index 157abd36..5aba7aa7 100644 --- a/web/js/admin.js +++ b/web/js/admin.js @@ -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(() => { diff --git a/web/js/comfy-metadata.js b/web/js/comfy-metadata.js index 8e35fc1d..234f66a5 100644 --- a/web/js/comfy-metadata.js +++ b/web/js/comfy-metadata.js @@ -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"], }; diff --git a/web/js/focused-node.js b/web/js/focused-node.js index 138904c8..f2eed6f6 100644 --- a/web/js/focused-node.js +++ b/web/js/focused-node.js @@ -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) { diff --git a/web/js/gallery.js b/web/js/gallery.js index c94e7368..56596f7e 100644 --- a/web/js/gallery.js +++ b/web/js/gallery.js @@ -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 = `
@@ -1309,6 +1316,8 @@
+ ${textSourceLabel ? `

${escapeHtml(textSourceLabel)}

` : ''} +
diff --git a/web/js/prompt-editor.css b/web/js/prompt-editor.css index 0698e786..ba6fbf4a 100644 --- a/web/js/prompt-editor.css +++ b/web/js/prompt-editor.css @@ -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; } +} diff --git a/web/js/prompt-library.js b/web/js/prompt-library.js index b3059854..43ff249f 100644 --- a/web/js/prompt-library.js +++ b/web/js/prompt-library.js @@ -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()); diff --git a/web/js/scan-progress.js b/web/js/scan-progress.js index 567f0635..8105b9fa 100644 --- a/web/js/scan-progress.js +++ b/web/js/scan-progress.js @@ -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,