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

"""The suggestion read-path: raw predictions + centroids -> alias-resolved,
threshold-filtered, category-grouped, ranked suggestions for one image.
"""

from dataclasses import dataclass, field

from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession

from ...models import (
    ImageRecord,
    MLSettings,
    Tag,
    TagSuggestionRejection,
)
from ...models.tag import image_tag
from .aliases import AliasService
from .centroids import CentroidService
from .tagger import SURFACED_CATEGORIES


@dataclass(frozen=True)
class Suggestion:
    # canonical_tag_id is None when this is a raw Camie tag with no alias and
    # no existing Tag row — accepting it will create the tag.
    canonical_tag_id: int | None
    display_name: str
    category: str
    score: float
    source: str  # 'tagger' | 'centroid' | 'both'
    creates_new_tag: bool


@dataclass
class SuggestionList:
    by_category: dict[str, list[Suggestion]] = field(default_factory=dict)


class SuggestionService:
    def __init__(self, session: AsyncSession):
        self.session = session
        self.aliases = AliasService(session)
        self.centroids = CentroidService(session)

    async def _settings(self) -> MLSettings:
        return (
            await self.session.execute(select(MLSettings).where(MLSettings.id == 1))
        ).scalar_one()

    def _threshold_for(self, s: MLSettings, category: str) -> float:
        return {
            "artist": s.suggestion_threshold_artist,
            "character": s.suggestion_threshold_character,
            "copyright": s.suggestion_threshold_copyright,
            "general": s.suggestion_threshold_general,
        }.get(category, 1.01)  # 1.01 => never surfaces (unsurfaced category)

    async def for_image(self, image_id: int) -> SuggestionList:
        img = await self.session.get(ImageRecord, image_id)
        if img is None:
            return SuggestionList()

        settings = await self._settings()
        predictions: dict = img.tagger_predictions or {}

        applied = set(
            (
                await self.session.execute(
                    select(image_tag.c.tag_id).where(
                        image_tag.c.image_record_id == image_id
                    )
                )
            ).scalars().all()
        )
        rejected = set(
            (
                await self.session.execute(
                    select(TagSuggestionRejection.tag_id).where(
                        TagSuggestionRejection.image_record_id == image_id
                    )
                )
            ).scalars().all()
        )

        # --- Camie predictions ---
        candidates: list[tuple[str, str, float]] = []
        for name, p in predictions.items():
            category = p.get("category", "general")
            if category not in SURFACED_CATEGORIES:
                continue
            conf = float(p.get("confidence", 0.0))
            if conf < self._threshold_for(settings, category):
                continue
            candidates.append((name, category, conf))

        alias_map = await self.aliases.resolve_many(
            [(n, c) for n, c, _ in candidates]
        )

        merged: dict[object, Suggestion] = {}

        def _merge(key, sug: Suggestion):
            existing = merged.get(key)
            if existing is None:
                merged[key] = sug
            elif sug.score > existing.score:
                merged[key] = Suggestion(
                    canonical_tag_id=existing.canonical_tag_id,
                    display_name=existing.display_name,
                    category=existing.category,
                    score=sug.score,
                    source="both"
                    if existing.source != sug.source
                    else existing.source,
                    creates_new_tag=existing.creates_new_tag,
                )

        for name, category, conf in candidates:
            canonical = alias_map.get((name, category))
            if canonical is not None:
                if canonical.id in applied or canonical.id in rejected:
                    continue
                _merge(
                    canonical.id,
                    Suggestion(
                        canonical_tag_id=canonical.id,
                        display_name=canonical.name,
                        category=category,
                        score=conf,
                        source="tagger",
                        creates_new_tag=False,
                    ),
                )
            else:
                existing_tag = (
                    await self.session.execute(
                        select(Tag).where(Tag.name == name)
                    )
                ).scalars().first()
                if existing_tag is not None:
                    if (
                        existing_tag.id in applied
                        or existing_tag.id in rejected
                    ):
                        continue
                    _merge(
                        existing_tag.id,
                        Suggestion(
                            canonical_tag_id=existing_tag.id,
                            display_name=existing_tag.name,
                            category=category,
                            score=conf,
                            source="tagger",
                            creates_new_tag=False,
                        ),
                    )
                else:
                    _merge(
                        f"raw:{name}:{category}",
                        Suggestion(
                            canonical_tag_id=None,
                            display_name=name,
                            category=category,
                            score=conf,
                            source="tagger",
                            creates_new_tag=True,
                        ),
                    )

        # --- Centroid augmentation ---
        hits = await self.centroids.find_similar_tags(image_id, limit=30)
        for hit in hits:
            if hit.similarity < settings.centroid_similarity_threshold:
                continue
            if hit.tag_id in applied or hit.tag_id in rejected:
                continue
            tag = await self.session.get(Tag, hit.tag_id)
            if tag is None:
                continue
            cat = tag.kind.value if hasattr(tag.kind, "value") else str(tag.kind)
            display_cat = cat if cat in SURFACED_CATEGORIES else "general"
            _merge(
                tag.id,
                Suggestion(
                    canonical_tag_id=tag.id,
                    display_name=tag.name,
                    category=display_cat,
                    score=hit.similarity,
                    source="centroid",
                    creates_new_tag=False,
                ),
            )

        result = SuggestionList()
        for sug in merged.values():
            result.by_category.setdefault(sug.category, []).append(sug)
        for cat in result.by_category:
            result.by_category[cat].sort(key=lambda s: s.score, reverse=True)
        return result