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 ""