FabledCurator/tests/test_ccip.py

"""CCIP few-shot character matcher (#114). numpy cosine on stored vectors — no
model needed, so it runs in CI with synthetic CCIP vectors."""
import pytest
from sqlalchemy import select

from backend.app.models import (
    CcipPrototypeState,
    CharacterPrototype,
    ImageRecord,
    ImageRegion,
    TagKind,
)
from backend.app.models.tag import image_tag
from backend.app.services.ml.ccip import match_image
from backend.app.services.tag_service import TagService

pytestmark = pytest.mark.integration


def _ccip(slot: int) -> list[float]:
    v = [0.0] * 768
    v[slot] = 1.0
    return v


async def _img(db, sha) -> 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",
    )
    db.add(img)
    await db.flush()
    return img


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


async def _tag_image(db, image_id, tag_id):
    await db.execute(image_tag.insert().values(
        image_record_id=image_id, tag_id=tag_id, source="manual",
    ))


@pytest.mark.asyncio
async def test_matches_same_character_across_images(db):
    raven = await TagService(db).find_or_create("Raven", TagKind.character)
    ref = await _img(db, "a" * 64)               # a tagged example = a prototype
    await _figure(db, ref.id, _ccip(0))
    await _tag_image(db, ref.id, raven.id)
    query = await _img(db, "b" * 64)             # untagged, near-identical figure
    await _figure(db, query.id, _ccip(0))
    await db.commit()

    matches = await match_image(db, query.id)
    m = next(x for x in matches if x["tag_id"] == raven.id)
    assert m["source"] == "ccip" and m["category"] == "character"
    assert m["score"] > 0.9
    # #1206: the match grounds to the figure region that matched (hover → the
    # figure box lights up), so a character suggestion is localized too.
    assert m["grounding"]["bbox"] == pytest.approx([0.0, 0.0, 1.0, 1.0])
    assert m["grounding"]["kind"] == "figure"


@pytest.mark.asyncio
async def test_matches_from_prototype_store(db):
    # Once the prototype store is populated (#1317), match_image reads it (the
    # fast incremental path) instead of rebuilding references on the request —
    # same match + grounding as the legacy build.
    raven = await TagService(db).find_or_create("Raven", TagKind.character)
    db.add(CharacterPrototype(tag_id=raven.id, ccip_embedding=_ccip(0)))
    db.add(CcipPrototypeState(tag_id=raven.id, fingerprint="1:1"))
    query = await _img(db, "n" * 64)
    await _figure(db, query.id, _ccip(0))       # query figure matches the prototype
    await db.commit()

    matches = await match_image(db, query.id)
    m = next(x for x in matches if x["tag_id"] == raven.id)
    assert m["source"] == "ccip" and m["score"] > 0.9
    assert m["grounding"]["kind"] == "figure"


@pytest.mark.asyncio
async def test_no_match_for_different_character(db):
    raven = await TagService(db).find_or_create("Raven", TagKind.character)
    ref = await _img(db, "c" * 64)
    await _figure(db, ref.id, _ccip(0))
    await _tag_image(db, ref.id, raven.id)
    query = await _img(db, "d" * 64)
    await _figure(db, query.id, _ccip(5))        # orthogonal → not Raven
    await db.commit()
    assert await match_image(db, query.id) == []


@pytest.mark.asyncio
async def test_excludes_already_applied_character(db):
    raven = await TagService(db).find_or_create("Raven", TagKind.character)
    ref = await _img(db, "e" * 64)
    await _figure(db, ref.id, _ccip(0))
    await _tag_image(db, ref.id, raven.id)
    query = await _img(db, "f" * 64)
    await _figure(db, query.id, _ccip(0))
    await _tag_image(db, query.id, raven.id)     # already tagged → no re-suggest
    await db.commit()
    assert all(m["tag_id"] != raven.id for m in await match_image(db, query.id))


@pytest.mark.asyncio
async def test_no_figure_vectors_means_no_match(db):
    query = await _img(db, "g" * 64)
    await db.commit()
    assert await match_image(db, query.id) == []


@pytest.mark.asyncio
async def test_threshold_gates_borderline_match(db):
    # A figure ~0.9 cosine from the reference: matched at 0.85, dropped at 0.95.
    raven = await TagService(db).find_or_create("Raven", TagKind.character)
    ref = await _img(db, "h" * 64)
    await _figure(db, ref.id, _ccip(0))                 # e0
    await _tag_image(db, ref.id, raven.id)
    near = [0.0] * 768
    near[0], near[1] = 0.9, 0.4359                       # |·|=1, cos(e0)=0.9
    query = await _img(db, "i" * 64)
    await _figure(db, query.id, near)
    await db.commit()

    assert any(m["tag_id"] == raven.id for m in await match_image(db, query.id, 0.85))
    assert await match_image(db, query.id, 0.95) == []


@pytest.mark.asyncio
async def test_multi_character_image_not_used_as_reference(db):
    # A figure on a 2-character image is ambiguous (tag is image-level), so it
    # must NOT seed either character's prototypes — else it'd match both.
    raven = await TagService(db).find_or_create("Raven", TagKind.character)
    daphne = await TagService(db).find_or_create("Daphne", TagKind.character)
    multi = await _img(db, "j" * 64)
    await _figure(db, multi.id, _ccip(0))
    await _tag_image(db, multi.id, raven.id)
    await _tag_image(db, multi.id, daphne.id)
    query = await _img(db, "k" * 64)
    await _figure(db, query.id, _ccip(0))            # identical to the ambiguous figure
    await db.commit()
    assert await match_image(db, query.id) == []     # no clean references → nothing


@pytest.mark.asyncio
async def test_auto_apply_tags_confident_match(db):
    raven = await TagService(db).find_or_create("Raven", TagKind.character)
    ref = await _img(db, "l" * 64)
    await _figure(db, ref.id, _ccip(0))
    await _tag_image(db, ref.id, raven.id)           # single-character reference
    query = await _img(db, "m" * 64)
    await _figure(db, query.id, _ccip(0))            # identical → cosine 1.0
    await db.commit()

    from backend.app.tasks.ml import scheduled_ccip_auto_apply
    assert "applied=" in scheduled_ccip_auto_apply()   # sync task, own session

    rows = (await db.execute(
        select(image_tag.c.tag_id, image_tag.c.source).where(
            image_tag.c.image_record_id == query.id
        )
    )).all()
    assert (raven.id, "ccip_auto") in [(t, s) for t, s in rows]