feat(library): expose run count, last use and variant totals

This commit is contained in:
Vito Sansevero
2026-10-19 23:05:34 -07:00
parent d2b5b25a42
commit 73d7b9490b
6 changed files with 295 additions and 35 deletions
+35 -3
View File
@@ -1,5 +1,6 @@
"""Read-only positive-parent views over existing prompt and pair tables."""
import datetime
import sqlite3
from collections.abc import Iterator
from contextlib import contextmanager
@@ -46,6 +47,26 @@ def display_title(stored_titles: tuple, text: object) -> str | None:
return titled(stored_titles, text)[0]
def utc_iso(value: object) -> str | None:
"""A stored timestamp as UTC ISO text, or None when absent or unreadable.
Usage writes are already UTC ISO; a naive value (SQLite's CURRENT_TIMESTAMP
form) is read as UTC, and an offset is converted to UTC.
"""
if not isinstance(value, str):
return None
text = value.strip()
if text.endswith(("Z", "z")):
text = text[:-1] + "+00:00"
try:
stamp = datetime.datetime.fromisoformat(text)
except ValueError:
return None
if stamp.tzinfo is None:
return stamp.replace(tzinfo=datetime.timezone.utc).isoformat()
return stamp.astimezone(datetime.timezone.utc).isoformat()
_NAME = (
f"LOWER(CASE WHEN TRIM(COALESCE(ls.title,''),{_BLANK})!='' "
"THEN ls.title ELSE p.text END)"
@@ -174,6 +195,7 @@ class LibraryOperations:
) -> list[dict]:
rows = conn.execute(
"SELECT p.id,p.text,ls.title,p.category,p.rating,p.notes,p.created_at,"
"p.run_count,p.last_used_at,"
+ (f"{sort_expression} AS _sort," if sort_expression else "")
+ "(SELECT COUNT(*) FROM prompt_pairs v WHERE v.positive_prompt_id=p.id "
"AND v.archived=0) AS variant_count,"
@@ -194,6 +216,9 @@ class LibraryOperations:
):
by_id[row["prompt_id"]]["tags"].append(row["name"])
for row in result:
# Usage is shown, never compared: the token hashes owner fields only.
row["run_count"] = row["run_count"] if type(row["run_count"]) is int else 0
row["last_used_at"] = utc_iso(row["last_used_at"])
row["edit_token"] = library_edit_token(epoch, row)
return result
@@ -225,9 +250,15 @@ class LibraryOperations:
ceiling = after["ceiling"]
where += " AND p.id<=?"
params.append(ceiling)
total = conn.execute(
"SELECT COUNT(*) FROM " + _FROM + " WHERE " + where, params
).fetchone()[0]
# One SELECT counts the filtered parents and their active variants.
total, variant_total = conn.execute(
"SELECT COUNT(*),COALESCE(SUM((SELECT COUNT(*) FROM prompt_pairs v "
"WHERE v.positive_prompt_id=p.id AND v.archived=0)),0) FROM "
+ _FROM
+ " WHERE "
+ where,
params,
).fetchone()
expression, direction = _ORDERS[sort]
if after is not None:
operator = ">" if direction == "ASC" else "<"
@@ -259,6 +290,7 @@ class LibraryOperations:
"library_epoch": epoch,
"items": rows,
"total": total,
"variant_total": variant_total,
"has_more": has_more,
"next_key": next_key,
}
+10 -8
View File
@@ -203,7 +203,7 @@ WHERE gi.image_path IN ({marks}) OR gi.file_path IN ({marks}){fallback}"""
# 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}
_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
@@ -2330,23 +2330,25 @@ class PromptDatabase(
*,
every_link: bool = False,
wanted: Optional[Set[str]] = None,
) -> Dict[str, Tuple[bool, Optional[int]]]:
"""Map normalized paths of unambiguously linked images to (hit, rating).
) -> Dict[str, Tuple[bool, Optional[int], int]]:
"""Map normalized paths of unambiguously linked images to
(hit, rating, positive id).
Rows are limited to images whose linked positive has *category* and
every tag in *tags* (exact names). ``hit`` says whether *q* occurs in
the positive's text, Library title or a tag name, or in the pair title.
Without category or tags only hits are returned, unless *every_link*
asks for every linked image (the rating sort). One SELECT whose rows
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]]] = {}
matches: Dict[str, Tuple[bool, Optional[int], int]] = {}
with self.model.get_connection() as conn:
for path, rating, matched in conn.execute(sql, params):
for path, rating, matched, positive in conn.execute(sql, params):
if wanted is None or path in wanted:
matches[path] = (bool(matched), rating)
matches[path] = (bool(matched), rating, positive)
return matches
def get_image_paths_matching(
@@ -2359,7 +2361,7 @@ class PromptDatabase(
"""
return {
path
for path, (hit, _rating) in self.get_image_link_matches(
for path, (hit, _rating, _positive) in self.get_image_link_matches(
q, category, tags
).items()
if hit or not q
+36 -20
View File
@@ -561,18 +561,20 @@ def _workspace_filters(query: Mapping) -> dict:
def _filtered_page(
entries: list, start: int, size: int, filters: dict, links: dict | None
) -> tuple[int, list]:
) -> tuple[int, list, int | None]:
"""Count the *entries* the filters admit and take *size* from *start* on.
*links* maps normalized paths to (hit, rating) for linked images (see
PromptDatabase.get_image_link_matches); a category or tag filter admits
only those, and ``q`` admits a filename match or a linked-positive hit.
*links* maps normalized paths to (hit, rating, positive id) for every
unambiguously linked image (see PromptDatabase.get_image_link_matches); a
category or tag filter admits only those, and ``q`` admits a filename match
or a linked-positive hit. Returns the total, the page and the number of
distinct linked positives among all admitted entries (None without links).
"""
query, media, folder = filters["query"], filters["media"], filters["folder"]
members_only = filters["category"] is not None or bool(filters["tags"])
native_folder = folder.replace("/", os.sep)
nested = native_folder + os.sep
total, page = 0, []
total, page, positives = 0, [], set()
for position, entry in enumerate(entries):
if folder and not (
entry.folder == native_folder or entry.folder.startswith(nested)
@@ -580,16 +582,17 @@ def _filtered_page(
continue
if media != "all" and (media == "video") != entry.video:
continue
if links is not None and (members_only or query):
match = links.get(entry.normalized)
if members_only and match is None:
continue
if query and query not in entry.folded and not (match and match[0]):
continue
match = None if links is None else links.get(entry.normalized)
if members_only and match is None:
continue
if query and query not in entry.folded and not (match and match[0]):
continue
total += 1
if match is not None:
positives.add(match[2])
if position >= start and len(page) < size:
page.append(entry)
return total, page
return total, page, None if links is None else len(positives)
def _in_subfolder(rel_dir, subfolder):
@@ -969,8 +972,10 @@ class ImageRoutesMixin:
for index, root in enumerate(output_dirs)
if root_index is None or index == root_index
]
page, total, has_more, next_cursor = await self._run_in_executor(
self._select_workspace_page, listings, filters, after, limit, scope
page, total, has_more, next_cursor, positive_count = (
await self._run_in_executor(
self._select_workspace_page, listings, filters, after, limit, scope
)
)
images = await self._run_in_executor(self._format_workspace_page, page)
return web.json_response(
@@ -978,6 +983,7 @@ class ImageRoutesMixin:
"success": True,
"images": images,
"total": total,
"positive_count": positive_count,
"limit": limit,
"has_more": has_more,
"next_cursor": next_cursor,
@@ -1061,19 +1067,26 @@ class ImageRoutesMixin:
) -> tuple:
"""One keyset page of the filtered listing (blocking; run in executor).
At most one SELECT (get_image_link_matches) when ``q``, a category,
tags or the rating sort need the linked positives.
One SELECT (get_image_link_matches) reads every linked positive of the
listing: the filters, the rating sort and the positive count share it.
Without a filter that needs it, a failed read leaves the listing
usable and the positive count unknown (None).
"""
order, category, tags = filters["order"], filters["category"], filters["tags"]
links = None
if filters["query"] or category is not None or tags or order == "rating":
needed = bool(filters["query"]) or category is not None or bool(tags)
try:
links = self.db.get_image_link_matches(
filters["text"],
category,
tags,
every_link=order == "rating",
every_link=True,
wanted=self._workspace_wanted(listings),
)
except Exception:
if needed or order == "rating":
raise
self.logger.warning("Workspace positive count could not be read")
links = None
if order == "rating":
def rank(entry):
@@ -1102,7 +1115,9 @@ class ImageRoutesMixin:
start = len(entries) - bisect_left(ascending, after, key=key)
else:
start = bisect_right(entries, after, key=key)
total, page = _filtered_page(entries, start, limit + 1, filters, links)
total, page, positive_count = _filtered_page(
entries, start, limit + 1, filters, links
)
has_more = len(page) > limit
page = page[:limit]
next_cursor = _workspace_cursor(key(page[-1]), scope) if has_more else None
@@ -1111,6 +1126,7 @@ class ImageRoutesMixin:
total,
has_more,
next_cursor,
positive_count,
)
async def serve_image(self, request):
+30
View File
@@ -188,6 +188,10 @@ class TestLibraryAPI(ImageAPITestCase):
async def test_maximum_parent_page_uses_bounded_sql_queries_for_previews(self):
for index in range(200):
positive = self._save_prompt(f"Parent {index}", tags=[f"tag {index}"])
if index % 3 == 0:
self.api.db.save_prompt_pair(f"Parent {index}", f"Negative {index}")
if index % 4 == 0:
self.api.db.record_prompt_use(positive, times=index + 1)
make_png(self.output_dir / f"{index}.png")
self.api.db.link_image_to_prompt(
positive, str(self.output_dir / f"{index}.png")
@@ -232,6 +236,32 @@ class TestLibraryAPI(ImageAPITestCase):
for item in result["items"]
)
)
# Usage fields and the variant total ride on the same five SELECTs.
self.assertEqual(result["variant_total"], 67)
self.assertEqual(
result["variant_total"],
sum(item["variant_count"] for item in result["items"]),
)
used = {
int(item["text"].split()[-1]): item
for item in result["items"]
if item["last_used_at"] is not None
}
self.assertEqual(sorted(used), list(range(0, 200, 4)))
self.assertTrue(
all(
item["run_count"] == index + 1
and item["last_used_at"].endswith("+00:00")
for index, item in used.items()
)
)
self.assertTrue(
all(
type(item["run_count"]) is int
for item in result["items"]
if item["last_used_at"] is None
)
)
async def test_long_name_cursor_stays_small_and_preserves_full_sorting(self):
prefix = "猫" * 10_000
+85 -4
View File
@@ -560,7 +560,9 @@ class TestWorkspaceImages(ImageAPITestCase):
for query in statements
if query.lstrip().upper().startswith(("SELECT", "WITH"))
]
self.assertEqual(len(selects), 1, selects)
# The page's link summary, and the positive count's one link read.
self.assertEqual(len(selects), 2, selects)
self.assertEqual(data["positive_count"], 1) # six files, one prompt
async def test_link_summary_runs_in_executor_and_survives_database_errors(self):
import threading
@@ -827,6 +829,79 @@ class TestWorkspaceImages(ImageAPITestCase):
db.link_image_to_prompt(owner, str(files["ambiguous"]))
return files
async def test_positive_count_counts_distinct_linked_positives(self):
files = self._filter_fixtures()
db = self.api.db
night_ab = db.get_prompt_id_for_image(str(files["night_ab"]))
# A second file of one positive still counts that positive once.
again = self.media("again.png", mtime=50)
db.link_image_to_prompt(night_ab, str(again))
cases = (
# Six unambiguous positives; the ambiguous file and unlinked count none.
({}, 10, 6),
# Text, title, tag, pair title (its positive) and the filename match.
({"q": "ZEBRA"}, 6, 4),
({"category": "Night"}, 4, 3),
({"tags": "a,b"}, 3, 2),
({"q": "plain_unlinked"}, 1, 0),
({"media": "video"}, 0, 0),
({"sort": "rating"}, 10, 6),
({"sort": "name", "category": "Day"}, 3, 3),
)
for params, total, positives in cases:
with self.subTest(**params):
response, data = await self.get(**params)
self.assertEqual(response.status, 200, data)
self.assertEqual(data["total"], total)
self.assertEqual(data["positive_count"], positives)
# The count covers the whole matching set, on every cursor page.
response, first = await self.get(limit=2, q="ZEBRA")
self.assertEqual(first["positive_count"], 4)
response, second = await self.get(
limit=2, q="ZEBRA", cursor=first["next_cursor"]
)
self.assertEqual(response.status, 200, second)
self.assertEqual(second["positive_count"], 4)
# One SELECT counts the positives; the page's link summary is the other.
statements = []
original = db.model.get_connection
def trace_connection():
connection = original()
connection.set_trace_callback(statements.append)
return connection
with patch.object(db.model, "get_connection", side_effect=trace_connection):
response, data = await self.get()
self.assertEqual(response.status, 200)
selects = [
query
for query in statements
if query.lstrip().upper().startswith(("SELECT", "WITH"))
]
self.assertEqual(len(selects), 2, selects)
# Without filters the listing outlives a database failure; the count
# is then unknown rather than wrong.
with patch.object(
db, "get_image_link_matches", side_effect=RuntimeError("database gone")
):
response, data = await self.get()
self.assertEqual(response.status, 200, data)
self.assertEqual(data["total"], 10)
self.assertIsNone(data["positive_count"])
self.assertNotIn("database gone", json.dumps(data))
# A filter or the rating sort needs the read, so its failure is an error.
for params in ({"q": "zebra"}, {"category": "Night"}, {"sort": "rating"}):
with (
self.subTest(failing=params),
patch.object(
db, "get_image_link_matches", side_effect=RuntimeError("gone")
),
):
response, data = await self.get(**params)
self.assertEqual(response.status, 500, data)
self.assertNotIn("images", data)
async def test_search_matches_filename_prompt_text_and_tags(self):
from database.models import normalize_image_path
@@ -1026,13 +1101,18 @@ class TestWorkspaceImages(ImageAPITestCase):
self._names(data), ["night_a.png", "night_ab.png", "paired.png"]
)
wanted = {normalize_image_path(str(files["night_ab"])), "/elsewhere/7.png"}
elsewhere = self.api.db.get_prompt_id_for_image("/elsewhere/7.png")
self.assertEqual(
self.api.db.get_image_link_matches(
"", "Night", [], every_link=True, wanted=wanted
),
{
normalize_image_path(str(files["night_ab"])): (True, 5),
"/elsewhere/7.png": (True, 5),
normalize_image_path(str(files["night_ab"])): (
True,
5,
self.api.db.get_prompt_id_for_image(str(files["night_ab"])),
),
"/elsewhere/7.png": (True, 5, elsewhere),
},
)
@@ -1157,7 +1237,8 @@ class TestWorkspaceImages(ImageAPITestCase):
response, _ = await self.get(limit=10)
self.assertEqual(response.status, 200)
self.assertEqual(len(built), len(listed))
self.assertEqual(selects(), []) # no filter, no SELECT
# No filter still reads the linked positives once, for the count.
self.assertEqual(len(selects()), 1)
newest = self.api._workspace_order_cache[((0, str(self.output_dir)),)]
newest = newest[1]["newest"]
for params in self._BUDGET_REQUESTS:
+99
View File
@@ -487,6 +487,105 @@ class TestLibrary(unittest.TestCase):
rotate_history_epoch(self.db.model.get_connection())
self.assertNotEqual(token(), before)
def test_library_items_expose_run_count_and_last_used(self):
import datetime
used = self.save("Used")
never = self.save("Never")
self.assertTrue(self.db.record_prompt_use(used))
self.assertTrue(self.db.record_prompt_use(used, times=2))
legacy = self.save("Legacy naive")
offset = self.save("Offset")
broken = self.save("Broken")
zulu = self.save("Zulu")
with self.db.model.get_connection() as connection:
for prompt_id, value in (
(legacy, "2026-10-07 18:39:06"),
(zulu, "2026-10-07T18:39:06Z"),
(offset, "2026-10-07T20:39:06.250+02:00"),
(broken, "yesterday"),
):
connection.execute(
"UPDATE prompts SET last_used_at=?,run_count=1 WHERE id=?",
(value, prompt_id),
)
items = {row["id"]: row for row in self.db.list_library()["items"]}
self.assertEqual(items[used]["run_count"], 3)
stamp = datetime.datetime.fromisoformat(items[used]["last_used_at"])
self.assertEqual(stamp.utcoffset(), datetime.timedelta(0))
self.assertLess(
abs(datetime.datetime.now(datetime.timezone.utc) - stamp),
datetime.timedelta(minutes=5),
)
self.assertEqual(items[never]["run_count"], 0)
self.assertIsNone(items[never]["last_used_at"])
self.assertEqual(items[legacy]["last_used_at"], "2026-10-07T18:39:06+00:00")
self.assertEqual(
items[offset]["last_used_at"], "2026-10-07T18:39:06.250000+00:00"
)
self.assertIsNone(items[broken]["last_used_at"])
self.assertEqual(items[zulu]["last_used_at"], "2026-10-07T18:39:06+00:00")
self.assertEqual(type(items[broken]["run_count"]), int)
detail = self.db.get_library_positive(used)["positive"]
self.assertEqual(detail["run_count"], 3)
self.assertEqual(detail["last_used_at"], items[used]["last_used_at"])
def test_variant_total_matches_filtered_cards(self):
archived = self.db.save_prompt_pair("Cat", "Blur", category="A")
self.db.save_prompt_pair("Cat", "Noise", category="A")
self.db.save_prompt_pair("Cat", "Grain", category="B")
self.db.archive_prompt_pair(archived)
self.db.save_prompt_pair("Dog", "Blur", category="B")
self.save("Legacy")
everything = self.db.list_library()
self.assertEqual(everything["variant_total"], 3)
self.assertEqual(
everything["variant_total"],
sum(row["variant_count"] for row in everything["items"]),
)
cats = self.db.list_library(text="Cat")
self.assertEqual(
(cats["total"], cats["variant_total"]),
(1, 2),
)
self.assertEqual(self.db.list_library(category="B")["variant_total"], 3)
self.assertEqual(self.db.list_library(text="nothing")["variant_total"], 0)
first = self.db.list_library(limit=1)
self.assertEqual(first["variant_total"], 3)
# A later parent is past the cursor's ceiling, like the total.
self.db.save_prompt_pair("Bird", "Blur")
second = self.db.list_library(limit=1, after=first["next_key"])
self.assertEqual(
(second["total"], second["variant_total"]),
(first["total"], 3),
)
def test_edit_token_unchanged_by_usage_fields(self):
positive = self.save("Cat", category="A", tags=["a"], rating=3)
def parent():
return self.db.get_library_positive(positive)["positive"]
before = parent()
listed = self.db.list_library()["items"][0]
self.assertEqual(listed["edit_token"], before["edit_token"])
self.db.record_prompt_use(positive, times=4)
with self.db.model.get_connection() as connection:
connection.execute(
"UPDATE prompts SET last_used_at='2001-01-01T00:00:00+00:00' "
"WHERE id=?",
(positive,),
)
after = parent()
self.assertNotEqual(
(after["run_count"], after["last_used_at"]),
(before["run_count"], before["last_used_at"]),
)
self.assertEqual(after["edit_token"], before["edit_token"])
self.assertEqual(
self.db.list_library()["items"][0]["edit_token"], before["edit_token"]
)
def test_library_snapshot_reads_epoch_and_ceiling_in_one_select_and_ends(self):
self.save("Cat")
connection = self.db.model.get_connection()