From 66f61de0dbca0f952775699ced37fa5fe6d8fba1 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Thu, 22 Oct 2026 09:31:58 -0700 Subject: [PATCH] feat(images): match image tags in the Tags filter, search and inspector The workspace Images listing now admits a file under a Tags filter when every requested tag is carried by its linked positive or its current image tag set, and q also matches image tag names (case-folded substring). Image tags count only while a byte record is verified at the file's present (size, mtime_ns) (spec amendment S1): the tagged records are streamed once, only listed files carrying a filter tag or q hit are stat'ed, and their linked positives are read through the file_path index so tags, category, rating and the positive count combine. The rooted inspector envelope gains a validated image_tags facet (current / none / unavailable) read at the inspection's own stable stat. A failed image-tag read fails a Tags filter like a failed link read and degrades a q search to no image tags; plain listings read none. Legacy gallery and search responses are unchanged. --- TESTING.md | 6 + database/operations.py | 57 +++++ py/api/images.py | 164 +++++++++++++- py/api/inspection_service.py | 18 +- tests/test_api_annotations.py | 3 + tests/test_api_prompt_inspector.py | 79 +++++++ tests/test_api_workspace_images.py | 332 ++++++++++++++++++++++++++++- tests/test_inspection_response.py | 65 ++++++ utils/inspection_response.py | 43 +++- 9 files changed, 759 insertions(+), 8 deletions(-) diff --git a/TESTING.md b/TESTING.md index 14a10c0b..5227273a 100644 --- a/TESTING.md +++ b/TESTING.md @@ -1389,6 +1389,12 @@ impact and command receipts report `image_count` / `images` / (spec amendment S1) use a byte record only while the file's current `(size, mtime_ns)` equals its verified fingerprint; an in-place rewrite that preserves both keeps the old set matching until the file is next observed. The +Images Tags filter (a file matches when every tag is on its linked positive or +its current image tag set) and `q` search stream the tagged records once and +stat only the listed files whose records carry a filter tag or a `q` hit; the +inspector's `image_tags` facet (`current` / `none` / `unavailable`) uses the +inspection's own stable stat (`tests/test_api_workspace_images.py` +`TestWorkspaceImageTags`, `tests/test_api_prompt_inspector.py`). The browser fixtures clear `media_image_tags`, `media_annotation_bytes` and `media_annotation_files` on reset and never the metadata row. diff --git a/database/operations.py b/database/operations.py index 8fa24dd8..28bf35db 100644 --- a/database/operations.py +++ b/database/operations.py @@ -232,6 +232,23 @@ 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). @@ -2450,6 +2467,46 @@ class PromptDatabase( 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]: diff --git a/py/api/images.py b/py/api/images.py index 415cde9e..4898a995 100644 --- a/py/api/images.py +++ b/py/api/images.py @@ -14,7 +14,7 @@ import urllib.parse from bisect import bisect_left, bisect_right from collections.abc import Mapping from pathlib import Path -from typing import NamedTuple +from typing import Callable, NamedTuple from aiohttp import web from PIL import Image, UnidentifiedImageError @@ -418,6 +418,29 @@ def _listed(path: Path, mtime: float, root: Path, index: int) -> _Listed: ) +class _ImageTagMatch(NamedTuple): + """How one listed file's current image tags meet the filters (#249 B3).""" + + tags: bool # every filter tag is in its image tags or linked positive tags + query: bool # ``q`` is a case-folded substring of one image tag name + link: tuple | None # (hit, rating, positive id) of its unambiguous link + + +def _current_fingerprint(path: Path) -> tuple[int, int] | None: + """The file's present ``(st_size, st_mtime_ns)``; None when it is gone.""" + try: + observed = os.stat(path) + except OSError: + return None + return observed.st_size, observed.st_mtime_ns + + +def _image_tags_relevant(names: list, wanted: set, query: str) -> bool: + if wanted and not wanted.isdisjoint(names): + return True + return bool(query) and any(query in name.casefold() for name in names) + + def _workspace_key(entry: tuple, order: str) -> tuple: """The keyset sort key of *entry* (a _Listed or a 4-tuple entry).""" if not isinstance(entry, _Listed): @@ -752,6 +775,7 @@ def _filtered_page( filters: dict, links: dict | None, favorites: frozenset | None = None, + image_tags: dict | None = None, ) -> tuple[int, list, int | None]: """Count the *entries* the filters admit and take *size* from *start* on. @@ -761,6 +785,9 @@ def _filtered_page( or a linked-positive hit. Returns the total, the page and the number of distinct linked positives among all admitted entries (None without links). *favorites*, when given, admits only those listed paths (``favorites=1``). + *image_tags* maps normalized paths to their _ImageTagMatch: a tag filter + also admits a file whose current image tags (with its linked positive's + tags) carry every tag, and ``q`` also admits an image tag name hit. """ query, media, folder = filters["query"], filters["media"], filters["folder"] members_only = filters["category"] is not None or bool(filters["tags"]) @@ -780,9 +807,15 @@ def _filtered_page( if favorites is not None and entry.path not in favorites: continue match = None if links is None else links.get(entry.normalized) - if members_only and match is None: + image = None if image_tags is None else image_tags.get(entry.normalized) + if members_only and match is None and not (image and image.tags): continue - if query and query not in entry.folded and not (match and match[0]): + if ( + query + and query not in entry.folded + and not (match and match[0]) + and not (image and image.query) + ): continue total += 1 if match is not None: @@ -1276,6 +1309,80 @@ class ImageRoutesMixin: ) return orders[order] + def _workspace_locators(self, listings: list[tuple]) -> dict: + """The selected listing's private locators -> _Listed (cached). + + A locator is ``normcase(join(realpath(root), relative))``, the key + image tags are stored under (database/annotations.py); the listing + never follows symlinks, so it equals the listed file's own locator. + """ + orders = self._workspace_orders(listings) + reals = tuple(os.path.realpath(root) for _index, root, _files in listings) + key = ("locators", reals) + if key not in orders: + located = {} + for (index, root, files), real in zip(listings, reals): + prefix = os.path.join(str(root), "") + for entry in self._workspace_listed(index, root, files): + text = str(entry.path) + relative = ( + text[len(prefix) :] + if text.startswith(prefix) + else str(entry.path.relative_to(root)) + ) + located[os.path.normcase(os.path.join(real, relative))] = entry + orders[key] = located + return orders[key] + + def _image_tag_matches(self, listings: list[tuple], filters: dict) -> dict: + """Normalized path -> _ImageTagMatch for files whose image tags count. + + Streams the tagged byte records once and keeps only listed locators + carrying a filter tag or a ``q`` hit; only those candidates are + stat'ed (never every listed file), and a candidate's tags count only + while a byte record is verified at its present ``(size, mtime_ns)`` + (spec amendment S1). Under a Tags filter the candidates' linked + positives are read too, so their tags and category combine. + """ + tags, query, category = filters["tags"], filters["query"], filters["category"] + wanted = set(tags) + located = self._workspace_locators(listings) + candidates = {} + records = self.db.iter_image_tag_records() + try: + for locator, _size, _mtime_ns, names in records: + entry = located.get(locator) + if entry is not None and locator not in candidates: + if _image_tags_relevant(names, wanted, query): + candidates[locator] = entry + finally: + records.close() + stats = {} + for locator, entry in candidates.items(): + fingerprint = _current_fingerprint(entry.path) + if fingerprint is not None: + stats[locator] = fingerprint + current = self.db.current_image_tags(stats) if stats else {} + facts = {} + if wanted and current: + facts = self.db.get_image_link_facts( + [candidates[locator].normalized for locator in current], + filters["text"], + ) + matches = {} + for locator, names in current.items(): + entry = candidates[locator] + link = facts.get(entry.normalized) + carried = set(names).union(link[4] if link else ()) + matches[entry.normalized] = _ImageTagMatch( + bool(wanted) + and wanted <= carried + and (category is None or (link is not None and link[3] == category)), + bool(query) and any(query in name.casefold() for name in names), + link[:3] if link else None, + ) + return matches + def _select_workspace_page( self, listings: list[tuple], @@ -1328,6 +1435,30 @@ class ImageRoutesMixin: ), exc_info=True, ) + image_tags = None + if filters["query"] or tags: + try: + image_tags = self._image_tag_matches(listings, filters) + except Exception: + # A Tags filter cannot be served without them (as without the + # links); a search then matches without image tags. + if tags: + raise + self.logger.warning( + "Workspace image tags could not be read; searching without " + "image tags", + exc_info=True, + ) + if image_tags and links is not None: + # A linked file admitted through its image tags keeps its rating + # and counts its positive. + admitted = { + path: match.link + for path, match in image_tags.items() + if match.tags and match.link is not None and path not in links + } + if admitted: + links = {**links, **admitted} if order == "rating": def rank(entry): @@ -1360,7 +1491,7 @@ class ImageRoutesMixin: self._workspace_favorite_paths(listings) if filters["favorites"] else None ) total, page, positive_count = _filtered_page( - entries, start, limit + 1, filters, links, favorites + entries, start, limit + 1, filters, links, favorites, image_tags ) if not needed and after is not None: positive_count = _COUNT_OMITTED @@ -2467,6 +2598,9 @@ class ImageRoutesMixin: library=self._inspection_library( saved, unavailable=unavailable ), + image_tags=self._inspection_image_tags( + roots[index], Path(relative).parts + ), ) except CurrentFileChanged: return None @@ -2505,6 +2639,28 @@ class ImageRoutesMixin: {"success": False, "error": "Prompt inspection failed"}, status=500 ) + def _inspection_image_tags(self, root: Path, parts: tuple) -> Callable: + """The inspector's ``image_tags`` reader for one rooted file (S1). + + The returned function takes the inspection's own stable stat and + answers ``current`` only for a byte record verified at exactly that + ``(size, mtime_ns)``; a failed read is ``unavailable``, never an error. + """ + locator = os.path.normcase(os.path.join(os.path.realpath(root), *parts)) + + def read(size: int, mtime_ns: int) -> dict: + try: + found = self.db.current_image_tags({locator: (size, mtime_ns)}) + except Exception: + self.logger.warning("Inspection image tags could not be read") + return {"state": "unavailable", "tags": []} + tags = found.get(locator) + if tags is None: + return {"state": "none", "tags": []} + return {"state": "current", "tags": list(tags)} + + return read + def _inspection_library(self, saved: dict | None, *, unavailable: bool) -> dict: """The inspector's Library facts; the route sets the title afterwards.""" result = { diff --git a/py/api/inspection_service.py b/py/api/inspection_service.py index 787025c7..ee8fa6c3 100644 --- a/py/api/inspection_service.py +++ b/py/api/inspection_service.py @@ -3,6 +3,7 @@ import os import stat from pathlib import Path +from typing import Callable try: from ...utils.prompt_inspector import ( @@ -93,8 +94,19 @@ def _in_host_temp(path: Path) -> bool: def inspect_current_file( - path: Path, *, source_label: str, saved: dict | None, library: dict + path: Path, + *, + source_label: str, + saved: dict | None, + library: dict, + image_tags: Callable[[int, int], dict] | None = None, ) -> dict: + """Current-file evidence read between two equal stats of *path*. + + *image_tags*, when given, is called once with that stable + ``(st_size, st_mtime_ns)`` and its facet is returned as ``image_tags`` + (spec amendment S1: tags of a byte record verified at exactly that stat). + """ before = _current_file_facts(path) try: inspection = read_inspection_file(path) @@ -104,6 +116,9 @@ def inspect_current_file( fields = ("st_dev", "st_ino", "st_mode", "st_size", "st_mtime_ns", "st_ctime_ns") if any(getattr(before, key) != getattr(after, key) for key in fields): raise CurrentFileChanged + facets = {} + if image_tags is not None: + facets["image_tags"] = image_tags(before.st_size, before.st_mtime_ns) if inspection is None: evidence = unreadable_evidence() else: @@ -132,4 +147,5 @@ def inspect_current_file( **evidence, "inspection": inspection, "library": library, + **facets, } diff --git a/tests/test_api_annotations.py b/tests/test_api_annotations.py index 624b41b2..dca5fea8 100644 --- a/tests/test_api_annotations.py +++ b/tests/test_api_annotations.py @@ -930,11 +930,13 @@ class TestEnvelopes(AnnotationAPITestCase): ) self.assertEqual(resp.status, 200) data = await self.json_of(resp) + # Additive only: B3 adds the image_tags facet beside the Library facts. self.assertEqual( sorted(data), [ "file", "image", + "image_tags", "inspection_version", "library", "metadata", @@ -945,6 +947,7 @@ class TestEnvelopes(AnnotationAPITestCase): ], ) self.assertEqual(sorted(data["image"]), ["filename", "path", "root_index"]) + self.assertEqual(data["image_tags"], {"state": "none", "tags": []}) self.assertEqual( sorted(data["library"]), [ diff --git a/tests/test_api_prompt_inspector.py b/tests/test_api_prompt_inspector.py index fcd92613..027b1453 100644 --- a/tests/test_api_prompt_inspector.py +++ b/tests/test_api_prompt_inspector.py @@ -1642,3 +1642,82 @@ class TestPromptInspector(ImageAPITestCase): lookup.assert_called_once_with( str(path), False, report_ambiguity=True ) + + +class TestInspectorImageTags(ImageAPITestCase): + """B3: the inspector envelope's ``image_tags`` facet (spec amendment S1).""" + + inspect = TestPromptInspector.inspect + + def tag(self, path, tags): + import hashlib + import os + + stat = path.stat() + return self.api.db.replace_image_tags( + locator=os.path.normcase(os.path.realpath(path)), + sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + size=stat.st_size, + mtime_ns=stat.st_mtime_ns, + expected_epoch=self.api.db.get_execution_epoch(), + expected_tag_revision=0, + tags=tags, + ) + + async def test_inspector_reports_current_none_and_unavailable_image_tags(self): + import os + + tagged = png(self.output_dir / "sub" / "tagged.png", prompt=graph()) + changed = png(self.output_dir / "changed.png", prompt=graph()) + png(self.output_dir / "untagged.png", prompt=graph()) + empty = png(self.output_dir / "empty.png", prompt=graph()) + self.tag(tagged, ["dunes", "warm_light"]) + self.tag(changed, ["dunes"]) + self.tag(empty, []) + stat = changed.stat() + os.utime(changed, ns=(stat.st_atime_ns, stat.st_mtime_ns + 1_000_000)) + + response, envelope = await self.inspect("sub/tagged.png") + self.assertEqual(response.status, 200, envelope) + self.assertEqual( + envelope.get("image_tags"), + {"state": "current", "tags": ["dunes", "warm_light"]}, + ) + # The tags never change the Library facts of the prompt link. + self.assertEqual(envelope["library"]["tags"], []) + _, changed_body = await self.inspect("changed.png") + self.assertEqual(changed_body.get("image_tags"), {"state": "none", "tags": []}) + _, untagged = await self.inspect("untagged.png") + self.assertEqual(untagged.get("image_tags"), {"state": "none", "tags": []}) + _, saved_empty = await self.inspect("empty.png") + self.assertEqual( + saved_empty.get("image_tags"), {"state": "current", "tags": []} + ) + with patch.object( + self.api.db, "current_image_tags", side_effect=RuntimeError("boom") + ): + with self.assertLogs("prompt_manager", level="WARNING"): + response, failed = await self.inspect("sub/tagged.png") + self.assertEqual(response.status, 200, failed) + self.assertEqual(failed.get("image_tags"), {"state": "unavailable", "tags": []}) + self.assertEqual(failed["prompts"], envelope["prompts"]) + self.assertNotIn(str(self.comfy_dir), json.dumps(envelope)) + + async def test_inspector_image_tags_use_the_inspection_stable_stat(self): + tagged = png(self.output_dir / "tagged.png", prompt=graph()) + self.tag(tagged, ["dunes"]) + stat = tagged.stat() + seen = [] + real = self.api.db.current_image_tags + + def spy(stats): + seen.append(dict(stats)) + return real(stats) + + with patch.object(self.api.db, "current_image_tags", side_effect=spy): + _, body = await self.inspect("tagged.png") + self.assertEqual( + body.get("image_tags"), {"state": "current", "tags": ["dunes"]} + ) + self.assertEqual(list(seen[0].values()), [(stat.st_size, stat.st_mtime_ns)]) + self.assertEqual(len(seen), 1) diff --git a/tests/test_api_workspace_images.py b/tests/test_api_workspace_images.py index 18c71f4c..80d26558 100644 --- a/tests/test_api_workspace_images.py +++ b/tests/test_api_workspace_images.py @@ -1361,8 +1361,10 @@ class TestWorkspaceImages(ImageAPITestCase): for query in statements if query.lstrip().upper().startswith(("SELECT", "WITH")) ] - self.assertEqual(len(selects), 3, selects) + # Plus the one stream of tagged image records q and tags read (B3). + self.assertEqual(len(selects), 4, selects) self.assertEqual(sum("image_favorites" in q for q in selects), 1, selects) + self.assertEqual(sum("media_image_tags" in q for q in selects), 1, selects) async def test_rating_sort_orders_linked_then_unlinked(self): self._filter_fixtures() @@ -1626,7 +1628,15 @@ class TestWorkspaceImages(ImageAPITestCase): response, data = await self.get(limit=10, **params) self.assertEqual(response.status, 200, data) self.assertEqual(built, []) - self.assertEqual(len(selects()), 1) + found = selects() + # q and tags also stream the tagged image records once + # (#249 B3); none are tagged here, so nothing is stat'ed. + streamed = "q" in params or "tags" in params + self.assertEqual(len(found), 2 if streamed else 1, found) + self.assertEqual( + sum("media_image_tags" in query for query in found), + int(streamed), + ) self.assertIs( self.api._workspace_order_cache[((0, str(self.output_dir)),)][1][ "newest" @@ -1773,3 +1783,321 @@ class TestWorkspaceCursorEncoding(unittest.TestCase): token = _workspace_cursor(key, "scope") self.assertLessEqual(len(token), 4096) self.assertEqual(_workspace_cursor_key(token, "scope", order), key) + + +class TestWorkspaceImageTags(ImageAPITestCase): + """B3: image tags in the Tags filter and search (spec amendment S1).""" + + media = TestWorkspaceImages.media + get = TestWorkspaceImages.get + + async def tearDownAsync(self): + self.api.db.close_all() + assert_closed(self, self._temp_db.name, self._output_tmp.name) + await super().tearDownAsync() + + def tag(self, path, tags): + """Save *tags* for the current bytes of *path*, as the Tag dialog does.""" + import hashlib + + stat = path.stat() + return self.api.db.replace_image_tags( + locator=os.path.normcase(os.path.realpath(path)), + sha256=hashlib.sha256(path.read_bytes()).hexdigest(), + size=stat.st_size, + mtime_ns=stat.st_mtime_ns, + expected_epoch=self.api.db.get_execution_epoch(), + expected_tag_revision=0, + tags=tags, + ) + + def names(self, data): + return [item["path"] for item in data["images"]] + + async def test_tags_filter_admits_unlinked_image_with_current_image_tags(self): + tagged = self.media("unlinked_tagged.png", mtime=103) + self.media("plain.png", mtime=102) + linked = self.media("linked.png", mtime=101) + self.api.db.link_image_to_prompt( + self._save_prompt("Linked", tags=["other"]), str(linked) + ) + self.tag(tagged, ["dunes", "warm_light"]) + for tags in ("dunes", "dunes,warm_light", "warm_light"): + with self.subTest(tags=tags): + response, body = await self.get(tags=tags) + self.assertEqual(response.status, 200, body) + self.assertEqual( + [i["path"] for i in body["images"]], ["unlinked_tagged.png"] + ) + self.assertEqual(body["total"], 1) + self.assertEqual(body["positive_count"], 0) + # Exact names, as for prompt tags; a missing tag excludes the file. + for tags in ("Dunes", "dunes,missing", "dune"): + with self.subTest(tags=tags): + _, body = await self.get(tags=tags) + self.assertEqual(self.names(body), []) + # A category filter reads only the linked positive. + _, body = await self.get(tags="dunes", category="Night") + self.assertEqual(self.names(body), []) + + async def test_tags_filter_combines_image_and_positive_tags(self): + both = self._save_prompt("Both", category="Night", tags=["a", "b"], rating=2) + only_a = self._save_prompt("Only a", category="Night", tags=["a"], rating=5) + day_a = self._save_prompt("Day a", category="Day", tags=["a"], rating=4) + files = { + name: self.media(f"{name}.png", mtime=100 + index) + for index, name in enumerate( + ("positive_ab", "positive_a_image_b", "image_ab", "image_a", "day") + ) + } + self.api.db.link_image_to_prompt(both, str(files["positive_ab"])) + self.api.db.link_image_to_prompt(only_a, str(files["positive_a_image_b"])) + self.api.db.link_image_to_prompt(day_a, str(files["day"])) + self.tag(files["positive_a_image_b"], ["b"]) + self.tag(files["image_ab"], ["a", "b"]) + self.tag(files["image_a"], ["a"]) + self.tag(files["day"], ["b"]) + _, body = await self.get(tags="a,b", sort="name") + self.assertEqual( + self.names(body), + ["day.png", "image_ab.png", "positive_a_image_b.png", "positive_ab.png"], + ) + self.assertEqual(body["total"], 4) + self.assertEqual(body["positive_count"], 3) + _, body = await self.get(tags="a,b", category="Night", sort="name") + self.assertEqual( + self.names(body), ["positive_a_image_b.png", "positive_ab.png"] + ) + # The rating sort ranks an image-admitted linked file by its positive. + _, body = await self.get(tags="a,b", sort="rating") + self.assertEqual( + self.names(body), + ["positive_a_image_b.png", "day.png", "positive_ab.png", "image_ab.png"], + ) + # With q, a linked positive hit admits as before. + _, body = await self.get(tags="b", q="only", sort="name") + self.assertEqual(self.names(body), ["positive_a_image_b.png"]) + + async def test_q_matches_image_tag_names(self): + tagged = self.media("unlinked_tagged.png", mtime=102) + self.media("plain.png", mtime=101) + self.tag(tagged, ["Sand Dunes", "warm_light"]) + for query in ("dune", "DUNES", "sand d", "warm_"): + with self.subTest(q=query): + response, body = await self.get(q=query) + self.assertEqual(response.status, 200, body) + self.assertEqual(self.names(body), ["unlinked_tagged.png"]) + _, body = await self.get(q="plain") + self.assertEqual(self.names(body), ["plain.png"]) + _, body = await self.get(q="nothing") + self.assertEqual(self.names(body), []) + + async def test_filter_ignores_tags_of_a_changed_fingerprint(self): + moved = self.media("touched.png", mtime=102) + grown = self.media("grown.png", mtime=101) + kept = self.media("kept.png", mtime=100) + for path in (moved, grown, kept): + self.tag(path, ["dunes"]) + _, body = await self.get(tags="dunes") + self.assertEqual(self.names(body), ["touched.png", "grown.png", "kept.png"]) + os.utime(moved, (500, 500)) + with grown.open("ab") as handle: + handle.write(b"\0") + os.utime(grown, (101, 101)) + self.api.invalidate_gallery_cache() + _, body = await self.get(tags="dunes") + self.assertEqual(self.names(body), ["kept.png"]) + _, body = await self.get(q="dunes") + self.assertEqual(self.names(body), ["kept.png"]) + + async def test_stat_calls_equal_tagged_candidates_not_listing_size(self): + from py.api import images as images_module + + listed = [ + (self.output_dir / f"day-{i % 50:02}" / f"img_{i:05}.png", 1_000.0 + i) + for i in range(1_997) + ] + tagged = [self.media(f"tagged_{i}.png", mtime=5_000 + i) for i in range(3)] + listed += [(path, path.stat().st_mtime) for path in tagged] + # Tagged but outside the listing, and listed with a non-matching tag. + outside = self.media("elsewhere.png", root=self.comfy_dir) + unrelated = self.media("unrelated.png", mtime=4_000) + listed.append((unrelated, unrelated.stat().st_mtime)) + for path in tagged: + self.tag(path, ["dunes"]) + self.tag(outside, ["dunes"]) + self.tag(unrelated, ["other"]) + self.assertEqual(len(listed), 2_001) + stat_calls = 0 + real = getattr(images_module, "_current_fingerprint", None) + self.assertIsNotNone(real, "the candidate stat helper is not implemented") + + def counting(path): + nonlocal stat_calls + stat_calls += 1 + return real(path) + + with ( + patch.object( + self.api, "_get_gallery_files", AsyncMock(return_value=listed) + ), + patch.object( + self.api, "_format_workspace_page", TestWorkspaceImages._stub_page + ), + patch.object(images_module, "_current_fingerprint", counting), + ): + response, body = await self.get(tags="dunes", limit=10) + self.assertEqual(response.status, 200, body) + self.assertEqual( + sorted(i["filename"] for i in body["images"]), + ["tagged_0.png", "tagged_1.png", "tagged_2.png"], + ) + self.assertEqual(stat_calls, 3) + + async def test_plain_listing_reads_no_image_tags(self): + prompt = self._save_prompt("Linked", category="Night", tags=["a"]) + for index in range(4): + path = self.media(f"{index}.png", mtime=100 + index) + self.tag(path, ["dunes"]) + self.api.db.link_image_to_prompt(prompt, str(path)) + statements = [] + original = self.api.db.model.get_connection + + def trace_connection(): + connection = original() + connection.set_trace_callback(statements.append) + return connection + + with patch.object( + self.api.db.model, "get_connection", side_effect=trace_connection + ): + _, first = await self.get(limit=2) + _, rest = await self.get(limit=2, cursor=first["next_cursor"]) + for params in ( + {"sort": "rating"}, + {"category": "Night"}, + {"media": "image"}, + ): + response, _ = await self.get(**params) + self.assertEqual(response.status, 200) + self.assertEqual(len(first["images"] + rest["images"]), 4) + self.assertFalse( + [q for q in statements if "media_" in q or "annotation" in q], statements + ) + + async def test_image_tag_read_failure_fails_tags_and_degrades_search(self): + tagged = self.media("unlinked_tagged.png", mtime=102) + self.media("dunes_by_name.png", mtime=101) + self.tag(tagged, ["dunes"]) + with patch.object( + self.api.db, "iter_image_tag_records", side_effect=RuntimeError("boom") + ): + # As a failed link read does under a Tags filter. + response, body = await self.get(tags="dunes") + self.assertEqual((response.status, body["success"]), (500, False)) + with self.assertLogs("prompt_manager", level="WARNING") as logged: + response, body = await self.get(q="dunes") + self.assertEqual(response.status, 200, body) + self.assertEqual(self.names(body), ["dunes_by_name.png"]) + self.assertIn("image tags", "\n".join(logged.output)) + + async def test_legacy_gallery_and_search_shapes_unchanged(self): + path = self.media("legacy.png", mtime=100) + self.media("other.png", mtime=99) + self._link_image(path, "dunes at dusk") + + async def legacy(): + bodies = [] + for url, params in ( + ("/prompt_manager/images/output", {}), + ("/prompt_manager/images/search", {"q": "dunes"}), + ("/prompt_manager/images/all", {}), + ): + response = await self.client.get(url, params=params) + self.assertEqual(response.status, 200, url) + bodies.append(await response.json()) + return bodies + + before = await legacy() + self.tag(path, ["dunes", "warm_light"]) + self.tag(self.output_dir / "other.png", ["dunes"]) + self.api.invalidate_gallery_cache() + self.assertEqual(await legacy(), before) + + def test_link_facts_resolve_pairs_skip_ambiguous_and_chunk(self): + from database.models import normalize_image_path + + db = self.api.db + plain = self._save_prompt("Plain dunes", category="Day", tags=["a", "b"]) + positive = self._save_prompt("Pair positive", category="Night", tags=["p"]) + pair = db.save_prompt_pair("Pair positive", "Noise", title="Pair") + owners = (self._save_prompt("Owner one"), self._save_prompt("Owner two")) + paths = { + name: str(self.output_dir / f"{name}.png") + for name in ("plain", "paired", "ambiguous") + } + db.link_image_to_prompt(plain, paths["plain"]) + image = db.link_image_to_prompt(positive, paths["paired"]) + db.link_image_to_pair(pair, image, {"source": "fixture"}) + for owner in owners: + db.link_image_to_prompt(owner, paths["ambiguous"]) + keys = {name: normalize_image_path(path) for name, path in paths.items()} + padding = [f"/nowhere/{index}.png" for index in range(600)] + facts = db.get_image_link_facts([*padding, *keys.values(), keys["plain"]]) + self.assertEqual( + facts, + { + keys["plain"]: (True, None, plain, "Day", ("a", "b")), + keys["paired"]: (True, None, positive, "Night", ("p",)), + }, + ) + searched = db.get_image_link_facts(list(keys.values()), "dunes") + self.assertEqual( + {path: row[0] for path, row in searched.items()}, + {keys["plain"]: True, keys["paired"]: False}, + ) + self.assertEqual(db.get_image_link_facts([]), {}) + + def test_new_link_fact_reads_use_indexes(self): + from database.operations import ( + _LINK_FACTS_SQL, + _LINK_FILTER_HIT, + _POSITIVE_TAG_NAMES_SQL, + ) + + with self.api.db.model.get_connection() as conn: + conn.executemany( + "INSERT INTO prompts(id,text,hash) VALUES(?,?,?)", + [(i, f"prompt {i}", f"hash-{i}") for i in range(1, 3_001)], + ) + conn.executemany( + "INSERT INTO generated_images(prompt_id,image_path,filename," + "file_path) VALUES(?,?,?,?)", + [ + (i, f"/x/{i}.png", f"{i}.png", f"/x/{i}.png") + for i in range(1, 3_001) + ], + ) + for hit, params in ( + ("1", {}), + (_LINK_FILTER_HIT, {"pattern": "%prompt 1%"}), + ): + sql = _LINK_FACTS_SQL.format(hit=hit, marks=":path0,:path1") + bound = {**params, "path0": "/x/1.png", "path1": "/x/2.png"} + plan = [ + row[3] for row in conn.execute("EXPLAIN QUERY PLAN " + sql, bound) + ] + self.assertIn( + "SEARCH gi USING INDEX idx_image_file_path (file_path=?)", plan + ) + self.assertFalse( + [s for s in plan if s.startswith(("SCAN gi", "SCAN other"))], plan + ) + plan = [ + row[3] + for row in conn.execute( + "EXPLAIN QUERY PLAN " + _POSITIVE_TAG_NAMES_SQL.format(marks="?,?"), + (1, 2), + ) + ] + self.assertFalse([step for step in plan if step.startswith("SCAN")], plan) diff --git a/tests/test_inspection_response.py b/tests/test_inspection_response.py index 936caa68..c3f33551 100644 --- a/tests/test_inspection_response.py +++ b/tests/test_inspection_response.py @@ -823,6 +823,71 @@ def _png_with_chunks(path, chunks): return path +class TestImageTagsFacet(unittest.TestCase): + """B3: the rooted envelope's ``image_tags`` facet is validated.""" + + def encoded(self, facet): + import inspect as signature_of + + parameters = signature_of.signature(encode_inspection_response).parameters + if "image_tags" not in parameters: + self.fail("the image_tags facet is not implemented") + body = encode_inspection_response( + image=IMAGE, + root_scope="a" * 64, + file={ + "observation": "current-file", + "source_label": "output", + "media_type": "image", + "extension": ".png", + "size": 100, + "modified_time": 100, + }, + prompts=PROMPTS, + resolution={"state": "resolved", "sampler_id": "3", "reason": ""}, + inspection=inspection({"prompt": graph_text()}), + library={"status": "none"}, + image_tags=facet, + ) + return json.loads(body) + + def test_valid_facets_pass_through(self): + for facet in ( + {"state": "current", "tags": ["dunes", "warm_light"]}, + {"state": "current", "tags": []}, + {"state": "none", "tags": []}, + {"state": "unavailable", "tags": []}, + {"state": "current", "tags": [f"tag {i}" for i in range(100)]}, + ): + with self.subTest(facet=facet): + self.assertEqual(self.encoded(facet)["image_tags"], facet) + + def test_envelope_rejects_malformed_image_tags(self): + unavailable = {"state": "unavailable", "tags": []} + for facet in ( + "current", + ["dunes"], + {"state": "current"}, + {"state": "stale", "tags": []}, + {"state": "current", "tags": "dunes"}, + {"state": "current", "tags": ["dunes", 3]}, + {"state": "current", "tags": [""]}, + {"state": "current", "tags": [" dunes"]}, + {"state": "current", "tags": ["a,b"]}, + {"state": "current", "tags": ["x" * 51]}, + {"state": "current", "tags": [f"tag {i}" for i in range(101)]}, + {"state": "none", "tags": ["dunes"]}, + {"state": "unavailable", "tags": ["dunes"]}, + {"state": "current", "tags": [], "locator": "/abs/path.png"}, + ): + with self.subTest(facet=facet): + self.assertEqual(self.encoded(facet)["image_tags"], unavailable) + + def test_facet_is_absent_unless_given(self): + self.assertNotIn("image_tags", json.loads(encode(inspection({})))) + self.assertNotIn("image_tags", self.encoded(None)) + + class TestVideoProbeBoundsMatchPngBudgets(unittest.TestCase): """ffprobe escapes tag text; the stdout bounds must never gate a record that the identical PNG reads within its decoded budget (I3).""" diff --git a/utils/inspection_response.py b/utils/inspection_response.py index 84d4530b..fae4c35b 100644 --- a/utils/inspection_response.py +++ b/utils/inspection_response.py @@ -6,12 +6,16 @@ from decimal import Decimal, InvalidOperation from typing import TYPE_CHECKING, Any, Literal, Optional from .inspection_outcomes import MESSAGE_LENGTH, project_resolution +from .validators import validate_tags if TYPE_CHECKING: from .prompt_inspector import InspectionRead MAX_RESPONSE_BYTES = 8 * 1024 * 1024 MAX_LIBRARY_BYTES = 32 * 1024 +# The image tag set cap (database/annotations.py MAX_IMAGE_TAGS). +MAX_IMAGE_TAGS = 100 +_IMAGE_TAG_STATES = frozenset({"current", "none", "unavailable"}) _MAX_SAFE_INTEGER = 2**53 - 1 _RECORD_KEYS = frozenset( {"prompt", "workflow", "parameters", "comfyassets_prompt_pairs"} @@ -363,6 +367,35 @@ def _library_projection(library: dict) -> dict: return result +def _image_tag_name(value: object) -> bool: + if not isinstance(value, str) or value != value.strip() or "," in value: + return False + try: + return validate_tags([value]) + except ValueError: + return False + + +def project_image_tags(facet: object) -> dict: + """The ``image_tags`` facet as the client accepts it, else unavailable. + + ``current`` carries the byte record's names (at most 100 valid vocabulary + names); ``none`` and ``unavailable`` carry no names. Anything else, extra + keys included, is reported as ``unavailable`` rather than passed on. + """ + unavailable = {"state": "unavailable", "tags": []} + if not isinstance(facet, dict) or set(facet) != {"state", "tags"}: + return unavailable + state, tags = facet["state"], facet["tags"] + if state not in _IMAGE_TAG_STATES or not isinstance(tags, list): + return unavailable + if state != "current" and tags: + return unavailable + if len(tags) > MAX_IMAGE_TAGS or not all(_image_tag_name(tag) for tag in tags): + return unavailable + return {"state": state, "tags": list(tags)} + + def _bounded_message(value: Any) -> Optional[str]: """A role line as the client accepts it: 1 to 512 UTF-16 code units.""" if not isinstance(value, str) or not value: @@ -394,6 +427,7 @@ def _encode_inspection( resolution: dict, inspection: "InspectionRead | None", library: dict, + image_tags: dict | None = None, ) -> bytes: """Serialize exactly once per budget step, preserving identity and roles.""" metadata = _metadata(inspection) @@ -407,6 +441,8 @@ def _encode_inspection( "metadata": metadata, "library": _library_projection(library), } + if image_tags is not None: + response["image_tags"] = project_image_tags(image_tags) body = _bytes(response) # Omission order keeps prompt evidence longest: the workflow graph first, # then the prompt graph, then all raw records atomically. @@ -438,8 +474,12 @@ def encode_inspection_response( resolution: dict, inspection: "InspectionRead | None", library: dict, + image_tags: dict | None = None, ) -> bytes: - """Preserve the existing rooted wire and its bounded metadata projection.""" + """Preserve the existing rooted wire and its bounded metadata projection. + + *image_tags*, when given, adds the validated ``image_tags`` facet (B3). + """ return _encode_inspection( identity={"root_scope": root_scope, "image": image}, file=file, @@ -447,6 +487,7 @@ def encode_inspection_response( resolution=resolution, inspection=inspection, library=library, + image_tags=image_tags, )