From 04fdd254f965c35a24c2c32ee2fa00a040f62b5a Mon Sep 17 00:00:00 2001
From: Nicolas Mowen <nickmowen213@gmail.com>
Date: Thu, 5 Jun 2025 07:59:22 -0600
Subject: [PATCH] Improve toasts and model state

---
 .../locales/en/views/classificationModel.json |  7 +++-
 .../classification/ModelTrainingView.tsx      | 39 ++++++++++++++++++-
 2 files changed, 42 insertions(+), 4 deletions(-)

diff --git a/web/public/locales/en/views/classificationModel.json b/web/public/locales/en/views/classificationModel.json
index eb09ecaa0..0af0179b9 100644
--- a/web/public/locales/en/views/classificationModel.json
+++ b/web/public/locales/en/views/classificationModel.json
@@ -9,12 +9,15 @@
     "success": {
       "deletedCategory": "Deleted Class",
       "deletedImage": "Deleted Images",
-      "categorizedImage": "Successfully Classified Image"
+      "categorizedImage": "Successfully Classified Image",
+      "trainedModel": "Successfully trained model.",
+      "trainingModel": "Successfully started model training."
     },
     "error": {
       "deleteImageFailed": "Failed to delete: {{errorMessage}}",
       "deleteCategoryFailed": "Failed to delete class: {{errorMessage}}",
-      "categorizeFailed": "Failed to categorize image: {{errorMessage}}"
+      "categorizeFailed": "Failed to categorize image: {{errorMessage}}",
+      "trainingFailed": "Failed to start model training: {{errorMessage}}"
     }
   },
   "deleteCategory": {
diff --git a/web/src/views/classification/ModelTrainingView.tsx b/web/src/views/classification/ModelTrainingView.tsx
index 197768a95..1f62a4f53 100644
--- a/web/src/views/classification/ModelTrainingView.tsx
+++ b/web/src/views/classification/ModelTrainingView.tsx
@@ -59,6 +59,7 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) {
 
   // model state
 
+  const [wasTraining, setWasTraining] = useState(false);
   const { payload: lastModelState } = useModelState(model.name, true);
   const modelState = useMemo<ModelState>(() => {
     if (!lastModelState || lastModelState == "downloaded") {
@@ -68,6 +69,21 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) {
     return lastModelState;
   }, [lastModelState]);
 
+  useEffect(() => {
+    if (!wasTraining) {
+      return;
+    }
+
+    if (modelState == "complete") {
+      toast.success(t("toast.success.trainedModel"), {
+        position: "top-center",
+      });
+      setWasTraining(false);
+    }
+    // only refresh when modelState changes
+    // eslint-disable-next-line react-hooks/exhaustive-deps
+  }, [modelState]);
+
   // dataset
 
   const { data: trainImages, mutate: refreshTrain } = useSWR<string[]>(
@@ -115,8 +131,27 @@ export default function ModelTrainingView({ model }: ModelTrainingViewProps) {
   // actions
 
   const trainModel = useCallback(() => {
-    axios.post(`classification/${model.name}/train`);
-  }, [model]);
+    axios
+      .post(`classification/${model.name}/train`)
+      .then((resp) => {
+        if (resp.status == 200) {
+          setWasTraining(true);
+          toast.success(t("toast.success.trainingModel"), {
+            position: "top-center",
+          });
+        }
+      })
+      .catch((error) => {
+        const errorMessage =
+          error.response?.data?.message ||
+          error.response?.data?.detail ||
+          "Unknown error";
+
+        toast.error(t("toast.error.trainingFailed", { errorMessage }), {
+          position: "top-center",
+        });
+      });
+  }, [model, t]);
 
   const [deleteDialogOpen, setDeleteDialogOpen] = useState<string[] | null>(
     null,