diff --git a/backend/app/services/ml/ccip.py b/backend/app/services/ml/ccip.py index 0a30def..bf09a09 100644 --- a/backend/app/services/ml/ccip.py +++ b/backend/app/services/ml/ccip.py @@ -16,7 +16,14 @@ 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 import ( + CcipPrototypeState, + CharacterPrototype, + ImageRegion, + MLSettings, + Tag, + TagKind, +) from ...models.tag import image_tag # Cosine-similarity floor to call a figure the same character. The live setting @@ -148,6 +155,60 @@ async def _tag_names(session: AsyncSession, tag_ids: list[int]) -> dict[int, str ) +# Per-character normalized prototype matrices, cached per process and refreshed +# INCREMENTALLY: only characters whose ccip_prototype_state.updated_at advanced +# are reloaded. This replaces the request-path rebuild of the ENTIRE reference +# blob (the ~4s stall, #1317) — the prototypes are precomputed off the request +# path by services.ml.character_prototypes (a beat + after each retrain). +_PROTO_CACHE: dict = {"mats": {}, "ver": {}} + + +async def _load_prototypes(session: AsyncSession) -> dict: + """{tag_id: (P, D) L2-normalized prototype matrix} from character_prototype, + served from the in-process cache and reloading ONLY the characters whose + updated_at changed. Empty dict when the store isn't populated yet (cold start + → match_image falls back to the legacy on-the-fly reference build).""" + import numpy as np + + versions = dict( + ( + await session.execute( + select(CcipPrototypeState.tag_id, CcipPrototypeState.updated_at) + ) + ).all() + ) + mats = _PROTO_CACHE["mats"] + ver = _PROTO_CACHE["ver"] + # Forget characters that no longer have prototypes. + for tag_id in [t for t in mats if t not in versions]: + mats.pop(tag_id, None) + ver.pop(tag_id, None) + # Reload only the characters whose prototypes changed since we cached them. + stale = [t for t, u in versions.items() if ver.get(t) != u] + if stale: + rows = ( + await session.execute( + select( + CharacterPrototype.tag_id, CharacterPrototype.ccip_embedding + ).where(CharacterPrototype.tag_id.in_(stale)) + ) + ).all() + by_tag: dict[int, list] = {} + for tag_id, vec in rows: + by_tag.setdefault(tag_id, []).append( + np.asarray(vec, dtype=np.float32) + ) + for tag_id in stale: + vecs = by_tag.get(tag_id) + if vecs: + mats[tag_id] = _l2norm(np.vstack(vecs), np) + ver[tag_id] = versions[tag_id] + else: + mats.pop(tag_id, None) + ver.pop(tag_id, None) + return mats + + async def match_image( session: AsyncSession, image_id: int, threshold: float | None = None ) -> list[dict]: @@ -178,7 +239,13 @@ async def match_image( ).all() if not fig_rows: return [] - refs = await character_references(session) + # Prefer the precomputed prototype store (fast, incremental). On a cold start + # (store not yet populated post-deploy) fall back to the legacy on-the-fly + # reference build so character suggestions work immediately — the background + # refresh populates the store within ~15 min, after which this path is used + # and the per-accept ~4s rebuild is gone (#1317). + protos = await _load_prototypes(session) + refs = protos if protos else await character_references(session) if not refs: return [] applied = set( @@ -202,7 +269,11 @@ async def match_image( 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) + # Prototype matrices are already L2-normalized; legacy refs are raw + # vector lists that still need stacking + normalizing. + R = vecs if protos else _l2norm( + np.vstack([np.asarray(v, dtype=np.float32) for v in vecs]), np + ) sims = Q @ R.T # (n_query_figures, n_references) per_figure = sims.max(axis=1) # best reference cosine per figure best_figure = int(per_figure.argmax()) diff --git a/tests/test_ccip.py b/tests/test_ccip.py index 8231c4b..df52378 100644 --- a/tests/test_ccip.py +++ b/tests/test_ccip.py @@ -3,7 +3,13 @@ model needed, so it runs in CI with synthetic CCIP vectors.""" import pytest from sqlalchemy import select -from backend.app.models import ImageRecord, ImageRegion, TagKind +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 @@ -61,6 +67,24 @@ async def test_matches_same_character_across_images(db): 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)