make template type enum

Agenta-AI · Jun 6, 2024 · 256c249 · 256c249
1 parent 99f032b
commit 256c249
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/agenta-backend/agenta_backend/models/db_models.py b/agenta-backend/agenta_backend/models/db_models.py
@@ -9,11 +9,14 @@
     Boolean,
     ForeignKey,
     Float,
+    Enum,
 )
 from sqlalchemy.orm import relationship, declarative_base
 import uuid_utils.compat as uuid
 from sqlalchemy.dialects.postgresql import UUID, JSONB
 
+from agenta_backend.models.shared_models import TemplateType
+
 Base = declarative_base()
 
 
@@ -261,7 +264,7 @@ class TemplateDB(Base):
         unique=True,
         nullable=False,
     )
-    type = Column(String, default="image")
+    type = Column(Enum(TemplateType), default=TemplateType.IMAGE, nullable=False)
     template_uri = Column(String)
     tag_id = Column(Integer)
     name = Column(String, unique=True)

diff --git a/agenta-backend/agenta_backend/models/shared_models.py b/agenta-backend/agenta_backend/models/shared_models.py
@@ -1,5 +1,6 @@
 from pydantic import BaseModel, Field
 from typing import Any, Dict, List, Optional
+import enum
 
 
 class ConfigDB(BaseModel):
@@ -47,3 +48,8 @@ class HumanEvaluationScenarioInput(BaseModel):
 class HumanEvaluationScenarioOutput(BaseModel):
     variant_id: str
     variant_output: str
+
+
+class TemplateType(enum.Enum):
+    IMAGE = "image"
+    ZIP = "zip"