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