diff --git a/app/models/types.py b/app/models/types.py
index 7bb978d..b463d52 100644
--- a/app/models/types.py
+++ b/app/models/types.py
@@ -1,21 +1,24 @@
-from datetime import timezone
+from datetime import datetime, timezone
+from typing import Optional
 
-from sqlalchemy import DateTime, TypeDecorator
+from sqlalchemy.engine import Dialect
+from sqlalchemy.types import TypeDecorator, DateTime
 
 
-class AwareDateTime(TypeDecorator):
+class AwareDateTime(TypeDecorator[datetime]):
     impl = DateTime(timezone=True)
-
     cache_ok = True
 
-    def process_bind_param(self, value, dialect):
-        # Ensure the value is aware. If it's naive, assume UTC.
+    def process_bind_param(
+        self, value: Optional[datetime], dialect: Dialect
+    ) -> Optional[datetime]:
         if value is not None and value.tzinfo is None:
             value = value.replace(tzinfo=timezone.utc)
         return value
 
-    def process_result_value(self, value, dialect):
-        # Ensure the value is aware. If it's naive, assume UTC.
+    def process_result_value(
+        self, value: Optional[datetime], dialect: Dialect
+    ) -> Optional[datetime]:
         if value is not None and value.tzinfo is None:
             value = value.replace(tzinfo=timezone.utc)
         return value