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

"""CCIP few-shot character matcher (#114) — server-side, numpy on stored vectors.

CCIP is a FROZEN identity embedding; we don't train it. Instead the operator's
tagged characters become reference prototypes: a character tag's references are
the CCIP vectors of figure/face regions on images carrying that tag. To suggest
characters for a new image, we compare its figure-region CCIP vectors to every
character's references (multi-prototype: best match over a character's examples)
and surface the ones that clear a similarity threshold. No GPU here — the agent
already produced the vectors; this is cosine matching on what's stored.

v1 uses cosine similarity on the raw CCIP vectors with a tunable threshold; the
exact CCIP difference metric/threshold gets validated against the model during
the hands-on eval. numpy is imported lazily (API worker has it via pgvector).
"""

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

from ...models import ImageRegion, MLSettings, Tag, TagKind
from ...models.tag import image_tag

# Cosine-similarity floor to call a figure the same character. The live setting
# (ml_settings.ccip_match_threshold) drives it; this is only the fallback when no
# threshold is supplied AND no settings row exists.
DEFAULT_SIM_THRESHOLD = 0.85
_FIGURE_KINDS = ("face", "figure")


async def _settings_threshold(session: AsyncSession) -> float:
    val = (
        await session.execute(
            select(MLSettings.ccip_match_threshold).where(MLSettings.id == 1)
        )
    ).scalar_one_or_none()
    return float(val) if val is not None else DEFAULT_SIM_THRESHOLD


def _l2norm(mat, np):
    n = np.linalg.norm(mat, axis=1, keepdims=True)
    n[n == 0] = 1.0
    return mat / n


# Single-shot cache of the (expensive) reference load, keyed on a cheap
# signature that changes exactly when references could: a character tag added/
# removed (n_char_tags) or a figure embedded (max/ n of ccip regions). Shared by
# the live matcher (every modal open) and the auto-apply sweep.
_REF_CACHE: dict = {"sig": None, "refs": None}


def _single_character_images():
    """Subquery of image ids carrying EXACTLY ONE character tag. References come
    only from these — on a multi-character image the tag is image-level, so every
    figure would otherwise pollute each character's prototype set (a 2-character
    image tagged 'Velma' would make Daphne's figure a Velma reference)."""
    return (
        select(image_tag.c.image_record_id)
        .join(Tag, Tag.id == image_tag.c.tag_id)
        .where(Tag.kind == TagKind.character)
        .group_by(image_tag.c.image_record_id)
        .having(func.count() == 1)
    )


async def _ref_signature(session: AsyncSession) -> tuple:
    n_tags = (
        await session.execute(
            select(func.count())
            .select_from(image_tag)
            .join(Tag, Tag.id == image_tag.c.tag_id)
            .where(Tag.kind == TagKind.character)
        )
    ).scalar_one()
    n_regs, max_id = (
        await session.execute(
            select(func.count(), func.max(ImageRegion.id)).where(
                ImageRegion.kind.in_(_FIGURE_KINDS),
                ImageRegion.ccip_embedding.is_not(None),
            )
        )
    ).one()
    return (n_tags, n_regs, max_id)


async def character_references(session: AsyncSession) -> dict[int, list]:
    """Per character-tag CCIP reference vectors: figure/face-region CCIP
    embeddings on UNAMBIGUOUS (single-character) images carrying that tag.
    Multi-prototype — several vectors per character. Cached on a cheap signature."""
    sig = await _ref_signature(session)
    if _REF_CACHE["sig"] == sig and _REF_CACHE["refs"] is not None:
        return _REF_CACHE["refs"]
    rows = (
        await session.execute(
            select(image_tag.c.tag_id, ImageRegion.ccip_embedding)
            .select_from(ImageRegion)
            .join(
                image_tag,
                image_tag.c.image_record_id == ImageRegion.image_record_id,
            )
            .join(Tag, Tag.id == image_tag.c.tag_id)
            .where(Tag.kind == TagKind.character)
            .where(ImageRegion.kind.in_(_FIGURE_KINDS))
            .where(ImageRegion.ccip_embedding.is_not(None))
            .where(ImageRegion.image_record_id.in_(_single_character_images()))
        )
    ).all()
    refs: dict[int, list] = {}
    for tag_id, vec in rows:
        refs.setdefault(tag_id, []).append(vec)
    _REF_CACHE.update(sig=sig, refs=refs)
    return refs


async def _tag_names(session: AsyncSession, tag_ids: list[int]) -> dict[int, str]:
    if not tag_ids:
        return {}
    return dict(
        (
            await session.execute(
                select(Tag.id, Tag.name).where(Tag.id.in_(tag_ids))
            )
        ).all()
    )


async def match_image(
    session: AsyncSession, image_id: int, threshold: float | None = None
) -> list[dict]:
    """Character suggestions for one image from its figure-region CCIP vectors:
    [{tag_id, name, category:'character', score, source:'ccip'}], ranked.
    Already-applied character tags are excluded. Empty if the image has no figure
    CCIP vectors or no character references exist yet. threshold defaults to the
    live ml_settings.ccip_match_threshold."""
    import numpy as np

    if threshold is None:
        threshold = await _settings_threshold(session)

    qvecs = (
        await session.execute(
            select(ImageRegion.ccip_embedding).where(
                ImageRegion.image_record_id == image_id,
                ImageRegion.kind.in_(_FIGURE_KINDS),
                ImageRegion.ccip_embedding.is_not(None),
            )
        )
    ).scalars().all()
    if not qvecs:
        return []
    refs = await character_references(session)
    if not refs:
        return []
    applied = set(
        (
            await session.execute(
                select(image_tag.c.tag_id).where(
                    image_tag.c.image_record_id == image_id
                )
            )
        ).scalars()
    )
    names = await _tag_names(session, [t for t in refs if t not in applied])

    Q = _l2norm(np.vstack([np.asarray(v, dtype=np.float32) for v in qvecs]), np)
    out = []
    for tag_id, vecs in refs.items():
        if tag_id in applied:
            continue
        R = _l2norm(np.vstack([np.asarray(v, dtype=np.float32) for v in vecs]), np)
        best = float((Q @ R.T).max())   # best (query figure, reference) cosine
        if best >= threshold:
            out.append({
                "tag_id": tag_id,
                "name": names.get(tag_id, str(tag_id)),
                "category": "character",
                "score": round(best, 4),
                "source": "ccip",
            })
    out.sort(key=lambda d: d["score"], reverse=True)
    return out