From 7860b86a13c19a8cd2ed483a8ea930b7e70d7485 Mon Sep 17 00:00:00 2001 From: Bryan Van Deusen Date: Fri, 15 May 2026 07:38:33 -0400 Subject: [PATCH] =?UTF-8?q?feat(fc2b):=20add=20SuggestionService=20?= =?UTF-8?q?=E2=80=94=20alias-resolved,=20threshold-filtered,=20ranked?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The read path: load tagger_predictions, drop unsurfaced categories (rating/meta/year), apply per-category thresholds, batch-resolve aliases, skip applied + rejected, augment with centroid hits above the similarity threshold, merge duplicate signals (take max score, mark source 'both'), group by category, sort by score DESC. Tests marked integration. Co-Authored-By: Claude Opus 4.7 (1M context) --- backend/app/services/ml/suggestions.py | 199 +++++++++++++++++++++++++ tests/test_ml_suggestions.py | 105 +++++++++++++ 2 files changed, 304 insertions(+) create mode 100644 backend/app/services/ml/suggestions.py create mode 100644 tests/test_ml_suggestions.py diff --git a/backend/app/services/ml/suggestions.py b/backend/app/services/ml/suggestions.py new file mode 100644 index 0000000..b725d9c --- /dev/null +++ b/backend/app/services/ml/suggestions.py @@ -0,0 +1,199 @@ +"""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 select +from sqlalchemy.ext.asyncio import AsyncSession + +from ...models import ( + ImageRecord, + MLSettings, + Tag, + TagSuggestionRejection, +) +from ...models.tag import image_tag +from .aliases import AliasService +from .centroids import CentroidService +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 + + +@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() + + def _threshold_for(self, s: MLSettings, category: str) -> float: + return { + "artist": s.suggestion_threshold_artist, + "character": s.suggestion_threshold_character, + "copyright": s.suggestion_threshold_copyright, + "general": s.suggestion_threshold_general, + }.get(category, 1.01) # 1.01 => never surfaces (unsurfaced category) + + async def for_image(self, image_id: int) -> SuggestionList: + img = await self.session.get(ImageRecord, image_id) + if img is None: + return SuggestionList() + + settings = await self._settings() + predictions: dict = img.tagger_predictions or {} + + 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: list[tuple[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): + continue + candidates.append((name, category, conf)) + + alias_map = await self.aliases.resolve_many( + [(n, c) for n, c, _ 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, + ) + + for name, category, conf in candidates: + canonical = alias_map.get((name, 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, + ), + ) + else: + existing_tag = ( + await self.session.execute( + select(Tag).where(Tag.name == name) + ) + ).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, + ), + ) + else: + _merge( + f"raw:{name}:{category}", + Suggestion( + canonical_tag_id=None, + display_name=name, + category=category, + score=conf, + source="tagger", + creates_new_tag=True, + ), + ) + + # --- 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 diff --git a/tests/test_ml_suggestions.py b/tests/test_ml_suggestions.py new file mode 100644 index 0000000..0154eed --- /dev/null +++ b/tests/test_ml_suggestions.py @@ -0,0 +1,105 @@ +import pytest + +from backend.app.models import ImageRecord, TagKind +from backend.app.models.tag import image_tag +from backend.app.services.ml.aliases import AliasService +from backend.app.services.ml.suggestions import SuggestionService +from backend.app.services.tag_service import TagService + +pytestmark = pytest.mark.integration + + +def _img(sha: str, predictions: dict) -> 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", + tagger_predictions=predictions, + ) + + +@pytest.mark.asyncio +async def test_threshold_filters_low_confidence_general(db): + img = _img( + "a" * 64, + { + "smile": {"category": "general", "confidence": 0.80}, + "sword": {"category": "general", "confidence": 0.97}, + }, + ) + db.add(img) + await db.flush() + sl = await SuggestionService(db).for_image(img.id) + names = [s.display_name for s in sl.by_category.get("general", [])] + assert "sword" in names + assert "smile" not in names + + +@pytest.mark.asyncio +async def test_unsurfaced_category_dropped(db): + img = _img( + "b" * 64, + {"safe": {"category": "rating", "confidence": 0.99}}, + ) + db.add(img) + await db.flush() + 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 = _img( + "c" * 64, + {"uchiha_sasuke": {"category": "character", "confidence": 0.96}}, + ) + db.add(img) + await db.flush() + 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 + + +@pytest.mark.asyncio +async def test_raw_tag_creates_new(db): + img = _img( + "d" * 64, + {"brand_new_tag": {"category": "character", "confidence": 0.96}}, + ) + db.add(img) + await db.flush() + sl = await SuggestionService(db).for_image(img.id) + chars = sl.by_category["character"] + assert chars[0].display_name == "brand_new_tag" + assert chars[0].creates_new_tag is True + assert chars[0].canonical_tag_id is None + + +@pytest.mark.asyncio +async def test_applied_tag_not_suggested(db): + tags = TagService(db) + tag = await tags.find_or_create("alreadyhere", TagKind.character) + img = _img( + "e" * 64, + {"alreadyhere": {"category": "character", "confidence": 0.96}}, + ) + db.add(img) + await db.flush() + await db.execute( + image_tag.insert().values( + image_record_id=img.id, tag_id=tag.id, source="manual" + ) + ) + sl = await SuggestionService(db).for_image(img.id) + assert "character" not in sl.by_category or not sl.by_category["character"]