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

"""The suggestion read-path: trained HEADS score one image's frozen embedding
into alias-resolved, category-grouped, ranked suggestions.

Tagging-v2 (#114): suggestions now come from the per-concept heads that LEARN
from the operator's tags (services/ml/heads.py) — the Camie prediction source
and the per-tag SigLIP centroid have been REMOVED. A head exists only for an
existing concept tag, so every suggestion is a canonical tag (no raw model key,
no alias remap, no creates-new). Rejected tags stay in the list FLAGGED (not
dropped) so the rail can show + reverse a dismissal.
"""

from dataclasses import dataclass, field

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

from ...models import ImageRecord, TagSuggestionRejection
from ...models.tag import image_tag
from .heads import score_image


@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  # 'head' (Camie 'tagger'/'centroid' sources removed in v2)
    creates_new_tag: bool
    # raw_name = the booru model vocab key behind this suggestion. It's the key
    # an alias MUST be stored under (resolution looks up the raw key), so the
    # modal needs it to author an alias correctly. None for centroid-only hits
    # (no underlying prediction → nothing to alias).
    raw_name: str | None = None
    # via_alias = this suggestion was surfaced because an operator alias remapped
    # the raw prediction to this canonical tag. Lets the UI mark it + offer undo.
    via_alias: bool = False
    # rejected = the operator dismissed this tag for this image (a stored
    # TagSuggestionRejection). It stays in the list — flagged, not dropped — so
    # the rejection is VISIBLE and REVERSIBLE in the rail (misclick recovery,
    # operator-asked 2026-06-27) instead of silently vanishing or re-suggesting.
    rejected: bool = False


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


class SuggestionService:
    def __init__(self, session: AsyncSession):
        self.session = session

    async def for_image(
        self, image_id: int, threshold_override: float | None = None,
    ) -> SuggestionList:
        """Head-scored suggestions for one image, grouped by category and ranked.

        Each trained head scores the image's frozen embedding; a concept surfaces
        when its score clears the head's own suggest threshold. threshold_override
        (used by the typed tag-input dropdown's "show everything" mode) replaces
        that per-head cut with a flat floor (0 → every head), so a low-scoring
        concept can still be typed + picked in canonical formatting.

        Already-applied tags are dropped; rejected tags stay FLAGGED and sink to
        the bottom of their category so a dismissal is visible + reversible."""
        img = await self.session.get(ImageRecord, image_id)
        if img is None:
            return SuggestionList()

        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()
        )

        hits = await score_image(
            self.session, image_id, threshold_override=threshold_override
        )

        result = SuggestionList()
        for h in hits:
            tag_id = h["tag_id"]
            if tag_id in applied:
                continue
            result.by_category.setdefault(h["category"], []).append(
                Suggestion(
                    canonical_tag_id=tag_id,
                    display_name=h["name"],
                    category=h["category"],
                    score=h["score"],
                    source="head",
                    creates_new_tag=False,
                    rejected=tag_id in rejected,
                )
            )
        for cat in result.by_category:
            # Live suggestions first (by score), rejected ones sink to the
            # bottom of the category — visible for recovery, out of the way.
            result.by_category[cat].sort(key=lambda s: (s.rejected, -s.score))
        return result

    async def for_selection(
        self,
        image_ids: list[int],
        threshold: float = 0.8,
        top_k: int = 10,
    ) -> dict[str, list[dict]]:
        """Consensus suggestions across image_ids. A tag is included iff it
        was suggested for (or already applied to) >= threshold fraction of
        the selection AND was acceptable on >= 1 image. Confidence is the
        mean over images where it was suggested. Aggregated by
        canonical_tag_id; creates-new (no canonical id) suggestions are
        skipped (bulk Accept applies by tag id)."""
        if not image_ids:
            return {}
        threshold = min(1.0, max(0.0, threshold))
        total = len(image_ids)

        stats: dict[int, dict] = {}
        for image_id in image_ids:
            sl = await self.for_image(image_id)
            for category, items in sl.by_category.items():
                for s in items:
                    if s.canonical_tag_id is None or s.creates_new_tag:
                        continue
                    # for_image keeps rejected tags (flagged) for the rail;
                    # bulk consensus must still ignore them — a tag dismissed on
                    # an image isn't a suggestion for that image.
                    if s.rejected:
                        continue
                    st = stats.get(s.canonical_tag_id)
                    if st is None:
                        st = {
                            "tag_id": s.canonical_tag_id,
                            "name": s.display_name,
                            "category": category,
                            "source": s.source,
                            "suggested_count": 0,
                            "sum_score": 0.0,
                        }
                        stats[s.canonical_tag_id] = st
                    st["suggested_count"] += 1
                    st["sum_score"] += s.score

        rows = (
            await self.session.execute(
                select(
                    image_tag.c.image_record_id, image_tag.c.tag_id
                ).where(image_tag.c.image_record_id.in_(image_ids))
            )
        ).all()
        applied_by_tag: dict[int, set[int]] = {}
        for iid, tid in rows:
            applied_by_tag.setdefault(tid, set()).add(iid)

        result: dict[str, list[dict]] = {}
        for st in stats.values():
            existing_count = len(applied_by_tag.get(st["tag_id"], set()))
            covered = st["suggested_count"] + existing_count
            coverage = covered / total
            if coverage < threshold or st["suggested_count"] < 1:
                continue
            result.setdefault(st["category"], []).append(
                {
                    "canonical_tag_id": st["tag_id"],
                    "name": st["name"],
                    "category": st["category"],
                    "confidence": round(
                        st["sum_score"] / st["suggested_count"], 4
                    ),
                    "coverage": round(coverage, 4),
                    "covered_count": covered,
                    "source": st["source"],
                }
            )
        for cat in result:
            result[cat].sort(key=lambda x: x["confidence"], reverse=True)
            result[cat] = result[cat][:top_k]
        return result