diff --git a/backend/app/services/ml/allowlist.py b/backend/app/services/ml/allowlist.py new file mode 100644 index 0000000..dc8aeab --- /dev/null +++ b/backend/app/services/ml/allowlist.py @@ -0,0 +1,126 @@ +"""Allowlist semantics: accepting a suggestion adds the canonical tag to +image_tag AND to tag_allowlist; per-image removal/dismiss writes a rejection. +""" + +from collections.abc import Sequence +from dataclasses import dataclass + +from sqlalchemy import delete, select +from sqlalchemy.dialects.postgresql import insert +from sqlalchemy.ext.asyncio import AsyncSession + +from ...models import Tag, TagAllowlist, TagSuggestionRejection +from ...models.tag import image_tag +from .aliases import AliasService + + +@dataclass(frozen=True) +class AllowlistRow: + tag_id: int + tag_name: str + tag_kind: str + min_confidence: float + + +class AllowlistService: + def __init__(self, session: AsyncSession): + self.session = session + self.aliases = AliasService(session) + + async def _apply_image_tag(self, image_id: int, tag_id: int, source: str): + stmt = insert(image_tag).values( + image_record_id=image_id, tag_id=tag_id, source=source + ) + stmt = stmt.on_conflict_do_nothing( + index_elements=["image_record_id", "tag_id"] + ) + await self.session.execute(stmt) + + async def _add_to_allowlist(self, tag_id: int) -> bool: + """Returns True if newly added (caller should kick off retro-apply).""" + exists = await self.session.get(TagAllowlist, tag_id) + if exists is not None: + return False + self.session.add(TagAllowlist(tag_id=tag_id)) + await self.session.flush() + return True + + async def _clear_rejection(self, image_id: int, tag_id: int): + await self.session.execute( + delete(TagSuggestionRejection) + .where(TagSuggestionRejection.image_record_id == image_id) + .where(TagSuggestionRejection.tag_id == tag_id) + ) + + async def accept(self, image_id: int, tag_id: int) -> bool: + """Accept a suggestion. Returns True if the tag was newly added to + the allowlist (the API layer enqueues apply_allowlist_tags then).""" + await self._apply_image_tag(image_id, tag_id, source="ml_accepted") + await self._clear_rejection(image_id, tag_id) + return await self._add_to_allowlist(tag_id) + + async def add_alias_and_accept( + self, + image_id: int, + alias_string: str, + alias_category: str, + canonical_tag_id: int, + ) -> bool: + await self.aliases.create( + alias_string, alias_category, canonical_tag_id + ) + return await self.accept(image_id, canonical_tag_id) + + async def dismiss(self, image_id: int, tag_id: int) -> None: + stmt = insert(TagSuggestionRejection).values( + image_record_id=image_id, tag_id=tag_id + ) + stmt = stmt.on_conflict_do_nothing( + index_elements=["image_record_id", "tag_id"] + ) + await self.session.execute(stmt) + + async def reject_applied_tag(self, image_id: int, tag_id: int) -> None: + """Operator removed an applied tag from an image. Remove the + image_tag row AND record a rejection so the allowlist won't + re-apply it on the next maintenance sweep.""" + await self.session.execute( + image_tag.delete() + .where(image_tag.c.image_record_id == image_id) + .where(image_tag.c.tag_id == tag_id) + ) + await self.dismiss(image_id, tag_id) + + async def update_threshold( + self, tag_id: int, min_confidence: float + ) -> None: + row = await self.session.get(TagAllowlist, tag_id) + if row is not None: + row.min_confidence = min_confidence + + async def remove(self, tag_id: int) -> None: + await self.session.execute( + delete(TagAllowlist).where(TagAllowlist.tag_id == tag_id) + ) + + async def list_all(self) -> Sequence[AllowlistRow]: + stmt = ( + select( + TagAllowlist.tag_id, + Tag.name, + Tag.kind, + TagAllowlist.min_confidence, + ) + .join(Tag, Tag.id == TagAllowlist.tag_id) + .order_by(Tag.name.asc()) + ) + rows = (await self.session.execute(stmt)).all() + return [ + AllowlistRow( + tag_id=r[0], + tag_name=r[1], + tag_kind=r[2].value if hasattr(r[2], "value") else str(r[2]), + min_confidence=r[3], + ) + for r in rows + ] diff --git a/backend/app/services/tag_service.py b/backend/app/services/tag_service.py index 1469f86..85a72c9 100644 --- a/backend/app/services/tag_service.py +++ b/backend/app/services/tag_service.py @@ -162,3 +162,39 @@ class TagService: .order_by(Tag.kind.asc(), Tag.name.asc()) ) return (await self.session.execute(stmt)).scalars().all() + + async def rename(self, tag_id: int, new_name: str) -> Tag: + """Rename a tag. Raises TagValidationError if the new name collides + with an existing tag of the same (kind, fandom_id). + + Tag merge (consolidating two tags) lands in FC-2c; until then, + rename refuses on collision. + """ + new_name = new_name.strip() + if not new_name: + raise TagValidationError("Tag name cannot be empty") + tag = await self.session.get(Tag, tag_id) + if tag is None: + raise TagValidationError(f"Tag {tag_id} not found") + + clash_stmt = ( + select(Tag) + .where(Tag.name == new_name) + .where(Tag.kind == tag.kind) + .where( + Tag.fandom_id.is_(None) + if tag.fandom_id is None + else Tag.fandom_id == tag.fandom_id + ) + .where(Tag.id != tag_id) + ) + clash = (await self.session.execute(clash_stmt)).scalar_one_or_none() + if clash is not None: + raise TagValidationError( + f"A {tag.kind} tag named {new_name!r} already exists " + f"(merge lands in FC-2c)" + ) + + tag.name = new_name + await self.session.flush() + return tag diff --git a/tests/test_ml_allowlist.py b/tests/test_ml_allowlist.py new file mode 100644 index 0000000..a4c94bd --- /dev/null +++ b/tests/test_ml_allowlist.py @@ -0,0 +1,106 @@ +import pytest +from sqlalchemy import select + +from backend.app.models import TagAllowlist, TagKind, TagSuggestionRejection +from backend.app.models.tag import image_tag +from backend.app.services.ml.allowlist import AllowlistService +from backend.app.services.tag_service import TagService + +pytestmark = pytest.mark.integration + + +async def _make_image(db): + from backend.app.models import ImageRecord + img = ImageRecord( + path="/images/x.jpg", sha256="x" * 64, size_bytes=1, + mime="image/jpeg", width=1, height=1, + origin="imported_filesystem", integrity_status="unknown", + ) + db.add(img) + await db.flush() + return img + + +@pytest.mark.asyncio +async def test_accept_applies_and_allowlists(db): + img = await _make_image(db) + tag = await TagService(db).find_or_create("Hero", TagKind.character) + svc = AllowlistService(db) + newly_added = await svc.accept(img.id, tag.id) + assert newly_added is True + + applied = ( + await db.execute( + select(image_tag.c.source) + .where(image_tag.c.image_record_id == img.id) + .where(image_tag.c.tag_id == tag.id) + ) + ).scalar_one() + assert applied == "ml_accepted" + assert await db.get(TagAllowlist, tag.id) is not None + + +@pytest.mark.asyncio +async def test_accept_idempotent_allowlist(db): + img = await _make_image(db) + tag = await TagService(db).find_or_create("Hero2", TagKind.character) + svc = AllowlistService(db) + assert await svc.accept(img.id, tag.id) is True + assert await svc.accept(img.id, tag.id) is False + + +@pytest.mark.asyncio +async def test_reject_applied_tag_records_rejection(db): + img = await _make_image(db) + tag = await TagService(db).find_or_create("Removeme", TagKind.general) + svc = AllowlistService(db) + await svc.accept(img.id, tag.id) + await svc.reject_applied_tag(img.id, tag.id) + + still_applied = ( + await db.execute( + select(image_tag.c.tag_id) + .where(image_tag.c.image_record_id == img.id) + .where(image_tag.c.tag_id == tag.id) + ) + ).scalar_one_or_none() + assert still_applied is None + rej = await db.get(TagSuggestionRejection, (img.id, tag.id)) + assert rej is not None + + +@pytest.mark.asyncio +async def test_dismiss_records_rejection(db): + img = await _make_image(db) + tag = await TagService(db).find_or_create("Dismissme", TagKind.general) + await AllowlistService(db).dismiss(img.id, tag.id) + assert await db.get(TagSuggestionRejection, (img.id, tag.id)) is not None + + +@pytest.mark.asyncio +async def test_add_alias_and_accept(db): + img = await _make_image(db) + canonical = await TagService(db).find_or_create( + "Canonical Char", TagKind.character + ) + svc = AllowlistService(db) + await svc.add_alias_and_accept( + img.id, "model_char_name", "character", canonical.id + ) + from backend.app.services.ml.aliases import AliasService + resolved = await AliasService(db).resolve("model_char_name", "character") + assert resolved.id == canonical.id + assert await db.get(TagAllowlist, canonical.id) is not None + + +@pytest.mark.asyncio +async def test_update_threshold_and_remove(db): + tag = await TagService(db).find_or_create("Thr", TagKind.general) + svc = AllowlistService(db) + img = await _make_image(db) + await svc.accept(img.id, tag.id) + await svc.update_threshold(tag.id, 0.80) + row = await db.get(TagAllowlist, tag.id) + assert abs(row.min_confidence - 0.80) < 1e-6 + await svc.remove(tag.id) + assert await db.get(TagAllowlist, tag.id) is None diff --git a/tests/test_tag_rename.py b/tests/test_tag_rename.py new file mode 100644 index 0000000..fcb1ee2 --- /dev/null +++ b/tests/test_tag_rename.py @@ -0,0 +1,41 @@ +import pytest + +from backend.app.models import TagKind +from backend.app.services.tag_service import TagService, TagValidationError + +pytestmark = pytest.mark.integration + + +@pytest.mark.asyncio +async def test_rename_success(db): + svc = TagService(db) + t = await svc.find_or_create("uchiha_sasuke", TagKind.character) + renamed = await svc.rename(t.id, "Sasuke Uchiha") + assert renamed.id == t.id + assert renamed.name == "Sasuke Uchiha" + + +@pytest.mark.asyncio +async def test_rename_collision_raises(db): + svc = TagService(db) + await svc.find_or_create("Existing", TagKind.character) + other = await svc.find_or_create("Other", TagKind.character) + with pytest.raises(TagValidationError, match="already exists"): + await svc.rename(other.id, "Existing") + + +@pytest.mark.asyncio +async def test_rename_same_name_different_kind_ok(db): + svc = TagService(db) + await svc.find_or_create("Shared", TagKind.character) + artist = await svc.find_or_create("ArtistTag", TagKind.artist) + renamed = await svc.rename(artist.id, "Shared") + assert renamed.name == "Shared" + + +@pytest.mark.asyncio +async def test_rename_empty_raises(db): + svc = TagService(db) + t = await svc.find_or_create("Whatever", TagKind.general) + with pytest.raises(TagValidationError): + await svc.rename(t.id, " ")