Added ability to store CNN models

model_service/model/aimodel.py

@@ -3,16 +3,14 @@ from .db import db
 from sqlalchemy.sql import func
 from sqlalchemy.dialects.postgresql import UUID
 import uuid
+import enum
 
+class AIModelType(enum.Enum):
+    SVM = 1
+    CNN = 2
 
 class AIModel(db.Model):
     id = db.Column(UUID(as_uuid=True), primary_key=True, default=uuid.uuid4, unique=True, nullable=False)
     timestamp = db.Column(db.TIMESTAMP, nullable=False, server_default=func.now())
 
-    # details
-    mid_window = db.Column(db.Float)
-    mid_step = db.Column(db.Float)
-    short_window = db.Column(db.Float)
-    short_step = db.Column(db.Float)
-    compute_beat = db.Column(db.Boolean)
-    type = db.Column(db.String(15))
+    type = db.Column(db.Enum(AIModelType), nullable=False)