diff --git a/api/models/task.py b/api/models/task.py
index aa0e36dc95..c17b0c3656 100644
--- a/api/models/task.py
+++ b/api/models/task.py
@@ -1,5 +1,7 @@
+from datetime import datetime
+from typing import Optional
 from celery import states  # type: ignore
-from sqlalchemy.orm import mapped_column
+from sqlalchemy.orm import mapped_column, Mapped
 
 from libs.datetime_utils import naive_utc_now
 from models.base import Base
@@ -36,7 +38,9 @@ class CeleryTaskSet(Base):
 
     __tablename__ = "celery_tasksetmeta"
 
-    id = db.Column(db.Integer, db.Sequence("taskset_id_sequence"), autoincrement=True, primary_key=True)
-    taskset_id = db.Column(db.String(155), unique=True)
-    result = db.Column(db.PickleType, nullable=True)
-    date_done = db.Column(db.DateTime, default=lambda: naive_utc_now(), nullable=True)
+    id: Mapped[int] = mapped_column(
+        db.Integer, db.Sequence("taskset_id_sequence"), autoincrement=True, primary_key=True
+    )
+    taskset_id = mapped_column(db.String(155), unique=True)
+    result = mapped_column(db.PickleType, nullable=True)
+    date_done: Mapped[Optional[datetime]] = mapped_column(db.DateTime, default=lambda: naive_utc_now(), nullable=True)