From f3366a13f896a086aa7521adc488fee7da757420 Mon Sep 17 00:00:00 2001 From: Vito Sansevero Date: Thu, 22 Oct 2026 15:56:58 -0700 Subject: [PATCH] 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. --- TESTING.md | 24 ++ __init__.py | 14 + database/operations.py | 2 + database/review.py | 157 +++++++++ py/api/admin.py | 102 +++++- py/api/annotations.py | 20 +- tests/browser/server.py | 2 + tests/test_api_admin.py | 10 +- tests/test_review.py | 727 ++++++++++++++++++++++++++++++++++++++++ utils/image_monitor.py | 38 ++- utils/review_capture.py | 199 +++++++++++ 11 files changed, 1276 insertions(+), 19 deletions(-) create mode 100644 database/review.py create mode 100644 tests/test_review.py create mode 100644 utils/review_capture.py diff --git a/TESTING.md b/TESTING.md index 86931b7..b7342c0 100644 --- a/TESTING.md +++ b/TESTING.md @@ -1610,6 +1610,30 @@ every evidence document in one streamed pass at startup and on restore (about `media_review_decisions` and `media_review_items` before the annotation tables, never the metadata row. +Needs review items (#249, `utils/review_capture.py`, `database/review.py`, +`tests/test_review.py`): the output scan (Activity's Scan for metadata and the +legacy `/scan`) and the live monitor hand every file the ingestion resolver +finds `ambiguous` to `record_ambiguous_file`, whether or not the browser parser +then links it (`conflict` and `unresolved` files are not queued). It captures +the file (scope `needs-review`), reads its metadata again from the capture, +resolves it with the default profile and stores only an ambiguous result, +bound to the capture's byte record (created or refreshed like an observe) and +fenced by the capture's storage epoch. Evidence is the resolver reason plus one +`{sampler_id, state, reason, positive, negative, hash}` per selectable +candidate (nested ambiguous stages flattened, duplicates dropped; each side +`{text, storage_text, provenance}`; `hash` = SHA-256 of the canonical candidate +without it). No candidate, more than 16, or evidence over 8 MiB is a `limited` +item without candidates. Recording is idempotent per `(byte_id, +resolver_version, evidence_hash)` and never touches an existing item or its +decisions; changed bytes add a new pending item and keep the old one. The +capture is released before the write; no model is loaded and no pixel is +decoded. Scan statuses, counters, outcome rows and the legacy SSE events are +unchanged; the job's `scope` gains `review_items` (new items) and +`review_errors`, and the log ends with `? N images added to Needs review` (and +`! N images could not be added to Needs review` on failures, e.g. capture +capacity). The monitor records after its fallback link, best effort: a failure +is logged and the link stays. A new item moves `review_index_revision`. + Capture store (#249, `utils/media_capture.py`, `tests/test_media_capture.py`): one explicit action captures one rooted media file into an owner-only `prompt-manager-captures-*` temp directory, streaming it in 1 MiB chunks while diff --git a/__init__.py b/__init__.py index 140f7ea..ad4e504 100644 --- a/__init__.py +++ b/__init__.py @@ -122,6 +122,20 @@ try: _global_prompt_tracker = get_prompt_tracker(_global_db) _global_image_monitor = get_image_monitor(_global_db, _global_prompt_tracker) + # Ambiguous new outputs go to Needs review through the API's capture store + # (the one the scan and the Tag dialog use). Recording stops before that + # store is closed at shutdown. + _review_api = globals().get("api") + if _review_api is not None: + _global_image_monitor.set_review_recorder(_review_api.record_review_file) + _review_host = getattr(config.server_instance, "app", None) + if _review_host is not None and hasattr(_review_host, "on_cleanup"): + + async def _stop_review_recording(app, monitor=_global_image_monitor): + monitor.set_review_recorder(None) + + _review_host.on_cleanup.insert(0, _stop_review_recording) + # Start monitoring _global_image_monitor.start_monitoring() _init_logger.info("Global image monitoring system started successfully") diff --git a/database/operations.py b/database/operations.py index 28bf35d..bce3446 100644 --- a/database/operations.py +++ b/database/operations.py @@ -13,6 +13,7 @@ from .pairs import PairDependencyError, PairOperations from .library import LibraryOperations, titled from .image_favorites import ImageFavoriteOperations from .annotations import AnnotationOperations +from .review import ReviewOperations from .saved_views import SavedViewOperations from .library_commands import LibraryCommandOperations from .vocabulary import VocabularyOperations @@ -341,6 +342,7 @@ class PromptDatabase( HistoryReadOperations, ImageFavoriteOperations, AnnotationOperations, + ReviewOperations, SavedViewOperations, ): """Database operations class for managing prompts.""" diff --git a/database/review.py b/database/review.py new file mode 100644 index 0000000..73428cf --- /dev/null +++ b/database/review.py @@ -0,0 +1,157 @@ +"""Needs review items (#249): persist and read capture-bound ambiguous evidence. + +An item is one ambiguous resolver result for one annotation byte record (the +whole-file SHA-256 of one physical file at one locator), keyed by +``(byte_id, resolver_version, evidence_hash)``. Recording is idempotent: the +same evidence for the same bytes is a duplicate and never touches an existing +item, its state, revision or decisions. Changed bytes are a new byte record +and therefore a new pending item; the old one stays. + +Callers pass verified facts (the capture's digest, size and ``mtime_ns``); +nothing here reads the filesystem. Writes run in one ``_pair_transaction`` +fenced by the storage epoch the capture was taken under. +""" + +import json +import sqlite3 +from typing import Iterator + +from .annotation_schema import ( + BYTES_TABLE, + FILES_TABLE, + METADATA_TABLE as ANNOTATION_METADATA_TABLE, +) +from .annotations import _new_id, _validate_locator, _validate_target +from .models import utc_now_iso +from .review_schema import ( + ITEMS_TABLE, + MAX_CANDIDATES, + METADATA_TABLE, + canonical_evidence, + evidence_digest, +) + +# Only this module's table constants are interpolated into the SQL below. +_INSERT_ITEM = ( + f"INSERT OR IGNORE INTO {ITEMS_TABLE}(id,byte_id,resolver_version," # nosec B608 + "evidence_hash,evidence,candidate_count,limited,state,revision,created_at," + "updated_at) VALUES(?,?,?,?,?,?,?,'pending',0,?,?)" +) +_BUMP_REVIEW_INDEX = ( + f"UPDATE {METADATA_TABLE} SET " # nosec B608 + "review_index_revision=review_index_revision+1 WHERE singleton_id=1" +) +_REVIEW_INDEX = ( + f"SELECT review_index_revision FROM {METADATA_TABLE} " # nosec B608 + "WHERE singleton_id=1" +) +_BUMP_IMAGE_INDEX = ( + f"UPDATE {ANNOTATION_METADATA_TABLE} SET " # nosec B608 + "image_index_revision=image_index_revision+1 WHERE singleton_id=1" +) +_ITEMS_OF_LOCATOR = ( + "SELECT i.id,i.byte_id,i.resolver_version,i.evidence_hash," # nosec B608 + "i.evidence,i.candidate_count,i.limited,i.state,i.revision,i.created_at," + f"i.updated_at,b.sha256 FROM {FILES_TABLE} f " + f"JOIN {BYTES_TABLE} b ON b.file_id=f.id " + f"JOIN {ITEMS_TABLE} i ON i.byte_id=b.id " + "WHERE f.locator=? ORDER BY i.created_at,i.id" +) +_ITEM_COLUMNS = ( + "id", + "byte_id", + "resolver_version", + "evidence_hash", + "evidence", + "candidate_count", + "limited", + "state", + "revision", + "created_at", + "updated_at", + "sha256", +) + + +class ReviewOperations: + """PromptDatabase mixin: insert and read Needs review items.""" + + def record_review_item( + self, + *, + locator: str, + sha256: str, + size: int, + mtime_ns: int, + epoch: str, + resolver_version: int, + evidence: dict, + candidate_count: int, + limited: bool, + ) -> bool: + """Store one ambiguous result for these exact bytes; True when new. + + The byte record is created (or its verified fingerprint refreshed) + like an observe. Raises ``LibraryCommandError("storage-changed")`` when + the storage epoch is no longer ``epoch`` (a restore since the capture) + and ``ValueError`` for invalid facts or evidence. + """ + _validate_target(locator, sha256, size, mtime_ns) + if type(resolver_version) is not int or resolver_version < 1: + raise ValueError("resolver_version must be a positive integer") + if ( + type(candidate_count) is not int + or not 0 <= candidate_count <= MAX_CANDIDATES + ): + raise ValueError(f"candidate_count must be between 0 and {MAX_CANDIDATES}") + if not isinstance(limited, bool) or (limited and candidate_count): + raise ValueError("A limited item carries no candidates") + text = canonical_evidence(evidence) + evidence_hash = evidence_digest(text) + with self._pair_transaction() as conn: + self._library_epoch(conn, epoch) + now = utc_now_iso() + file_id = self._annotation_owner(conn, locator, now) + byte_id, _revision, created, moved = self._annotation_byte( + conn, file_id, sha256, size, mtime_ns, now + ) + if created or moved: + conn.execute(_BUMP_IMAGE_INDEX) + cursor = conn.execute( + _INSERT_ITEM, + ( + _new_id(), + byte_id, + resolver_version, + evidence_hash, + text, + candidate_count, + int(limited), + now, + now, + ), + ) + inserted = cursor.rowcount > 0 + if inserted: + conn.execute(_BUMP_REVIEW_INDEX) + return inserted + + def review_index_revision(self) -> int: + """The queue's revision (moved by every queue change).""" + row = self.model.get_connection().execute(_REVIEW_INDEX).fetchone() + if row is None: + raise sqlite3.DatabaseError("Missing review queue metadata") + return row[0] + + def review_items_for(self, locator: str) -> list: + """Every item of the file at ``locator``, oldest first, evidence parsed.""" + _validate_locator(locator) + return list(self._iter_review_items(locator)) + + def _iter_review_items(self, locator: str) -> Iterator[dict]: + conn = self.model.get_connection() + for row in conn.execute(_ITEMS_OF_LOCATOR, (locator,)): + item = dict(zip(_ITEM_COLUMNS, tuple(row))) + item["evidence"] = json.loads(item["evidence"]) + item["limited"] = bool(item["limited"]) + yield item diff --git a/py/api/admin.py b/py/api/admin.py index bc900ee..01d1a4b 100644 --- a/py/api/admin.py +++ b/py/api/admin.py @@ -43,6 +43,7 @@ try: parse_js_json, ) from ...utils.prompt_pair_resolver import resolve_prompt_pair + from ...utils.review_capture import record_ambiguous_file except ImportError: from database.operations import PromptDatabase from database.pair_ingestion import ingest_resolved_pair @@ -53,6 +54,7 @@ except ImportError: parse_js_json, ) from utils.prompt_pair_resolver import resolve_prompt_pair + from utils.review_capture import record_ambiguous_file try: @@ -1585,10 +1587,13 @@ class AdminRoutesMixin: known = set() if media_files and job.scope.get("files") != "all": known = await self._known_scan_files(job, media_files) + review = {"recorded": 0, "errors": 0} + job.scope.update(review_items=0, review_errors=0) async for event in self._scan_batches( - media_files, counts, job, roots, known + media_files, counts, job, roots, known, review ): yield event + self._log_review_items(job, review) self.logger.info( f"Scan completed: processed={counts['processed']}, " @@ -1625,12 +1630,16 @@ class AdminRoutesMixin: sizes = {str(media_file): _file_size(media_file) for media_file in media_files} return self.db.get_known_resolved_paths(sizes) - async def _scan_batches(self, media_files, counts, job, roots=(), known=()): + async def _scan_batches( + self, media_files, counts, job, roots=(), known=(), review=None + ): """Process ``media_files`` in executor batches, updating ``counts`` in place. Records each file's outcome in ``job`` and yields one progress event per batch. Files in ``known`` (path strings) are not read: they count - as processed and unchanged and are recorded as skipped. + as processed and unchanged and are recorded as skipped. ``review`` + counts the Needs review items the batches record (``job.scope`` + ``review_items`` / ``review_errors`` follow it after every batch). """ total = len(media_files) started = time.monotonic() @@ -1653,9 +1662,15 @@ class AdminRoutesMixin: batch_results, roots=roots, outcomes=outcomes, + review=review, ) for key in counts: counts[key] += batch_counts[key] + if review is not None: + job.scope.update( + review_items=review["recorded"], + review_errors=review["errors"], + ) for group, row, log_line in outcomes: job.record(group, row, log_line) @@ -1705,12 +1720,17 @@ class AdminRoutesMixin: return map_parallel(extract, file_batch, PromptManagerConfig.WORKER_THREADS) - def _ingest_scan_batch_sync(self, batch_results, roots=(), outcomes=None): + def _ingest_scan_batch_sync( + self, batch_results, roots=(), outcomes=None, review=None + ): """Store the prompts found in a metadata batch (blocking DB work). Returns the legacy counts. When ``outcomes`` is a list, every processed file appends its ``(group, row, log_line)`` for the scan job, with the - path relative to its root in ``roots`` (see ``_scan_roots``). + path relative to its root in ``roots`` (see ``_scan_roots``). A file the + resolver finds ambiguous is also recorded for Needs review (best + effort, counted in ``review`` when given); statuses and counts do not + depend on it. """ counts = dict.fromkeys(SCAN_COUNT_KEYS, 0) @@ -1722,13 +1742,18 @@ class AdminRoutesMixin: if self._scan_job and self._scan_job.cancelled: break counts["processed"] += 1 + ambiguous = [] try: - outcome = self._ingest_scanned_file(media_file, metadata, detailed=True) + outcome = self._ingest_scanned_file( + media_file, metadata, detailed=True, ambiguous=ambiguous + ) except Exception as e: self.logger.error(f"Error processing {media_file.name}: {e}") counts["failed"] += 1 note(media_file, "error") continue + if ambiguous: + self._record_scanned_review(media_file, roots, review) if outcome is None: note(media_file, None) continue @@ -1759,13 +1784,17 @@ class AdminRoutesMixin: return None return positive.strip() or None - def _ingest_scanned_file(self, media_file, metadata, detailed=False): + def _ingest_scanned_file( + self, media_file, metadata, detailed=False, ambiguous=None + ): """Enrich exact pairs; retain legacy positive-only import for absent pairs. Returns the file's scan status, or None without prompt metadata. With ``detailed``, a stored or linked file returns a dict instead: its ``status``, the prompt's display ``title`` and, for a resolved pair, - the pair counts. + the pair counts. When the resolver finds the file ambiguous (whether + or not the browser parser then links it) and ``ambiguous`` is a list, + ``media_file`` is appended to it. """ if not metadata: return None @@ -1781,6 +1810,8 @@ class AdminRoutesMixin: return None _with_file_size(canonical, media_file) resolution = resolve_prompt_pair(canonical) + if resolution["state"] == "ambiguous" and ambiguous is not None: + ambiguous.append(media_file) if resolution["state"] == "resolved": result = ingest_resolved_pair( self.db, str(media_file), canonical, resolution @@ -1823,6 +1854,61 @@ class AdminRoutesMixin: return status return {"status": status, "title": self._scan_title(media_file, prompt_text)} + # ── Needs review items (the scan's and the monitor's ambiguous files) ── + + def _record_scanned_review(self, media_file, roots, review=None): + """Record one ambiguous scanned file; a failure is logged and counted.""" + try: + outcome = self._record_review_sync(media_file, roots) + except Exception: + self.logger.warning( + f"{media_file.name} could not be added to Needs review", + exc_info=True, + ) + if review is not None: + review["errors"] += 1 + return + if review is not None and outcome in ("recorded", "limited"): + review["recorded"] += 1 + + def _record_review_sync(self, media_file, roots): + """``record_ambiguous_file`` for a file under one of ``roots`` (blocking).""" + relative, _label, index = _scan_location(media_file, roots) + root = next((r for r, _l, i in roots if i == index), None) + if index is None or root is None: + return "skipped" + return record_ambiguous_file( + self.db, self._capture_store(), Path(root), index, relative + ) + + def record_review_file(self, image_path): + """Needs review recording for the live monitor (its worker thread). + + Maps ``image_path`` to the gallery root that lists it (the listing's + ``root_index``); a file under no current root is ``skipped``. Errors + propagate: the monitor logs them and keeps its link. + """ + full = os.path.abspath(str(image_path)) + located = Path(os.path.realpath(os.path.dirname(full))) / os.path.basename(full) + return self._record_review_sync( + located, _scan_roots(self._get_all_output_dirs()) + ) + + @staticmethod + def _log_review_items(job, review): + """The scan log's Needs review summary lines (none when nothing moved).""" + if review["recorded"]: + job.log_line( + f"{SCAN_LOG_PREFIXES['review']} " + f"{_plural(review['recorded'], 'image')} added to Needs review" + ) + if review["errors"]: + job.log_line( + f"{SCAN_LOG_PREFIXES['errors']} " + f"{_plural(review['errors'], 'image')} could not be added to " + "Needs review" + ) + def _scan_title(self, media_file, text): """The linked prompt's Library title (its stored title, else ``text``'s first line), as the Library and the inspector name it.""" diff --git a/py/api/annotations.py b/py/api/annotations.py index 1e08e8c..caadb24 100644 --- a/py/api/annotations.py +++ b/py/api/annotations.py @@ -13,6 +13,7 @@ import datetime import functools import os import re +import threading from pathlib import Path from aiohttp import web @@ -73,6 +74,9 @@ _HEX64 = re.compile(r"[0-9a-f]{64}") _UUID = re.compile(r"[0-9a-f]{8}(?:-[0-9a-f]{4}){3}-[0-9a-f]{12}") # secrets.token_urlsafe(32) is 43 characters; the bound only rejects junk. _TOKEN = re.compile(r"[A-Za-z0-9_-]{1,128}") +# The store is created on first use by the event loop, the scan's executor or +# the monitor's worker (Needs review), so its creation is serialized. +_CAPTURE_STORE_LOCK = threading.Lock() _PREVIEW_TYPES = { ".png": "image/png", @@ -259,14 +263,16 @@ class AnnotationRoutesMixin: Created lazily so an API object that never observes a file (most tests, a server that never opens the Tag dialog) owns no temp directory. A restore of the database dooms every capture (``invalidate_all``). + Thread-safe: Needs review recording calls it off the event loop. """ - if self._captures is None: - self._captures = CaptureStore() - if self._captures_unsubscribe is None: - self._captures_unsubscribe = subscribe_history_restore( - self.db.model.db_path, self._invalidate_captures - ) - return self._captures + with _CAPTURE_STORE_LOCK: + if self._captures is None: + self._captures = CaptureStore() + if self._captures_unsubscribe is None: + self._captures_unsubscribe = subscribe_history_restore( + self.db.model.db_path, self._invalidate_captures + ) + return self._captures def _invalidate_captures(self) -> None: store = self._captures diff --git a/tests/browser/server.py b/tests/browser/server.py index 6cf7fe5..0bb5230 100644 --- a/tests/browser/server.py +++ b/tests/browser/server.py @@ -548,6 +548,8 @@ def create_app(root: Path) -> web.Application: return web.json_response("dark") async def cleanup(application: web.Application) -> None: + # A scan of an ambiguous fixture file captures it for Needs review. + api.close_annotation_captures() api.db.close_all() seed() diff --git a/tests/test_api_admin.py b/tests/test_api_admin.py index bddcc39..ec75051 100644 --- a/tests/test_api_admin.py +++ b/tests/test_api_admin.py @@ -88,6 +88,8 @@ class AdminAPITestCase(AioHTTPTestCase): sys.modules.pop("folder_paths", None) else: sys.modules["folder_paths"] = self._orig_folder_paths + # A scan of an ambiguous file captures it for Needs review. + self.api.close_annotation_captures() PromptModel.close_all_instances() # restore may have swapped api.db assert_closed(self, self.tmpdir) shutil.rmtree(self.tmpdir, ignore_errors=True) @@ -2044,7 +2046,7 @@ class TestActivityJobs(AdminAPITestCase): for name in names: (self.output_dir / name).write_bytes(b"not an image") - def fake_ingest(media_file, metadata, detailed=False): + def fake_ingest(media_file, metadata, detailed=False, ambiguous=None): key = Path(media_file).relative_to(self.output_dir).as_posix() outcome = names[key] if isinstance(outcome, Exception): @@ -2173,8 +2175,10 @@ class TestActivityJobs(AdminAPITestCase): first_done, release = threading.Event(), threading.Event() original = self.api._ingest_scan_batch_sync - def gated(batch_results, roots=(), outcomes=None): - counts = original(batch_results, roots=roots, outcomes=outcomes) + def gated(batch_results, roots=(), outcomes=None, review=None): + counts = original( + batch_results, roots=roots, outcomes=outcomes, review=review + ) first_done.set() release.wait(5) return counts diff --git a/tests/test_review.py b/tests/test_review.py new file mode 100644 index 0000000..a9ebdd1 --- /dev/null +++ b/tests/test_review.py @@ -0,0 +1,727 @@ +"""Needs review items (#249 E2): capture-bound ambiguous evidence. + +The scan and the live monitor persist one review item per ambiguous file: +the file is captured, its metadata is read again from the capture, the +default (ingestion) resolver runs on it, and only an ``ambiguous`` result is +stored, bound to the capture's byte record. Statuses, counters and links of +both producers stay what they were. +""" + +import asyncio +import hashlib +import json +import logging +import os +import sqlite3 +import unittest +from pathlib import Path +from unittest.mock import patch + +try: + from tests.open_handles import assert_closed + from tests.test_api_admin import AdminAPITestCase, _sse_payloads +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed + from test_api_admin import AdminAPITestCase, _sse_payloads + +from PIL import Image, PngImagePlugin + +from database.operations import PromptDatabase +from database.review_schema import canonical_evidence, evidence_digest +from py.api import PromptManagerAPI +from py.config import GalleryConfig +from utils.hashing import generate_prompt_hash +from utils.image_monitor import ImageGenerationHandler +from utils.media_capture import CaptureCapacity, CaptureStore +from utils.prompt_tracker import PromptTracker + +try: # guarded so a missing module fails assertions, not test loading + from utils import review_capture +except ImportError: + review_capture = None + +NOW = "2026-10-22T00:00:00+00:00" +# Two negative prompt delimiters: the A1111 record is ambiguous. +DELIMITERS = ( + "cat\nNegative prompt: blur\nNegative prompt: noise\nSteps: 20, Sampler: Euler" +) + + +def _sampler(positive, negative="bad"): + return {"positive": positive, "negative": negative} + + +def ambiguous_graph(*positives): + """One KSampler per positive text, each with its own encoder: ambiguous.""" + graph = { + "1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "m"}}, + "2": {"class_type": "CLIPTextEncode", "inputs": {"text": "bad"}}, + } + for index, text in enumerate(positives or ("cat", "dog")): + encoder, sampler = str(10 + 2 * index), str(11 + 2 * index) + graph[encoder] = {"class_type": "CLIPTextEncode", "inputs": {"text": text}} + graph[sampler] = { + "class_type": "KSampler", + "inputs": { + "model": ["1", 0], + "positive": [encoder, 0], + "negative": ["2", 0], + }, + } + return graph + + +def resolved_graph(text="a red square"): + return ambiguous_graph(text) + + +def write_png(path, *, prompt=None, parameters=None, color=(255, 0, 0)): + info = PngImagePlugin.PngInfo() + if prompt is not None: + info.add_text("prompt", json.dumps(prompt)) + if parameters is not None: + info.add_text("parameters", parameters) + Image.new("RGB", (2, 2), color).save(str(path), pnginfo=info) + + +def digest(path): + return hashlib.sha256(Path(path).read_bytes()).hexdigest() + + +class ReviewTestCase(AdminAPITestCase): + """Scan helpers over a capture store kept inside the test directory.""" + + async def get_application(self): + app = await super().get_application() + self.capture_parent = Path(self.tmpdir) / "captures" + self.api._captures = CaptureStore(temp_parent=self.capture_parent) + return app + + async def tearDownAsync(self): + self.api.close_annotation_captures() + await super().tearDownAsync() + + def _roots(self): + GalleryConfig.MONITORING_DIRECTORIES = [self._root(self.output_dir)] + + async def _scan(self, scope="new"): + resp = await self.client.request( + "POST", "/prompt_manager/jobs/scan", json={"scope": scope} + ) + data = await resp.json() + self.assertEqual(resp.status, 202, data) + await asyncio.wait_for(self.api._scan_job.task, 20) + detail = await self.client.request( + "GET", f"/prompt_manager/jobs/{data['job']['id']}" + ) + return (await detail.json())["job"] + + def _items(self): + conn = self.api.db.model.get_connection() + return [ + dict(row) + for row in conn.execute( + "SELECT i.*, b.sha256, b.size, b.mtime_ns, f.locator " + "FROM media_review_items i " + "JOIN media_annotation_bytes b ON b.id=i.byte_id " + "JOIN media_annotation_files f ON f.id=b.file_id " + "ORDER BY i.created_at, i.id" + ) + ] + + def _review_revision(self): + conn = self.api.db.model.get_connection() + return conn.execute( + "SELECT review_index_revision FROM media_review_metadata" + ).fetchone()[0] + + def _capture_files(self): + if not self.capture_parent.exists(): + return [] + return [p for p in self.capture_parent.rglob("*") if p.is_file()] + + def _require_module(self): + self.assertIsNotNone(review_capture, "utils.review_capture is not implemented") + + +class TestScanRecording(ReviewTestCase): + async def test_scan_persists_only_ambiguous_files_with_capture_digest(self): + self._require_module() + self._roots() + ambiguous = self.output_dir / "sub" / "two.png" + ambiguous.parent.mkdir() + write_png(ambiguous, prompt=ambiguous_graph("cat", "dog")) + write_png(self.output_dir / "resolved.png", prompt=resolved_graph()) + write_png(self.output_dir / "plain.png") + write_png( + self.output_dir / "a1111.png", + parameters="cat\nNegative prompt: blur\nSteps: 20, Sampler: Euler", + ) + + job = await self._scan() + + items = self._items() + self.assertEqual(len(items), 1, items) + item = items[0] + self.assertEqual(item["sha256"], digest(ambiguous)) + self.assertEqual(item["size"], ambiguous.stat().st_size) + self.assertEqual(item["mtime_ns"], ambiguous.stat().st_mtime_ns) + self.assertEqual( + item["locator"], os.path.normcase(os.path.realpath(str(ambiguous))) + ) + self.assertEqual( + (item["state"], item["revision"], item["limited"]), ("pending", 0, 0) + ) + self.assertEqual(item["resolver_version"], 1) + self.assertEqual(item["candidate_count"], 2) + self.assertEqual(item["evidence_hash"], evidence_digest(item["evidence"])) + evidence = json.loads(item["evidence"]) + self.assertEqual(canonical_evidence(evidence), item["evidence"]) + self.assertEqual(evidence["resolver_version"], 1) + self.assertEqual( + sorted(c["positive"]["text"] for c in evidence["candidates"]), + ["cat", "dog"], + ) + for candidate in evidence["candidates"]: + self.assertEqual( + set(candidate), + {"hash", "sampler_id", "state", "reason", "positive", "negative"}, + ) + self.assertEqual( + set(candidate["positive"]), {"text", "storage_text", "provenance"} + ) + self.assertEqual(candidate["negative"]["text"], "bad") + body = {k: v for k, v in candidate.items() if k != "hash"} + self.assertEqual( + candidate["hash"], + hashlib.sha256(canonical_evidence(body).encode()).hexdigest(), + ) + self.assertEqual(self._review_revision(), 1) + self.assertEqual(job["scope"]["review_items"], 1) + self.assertIn("? 1 image added to Needs review", "\n".join(job["log"])) + self.assertEqual(self._capture_files(), []) + self.assertEqual(self.api._captures.stats()["captures"], 0) + + async def test_scan_statuses_counters_and_outcome_rows_unchanged(self): + self._require_module() + self._roots() + write_png(self.output_dir / "two.png", prompt=ambiguous_graph("cat", "dog")) + write_png(self.output_dir / "one.png", prompt=resolved_graph()) + write_png( + self.output_dir / "delims.png", + parameters=DELIMITERS, + ) + write_png(self.output_dir / "none.png") + files = sorted(self.output_dir.iterdir()) + roots = self._scan_roots() + + def run(api, recorder=None): + batch = api._extract_batch_metadata_sync(files) + outcomes = [] + if recorder is None: + counts = api._ingest_scan_batch_sync(batch, roots, outcomes) + else: + from py.api import admin + + with patch.object(admin, "record_ambiguous_file", recorder): + counts = api._ingest_scan_batch_sync(batch, roots, outcomes) + prompts = api.db.get_recent_prompts()["prompts"] + return counts, outcomes, sorted(p["text"] for p in prompts) + + other_db = PromptDatabase(os.path.join(self.tmpdir, "other.db")) + other = PromptManagerAPI(db=other_db) + + def broken(*args, **kwargs): + raise RuntimeError("recording failed") + + recorded = run(self.api) + with self.assertLogs("prompt_manager.api", logging.WARNING) as logs: + without = run(other, broken) + self.assertEqual(recorded, without) + self.assertEqual(recorded[0]["ambiguous"], 1) + self.assertEqual(recorded[0]["added"], 2) + self.assertTrue(any("Needs review" in line for line in logs.output)) + self.assertEqual(len(self._items()), 2) + other.close_annotation_captures() + other_db.close_all() + + def _scan_roots(self): + from py.api.admin import _scan_roots + + return _scan_roots(self.api._get_all_output_dirs()) + + async def test_new_scope_rereads_ambiguous_files_and_stays_idempotent(self): + self._require_module() + self._roots() + write_png(self.output_dir / "two.png", prompt=ambiguous_graph("cat", "dog")) + write_png(self.output_dir / "one.png", prompt=resolved_graph()) + first = await self._scan("new") + revision = self._review_revision() + reads = [] + original = self.api._extract_comfyui_metadata + + def spy(path): + reads.append(os.path.basename(path)) + return original(path) + + self.api._extract_comfyui_metadata = spy + second = await self._scan("new") + + self.assertEqual(reads, ["two.png"]) + self.assertEqual(first["scope"]["review_items"], 1) + self.assertEqual(second["scope"]["review_items"], 0) + self.assertNotIn("added to Needs review", "\n".join(second["log"])) + self.assertEqual(len(self._items()), 1) + self.assertEqual(self._review_revision(), revision) + self.assertEqual(self._capture_files(), []) + + async def test_repeat_scan_keeps_decisions(self): + self._require_module() + self._roots() + write_png(self.output_dir / "two.png", prompt=ambiguous_graph("cat", "dog")) + await self._scan("all") + item = self._items()[0] + candidate = json.loads(item["evidence"])["candidates"][0]["hash"] + conn = self.api.db.model.get_connection() + with conn: + conn.execute( + "UPDATE media_review_items SET state='selected', revision=1 WHERE id=?", + (item["id"],), + ) + conn.execute( + "INSERT INTO media_review_decisions(id,item_id,decision_revision," + "created_at," + "outcome,candidate_hash,evidence) VALUES(?,?,?,?,?,?,?)", + ( + "00000000-0000-4000-8000-0000000000d1", + item["id"], + 1, + NOW, + "selected", + candidate, + '{"kind":"user-selection"}', + ), + ) + + job = await self._scan("all") + + items = self._items() + self.assertEqual(len(items), 1) + self.assertEqual((items[0]["state"], items[0]["revision"]), ("selected", 1)) + self.assertEqual(items[0]["updated_at"], item["updated_at"]) + decisions = conn.execute("SELECT id FROM media_review_decisions").fetchall() + self.assertEqual(len(decisions), 1) + self.assertEqual(job["scope"]["review_items"], 0) + + async def test_changed_bytes_create_new_pending_item_and_keep_old(self): + self._require_module() + self._roots() + path = self.output_dir / "two.png" + write_png(path, prompt=ambiguous_graph("cat", "dog")) + await self._scan() + old = self._items()[0] + write_png(path, prompt=ambiguous_graph("cat", "dog"), color=(0, 0, 255)) + self.assertNotEqual(digest(path), old["sha256"]) + + job = await self._scan() + + items = self._items() + self.assertEqual(len(items), 2) + self.assertEqual(items[0]["id"], old["id"]) + self.assertEqual(items[0]["sha256"], old["sha256"]) + self.assertEqual(items[1]["sha256"], digest(path)) + self.assertEqual(items[1]["state"], "pending") + self.assertEqual(items[0]["locator"], items[1]["locator"]) + self.assertEqual(job["scope"]["review_items"], 1) + + async def test_parameters_delimiter_ambiguity_is_limited_without_candidates(self): + self._require_module() + self._roots() + write_png( + self.output_dir / "delims.png", + parameters=DELIMITERS, + ) + many = [f"prompt {n}" for n in range(17)] + write_png(self.output_dir / "many.png", prompt=ambiguous_graph(*many)) + + job = await self._scan() + + items = {Path(item["locator"]).name: item for item in self._items()} + self.assertEqual(set(items), {"delims.png", "many.png"}) + for item in items.values(): + self.assertEqual((item["limited"], item["candidate_count"]), (1, 0)) + evidence = json.loads(item["evidence"]) + self.assertEqual(evidence["candidates"], []) + self.assertIs(evidence["limited"], True) + self.assertEqual( + json.loads(items["many.png"]["evidence"])["candidate_total"], 17 + ) + self.assertEqual(job["scope"]["review_items"], 2) + + async def test_capacity_errors_are_logged_and_counted_not_fatal(self): + self._require_module() + self._roots() + write_png(self.output_dir / "two.png", prompt=ambiguous_graph("cat", "dog")) + + def full(*args, **kwargs): + raise CaptureCapacity("Too many captures are in progress") + + with patch.object(self.api._captures, "capture", side_effect=full): + job = await self._scan() + + self.assertEqual(job["state"], "finished") + # The browser parser links this ambiguous file: Activity counts it Added. + self.assertEqual(job["counters"]["added"], 1) + self.assertEqual(job["scope"]["review_items"], 0) + self.assertEqual(job["scope"]["review_errors"], 1) + self.assertIn( + "! 1 image could not be added to Needs review", "\n".join(job["log"]) + ) + self.assertEqual(self._items(), []) + + async def test_legacy_scan_events_are_unchanged(self): + self._require_module() + self._roots() + write_png(self.output_dir / "two.png", prompt=ambiguous_graph("cat", "dog")) + resp = await self.client.request("POST", "/prompt_manager/scan") + payloads = _sse_payloads(await resp.text()) + final = payloads[-1] + self.assertEqual(final["type"], "complete") + self.assertNotIn("review_items", final) + self.assertEqual(len(self._items()), 1) + + +class TestRecordAmbiguousFile(ReviewTestCase): + def _record(self, relative): + return review_capture.record_ambiguous_file( + self.api.db, self.api._captures, self.output_dir, 0, relative + ) + + async def test_record_outcomes(self): + self._require_module() + write_png(self.output_dir / "two.png", prompt=ambiguous_graph("cat", "dog")) + write_png(self.output_dir / "one.png", prompt=resolved_graph()) + write_png( + self.output_dir / "delims.png", + parameters=DELIMITERS, + ) + self.assertEqual(self._record("two.png"), "recorded") + self.assertEqual(self._record("two.png"), "duplicate") + self.assertEqual(self._record("one.png"), "skipped") + self.assertEqual(self._record("missing.png"), "skipped") + self.assertEqual(self._record("delims.png"), "limited") + self.assertEqual(self._record("delims.png"), "duplicate") + self.assertEqual(self._review_revision(), 2) + + async def test_restore_between_capture_and_write_skips(self): + self._require_module() + write_png(self.output_dir / "two.png", prompt=ambiguous_graph("cat", "dog")) + with patch.object( + self.api.db, "get_execution_epoch", return_value="stale-epoch" + ): + self.assertEqual(self._record("two.png"), "skipped") + self.assertEqual(self._items(), []) + + async def test_oversize_evidence_is_limited(self): + self._require_module() + big = "x" * (4 * 1024 * 1024 + 10) + side = {"text": big, "storage_text": big, "provenance": {}} + resolution = { + "state": "ambiguous", + "reason": "Multiple sampler branches", + "candidates": [ + { + "state": "resolved", + "sampler_id": "3", + "reason": "", + "positive": side, + "negative": {"text": "", "storage_text": "", "provenance": {}}, + "candidates": [], + } + ], + } + evidence, count, limited = review_capture.review_evidence(resolution) + self.assertEqual((count, limited), (0, True)) + self.assertEqual(evidence["candidates"], []) + self.assertLess(len(canonical_evidence(evidence)), 1024) + + async def test_nested_ambiguous_candidates_are_flattened_and_deduplicated(self): + self._require_module() + + def side(text): + return {"text": text, "storage_text": text, "provenance": {"n": text}} + + def leaf(sampler, text): + return { + "state": "resolved", + "sampler_id": sampler, + "reason": "", + "positive": side(text), + "negative": side("bad"), + "candidates": [], + "reason_code": None, + } + + resolution = { + "state": "ambiguous", + "reason": "Multiple sampler branches", + "candidates": [ + leaf("3", "cat"), + { + "state": "ambiguous", + "sampler_id": None, + "reason": "stages", + "positive": None, + "negative": None, + "candidates": [leaf("5", "dog"), leaf("3", "cat")], + }, + ], + } + evidence, count, limited = review_capture.review_evidence(resolution) + self.assertEqual((count, limited), (2, False)) + self.assertEqual( + [c["positive"]["text"] for c in evidence["candidates"]], ["cat", "dog"] + ) + + async def test_no_inference_or_model_access_during_recording(self): + self._require_module() + write_png(self.output_dir / "two.png", prompt=ambiguous_graph("cat", "dog")) + from py import autotag + + with ( + patch.object(autotag.AutoTagService, "load_model") as load, + patch.object(autotag.AutoTagService, "download_model") as download, + patch.object(autotag, "get_autotag_service") as service, + patch.object( + Image.Image, "load", autospec=True, side_effect=AssertionError("decode") + ) as decode, + ): + self.assertEqual(self._record("two.png"), "recorded") + load.assert_not_called() + download.assert_not_called() + service.assert_not_called() + decode.assert_not_called() + + async def test_capture_released_and_handles_closed(self): + self._require_module() + write_png(self.output_dir / "two.png", prompt=ambiguous_graph("cat", "dog")) + self.assertEqual(self._record("two.png"), "recorded") + self.assertEqual(self.api._captures.stats()["captures"], 0) + self.assertEqual(self._capture_files(), []) + write_png(self.output_dir / "three.png", prompt=ambiguous_graph("a", "b")) + with patch.object( + review_capture, + "resolve_prompt_pair", + side_effect=RuntimeError("resolver failed"), + ): + with self.assertRaises(RuntimeError): + self._record("three.png") + self.assertEqual(self.api._captures.stats()["captures"], 0) + self.assertEqual(self._capture_files(), []) + assert_closed(self, self.output_dir, self.capture_parent) + + +class TestReviewEdges(ReviewTestCase): + def _side(self, text): + return {"text": text, "storage_text": text, "provenance": {}} + + def _leaf(self, sampler, text): + return { + "state": "resolved", + "sampler_id": sampler, + "reason": "", + "positive": self._side(text), + "negative": self._side(""), + "candidates": [], + } + + async def test_evidence_over_the_bound_together_is_limited(self): + self._require_module() + half = "y" * (2 * 1024 * 1024 + 10) # twice per side, twice per item + resolution = { + "state": "ambiguous", + "reason": "Multiple sampler branches", + "candidates": [self._leaf("3", half), self._leaf("5", half + "z")], + } + evidence, count, limited = review_capture.review_evidence(resolution) + self.assertEqual((count, limited, evidence["candidate_total"]), (0, True, 2)) + + async def test_malformed_candidates_are_ignored_or_limited(self): + self._require_module() + self.assertEqual( + review_capture.review_evidence({"state": "ambiguous", "candidates": None}), + ( + { + "resolver_version": 1, + "reason": "", + "limited": True, + "candidate_total": 0, + "candidates": [], + }, + 0, + True, + ), + ) + resolution = { + "state": "ambiguous", + "reason": "x", + "candidates": ["junk", self._leaf("3", "cat"), self._leaf("5", "dog")], + } + self.assertEqual(review_capture.review_evidence(resolution)[1:], (2, False)) + broken = self._leaf("7", "eel") + broken["positive"] = {"storage_text": "no text"} + resolution["candidates"].append(broken) + self.assertEqual(review_capture.review_evidence(resolution)[1:], (0, True)) + + async def test_unreadable_or_promptless_captures_are_skipped(self): + self._require_module() + record = review_capture.record_ambiguous_file + (self.output_dir / "junk.png").write_bytes(b"not a png at all") + write_png(self.output_dir / "plain.png") + info = PngImagePlugin.PngInfo() + info.add_text("prompt", "{not json") + Image.new("RGB", (2, 2)).save(str(self.output_dir / "bad.png"), pnginfo=info) + for name in ("junk.png", "plain.png", "bad.png"): + self.assertEqual( + record(self.api.db, self.api._captures, self.output_dir, 0, name), + "skipped", + name, + ) + write_png(self.output_dir / "two.png", prompt=ambiguous_graph()) + with patch.object(self.api._captures, "acquire", return_value=None): + self.assertEqual( + record(self.api.db, self.api._captures, self.output_dir, 0, "two.png"), + "skipped", + ) + self.assertEqual(self._items(), []) + self.assertEqual(self._capture_files(), []) + + async def test_record_review_item_validates_its_facts(self): + self._require_module() + epoch = self.api.db.get_execution_epoch() + good = { + "locator": "/out/a.png", + "sha256": "a" * 64, + "size": 1, + "mtime_ns": 2, + "epoch": epoch, + "resolver_version": 1, + "evidence": {"candidates": []}, + "candidate_count": 0, + "limited": True, + } + for change in ( + {"resolver_version": 0}, + {"resolver_version": True}, + {"candidate_count": 17}, + {"candidate_count": -1}, + {"limited": 1}, + {"limited": True, "candidate_count": 2}, + {"sha256": "A" * 64}, + {"evidence": ["not", "an", "object"]}, + ): + with self.assertRaises(ValueError, msg=change): + self.api.db.record_review_item(**{**good, **change}) + self.assertTrue(self.api.db.record_review_item(**good)) + self.assertFalse(self.api.db.record_review_item(**good)) + items = self.api.db.review_items_for("/out/a.png") + self.assertEqual(len(items), 1) + self.assertEqual(items[0]["evidence"], {"candidates": []}) + self.assertIs(items[0]["limited"], True) + self.assertEqual(items[0]["sha256"], "a" * 64) + self.assertEqual(self.api.db.review_items_for("/out/b.png"), []) + self.assertEqual(self.api.db.review_index_revision(), 1) + with self.assertRaises(ValueError): + self.api.db.review_items_for("") + conn = self.api.db.model.get_connection() + with conn: + conn.execute("DELETE FROM media_review_metadata") + with self.assertRaises(sqlite3.DatabaseError): + self.api.db.review_index_revision() + + +class TestMonitorRecording(ReviewTestCase): + def _handler(self): + handler = ImageGenerationHandler(self.api.db, PromptTracker(self.api.db)) + return handler + + def _linked_prompt(self, path): + conn = self.api.db.model.get_connection() + row = conn.execute( + "SELECT prompt_id FROM generated_images WHERE image_path=?", (str(path),) + ).fetchone() + return row[0] if row else None + + async def test_monitor_keeps_fallback_link_and_records_item_best_effort(self): + self._require_module() + self._roots() + path = self.output_dir / "two.png" + write_png(path, prompt=ambiguous_graph("cat", "dog")) + prompt_id = self.api.db.save_prompt( + text="latest", prompt_hash=generate_prompt_hash("latest") + ) + metadata = {"prompt": ambiguous_graph("cat", "dog")} + context = {"id": prompt_id, "text": "latest", "fallback": True} + + handler = self._handler() + + def broken(image_path): + raise sqlite3.OperationalError("database is locked") + + handler.review_recorder = broken + with self.assertLogs("prompt_manager.image_monitor", logging.WARNING) as logs: + handler.link_image_to_prompt(str(path), context, metadata) + self.assertEqual(self._linked_prompt(path), prompt_id) + self.assertTrue(any("Needs review" in line for line in logs.output)) + self.assertEqual(self._items(), []) + + other = self.output_dir / "three.png" + write_png(other, prompt=ambiguous_graph("cat", "dog")) + handler.review_recorder = self.api.record_review_file + handler.link_image_to_prompt(str(other), context, metadata) + self.assertEqual(self._linked_prompt(other), prompt_id) + items = self._items() + self.assertEqual(len(items), 1) + self.assertEqual(items[0]["sha256"], digest(other)) + self.assertEqual(self._capture_files(), []) + + async def test_monitor_skips_resolved_and_files_outside_the_roots(self): + self._require_module() + self._roots() + prompt_id = self.api.db.save_prompt( + text="latest", prompt_hash=generate_prompt_hash("latest") + ) + context = {"id": prompt_id, "text": "latest"} + calls = [] + handler = self._handler() + handler.review_recorder = lambda image_path: calls.append(image_path) + resolved = self.output_dir / "one.png" + write_png(resolved, prompt=resolved_graph()) + handler.link_image_to_prompt( + str(resolved), context, {"prompt": resolved_graph()} + ) + self.assertEqual(calls, []) + outside = self.outside_dir / "two.png" + write_png(outside, prompt=ambiguous_graph()) + self.assertEqual(self.api.record_review_file(str(outside)), "skipped") + self.assertEqual(self._items(), []) + + async def test_monitor_wiring_reaches_new_and_existing_handlers(self): + self._require_module() + from utils.image_monitor import ImageMonitor + + monitor = ImageMonitor(self.api.db, PromptTracker(self.api.db)) + recorder = self.api.record_review_file + monitor.set_review_recorder(recorder) + self.assertIs(monitor.review_recorder, recorder) + monitor.start_monitoring([str(self.output_dir)]) + try: + self.assertIs(monitor.handler.review_recorder, recorder) + monitor.set_review_recorder(None) + self.assertIsNone(monitor.handler.review_recorder) + finally: + monitor.stop_monitoring() + + +if __name__ == "__main__": + unittest.main() diff --git a/utils/image_monitor.py b/utils/image_monitor.py index a42504e..0af5016 100644 --- a/utils/image_monitor.py +++ b/utils/image_monitor.py @@ -27,7 +27,7 @@ import os import queue import time import threading -from typing import Optional, Dict, Any, Tuple +from typing import Callable, Optional, Dict, Any, Tuple from watchdog.observers import Observer from watchdog.events import FileSystemEventHandler @@ -75,6 +75,9 @@ class ImageGenerationHandler(FileSystemEventHandler): self.prompt_tracker = prompt_tracker self.metadata_extractor = ComfyUIMetadataExtractor() self.logger = get_logger("prompt_manager.image_monitor") + # Needs review recording for ambiguous outputs (the API's + # ``record_review_file``, wired at start-up); None records nothing. + self.review_recorder: Optional[Callable[[str], str]] = None # Read from GalleryConfig if available, otherwise use defaults try: @@ -579,6 +582,28 @@ class ImageGenerationHandler(FileSystemEventHandler): ) except Exception as e: self.logger.error(f"Failed to link image to prompt: {e}") + return + if resolution["state"] == "ambiguous": + self._record_review(image_path) + + def _record_review(self, image_path: str) -> None: + """Queue an ambiguous output for Needs review, best effort. + + The fallback link above stays whatever happens here: a failure is + logged, never raised. + """ + recorder = getattr(self, "review_recorder", None) + if recorder is None: + return + try: + outcome = recorder(image_path) + except Exception: + self.logger.warning( + f"{os.path.basename(image_path)} could not be added to Needs review", + exc_info=True, + ) + return + self.logger.debug(f"Needs review {outcome}: {os.path.basename(image_path)}") class ImageMonitor: @@ -608,6 +633,16 @@ class ImageMonitor: self.handler = None self.monitored_directories = [] self.logger = get_logger("prompt_manager.image_monitor") + self.review_recorder: Optional[Callable[[str], str]] = None + + def set_review_recorder(self, recorder: Optional[Callable[[str], str]]) -> None: + """Record ambiguous outputs through ``recorder`` (None stops recording). + + Applies to the running handler and to every handler started later. + """ + self.review_recorder = recorder + if self.handler is not None: + self.handler.review_recorder = recorder def start_monitoring(self, output_directories: Optional[list] = None): """ @@ -665,6 +700,7 @@ class ImageMonitor: # Only a started observer becomes self.observer, so a failed start never # leaves the monitor stuck "already running" with nothing to stop. handler = ImageGenerationHandler(self.db_manager, self.prompt_tracker) + handler.review_recorder = self.review_recorder observer = Observer() scheduled = [] for output_dir in valid_directories: diff --git a/utils/review_capture.py b/utils/review_capture.py new file mode 100644 index 0000000..0ab9c34 --- /dev/null +++ b/utils/review_capture.py @@ -0,0 +1,199 @@ +"""Capture-bound Needs review items for ambiguous files (#249 E2). + +The output scan and the live monitor hand every file the ingestion resolver +finds ``ambiguous`` to :func:`record_ambiguous_file`. It captures the file +(the bounded private copy of ``utils.media_capture``, hashed while copied), +reads the metadata again from the capture, resolves it with the default +(ingestion) profile and stores only an ``ambiguous`` result, bound to the +capture's byte record. The evidence a reviewer chooses from is therefore the +evidence of exactly those bytes. Nothing here links a file, loads a model or +decodes pixels; the capture is released before the database write. + +Evidence is small on purpose (startup validates every row): the resolver's +reason and, per candidate sampler, ``{sampler_id, state, reason, positive, +negative}`` where each side is ``{text, storage_text, provenance}``. A +candidate's ``hash`` is the SHA-256 of its canonical JSON without the hash. +More than ``MAX_CANDIDATES`` candidates, none at all, or evidence over the +8 MiB bound make a ``limited`` item without candidates. +""" + +from __future__ import annotations + +import hashlib +import json +import logging +from pathlib import Path +from typing import Any + +try: + from ..database.library_commands import LibraryCommandError + from ..database.review_schema import MAX_CANDIDATES, canonical_evidence + from .media_capture import CaptureCapacity, CaptureError + from .prompt_inspector import MetadataReadError, read_inspection_file + from .prompt_pair_resolver import resolve_prompt_pair +except ImportError: # imported with the repo root on sys.path + from database.library_commands import LibraryCommandError + from database.review_schema import MAX_CANDIDATES, canonical_evidence + from utils.media_capture import CaptureCapacity, CaptureError + from utils.prompt_inspector import MetadataReadError, read_inspection_file + from utils.prompt_pair_resolver import resolve_prompt_pair + +logger = logging.getLogger("prompt_manager.review_capture") + +# The ingestion resolver's provenance version (``database.pair_ingestion``). +RESOLVER_VERSION = 1 +# Capture scope label: review captures are never handed to a client. +REVIEW_SCOPE = "needs-review" +# Nested ambiguous stages are flattened at most this deep. +MAX_CANDIDATE_DEPTH = 8 + + +def record_ambiguous_file(db, store, root, root_index: int, relative: str) -> str: + """Persist the ambiguous evidence of one rooted file. + + Returns ``recorded`` (new item with candidates), ``limited`` (new item + without candidates), ``duplicate`` (these bytes already carry this + evidence) or ``skipped`` (the file changed, vanished, is not ambiguous in + its captured bytes, or storage was restored meanwhile). + ``CaptureCapacity`` and database errors propagate to the caller, which + logs and counts them. + """ + epoch = db.get_execution_epoch() + try: + capture = store.capture( + root, relative, epoch=epoch, scope=REVIEW_SCOPE, root_index=root_index + ) + except CaptureCapacity: + raise + except CaptureError as error: + logger.debug("Review capture skipped (%s): %s", error.code, error) + return "skipped" + try: + resolution = _captured_resolution(store, capture) + finally: + store.release(capture.token) + if resolution is None or resolution.get("state") != "ambiguous": + return "skipped" + evidence, count, limited = review_evidence(resolution) + try: + inserted = db.record_review_item( + locator=capture.locator, + sha256=capture.sha256, + size=capture.size, + mtime_ns=capture.mtime_ns, + epoch=capture.epoch, + resolver_version=RESOLVER_VERSION, + evidence=evidence, + candidate_count=count, + limited=limited, + ) + except LibraryCommandError: + return "skipped" + if not inserted: + return "duplicate" + return "limited" if limited else "recorded" + + +def _captured_resolution(store, capture) -> dict | None: + """The ingestion resolver's result for the capture's own metadata.""" + if store.acquire(capture.token) is None: + return None + try: + try: + metadata = read_inspection_file(Path(capture.path)).metadata + except MetadataReadError: + return None + finally: + store.release_use(capture.token) + canonical = dict(metadata) + for key in ("prompt", "workflow"): + if isinstance(canonical.get(key), str): + try: + canonical[key] = json.loads(canonical[key]) + except (ValueError, TypeError, RecursionError): + # Malformed metadata stays as text for the resolver to reject. + pass + if not canonical.get("prompt") and not canonical.get("parameters"): + return None + return resolve_prompt_pair(canonical) + + +def review_evidence(resolution: dict) -> tuple[dict, int, bool]: + """``(evidence, candidate_count, limited)`` for one ambiguous result.""" + base = { + "resolver_version": RESOLVER_VERSION, + "reason": _text(resolution.get("reason")), + } + candidates: list[dict] = [] + seen: set[str] = set() + total = 0 + for candidate in _flattened(resolution.get("candidates"), 0): + total += 1 + try: + projected = _projection(candidate) + digest = hashlib.sha256( + canonical_evidence(projected).encode("utf-8") + ).hexdigest() + except (ValueError, TypeError, KeyError): + return _limited(base, total), 0, True + if digest in seen: + total -= 1 + continue + seen.add(digest) + candidates.append({**projected, "hash": digest}) + if not candidates or len(candidates) > MAX_CANDIDATES: + return _limited(base, total), 0, True + evidence = {**base, "candidates": candidates} + try: + canonical_evidence(evidence) + except ValueError: + return _limited(base, total), 0, True + return evidence, len(candidates), False + + +def _limited(base: dict, total: int) -> dict: + return {**base, "limited": True, "candidate_total": total, "candidates": []} + + +def _flattened(candidates: Any, depth: int): + """Selectable candidates; a nested ambiguous stage yields its own.""" + if not isinstance(candidates, list): + return + for candidate in candidates: + if not isinstance(candidate, dict): + continue + nested = candidate.get("candidates") + if ( + candidate.get("state") == "ambiguous" + and isinstance(nested, list) + and nested + and depth < MAX_CANDIDATE_DEPTH + ): + yield from _flattened(nested, depth + 1) + else: + yield candidate + + +def _projection(candidate: dict) -> dict: + return { + "sampler_id": candidate.get("sampler_id"), + "state": _text(candidate.get("state")), + "reason": _text(candidate.get("reason")), + "positive": _side(candidate.get("positive")), + "negative": _side(candidate.get("negative")), + } + + +def _side(side: Any) -> dict | None: + if not isinstance(side, dict): + return None + text = side["text"] + return { + "text": text, + "storage_text": side.get("storage_text", text), + "provenance": side.get("provenance"), + } + + +def _text(value: Any) -> str: + return value if isinstance(value, str) else ""