FabledCurator/backend/app/services/ml/allowlist.py

"""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 MLSettings, 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 _store_floor(self) -> float:
        return (
            await self.session.execute(
                select(MLSettings.tagger_store_floor).where(MLSettings.id == 1)
            )
        ).scalar_one()

    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:
            # An allowlist tag can't auto-apply more permissively than the
            # ingest store floor — predictions below tagger_store_floor aren't
            # stored, so a lower min_confidence would behave identically to the
            # floor. Clamp so the stored threshold matches actual behavior
            # (#764).
            floor = await self._store_floor()
            row.min_confidence = max(min_confidence, floor)

    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
        ]