* feat(review): add review queue schema with restore validation
Add the Needs review schema group (media_review_metadata,
media_review_items, media_review_decisions) beside the image annotation
group it cites. review_schema_error validates the catalog, columns,
foreign keys, row values and canonical 8 MiB evidence read-only;
migrate_review_state installs it in a SAVEPOINT after the annotation
group, and verify_database_file requires it under require_library_state.
A review group without a valid annotation group is refused. Decisions
null their pair and image navigation through indexes when those are
deleted. Old backups and annotation-era backups restore and gain the
group; browser fixtures clear the new tables on reset.
* feat(review): persist ambiguous sampler evidence from scans and new outputs
The output scan and the live monitor now queue every file the ingestion
resolver finds ambiguous for Needs review. The file is captured, its metadata
is read again from the capture and resolved with the default profile, and only
an ambiguous result is stored, bound to the capture's annotation byte record
and fenced by its storage epoch. Evidence holds one hashed projection per
selectable candidate; none, more than 16 or over 8 MiB is a limited item.
Recording is idempotent per bytes and evidence and never touches a decision.
Scan statuses, counters, outcome rows and legacy SSE events are unchanged; the
job scope gains review_items/review_errors and the log a Needs review line.
The monitor records after its fallback link, best effort.
* feat(review): explicit association decisions with compatible linking
Needs review gains its queue reads and decisions:
- GET /prompt_manager/review/items lists pending items oldest first with
the exact pending total and an HMAC keyset cursor bound to the storage
epoch and the review index revision.
- GET /prompt_manager/review/items/{id} adds the candidates (role texts,
label, selectability), the dependency fingerprint and the digest.
- POST /prompt_manager/review/items/{id}/decision applies a candidate only
for a capture of the item's bytes that the current file still hashes to
and whose evidence is unchanged, through the ingester's compatible-only
rules in the same transaction as the decision row; a conflict writes
nothing. Leave unresolved records the decision and links nothing.
ingest_resolved_pair keeps its contract; its body moves into
_ingest_resolved_pair_in, which runs in the caller's transaction.
Items keep the file's recorded spelling so targets keep their case on
Windows, and a rescan skips the copy and hash of a file whose size and
mtime match the bytes an item already cites.
* fix(review): build the spelling validation test from native paths
On Windows the locator is case-folded and uses backslashes, so the fixed
POSIX pair /out/a.png no longer folded to itself and the new spelling
check refused the good case. The test now spells a mixed-case native path
and folds it with os.path.normcase, and also refuses a locator of another
file.
* fix(review): Apply sees the file under every spelling and rechecks only its choice
Apply now finds the legacy rows of the same physical file however its root
is reached. The roots are real paths while the monitor links under the
directory it watches, so a symlinked or junctioned root, subst or a mapped
drive used to hide the monitor's row and Apply added a second, differently
linked row. Same-name rows are resolved to their physical identity
(normcase(realpath), else samefile); the dependency fingerprint covers all
of them, two rows refuse the choice, one row is enriched under the
ingester's compatible-only rules, and a same-name row the request did not
classify is a changed association.
Apply re-verifies only what the decision depends on: the bytes are still
ambiguous and still offer the chosen candidate. A resolver change that
rewords the reason or reorders candidates leaves pending items applicable.
The resolver version is one constant in prompt_pair_resolver, shared by
ingestion provenance and Needs review.
The transaction's byte check receives the re-hashed digest (or the body's
for Leave unresolved), every Apply refusal branch is tested with zero
writes, and the root listing and realpath calls run on a worker thread.
* fix(review): validate stored review evidence once per startup
Startup used to run the full review validation twice (the early guard and
the install) and the annotation group's row checks inside each, so every
evidence document was parsed and re-encoded twice. The early guard is now
catalog-only (a review group without annotation tables), and the install
validates the group once, trusting the annotation group it has just
validated. The evidence predicate no longer pattern-scans every document in
SQL, because the streamed canonical pass proves it, and the codec is built
once.
Measured on a database with 40,000 items of 4 KiB evidence: startup open
went from 7.7 s to 2.5 s. A test pins one evidence pass per reopen.
Restore and verify_database_file keep the full validation.
* fix(review): keep one capture worker for interactive requests
Needs review recording from the scan and the monitor shared the capture
store's two workers with the Tag dialog and Apply, so a busy scan could
answer a user's preview or decision with "Too many previews are open".
Background captures now never take the last worker or capture slot, and
wait up to two seconds for room before counting the file as not added.
Interactive captures are admitted or refused at once, as before.
* fix(review): refresh the recorded spelling when a file is seen under a new one
A case-only rename on Windows or macOS keeps size and mtime_ns, so the
rescan skipped the file and the item kept the old spelling. Its target was
then formatted from a name the directory no longer lists, the capture
refused it, and the item stayed pending with no way to apply it. The skip
path and a duplicate insert now record the current spelling for every item
at the locator and move the review index revision; item revisions do not
change, so an open decision stays valid. Tested under case-folding
semantics on every platform.
* fix(review): final capture store close and recoverable stale cursors
Closing the capture store at shutdown now runs under the store lock and is
final: a monitor recording or scan batch still in flight gets a capture
error instead of creating a new temp directory and restore subscription
after the cleanup.
A well-formed queue cursor this process cannot verify (opened before a
ComfyUI restart, whose cursor secret is per process) now answers 409
review-conflict with the current total, like any other stale page, instead
of an uncoded 400. Malformed cursors stay 400.
* fix(review): make only the shutdown cleanup final
The previous commit made every close of the capture store final, which
also refused the explicit closes that swap in a fresh store. Only the host's
on_cleanup hook is final now; an observe after it answers a coded capture
error instead of an unhandled 500.
* refactor(admin): move Needs review recording into its own mixin
_record_scanned_review, _record_review_sync, record_review_file and
_log_review_items move unchanged from py/api/admin.py to
py/api/review_recording.py (ReviewRecordingMixin), keeping admin.py from
growing further past the size ceiling. No behaviour change.
* fix(review): record the review fix round and the skip's residual risks in TESTING.md
The Needs review paragraphs now state the rescan skip's residual risks (an
in-place rewrite that keeps size and mtime_ns is never re-queued; a deleted
or moved file leaves its pending item counted), the physical identity rule
for Apply, the candidate-only recheck and shared resolver version, the
spelling refresh, background capture admission, the final shutdown close,
the stale-cursor answer and the measured startup validation cost.
* fix(review): refuse monitor recording beside record-changing jobs
System S1 (#289) gives every record change a registry admission class. The
scan already records Needs review items under its own records job. The
monitor's recording lasts one file and is not a job, so it now checks the
records admission: while a records or guarded job or hold runs it records
nothing, logs that the file waits for the next scan, and answers deferred.
The fallback link stands and the next "new" scan reads the ambiguous file
again, because a fallback link is not a resolved one.
JobRegistry.blocker and active iterate a snapshot of the jobs, since the
monitor's worker thread now asks them too.
* test(backup): wait for the restore, not a 5 s timer, in the thread test
test_database_is_usable_from_every_thread_after_restore let its worker
wait 5 s for the restore. On a stalled Windows runner (py3.11, where the
previous restore test also took 25 s instead of 0.3 s) the restore took
longer, the wait timed out, and the worker read the old database before
the restore ran. The worker now waits for the event under a 120 s bound
that only a hung test reaches, reports a timeout as an error, and the
main thread releases it even when the restore raises. With a restore
delayed by 6 s the old test fails and the new one passes.
* fix(review): store non-finite evidence as null and count pending items
canonical_evidence now follows main's non-finite policy (#303): NaN and
±Infinity are written as null instead of raising, so a NaN-stamped graph
never turns a review item limited. The System Overview reports the review
queue's pending total in needs_review.pending, read in the same snapshot
as the Library counts.
* fix(review): show the pending count on the System Overview review row
The Overview's review row now reads "N ambiguous image links" with the
concept's line, from needs_review.pending. Its Review metadata action
stays aria-disabled until the Needs review page ships. The system specs
and their concept comparison follow the real count.
* fix(system): answer the Overview and Statistics only after every read ended
asyncio.gather answered at the first failed read while its sibling
executor reads were still inside SQLite. In the suite the test teardown
then deleted the database under the running read, and py3.10 on Ubuntu
segfaulted in get_statistics (test_statistics_failure_is_an_error_envelope).
Both routes now wait for all their reads before answering.
* fix(review): queue only ambiguous files that nothing links, and heal the queue
The maintainer's live scan (19,908 files, scope "new") reported 4,677 files
under Needs review, but the queue held 10,232 items: every file the strict
resolver found ambiguous was queued, including the ones the legacy
browser-parser fallback then linked (all sampled "added" rows) and the ones
already linked to a known prompt (135 of 200 sampled "known" rows).
A file is now queued only when its final outcome is ambiguous and unlinked:
the resolver finds it ambiguous, the fallback links nothing and no saved row
links it already. An ambiguous file that a saved row links is reported as
unchanged ("Already linked to its prompt"), not as Needs review.
The queue heals itself. After each batch, the scan removes the pending items
of every listed file (read, or skipped as known) that a saved row now links.
The monitor does the same for each ambiguous output it links, and never
queues one. Only an undecided item is removed (pending, revision 0, no
decision row), so a decision is never deleted or changed and the schema
needs no new state.
The scan job's scope.review_items now counts this scan's files that are in
the pending queue after it, new or already there, so it agrees with the
review group's ambiguous rows and with the queue total. review_added and
review_settled are new, and the log gains "already in Needs review" and
"left Needs review · now linked" lines. A duplicate whose item is decided is
reported as "decided" and is not counted.
* fix(review): never queue an ambiguous file with nothing to choose
The maintainer's live queue held 39 items whose candidates all read
"unresolved" ("No positive / No negative"). An ambiguous result is now
reviewable only when at least two candidates are resolved with non-blank
positive text. With none the file could not be traced; with exactly one,
the untraced branches may still have made the file, so that text is not
certain either. Both are reported unresolved, never linked and never
queued (also the A1111 two-delimiter case, which has no candidates).
A rescan removes the pending items of such files whose own evidence
offers no choice, with the same safe delete as the linked settle
(pending, revision 0, no decision row). A limited item (over 16 candidates) that holds at least two
candidates is kept. The job scope gains review_untraced and the log
"= N images left Needs review · nothing to choose".
* test(scan): pin only the per-source scope keys beside the review counters
* test(system): match the scan scope keys beside the review counters
2614 lines
99 KiB
Python
2614 lines
99 KiB
Python
"""
|
|
Database operations for KikoTextEncode prompt storage and retrieval.
|
|
"""
|
|
|
|
import sqlite3
|
|
import json
|
|
import math
|
|
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 .annotations import AnnotationOperations
|
|
from .review import ReviewOperations
|
|
from .review_decisions import ReviewDecisionOperations
|
|
from .saved_views import SavedViewOperations
|
|
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
|
|
from .health import HealthOperations, iter_consistency_issues
|
|
|
|
# Import logging system
|
|
try:
|
|
from ..utils.logging_config import get_logger
|
|
from ..utils.hashing import generate_prompt_hash
|
|
from ..utils.json_strict import dumps_strict
|
|
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
|
|
from utils.json_strict import dumps_strict
|
|
|
|
|
|
# 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)
|
|
"""
|
|
|
|
|
|
# The scan's known set (#249 Activity-2): every image row, with its stored file
|
|
# size and whether its pair link is resolved: exactly one distinct pair, the rule
|
|
# _LINK_SUMMARY_SQL and get_image_prompt_info use. A file is known and resolved
|
|
# when it has exactly one image row, that row is resolved, and the row's stored
|
|
# size equals the file's size now; several rows (one file linked to more than
|
|
# one prompt) are ambiguous, and an unknown (NULL) stored size is not known. The
|
|
# rows are streamed in fetches of KNOWN_PATHS_FETCH_SIZE and only the scanned
|
|
# paths are kept, so memory follows the scan (at most MAX_SCAN_FILES paths), not
|
|
# the table.
|
|
_KNOWN_PATHS_SQL = """SELECT gi.file_path, gi.image_path, gi.file_size,
|
|
(SELECT COUNT(DISTINCT ip.pair_id) FROM image_prompt_pairs ip
|
|
WHERE ip.image_id = gi.id) = 1
|
|
FROM generated_images gi"""
|
|
KNOWN_PATHS_FETCH_SIZE = 1000
|
|
|
|
# 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 (text, Library title, category, rating)
|
|
# then describes the card. The pair's own title names its negative variant, so
|
|
# it is not read (D10).
|
|
_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, 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 linked positive of each given path (one unambiguous link, pair scope as
|
|
# in _LINK_FILTER_SQL) with its category; {hit} as in link_filter_query. Used
|
|
# for the few files whose current image tags meet a Tags filter (#249 B3).
|
|
_LINK_FACTS_SQL = """SELECT gi.file_path, p.rating, {hit}, p.id, p.category
|
|
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 IN ({marks}) AND NOT EXISTS (
|
|
SELECT 1 FROM generated_images other
|
|
WHERE other.file_path = gi.file_path AND other.id <> gi.id)"""
|
|
_POSITIVE_TAG_NAMES_SQL = (
|
|
"SELECT pt.prompt_id, t.name FROM prompt_tags pt JOIN tags t "
|
|
"ON t.id = pt.tag_id WHERE pt.prompt_id IN ({marks})"
|
|
)
|
|
_FACTS_CHUNK = 500
|
|
# The positive's text, its Library title or one of its tag names contains
|
|
# :pattern (LIKE, so ASCII case is ignored as in Library search). A pair's own
|
|
# title names its negative variant and is never searched (#249 D12).
|
|
_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 '\\'))"""
|
|
_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[8], row[9], row[10]
|
|
# The linked positive's Library title (D10/D20).
|
|
titles = (row[11],)
|
|
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,
|
|
AnnotationOperations,
|
|
ReviewOperations,
|
|
ReviewDecisionOperations,
|
|
SavedViewOperations,
|
|
HealthOperations,
|
|
):
|
|
"""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 _bump_image_tag_bytes(conn: sqlite3.Connection, tag_id: int) -> int:
|
|
"""Move the tag revision of every image set that carries *tag_id*.
|
|
|
|
Vocabulary rename, merge and delete change what an image's set reads
|
|
as, so a Tag dialog opened before the change must see ``tag-conflict``
|
|
on Save instead of recreating the old name. Bumps the image index
|
|
revision once when any set was affected; returns that count.
|
|
"""
|
|
cursor = conn.execute(
|
|
"UPDATE media_annotation_bytes SET tag_revision = tag_revision + 1 "
|
|
"WHERE id IN (SELECT byte_id FROM media_image_tags WHERE tag_id = ?)",
|
|
(tag_id,),
|
|
)
|
|
if cursor.rowcount > 0:
|
|
conn.execute(
|
|
"UPDATE media_annotation_metadata SET "
|
|
"image_index_revision = image_index_revision + 1 WHERE singleton_id = 1"
|
|
)
|
|
return cursor.rowcount
|
|
|
|
@classmethod
|
|
def _move_tag_links(
|
|
cls, conn: sqlite3.Connection, source_id: int, target_id: int
|
|
) -> None:
|
|
"""Repoint component, pair and image links from source to target.
|
|
|
|
Owners that already carry the target keep that one link; the source's
|
|
leftover duplicates are removed (an image set keeps the target's
|
|
position). The source registry row is left for the caller to delete.
|
|
Image sets that carried the source get a new tag revision first.
|
|
"""
|
|
cls._bump_image_tag_bytes(conn, source_id)
|
|
for junction in ("prompt_tags", "prompt_pair_tags", "media_image_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
|
|
self._bump_image_tag_bytes(conn, old_tag_id)
|
|
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
|
|
self._bump_image_tag_bytes(conn, tag_id)
|
|
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]
|
|
image_count = conn.execute(
|
|
"SELECT COUNT(*) FROM media_image_tags WHERE tag_id = ?", (src_id,)
|
|
).fetchone()[0]
|
|
if src_count or pair_count or image_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,
|
|
dumps_strict(metadata.get("workflow", {})),
|
|
dumps_strict(metadata.get("prompt", {})),
|
|
dumps_strict(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 and Infinity 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 not math.isfinite(obj):
|
|
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_known_resolved_paths(self, files: Dict[str, Optional[int]]) -> Set[str]:
|
|
"""The paths in ``files`` the Library already links with a resolved link.
|
|
|
|
``files`` maps each scanned path to its size on disk (None when it could
|
|
not be read). A path qualifies when exactly one ``generated_images`` row
|
|
claims the file, that row has exactly one pair link
|
|
(``image_prompt_pairs``, written only by resolved ingestion), and its
|
|
stored ``file_size`` is not NULL and equals the given size, so a file
|
|
replaced at the same path is read again. Files linked without a pair
|
|
(the monitor's fallback and FIFO links, the legacy positive-only
|
|
import), linked to several prompts or pairs, or not linked at all do
|
|
not qualify. Both sides go through ``normalize_image_path`` (absolute,
|
|
normalized, ``os.path.normcase``), so a Windows row stored in another
|
|
case or spelling still matches. One query, streamed; returns the
|
|
matching keys of ``files``.
|
|
"""
|
|
wanted: Dict[str, List[Tuple[str, Optional[int]]]] = {}
|
|
for path, size in files.items():
|
|
key = normalize_image_path(str(path))
|
|
wanted.setdefault(key, []).append((str(path), size))
|
|
if not wanted:
|
|
return set()
|
|
rows_seen: Dict[str, int] = {}
|
|
stored_sizes: Dict[str, int] = {}
|
|
with self.model.get_connection() as conn:
|
|
cursor = conn.execute(_KNOWN_PATHS_SQL)
|
|
while True:
|
|
rows = cursor.fetchmany(KNOWN_PATHS_FETCH_SIZE)
|
|
if not rows:
|
|
break
|
|
for file_path, image_path, file_size, is_resolved in rows:
|
|
# image_path is NOT NULL; file_path is NULL only on old rows.
|
|
key = normalize_image_path(file_path or image_path)
|
|
if key not in wanted:
|
|
continue
|
|
rows_seen[key] = rows_seen.get(key, 0) + 1
|
|
if is_resolved and file_size is not None:
|
|
stored_sizes[key] = file_size
|
|
return {
|
|
path
|
|
for key, stored_size in stored_sizes.items()
|
|
if rows_seen[key] == 1
|
|
for path, size in wanted[key]
|
|
if size is not None and size == stored_size
|
|
}
|
|
|
|
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; a pair's own title
|
|
names its negative variant and is not searched.
|
|
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_link_facts(
|
|
self, paths: List[str], q: str = ""
|
|
) -> Dict[str, Tuple[bool, Optional[int], int, Optional[str], Tuple[str, ...]]]:
|
|
"""Map each normalized path with one unambiguous link to its positive's
|
|
(hit, rating, positive id, category, tag names).
|
|
|
|
The link is resolved as in get_image_link_matches, and ``hit`` says
|
|
whether *q* occurs in the positive's text, Library title or a tag
|
|
name (always true without *q*). Paths go in chunks of 500 through the
|
|
``file_path`` index; the tag names are read for those positives only.
|
|
"""
|
|
hit, params = "1", {}
|
|
if q:
|
|
hit = _LINK_FILTER_HIT
|
|
params["pattern"] = f"%{escape_like(q)}%"
|
|
rows: Dict[str, Tuple[bool, Optional[int], int, Optional[str]]] = {}
|
|
names: Dict[int, List[str]] = {}
|
|
unique = list(dict.fromkeys(paths))
|
|
with self.model.get_connection() as conn:
|
|
for start in range(0, len(unique), _FACTS_CHUNK):
|
|
chunk = unique[start : start + _FACTS_CHUNK]
|
|
bound = {f"path{i}": path for i, path in enumerate(chunk)}
|
|
marks = ",".join(f":{key}" for key in bound)
|
|
sql = _LINK_FACTS_SQL.format(hit=hit, marks=marks)
|
|
for path, rating, matched, positive, category in conn.execute(
|
|
sql, {**params, **bound}
|
|
):
|
|
rows[path] = (bool(matched), rating, positive, category)
|
|
positives = sorted({row[2] for row in rows.values()})
|
|
for start in range(0, len(positives), _FACTS_CHUNK):
|
|
chunk = positives[start : start + _FACTS_CHUNK]
|
|
marks = ",".join("?" * len(chunk))
|
|
for positive, name in conn.execute(
|
|
_POSITIVE_TAG_NAMES_SQL.format(marks=marks), chunk
|
|
):
|
|
names.setdefault(positive, []).append(name)
|
|
return {
|
|
path: (*row, tuple(names.get(row[2], ()))) for path, row in rows.items()
|
|
}
|
|
|
|
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.
|
|
|
|
The same queries back the System health check's Consistency row
|
|
(database/health.py).
|
|
"""
|
|
with self.model.get_connection() as conn:
|
|
return list(iter_consistency_issues(conn))
|
|
|
|
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
|