From 73d7b9490b753a0953ceabc5c647cc31e096e0e1 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Mon, 19 Oct 2026 22:22:03 -0700 Subject: [PATCH] feat(library): expose run count, last use and variant totals --- database/library.py | 38 +++++++++++- database/operations.py | 18 +++--- py/api/images.py | 56 +++++++++++------ tests/test_api_library.py | 30 +++++++++ tests/test_api_workspace_images.py | 89 +++++++++++++++++++++++++-- tests/test_library.py | 99 ++++++++++++++++++++++++++++++ 6 files changed, 295 insertions(+), 35 deletions(-) diff --git a/database/library.py b/database/library.py index b961f27..232e4a1 100644 --- a/database/library.py +++ b/database/library.py @@ -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, } diff --git a/database/operations.py b/database/operations.py index 1bfed70..6acb86b 100644 --- a/database/operations.py +++ b/database/operations.py @@ -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 diff --git a/py/api/images.py b/py/api/images.py index 5d70ce0..42fad5c 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -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): diff --git a/tests/test_api_library.py b/tests/test_api_library.py index cd30b55..e8c3f63 100644 --- a/tests/test_api_library.py +++ b/tests/test_api_library.py @@ -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 diff --git a/tests/test_api_workspace_images.py b/tests/test_api_workspace_images.py index f67fd49..da33b9f 100644 --- a/tests/test_api_workspace_images.py +++ b/tests/test_api_workspace_images.py @@ -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: diff --git a/tests/test_library.py b/tests/test_library.py index 5ddfd00..f204061 100644 --- a/tests/test_library.py +++ b/tests/test_library.py @@ -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()