feat(sqlalchemy): Replace peewee with sqlalchemy

backend/apps/webui/models/models.py

@@ -2,13 +2,11 @@ import json
 import logging
 from typing import Optional
 
-import peewee as pw
-from peewee import *
-
-from playhouse.shortcuts import model_to_dict
 from pydantic import BaseModel, ConfigDict
+from sqlalchemy import String, Column, BigInteger
+from sqlalchemy.orm import Session
 
-from apps.webui.internal.db import DB, JSONField
+from apps.webui.internal.db import Base, JSONField
 
 from typing import List, Union, Optional
 from config import SRC_LOG_LEVELS
@@ -46,41 +44,42 @@ class ModelMeta(BaseModel):
     pass
 
 
-class Model(pw.Model):
-    id = pw.TextField(unique=True)
+class Model(Base):
+    __tablename__ = "model"
+
+    id = Column(String, primary_key=True)
     """
         The model's id as used in the API. If set to an existing model, it will override the model.
     """
-    user_id = pw.TextField()
+    user_id = Column(String)
 
-    base_model_id = pw.TextField(null=True)
+    base_model_id = Column(String, nullable=True)
     """
         An optional pointer to the actual model that should be used when proxying requests.
     """
 
-    name = pw.TextField()
+    name = Column(String)
     """
         The human-readable display name of the model.
     """
 
-    params = JSONField()
+    params = Column(JSONField)
     """
         Holds a JSON encoded blob of parameters, see `ModelParams`.
     """
 
-    meta = JSONField()
+    meta = Column(JSONField)
     """
         Holds a JSON encoded blob of metadata, see `ModelMeta`.
     """
 
-    updated_at = BigIntegerField()
-    created_at = BigIntegerField()
-
-    class Meta:
-        database = DB
+    updated_at = Column(BigInteger)
+    created_at = Column(BigInteger)
 
 
 class ModelModel(BaseModel):
+    model_config = ConfigDict(from_attributes=True)
+
     id: str
     user_id: str
     base_model_id: Optional[str] = None
@@ -115,15 +114,9 @@ class ModelForm(BaseModel):
 
 
 class ModelsTable:
-    def __init__(
-        self,
-        db: pw.SqliteDatabase | pw.PostgresqlDatabase,
-    ):
-        self.db = db
-        self.db.create_tables([Model])
 
     def insert_new_model(
-        self, form_data: ModelForm, user_id: str
+        self, db: Session, form_data: ModelForm, user_id: str
     ) -> Optional[ModelModel]:
         model = ModelModel(
             **{
@@ -134,46 +127,50 @@ class ModelsTable:
             }
         )
         try:
-            result = Model.create(**model.model_dump())
+            result = Model(**model.dict())
+            db.add(result)
+            db.commit()
+            db.refresh(result)
 
             if result:
-                return model
+                return ModelModel.model_validate(result)
             else:
                 return None
         except Exception as e:
             print(e)
             return None
 
-    def get_all_models(self) -> List[ModelModel]:
-        return [ModelModel(**model_to_dict(model)) for model in Model.select()]
+    def get_all_models(self, db: Session) -> List[ModelModel]:
+        return [ModelModel.model_validate(model) for model in db.query(Model).all()]
 
-    def get_model_by_id(self, id: str) -> Optional[ModelModel]:
+    def get_model_by_id(self, db: Session, id: str) -> Optional[ModelModel]:
         try:
-            model = Model.get(Model.id == id)
-            return ModelModel(**model_to_dict(model))
+            model = db.get(Model, id)
+            return ModelModel.model_validate(model)
         except:
             return None
 
-    def update_model_by_id(self, id: str, model: ModelForm) -> Optional[ModelModel]:
+    def update_model_by_id(
+        self, db: Session, id: str, model: ModelForm
+    ) -> Optional[ModelModel]:
         try:
             # update only the fields that are present in the model
-            query = Model.update(**model.model_dump()).where(Model.id == id)
-            query.execute()
-
-            model = Model.get(Model.id == id)
-            return ModelModel(**model_to_dict(model))
+            model = db.query(Model).get(id)
+            model.update(**model.model_dump())
+            db.commit()
+            db.refresh(model)
+            return ModelModel.model_validate(model)
         except Exception as e:
             print(e)
 
             return None
 
-    def delete_model_by_id(self, id: str) -> bool:
+    def delete_model_by_id(self, db: Session, id: str) -> bool:
         try:
-            query = Model.delete().where(Model.id == id)
-            query.execute()
+            db.query(Model).filter_by(id=id).delete()
             return True
         except:
             return False
 
 
-Models = ModelsTable(DB)
+Models = ModelsTable()