Merge pull request #377 from hotosm/hotfix/notification-training-id

 HOTFix : Training Instance in Notification

Author: kshitijrajsharma

backend/core/models.py

@@ -187,6 +187,7 @@ class UserNotification(models.Model):
     created_at = models.DateTimeField(default=timezone.now)
     read_at = models.DateTimeField(null=True, blank=True)
     message = models.TextField(max_length=500)
+    training = models.ForeignKey(Training, to_field="id", on_delete=models.DO_NOTHING)
 
     def mark_as_read(self):
         if not self.is_read:

backend/core/serializers.py

@@ -496,12 +496,27 @@ def get_unread_notifications_count(self, obj):
 
 
 class UserNotificationSerializer(serializers.ModelSerializer):
+    training_model = serializers.SerializerMethodField()
+
     class Meta:
         model = UserNotification
-        fields = ("id", "is_read", "created_at", "read_at", "message")
+        fields = (
+            "id",
+            "is_read",
+            "created_at",
+            "read_at",
+            "message",
+            "training_model",
+        )
         read_only_fields = (
             "id",
             "created_at",
             "read_at",
             "message",
-        )
+            "training_model",
+        )
+
+    def get_training_model(self, obj):
+        if obj.training and obj.training.model:
+            return obj.training.model.id
+        return None

backend/core/tasks.py

@@ -30,11 +30,10 @@
 from django.conf import settings
 from django.contrib.gis.db.models.aggregates import Extent
 from django.contrib.gis.geos import GEOSGeometry
-from django.core.mail import send_mail
 from django.shortcuts import get_object_or_404
 from django.utils import timezone
 
