diff --git a/backend/app/api/allowlist.py b/backend/app/api/allowlist.py index 53f38bb..31241de 100644 --- a/backend/app/api/allowlist.py +++ b/backend/app/api/allowlist.py @@ -20,12 +20,37 @@ async def list_allowlist(): "tag_name": r.tag_name, "tag_kind": r.tag_kind, "min_confidence": r.min_confidence, + "applied_count": r.applied_count, + "coverage_count": r.coverage_count, } for r in rows ] ) +@allowlist_bp.route("/tags//allowlist/coverage", methods=["GET"]) +async def coverage(tag_id: int): + """Live "at threshold T, a sweep would cover ~N images" projection for the + allowlist tuning dashboard. Defaults to the tag's stored threshold.""" + raw = request.args.get("threshold") + async with get_session() as session: + svc = AllowlistService(session) + if raw is not None: + try: + threshold = float(raw) + except ValueError: + return jsonify({"error": "threshold must be a float"}), 400 + if not (0 < threshold <= 1): + return jsonify({"error": "threshold must be in (0, 1]"}), 400 + else: + row = await session.get(TagAllowlist, tag_id) + if row is None: + return jsonify({"error": "not on allowlist"}), 404 + threshold = row.min_confidence + count = await svc.coverage(tag_id, threshold) + return jsonify({"count": count, "threshold": threshold}) + + @allowlist_bp.route("/tags//allowlist", methods=["GET"]) async def get_one(tag_id: int): async with get_session() as session: diff --git a/backend/app/api/suggestions.py b/backend/app/api/suggestions.py index a8a9cf4..46fa2e5 100644 --- a/backend/app/api/suggestions.py +++ b/backend/app/api/suggestions.py @@ -3,12 +3,31 @@ from quart import Blueprint, jsonify, request from ..extensions import get_session +from ..models import Tag, TagAllowlist from ..services.ml.allowlist import AllowlistService from ..services.ml.suggestions import SuggestionService suggestions_bp = Blueprint("suggestions", __name__, url_prefix="/api") +async def _accept_payload(session, svc, newly_added: bool, tag_id: int) -> dict: + """Shape the accept/alias response. When accepting newly allowlists a tag, + include the coverage PROJECTION (at the tag's threshold) so the UI can show + a non-blocking "auto-applying to ~N images" toast — the actual apply runs + async via apply_allowlist_tags, so this is an estimate, not a post-hoc + count (#7).""" + payload = {"allowlisted": newly_added} + if newly_added: + tag = await session.get(Tag, tag_id) + row = await session.get(TagAllowlist, tag_id) + payload["tag_id"] = tag_id + payload["tag_name"] = tag.name if tag is not None else None + payload["projected_count"] = await svc.coverage( + tag_id, row.min_confidence if row is not None else 0.90, + ) + return payload + + @suggestions_bp.route("/images//suggestions", methods=["GET"]) async def get_suggestions(image_id: int): # ?min= overrides the configured per-category thresholds so the typed @@ -60,13 +79,15 @@ async def accept_suggestion(image_id: int): return jsonify({"error": "tag_id required"}), 400 tag_id = body["tag_id"] async with get_session() as session: - newly_added = await AllowlistService(session).accept(image_id, tag_id) + svc = AllowlistService(session) + newly_added = await svc.accept(image_id, tag_id) + payload = await _accept_payload(session, svc, newly_added, tag_id) await session.commit() if newly_added: from ..tasks.ml import apply_allowlist_tags apply_allowlist_tags.delay(tag_id=tag_id) - return "", 204 + return jsonify(payload) @suggestions_bp.route( @@ -77,19 +98,24 @@ async def alias_suggestion(image_id: int): required = {"alias_string", "alias_category", "canonical_tag_id"} if not body or not required.issubset(body): return jsonify({"error": f"required: {sorted(required)}"}), 400 + canonical_tag_id = body["canonical_tag_id"] async with get_session() as session: - newly_added = await AllowlistService(session).add_alias_and_accept( + svc = AllowlistService(session) + newly_added = await svc.add_alias_and_accept( image_id, body["alias_string"], body["alias_category"], - body["canonical_tag_id"], + canonical_tag_id, + ) + payload = await _accept_payload( + session, svc, newly_added, canonical_tag_id, ) await session.commit() if newly_added: from ..tasks.ml import apply_allowlist_tags - apply_allowlist_tags.delay(tag_id=body["canonical_tag_id"]) - return "", 204 + apply_allowlist_tags.delay(tag_id=canonical_tag_id) + return jsonify(payload) @suggestions_bp.route( diff --git a/backend/app/services/ml/allowlist.py b/backend/app/services/ml/allowlist.py index 08a6ac7..8f7d14e 100644 --- a/backend/app/services/ml/allowlist.py +++ b/backend/app/services/ml/allowlist.py @@ -5,11 +5,18 @@ image_tag AND to tag_allowlist; per-image removal/dismiss writes a rejection. from collections.abc import Sequence from dataclasses import dataclass -from sqlalchemy import delete, select +from sqlalchemy import and_, delete, distinct, func, or_, select from sqlalchemy.dialects.postgresql import insert from sqlalchemy.ext.asyncio import AsyncSession -from ...models import MLSettings, Tag, TagAllowlist, TagSuggestionRejection +from ...models import ( + ImagePrediction, + MLSettings, + Tag, + TagAlias, + TagAllowlist, + TagSuggestionRejection, +) from ...models.tag import image_tag from .aliases import AliasService @@ -20,6 +27,8 @@ class AllowlistRow: tag_name: str tag_kind: str min_confidence: float + applied_count: int # image_tag rows currently carrying this tag + coverage_count: int # images a sweep WOULD cover at min_confidence class AllowlistService: @@ -116,6 +125,44 @@ class AllowlistService: delete(TagAllowlist).where(TagAllowlist.tag_id == tag_id) ) + async def _coverage_match(self, tag: Tag): + """The predicate over image_prediction rows that resolve to `tag`, + mirroring tasks.ml._confidence_for_tag's resolution: a prediction whose + raw_name equals the tag name (any category), OR an alias maps + (raw_name, category) -> this tag. Returns a SQLAlchemy boolean clause. + """ + alias_rows = ( + await self.session.execute( + select(TagAlias.alias_string, TagAlias.alias_category).where( + TagAlias.canonical_tag_id == tag.id + ) + ) + ).all() + name_clause = ImagePrediction.raw_name == tag.name + alias_clauses = [ + and_( + ImagePrediction.raw_name == a, + ImagePrediction.category == c, + ) + for a, c in alias_rows + ] + return or_(name_clause, *alias_clauses) if alias_clauses else name_clause + + async def coverage(self, tag_id: int, threshold: float) -> int: + """How many distinct images a sweep WOULD cover for this tag at + `threshold`: images with a resolving prediction scoring >= threshold. + The gross candidate pool (NOT minus already-applied/rejected) — it's + the tuning signal for "lower the threshold and ~N more images qualify". + """ + tag = await self.session.get(Tag, tag_id) + if tag is None: + return 0 + match = await self._coverage_match(tag) + stmt = select( + func.count(distinct(ImagePrediction.image_record_id)) + ).where(ImagePrediction.score >= threshold, match) + return (await self.session.execute(stmt)).scalar_one() + async def list_all(self) -> Sequence[AllowlistRow]: stmt = ( select( @@ -128,12 +175,33 @@ class AllowlistService: .order_by(Tag.name.asc()) ) rows = (await self.session.execute(stmt)).all() - return [ - AllowlistRow( - tag_id=r[0], - tag_name=r[1], - tag_kind=r[2].value if hasattr(r[2], "value") else str(r[2]), - min_confidence=r[3], + tag_ids = [r[0] for r in rows] + + # Applied counts in ONE grouped query (vs N per-row counts). + applied: dict[int, int] = {} + if tag_ids: + applied = dict( + ( + await self.session.execute( + select(image_tag.c.tag_id, func.count()) + .where(image_tag.c.tag_id.in_(tag_ids)) + .group_by(image_tag.c.tag_id) + ) + ).all() ) - for r in rows - ] + + result = [] + for r in rows: + # Coverage is per-tag (alias set differs); allowlist is small. + cov = await self.coverage(r[0], r[3]) + result.append( + AllowlistRow( + tag_id=r[0], + tag_name=r[1], + tag_kind=r[2].value if hasattr(r[2], "value") else str(r[2]), + min_confidence=r[3], + applied_count=applied.get(r[0], 0), + coverage_count=cov, + ) + ) + return result diff --git a/tests/test_api_allowlist.py b/tests/test_api_allowlist.py index 3b4f45b..01539ae 100644 --- a/tests/test_api_allowlist.py +++ b/tests/test_api_allowlist.py @@ -1,6 +1,6 @@ import pytest -from backend.app.models import TagAllowlist, TagKind +from backend.app.models import ImagePrediction, ImageRecord, TagAllowlist, TagKind from backend.app.services.tag_service import TagService pytestmark = pytest.mark.integration @@ -39,3 +39,50 @@ async def test_patch_rejects_out_of_range(client, db): f"/api/tags/{tag.id}/allowlist", json={"min_confidence": 1.5} ) assert resp.status_code == 400 + + +@pytest.mark.asyncio +async def test_coverage_endpoint(client, db): + tag = await TagService(db).find_or_create("Cover", TagKind.general) + db.add(TagAllowlist(tag_id=tag.id, min_confidence=0.90)) + for i, score in enumerate((0.95, 0.60)): + img = ImageRecord( + path=f"/images/cov{i}.jpg", sha256=f"cv{i:062d}", size_bytes=1, + mime="image/jpeg", width=1, height=1, + origin="imported_filesystem", integrity_status="unknown", + ) + db.add(img) + await db.flush() + db.add(ImagePrediction( + image_record_id=img.id, raw_name="Cover", + category="general", score=score, + )) + await db.commit() + + # Explicit threshold. + resp = await client.get( + f"/api/tags/{tag.id}/allowlist/coverage?threshold=0.90" + ) + assert resp.status_code == 200 + assert (await resp.get_json())["count"] == 1 + # Lower what-if threshold widens coverage. + resp = await client.get( + f"/api/tags/{tag.id}/allowlist/coverage?threshold=0.50" + ) + assert (await resp.get_json())["count"] == 2 + # No threshold → uses the stored min_confidence (0.90). + resp = await client.get(f"/api/tags/{tag.id}/allowlist/coverage") + body = await resp.get_json() + assert body["count"] == 1 + assert body["threshold"] == pytest.approx(0.90) + + +@pytest.mark.asyncio +async def test_coverage_rejects_bad_threshold(client, db): + tag = await TagService(db).find_or_create("Cover2", TagKind.general) + db.add(TagAllowlist(tag_id=tag.id)) + await db.commit() + resp = await client.get( + f"/api/tags/{tag.id}/allowlist/coverage?threshold=2.0" + ) + assert resp.status_code == 400 diff --git a/tests/test_api_suggestions.py b/tests/test_api_suggestions.py index 84b95b0..36c0f89 100644 --- a/tests/test_api_suggestions.py +++ b/tests/test_api_suggestions.py @@ -14,11 +14,11 @@ def eager(): celery.conf.task_always_eager = False -async def _img(db, preds): +async def _img(db, preds, sha="s" * 64): from tests._prediction_helpers import seed_predictions img = ImageRecord( - path="/images/s.jpg", sha256="s" * 64, size_bytes=1, + path=f"/images/{sha[:8]}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg", width=1, height=1, origin="imported_filesystem", integrity_status="unknown", ) @@ -57,7 +57,31 @@ async def test_accept_then_applied(client, db): resp = await client.post( f"/api/images/{img.id}/suggestions/accept", json={"tag_id": tag.id} ) - assert resp.status_code == 204 + assert resp.status_code == 200 + body = await resp.get_json() + # #7b: a fresh accept newly-allowlists → projection payload for the toast. + assert body["allowlisted"] is True + assert body["tag_id"] == tag.id + assert body["tag_name"] == "AcceptMe" + assert "projected_count" in body + + +@pytest.mark.asyncio +async def test_accept_already_allowlisted_reports_not_new(client, db): + img1 = await _img(db, {}, sha="c" * 64) + img2 = await _img(db, {}, sha="d" * 64) + tag = await TagService(db).find_or_create("Twice", TagKind.character) + await db.commit() + first = await client.post( + f"/api/images/{img1.id}/suggestions/accept", json={"tag_id": tag.id} + ) + assert (await first.get_json())["allowlisted"] is True + second = await client.post( + f"/api/images/{img2.id}/suggestions/accept", json={"tag_id": tag.id} + ) + body = await second.get_json() + assert body["allowlisted"] is False # already on the allowlist + assert "projected_count" not in body @pytest.mark.asyncio @@ -127,7 +151,8 @@ async def test_alias_roundtrip_resolves_by_raw_key(client, db): "canonical_tag_id": canonical.id, }, ) - assert resp.status_code == 204 + 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 diff --git a/tests/test_ml_allowlist.py b/tests/test_ml_allowlist.py index ad3298f..3d853b1 100644 --- a/tests/test_ml_allowlist.py +++ b/tests/test_ml_allowlist.py @@ -1,7 +1,13 @@ import pytest from sqlalchemy import select -from backend.app.models import TagAllowlist, TagKind, TagSuggestionRejection +from backend.app.models import ( + ImagePrediction, + TagAlias, + TagAllowlist, + TagKind, + TagSuggestionRejection, +) from backend.app.models.tag import image_tag from backend.app.services.ml.allowlist import AllowlistService from backend.app.services.tag_service import TagService @@ -9,10 +15,10 @@ from backend.app.services.tag_service import TagService pytestmark = pytest.mark.integration -async def _make_image(db): +async def _make_image(db, sha: str = "x" * 64): from backend.app.models import ImageRecord img = ImageRecord( - path="/images/x.jpg", sha256="x" * 64, size_bytes=1, + path=f"/images/{sha[:8]}.jpg", sha256=sha, size_bytes=1, mime="image/jpeg", width=1, height=1, origin="imported_filesystem", integrity_status="unknown", ) @@ -21,6 +27,14 @@ async def _make_image(db): return img +async def _add_pred(db, image_id, raw_name, score, category="general"): + db.add(ImagePrediction( + image_record_id=image_id, raw_name=raw_name, + category=category, score=score, + )) + await db.flush() + + @pytest.mark.asyncio async def test_accept_applies_and_allowlists(db): img = await _make_image(db) @@ -106,6 +120,52 @@ async def test_update_threshold_and_remove(db): assert await db.get(TagAllowlist, tag.id) is None +@pytest.mark.asyncio +async def test_coverage_by_threshold_direct_name(db): + tag = await TagService(db).find_or_create("Cov", TagKind.general) + svc = AllowlistService(db) + for i, score in enumerate((0.95, 0.80, 0.60)): + img = await _make_image(db, sha=f"c{i:063d}") + await _add_pred(db, img.id, "Cov", score) + assert await svc.coverage(tag.id, 0.90) == 1 + assert await svc.coverage(tag.id, 0.70) == 2 + assert await svc.coverage(tag.id, 0.50) == 3 + + +@pytest.mark.asyncio +async def test_coverage_via_alias_respects_category(db): + tag = await TagService(db).find_or_create("Aliased", TagKind.character) + db.add(TagAlias( + alias_string="model_key", alias_category="character", + canonical_tag_id=tag.id, + )) + await db.flush() + svc = AllowlistService(db) + hit = await _make_image(db, sha=f"a{0:063d}") + await _add_pred(db, hit.id, "model_key", 0.92, category="character") + # Same alias string but wrong category must NOT resolve to the tag. + miss = await _make_image(db, sha=f"a{1:063d}") + await _add_pred(db, miss.id, "model_key", 0.99, category="general") + assert await svc.coverage(tag.id, 0.90) == 1 + + +@pytest.mark.asyncio +async def test_list_all_reports_applied_and_coverage(db): + tag = await TagService(db).find_or_create("Both", TagKind.general) + svc = AllowlistService(db) + applied_img = await _make_image(db, sha=f"b{0:063d}") + await svc.accept(applied_img.id, tag.id) # applies + allowlists + await _add_pred(db, applied_img.id, "Both", 0.95) + # A second image only has a qualifying prediction (covered, not applied). + cov_img = await _make_image(db, sha=f"b{1:063d}") + await _add_pred(db, cov_img.id, "Both", 0.95) + + rows = await svc.list_all() + row = next(r for r in rows if r.tag_id == tag.id) + assert row.applied_count == 1 # only the accepted image + assert row.coverage_count == 2 # both have a ≥threshold pred + + @pytest.mark.asyncio async def test_update_threshold_clamped_to_store_floor(db): # A min_confidence below the store floor (default 0.70) is clamped up —