diff --git a/database/pair_ingestion.py b/database/pair_ingestion.py index 50011947..42c45a08 100644 --- a/database/pair_ingestion.py +++ b/database/pair_ingestion.py @@ -1,7 +1,14 @@ -"""Atomic metadata enrichment for a single known generation context.""" +"""Atomic metadata enrichment for a single known generation context. + +Also the helpers both ingestion store paths share (#247): the file's link +rows and the one-prompt-per-file rule (``file_link_rows``, +``link_conflict``) and the marks a prompt created by ingestion carries +(``Marks``, ``mark_created``). +""" import os -from typing import Any, Dict +from dataclasses import dataclass +from typing import Any, Dict, List, Optional, Tuple from .models import normalize_image_path from .operations import _coerce_file_info @@ -37,16 +44,69 @@ def blank_positive(resolution: Dict[str, Any]) -> bool: return not (isinstance(text, str) and text.strip()) +@dataclass(frozen=True) +class Marks: + """What a prompt ingestion creates carries, so it can be found and removed. + + ``note`` is formatted with ``name``, the file's name. + """ + + category: str + tags: Tuple[str, ...] + note: str + + +SCAN_MARKS = Marks("scanned", ("auto-scanned",), "Auto-scanned from {name}") +RUN_MARKS = Marks("scanned", ("auto-scanned",), "Saved from a workflow run: {name}") + + +def mark_created(conn: Any, db: Any, prompt_id: int, marks: Marks, name: str) -> None: + """Apply ``marks`` to a prompt row the caller just created on ``conn``.""" + conn.execute( + "UPDATE prompts SET category = ?, notes = ? WHERE id = ?", + (marks.category, marks.note.format(name=name), prompt_id), + ) + db._sync_prompt_tags(conn, prompt_id, list(marks.tags)) + + +def file_link_rows(conn: Any, image_path: str) -> List[Any]: + """The file's link rows (either spelling), each with its prompt's hash.""" + return conn.execute( + "SELECT gi.*,p.hash AS positive_hash FROM generated_images gi " + "JOIN prompts p ON p.id=gi.prompt_id " + "WHERE gi.file_path=? OR gi.image_path=?", + (normalize_image_path(image_path), image_path), + ).fetchall() + + +def link_conflict(rows: List[Any], prompt_hash: str) -> bool: + """One prompt per file: more than one row, or a row of another prompt.""" + return len(rows) > 1 or bool(rows and rows[0]["positive_hash"] != prompt_hash) + + def ingest_resolved_pair( - db: Any, image_path: str, metadata: Dict[str, Any], resolution: Dict[str, Any] + db: Any, + image_path: str, + metadata: Dict[str, Any], + resolution: Dict[str, Any], + *, + marks: Optional[Marks] = None, ) -> Dict[str, Any]: + """Store a resolved pair and link the file to it in one transaction. + + With ``marks`` (ingestion from files), each pair component this call + creates carries them; without (the PromptManager node path) components + are created bare, as always. + """ state = resolution.get("state", "unresolved") if state == "resolved" and blank_positive(resolution): state = "unresolved" if state != "resolved": return {"status": state, "pairs_added": 0, "links_added": 0} with db._pair_transaction() as conn: - return _ingest_resolved_pair_in(conn, db, image_path, metadata, resolution) + return _ingest_resolved_pair_in( + conn, db, image_path, metadata, resolution, marks=marks + ) def _conflict() -> Dict[str, Any]: @@ -61,6 +121,7 @@ def _ingest_resolved_pair_in( resolution: Dict[str, Any], *, provenance: Dict[str, Any] | None = None, + marks: Optional[Marks] = None, ) -> Dict[str, Any]: """``ingest_resolved_pair`` inside the caller's open ``_pair_transaction``. @@ -69,7 +130,8 @@ def _ingest_resolved_pair_in( negative is a ``conflict`` and nothing is written. The caller owns the transaction, so its own later failure rolls the link back too. ``provenance`` replaces the link's default resolver provenance (Needs - review records the user's selection there). + review records the user's selection there). ``marks`` as in + ``ingest_resolved_pair``. """ positive, negative = resolution["positive"], resolution["negative"] positive_text = positive.get("storage_text", positive["text"]) @@ -82,15 +144,8 @@ def _ingest_resolved_pair_in( "negative": negative, } encoded = dumps_strict(provenance, sort_keys=True, separators=(",", ":")) - images = conn.execute( - "SELECT gi.*,p.hash AS positive_hash FROM generated_images gi " - "JOIN prompts p ON p.id=gi.prompt_id " - "WHERE gi.file_path=? OR gi.image_path=?", - (normalized, image_path), - ).fetchall() - if len(images) > 1 or ( - images and images[0]["positive_hash"] != generate_prompt_hash(positive_text) - ): + images = file_link_rows(conn, image_path) + if link_conflict(images, generate_prompt_hash(positive_text)): return _conflict() existing = images[0] if images else None if existing: @@ -115,7 +170,12 @@ def _ingest_resolved_pair_in( generate_prompt_hash(negative_text) if negative_text.strip() else None, ), ).fetchone() - pair_id = db._save_pair(conn, positive_text, negative_text) + if marks is None: + pair_id = db._save_pair(conn, positive_text, negative_text) + else: + pair_id = _save_marked_pair( + conn, db, positive_text, negative_text, marks, os.path.basename(image_path) + ) pair = conn.execute("SELECT * FROM prompt_pairs WHERE id=?", (pair_id,)).fetchone() pairs_added = int(prior_pair is None) file_size, width, height, image_format = _coerce_file_info(metadata) @@ -197,6 +257,26 @@ def _ingest_resolved_pair_in( } +def _save_marked_pair( + conn: Any, + db: Any, + positive_text: str, + negative_text: str, + marks: Marks, + name: str, +) -> int: + """``_save_pair`` that marks each component it creates.""" + texts = [positive_text] + ([negative_text] if negative_text.strip() else []) + ids = [] + for text in texts: + prompt_id, created = db._pair_component_outcome(conn, text) + if created: + mark_created(conn, db, prompt_id, marks, name) + ids.append(prompt_id) + negative_id = ids[1] if len(ids) > 1 else None + return db._save_pair_components(conn, ids[0], negative_id)[0] + + def _document_form(conn: Any, values: Dict[str, Any]) -> Dict[str, Any]: """``values`` with the workflow and prompt graph texts stored once. diff --git a/database/run_attribution.py b/database/run_attribution.py index 50a671db..6a2c8f2c 100644 --- a/database/run_attribution.py +++ b/database/run_attribution.py @@ -164,7 +164,7 @@ def _runs_by_location(db: Any, names: List[str]) -> Optional[Dict[tuple, dict]]: return found -def _recorded_runs( +def recorded_runs( db: Any, image_paths: Iterable[str], output_root: Optional[str], @@ -201,7 +201,8 @@ def _recorded_runs( return attributed -def _named_prompt(db: Any, run_id: str) -> Optional[int]: +def named_prompt(db: Any, run_id: str) -> Optional[int]: + """The one stored prompt every positive context of ``run_id`` names, else None.""" contexts = db.get_execution_contexts(run_id).get("contexts") or [] named = { context.get("positive_prompt_id") @@ -228,11 +229,11 @@ def stored_run_prompts( paths = list(dict.fromkeys(path for path in image_paths if path)) result: Dict[str, Optional[int]] = dict.fromkeys(paths) prompts: Dict[str, Optional[int]] = {} - for path, (run_id, _nodes) in _recorded_runs( + for path, (run_id, _nodes) in recorded_runs( db, paths, output_root, path_style ).items(): if run_id not in prompts: - prompts[run_id] = _named_prompt(db, run_id) + prompts[run_id] = named_prompt(db, run_id) result[path] = prompts[run_id] return result @@ -275,7 +276,7 @@ def recorded_output( ) -> Optional[RecordedOutput]: """The recorded run facts for ``image_path``, else None (#386). - The file must be attributed to one run (``_recorded_runs``: one run + The file must be attributed to one run (``recorded_runs``: one run reported its location and it was written during that run). Read-only; it never links anything. """ @@ -286,7 +287,7 @@ def recorded_output( if own_transaction: # outputs and texts from one state of history conn.execute("BEGIN") try: - found = _recorded_runs(db, [image_path], output_root, path_style) + found = recorded_runs(db, [image_path], output_root, path_style) if image_path not in found: return None run_id, nodes = found[image_path] diff --git a/database/scanned_prompts.py b/database/scanned_prompts.py new file mode 100644 index 00000000..aa34a11f --- /dev/null +++ b/database/scanned_prompts.py @@ -0,0 +1,57 @@ +"""Link an output file to the prompt of one text, creating it when new (#247). + +The store step of ingestion from file metadata when the decision is a +positive text (``utils.file_prompt``). The prompt insert and the link run in +one ``BEGIN IMMEDIATE`` transaction: two writers never store one text twice, +a failed link leaves no new prompt behind, and a file stays linked to one +prompt (a file already linked to another prompt is a ``conflict``). +""" + +import logging +import os +from typing import Any, Dict, Optional, Tuple + +from .pair_ingestion import Marks, file_link_rows, link_conflict, mark_created + +try: + from ..utils.hashing import generate_prompt_hash +except ImportError: + from utils.hashing import generate_prompt_hash + +_LOGGER = logging.getLogger("prompt_manager.database.scanned_prompts") + + +def store_scanned_prompt_or_raise( + db: Any, path: str, canonical: Dict[str, Any], text: str, marks: Marks +) -> Tuple[str, Optional[int]]: + """``(status, prompt id)`` after linking ``path`` to ``text``'s prompt. + + ``added`` (prompt created, carrying ``marks``), ``linked`` (stored text + reused), ``unchanged`` (already linked to it) or ``conflict`` (linked to + another prompt; nothing written, id None). Database errors are raised + and roll the whole store back, so a caller can run it again. + """ + with db._pair_transaction() as conn: + rows = file_link_rows(conn, path) + if link_conflict(rows, generate_prompt_hash(text)): + return "conflict", None + if rows: + return "unchanged", rows[0]["prompt_id"] + prompt_id, created = db._pair_component_outcome(conn, text) + if created: + mark_created(conn, db, prompt_id, marks, os.path.basename(path)) + db._insert_image_link(conn, prompt_id, path, canonical) + return ("added" if created else "linked"), prompt_id + + +def store_scanned_prompt( + db: Any, path: str, canonical: Dict[str, Any], text: str, marks: Marks +) -> Tuple[str, Optional[int]]: + """The scan's form: an error is logged and reported as ``("found", None)``.""" + try: + return store_scanned_prompt_or_raise(db, path, canonical, text, marks) + except Exception as error: + _LOGGER.error( + f"Failed to store the prompt of {os.path.basename(path)}: {error}" + ) + return "found", None diff --git a/tests/test_run_attribution.py b/tests/test_run_attribution.py index 050d710a..4b0da18c 100644 --- a/tests/test_run_attribution.py +++ b/tests/test_run_attribution.py @@ -326,7 +326,7 @@ class TestRecordedOutputNode(RecordedRuns, unittest.TestCase): self.assertIsNone(recorded_output_lookup(broken)(self.image)) with ( mock.patch( - "database.run_attribution._recorded_runs", + "database.run_attribution.recorded_runs", side_effect=sqlite3.OperationalError("locked"), ), self.assertLogs("prompt_manager.run_attribution", "WARNING"), diff --git a/tests/test_scanned_prompts.py b/tests/test_scanned_prompts.py new file mode 100644 index 00000000..0cabf9b1 --- /dev/null +++ b/tests/test_scanned_prompts.py @@ -0,0 +1,218 @@ +"""Storing a file's prompt from its text: one transaction, one prompt each (#247).""" + +import os +import sqlite3 +import tempfile +import threading +import unittest +from concurrent.futures import ThreadPoolExecutor +from pathlib import Path +from unittest import mock + +try: + from tests.open_handles import assert_closed + from tests import test_execution_history_database as history_tests +except ImportError: # discovered with tests/ as the top-level directory + from open_handles import assert_closed + import test_execution_history_database as history_tests + +from database.operations import PromptDatabase +from database.pair_ingestion import ( + RUN_MARKS, + SCAN_MARKS, + Marks, + ingest_resolved_pair, +) +from database.scanned_prompts import store_scanned_prompt, store_scanned_prompt_or_raise +from utils.hashing import generate_prompt_hash + +META = {"prompt": {"1": {"class_type": "CLIPTextEncode", "inputs": {"text": "x"}}}} + + +def resolved(positive="cat", negative="blur"): + return { + "state": "resolved", + "positive": {"text": positive, "provenance": {"source": "api"}}, + "negative": {"text": negative, "provenance": {"source": "api"}}, + } + + +class StoreCase(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.db = PromptDatabase(str(Path(self.temp.name) / "test.db")) + self.path = str(Path(self.temp.name) / "gen.png") + + def tearDown(self): + self.db.close_all() + assert_closed(self, self.temp.name) + self.temp.cleanup() + + def prompts(self): + rows = self.db.model.get_connection().execute( + "SELECT id, text FROM prompts ORDER BY id" + ) + return [tuple(row) for row in rows] + + def links(self): + rows = self.db.model.get_connection().execute( + "SELECT prompt_id, image_path FROM generated_images ORDER BY id" + ) + return [tuple(row) for row in rows] + + def save(self, text, **fields): + return self.db.save_prompt( + text, prompt_hash=generate_prompt_hash(text), **fields + ) + + +class TestStoreScannedPrompt(StoreCase): + def test_a_new_text_is_created_with_the_scan_marks_and_linked(self): + status, prompt_id = store_scanned_prompt_or_raise( + self.db, self.path, META, " a cat ", SCAN_MARKS + ) + self.assertEqual(status, "added") + prompt = self.db.get_prompt_by_id(prompt_id) + self.assertEqual(prompt["text"], "a cat") + self.assertEqual(prompt["category"], "scanned") + self.assertEqual(prompt["tags"], ["auto-scanned"]) + self.assertEqual(prompt["notes"], "Auto-scanned from gen.png") + self.assertEqual(prompt["run_count"], 0) + self.assertEqual(self.links(), [(prompt_id, self.path)]) + + def test_a_run_created_prompt_names_the_workflow_run(self): + _status, prompt_id = store_scanned_prompt_or_raise( + self.db, self.path, META, "a cat", RUN_MARKS + ) + prompt = self.db.get_prompt_by_id(prompt_id) + self.assertEqual(prompt["notes"], "Saved from a workflow run: gen.png") + self.assertEqual(prompt["tags"], ["auto-scanned"]) + + def test_a_stored_text_is_reused_by_hash_and_keeps_its_fields(self): + existing = self.save("A Cat", category="mine", notes="kept") + status, prompt_id = store_scanned_prompt_or_raise( + self.db, self.path, META, "a cat ", SCAN_MARKS + ) + self.assertEqual((status, prompt_id), ("linked", existing)) + prompt = self.db.get_prompt_by_id(existing) + self.assertEqual((prompt["category"], prompt["notes"]), ("mine", "kept")) + self.assertEqual(prompt["tags"], []) + self.assertEqual(len(self.prompts()), 1) + + def test_a_repeat_is_unchanged(self): + first = store_scanned_prompt_or_raise( + self.db, self.path, META, "a cat", SCAN_MARKS + ) + again = store_scanned_prompt_or_raise( + self.db, self.path, META, "a cat", SCAN_MARKS + ) + self.assertEqual(again, ("unchanged", first[1])) + self.assertEqual(len(self.links()), 1) + + def test_a_file_linked_to_another_prompt_is_a_conflict(self): + other = self.save("a dog") + # Linked under another spelling of the same file. + spelled = os.path.join(self.temp.name, ".", "gen.png") + self.db.link_image_to_prompt(other, spelled) + result = store_scanned_prompt_or_raise( + self.db, self.path, META, "a cat", SCAN_MARKS + ) + self.assertEqual(result, ("conflict", None)) + self.assertEqual(self.prompts(), [(other, "a dog")]) + self.assertEqual(self.links(), [(other, spelled)]) + + def test_a_failed_link_raises_and_leaves_no_prompt_behind(self): + with mock.patch.object( + self.db, "_insert_image_link", side_effect=sqlite3.OperationalError("x") + ): + with self.assertRaises(sqlite3.OperationalError): + store_scanned_prompt_or_raise( + self.db, self.path, META, "a cat", SCAN_MARKS + ) + with self.assertLogs("prompt_manager", "ERROR"): + result = store_scanned_prompt( + self.db, self.path, META, "a cat", SCAN_MARKS + ) + self.assertEqual(result, ("found", None)) + self.assertEqual(self.prompts(), []) + self.assertEqual(self.links(), []) + + def test_two_concurrent_writers_store_one_prompt_and_one_link(self): + hold = history_tests.FirstWriterHold() + outcome = self.db._pair_component_outcome + started = threading.Barrier(2) + + def held_open(*args, **kwargs): + hold() + return outcome(*args, **kwargs) + + def write(_index): + started.wait(timeout=30) + try: + return store_scanned_prompt_or_raise( + self.db, self.path, META, "a cat", SCAN_MARKS + )[0] + finally: + self.db.model.close() + + with ( + mock.patch.object( + self.db, "_pair_component_outcome", side_effect=held_open + ), + ThreadPoolExecutor(max_workers=2) as pool, + ): + statuses = sorted(pool.map(write, range(2))) + self.assertEqual(statuses, ["added", "unchanged"]) + self.assertEqual(len(self.prompts()), 1) + self.assertEqual(len(self.links()), 1) + + +class TestPairMarks(StoreCase): + def test_marks_reach_only_created_components(self): + kept = self.save("cat", category="mine") + result = ingest_resolved_pair( + self.db, self.path, META, resolved(), marks=RUN_MARKS + ) + self.assertEqual(result["status"], "added") + positive = self.db.get_prompt_by_id(kept) + self.assertEqual((positive["category"], positive["tags"]), ("mine", [])) + negative_id = next(pid for pid, text in self.prompts() if text == "blur") + negative = self.db.get_prompt_by_id(negative_id) + self.assertEqual(negative["category"], "scanned") + self.assertEqual(negative["tags"], ["auto-scanned"]) + self.assertEqual(negative["notes"], "Saved from a workflow run: gen.png") + + def test_without_marks_created_components_stay_bare(self): + """Regression guard (passes on the base): the node path is unchanged.""" + ingest_resolved_pair(self.db, self.path, META, resolved()) + for prompt_id, _text in self.prompts(): + prompt = self.db.get_prompt_by_id(prompt_id) + self.assertIsNone(prompt["category"]) + self.assertEqual(prompt["tags"], []) + + def test_a_blank_negative_creates_only_the_positive(self): + result = ingest_resolved_pair( + self.db, self.path, META, resolved(negative=""), marks=SCAN_MARKS + ) + self.assertEqual(result["status"], "added") + self.assertEqual([text for _id, text in self.prompts()], ["cat"]) + + def test_marks_are_a_frozen_value(self): + self.assertEqual( + SCAN_MARKS, Marks("scanned", ("auto-scanned",), SCAN_MARKS.note) + ) + with self.assertRaises(AttributeError): + SCAN_MARKS.category = "x" + + +class TestPublicRunAttribution(unittest.TestCase): + def test_recorded_runs_and_named_prompt_are_public(self): + from database import run_attribution + + self.assertTrue(callable(run_attribution.recorded_runs)) + self.assertTrue(callable(run_attribution.named_prompt)) + self.assertFalse(hasattr(run_attribution, "_recorded_runs")) + + +if __name__ == "__main__": + unittest.main()