75 lines
2.3 KiB
Python
75 lines
2.3 KiB
Python
"""ML admin API: settings, backfill trigger, centroid recompute trigger."""
|
|
|
|
from quart import Blueprint, jsonify, request
|
|
|
|
from ..extensions import get_session
|
|
from ..models import MLSettings
|
|
|
|
ml_admin_bp = Blueprint("ml_admin", __name__, url_prefix="/api/ml")
|
|
|
|
|
|
_EDITABLE = (
|
|
"suggestion_threshold_artist",
|
|
"suggestion_threshold_character",
|
|
"suggestion_threshold_copyright",
|
|
"suggestion_threshold_general",
|
|
"centroid_similarity_threshold",
|
|
"min_reference_images",
|
|
)
|
|
|
|
|
|
@ml_admin_bp.route("/settings", methods=["GET"])
|
|
async def get_settings():
|
|
from sqlalchemy import select
|
|
|
|
async with get_session() as session:
|
|
s = (
|
|
await session.execute(select(MLSettings).where(MLSettings.id == 1))
|
|
).scalar_one()
|
|
return jsonify(
|
|
{
|
|
"suggestion_threshold_artist": s.suggestion_threshold_artist,
|
|
"suggestion_threshold_character": s.suggestion_threshold_character,
|
|
"suggestion_threshold_copyright": s.suggestion_threshold_copyright,
|
|
"suggestion_threshold_general": s.suggestion_threshold_general,
|
|
"centroid_similarity_threshold": s.centroid_similarity_threshold,
|
|
"min_reference_images": s.min_reference_images,
|
|
"tagger_model_version": s.tagger_model_version,
|
|
"embedder_model_version": s.embedder_model_version,
|
|
}
|
|
)
|
|
|
|
|
|
@ml_admin_bp.route("/settings", methods=["PATCH"])
|
|
async def patch_settings():
|
|
from sqlalchemy import select
|
|
|
|
body = await request.get_json()
|
|
if not isinstance(body, dict):
|
|
return jsonify({"error": "body must be an object"}), 400
|
|
async with get_session() as session:
|
|
s = (
|
|
await session.execute(select(MLSettings).where(MLSettings.id == 1))
|
|
).scalar_one()
|
|
for field in _EDITABLE:
|
|
if field in body:
|
|
setattr(s, field, body[field])
|
|
await session.commit()
|
|
return await get_settings()
|
|
|
|
|
|
@ml_admin_bp.route("/backfill", methods=["POST"])
|
|
async def trigger_backfill():
|
|
from ..tasks.ml import backfill
|
|
|
|
r = backfill.delay()
|
|
return jsonify({"celery_task_id": r.id}), 202
|
|
|
|
|
|
@ml_admin_bp.route("/recompute-centroids", methods=["POST"])
|
|
async def trigger_recompute():
|
|
from ..tasks.ml import recompute_centroids
|
|
|
|
r = recompute_centroids.delay()
|
|
return jsonify({"celery_task_id": r.id}), 202
|