Files
FabledCurator/backend/app/services/ml/suggestions.py
T
bvandeusen 5c3f8ebd70
CI / lint (push) Successful in 3s
CI / frontend-build (push) Successful in 22s
CI / backend-lint-and-test (push) Successful in 35s
CI / integration (push) Successful in 3m7s
fix(aliases): store modal alias under raw model key + make aliases visible/manageable
The headline bug: aliases created from the modal NEVER resolved. Create
sent the normalized display name ('Sword', 'Uchiha Sasuke') while
resolution keys on the raw booru model key ('sword', 'uchiha_sasuke',
case-sensitive) — so the mapping was stored under a key nothing looks up,
and the prediction kept reappearing unaliased. The raw key wasn't even in
the /suggestions response, so the modal couldn't send it.

- Suggestion now carries raw_name (the model key an alias must use) and
  via_alias (surfaced via an operator alias); both serialized by the API.
- Modal alias-create sends raw_name, not display_name (the fix). Aliased
  suggestions show an 'alias' badge and a 'Remove alias' action; 'Treat as
  alias for…' is hidden for centroid hits (no model key) and already-aliased
  rows.
- Tag-side management: TagCard ⋮ → 'Aliases…' opens a dialog listing the
  model keys that fold into a tag, with remove (GET /api/tags/<id>/aliases +
  AliasService.list_for_tag). Creation stays in the modal suggestion flow.

Tests: full API round-trip locking the raw-key contract (raw_name exposed →
alias authored with it → resolves + via_alias on a later image);
list_for_tag (service + API); via_alias/raw_name on the existing service
suggestion tests. No migration.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
2026-06-12 13:05:58 -04:00

349 lines
14 KiB
Python

"""The suggestion read-path: raw predictions + centroids -> alias-resolved,
threshold-filtered, category-grouped, ranked suggestions for one image.
"""
from dataclasses import dataclass, field
from sqlalchemy import func, select
from sqlalchemy.ext.asyncio import AsyncSession
from ...models import (
ImagePrediction,
ImageRecord,
MLSettings,
Tag,
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
@dataclass(frozen=True)
class Suggestion:
# canonical_tag_id is None when this is a raw Camie tag with no alias and
# no existing Tag row — accepting it will create the tag.
canonical_tag_id: int | None
display_name: str
category: str
score: float
source: str # 'tagger' | 'centroid' | 'both'
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
# modal needs it to author an alias correctly. None for centroid-only hits
# (no underlying prediction → nothing to alias).
raw_name: str | None = None
# via_alias = this suggestion was surfaced because an operator alias remapped
# the raw prediction to this canonical tag. Lets the UI mark it + offer undo.
via_alias: bool = False
@dataclass
class SuggestionList:
by_category: dict[str, list[Suggestion]] = field(default_factory=dict)
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,
) -> SuggestionList:
"""Ranked suggestions for one image.
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."""
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(
select(image_tag.c.tag_id).where(
image_tag.c.image_record_id == image_id
)
)
).scalars().all()
)
rejected = set(
(
await self.session.execute(
select(TagSuggestionRejection.tag_id).where(
TagSuggestionRejection.image_record_id == image_id
)
)
).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]
)
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,
)
for raw, display, category, conf in candidates:
canonical = alias_map.get((raw, category))
if canonical is not None:
if canonical.id in applied or canonical.id in rejected:
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,
),
)
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
or existing_tag.id in rejected
):
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,
),
)
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 or hit.tag_id in rejected:
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,
),
)
result = SuggestionList()
for sug in merged.values():
result.by_category.setdefault(sug.category, []).append(sug)
for cat in result.by_category:
result.by_category[cat].sort(key=lambda s: s.score, reverse=True)
return result
async def for_selection(
self,
image_ids: list[int],
threshold: float = 0.8,
top_k: int = 10,
) -> dict[str, list[dict]]:
"""Consensus suggestions across image_ids. A tag is included iff it
was suggested for (or already applied to) >= threshold fraction of
the selection AND was acceptable on >= 1 image. Confidence is the
mean over images where it was suggested. Aggregated by
canonical_tag_id; creates-new (no canonical id) suggestions are
skipped (bulk Accept applies by tag id)."""
if not image_ids:
return {}
threshold = min(1.0, max(0.0, threshold))
total = len(image_ids)
stats: dict[int, dict] = {}
for image_id in image_ids:
sl = await self.for_image(image_id)
for category, items in sl.by_category.items():
for s in items:
if s.canonical_tag_id is None or s.creates_new_tag:
continue
st = stats.get(s.canonical_tag_id)
if st is None:
st = {
"tag_id": s.canonical_tag_id,
"name": s.display_name,
"category": category,
"source": s.source,
"suggested_count": 0,
"sum_score": 0.0,
}
stats[s.canonical_tag_id] = st
st["suggested_count"] += 1
st["sum_score"] += s.score
rows = (
await self.session.execute(
select(
image_tag.c.image_record_id, image_tag.c.tag_id
).where(image_tag.c.image_record_id.in_(image_ids))
)
).all()
applied_by_tag: dict[int, set[int]] = {}
for iid, tid in rows:
applied_by_tag.setdefault(tid, set()).add(iid)
result: dict[str, list[dict]] = {}
for st in stats.values():
existing_count = len(applied_by_tag.get(st["tag_id"], set()))
covered = st["suggested_count"] + existing_count
coverage = covered / total
if coverage < threshold or st["suggested_count"] < 1:
continue
result.setdefault(st["category"], []).append(
{
"canonical_tag_id": st["tag_id"],
"name": st["name"],
"category": st["category"],
"confidence": round(
st["sum_score"] / st["suggested_count"], 4
),
"coverage": round(coverage, 4),
"covered_count": covered,
"source": st["source"],
}
)
for cat in result:
result[cat].sort(key=lambda x: x["confidence"], reverse=True)
result[cat] = result[cat][:top_k]
return result