-from .utils import S3Uploader
+from .utils import S3Uploader, send_notification
 
 logging.basicConfig(
     format="%(asctime)s - %(levelname)s - %(message)s", level=logging.INFO
@@ -516,8 +515,8 @@ def train_model(
 
     training_instance = get_object_or_404(Training, id=training_id)
     model_instance = get_object_or_404(Model, id=training_instance.model.id)
-    
-    send_notification(training_instance,"Started")
+
+    send_notification(training_instance, "Started")
 
     training_instance.status = "RUNNING"
     training_instance.started_at = timezone.now()
@@ -579,46 +578,3 @@ def train_model(
         training_instance.save()
         send_notification(training_instance, "Failed")
         raise ex
-
-def get_email_message(training_instance,status):
-    
-    hostname = settings.FRONTEND_URL  
-    training_model_url = f"{hostname}/ai-models/{training_instance.model.id}"
-
-    message_template = (
-        "Hi {username},\n\n"
-        "Your training task (ID: {training_id}) of model {model_name} has {status}. You can view the details here:\n"
-        "{training_model_url}\n\n"
-        "Thank you for using fAIr - AI Assisted Mapping Tool.\n\n"
-        "Best regards,\n"
-        "The fAIr Dev Team\n\n"
-        "Get Involved : https://www.hotosm.org/get-involved/\n"
-        "https://github.com/hotosm/fAIr/"
-    )
-
-    message = message_template.format(
-        username=training_instance.user.username,
-        training_id=training_instance.id,
-        model_name=training_instance.model.name,
-        status=status.lower(),
-        training_model_url=training_model_url,
-        hostname=hostname,
-
-    )
-    subject = f"fAIr : Training {training_instance.id} {status.capitalize()}"
-    return message, subject
-
-
-def send_notification(training_instance,status):
-    if any(method in training_instance.user.notifications_delivery_methods for method in ["web", "email"]):
-        UserNotification.objects.create(user=training_instance.user, message=f"Training {training_instance.id} has {status}.")
-    if "email" in training_instance.user.notifications_delivery_methods:
-        if training_instance.user.email and training_instance.user.email != '':
-            message,subject=get_email_message(training_instance,status)
-            send_mail(
-                subject=subject,
-                message=message,
-                from_email=settings.DEFAULT_FROM_EMAIL,
-                recipient_list=[training_instance.user.email],
-                fail_silently=False,
-            )

backend/core/utils.py

@@ -15,11 +15,12 @@
 import requests
 from botocore.exceptions import ClientError, NoCredentialsError
 from django.conf import settings
+from django.core.mail import send_mail
 from django.http import HttpResponseRedirect
 from gpxpy.gpx import GPX, GPXTrack, GPXTrackSegment, GPXWaypoint
 from tqdm import tqdm
 
-from .models import AOI, FeedbackAOI, FeedbackLabel, Label
+from .models import AOI, FeedbackAOI, FeedbackLabel, Label, UserNotification
 from .serializers import FeedbackLabelSerializer, LabelSerializer
 
 
@@ -450,3 +451,53 @@ def _upload_directory(self, directory_path, bucket_name):
             "total_files_uploaded": total_files,
             "s3_path": f"s3://{bucket_name}/{self.parent}/",
         }
+
+
+def get_email_message(training_instance, status):
+
+    hostname = settings.FRONTEND_URL
+    training_model_url = f"{hostname}/ai-models/{training_instance.model.id}"
+
+    message_template = (
+        "Hi {username},\n\n"
+        "Your training task (ID: {training_id}) of model {model_name} has {status}. You can view the details here:\n"
+        "{training_model_url}\n\n"
+        "Thank you for using fAIr - AI Assisted Mapping Tool.\n\n"
+        "Best regards,\n"
+        "The fAIr Dev Team\n\n"
+        "Get Involved : https://www.hotosm.org/get-involved/\n"
+        "https://github.com/hotosm/fAIr/"
+    )
+
+    message = message_template.format(
+        username=training_instance.user.username,
+        training_id=training_instance.id,
+        model_name=training_instance.model.name,
+        status=status.lower(),
+        training_model_url=training_model_url,
+        hostname=hostname,
+    )
+    subject = f"fAIr : Training {training_instance.id} {status.capitalize()}"
+    return message, subject
+
+
+def send_notification(training_instance, status):
+    if any(
+        method in training_instance.user.notifications_delivery_methods
+        for method in ["web", "email"]
+    ):
+        UserNotification.objects.create(
+            user=training_instance.user,
+            message=f"Training {training_instance.id} has {status}.",
+            training=training_instance,
+        )
+    if "email" in training_instance.user.notifications_delivery_methods:
+        if training_instance.user.email and training_instance.user.email != "":
+            message, subject = get_email_message(training_instance, status)
+            send_mail(
+                subject=subject,
+                message=message,
+                from_email=settings.DEFAULT_FROM_EMAIL,
+                recipient_list=[training_instance.user.email],
+                fail_silently=False,
+            )

backend/core/views.py

@@ -78,6 +78,7 @@
     process_rawdata,
     request_rawdata,
     s3_object_exists,
+    send_notification,
 )
 
 if settings.ENABLE_PREDICTION_API:
@@ -983,13 +984,14 @@ class UserNotificationViewSet(ReadOnlyModelViewSet):
     permission_classes = [IsOsmAuthenticated]
     serializer_class = UserNotificationSerializer
     filter_backends = [DjangoFilterBackend, filters.OrderingFilter]
-    filterset_fields = ['is_read']  
-    ordering = ['-created_at'] 
-    ordering_fields = ['created_at', 'read_at','is_read'] 
+    filterset_fields = ["is_read"]
+    ordering = ["-created_at"]
+    ordering_fields = ["created_at", "read_at", "is_read"]
 
     def get_queryset(self):
         return UserNotification.objects.filter(user=self.request.user)
 
+
 class MarkNotificationAsRead(APIView):
     authentication_classes = [OsmAuthentication]
     permission_classes = [IsOsmAuthenticated]
@@ -999,19 +1001,29 @@ class MarkNotificationAsRead(APIView):
     )
     def post(self, request, notification_id, format=None):
         try:
-            notification = UserNotification.objects.get(id=notification_id, user=request.user)
+            notification = UserNotification.objects.get(
+                id=notification_id, user=request.user
+            )
 
             if notification.is_read:
