* feat(images): link summary on workspace image items
* fix(images): carry title source and cut first-line titles like the Library
Summaries gain title_source (stored or text). The inspector always shows a
stored title and confirms only a first-line title against the file, so a
stored title equal to the first line is no longer withheld. A link-title read
failure no longer marks a readable link unavailable. First-line titles over
80 characters end with an ellipsis, as the Library's title() does. The card
gate for unconfirmed first-line titles is pinned as an expected failure for B2.
* feat(images): unambiguous seed on workspace items
Workspace items gain `seed`, read with the shared generation-parameter
parser from the prompt/workflow records the header open already loaded
(each at most 2 MiB; no extra file reads, videos give None). A
placeholder or non-integer seed gives None. JSON integers beyond 2**53
are recovered exactly, so the card prints the digits the inspector shows.
The same parse confirms first-line card titles: a `text` title shows only
when the file's parsed positive gives the same title. One helper applies
that rule on the card and in the inspector, which now reads the positive
from its own inspection records with the same parser.
* fix(images): a deeply nested record no longer fails the workspace listing
The Python parser recurses over nested arrays where JavaScript String()
does not, so a crafted 2 KB PNG turned every workspace page holding it
(and its inspect route) into a 500. Card facts now treat RecursionError
like any record the parser cannot use: no seed and no confirmed title.
Also: Library previews skip the card parse they never show; the parse
runs after the file handle closes; the exact-seed read uses the one JSON
reader (parse_js_json(..., exact_integers=True)) and the public
number_string; record size is decided from the text length unless only
the UTF-8 byte count can tell.
* feat(images): search and filter workspace images by linked prompts
The workspace listing now takes category (exact), tags (comma list, all
required) and sort=rating, and q also matches the linked positive's text,
Library title, pair title and tag names. One bounded, parameterized SELECT
(get_image_link_matches / get_image_paths_matching) returns normalized
paths of unambiguously linked images; the listing is intersected with it.
Unlinked and ambiguous files are excluded by a category or tag filter and
sort last by rating, with a newest-first tie-break. The cursor scope covers
the new filters and sort; unknown values are still rejected.
The derived listing values and sort orders are cached per gallery scan so a
50,000-file listing with 20,000 links pages well under 150 ms at p95.
* test(images): count the filter budget everywhere, time it in one job
The 150 ms wall-clock budget failed under the CI coverage tracer, so it
is split. test_filter_work_is_bounded runs on every job and under
coverage: one filter SELECT per request, the ambiguity probe is an index
search (EXPLAIN QUERY PLAN), the per-scan values are built on the cold
request only, and a cursor page calls the sort key at most
ceil(log2 n)+3 times. test_filter_budget runs only with
PROMPTMANAGER_PERF_BUDGET=1 and no tracer, asserts the best of three
rounds' p95, and has its own non-required "perf budget" job. TESTING.md
records the convention.
* fix(images): rating sort works however many link rows the database holds
The filter SELECT counted every generated_images row against a 500,000
row cap and materialized all of them first, so a large history made
sort=rating answer 400 with nothing to narrow. The rows are now streamed
and only paths in the listing's normalized set are kept, so memory is
bounded by the listing; the cap and LinkFilterTooBroad are gone.
* fix(images): rebuild workspace listing values only where the scan changed
A TTL rescan that finds the same files reuses the previous values and
sort orders; one that finds a few changes rebuilds only those files.
Gallery invalidation and a gallery-roots change drop the derived values,
so old roots no longer hold memory. category is trimmed like tags.
* fix(images): a monitor invalidation rebuilds only the new files' values
The image monitor invalidates the gallery whenever it sees a new file,
and that dropped every derived workspace value, so the next request paid
a full cold build. Invalidation now keeps the per-file values and drops
only the sort orders; the rescan reuses every unchanged file and builds
only new or changed ones. A gallery-roots change still drops everything.
* feat(library): expose run count, last use and variant totals
* fix(library): an out-of-range last use reads as unknown, not a failed list
* perf(images): read positive counts only on a plain listing's first page
* feat(images): persist image favorites
Favorites are stored in image_favorites keyed by the gallery root label
and the relative path, set and cleared through PUT/DELETE
/prompt_manager/images/favorite with a {root_index, path} body validated
against the rooted listing. Workspace image cards carry favorite: bool.
The table is created by migration, travels with backups, is optional in
old backups and required after a restore's migration.
* fix(images): key favorites by the real root and case-normalized path
Folder names are not unique across gallery roots, so two roots shown as
"output" shared each other's stars. Favorites are now keyed by the
root's real, case-normalized path (server-side only) and a
case-normalized relative path, so a Windows case-only rename keeps its
star while distinct Linux files stay distinct. Toggles look the file up
in a per-scan dictionary instead of walking the listing.
* fix(images): favorites lookup follows monitor invalidation and the stored key bound
A monitor invalidation now also drops the favorites lookup, which held the
previous scan list until the next toggle. PUT refuses a favorite whose root
or relative key would exceed the schema's 4096-character bound, so the live
database keeps passing startup and restore validation.
* docs(testing): workspace data fields for the concept pages
* fix(images): card facts read the inspector's projected graphs
The inspector hands its browser parser only graphs the inspection
projection parsed (strict JSON, depth 64, node ceiling); a deeper record
is "limited" there, so the legacy fill never runs and no seed shows. The
card parsed the raw records instead and showed a seed the inspector did
not, e.g. for a 900-deep class_type beside a normal sampler.
project_graph() is now the one projection step, used by the inspection
response and by the card. The card runs the shared parser on the
projected graphs, the same values the browser receives: big integers
arrive as their exact digits, so the separate exact-integer re-parse and
the parse_js_json flag it needed are gone.
* fix(images): send the card seed as a decimal string
JSON.parse rounds a number past 2**53, so a listing seed such as
9223372036854788153 reached the page as 9223372036854788000 while the
inspector showed the file's digits. Every non-null card seed is now its
decimal string.
The Node parity harness now JSON.parses the real listing response and
compares String(seed) with the seed the inspector shows, including a
seed past 2**53. TESTING.md describes the string contract and drops the
removed exact re-read.
* fix(images): bound digit-string seeds and isolate each card's parse
A seed recorded as a 5000-digit string reached int(), which raises
ValueError past Python's integer string limit, and the listing answered
500. A digit-string seed now has at most 20 digits (2**64 fits), and the
conversion falls back to no seed on ValueError.
Each card's parse is also guarded in the page builder: an unexpected
error in one file's records logs a warning and gives that card no seed
and no first-line title, and the rest of the page still lists.
* fix(images): filename search survives a failed prompt-link read
On main a workspace q matched only file paths, so search worked with a
broken database. Since search also matches linked prompts, any q turned
a failed link read into a 500. A q without category, tags or the rating
sort now falls back to filename matches with a null positive_count and
logs the failure with its traceback; those linked filters still fail.
* fix(images): card reads PNG records with the inspector's reader
The card read PNG text through Pillow's header info, where a later
duplicate prompt chunk silently wins and text after IDAT is missed. The
inspector's bounded reader refuses conflicting duplicates and reads text
after IDAT, so a file with two different prompt chunks showed a card
seed the inspector refused.
PNG cards now read their records with read_inspection_file, the reader
the inspector uses. When it refuses a file the card keeps its
dimensions from Pillow and shows no seed. Other formats still use
Pillow's header info, as the inspector's header read does.
* perf(images): read each card's header facts once per gallery scan
Every page parsed each card's prompt and workflow records again, about
3.4 ms per 50 KB record set and 25 ms per 400 KB workflow, inside
ComfyUI's server process, on every scroll-back and refresh.
Dimensions, the seed and the parsed positive's display title are now
cached per root by (path, mtime), beside the scan's per-file values and
invalidated the same way: a rescan keeps unchanged files and prunes the
rest, a roots change drops them all, and only files of the current scan
are stored. An unreadable header is not stored. Links and favorites
are still read on every page.
* ci: run the perf budget on main pushes and on demand only
The wall-clock budget has under 2x margin on shared ubuntu runners, and
a red exact-head check blocks a merge that must never be rerun to green.
The job now runs on each push to main and on workflow_dispatch, not on
pull requests. TESTING.md says a red run on main means fix the code.
* fix(library): cut display titles at 80 UTF-16 units, never inside a pair
The server counted code points while the Library's JavaScript counts UTF-16
units, so a first line of emoji kept 80 emoji on cards but showed 40 in the
Library. Both now cut at 80 UTF-16 units and drop a character that would
straddle the limit instead of leaving half a surrogate pair.
2471 lines
91 KiB
Python
2471 lines
91 KiB
Python
"""
|
|
Database operations for KikoTextEncode prompt storage and retrieval.
|
|
"""
|
|
|
|
import sqlite3
|
|
import json
|
|
import datetime
|
|
import os
|
|
from typing import Optional, List, Dict, Any, Set, Tuple
|
|
|
|
from .models import PromptModel, normalize_image_path, utc_now_iso
|
|
from .pairs import PairDependencyError, PairOperations
|
|
from .library import LibraryOperations, titled
|
|
from .image_favorites import ImageFavoriteOperations
|
|
from .library_commands import LibraryCommandOperations
|
|
from .vocabulary import VocabularyOperations
|
|
from .execution_history import ExecutionHistoryOperations
|
|
from .execution_contexts import ExecutionContextOperations, LIVE_HINT_CTES
|
|
from .execution_outputs import ExecutionOutputOperations
|
|
from .history_reads import HistoryReadOperations
|
|
|
|
# Import logging system
|
|
try:
|
|
from ..utils.logging_config import get_logger
|
|
from ..utils.hashing import generate_prompt_hash
|
|
except ImportError:
|
|
import sys
|
|
|
|
current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
sys.path.insert(0, current_dir)
|
|
from utils.logging_config import get_logger
|
|
from utils.hashing import generate_prompt_hash
|
|
|
|
|
|
# Subquery to fetch tags from junction table, embedded in SELECT statements
|
|
TAG_SUBQUERY = (
|
|
"(SELECT GROUP_CONCAT(t.name, '|||') "
|
|
"FROM prompt_tags pt JOIN tags t ON pt.tag_id = t.id "
|
|
"WHERE pt.prompt_id = prompts.id) AS _tag_list"
|
|
)
|
|
|
|
# Server-side sort orders, keyed by the value the UI sends. Only these strings ever
|
|
# reach ORDER BY; unknown keys fall back to DEFAULT_SORT.
|
|
# last_used_at is NULL until the prompt has run, so "IS NULL" first keeps
|
|
# never-run prompts after every run one (SQLite sorts NULL first in DESC).
|
|
SORT_ORDERS = {
|
|
"last_used_desc": "last_used_at IS NULL, last_used_at DESC, id DESC",
|
|
"created_desc": "created_at DESC, id DESC",
|
|
"created_asc": "created_at ASC, id ASC",
|
|
"run_count_desc": (
|
|
"run_count DESC, last_used_at IS NULL, last_used_at DESC, id DESC"
|
|
),
|
|
"rating_desc": "rating IS NULL, rating DESC, id DESC",
|
|
"rating_asc": "rating IS NULL, rating ASC, id DESC",
|
|
"text_asc": "text COLLATE NOCASE ASC, id ASC",
|
|
"text_desc": "text COLLATE NOCASE DESC, id DESC",
|
|
}
|
|
DEFAULT_SORT = "created_desc"
|
|
|
|
|
|
def order_by_clause(sort: Optional[str]) -> str:
|
|
"""Return a whitelisted ORDER BY expression for a UI sort key."""
|
|
return SORT_ORDERS.get(sort or DEFAULT_SORT, SORT_ORDERS[DEFAULT_SORT])
|
|
|
|
|
|
LIKE_ESCAPE = "\\"
|
|
|
|
|
|
def escape_like(term: str) -> str:
|
|
"""Escape LIKE wildcards in user text; use with ``LIKE ? ESCAPE '\\'``."""
|
|
return (
|
|
term.replace(LIKE_ESCAPE, LIKE_ESCAPE + LIKE_ESCAPE)
|
|
.replace("%", LIKE_ESCAPE + "%")
|
|
.replace("_", LIKE_ESCAPE + "_")
|
|
)
|
|
|
|
|
|
def _resolve_db_path(db_path: Optional[str] = None) -> str:
|
|
"""Resolve the database path from config, falling back to defaults.
|
|
|
|
When no explicit path is given, reads PromptManagerConfig.DEFAULT_DB_PATH
|
|
(set via config.json). Relative paths are resolved against the extension
|
|
root directory so the database location is predictable regardless of the
|
|
process working directory (important for Docker deployments).
|
|
|
|
Args:
|
|
db_path: Explicit database path, or None to use config/default.
|
|
|
|
Returns:
|
|
Absolute path to the SQLite database file.
|
|
"""
|
|
extension_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
|
|
if db_path is None:
|
|
try:
|
|
try:
|
|
from ..py.config import PromptManagerConfig
|
|
except ImportError:
|
|
from py.config import PromptManagerConfig
|
|
|
|
db_path = PromptManagerConfig.DEFAULT_DB_PATH
|
|
except Exception:
|
|
db_path = "prompts.db"
|
|
|
|
if not os.path.isabs(db_path):
|
|
db_path = os.path.normpath(os.path.join(extension_root, db_path))
|
|
|
|
# Ensure parent directory exists for custom paths
|
|
parent = os.path.dirname(db_path)
|
|
if parent:
|
|
os.makedirs(parent, exist_ok=True)
|
|
|
|
return db_path
|
|
|
|
|
|
# Bounds for the file_info a link request may attach to an image. Anything
|
|
# outside them (or of the wrong type) is stored as NULL instead of being
|
|
# trusted into an INTEGER or TEXT column and echoed back to the gallery.
|
|
MAX_IMAGE_FILE_SIZE = 10**12
|
|
MAX_IMAGE_DIMENSION = 65535
|
|
IMAGE_FORMATS = frozenset(
|
|
{"PNG", "JPEG", "JPG", "WEBP", "GIF", "BMP", "TIFF", "MP4", "WEBM", "MOV"}
|
|
)
|
|
|
|
|
|
def _bounded_int(value: Any, upper: int) -> Optional[int]:
|
|
"""``value`` as an int within ``[0, upper]``; None for anything else."""
|
|
if isinstance(value, bool):
|
|
return None
|
|
try:
|
|
number = int(value)
|
|
except (TypeError, ValueError, OverflowError):
|
|
return None
|
|
return number if 0 <= number <= upper else None
|
|
|
|
|
|
def _known_image_format(value: Any) -> Optional[str]:
|
|
"""Upper-cased *value* when it names a supported media format, else None."""
|
|
if not isinstance(value, str):
|
|
return None
|
|
name = value.strip().upper()
|
|
return name if name in IMAGE_FORMATS else None
|
|
|
|
|
|
def _coerce_file_info(metadata: Any) -> tuple:
|
|
"""``(size, width, height, format)`` from request metadata, each validated."""
|
|
file_info = metadata.get("file_info") if isinstance(metadata, dict) else None
|
|
if not isinstance(file_info, dict):
|
|
return None, None, None, None
|
|
dimensions = file_info.get("dimensions")
|
|
if not isinstance(dimensions, (list, tuple)) or len(dimensions) < 2:
|
|
dimensions = (None, None)
|
|
return (
|
|
_bounded_int(file_info.get("size"), MAX_IMAGE_FILE_SIZE),
|
|
_bounded_int(dimensions[0], MAX_IMAGE_DIMENSION),
|
|
_bounded_int(dimensions[1], MAX_IMAGE_DIMENSION),
|
|
_known_image_format(file_info.get("format")),
|
|
)
|
|
|
|
|
|
# Exempt references: images, pairs, the __protected__ tag, Library parent
|
|
# state, and live recorded-run hints (computed once, see LIVE_HINT_CTES).
|
|
ORPHANED_PROMPTS_SQL = "WITH " + LIVE_HINT_CTES + """
|
|
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
|
|
WHERE pt.prompt_id = p.id AND t.name = '__protected__'
|
|
)
|
|
AND NOT EXISTS (
|
|
SELECT 1 FROM prompt_library_state ls WHERE ls.prompt_id = p.id
|
|
)
|
|
AND p.id NOT IN (SELECT prompt_id FROM live_hint_prompts)
|
|
"""
|
|
|
|
|
|
# One SELECT reads the link summaries of at most this many paths: three bound
|
|
# parameters per path stays under SQLite's historical 999-variable limit.
|
|
LINK_SUMMARY_BATCH = 200
|
|
# A pair link counts only when an image has exactly one distinct pair, as in
|
|
# get_image_prompt_info; its positive, title and metadata then describe the card.
|
|
_LINK_SUMMARY_SQL = """SELECT gi.image_path, gi.file_path, gi.prompt_id,
|
|
p.text, p.category, p.rating, ls.title,
|
|
pp.positive_prompt_id, pp.title, pt.text, pt.category, pt.rating, pls.title
|
|
FROM generated_images gi
|
|
JOIN prompts p ON gi.prompt_id = p.id
|
|
LEFT JOIN prompt_library_state ls ON ls.prompt_id = p.id
|
|
LEFT JOIN prompt_pairs pp ON pp.id = (
|
|
SELECT CASE WHEN COUNT(DISTINCT ip.pair_id) = 1 THEN MIN(ip.pair_id) END
|
|
FROM image_prompt_pairs ip WHERE ip.image_id = gi.id)
|
|
LEFT JOIN prompts pt ON pt.id = pp.positive_prompt_id
|
|
LEFT JOIN prompt_library_state pls ON pls.prompt_id = pp.positive_prompt_id
|
|
WHERE gi.image_path IN ({marks}) OR gi.file_path IN ({marks}){fallback}"""
|
|
# One image row describes a card only when no other row claims the same file
|
|
# (as in get_image_link_summaries, several rows are ambiguous). The pair scope
|
|
# and the linked positive follow _LINK_SUMMARY_SQL; category, rating, tags and
|
|
# the searchable text and titles come from that positive (spec D12/D13). Every
|
|
# writer since the 3.2.4 migration stores file_path, so it is the path key.
|
|
# The prompt-level conditions are uncorrelated IN sets SQLite builds once, and
|
|
# the CASE evaluates the per-row ambiguity probe only for admitted rows.
|
|
_LINK_FILTER_SQL = """SELECT gi.file_path, p.rating, {hit}, p.id
|
|
FROM generated_images gi
|
|
LEFT JOIN prompt_pairs pp ON pp.id = (
|
|
SELECT CASE WHEN COUNT(DISTINCT ip.pair_id) = 1 THEN MIN(ip.pair_id) END
|
|
FROM image_prompt_pairs ip WHERE ip.image_id = gi.id)
|
|
JOIN prompts p ON p.id = COALESCE(pp.positive_prompt_id, gi.prompt_id)
|
|
WHERE gi.file_path IS NOT NULL AND CASE WHEN {admit} THEN NOT EXISTS (
|
|
SELECT 1 FROM generated_images other
|
|
WHERE other.file_path = gi.file_path AND other.id <> gi.id) ELSE 0 END"""
|
|
# The positive's text, its Library title or one of its tag names, or the pair
|
|
# title, contains :pattern (LIKE, so ASCII case is ignored as in Library search).
|
|
_LINK_FILTER_HIT = """(p.id IN (
|
|
SELECT hp.id FROM prompts hp
|
|
LEFT JOIN prompt_library_state hls ON hls.prompt_id = hp.id
|
|
WHERE hp.text LIKE :pattern ESCAPE '\\' OR hls.title LIKE :pattern ESCAPE '\\'
|
|
OR hp.id IN (SELECT ht.prompt_id FROM prompt_tags ht JOIN tags hn
|
|
ON hn.id = ht.tag_id WHERE hn.name LIKE :pattern ESCAPE '\\'))
|
|
OR pp.title LIKE :pattern ESCAPE '\\')"""
|
|
_LINK_FILTER_TAG = (
|
|
"p.id IN (SELECT ft.prompt_id FROM prompt_tags ft JOIN tags fn "
|
|
"ON fn.id = ft.tag_id WHERE fn.name = :tag{index})"
|
|
)
|
|
|
|
|
|
def link_filter_query(
|
|
q: str, category: Optional[str], tags: List[str], *, every_link: bool = False
|
|
) -> Tuple[str, Dict[str, Any]]:
|
|
"""The workspace filter SELECT and its named parameters (all bound).
|
|
|
|
Only code constants and enumerate indices are formatted into the SQL; see
|
|
PromptDatabase.get_image_link_matches for the semantics.
|
|
"""
|
|
params: Dict[str, Any] = {}
|
|
hit = "1"
|
|
if q:
|
|
hit = _LINK_FILTER_HIT
|
|
params["pattern"] = f"%{escape_like(q)}%"
|
|
admitted = []
|
|
if category is not None:
|
|
admitted.append("p.category = :category")
|
|
params["category"] = category
|
|
for index, tag in enumerate(tags):
|
|
admitted.append(_LINK_FILTER_TAG.format(index=index))
|
|
params[f"tag{index}"] = tag
|
|
if q and not admitted and not every_link:
|
|
admitted.append(hit)
|
|
sql = _LINK_FILTER_SQL.format(hit=hit, admit=" AND ".join(admitted) or "1")
|
|
return sql, params
|
|
|
|
|
|
_ASCII_FOLD = str.maketrans("ABCDEFGHIJKLMNOPQRSTUVWXYZ", "abcdefghijklmnopqrstuvwxyz")
|
|
|
|
|
|
def _like_suffix(value: object, suffix: str) -> bool:
|
|
"""Python form of SQLite's ``value LIKE '%' || suffix``: ASCII-only case."""
|
|
return isinstance(value, str) and value.translate(_ASCII_FOLD).endswith(
|
|
suffix.translate(_ASCII_FOLD)
|
|
)
|
|
|
|
|
|
def _empty_link_summary(status: str) -> Dict[str, Any]:
|
|
return {
|
|
"status": status,
|
|
"positive_id": None,
|
|
"title": None,
|
|
"title_source": None,
|
|
"category": None,
|
|
"rating": None,
|
|
}
|
|
|
|
|
|
def _link_summary(row: tuple) -> Dict[str, Any]:
|
|
if row[7] is not None:
|
|
positive, text, category, rating = row[7], row[9], row[10], row[11]
|
|
titles = (row[8], row[12])
|
|
else:
|
|
positive, text, category, rating = row[2], row[3], row[4], row[5]
|
|
titles = (row[6],)
|
|
title, source = titled(titles, text)
|
|
return {
|
|
"status": "linked",
|
|
"positive_id": positive,
|
|
"title": title,
|
|
"title_source": source,
|
|
"category": category,
|
|
"rating": rating,
|
|
}
|
|
|
|
|
|
class PromptDatabase(
|
|
PairOperations,
|
|
LibraryOperations,
|
|
LibraryCommandOperations,
|
|
VocabularyOperations,
|
|
ExecutionHistoryOperations,
|
|
ExecutionContextOperations,
|
|
ExecutionOutputOperations,
|
|
HistoryReadOperations,
|
|
ImageFavoriteOperations,
|
|
):
|
|
"""Database operations class for managing prompts."""
|
|
|
|
def __init__(self, db_path: Optional[str] = None):
|
|
"""
|
|
Initialize the database operations.
|
|
|
|
Args:
|
|
db_path: Path to the SQLite database file. If None, resolves from
|
|
config.json or defaults to prompts.db in the extension root.
|
|
"""
|
|
self.logger = get_logger("prompt_manager.database")
|
|
db_path = _resolve_db_path(db_path)
|
|
self.logger.debug(f"Initializing database operations with path: {db_path}")
|
|
self.model = PromptModel(db_path)
|
|
self.logger.debug("Database operations initialized successfully")
|
|
|
|
def close(self) -> None:
|
|
"""Close the calling thread's database connection (see PromptModel.close)."""
|
|
self.model.close()
|
|
|
|
@staticmethod
|
|
def close_all_instances() -> None:
|
|
"""Close the connections of every PromptDatabase/PromptModel in the process."""
|
|
PromptModel.close_all_instances()
|
|
|
|
def close_all(self) -> None:
|
|
"""Close every thread's connection (see PromptModel.close_all)."""
|
|
self.model.close_all()
|
|
|
|
def save_prompt(
|
|
self,
|
|
text: str,
|
|
category: Optional[str] = None,
|
|
tags: Optional[List[str]] = None,
|
|
rating: Optional[int] = None,
|
|
notes: Optional[str] = None,
|
|
prompt_hash: Optional[str] = None,
|
|
) -> int:
|
|
"""
|
|
Save a new prompt to the database.
|
|
|
|
Args:
|
|
text: The prompt text
|
|
category: Optional category
|
|
tags: List of tags
|
|
rating: Rating 1-5
|
|
notes: Optional notes
|
|
prompt_hash: SHA256 hash of the prompt
|
|
|
|
Returns:
|
|
int: The ID of the saved prompt
|
|
|
|
Raises:
|
|
ValueError: If required parameters are invalid
|
|
sqlite3.Error: If database operation fails
|
|
"""
|
|
if not text or not text.strip():
|
|
raise ValueError("Prompt text cannot be empty")
|
|
|
|
if rating is not None and (rating < 1 or rating > 5):
|
|
raise ValueError("Rating must be between 1 and 5")
|
|
|
|
self.logger.debug(
|
|
f"Saving prompt: text_length={len(text)}, category={category}, "
|
|
f"tags={tags}, rating={rating}"
|
|
)
|
|
|
|
try:
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
"""
|
|
INSERT INTO prompts (
|
|
text, category, tags, rating, notes, hash, created_at,
|
|
updated_at
|
|
) VALUES (?, ?, NULL, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
text.strip(),
|
|
category,
|
|
rating,
|
|
notes,
|
|
prompt_hash,
|
|
datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
|
datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
|
),
|
|
)
|
|
prompt_id = cursor.lastrowid
|
|
if tags:
|
|
self._sync_prompt_tags(conn, prompt_id, tags)
|
|
conn.commit()
|
|
self.logger.debug(f"Successfully saved prompt with ID: {prompt_id}")
|
|
return prompt_id
|
|
except sqlite3.IntegrityError:
|
|
# Another caller inserted the same hash between the caller's
|
|
# duplicate check and this INSERT (the with-block rolled back).
|
|
existing = self.get_prompt_by_hash(prompt_hash) if prompt_hash else None
|
|
if existing is None:
|
|
raise
|
|
self.logger.debug(
|
|
f"Prompt with hash {prompt_hash} was saved concurrently; "
|
|
f"returning existing ID {existing['id']}"
|
|
)
|
|
return existing["id"]
|
|
|
|
def record_prompt_use(self, prompt_id: int, times: int = 1) -> bool:
|
|
"""
|
|
Count runs of a prompt and mark it as most recently used.
|
|
|
|
Args:
|
|
prompt_id: The prompt ID
|
|
times: Number of runs to add (queued runs counted on first save)
|
|
|
|
Returns:
|
|
True if the prompt exists and was updated
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
"UPDATE prompts SET run_count = run_count + ?, last_used_at = ?"
|
|
" WHERE id = ?",
|
|
(max(1, int(times)), utc_now_iso(), prompt_id),
|
|
)
|
|
conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
def get_prompt_by_id(self, prompt_id: int) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Get a prompt by its ID.
|
|
|
|
Args:
|
|
prompt_id: The prompt ID
|
|
|
|
Returns:
|
|
Dict containing prompt data or None if not found
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts WHERE id = ?",
|
|
(prompt_id,),
|
|
)
|
|
row = cursor.fetchone()
|
|
return self._row_to_dict(row) if row else None
|
|
|
|
def get_prompt_by_hash(self, prompt_hash: str) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Get a prompt by its hash.
|
|
|
|
Args:
|
|
prompt_hash: SHA256 hash of the prompt
|
|
|
|
Returns:
|
|
Dict containing prompt data or None if not found
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts WHERE hash = ?",
|
|
(prompt_hash,),
|
|
)
|
|
row = cursor.fetchone()
|
|
return self._row_to_dict(row) if row else None
|
|
|
|
def search_prompts(
|
|
self,
|
|
text: Optional[str] = None,
|
|
category: Optional[str] = None,
|
|
tags: Optional[List[str]] = None,
|
|
rating_min: Optional[int] = None,
|
|
rating_max: Optional[int] = None,
|
|
date_from: Optional[str] = None,
|
|
date_to: Optional[str] = None,
|
|
folder: Optional[str] = None,
|
|
limit: int = 100,
|
|
offset: int = 0,
|
|
tag_partial: bool = False,
|
|
sort: Optional[str] = None,
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Search prompts with various filters.
|
|
|
|
Args:
|
|
text: Text to search for in prompt content
|
|
category: Filter by category
|
|
tags: Filter by tags (must contain all specified tags)
|
|
rating_min: Minimum rating filter
|
|
rating_max: Maximum rating filter
|
|
date_from: Start date filter (ISO format)
|
|
date_to: End date filter (ISO format)
|
|
folder: Filter by subfolder name in generated image paths
|
|
limit: Maximum number of results
|
|
offset: Number of results to skip
|
|
tag_partial: Use LIKE matching for tags instead of exact match
|
|
sort: Sort key from SORT_ORDERS (default: newest first)
|
|
|
|
Returns:
|
|
List of dictionaries containing prompt data
|
|
"""
|
|
query_parts = [f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts WHERE 1=1"]
|
|
params = []
|
|
|
|
if text:
|
|
query_parts.append("AND text LIKE ? ESCAPE '\\'")
|
|
params.append(f"%{escape_like(text)}%")
|
|
|
|
if category:
|
|
query_parts.append("AND category = ?")
|
|
params.append(category)
|
|
|
|
if tags:
|
|
for tag in tags:
|
|
if tag_partial:
|
|
query_parts.append(
|
|
"AND prompts.id IN ("
|
|
" SELECT pt.prompt_id FROM prompt_tags pt"
|
|
" JOIN tags t ON pt.tag_id = t.id"
|
|
" WHERE t.name LIKE ? ESCAPE '\\')"
|
|
)
|
|
params.append(f"%{escape_like(tag)}%")
|
|
else:
|
|
query_parts.append(
|
|
"AND prompts.id IN ("
|
|
" SELECT pt.prompt_id FROM prompt_tags pt"
|
|
" JOIN tags t ON pt.tag_id = t.id WHERE t.name = ?)"
|
|
)
|
|
params.append(tag)
|
|
|
|
if folder:
|
|
# Escape LIKE wildcards in the folder name
|
|
safe_folder = (
|
|
folder.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
|
)
|
|
query_parts.append(
|
|
"AND prompts.id IN ("
|
|
" SELECT DISTINCT prompt_id FROM generated_images"
|
|
" WHERE image_path LIKE ? ESCAPE '\\'"
|
|
" OR image_path LIKE ? ESCAPE '\\'"
|
|
" OR image_path LIKE ? ESCAPE '\\'"
|
|
")"
|
|
)
|
|
# Pattern 1: folder between forward slashes (Unix + normalized Windows)
|
|
params.append(f"%/{safe_folder}/%")
|
|
# Pattern 2: folder between backslashes (Windows native)
|
|
params.append(f"%\\\\{safe_folder}\\\\%")
|
|
# Pattern 3: folder at start of a relative path
|
|
params.append(f"{safe_folder}/%")
|
|
|
|
if rating_min is not None:
|
|
query_parts.append("AND rating >= ?")
|
|
params.append(rating_min)
|
|
|
|
if rating_max is not None:
|
|
query_parts.append("AND rating <= ?")
|
|
params.append(rating_max)
|
|
|
|
if date_from:
|
|
query_parts.append("AND created_at >= ?")
|
|
params.append(date_from)
|
|
|
|
if date_to:
|
|
query_parts.append("AND created_at <= ?")
|
|
params.append(date_to)
|
|
|
|
query_parts.append(f"ORDER BY {order_by_clause(sort)} LIMIT ? OFFSET ?")
|
|
params.extend([limit, offset])
|
|
|
|
query = " ".join(query_parts)
|
|
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(query, params)
|
|
rows = cursor.fetchall()
|
|
prompts = [self._row_to_dict(row) for row in rows]
|
|
|
|
if prompts:
|
|
prompt_ids = [p["id"] for p in prompts]
|
|
self._attach_preview_images(conn, prompts, prompt_ids)
|
|
|
|
return prompts
|
|
|
|
def get_recent_prompts(
|
|
self, limit: int = 10, offset: int = 0, sort: Optional[str] = None
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Get the most recent prompts with pagination support.
|
|
|
|
Args:
|
|
limit: Maximum number of prompts to return
|
|
offset: Number of prompts to skip (for pagination)
|
|
sort: Sort key from SORT_ORDERS (default: newest first)
|
|
|
|
Returns:
|
|
Dictionary containing prompt data and pagination info
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
# Get total count
|
|
cursor = conn.execute("SELECT COUNT(*) FROM prompts")
|
|
row = cursor.fetchone()
|
|
total_count = (row[0] if row else 0) or 0
|
|
|
|
# Get paginated results
|
|
cursor = conn.execute(
|
|
f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts"
|
|
f" ORDER BY {order_by_clause(sort)} LIMIT ? OFFSET ?",
|
|
(limit, offset),
|
|
)
|
|
rows = cursor.fetchall()
|
|
prompts = [self._row_to_dict(row) for row in rows]
|
|
|
|
if prompts:
|
|
prompt_ids = [p["id"] for p in prompts]
|
|
self._attach_preview_images(conn, prompts, prompt_ids)
|
|
|
|
return {
|
|
"prompts": prompts,
|
|
"total": total_count,
|
|
"limit": limit,
|
|
"offset": offset,
|
|
"has_more": (offset + limit) < total_count,
|
|
"page": (offset // limit) + 1,
|
|
"total_pages": (total_count + limit - 1) // limit, # Ceiling division
|
|
}
|
|
|
|
def get_prompts_by_category(
|
|
self, category: str, limit: int = 100
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Get all prompts in a specific category.
|
|
|
|
Args:
|
|
category: The category name
|
|
limit: Maximum number of prompts to return
|
|
|
|
Returns:
|
|
List of dictionaries containing prompt data
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts WHERE category = ? "
|
|
"ORDER BY created_at DESC LIMIT ?",
|
|
(category, limit),
|
|
)
|
|
rows = cursor.fetchall()
|
|
return [self._row_to_dict(row) for row in rows]
|
|
|
|
def get_top_rated_prompts(self, limit: int = 10) -> List[Dict[str, Any]]:
|
|
"""
|
|
Get the highest rated prompts.
|
|
|
|
Args:
|
|
limit: Maximum number of prompts to return
|
|
|
|
Returns:
|
|
List of dictionaries containing prompt data
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
f"""
|
|
SELECT prompts.*, {TAG_SUBQUERY} FROM prompts
|
|
WHERE rating IS NOT NULL
|
|
ORDER BY rating DESC, created_at DESC
|
|
LIMIT ?
|
|
""",
|
|
(limit,),
|
|
)
|
|
rows = cursor.fetchall()
|
|
return [self._row_to_dict(row) for row in rows]
|
|
|
|
def update_prompt_metadata(
|
|
self,
|
|
prompt_id: int,
|
|
category: Optional[str] = None,
|
|
tags: Optional[List[str]] = None,
|
|
rating: Optional[int] = None,
|
|
notes: Optional[str] = None,
|
|
) -> bool:
|
|
"""
|
|
Update metadata for an existing prompt.
|
|
|
|
Args:
|
|
prompt_id: The prompt ID
|
|
category: New category
|
|
tags: New tags list
|
|
rating: New rating
|
|
notes: New notes
|
|
|
|
Returns:
|
|
bool: True if update was successful, False otherwise
|
|
"""
|
|
updates = []
|
|
params = []
|
|
|
|
if category is not None:
|
|
updates.append("category = ?")
|
|
params.append(category)
|
|
|
|
if rating is not None:
|
|
if rating < 1 or rating > 5:
|
|
raise ValueError("Rating must be between 1 and 5")
|
|
updates.append("rating = ?")
|
|
params.append(rating)
|
|
|
|
if notes is not None:
|
|
updates.append("notes = ?")
|
|
params.append(notes)
|
|
|
|
if not updates and tags is None:
|
|
return False
|
|
|
|
with self.model.get_connection() as conn:
|
|
if updates:
|
|
updates.append("updated_at = ?")
|
|
params.append(datetime.datetime.now(datetime.timezone.utc).isoformat())
|
|
params.append(prompt_id)
|
|
query = f"UPDATE prompts SET {', '.join(updates)} WHERE id = ?"
|
|
conn.execute(query, params)
|
|
if tags is not None:
|
|
self._sync_prompt_tags(conn, prompt_id, tags)
|
|
conn.execute(
|
|
"UPDATE prompts SET updated_at = ? WHERE id = ?",
|
|
(
|
|
datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
|
prompt_id,
|
|
),
|
|
)
|
|
conn.commit()
|
|
return True
|
|
|
|
def delete_prompt(self, prompt_id: int) -> bool:
|
|
"""
|
|
Delete a prompt by its ID.
|
|
|
|
Args:
|
|
prompt_id: The prompt ID to delete
|
|
|
|
Returns:
|
|
bool: True if deletion was successful, False otherwise
|
|
"""
|
|
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,)
|
|
)
|
|
# Then delete the prompt
|
|
cursor = conn.execute("DELETE FROM prompts WHERE id = ?", (prompt_id,))
|
|
conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
def delete_prompts_by_category(self, category: str) -> int:
|
|
"""Delete all prompts with the given category.
|
|
|
|
Returns:
|
|
Number of prompts deleted.
|
|
"""
|
|
with self._pair_transaction() as conn:
|
|
# Get IDs first for cascade cleanup
|
|
ids = [
|
|
r[0]
|
|
for r in conn.execute(
|
|
"SELECT id FROM prompts WHERE category = ?", (category,)
|
|
).fetchall()
|
|
]
|
|
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})",
|
|
ids,
|
|
)
|
|
conn.execute(
|
|
f"DELETE FROM prompt_tags WHERE prompt_id IN ({placeholders})",
|
|
ids,
|
|
)
|
|
cursor = conn.execute(
|
|
f"DELETE FROM prompts WHERE id IN ({placeholders})",
|
|
ids,
|
|
)
|
|
conn.commit()
|
|
return cursor.rowcount
|
|
|
|
def get_all_categories(self) -> List[str]:
|
|
"""
|
|
Get all unique categories from the database.
|
|
|
|
Returns:
|
|
List of category names
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
"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()]
|
|
|
|
def get_prompt_subfolders(
|
|
self,
|
|
root_dirs: Optional[List[str]] = None,
|
|
include_ancestors: bool = False,
|
|
) -> List[str]:
|
|
"""
|
|
Get distinct subfolder paths from generated_images.
|
|
|
|
Extracts the directory portion of image_path, made relative to
|
|
root_dirs if provided. Returns sorted unique folder names.
|
|
|
|
Args:
|
|
root_dirs: Gallery root directories for computing relative paths
|
|
include_ancestors: If True, also include all ancestor path segments
|
|
so that hierarchical tree navigation works. For example, a path
|
|
``2026/08-Aug/2026-08-06`` also adds ``2026`` and ``2026/08-Aug``.
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
"SELECT DISTINCT image_path FROM generated_images "
|
|
"WHERE image_path IS NOT NULL AND image_path != ''"
|
|
)
|
|
|
|
folders = set()
|
|
for row in cursor.fetchall():
|
|
image_path = row["image_path"]
|
|
parent = os.path.dirname(image_path)
|
|
if not parent:
|
|
continue
|
|
|
|
made_relative = False
|
|
if root_dirs:
|
|
for root in root_dirs:
|
|
try:
|
|
rel = os.path.relpath(parent, root)
|
|
if not rel.startswith(".."):
|
|
if rel != ".":
|
|
# Clients split and filter on "/" whatever the OS
|
|
folders.add(rel.replace(os.sep, "/"))
|
|
made_relative = True
|
|
break
|
|
except ValueError:
|
|
continue
|
|
|
|
if not made_relative:
|
|
# Preserve full relative-style path instead of collapsing
|
|
# to just the basename, which loses hierarchy and creates
|
|
# ambiguity (e.g. foo/bar and baz/bar both become "bar").
|
|
folders.add(parent.replace(os.sep, "/"))
|
|
|
|
if include_ancestors:
|
|
ancestors = set()
|
|
for folder in folders:
|
|
parts = folder.replace("\\", "/").split("/")
|
|
for i in range(1, len(parts)):
|
|
ancestors.add("/".join(parts[:i]))
|
|
folders.update(ancestors)
|
|
|
|
return sorted(folders)
|
|
|
|
def get_all_tags(self) -> List[str]:
|
|
"""
|
|
Get every tag in the shared vocabulary, including unused registry rows.
|
|
|
|
Returns:
|
|
Sorted list of tag names
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute("SELECT t.name AS tag FROM tags t ORDER BY t.name")
|
|
return [row["tag"] for row in cursor.fetchall()]
|
|
|
|
def get_tags_with_counts(
|
|
self,
|
|
limit: int = 50,
|
|
offset: int = 0,
|
|
search: Optional[str] = None,
|
|
sort: str = "alpha_asc",
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Get all unique tags with their usage counts via junction table.
|
|
|
|
Args:
|
|
limit: Maximum number of tags to return
|
|
offset: Number of tags to skip
|
|
search: Optional case-insensitive substring filter
|
|
sort: Sort order - alpha_asc, alpha_desc, count_desc, count_asc
|
|
|
|
Returns:
|
|
Dict with tags list, total count, and pagination info
|
|
"""
|
|
search_clause = ""
|
|
params: list = []
|
|
if search:
|
|
search_clause = "HAVING t.name LIKE ? ESCAPE '\\'"
|
|
params.append(f"%{escape_like(search)}%")
|
|
|
|
sort_map = {
|
|
"alpha_desc": "tag COLLATE NOCASE DESC",
|
|
"count_desc": "count DESC, tag COLLATE NOCASE ASC",
|
|
"count_asc": "count ASC, tag COLLATE NOCASE ASC",
|
|
}
|
|
order = sort_map.get(sort, "tag COLLATE NOCASE ASC")
|
|
|
|
with self.model.get_connection() as conn:
|
|
count_sql = (
|
|
"SELECT COUNT(*) as total FROM ("
|
|
" SELECT t.name AS tag, COUNT(*) AS count"
|
|
" FROM prompt_tags pt JOIN tags t ON pt.tag_id = t.id"
|
|
f" GROUP BY t.name {search_clause}"
|
|
")"
|
|
)
|
|
cursor = conn.execute(count_sql, params)
|
|
row = cursor.fetchone()
|
|
total = (row[0] if row else 0) or 0
|
|
|
|
data_sql = (
|
|
"SELECT t.name AS tag, COUNT(*) AS count"
|
|
" FROM prompt_tags pt JOIN tags t ON pt.tag_id = t.id"
|
|
f" GROUP BY t.name {search_clause}"
|
|
f" ORDER BY {order}"
|
|
" LIMIT ? OFFSET ?"
|
|
)
|
|
cursor = conn.execute(data_sql, params + [limit, offset])
|
|
tags = [
|
|
{"name": row["tag"], "count": row["count"]} for row in cursor.fetchall()
|
|
]
|
|
|
|
return {
|
|
"tags": tags,
|
|
"total": total,
|
|
"limit": limit,
|
|
"offset": offset,
|
|
"has_more": (offset + limit) < total,
|
|
}
|
|
|
|
def suggest_tags(
|
|
self, prefix: str, with_tags: List[str], limit: int = 15
|
|
) -> List[Dict[str, Any]]:
|
|
"""Tag autocomplete narrowed by co-occurrence.
|
|
|
|
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. Counts include separate standalone
|
|
and pair references, with owner types kept distinct. Ordered by count/name.
|
|
Without context, unused vocabulary rows are included last with count 0;
|
|
with context they never appear.
|
|
"""
|
|
seen = set()
|
|
context = []
|
|
for raw in with_tags or []:
|
|
name = str(raw).strip()
|
|
if name and name.lower() not in seen:
|
|
seen.add(name.lower())
|
|
context.append(name)
|
|
|
|
where = ["t.name LIKE ? ESCAPE '\\'"]
|
|
params: list = [f"{escape_like(prefix or '')}%"]
|
|
if context:
|
|
placeholders = ",".join("?" * len(context))
|
|
where.append(
|
|
"(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.kind,pt2.owner_id"
|
|
" HAVING COUNT(DISTINCT lower(t2.name)) = ?)"
|
|
)
|
|
params.extend(context)
|
|
params.append(len(context))
|
|
where.append(f"t.name COLLATE NOCASE NOT IN ({placeholders})")
|
|
params.extend(context)
|
|
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(pt.tag_id) AS count"
|
|
" FROM tags t LEFT JOIN tag_refs pt ON t.id = pt.tag_id"
|
|
f" WHERE {' AND '.join(where)}"
|
|
" GROUP BY t.id"
|
|
" ORDER BY count DESC, t.name COLLATE NOCASE ASC"
|
|
" LIMIT ?"
|
|
)
|
|
with self.model.get_connection() as conn:
|
|
rows = conn.execute(sql, params).fetchall()
|
|
return [{"name": row["name"], "count": row["count"]} for row in rows]
|
|
|
|
def get_prompts_by_tags(
|
|
self, tags: List[str], mode: str = "and", limit: int = 20, offset: int = 0
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Get prompts that match the given tags with AND/OR filtering.
|
|
|
|
Uses junction table for precise tag matching and a batched image
|
|
fetch to avoid N+1 queries.
|
|
|
|
Args:
|
|
tags: List of tag names to filter by
|
|
mode: 'and' (must have all tags) or 'or' (must have any tag)
|
|
limit: Maximum number of prompts to return
|
|
offset: Number of prompts to skip
|
|
|
|
Returns:
|
|
Dict with prompts list (including preview images), total count, pagination
|
|
"""
|
|
if not tags:
|
|
return {
|
|
"prompts": [],
|
|
"total": 0,
|
|
"limit": limit,
|
|
"offset": offset,
|
|
"has_more": False,
|
|
}
|
|
|
|
with self.model.get_connection() as conn:
|
|
placeholders = ",".join(["?"] * len(tags))
|
|
if mode == "and":
|
|
where_clause = (
|
|
"prompts.id IN ("
|
|
" SELECT pt.prompt_id FROM prompt_tags pt"
|
|
" JOIN tags t ON pt.tag_id = t.id"
|
|
f" WHERE t.name IN ({placeholders})"
|
|
" GROUP BY pt.prompt_id"
|
|
f" HAVING COUNT(DISTINCT t.name) = ?"
|
|
")"
|
|
)
|
|
tag_params = list(tags) + [len(tags)]
|
|
else:
|
|
where_clause = (
|
|
"prompts.id IN ("
|
|
" SELECT pt.prompt_id FROM prompt_tags pt"
|
|
" JOIN tags t ON pt.tag_id = t.id"
|
|
f" WHERE t.name IN ({placeholders})"
|
|
")"
|
|
)
|
|
tag_params = list(tags)
|
|
|
|
cursor = conn.execute(
|
|
f"SELECT COUNT(*) FROM prompts WHERE {where_clause}",
|
|
tag_params,
|
|
)
|
|
row = cursor.fetchone()
|
|
total = (row[0] if row else 0) or 0
|
|
|
|
cursor = conn.execute(
|
|
f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts WHERE {where_clause} "
|
|
"ORDER BY created_at DESC LIMIT ? OFFSET ?",
|
|
tag_params + [limit, offset],
|
|
)
|
|
rows = cursor.fetchall()
|
|
prompts = [self._row_to_dict(row) for row in rows]
|
|
|
|
if prompts:
|
|
prompt_ids = [p["id"] for p in prompts]
|
|
self._attach_preview_images(conn, prompts, prompt_ids)
|
|
|
|
return {
|
|
"prompts": prompts,
|
|
"total": total,
|
|
"limit": limit,
|
|
"offset": offset,
|
|
"has_more": (offset + limit) < total,
|
|
}
|
|
|
|
@staticmethod
|
|
def _move_tag_links(
|
|
conn: sqlite3.Connection, source_id: int, target_id: int
|
|
) -> None:
|
|
"""Repoint component and pair links from source to target.
|
|
|
|
Owners that already carry the target keep that one link; the source's
|
|
leftover duplicates are removed. The source registry row is left for the
|
|
caller to delete.
|
|
"""
|
|
for junction in ("prompt_tags", "prompt_pair_tags"):
|
|
conn.execute(
|
|
f"UPDATE OR IGNORE {junction} SET tag_id = ? WHERE tag_id = ?",
|
|
(target_id, source_id),
|
|
)
|
|
conn.execute(f"DELETE FROM {junction} WHERE tag_id = ?", (source_id,))
|
|
|
|
def rename_tag_all_prompts(self, old_name: str, new_name: str) -> Dict[str, Any]:
|
|
"""
|
|
Rename a tag across all prompts that use it.
|
|
|
|
With junction tables this is a single UPDATE or a merge operation.
|
|
|
|
Args:
|
|
old_name: Current tag name
|
|
new_name: New tag name
|
|
|
|
Returns:
|
|
Dict with success status and affected_count
|
|
"""
|
|
if not old_name or not old_name.strip():
|
|
raise ValueError("Old tag name cannot be empty")
|
|
if not new_name or not new_name.strip():
|
|
raise ValueError("New tag name cannot be empty")
|
|
|
|
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,))
|
|
old_tag = cursor.fetchone()
|
|
if not old_tag:
|
|
return {"success": True, "affected_count": 0, "skipped_count": 0}
|
|
|
|
old_tag_id = old_tag["id"]
|
|
|
|
# Count affected prompts before the operation
|
|
cursor = conn.execute(
|
|
"SELECT COUNT(*) as c FROM prompt_tags WHERE tag_id = ?", (old_tag_id,)
|
|
)
|
|
affected = cursor.fetchone()["c"]
|
|
|
|
cursor = conn.execute("SELECT id FROM tags WHERE name = ?", (new_name,))
|
|
existing_new = cursor.fetchone()
|
|
|
|
if existing_new:
|
|
# Target tag exists — merge: move links, handle conflicts, delete old
|
|
self._move_tag_links(conn, old_tag_id, existing_new["id"])
|
|
conn.execute("DELETE FROM tags WHERE id = ?", (old_tag_id,))
|
|
else:
|
|
# Simple rename
|
|
conn.execute(
|
|
"UPDATE tags SET name = ? WHERE id = ?", (new_name, old_tag_id)
|
|
)
|
|
|
|
conn.commit()
|
|
|
|
self.logger.info(
|
|
f"Renamed tag '{old_name}' -> '{new_name}' in {affected} prompts"
|
|
)
|
|
return {"success": True, "affected_count": affected, "skipped_count": 0}
|
|
|
|
def delete_tag_all_prompts(self, tag_name: str) -> Dict[str, Any]:
|
|
"""
|
|
Remove a tag from all prompts that use it.
|
|
|
|
Deleting from tags cascades to standalone and pair tag relationships.
|
|
|
|
Args:
|
|
tag_name: Tag name to remove
|
|
|
|
Returns:
|
|
Dict with success status and affected_count
|
|
"""
|
|
if not tag_name or not tag_name.strip():
|
|
raise ValueError("Tag name cannot be empty")
|
|
|
|
tag_name = tag_name.strip()
|
|
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute("SELECT id FROM tags WHERE name = ?", (tag_name,))
|
|
tag_row = cursor.fetchone()
|
|
if not tag_row:
|
|
return {"success": True, "affected_count": 0, "skipped_count": 0}
|
|
|
|
tag_id = tag_row["id"]
|
|
cursor = conn.execute(
|
|
"SELECT COUNT(*) as c FROM prompt_tags WHERE tag_id = ?", (tag_id,)
|
|
)
|
|
affected = cursor.fetchone()["c"]
|
|
|
|
# CASCADE delete handles prompt_tags entries
|
|
conn.execute("DELETE FROM tags WHERE id = ?", (tag_id,))
|
|
conn.commit()
|
|
|
|
self.logger.info(f"Deleted tag '{tag_name}' from {affected} prompts")
|
|
return {"success": True, "affected_count": affected, "skipped_count": 0}
|
|
|
|
def merge_tags(self, source_tags: List[str], target_tag: str) -> Dict[str, Any]:
|
|
"""
|
|
Merge one or more source tags into a target tag across all prompts.
|
|
|
|
With junction tables, this is UPDATE + DELETE per source tag.
|
|
|
|
Args:
|
|
source_tags: List of tag names to merge from
|
|
target_tag: Tag name to merge into
|
|
|
|
Returns:
|
|
Dict with success status, affected_count, and tags_merged
|
|
"""
|
|
if not source_tags:
|
|
raise ValueError("Source tags list cannot be empty")
|
|
if not target_tag or not target_tag.strip():
|
|
raise ValueError("Target tag cannot be empty")
|
|
|
|
target_tag = target_tag.strip()
|
|
source_tags = [t.strip() for t in source_tags if t.strip()]
|
|
affected = 0
|
|
tags_merged = 0
|
|
|
|
with self.model.get_connection() as conn:
|
|
# Ensure target tag exists
|
|
conn.execute("INSERT OR IGNORE INTO tags (name) VALUES (?)", (target_tag,))
|
|
cursor = conn.execute("SELECT id FROM tags WHERE name = ?", (target_tag,))
|
|
target_id = cursor.fetchone()["id"]
|
|
|
|
for src_tag in source_tags:
|
|
cursor = conn.execute("SELECT id FROM tags WHERE name = ?", (src_tag,))
|
|
src_row = cursor.fetchone()
|
|
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"]
|
|
|
|
pair_count = conn.execute(
|
|
"SELECT COUNT(*) FROM prompt_pair_tags WHERE tag_id = ?", (src_id,)
|
|
).fetchone()[0]
|
|
if src_count or pair_count:
|
|
self._move_tag_links(conn, src_id, target_id)
|
|
tags_merged += 1
|
|
affected += src_count
|
|
|
|
# Delete the source tag
|
|
conn.execute("DELETE FROM tags WHERE id = ?", (src_id,))
|
|
|
|
conn.commit()
|
|
|
|
self.logger.info(
|
|
f"Merged {tags_merged} tags into '{target_tag}', affected {affected} "
|
|
"prompts"
|
|
)
|
|
return {
|
|
"success": True,
|
|
"affected_count": affected,
|
|
"tags_merged": tags_merged,
|
|
"skipped_count": 0,
|
|
}
|
|
|
|
def get_untagged_prompts_count(self) -> int:
|
|
"""
|
|
Get the count of prompts that have no tags.
|
|
|
|
Returns:
|
|
Number of untagged prompts
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
"SELECT COUNT(*) as total FROM prompts "
|
|
"WHERE NOT EXISTS (SELECT 1 FROM prompt_tags WHERE prompt_id = "
|
|
"prompts.id)"
|
|
)
|
|
return cursor.fetchone()["total"]
|
|
|
|
def get_untagged_prompts(self, limit: int = 20, offset: int = 0) -> Dict[str, Any]:
|
|
"""
|
|
Get prompts that have no tags, with pagination.
|
|
|
|
Uses batched image fetch to avoid N+1 queries.
|
|
|
|
Args:
|
|
limit: Maximum number of prompts to return
|
|
offset: Number of prompts to skip
|
|
|
|
Returns:
|
|
Dict with prompts list, total count, and pagination
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
where = (
|
|
"NOT EXISTS (SELECT 1 FROM prompt_tags WHERE prompt_id = prompts.id)"
|
|
)
|
|
|
|
cursor = conn.execute(
|
|
f"SELECT COUNT(*) as total FROM prompts WHERE {where}"
|
|
)
|
|
total = cursor.fetchone()["total"]
|
|
|
|
cursor = conn.execute(
|
|
f"SELECT prompts.*, {TAG_SUBQUERY} FROM prompts WHERE {where} "
|
|
"ORDER BY created_at DESC LIMIT ? OFFSET ?",
|
|
(limit, offset),
|
|
)
|
|
rows = cursor.fetchall()
|
|
prompts = [self._row_to_dict(row) for row in rows]
|
|
|
|
if prompts:
|
|
prompt_ids = [p["id"] for p in prompts]
|
|
self._attach_preview_images(conn, prompts, prompt_ids)
|
|
|
|
return {
|
|
"prompts": prompts,
|
|
"total": total,
|
|
"limit": limit,
|
|
"offset": offset,
|
|
"has_more": (offset + limit) < total,
|
|
}
|
|
|
|
def _attach_preview_images(
|
|
self,
|
|
conn: sqlite3.Connection,
|
|
prompts: List[Dict[str, Any]],
|
|
prompt_ids: List[int],
|
|
) -> None:
|
|
"""Batch-fetch up to 3 preview images + counts for a list of prompts."""
|
|
id_placeholders = ",".join(["?"] * len(prompt_ids))
|
|
|
|
img_cursor = conn.execute(
|
|
f"""SELECT * FROM (
|
|
SELECT id, prompt_id, image_path, filename,
|
|
generation_time, width, height, format,
|
|
ROW_NUMBER() OVER (
|
|
PARTITION BY prompt_id ORDER BY generation_time DESC
|
|
) AS rn
|
|
FROM generated_images
|
|
WHERE prompt_id IN ({id_placeholders})
|
|
) WHERE rn <= 3""",
|
|
prompt_ids,
|
|
)
|
|
images_by_prompt: Dict[int, list] = {}
|
|
for img_row in img_cursor.fetchall():
|
|
pid = img_row["prompt_id"]
|
|
images_by_prompt.setdefault(pid, []).append(dict(img_row))
|
|
|
|
cnt_cursor = conn.execute(
|
|
f"SELECT prompt_id, COUNT(*) as cnt FROM generated_images "
|
|
f"WHERE prompt_id IN ({id_placeholders}) GROUP BY prompt_id",
|
|
prompt_ids,
|
|
)
|
|
counts_by_prompt = {
|
|
row["prompt_id"]: row["cnt"] for row in cnt_cursor.fetchall()
|
|
}
|
|
|
|
for prompt in prompts:
|
|
pid = prompt["id"]
|
|
prompt["images"] = images_by_prompt.get(pid, [])
|
|
prompt["image_count"] = counts_by_prompt.get(pid, 0)
|
|
|
|
def _ensure_tags(
|
|
self, conn: sqlite3.Connection, tag_names: List[str]
|
|
) -> Dict[str, int]:
|
|
"""Ensure tag names exist in tags table, return name->id mapping."""
|
|
if not tag_names:
|
|
return {}
|
|
for name in tag_names:
|
|
conn.execute("INSERT OR IGNORE INTO tags (name) VALUES (?)", (name,))
|
|
placeholders = ",".join(["?"] * len(tag_names))
|
|
cursor = conn.execute(
|
|
f"SELECT id, name FROM tags WHERE name IN ({placeholders})",
|
|
list(tag_names),
|
|
)
|
|
return {row["name"]: row["id"] for row in cursor.fetchall()}
|
|
|
|
def _sync_prompt_tags(
|
|
self, conn: sqlite3.Connection, prompt_id: int, tags: List[str]
|
|
) -> None:
|
|
"""Replace all junction table entries for a prompt."""
|
|
conn.execute("DELETE FROM prompt_tags WHERE prompt_id = ?", (prompt_id,))
|
|
if tags:
|
|
tag_map = self._ensure_tags(conn, tags)
|
|
for tag_name in tags:
|
|
tag_id = tag_map.get(tag_name)
|
|
if tag_id:
|
|
conn.execute(
|
|
"INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) "
|
|
"VALUES (?, ?)",
|
|
(prompt_id, tag_id),
|
|
)
|
|
|
|
def _get_prompt_tags(self, conn: sqlite3.Connection, prompt_id: int) -> List[str]:
|
|
"""Get tag names for a prompt from junction table."""
|
|
cursor = conn.execute(
|
|
"SELECT t.name FROM prompt_tags pt "
|
|
"JOIN tags t ON pt.tag_id = t.id "
|
|
"WHERE pt.prompt_id = ?",
|
|
(prompt_id,),
|
|
)
|
|
return [row["name"] for row in cursor.fetchall()]
|
|
|
|
def _row_to_dict(self, row: sqlite3.Row) -> Dict[str, Any]:
|
|
"""
|
|
Convert a database row to a dictionary with parsed JSON fields.
|
|
|
|
Args:
|
|
row: SQLite row object
|
|
|
|
Returns:
|
|
Dictionary representation of the row
|
|
"""
|
|
data = dict(row)
|
|
|
|
# Parse tags from junction table subquery result (preferred)
|
|
if "_tag_list" in data and data["_tag_list"]:
|
|
data["tags"] = data["_tag_list"].split("|||")
|
|
del data["_tag_list"]
|
|
elif "_tag_list" in data:
|
|
# TAG_SUBQUERY present but NULL (no tags)
|
|
data["tags"] = []
|
|
del data["_tag_list"]
|
|
elif data.get("tags"):
|
|
# Fallback: parse legacy JSON column
|
|
try:
|
|
parsed = json.loads(data["tags"])
|
|
if isinstance(parsed, list):
|
|
data["tags"] = parsed
|
|
elif isinstance(parsed, str):
|
|
data["tags"] = [t.strip() for t in parsed.split(",") if t.strip()]
|
|
else:
|
|
data["tags"] = []
|
|
except (json.JSONDecodeError, TypeError):
|
|
data["tags"] = []
|
|
else:
|
|
data["tags"] = []
|
|
|
|
return data
|
|
|
|
def export_prompts(self, file_path: str, format: str = "json") -> bool:
|
|
"""
|
|
Export all prompts to a file.
|
|
|
|
Args:
|
|
file_path: Path to save the export file
|
|
format: Export format ("json" or "csv")
|
|
|
|
Returns:
|
|
bool: True if export was successful, False otherwise
|
|
"""
|
|
try:
|
|
prompts = self.search_prompts(limit=10000) # Get all prompts
|
|
|
|
if format.lower() == "json":
|
|
with open(file_path, "w", encoding="utf-8") as f:
|
|
json.dump(prompts, f, indent=2, ensure_ascii=False)
|
|
elif format.lower() == "csv":
|
|
import csv
|
|
|
|
if prompts:
|
|
with open(file_path, "w", newline="", encoding="utf-8") as f:
|
|
writer = csv.DictWriter(f, fieldnames=prompts[0].keys())
|
|
writer.writeheader()
|
|
for prompt in prompts:
|
|
# Convert lists to strings for CSV
|
|
row = prompt.copy()
|
|
if isinstance(row.get("tags"), list):
|
|
row["tags"] = ", ".join(row["tags"])
|
|
writer.writerow(row)
|
|
else:
|
|
raise ValueError(f"Unsupported export format: {format}")
|
|
|
|
return True
|
|
except Exception as e:
|
|
self.logger.error(f"Error exporting prompts: {e}")
|
|
return False
|
|
|
|
def cleanup_duplicates(self) -> int:
|
|
"""
|
|
Remove duplicate prompts based on text content, preserving all image links.
|
|
Merges metadata and transfers images to the retained prompt.
|
|
|
|
Returns:
|
|
int: Number of duplicates removed
|
|
"""
|
|
self.logger.info("Starting duplicate cleanup process with image preservation")
|
|
try:
|
|
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
|
|
cursor = conn.execute("""
|
|
SELECT LOWER(TRIM(text)) as normalized_text, COUNT(*) as count,
|
|
GROUP_CONCAT(id) as ids,
|
|
GROUP_CONCAT(created_at) as created_dates
|
|
FROM prompts
|
|
GROUP BY LOWER(TRIM(text))
|
|
HAVING COUNT(*) > 1
|
|
""")
|
|
|
|
duplicates = cursor.fetchall()
|
|
self.logger.debug(
|
|
f"Found {len(duplicates)} groups of duplicate prompts"
|
|
)
|
|
total_removed = 0
|
|
total_images_transferred = 0
|
|
|
|
for duplicate in duplicates:
|
|
ids = duplicate["ids"].split(",")
|
|
created_dates = duplicate["created_dates"].split(",")
|
|
|
|
# Sort IDs by created_at date to keep the oldest
|
|
id_date_pairs = list(zip(ids, created_dates))
|
|
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:]
|
|
|
|
self.logger.debug(
|
|
f"Merging duplicates: keeping {primary_id}, removing "
|
|
f"{duplicate_ids}"
|
|
)
|
|
|
|
# Get primary prompt details
|
|
cursor = conn.execute(
|
|
"SELECT * FROM prompts WHERE id = ?", (primary_id,)
|
|
)
|
|
primary_prompt = cursor.fetchone()
|
|
if not primary_prompt:
|
|
continue
|
|
|
|
# Collect and merge metadata from all duplicates
|
|
merged_metadata = self._merge_duplicate_metadata(
|
|
conn, primary_id, duplicate_ids
|
|
)
|
|
|
|
# Transfer all images from duplicates to primary prompt
|
|
images_transferred = self._transfer_images_to_primary(
|
|
conn, primary_id, duplicate_ids
|
|
)
|
|
total_images_transferred += images_transferred
|
|
|
|
# Update primary prompt with merged metadata
|
|
if merged_metadata:
|
|
self._update_primary_with_merged_metadata(
|
|
conn, primary_id, merged_metadata
|
|
)
|
|
|
|
# Delete duplicate prompts (images already transferred)
|
|
for duplicate_id in duplicate_ids:
|
|
conn.execute(
|
|
"DELETE FROM prompts WHERE id = ?", (duplicate_id,)
|
|
)
|
|
total_removed += 1
|
|
self.logger.debug(f"Removed duplicate prompt {duplicate_id}")
|
|
|
|
conn.commit()
|
|
|
|
if total_removed > 0:
|
|
self.logger.info(
|
|
f"Removed {total_removed} duplicate prompts, transferred "
|
|
f"{total_images_transferred} images"
|
|
)
|
|
else:
|
|
self.logger.info("No duplicate prompts found")
|
|
|
|
return total_removed
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error cleaning up duplicates: {e}", exc_info=True)
|
|
return 0
|
|
|
|
# Gallery-related methods
|
|
def link_image_to_prompt(
|
|
self, prompt_id: str, image_path: str, metadata: Optional[Dict[str, Any]] = None
|
|
) -> int:
|
|
"""
|
|
Link a generated image to a prompt.
|
|
|
|
Args:
|
|
prompt_id: ID of the prompt that generated this image
|
|
image_path: Full path to the image file
|
|
metadata: Optional metadata about the image
|
|
|
|
Returns:
|
|
int: The ID of the created image record
|
|
"""
|
|
# Validate that prompt_id is a valid integer and exists in prompts table
|
|
try:
|
|
# Convert prompt_id to integer if it's a string number
|
|
if isinstance(prompt_id, str) and prompt_id.isdigit():
|
|
prompt_id_int = int(prompt_id)
|
|
elif isinstance(prompt_id, int):
|
|
prompt_id_int = prompt_id
|
|
else:
|
|
# Handle temporary IDs like "temp_123456" - skip linking
|
|
if isinstance(prompt_id, str) and prompt_id.startswith("temp_"):
|
|
self.logger.debug(
|
|
f"Skipping image linking for temporary prompt ID: {prompt_id}"
|
|
)
|
|
return 0
|
|
else:
|
|
self.logger.warning(
|
|
f"Invalid prompt_id format: {prompt_id}, skipping image linking"
|
|
)
|
|
return 0
|
|
|
|
# Verify the prompt exists in the database
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
"SELECT id FROM prompts WHERE id = ?", (prompt_id_int,)
|
|
)
|
|
if not cursor.fetchone():
|
|
self.logger.warning(
|
|
f"Prompt ID {prompt_id_int} not found in database, skipping "
|
|
"image linking"
|
|
)
|
|
return 0
|
|
|
|
# Proceed with linking
|
|
if not isinstance(metadata, dict):
|
|
metadata = {}
|
|
filename = os.path.basename(image_path)
|
|
file_path = normalize_image_path(image_path)
|
|
file_size, width, height, image_format = _coerce_file_info(metadata)
|
|
|
|
# INSERT OR IGNORE skips a file already linked to this prompt
|
|
# (same prompt_id + normalised full path)
|
|
cursor = conn.execute(
|
|
"""
|
|
INSERT OR IGNORE INTO generated_images
|
|
(prompt_id, image_path, filename, file_path, file_size, width,
|
|
height, format, workflow_data, prompt_metadata, parameters)
|
|
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
(
|
|
prompt_id_int,
|
|
image_path,
|
|
filename,
|
|
file_path,
|
|
file_size,
|
|
width,
|
|
height,
|
|
image_format,
|
|
json.dumps(metadata.get("workflow", {})),
|
|
json.dumps(metadata.get("prompt", {})),
|
|
json.dumps(metadata.get("parameters", {})),
|
|
),
|
|
)
|
|
conn.commit()
|
|
|
|
if cursor.rowcount == 0:
|
|
self.logger.debug(
|
|
f"Image {file_path} already linked to prompt {prompt_id_int}"
|
|
)
|
|
return 0
|
|
|
|
return cursor.lastrowid
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error linking image to prompt {prompt_id}: {e}")
|
|
return 0
|
|
|
|
def get_prompt_images(self, prompt_id: str) -> List[Dict[str, Any]]:
|
|
"""
|
|
Get all images associated with a prompt.
|
|
|
|
Args:
|
|
prompt_id: The prompt ID
|
|
|
|
Returns:
|
|
List of image records
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
"""
|
|
SELECT * FROM generated_images
|
|
WHERE prompt_id = ?
|
|
ORDER BY generation_time DESC
|
|
""",
|
|
(prompt_id,),
|
|
)
|
|
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
|
|
|
|
def get_recent_images(
|
|
self, limit: int = 50, offset: int = 0
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Get recently generated images across all prompts, newest first.
|
|
|
|
Args:
|
|
limit: Maximum number of images to return
|
|
offset: Number of newest images to skip (paging)
|
|
|
|
Returns:
|
|
List of image records with prompt text
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
"""
|
|
SELECT gi.*, p.text as prompt_text
|
|
FROM generated_images gi
|
|
LEFT JOIN prompts p ON gi.prompt_id = p.id
|
|
ORDER BY gi.generation_time DESC
|
|
LIMIT ? OFFSET ?
|
|
""",
|
|
(int(limit), max(0, int(offset))),
|
|
)
|
|
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
|
|
|
|
def get_all_images(
|
|
self,
|
|
limit: int = 0,
|
|
offset: int = 0,
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Get generated images with their linked prompts.
|
|
|
|
Args:
|
|
limit: Maximum number of images to return (0 = all).
|
|
offset: Number of images to skip.
|
|
|
|
Returns:
|
|
List of image records with prompt text and tags
|
|
"""
|
|
sql = (
|
|
"SELECT gi.*, p.text as prompt_text, "
|
|
"(SELECT GROUP_CONCAT(t.name, '|||') FROM prompt_tags pt "
|
|
"JOIN tags t ON pt.tag_id = t.id WHERE pt.prompt_id = p.id) AS "
|
|
"_prompt_tags_list "
|
|
"FROM generated_images gi "
|
|
"INNER JOIN prompts p ON gi.prompt_id = p.id "
|
|
"WHERE gi.image_path IS NOT NULL AND gi.image_path != '' "
|
|
"ORDER BY gi.generation_time DESC"
|
|
)
|
|
params: list = []
|
|
if limit > 0:
|
|
sql += " LIMIT ? OFFSET ?"
|
|
params = [limit, offset]
|
|
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(sql, params)
|
|
rows = cursor.fetchall()
|
|
result = []
|
|
for row in rows:
|
|
data = self._image_row_to_dict(row)
|
|
if row["_prompt_tags_list"]:
|
|
data["prompt_tags"] = row["_prompt_tags_list"].split("|||")
|
|
else:
|
|
data["prompt_tags"] = []
|
|
result.append(data)
|
|
return result
|
|
|
|
def search_images_by_prompt(
|
|
self, search_term: str, limit: Optional[int] = None, offset: int = 0
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Search images by prompt text, newest first.
|
|
|
|
Args:
|
|
search_term: Text to search for in prompt content
|
|
limit: Maximum number of images to return (None for no limit)
|
|
offset: Number of newest matches to skip (paging)
|
|
|
|
Returns:
|
|
List of image records with prompt text
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
"""
|
|
SELECT gi.*, p.text as prompt_text
|
|
FROM generated_images gi
|
|
JOIN prompts p ON gi.prompt_id = p.id
|
|
WHERE p.text LIKE ? ESCAPE '\\'
|
|
ORDER BY gi.generation_time DESC
|
|
LIMIT ? OFFSET ?
|
|
""",
|
|
(
|
|
f"%{escape_like(search_term)}%",
|
|
-1 if limit is None else int(limit),
|
|
max(0, int(offset)),
|
|
),
|
|
)
|
|
return [self._image_row_to_dict(row) for row in cursor.fetchall()]
|
|
|
|
def get_image_by_id(self, image_id: int) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Get an image record by its ID.
|
|
|
|
Args:
|
|
image_id: The image ID
|
|
|
|
Returns:
|
|
Image record or None if not found
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
"SELECT * FROM generated_images WHERE id = ?", (image_id,)
|
|
)
|
|
row = cursor.fetchone()
|
|
return self._image_row_to_dict(row) if row else None
|
|
|
|
def delete_image(self, image_id: int) -> bool:
|
|
"""
|
|
Delete an image record by its ID.
|
|
|
|
Args:
|
|
image_id: The image ID to delete
|
|
|
|
Returns:
|
|
bool: True if deletion was successful
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
"DELETE FROM generated_images WHERE id = ?", (image_id,)
|
|
)
|
|
conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
def cleanup_missing_images(self) -> int:
|
|
"""
|
|
Remove image records where the actual file no longer exists.
|
|
|
|
Returns:
|
|
int: Number of orphaned records removed
|
|
"""
|
|
removed_count = 0
|
|
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute("SELECT id, image_path FROM generated_images")
|
|
images = cursor.fetchall()
|
|
|
|
for image in images:
|
|
if not os.path.exists(image["image_path"]):
|
|
conn.execute(
|
|
"DELETE FROM generated_images WHERE id = ?", (image["id"],)
|
|
)
|
|
removed_count += 1
|
|
|
|
conn.commit()
|
|
|
|
return removed_count
|
|
|
|
def _clean_nan_values(self, obj: Any) -> Any:
|
|
"""
|
|
Recursively clean NaN values from nested data structures.
|
|
|
|
Args:
|
|
obj: The object to clean (dict, list, or scalar)
|
|
|
|
Returns:
|
|
Cleaned object with NaN values replaced by None
|
|
"""
|
|
if isinstance(obj, dict):
|
|
return {key: self._clean_nan_values(value) for key, value in obj.items()}
|
|
elif isinstance(obj, list):
|
|
return [self._clean_nan_values(item) for item in obj]
|
|
elif isinstance(obj, float) and str(obj) == "nan":
|
|
return None
|
|
else:
|
|
return obj
|
|
|
|
def _merge_duplicate_metadata(
|
|
self, conn: sqlite3.Connection, primary_id: int, duplicate_ids: List[int]
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Merge metadata from duplicate prompts, prioritizing non-empty values.
|
|
|
|
Args:
|
|
conn: Database connection
|
|
primary_id: ID of the prompt to keep
|
|
duplicate_ids: List of duplicate prompt IDs
|
|
|
|
Returns:
|
|
Dict containing merged metadata
|
|
"""
|
|
try:
|
|
cursor = conn.execute(
|
|
"SELECT category, rating, notes FROM prompts WHERE id = ?",
|
|
(primary_id,),
|
|
)
|
|
primary_data = cursor.fetchone()
|
|
if not primary_data:
|
|
return {}
|
|
|
|
merged = {
|
|
"category": primary_data["category"],
|
|
"tags": self._get_prompt_tags(conn, primary_id),
|
|
"rating": primary_data["rating"],
|
|
"notes": primary_data["notes"] or "",
|
|
}
|
|
|
|
for dup_id in duplicate_ids:
|
|
cursor = conn.execute(
|
|
"SELECT category, rating, notes FROM prompts WHERE id = ?",
|
|
(dup_id,),
|
|
)
|
|
dup_data = cursor.fetchone()
|
|
if not dup_data:
|
|
continue
|
|
|
|
if not merged["category"] and dup_data["category"]:
|
|
merged["category"] = dup_data["category"]
|
|
|
|
dup_tags = self._get_prompt_tags(conn, dup_id)
|
|
for tag in dup_tags:
|
|
if tag not in merged["tags"]:
|
|
merged["tags"].append(tag)
|
|
|
|
if dup_data["rating"] and (
|
|
not merged["rating"] or dup_data["rating"] > merged["rating"]
|
|
):
|
|
merged["rating"] = dup_data["rating"]
|
|
|
|
if dup_data["notes"] and dup_data["notes"].strip():
|
|
if merged["notes"]:
|
|
merged["notes"] += f" | {dup_data['notes']}"
|
|
else:
|
|
merged["notes"] = dup_data["notes"]
|
|
|
|
return merged
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error merging metadata: {e}", exc_info=True)
|
|
return {}
|
|
|
|
def _transfer_images_to_primary(
|
|
self, conn: sqlite3.Connection, primary_id: int, duplicate_ids: List[int]
|
|
) -> int:
|
|
"""
|
|
Transfer all images from duplicate prompts to the primary prompt.
|
|
|
|
Args:
|
|
conn: Database connection
|
|
primary_id: ID of the prompt to keep
|
|
duplicate_ids: List of duplicate prompt IDs
|
|
|
|
Returns:
|
|
Number of images transferred
|
|
"""
|
|
transferred_count = 0
|
|
try:
|
|
for dup_id in duplicate_ids:
|
|
# Update all images to point to primary prompt
|
|
cursor = conn.execute(
|
|
"UPDATE generated_images SET prompt_id = ? WHERE prompt_id = ?",
|
|
(primary_id, dup_id),
|
|
)
|
|
transferred_count += cursor.rowcount
|
|
|
|
if cursor.rowcount > 0:
|
|
self.logger.debug(
|
|
f"Transferred {cursor.rowcount} images from prompt {dup_id} "
|
|
f"to {primary_id}"
|
|
)
|
|
|
|
return transferred_count
|
|
|
|
except Exception as e:
|
|
self.logger.error(f"Error transferring images: {e}", exc_info=True)
|
|
return 0
|
|
|
|
def _update_primary_with_merged_metadata(
|
|
self, conn: sqlite3.Connection, primary_id: int, merged_metadata: Dict[str, Any]
|
|
) -> None:
|
|
"""
|
|
Update the primary prompt with merged metadata.
|
|
|
|
Args:
|
|
conn: Database connection
|
|
primary_id: ID of the prompt to update
|
|
merged_metadata: Merged metadata dictionary
|
|
"""
|
|
try:
|
|
conn.execute(
|
|
"""
|
|
UPDATE prompts
|
|
SET category = ?, rating = ?, notes = ?, updated_at = ?
|
|
WHERE id = ?
|
|
""",
|
|
(
|
|
merged_metadata.get("category"),
|
|
merged_metadata.get("rating"),
|
|
merged_metadata.get("notes"),
|
|
datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
|
primary_id,
|
|
),
|
|
)
|
|
self._sync_prompt_tags(conn, primary_id, merged_metadata.get("tags", []))
|
|
self.logger.debug(
|
|
f"Updated primary prompt {primary_id} with merged metadata"
|
|
)
|
|
|
|
except Exception as e:
|
|
self.logger.error(
|
|
f"Error updating primary prompt metadata: {e}", exc_info=True
|
|
)
|
|
|
|
# ------------------------------------------------------------------
|
|
# Methods extracted from api.py raw SQL (Fix 3.1)
|
|
# ------------------------------------------------------------------
|
|
|
|
def get_statistics(self) -> Dict[str, Any]:
|
|
"""Get comprehensive database statistics."""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute("SELECT COUNT(*) FROM prompts")
|
|
row = cursor.fetchone()
|
|
total_prompts = (row[0] if row else 0) or 0
|
|
|
|
cursor = conn.execute(
|
|
"SELECT COUNT(DISTINCT TRIM(category)) FROM prompts "
|
|
"WHERE category IS NOT NULL AND TRIM(category) != ''"
|
|
)
|
|
row = cursor.fetchone()
|
|
total_categories = (row[0] if row else 0) or 0
|
|
|
|
cursor = conn.execute(
|
|
"SELECT AVG(rating) FROM prompts WHERE rating IS NOT NULL"
|
|
)
|
|
row = cursor.fetchone()
|
|
avg_rating = row[0] if row else None
|
|
|
|
cursor = conn.execute(
|
|
"SELECT COUNT(*) FROM tags t "
|
|
"WHERE EXISTS (SELECT 1 FROM prompt_tags pt WHERE pt.tag_id = t.id)"
|
|
)
|
|
row = cursor.fetchone()
|
|
total_tags = (row[0] if row else 0) or 0
|
|
|
|
cursor = conn.execute("SELECT COUNT(*) FROM generated_images")
|
|
row = cursor.fetchone()
|
|
total_images = (row[0] if row else 0) or 0
|
|
|
|
cursor = conn.execute(
|
|
"SELECT COUNT(DISTINCT prompt_id) FROM generated_images"
|
|
)
|
|
row = cursor.fetchone()
|
|
images_with_prompts = (row[0] if row else 0) or 0
|
|
|
|
avg = round(avg_rating, 2) if avg_rating else None
|
|
return {
|
|
"total_prompts": total_prompts,
|
|
"total_categories": total_categories,
|
|
"average_rating": avg,
|
|
"avg_rating": avg,
|
|
"total_tags": total_tags,
|
|
"total_images": total_images,
|
|
"images_with_prompts": images_with_prompts,
|
|
}
|
|
|
|
def update_prompt_text(self, prompt_id: int, new_text: str) -> bool:
|
|
"""Update prompt text and its deduplication hash atomically.
|
|
|
|
Returns:
|
|
True if a row was updated.
|
|
|
|
Raises:
|
|
sqlite3.IntegrityError: Another prompt already has the new hash.
|
|
"""
|
|
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 = ?",
|
|
(
|
|
new_text,
|
|
generate_prompt_hash(new_text),
|
|
datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
|
prompt_id,
|
|
),
|
|
)
|
|
conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
def update_prompt_rating(self, prompt_id: int, rating: int) -> bool:
|
|
"""Update only the rating of a prompt.
|
|
|
|
Returns:
|
|
True if a row was updated.
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
"UPDATE prompts SET rating = ?, updated_at = ? WHERE id = ?",
|
|
(
|
|
rating,
|
|
datetime.datetime.now(datetime.timezone.utc).isoformat(),
|
|
prompt_id,
|
|
),
|
|
)
|
|
conn.commit()
|
|
return cursor.rowcount > 0
|
|
|
|
def set_prompt_tags(self, prompt_id: int, tags: List[str]) -> None:
|
|
"""Overwrite the tags list for a prompt via junction table."""
|
|
with self.model.get_connection() as conn:
|
|
self._sync_prompt_tags(conn, prompt_id, tags)
|
|
conn.execute(
|
|
"UPDATE prompts SET updated_at = ? WHERE id = ?",
|
|
(datetime.datetime.now(datetime.timezone.utc).isoformat(), prompt_id),
|
|
)
|
|
conn.commit()
|
|
|
|
def bulk_delete_prompts(self, prompt_ids: List[int]) -> int:
|
|
"""Delete multiple prompts and their associated images.
|
|
|
|
Returns:
|
|
Number of prompts actually deleted.
|
|
"""
|
|
count = 0
|
|
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,))
|
|
if cursor.rowcount > 0:
|
|
count += 1
|
|
conn.commit()
|
|
return count
|
|
|
|
def bulk_add_tags(self, prompt_ids: List[int], new_tags: List[str]) -> int:
|
|
"""Add tags to multiple prompts via junction table (skipping duplicates).
|
|
|
|
Returns:
|
|
Number of prompts that were actually modified.
|
|
"""
|
|
count = 0
|
|
with self.model.get_connection() as conn:
|
|
tag_map = self._ensure_tags(conn, new_tags)
|
|
now = datetime.datetime.now(datetime.timezone.utc).isoformat()
|
|
for pid in prompt_ids:
|
|
cursor = conn.execute("SELECT id FROM prompts WHERE id = ?", (pid,))
|
|
if not cursor.fetchone():
|
|
continue
|
|
added = False
|
|
for tag_name in new_tags:
|
|
tag_id = tag_map.get(tag_name)
|
|
if tag_id:
|
|
cursor = conn.execute(
|
|
"INSERT OR IGNORE INTO prompt_tags (prompt_id, tag_id) "
|
|
"VALUES (?, ?)",
|
|
(pid, tag_id),
|
|
)
|
|
if cursor.rowcount > 0:
|
|
added = True
|
|
if added:
|
|
conn.execute(
|
|
"UPDATE prompts SET updated_at = ? WHERE id = ?", (now, pid)
|
|
)
|
|
count += 1
|
|
conn.commit()
|
|
return count
|
|
|
|
def bulk_set_category(self, prompt_ids: List[int], category: str) -> int:
|
|
"""Set category on multiple prompts.
|
|
|
|
Returns:
|
|
Number of prompts actually updated.
|
|
"""
|
|
count = 0
|
|
now = datetime.datetime.now(datetime.timezone.utc).isoformat()
|
|
with self.model.get_connection() as conn:
|
|
for pid in prompt_ids:
|
|
cursor = conn.execute(
|
|
"UPDATE prompts SET category = ?, updated_at = ? WHERE id = ?",
|
|
(category, now, pid),
|
|
)
|
|
if cursor.rowcount > 0:
|
|
count += 1
|
|
conn.commit()
|
|
return count
|
|
|
|
def get_image_prompt_info(
|
|
self,
|
|
image_path: str,
|
|
allow_filename_fallback: bool = True,
|
|
*,
|
|
report_ambiguity: bool = False,
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Look up prompt + image metadata for a given image path.
|
|
|
|
Returns a dict with prompt_id, text, category, tags, rating, notes,
|
|
workflow_data, prompt_metadata, generation_time, parameters — or None.
|
|
Opting into report_ambiguity returns only an association_status marker
|
|
for multiple matches; default callers still receive None.
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
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.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 = ?
|
|
OR (? AND gi.image_path LIKE ? ESCAPE '\\')""",
|
|
(
|
|
image_path,
|
|
normalize_image_path(image_path),
|
|
allow_filename_fallback,
|
|
f"%{escape_like(os.path.basename(image_path))}",
|
|
),
|
|
)
|
|
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:
|
|
if report_ambiguity and matches:
|
|
return {"association_status": "ambiguous"}
|
|
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
|
|
except (json.JSONDecodeError, TypeError):
|
|
workflow = None
|
|
try:
|
|
prompt_meta = json.loads(row[6]) if row[6] else None
|
|
except (json.JSONDecodeError, TypeError):
|
|
prompt_meta = None
|
|
return {
|
|
"pair": pair,
|
|
"pair_context": context,
|
|
"prompt_id": row[0],
|
|
"text": row[1],
|
|
"category": row[2],
|
|
"tags": tags,
|
|
"rating": row[3],
|
|
"notes": row[4],
|
|
"workflow_data": workflow,
|
|
"prompt_metadata": prompt_meta,
|
|
"generation_time": row[7],
|
|
"parameters": row[8] if len(row) > 8 else None,
|
|
}
|
|
|
|
def get_image_link_summaries(
|
|
self, paths: List[str], allow_filename_fallback: bool = True
|
|
) -> Dict[str, Dict[str, Any]]:
|
|
"""Summarize each path's linked saved positive for an image card.
|
|
|
|
Matching equals get_image_prompt_info: exact image_path or normalized
|
|
file_path rows win; otherwise, with the fallback, rows whose image_path
|
|
ends with the basename. One match is linked, several are ambiguous.
|
|
One SELECT reads up to LINK_SUMMARY_BATCH paths.
|
|
"""
|
|
unique = list(dict.fromkeys(paths))
|
|
summaries: Dict[str, Dict[str, Any]] = {}
|
|
for start in range(0, len(unique), LINK_SUMMARY_BATCH):
|
|
batch = unique[start : start + LINK_SUMMARY_BATCH]
|
|
summaries.update(self._link_summary_batch(batch, allow_filename_fallback))
|
|
return summaries
|
|
|
|
def _link_summary_batch(
|
|
self, paths: List[str], allow_filename_fallback: bool
|
|
) -> Dict[str, Dict[str, Any]]:
|
|
normalized = {path: normalize_image_path(path) for path in paths}
|
|
basenames = {path: os.path.basename(path) for path in paths}
|
|
marks = ",".join("?" for _ in paths)
|
|
params: List[Any] = [*paths, *normalized.values()]
|
|
fallback = ""
|
|
if allow_filename_fallback:
|
|
fallback = "".join(" OR gi.image_path LIKE ? ESCAPE '\\'" for _ in paths)
|
|
params.extend(f"%{escape_like(basenames[path])}" for path in paths)
|
|
sql = _LINK_SUMMARY_SQL.format(marks=marks, fallback=fallback)
|
|
with self.model.get_connection() as conn:
|
|
rows = conn.execute(sql, params).fetchall()
|
|
summaries = {}
|
|
for path in paths:
|
|
exact = [
|
|
row for row in rows if row[0] == path or row[1] == normalized[path]
|
|
]
|
|
matches = exact or (
|
|
[row for row in rows if _like_suffix(row[0], basenames[path])]
|
|
if allow_filename_fallback
|
|
else []
|
|
)
|
|
if len(matches) == 1:
|
|
summaries[path] = _link_summary(matches[0])
|
|
else:
|
|
summaries[path] = _empty_link_summary(
|
|
"ambiguous" if matches else "none"
|
|
)
|
|
return summaries
|
|
|
|
def get_image_link_matches(
|
|
self,
|
|
q: str,
|
|
category: Optional[str],
|
|
tags: List[str],
|
|
*,
|
|
every_link: bool = False,
|
|
wanted: Optional[Set[str]] = None,
|
|
) -> Dict[str, Tuple[bool, Optional[int], int]]:
|
|
"""Map normalized paths of unambiguously linked images to
|
|
(hit, rating, positive id).
|
|
|
|
Rows are limited to images whose linked positive has *category* and
|
|
every tag in *tags* (exact names). ``hit`` says whether *q* occurs in
|
|
the positive's text, Library title or a tag name, or in the pair title.
|
|
Without category or tags only hits are returned, unless *every_link*
|
|
asks for every linked image (the rating sort and the positive count).
|
|
The positive id is the card's linked positive. One SELECT whose rows
|
|
are streamed; with *wanted* (the listing's normalized paths) only those
|
|
paths are kept, so memory is bounded by the listing, not the database.
|
|
"""
|
|
sql, params = link_filter_query(q, category, tags, every_link=every_link)
|
|
matches: Dict[str, Tuple[bool, Optional[int], int]] = {}
|
|
with self.model.get_connection() as conn:
|
|
for path, rating, matched, positive in conn.execute(sql, params):
|
|
if wanted is None or path in wanted:
|
|
matches[path] = (bool(matched), rating, positive)
|
|
return matches
|
|
|
|
def get_image_paths_matching(
|
|
self, q: str, category: Optional[str], tags: List[str]
|
|
) -> Set[str]:
|
|
"""Normalized paths whose linked positive matches *q*, *category* and *tags*.
|
|
|
|
See get_image_link_matches; an empty *q* matches every linked image
|
|
the category and tags admit.
|
|
"""
|
|
return {
|
|
path
|
|
for path, (hit, _rating, _positive) in self.get_image_link_matches(
|
|
q, category, tags
|
|
).items()
|
|
if hit or not q
|
|
}
|
|
|
|
def get_prompt_id_for_image(self, image_path: str) -> Optional[int]:
|
|
"""Return the prompt_id linked to an image, or None."""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
"SELECT prompt_id FROM generated_images"
|
|
" WHERE image_path = ? OR file_path = ?",
|
|
(image_path, normalize_image_path(image_path)),
|
|
)
|
|
row = cursor.fetchone()
|
|
return row[0] if row else None
|
|
|
|
def check_hash_duplicates(self) -> List[Dict[str, Any]]:
|
|
"""Find prompts sharing the same hash.
|
|
|
|
Returns list of dicts with 'hash' and 'count'.
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(
|
|
"SELECT hash, COUNT(*) as count FROM prompts "
|
|
"WHERE hash IS NOT NULL GROUP BY hash HAVING COUNT(*) > 1"
|
|
)
|
|
return [dict(row) for row in cursor.fetchall()]
|
|
|
|
def prune_orphaned_prompts(self) -> int:
|
|
"""Delete prompts nothing references.
|
|
|
|
A prompt is kept when it has a linked image, belongs to a pair, carries
|
|
the ``__protected__`` tag, has Library parent state (a titled or
|
|
explicit parent), or is the currently valid navigation hint of a
|
|
recorded run: same storage epoch and exact raw snapshot text matching
|
|
only that prompt, the conditions under which history reads resolve it.
|
|
|
|
Returns number of prompts removed.
|
|
"""
|
|
with self._pair_transaction() as conn:
|
|
cursor = conn.execute(ORPHANED_PROMPTS_SQL)
|
|
orphaned = [row["id"] for row in cursor.fetchall()]
|
|
if not orphaned:
|
|
return 0
|
|
placeholders = ",".join(["?"] * len(orphaned))
|
|
cursor = conn.execute(
|
|
f"DELETE FROM prompts WHERE id IN ({placeholders})", orphaned
|
|
)
|
|
conn.commit()
|
|
return cursor.rowcount
|
|
|
|
def check_consistency(self) -> List[str]:
|
|
"""Run consistency checks and return a list of issue descriptions."""
|
|
issues: List[str] = []
|
|
with self.model.get_connection() as conn:
|
|
# Check for orphaned prompt_tags entries
|
|
cursor = conn.execute("""
|
|
SELECT pt.prompt_id, pt.tag_id FROM prompt_tags pt
|
|
LEFT JOIN prompts p ON pt.prompt_id = p.id
|
|
WHERE p.id IS NULL
|
|
""")
|
|
for ref in cursor.fetchall():
|
|
issues.append(
|
|
"prompt_tags entry references non-existent prompt "
|
|
f"{ref['prompt_id']}"
|
|
)
|
|
|
|
# Check for orphaned image entries
|
|
cursor = conn.execute("""
|
|
SELECT gi.id, gi.prompt_id FROM generated_images gi
|
|
LEFT JOIN prompts p ON gi.prompt_id = p.id
|
|
WHERE p.id IS NULL
|
|
""")
|
|
for ref in cursor.fetchall():
|
|
issues.append(
|
|
f"Image {ref['id']} references non-existent prompt "
|
|
f"{ref['prompt_id']}"
|
|
)
|
|
return issues
|
|
|
|
def _image_row_to_dict(self, row: sqlite3.Row) -> Dict[str, Any]:
|
|
"""
|
|
Convert an image database row to a dictionary with parsed JSON fields.
|
|
|
|
Args:
|
|
row: SQLite row object
|
|
|
|
Returns:
|
|
Dictionary representation of the row
|
|
"""
|
|
data = dict(row)
|
|
|
|
# Parse JSON fields
|
|
for field in ["workflow_data", "prompt_metadata", "parameters"]:
|
|
if data.get(field):
|
|
try:
|
|
parsed_data = json.loads(data[field])
|
|
data[field] = self._clean_nan_values(parsed_data)
|
|
except (json.JSONDecodeError, TypeError):
|
|
data[field] = {}
|
|
else:
|
|
data[field] = {}
|
|
|
|
return data
|