FabledCurator/tests/test_ml_suggestions.py

"""Suggestion read-path (tagging-v2): suggestions come from trained HEADS, not
Camie predictions or centroids. Heads are inserted directly (training needs
scikit-learn, ml image only); scoring is numpy-only (available via pgvector)."""
import pytest
from sqlalchemy import select

from backend.app.models import ImageRecord, ImageRegion, MLSettings, TagHead, TagKind
from backend.app.models.tag import image_tag
from backend.app.services.ml.allowlist import AllowlistService
from backend.app.services.ml.suggestions import SuggestionService
from backend.app.services.tag_service import TagService

pytestmark = pytest.mark.integration


def _emb(slot: int, val: float = 3.0) -> list[float]:
    """An embedding pointing along axis `slot` (so its L2-normalized form is the
    unit vector e_slot — a head with weights e_slot scores it sigmoid(1)≈0.73)."""
    v = [0.0] * 1152
    v[slot] = val
    return v


async def _img(db, sha: str, emb=None) -> ImageRecord:
    img = ImageRecord(
        path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg",
        width=1, height=1, origin="imported_filesystem",
        integrity_status="unknown", siglip_embedding=emb,
    )
    db.add(img)
    await db.flush()
    return img


async def _embver(db) -> str:
    s = (await db.execute(select(MLSettings).where(MLSettings.id == 1))).scalar_one()
    return s.embedder_model_version


async def _head(db, tag_id: int, slot: int, suggest_threshold: float = 0.5):
    weights = [0.0] * 1152
    weights[slot] = 1.0
    db.add(TagHead(
        tag_id=tag_id, embedding_version=await _embver(db),
        weights=weights, bias=0.0, suggest_threshold=suggest_threshold,
        auto_apply_threshold=None, n_pos=10, n_neg=30,
        ap=0.8, precision_cv=0.9, recall=0.6,
    ))


@pytest.mark.asyncio
async def test_head_suggestion_surfaces_for_matching_image(db):
    tag = await TagService(db).find_or_create("glasses", TagKind.general)
    img = await _img(db, "a" * 64, _emb(0))
    await _head(db, tag.id, slot=0)
    await db.commit()

    sl = await SuggestionService(db).for_image(img.id)
    general = sl.by_category["general"]
    assert len(general) == 1
    s = general[0]
    assert s.canonical_tag_id == tag.id
    assert s.source == "head"
    assert s.creates_new_tag is False
    assert s.via_alias is False and s.raw_name is None
    assert s.score > 0.5


@pytest.mark.asyncio
async def test_no_embedding_means_no_suggestions(db):
    img = await _img(db, "b" * 64, None)
    tag = await TagService(db).find_or_create("cat", TagKind.general)
    await _head(db, tag.id, slot=0)
    await db.commit()
    assert (await SuggestionService(db).for_image(img.id)).by_category == {}


@pytest.mark.asyncio
async def test_no_heads_means_no_suggestions(db):
    img = await _img(db, "c" * 64, _emb(0))
    await db.commit()  # no heads trained yet
    assert (await SuggestionService(db).for_image(img.id)).by_category == {}


@pytest.mark.asyncio
async def test_applied_tag_not_suggested(db):
    tag = await TagService(db).find_or_create("dog", TagKind.general)
    img = await _img(db, "d" * 64, _emb(0))
    await _head(db, tag.id, slot=0)
    await db.execute(
        image_tag.insert().values(
            image_record_id=img.id, tag_id=tag.id, source="manual"
        )
    )
    await db.commit()
    sl = await SuggestionService(db).for_image(img.id)
    assert "general" not in sl.by_category or not sl.by_category["general"]


@pytest.mark.asyncio
async def test_threshold_override_surfaces_below_cut(db):
    # A head with a high suggest_threshold won't surface on a so-so score, but
    # the dropdown's override=0 floor surfaces every head regardless.
    tag = await TagService(db).find_or_create("horse", TagKind.general)
    img = await _img(db, "e" * 64, _emb(1))   # orthogonal to the head → score 0.5
    await _head(db, tag.id, slot=0, suggest_threshold=0.6)
    await db.commit()
    svc = SuggestionService(db)
    assert svc and not (await svc.for_image(img.id)).by_category.get("general")
    flooded = await svc.for_image(img.id, threshold_override=0.0)
    assert any(s.canonical_tag_id == tag.id for s in flooded.by_category["general"])


@pytest.mark.asyncio
async def test_rejected_tag_surfaced_flagged_then_reversible(db):
    # A dismissed suggestion is NOT dropped: it stays flagged rejected so the
    # rail can show it + offer one-click un-reject (operator-asked 2026-06-27).
    tag = await TagService(db).find_or_create("goblin", TagKind.general)
    img = await _img(db, "f" * 64, _emb(0))
    await _head(db, tag.id, slot=0)
    await db.commit()
    await AllowlistService(db).dismiss(img.id, tag.id)
    await db.commit()

    sl = await SuggestionService(db).for_image(img.id)
    s = next(x for x in sl.by_category["general"] if x.canonical_tag_id == tag.id)
    assert s.rejected is True

    await AllowlistService(db).undismiss(img.id, tag.id)
    await db.commit()
    sl2 = await SuggestionService(db).for_image(img.id)
    s2 = next(x for x in sl2.by_category["general"] if x.canonical_tag_id == tag.id)
    assert s2.rejected is False


async def _figure(db, image_id, slot):
    v = [0.0] * 768
    v[slot] = 1.0
    db.add(ImageRegion(
        image_record_id=image_id, kind="figure",
        rx=0.0, ry=0.0, rw=1.0, rh=1.0,
        ccip_embedding=v, embedding_version="ccip-test",
    ))


@pytest.mark.asyncio
async def test_ccip_character_surfaces_in_rail(db):
    # A character with a CCIP reference (a tagged figure) is suggested on a new
    # image whose figure matches — overlaid into the rail alongside the heads.
    raven = await TagService(db).find_or_create("Raven", TagKind.character)
    ref = await _img(db, "0" * 64, None)         # the operator's tagged example
    await _figure(db, ref.id, slot=0)
    await db.execute(image_tag.insert().values(
        image_record_id=ref.id, tag_id=raven.id, source="manual",
    ))
    query = await _img(db, "1" * 64, None)        # untagged, matching figure
    await _figure(db, query.id, slot=0)
    await db.commit()

    sl = await SuggestionService(db).for_image(query.id)
    m = next(
        c for c in sl.by_category.get("character", [])
        if c.canonical_tag_id == raven.id
    )
    assert m.source == "ccip"