-                return Response({"detail": "Notification is already marked as read."}, status=status.HTTP_200_OK)
+                return Response(
+                    {"detail": "Notification is already marked as read."},
+                    status=status.HTTP_200_OK,
+                )
 
             notification.is_read = True
             notification.read_at = timezone.now()
             notification.save()
 
-            return Response({"detail": "Notification marked as read."}, status=status.HTTP_200_OK)
+            return Response(
+                {"detail": "Notification marked as read."}, status=status.HTTP_200_OK
+            )
 
         except UserNotification.DoesNotExist:
-            return Response({"detail": "Notification not found."}, status=status.HTTP_404_NOT_FOUND)
+            return Response(
+                {"detail": "Notification not found."}, status=status.HTTP_404_NOT_FOUND
+            )
+
 
 class MarkAllNotificationsAsRead(APIView):
     authentication_classes = [OsmAuthentication]
@@ -1023,20 +1035,29 @@ class MarkAllNotificationsAsRead(APIView):
             200: openapi.Response(
                 description="All unread notifications marked as read.",
                 examples={
-                    "application/json": {"detail": "All unread notifications marked as read."}
+                    "application/json": {
+                        "detail": "All unread notifications marked as read."
+                    }
                 },
             ),
         },
     )
     def post(self, request, format=None):
-        unread_notifications = UserNotification.objects.filter(user=request.user, is_read=False)
+        unread_notifications = UserNotification.objects.filter(
+            user=request.user, is_read=False
+        )
 
         if not unread_notifications.exists():
-            return Response({"detail": "No unread notifications found."}, status=status.HTTP_404_NOT_FOUND)
+            return Response(
+                {"detail": "No unread notifications found."},
+                status=status.HTTP_404_NOT_FOUND,
+            )
 
         unread_notifications.update(is_read=True, read_at=timezone.now())
-        return Response({"detail": "All unread notifications marked as read."}, status=status.HTTP_200_OK)
-    
+        return Response(
+            {"detail": "All unread notifications marked as read."},
+            status=status.HTTP_200_OK,
+        )
 
 
 class TerminateTrainingView(APIView):
@@ -1049,17 +1070,35 @@ def post(self, request, training_id, format=None):
 
             task_id = training_instance.task_id
             if not task_id:
-                return Response({"detail": "No task associated with this training."}, status=status.HTTP_400_BAD_REQUEST)
+                return Response(
+                    {"detail": "No task associated with this training."},
+                    status=status.HTTP_400_BAD_REQUEST,
+                )
 
-            task = AsyncResult(task_id,app=current_app)
-            if task.state in ["PENDING", "STARTED", "RETRY"]:
+            task = AsyncResult(task_id, app=current_app)
+            if (
+                task.state in ["PENDING", "STARTED", "RETRY", "FAILURE"]
+                and training_instance.status != "FAILED"
+            ):
                 current_app.control.revoke(task_id, terminate=True)
                 training_instance.status = "FAILED"
                 training_instance.finished_at = now()
                 training_instance.save()
-                return Response({"detail": "Training task cancelled successfully."}, status=status.HTTP_200_OK)
+                send_notification(training_instance, "Cancelled")
+                return Response(
+                    {"detail": "Training task cancelled successfully."},
+                    status=status.HTTP_200_OK,
+                )
             else:
-                return Response({"detail": f"Task cannot be cancelled. Current state: {task.state}"}, status=status.HTTP_400_BAD_REQUEST)
+                return Response(
+                    {
+                        "detail": f"Task cannot be cancelled. Current state: {task.state}"
+                    },
+                    status=status.HTTP_400_BAD_REQUEST,
+                )
 
         except Training.DoesNotExist:
-            return Response({"detail": "Training not found or do not belong to you"}, status=status.HTTP_404_NOT_FOUND)
+            return Response(
+                {"detail": "Training not found or do not belong to you"},
+                status=status.HTTP_404_NOT_FOUND,
+            )