feat: templates

Author: chanh-1

server/api/template.py

@@ -0,0 +1,41 @@
+import database
+from typing import Annotated
+from database import Template
+from sqlalchemy.orm import Session
+from fastapi import APIRouter, Depends, Query, Header
+from pydantic import BaseModel
+from commons.utils import get_user_from_jwt, verify_user
+
+template_router = APIRouter(prefix="", tags=["template"])
+
+
+class TemplateModel(BaseModel):
+    name: str = None
+    dag: dict = None
+    description: str = None
+
+
+@template_router.get("/templates", status_code=200)
+def get_all_templates(token: Annotated[str, Header()], db: Session = Depends(database.db_session)):
+    verify_user(get_user_from_jwt(token))
+    templates = db.query(Template).all()
+    response = {"msg": "success", "templates": [template.to_dict() for template in templates]}
+    return response
+
+
+@template_router.get("/template/{id}", status_code=200)
+def get_template_by_id(id: int, token: Annotated[str, Header()], db: Session = Depends(database.db_session)):
+    verify_user(get_user_from_jwt(token))
+    template: Template = db.query(Template).filter(Template.id == id).first()  # type: ignore
+    response = {"msg": "success", "template": template.to_dict()}
+    return response
+
+
+@template_router.post("/template", status_code=200)
+def create_template(inputs: TemplateModel, token: Annotated[str, Header()], db: Session = Depends(database.db_session)):
+    verify_user(get_user_from_jwt(token))
+    template = Template(name=inputs.name, description=inputs.description, dag=inputs.dag)
+    db.add(template)
+    db.commit()
+    response = {"msg": "success", "template": template.to_dict()}
+    return response

server/database.py

@@ -76,3 +76,15 @@ class IntermediateStep(Base):
 
     def __repr__(self):
         return f"Prompt(id={self.id}, name={self.name}, created_by={self.created_by}, dag={self.dag}, meta={self.meta})"
+
+
+class Template(Base):
+    __tablename__ = "template"
+    id = Column(Integer, primary_key=True)
+    name = Column(Text, nullable=False)
+    dag = Column(JSON, nullable=False)
+    description = Column(Text)
+    meta = Column(JSON)
+
+    def __repr__(self):
+        return f"Template(id={self.id}, name={self.name}, description={self.description}, dag={self.dag}, meta={self.meta})"