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

"""Tag centroids: the mean SigLIP embedding of a tag's member images.

Powers centroid-augmented suggestions (a tag whose centroid is close to an
image's embedding becomes a suggestion even if Camie didn't predict it).
"""

from dataclasses import dataclass

import numpy as np
from sqlalchemy import func, select
from sqlalchemy.dialects.postgresql import insert
from sqlalchemy.ext.asyncio import AsyncSession

from ...models import (
    ImageRecord,
    MLSettings,
    Tag,
    TagKind,
    TagReferenceEmbedding,
)
from ...models.tag import image_tag

ELIGIBLE_KINDS = {
    TagKind.character,
    TagKind.fandom,
    TagKind.general,
    TagKind.series,
}


@dataclass(frozen=True)
class CentroidHit:
    tag_id: int
    similarity: float


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

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

    async def _model_version(self) -> str:
        """Audit 2026-06-02: SigLIP model-version stamp comes from the
        DB row, not the env constant. tag_and_embed (tasks/ml.py:110)
        already reads from MLSettings.embedder_model_version, so by
        sourcing centroid stamps + drift checks from the same row, we
        eliminate the silent-drift case the audit flagged. env
        SIGLIP_MODEL_VERSION still drives which model embedder.py
        loads at runtime; the version stamp is purely the operator-
        controlled identifier."""
        return (
            await self.session.execute(
                select(MLSettings.embedder_model_version).where(MLSettings.id == 1)
            )
        ).scalar_one()

    async def recompute_for_tag(self, tag_id: int) -> bool:
        """Recompute one tag's centroid. Returns True if a centroid was
        written, False if skipped (ineligible kind or too few members)."""
        tag = await self.session.get(Tag, tag_id)
        if tag is None or tag.kind not in ELIGIBLE_KINDS:
            return False

        min_refs = await self._min_reference_images()

        stmt = (
            select(ImageRecord.siglip_embedding)
            .join(image_tag, image_tag.c.image_record_id == ImageRecord.id)
            .where(image_tag.c.tag_id == tag_id)
            .where(ImageRecord.siglip_embedding.is_not(None))
        )
        embeddings = [
            np.array(e, dtype=np.float32)
            for e in (await self.session.execute(stmt)).scalars().all()
        ]
        if len(embeddings) < min_refs:
            return False

        centroid = np.mean(np.stack(embeddings), axis=0).astype(np.float32)
        model_version = await self._model_version()

        stmt = insert(TagReferenceEmbedding).values(
            tag_id=tag_id,
            embedding=centroid.tolist(),
            reference_count=len(embeddings),
            model_version=model_version,
        )
        stmt = stmt.on_conflict_do_update(
            index_elements=["tag_id"],
            set_={
                "embedding": centroid.tolist(),
                "reference_count": len(embeddings),
                "model_version": model_version,
                "updated_at": func.now(),
            },
        )
        await self.session.execute(stmt)
        return True

    async def list_drifted(self) -> list[int]:
        """Tag ids whose centroid is stale: member count != reference_count,
        OR no centroid row, OR centroid built on a different SigLIP version.
        Only considers eligible-kind tags with embeddings present."""
        current_model_version = await self._model_version()
        member_counts = (
            select(
                image_tag.c.tag_id.label("tag_id"),
                func.count(image_tag.c.image_record_id).label("members"),
            )
            .join(ImageRecord, ImageRecord.id == image_tag.c.image_record_id)
            .where(ImageRecord.siglip_embedding.is_not(None))
            .group_by(image_tag.c.tag_id)
            .subquery()
        )
        stmt = (
            select(Tag.id)
            .join(member_counts, member_counts.c.tag_id == Tag.id)
            .outerjoin(
                TagReferenceEmbedding,
                TagReferenceEmbedding.tag_id == Tag.id,
            )
            .where(Tag.kind.in_(ELIGIBLE_KINDS))
            .where(
                (TagReferenceEmbedding.tag_id.is_(None))
                | (
                    TagReferenceEmbedding.reference_count
                    != member_counts.c.members
                )
                | (TagReferenceEmbedding.model_version != current_model_version)
            )
        )
        return list((await self.session.execute(stmt)).scalars().all())

    async def find_similar_tags(
        self, image_id: int, limit: int = 20
    ) -> list[CentroidHit]:
        """Cosine similarity between an image's embedding and stored
        centroids. Returns top-`limit` by similarity DESC. pgvector's
        cosine_distance gives 1 - cosine_similarity."""
        img = await self.session.get(ImageRecord, image_id)
        if img is None or img.siglip_embedding is None:
            return []
        emb = img.siglip_embedding
        distance = TagReferenceEmbedding.embedding.cosine_distance(emb)
        stmt = (
            select(
                TagReferenceEmbedding.tag_id,
                (1 - distance).label("similarity"),
            )
            .order_by(distance.asc())
            .limit(limit)
        )
        rows = (await self.session.execute(stmt)).all()
        return [
            CentroidHit(tag_id=r.tag_id, similarity=float(r.similarity))
            for r in rows
        ]