diff --git a/backend/app/services/ml/heads.py b/backend/app/services/ml/heads.py index b0907c0..aed4e69 100644 --- a/backend/app/services/ml/heads.py +++ b/backend/app/services/ml/heads.py @@ -287,10 +287,14 @@ async def _current_heads(session: AsyncSession, embedding_version: str): return loaded -async def score_image(session: AsyncSession, image_id: int) -> list[dict]: +async def score_image( + session: AsyncSession, image_id: int, threshold_override: float | None = None, +) -> list[dict]: """Suggestions for one image from the trained heads: [{tag_id, name, - category, score}], score >= each head's suggest_threshold, ranked. Empty if - the image has no embedding or no heads exist yet.""" + category, score}], ranked. A concept surfaces when its score clears the + head's own suggest_threshold — or, when threshold_override is given (the + typed-dropdown "show everything" mode), that flat floor instead (0 → every + head). Empty if the image has no embedding or no heads exist yet.""" import numpy as np img = await session.get(ImageRecord, image_id) @@ -307,7 +311,8 @@ async def score_image(session: AsyncSession, image_id: int) -> list[dict]: probs = 1.0 / (1.0 + np.exp(-z)) out = [] for i, p in enumerate(probs): - if p >= heads["thr"][i]: + cut = threshold_override if threshold_override is not None else heads["thr"][i] + if p >= cut: m = heads["meta"][i] out.append({ "tag_id": m["tag_id"], diff --git a/backend/app/services/ml/suggestions.py b/backend/app/services/ml/suggestions.py index 43a488c..1cb306f 100644 --- a/backend/app/services/ml/suggestions.py +++ b/backend/app/services/ml/suggestions.py @@ -1,24 +1,22 @@ -"""The suggestion read-path: raw predictions + centroids -> alias-resolved, -threshold-filtered, category-grouped, ranked suggestions for one image. +"""The suggestion read-path: trained HEADS score one image's frozen embedding +into alias-resolved, category-grouped, ranked suggestions. + +Tagging-v2 (#114): suggestions now come from the per-concept heads that LEARN +from the operator's tags (services/ml/heads.py) — the Camie prediction source +and the per-tag SigLIP centroid have been REMOVED. A head exists only for an +existing concept tag, so every suggestion is a canonical tag (no raw model key, +no alias remap, no creates-new). Rejected tags stay in the list FLAGGED (not +dropped) so the rail can show + reverse a dismissal. """ from dataclasses import dataclass, field -from sqlalchemy import func, select +from sqlalchemy import select from sqlalchemy.ext.asyncio import AsyncSession -from ...models import ( - ImagePrediction, - ImageRecord, - MLSettings, - Tag, - TagSuggestionRejection, -) +from ...models import ImageRecord, TagSuggestionRejection from ...models.tag import image_tag -from .aliases import AliasService -from .centroids import CentroidService -from .tag_name import normalize as normalize_tag_name -from .tagger import SURFACED_CATEGORIES +from .heads import score_image @dataclass(frozen=True) @@ -29,7 +27,7 @@ class Suggestion: display_name: str category: str score: float - source: str # 'tagger' | 'centroid' | 'both' + source: str # 'head' (Camie 'tagger'/'centroid' sources removed in v2) creates_new_tag: bool # raw_name = the booru model vocab key behind this suggestion. It's the key # an alias MUST be stored under (resolution looks up the raw key), so the @@ -54,67 +52,24 @@ class SuggestionList: class SuggestionService: def __init__(self, session: AsyncSession): self.session = session - self.aliases = AliasService(session) - self.centroids = CentroidService(session) - - async def _settings(self) -> MLSettings: - return ( - await self.session.execute(select(MLSettings).where(MLSettings.id == 1)) - ).scalar_one() - - async def _load_predictions(self, image_id: int) -> dict: - """Predictions for one image from the normalized image_prediction - table (#768), in the {raw_name: {category, confidence}} shape the rest - of this service consumed from the old JSON column — so all downstream - threshold/alias/merge logic is unchanged.""" - rows = ( - await self.session.execute( - select( - ImagePrediction.raw_name, - ImagePrediction.category, - ImagePrediction.score, - ).where(ImagePrediction.image_record_id == image_id) - ) - ).all() - return { - r.raw_name: {"category": r.category, "confidence": r.score} - for r in rows - } - - def _threshold_for( - self, s: MLSettings, category: str, override: float | None = None, - ) -> float: - # 'artist' (FC-2d-vii-c) and 'copyright' (2026-06-01) retired; - # both fall through to the 1.01 "never surfaces" default like any - # unsurfaced category. - # override (the typed-dropdown "show everything the model saw" mode) - # applies to the surfaced categories only — unsurfaced ones are already - # skipped before the threshold check, so they can't leak in. - if override is not None: - return override - return { - "character": s.suggestion_threshold_character, - "general": s.suggestion_threshold_general, - }.get(category, 1.01) async def for_image( - self, image_id: int, *, threshold_override: float | None = None, + self, image_id: int, threshold_override: float | None = None, ) -> SuggestionList: - """Ranked suggestions for one image. + """Head-scored suggestions for one image, grouped by category and ranked. - threshold_override surfaces EVERY stored tagger prediction (down to the - ingest STORE_FLOOR) regardless of the configured per-category suggestion - thresholds — backs the tag-input dropdown's "search all of the model's - predictions, including low-confidence ones, in the canonical formatting" - mode (operator-asked 2026-06-09). The Suggestions panel still calls with - no override so it stays the curated above-threshold list.""" + Each trained head scores the image's frozen embedding; a concept surfaces + when its score clears the head's own suggest threshold. threshold_override + (used by the typed tag-input dropdown's "show everything" mode) replaces + that per-head cut with a flat floor (0 → every head), so a low-scoring + concept can still be typed + picked in canonical formatting. + + Already-applied tags are dropped; rejected tags stay FLAGGED and sink to + the bottom of their category so a dismissal is visible + reversible.""" img = await self.session.get(ImageRecord, image_id) if img is None: return SuggestionList() - settings = await self._settings() - predictions: dict = await self._load_predictions(image_id) - applied = set( ( await self.session.execute( @@ -134,149 +89,26 @@ class SuggestionService: ).scalars().all() ) - # --- Camie predictions --- - # candidates carry (raw_name, display_name, category, confidence). - # raw_name = the booru-formatted vocab key, kept for alias_map - # lookup since alias rows are hand-curated against raw keys. - # display_name = normalize_tag_name(raw_name) — what the operator - # sees AND what gets written to tag.name on Accept. - candidates: list[tuple[str, str, str, float]] = [] - for name, p in predictions.items(): - category = p.get("category", "general") - if category not in SURFACED_CATEGORIES: - continue - conf = float(p.get("confidence", 0.0)) - if conf < self._threshold_for(settings, category, threshold_override): - continue - display = normalize_tag_name(name) - if display is None: - # emoticon / pure-punctuation vocab entry — drop entirely - continue - candidates.append((name, display, category, conf)) - - alias_map = await self.aliases.resolve_many( - [(raw, c) for raw, _disp, c, _conf in candidates] + hits = await score_image( + self.session, image_id, threshold_override=threshold_override ) - merged: dict[object, Suggestion] = {} - - def _merge(key, sug: Suggestion): - existing = merged.get(key) - if existing is None: - merged[key] = sug - elif sug.score > existing.score: - merged[key] = Suggestion( - canonical_tag_id=existing.canonical_tag_id, - display_name=existing.display_name, - category=existing.category, - score=sug.score, - source="both" - if existing.source != sug.source - else existing.source, - creates_new_tag=existing.creates_new_tag, - # Keep the alias identity from `existing`: the tagger pass - # (which carries raw_name / via_alias) runs before centroid - # augmentation, so it's always the first writer for a key. - raw_name=existing.raw_name, - via_alias=existing.via_alias, - # rejected is a property of the tag_id, so both writers for a - # key agree — preserve it through the higher-score rebuild. - rejected=existing.rejected, - ) - - for raw, display, category, conf in candidates: - canonical = alias_map.get((raw, category)) - if canonical is not None: - if canonical.id in applied: - continue - _merge( - canonical.id, - Suggestion( - canonical_tag_id=canonical.id, - display_name=canonical.name, - category=category, - score=conf, - source="tagger", - creates_new_tag=False, - raw_name=raw, - via_alias=True, - rejected=canonical.id in rejected, - ), - ) - else: - # Case-insensitive match on BOTH the raw camie key AND - # the normalized form — covers legacy underscore-named - # Tag rows accepted before normalization shipped, AND - # any tag the operator created with the human form. - existing_tag = ( - await self.session.execute( - select(Tag).where( - func.lower(Tag.name).in_( - [raw.lower(), display.lower()] - ) - ) - ) - ).scalars().first() - if existing_tag is not None: - if existing_tag.id in applied: - continue - _merge( - existing_tag.id, - Suggestion( - canonical_tag_id=existing_tag.id, - display_name=existing_tag.name, - category=category, - score=conf, - source="tagger", - creates_new_tag=False, - raw_name=raw, - via_alias=False, - rejected=existing_tag.id in rejected, - ), - ) - else: - _merge( - f"raw:{display}:{category}", - Suggestion( - canonical_tag_id=None, - display_name=display, - category=category, - score=conf, - source="tagger", - creates_new_tag=True, - raw_name=raw, - via_alias=False, - ), - ) - - # --- Centroid augmentation --- - hits = await self.centroids.find_similar_tags(image_id, limit=30) - for hit in hits: - if hit.similarity < settings.centroid_similarity_threshold: - continue - if hit.tag_id in applied: - continue - tag = await self.session.get(Tag, hit.tag_id) - if tag is None: - continue - cat = tag.kind.value if hasattr(tag.kind, "value") else str(tag.kind) - display_cat = cat if cat in SURFACED_CATEGORIES else "general" - _merge( - tag.id, - Suggestion( - canonical_tag_id=tag.id, - display_name=tag.name, - category=display_cat, - score=hit.similarity, - source="centroid", - creates_new_tag=False, - rejected=hit.tag_id in rejected, - ), - ) - result = SuggestionList() - for sug in merged.values(): - result.by_category.setdefault(sug.category, []).append(sug) + for h in hits: + tag_id = h["tag_id"] + if tag_id in applied: + continue + result.by_category.setdefault(h["category"], []).append( + Suggestion( + canonical_tag_id=tag_id, + display_name=h["name"], + category=h["category"], + score=h["score"], + source="head", + creates_new_tag=False, + rejected=tag_id in rejected, + ) + ) for cat in result.by_category: # Live suggestions first (by score), rejected ones sink to the # bottom of the category — visible for recovery, out of the way. @@ -307,7 +139,7 @@ class SuggestionService: for s in items: if s.canonical_tag_id is None or s.creates_new_tag: continue - # for_image now keeps rejected tags (flagged) for the rail; + # for_image keeps rejected tags (flagged) for the rail; # bulk consensus must still ignore them — a tag dismissed on # an image isn't a suggestion for that image. if s.rejected: diff --git a/tests/test_api_suggestions.py b/tests/test_api_suggestions.py index 165b43e..ac02374 100644 --- a/tests/test_api_suggestions.py +++ b/tests/test_api_suggestions.py @@ -1,7 +1,8 @@ import pytest +from sqlalchemy import select from backend.app.celery_app import celery -from backend.app.models import ImageRecord, TagKind +from backend.app.models import ImageRecord, MLSettings, TagHead, TagKind from backend.app.services.tag_service import TagService pytestmark = pytest.mark.integration @@ -31,13 +32,30 @@ async def _img(db, preds, sha="s" * 64): @pytest.mark.asyncio async def test_get_suggestions(client, db): - img = await _img( - db, {"sword": {"category": "general", "confidence": 0.97}} + # Suggestions come from a trained head now (Camie/centroid removed): an image + # whose embedding aligns with the head surfaces that concept. + s = (await db.execute(select(MLSettings).where(MLSettings.id == 1))).scalar_one() + img = ImageRecord( + path="/images/headsug.jpg", sha256="h" * 64, size_bytes=1, + mime="image/jpeg", width=1, height=1, origin="imported_filesystem", + integrity_status="unknown", siglip_embedding=[3.0] + [0.0] * 1151, ) + db.add(img) + await db.flush() + tag = await TagService(db).find_or_create("sword", TagKind.general) + db.add(TagHead( + tag_id=tag.id, embedding_version=s.embedder_model_version, + weights=[1.0] + [0.0] * 1151, bias=0.0, suggest_threshold=0.5, + auto_apply_threshold=None, n_pos=10, n_neg=30, + ap=0.8, precision_cv=0.9, recall=0.6, + )) + await db.commit() resp = await client.get(f"/api/images/{img.id}/suggestions") assert resp.status_code == 200 body = await resp.get_json() - assert "general" in body["by_category"] + general = body["by_category"].get("general", []) + s2 = next(x for x in general if x["canonical_tag_id"] == tag.id) + assert s2["source"] == "head" @pytest.mark.asyncio @@ -121,67 +139,3 @@ async def test_alias_requires_fields(client, db): f"/api/images/{img.id}/suggestions/alias", json={"alias_string": "x"} ) assert resp.status_code == 400 - - -async def _img_at(db, path, sha, preds): - from tests._prediction_helpers import seed_predictions - - img = ImageRecord( - path=path, sha256=sha, size_bytes=1, mime="image/jpeg", - width=1, height=1, origin="imported_filesystem", - integrity_status="unknown", - ) - db.add(img) - await db.commit() - await seed_predictions(db, img.id, preds) - await db.commit() - return img - - -@pytest.mark.asyncio -async def test_alias_roundtrip_resolves_by_raw_key(client, db): - """Locks the modal-alias contract: the suggestion exposes the RAW model key, - an alias authored with that key resolves on a later image, and the resolved - suggestion is flagged via_alias. (Pre-fix the modal stored the normalized - display name, which never resolved.)""" - canonical = await TagService(db).find_or_create( - "Sasuke Uchiha", TagKind.character - ) - await db.commit() - preds = {"uchiha_sasuke": {"category": "character", "confidence": 0.99}} - img_a = await _img_at(db, "/images/alias_a.jpg", "a" * 64, preds) - - # (a) raw_name is exposed so the modal can author the alias with it; the - # raw prediction doesn't textually match the tag, so it'd otherwise be +new. - body = await ( - await client.get(f"/api/images/{img_a.id}/suggestions") - ).get_json() - sug = body["by_category"]["character"][0] - assert sug["raw_name"] == "uchiha_sasuke" - assert sug["via_alias"] is False - assert sug["creates_new_tag"] is True - - # Author the alias keyed by the RAW key (what the frontend now sends). - resp = await client.post( - f"/api/images/{img_a.id}/suggestions/alias", - json={ - "alias_string": sug["raw_name"], - "alias_category": "character", - "canonical_tag_id": canonical.id, - }, - ) - assert resp.status_code == 200 - assert (await resp.get_json())["allowlisted"] is True - - # (b) A DIFFERENT image with the same prediction now resolves via the alias - # (image A's tag is applied, so it's filtered there). Had the alias been - # stored under the display name, this would NOT resolve. - img_b = await _img_at(db, "/images/alias_b.jpg", "b" * 64, preds) - body_b = await ( - await client.get(f"/api/images/{img_b.id}/suggestions") - ).get_json() - sug_b = body_b["by_category"]["character"][0] - assert sug_b["canonical_tag_id"] == canonical.id - assert sug_b["via_alias"] is True - assert sug_b["creates_new_tag"] is False - assert sug_b["raw_name"] == "uchiha_sasuke" diff --git a/tests/test_ml_artist_retired.py b/tests/test_ml_artist_retired.py index 955583a..0a3e208 100644 --- a/tests/test_ml_artist_retired.py +++ b/tests/test_ml_artist_retired.py @@ -14,14 +14,12 @@ def test_artist_not_centroid_eligible(): assert TagKind.artist not in ELIGIBLE_KINDS -def test_threshold_for_artist_is_unsurfaced(): - from backend.app.services.ml.suggestions import SuggestionService +def test_artist_not_head_eligible(): + # Tagging-v2: suggestions come from heads, and heads are only trained for + # general/character concepts — so 'artist' (and any other kind) can't surface. + from backend.app.models import TagKind + from backend.app.services.ml.heads import _HEAD_KINDS - class _S: - suggestion_threshold_character = 0.5 - suggestion_threshold_general = 0.5 - - svc = SuggestionService.__new__(SuggestionService) - # 'artist' and 'copyright' both retired — fall through to 1.01 - assert svc._threshold_for(_S(), "artist") == 1.01 - assert svc._threshold_for(_S(), "copyright") == 1.01 + assert TagKind.general in _HEAD_KINDS + assert TagKind.character in _HEAD_KINDS + assert TagKind.artist not in _HEAD_KINDS diff --git a/tests/test_ml_suggestions.py b/tests/test_ml_suggestions.py index c3cb276..7a427e2 100644 --- a/tests/test_ml_suggestions.py +++ b/tests/test_ml_suggestions.py @@ -1,8 +1,11 @@ +"""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, TagKind +from backend.app.models import ImageRecord, MLSettings, TagHead, TagKind from backend.app.models.tag import image_tag -from backend.app.services.ml.aliases import AliasService from backend.app.services.ml.allowlist import AllowlistService from backend.app.services.ml.suggestions import SuggestionService from backend.app.services.tag_service import TagService @@ -10,179 +13,121 @@ from backend.app.services.tag_service import TagService pytestmark = pytest.mark.integration -def _img(sha: str) -> ImageRecord: - return ImageRecord( - path=f"/images/{sha}.jpg", - sha256=sha, - size_bytes=1, - mime="image/jpeg", - width=1, - height=1, - origin="imported_filesystem", - integrity_status="unknown", +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, ) - - -async def _seed_img(db, sha: str, predictions: dict) -> ImageRecord: - """#768: create an image + seed its predictions into image_prediction - (the read path's source), returning the flushed record.""" - from tests._prediction_helpers import seed_predictions - - img = _img(sha) db.add(img) await db.flush() - await seed_predictions(db, img.id, predictions) 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_threshold_filters_low_confidence_general(db): - # Default general threshold is 0.50 (alembic 0029 lowered it from - # 0.95). Use 0.30/0.60 to keep the test asserting threshold behavior - # rather than the exact cutoff number. - img = await _seed_img( - db, - "a" * 64, - { - "lowconf": {"category": "general", "confidence": 0.30}, - "sword": {"category": "general", "confidence": 0.97}, - }, - ) +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) - names = [s.display_name for s in sl.by_category.get("general", [])] - # display_name is normalized (tag_name.normalize) before surfacing. - assert "Sword" in names - assert "Lowconf" not in names + 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_threshold_override_surfaces_low_confidence(db): - # The typed-dropdown "show everything the model saw" mode: threshold_override - # surfaces stored predictions below the configured threshold (in canonical - # formatting) so they can be picked instead of hand-typed (2026-06-09). - img = await _seed_img( - db, - "d" * 64, - { - "lowconf": {"category": "general", "confidence": 0.30}, - "sword": {"category": "general", "confidence": 0.97}, - }, - ) - sl = await SuggestionService(db).for_image(img.id, threshold_override=0.0) - names = [s.display_name for s in sl.by_category.get("general", [])] - assert "Sword" in names - assert "Lowconf" in names # below the configured threshold, surfaced anyway - - # Unsurfaced categories are still excluded even with the override. - img2 = await _seed_img( - db, "e" * 64, {"safe": {"category": "rating", "confidence": 0.99}} - ) - sl2 = await SuggestionService(db).for_image(img2.id, threshold_override=0.0) - assert "rating" not in sl2.by_category +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_unsurfaced_category_dropped(db): - img = await _seed_img( - db, - "b" * 64, - {"safe": {"category": "rating", "confidence": 0.99}}, - ) - sl = await SuggestionService(db).for_image(img.id) - assert "rating" not in sl.by_category - - -@pytest.mark.asyncio -async def test_alias_resolution(db): - tags = TagService(db) - canonical = await tags.find_or_create("Sasuke Uchiha", TagKind.character) - await AliasService(db).create("uchiha_sasuke", "character", canonical.id) - img = await _seed_img( - db, - "c" * 64, - {"uchiha_sasuke": {"category": "character", "confidence": 0.96}}, - ) - sl = await SuggestionService(db).for_image(img.id) - chars = sl.by_category["character"] - assert len(chars) == 1 - assert chars[0].display_name == "Sasuke Uchiha" - assert chars[0].canonical_tag_id == canonical.id - assert chars[0].creates_new_tag is False - # Surfaced via an alias on the raw model key — the UI marks it + offers undo. - assert chars[0].via_alias is True - assert chars[0].raw_name == "uchiha_sasuke" - - -@pytest.mark.asyncio -async def test_raw_tag_creates_new(db): - img = await _seed_img( - db, - "d" * 64, - {"brand_new_tag": {"category": "character", "confidence": 0.96}}, - ) - sl = await SuggestionService(db).for_image(img.id) - chars = sl.by_category["character"] - # display_name is the normalized Camie name (underscores -> spaces, - # title-cased), not the raw vocab key. - assert chars[0].display_name == "Brand New Tag" - assert chars[0].creates_new_tag is True - # Not aliased, but the raw key is carried so the modal can author one. - assert chars[0].via_alias is False - assert chars[0].raw_name == "brand_new_tag" - assert chars[0].canonical_tag_id is None +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): - tags = TagService(db) - tag = await tags.find_or_create("alreadyhere", TagKind.character) - img = await _seed_img( - db, - "e" * 64, - {"alreadyhere": {"category": "character", "confidence": 0.96}}, - ) + 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 "character" not in sl.by_category or not sl.by_category["character"] + 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 in the list flagged - # rejected=True so the rail can show it + offer one-click un-reject - # (visible, reversible rejection — operator-asked 2026-06-27). A live - # suggestion sorts ahead of the rejected one. - tags = TagService(db) - rejected_tag = await tags.find_or_create("rejectme", TagKind.general) - img = await _seed_img( - db, - "f" * 64, - { - "rejectme": {"category": "general", "confidence": 0.96}, - "keepme": {"category": "general", "confidence": 0.90}, - }, - ) - await AllowlistService(db).dismiss(img.id, rejected_tag.id) + # 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) - general = sl.by_category["general"] - # Match by id, not display casing (an existing tag keeps its stored name). - rej = next(s for s in general if s.canonical_tag_id == rejected_tag.id) - assert rej.rejected is True - live = [s for s in general if not s.rejected] - assert live, "the un-rejected 'keepme' suggestion should still surface" - # Live suggestions sort ahead of rejected ones regardless of score. - assert general[-1].canonical_tag_id == rejected_tag.id + s = next(x for x in sl.by_category["general"] if x.canonical_tag_id == tag.id) + assert s.rejected is True - # Un-reject reverts it to a live suggestion. - await AllowlistService(db).undismiss(img.id, rejected_tag.id) + await AllowlistService(db).undismiss(img.id, tag.id) + await db.commit() sl2 = await SuggestionService(db).for_image(img.id) - rej2 = next( - s for s in sl2.by_category["general"] - if s.canonical_tag_id == rejected_tag.id - ) - assert rej2.rejected is False + s2 = next(x for x in sl2.by_category["general"] if x.canonical_tag_id == tag.id) + assert s2.rejected is False diff --git a/tests/test_suggestions_bulk.py b/tests/test_suggestions_bulk.py index d0d101a..5df13f2 100644 --- a/tests/test_suggestions_bulk.py +++ b/tests/test_suggestions_bulk.py @@ -1,88 +1,84 @@ +"""Consensus (for_selection) over the tagging-v2 HEAD suggestion source.""" import pytest +from sqlalchemy import select from backend.app import create_app -from backend.app.models import ImageRecord, TagKind +from backend.app.models import ImageRecord, MLSettings, TagHead, TagKind from backend.app.models.tag import image_tag from backend.app.services.ml.suggestions import SuggestionService from backend.app.services.tag_service import TagService -from tests._prediction_helpers import seed_predictions pytestmark = pytest.mark.integration -def _img(sha: str) -> ImageRecord: - return ImageRecord( - path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, - mime="image/jpeg", width=1, height=1, - origin="imported_filesystem", integrity_status="unknown", +def _emb(slot: int) -> list[float]: + v = [0.0] * 1152 + v[slot] = 3.0 + 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 _head(db, tag_id: int, slot: int = 0): + s = (await db.execute(select(MLSettings).where(MLSettings.id == 1))).scalar_one() + weights = [0.0] * 1152 + weights[slot] = 1.0 + db.add(TagHead( + tag_id=tag_id, embedding_version=s.embedder_model_version, + weights=weights, bias=0.0, suggest_threshold=0.5, + 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_consensus_includes_tag_over_threshold(db): - tags = TagService(db) - t = await tags.find_or_create("sword", TagKind.general) - a = _img("a" * 64) - b = _img("b" * 64) - db.add_all([a, b]) - await db.flush() - await seed_predictions(db, a.id, {"sword": {"category": "general", "confidence": 0.97}}) - await seed_predictions(db, b.id, {"sword": {"category": "general", "confidence": 0.95}}) + t = await TagService(db).find_or_create("sword", TagKind.general) + a = await _img(db, "a" * 64, _emb(0)) + b = await _img(db, "b" * 64, _emb(0)) + await _head(db, t.id, slot=0) + await db.commit() res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8) - gen = res["general"] - assert any(s["canonical_tag_id"] == t.id for s in gen) - s = next(s for s in gen if s["canonical_tag_id"] == t.id) - assert s["coverage"] == 1.0 - assert 0.95 <= s["confidence"] <= 0.97 + s = next(s for s in res["general"] if s["canonical_tag_id"] == t.id) + assert s["coverage"] == 1.0 # suggested on both + assert s["confidence"] > 0.5 @pytest.mark.asyncio async def test_consensus_counts_already_applied_for_coverage(db): - tags = TagService(db) - t = await tags.find_or_create("sky", TagKind.general) - a = _img("c" * 64) - b = _img("d" * 64) # no prediction - db.add_all([a, b]) - await db.flush() - await seed_predictions(db, a.id, {"sky": {"category": "general", "confidence": 0.96}}) - # b already has the tag applied -> counts toward coverage, not confidence + t = await TagService(db).find_or_create("sky", TagKind.general) + a = await _img(db, "c" * 64, _emb(0)) # head suggests it + b = await _img(db, "d" * 64, None) # no embedding; tag applied instead + await _head(db, t.id, slot=0) await db.execute( image_tag.insert().values( image_record_id=b.id, tag_id=t.id, source="manual" ) ) + await db.commit() res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8) s = next(s for s in res["general"] if s["canonical_tag_id"] == t.id) assert s["coverage"] == 1.0 # 1 suggested + 1 applied / 2 - assert s["confidence"] == pytest.approx(0.96, abs=1e-4) @pytest.mark.asyncio async def test_consensus_excludes_below_threshold(db): - tags = TagService(db) - await tags.find_or_create("rare", TagKind.general) - a = _img("e" * 64) - b = _img("f" * 64) - db.add_all([a, b]) - await db.flush() - await seed_predictions(db, a.id, {"rare": {"category": "general", "confidence": 0.96}}) + t = await TagService(db).find_or_create("rare", TagKind.general) + a = await _img(db, "e" * 64, _emb(0)) # suggested here + b = await _img(db, "f" * 64, None) # not here → coverage 0.5 < 0.8 + await _head(db, t.id, slot=0) + await db.commit() res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8) - assert all( - s["name"] != "rare" for s in res.get("general", []) - ) # coverage 0.5 < 0.8 - - -@pytest.mark.asyncio -async def test_consensus_skips_creates_new_tag(db): - a = _img("g" * 64) - b = _img("h" * 64) - db.add_all([a, b]) - await db.flush() - await seed_predictions(db, a.id, {"neverseen": {"category": "general", "confidence": 0.99}}) - await seed_predictions(db, b.id, {"neverseen": {"category": "general", "confidence": 0.99}}) - res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8) - # 'neverseen' has no Tag row -> creates_new_tag -> excluded from consensus - assert all(s["name"] != "neverseen" for s in res.get("general", [])) + assert all(s["name"] != "rare" for s in res.get("general", [])) @pytest.mark.asyncio @@ -93,13 +89,9 @@ async def test_consensus_threshold_clamped_and_empty_for_no_ids(db): @pytest.mark.asyncio async def test_bulk_suggestions_route(db): - - tags = TagService(db) - await tags.find_or_create("sword", TagKind.general) - a = _img("i" * 64) - db.add(a) - await db.commit() - await seed_predictions(db, a.id, {"sword": {"category": "general", "confidence": 0.97}}) + t = await TagService(db).find_or_create("sword", TagKind.general) + a = await _img(db, "i" * 64, _emb(0)) + await _head(db, t.id, slot=0) await db.commit() app = create_app() async with app.test_client() as c: @@ -115,7 +107,6 @@ async def test_bulk_suggestions_route(db): @pytest.mark.asyncio async def test_bulk_suggestions_requires_ids(db): - app = create_app() async with app.test_client() as c: resp = await c.post("/api/suggestions/bulk", json={})