feat(suggestions): heads are the suggestion source — Camie + centroid removed (#114 C)
CI / lint (push) Successful in 2s
CI / frontend-build (push) Successful in 17s
CI / backend-lint-and-test (push) Successful in 26s
CI / integration (push) Successful in 3m19s

The rail's Suggestions now come from the trained per-concept heads. SuggestionService.for_image scores the image's frozen SigLIP embedding against
every head (heads.score_image) and surfaces concepts above each head's own
suggest threshold; the typed-dropdown's min=0 "show everything" mode maps to a
flat floor so any head-scored concept can still be picked. Already-applied tags
drop; rejected tags stay flagged + reversible (unchanged).

REMOVED from the suggestion path (rule 22, no fallback): the Camie
ImagePrediction candidate/alias/merge pipeline and the per-tag centroid
augmentation, plus the now-dead SuggestionService internals (_load_predictions,
_threshold_for, _settings, self.aliases, self.centroids). Head suggestions are
always canonical tags, so raw_name/via_alias are null/false and the rail's
alias kebab is inert by data (its removal + the Camie ingest-tagger rip are the
flagged follow-up). for_selection (bulk consensus) now aggregates head
suggestions unchanged.

Tests rewritten to the head path: test_ml_suggestions (surfaces/applied/
rejected-reversible/override/no-embedding/no-heads), test_suggestions_bulk
(consensus), test_api_suggestions (get + dropped the Camie-alias roundtrip),
and test_ml_artist_retired (artist not head-eligible via _HEAD_KINDS).

DEPLOY NOTE: after this lands, the rail is empty until you run Train heads
(Settings → Tagging → Concept heads) — deploy, train, then the rail populates.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Ttrj5P7upUTueSfoJcxEqa
This commit is contained in:
2026-06-28 11:20:11 -04:00
parent 06d5e83da4
commit ca1c17446c
6 changed files with 222 additions and 497 deletions
+9 -4
View File
@@ -287,10 +287,14 @@ async def _current_heads(session: AsyncSession, embedding_version: str):
return loaded 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, """Suggestions for one image from the trained heads: [{tag_id, name,
category, score}], score >= each head's suggest_threshold, ranked. Empty if category, score}], ranked. A concept surfaces when its score clears the
the image has no embedding or no heads exist yet.""" 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 import numpy as np
img = await session.get(ImageRecord, image_id) 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)) probs = 1.0 / (1.0 + np.exp(-z))
out = [] out = []
for i, p in enumerate(probs): 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] m = heads["meta"][i]
out.append({ out.append({
"tag_id": m["tag_id"], "tag_id": m["tag_id"],
+41 -209
View File
@@ -1,24 +1,22 @@
"""The suggestion read-path: raw predictions + centroids -> alias-resolved, """The suggestion read-path: trained HEADS score one image's frozen embedding
threshold-filtered, category-grouped, ranked suggestions for one image. 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 dataclasses import dataclass, field
from sqlalchemy import func, select from sqlalchemy import select
from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy.ext.asyncio import AsyncSession
from ...models import ( from ...models import ImageRecord, TagSuggestionRejection
ImagePrediction,
ImageRecord,
MLSettings,
Tag,
TagSuggestionRejection,
)
from ...models.tag import image_tag from ...models.tag import image_tag
from .aliases import AliasService from .heads import score_image
from .centroids import CentroidService
from .tag_name import normalize as normalize_tag_name
from .tagger import SURFACED_CATEGORIES
@dataclass(frozen=True) @dataclass(frozen=True)
@@ -29,7 +27,7 @@ class Suggestion:
display_name: str display_name: str
category: str category: str
score: float score: float
source: str # 'tagger' | 'centroid' | 'both' source: str # 'head' (Camie 'tagger'/'centroid' sources removed in v2)
creates_new_tag: bool creates_new_tag: bool
# raw_name = the booru model vocab key behind this suggestion. It's the key # 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 # an alias MUST be stored under (resolution looks up the raw key), so the
@@ -54,67 +52,24 @@ class SuggestionList:
class SuggestionService: class SuggestionService:
def __init__(self, session: AsyncSession): def __init__(self, session: AsyncSession):
self.session = session 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( async def for_image(
self, image_id: int, *, threshold_override: float | None = None, self, image_id: int, threshold_override: float | None = None,
) -> SuggestionList: ) -> 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 Each trained head scores the image's frozen embedding; a concept surfaces
ingest STORE_FLOOR) regardless of the configured per-category suggestion when its score clears the head's own suggest threshold. threshold_override
thresholds — backs the tag-input dropdown's "search all of the model's (used by the typed tag-input dropdown's "show everything" mode) replaces
predictions, including low-confidence ones, in the canonical formatting" that per-head cut with a flat floor (0 → every head), so a low-scoring
mode (operator-asked 2026-06-09). The Suggestions panel still calls with concept can still be typed + picked in canonical formatting.
no override so it stays the curated above-threshold list."""
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) img = await self.session.get(ImageRecord, image_id)
if img is None: if img is None:
return SuggestionList() return SuggestionList()
settings = await self._settings()
predictions: dict = await self._load_predictions(image_id)
applied = set( applied = set(
( (
await self.session.execute( await self.session.execute(
@@ -134,149 +89,26 @@ class SuggestionService:
).scalars().all() ).scalars().all()
) )
# --- Camie predictions --- hits = await score_image(
# candidates carry (raw_name, display_name, category, confidence). self.session, image_id, threshold_override=threshold_override
# 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]
) )
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() result = SuggestionList()
for sug in merged.values(): for h in hits:
result.by_category.setdefault(sug.category, []).append(sug) 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: for cat in result.by_category:
# Live suggestions first (by score), rejected ones sink to the # Live suggestions first (by score), rejected ones sink to the
# bottom of the category — visible for recovery, out of the way. # bottom of the category — visible for recovery, out of the way.
@@ -307,7 +139,7 @@ class SuggestionService:
for s in items: for s in items:
if s.canonical_tag_id is None or s.creates_new_tag: if s.canonical_tag_id is None or s.creates_new_tag:
continue 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 # bulk consensus must still ignore them — a tag dismissed on
# an image isn't a suggestion for that image. # an image isn't a suggestion for that image.
if s.rejected: if s.rejected:
+22 -68
View File
@@ -1,7 +1,8 @@
import pytest import pytest
from sqlalchemy import select
from backend.app.celery_app import celery 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 from backend.app.services.tag_service import TagService
pytestmark = pytest.mark.integration pytestmark = pytest.mark.integration
@@ -31,13 +32,30 @@ async def _img(db, preds, sha="s" * 64):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_get_suggestions(client, db): async def test_get_suggestions(client, db):
img = await _img( # Suggestions come from a trained head now (Camie/centroid removed): an image
db, {"sword": {"category": "general", "confidence": 0.97}} # 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") resp = await client.get(f"/api/images/{img.id}/suggestions")
assert resp.status_code == 200 assert resp.status_code == 200
body = await resp.get_json() 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 @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"} f"/api/images/{img.id}/suggestions/alias", json={"alias_string": "x"}
) )
assert resp.status_code == 400 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"
+8 -10
View File
@@ -14,14 +14,12 @@ def test_artist_not_centroid_eligible():
assert TagKind.artist not in ELIGIBLE_KINDS assert TagKind.artist not in ELIGIBLE_KINDS
def test_threshold_for_artist_is_unsurfaced(): def test_artist_not_head_eligible():
from backend.app.services.ml.suggestions import SuggestionService # 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: assert TagKind.general in _HEAD_KINDS
suggestion_threshold_character = 0.5 assert TagKind.character in _HEAD_KINDS
suggestion_threshold_general = 0.5 assert TagKind.artist not in _HEAD_KINDS
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
+91 -146
View File
@@ -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 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.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.allowlist import AllowlistService
from backend.app.services.ml.suggestions import SuggestionService from backend.app.services.ml.suggestions import SuggestionService
from backend.app.services.tag_service import TagService from backend.app.services.tag_service import TagService
@@ -10,179 +13,121 @@ from backend.app.services.tag_service import TagService
pytestmark = pytest.mark.integration pytestmark = pytest.mark.integration
def _img(sha: str) -> ImageRecord: def _emb(slot: int, val: float = 3.0) -> list[float]:
return ImageRecord( """An embedding pointing along axis `slot` (so its L2-normalized form is the
path=f"/images/{sha}.jpg", unit vector e_slot — a head with weights e_slot scores it sigmoid(1)≈0.73)."""
sha256=sha, v = [0.0] * 1152
size_bytes=1, v[slot] = val
mime="image/jpeg", return v
width=1,
height=1,
origin="imported_filesystem", async def _img(db, sha: str, emb=None) -> ImageRecord:
integrity_status="unknown", 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) db.add(img)
await db.flush() await db.flush()
await seed_predictions(db, img.id, predictions)
return img 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 @pytest.mark.asyncio
async def test_threshold_filters_low_confidence_general(db): async def test_head_suggestion_surfaces_for_matching_image(db):
# Default general threshold is 0.50 (alembic 0029 lowered it from tag = await TagService(db).find_or_create("glasses", TagKind.general)
# 0.95). Use 0.30/0.60 to keep the test asserting threshold behavior img = await _img(db, "a" * 64, _emb(0))
# rather than the exact cutoff number. await _head(db, tag.id, slot=0)
img = await _seed_img( await db.commit()
db,
"a" * 64,
{
"lowconf": {"category": "general", "confidence": 0.30},
"sword": {"category": "general", "confidence": 0.97},
},
)
sl = await SuggestionService(db).for_image(img.id) sl = await SuggestionService(db).for_image(img.id)
names = [s.display_name for s in sl.by_category.get("general", [])] general = sl.by_category["general"]
# display_name is normalized (tag_name.normalize) before surfacing. assert len(general) == 1
assert "Sword" in names s = general[0]
assert "Lowconf" not in names 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 @pytest.mark.asyncio
async def test_threshold_override_surfaces_low_confidence(db): async def test_no_embedding_means_no_suggestions(db):
# The typed-dropdown "show everything the model saw" mode: threshold_override img = await _img(db, "b" * 64, None)
# surfaces stored predictions below the configured threshold (in canonical tag = await TagService(db).find_or_create("cat", TagKind.general)
# formatting) so they can be picked instead of hand-typed (2026-06-09). await _head(db, tag.id, slot=0)
img = await _seed_img( await db.commit()
db, assert (await SuggestionService(db).for_image(img.id)).by_category == {}
"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
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_unsurfaced_category_dropped(db): async def test_no_heads_means_no_suggestions(db):
img = await _seed_img( img = await _img(db, "c" * 64, _emb(0))
db, await db.commit() # no heads trained yet
"b" * 64, assert (await SuggestionService(db).for_image(img.id)).by_category == {}
{"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
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_applied_tag_not_suggested(db): async def test_applied_tag_not_suggested(db):
tags = TagService(db) tag = await TagService(db).find_or_create("dog", TagKind.general)
tag = await tags.find_or_create("alreadyhere", TagKind.character) img = await _img(db, "d" * 64, _emb(0))
img = await _seed_img( await _head(db, tag.id, slot=0)
db,
"e" * 64,
{"alreadyhere": {"category": "character", "confidence": 0.96}},
)
await db.execute( await db.execute(
image_tag.insert().values( image_tag.insert().values(
image_record_id=img.id, tag_id=tag.id, source="manual" image_record_id=img.id, tag_id=tag.id, source="manual"
) )
) )
await db.commit()
sl = await SuggestionService(db).for_image(img.id) 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 @pytest.mark.asyncio
async def test_rejected_tag_surfaced_flagged_then_reversible(db): async def test_rejected_tag_surfaced_flagged_then_reversible(db):
# A dismissed suggestion is NOT dropped: it stays in the list flagged # A dismissed suggestion is NOT dropped: it stays flagged rejected so the
# rejected=True so the rail can show it + offer one-click un-reject # rail can show it + offer one-click un-reject (operator-asked 2026-06-27).
# (visible, reversible rejection — operator-asked 2026-06-27). A live tag = await TagService(db).find_or_create("goblin", TagKind.general)
# suggestion sorts ahead of the rejected one. img = await _img(db, "f" * 64, _emb(0))
tags = TagService(db) await _head(db, tag.id, slot=0)
rejected_tag = await tags.find_or_create("rejectme", TagKind.general) await db.commit()
img = await _seed_img( await AllowlistService(db).dismiss(img.id, tag.id)
db, await db.commit()
"f" * 64,
{
"rejectme": {"category": "general", "confidence": 0.96},
"keepme": {"category": "general", "confidence": 0.90},
},
)
await AllowlistService(db).dismiss(img.id, rejected_tag.id)
sl = await SuggestionService(db).for_image(img.id) sl = await SuggestionService(db).for_image(img.id)
general = sl.by_category["general"] s = next(x for x in sl.by_category["general"] if x.canonical_tag_id == tag.id)
# Match by id, not display casing (an existing tag keeps its stored name). assert s.rejected is True
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
# Un-reject reverts it to a live suggestion. await AllowlistService(db).undismiss(img.id, tag.id)
await AllowlistService(db).undismiss(img.id, rejected_tag.id) await db.commit()
sl2 = await SuggestionService(db).for_image(img.id) sl2 = await SuggestionService(db).for_image(img.id)
rej2 = next( s2 = next(x for x in sl2.by_category["general"] if x.canonical_tag_id == tag.id)
s for s in sl2.by_category["general"] assert s2.rejected is False
if s.canonical_tag_id == rejected_tag.id
)
assert rej2.rejected is False
+51 -60
View File
@@ -1,88 +1,84 @@
"""Consensus (for_selection) over the tagging-v2 HEAD suggestion source."""
import pytest import pytest
from sqlalchemy import select
from backend.app import create_app 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.models.tag import image_tag
from backend.app.services.ml.suggestions import SuggestionService from backend.app.services.ml.suggestions import SuggestionService
from backend.app.services.tag_service import TagService from backend.app.services.tag_service import TagService
from tests._prediction_helpers import seed_predictions
pytestmark = pytest.mark.integration pytestmark = pytest.mark.integration
def _img(sha: str) -> ImageRecord: def _emb(slot: int) -> list[float]:
return ImageRecord( v = [0.0] * 1152
path=f"/images/{sha}.jpg", sha256=sha, size_bytes=1, v[slot] = 3.0
mime="image/jpeg", width=1, height=1, return v
origin="imported_filesystem", integrity_status="unknown",
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 @pytest.mark.asyncio
async def test_consensus_includes_tag_over_threshold(db): async def test_consensus_includes_tag_over_threshold(db):
tags = TagService(db) t = await TagService(db).find_or_create("sword", TagKind.general)
t = await tags.find_or_create("sword", TagKind.general) a = await _img(db, "a" * 64, _emb(0))
a = _img("a" * 64) b = await _img(db, "b" * 64, _emb(0))
b = _img("b" * 64) await _head(db, t.id, slot=0)
db.add_all([a, b]) await db.commit()
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}})
res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8) res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8)
gen = res["general"] s = next(s for s in res["general"] if s["canonical_tag_id"] == t.id)
assert any(s["canonical_tag_id"] == t.id for s in gen) assert s["coverage"] == 1.0 # suggested on both
s = next(s for s in gen if s["canonical_tag_id"] == t.id) assert s["confidence"] > 0.5
assert s["coverage"] == 1.0
assert 0.95 <= s["confidence"] <= 0.97
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_consensus_counts_already_applied_for_coverage(db): async def test_consensus_counts_already_applied_for_coverage(db):
tags = TagService(db) t = await TagService(db).find_or_create("sky", TagKind.general)
t = await tags.find_or_create("sky", TagKind.general) a = await _img(db, "c" * 64, _emb(0)) # head suggests it
a = _img("c" * 64) b = await _img(db, "d" * 64, None) # no embedding; tag applied instead
b = _img("d" * 64) # no prediction await _head(db, t.id, slot=0)
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
await db.execute( await db.execute(
image_tag.insert().values( image_tag.insert().values(
image_record_id=b.id, tag_id=t.id, source="manual" 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) 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) 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["coverage"] == 1.0 # 1 suggested + 1 applied / 2
assert s["confidence"] == pytest.approx(0.96, abs=1e-4)
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_consensus_excludes_below_threshold(db): async def test_consensus_excludes_below_threshold(db):
tags = TagService(db) t = await TagService(db).find_or_create("rare", TagKind.general)
await tags.find_or_create("rare", TagKind.general) a = await _img(db, "e" * 64, _emb(0)) # suggested here
a = _img("e" * 64) b = await _img(db, "f" * 64, None) # not here → coverage 0.5 < 0.8
b = _img("f" * 64) await _head(db, t.id, slot=0)
db.add_all([a, b]) await db.commit()
await db.flush()
await seed_predictions(db, a.id, {"rare": {"category": "general", "confidence": 0.96}})
res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8) res = await SuggestionService(db).for_selection([a.id, b.id], threshold=0.8)
assert all( assert all(s["name"] != "rare" for s in res.get("general", []))
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", []))
@pytest.mark.asyncio @pytest.mark.asyncio
@@ -93,13 +89,9 @@ async def test_consensus_threshold_clamped_and_empty_for_no_ids(db):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_bulk_suggestions_route(db): async def test_bulk_suggestions_route(db):
t = await TagService(db).find_or_create("sword", TagKind.general)
tags = TagService(db) a = await _img(db, "i" * 64, _emb(0))
await tags.find_or_create("sword", TagKind.general) await _head(db, t.id, slot=0)
a = _img("i" * 64)
db.add(a)
await db.commit()
await seed_predictions(db, a.id, {"sword": {"category": "general", "confidence": 0.97}})
await db.commit() await db.commit()
app = create_app() app = create_app()
async with app.test_client() as c: async with app.test_client() as c:
@@ -115,7 +107,6 @@ async def test_bulk_suggestions_route(db):
@pytest.mark.asyncio @pytest.mark.asyncio
async def test_bulk_suggestions_requires_ids(db): async def test_bulk_suggestions_requires_ids(db):
app = create_app() app = create_app()
async with app.test_client() as c: async with app.test_client() as c:
resp = await c.post("/api/suggestions/bulk", json={}) resp = await c.post("/api/suggestions/bulk", json={})