feat(sqlalchemy): Replace peewee with sqlalchemy

backend/apps/webui/routers/models.py

@@ -5,6 +5,8 @@ from typing import List, Union, Optional
 from fastapi import APIRouter
 from pydantic import BaseModel
 import json
+
+from apps.webui.internal.db import get_db
 from apps.webui.models.models import Models, ModelModel, ModelForm, ModelResponse
 
 from utils.utils import get_verified_user, get_admin_user
@@ -18,8 +20,8 @@ router = APIRouter()
 
 
 @router.get("/", response_model=List[ModelResponse])
-async def get_models(user=Depends(get_verified_user)):
-    return Models.get_all_models()
+async def get_models(user=Depends(get_verified_user), db=Depends(get_db)):
+    return Models.get_all_models(db)
 
 
 ############################
@@ -29,7 +31,10 @@ async def get_models(user=Depends(get_verified_user)):
 
 @router.post("/add", response_model=Optional[ModelModel])
 async def add_new_model(
-    request: Request, form_data: ModelForm, user=Depends(get_admin_user)
+    request: Request,
+    form_data: ModelForm,
+    user=Depends(get_admin_user),
+    db=Depends(get_db),
 ):
     if form_data.id in request.app.state.MODELS:
         raise HTTPException(
@@ -37,7 +42,7 @@ async def add_new_model(
             detail=ERROR_MESSAGES.MODEL_ID_TAKEN,
         )
     else:
-        model = Models.insert_new_model(form_data, user.id)
+        model = Models.insert_new_model(db, form_data, user.id)
 
         if model:
             return model
@@ -53,9 +58,9 @@ async def add_new_model(
 ############################
 
 
-@router.get("/", response_model=Optional[ModelModel])
-async def get_model_by_id(id: str, user=Depends(get_verified_user)):
-    model = Models.get_model_by_id(id)
+@router.get("/{id}", response_model=Optional[ModelModel])
+async def get_model_by_id(id: str, user=Depends(get_verified_user), db=Depends(get_db)):
+    model = Models.get_model_by_id(db, id)
 
     if model:
         return model
@@ -73,15 +78,19 @@ async def get_model_by_id(id: str, user=Depends(get_verified_user)):
 
 @router.post("/update", response_model=Optional[ModelModel])
 async def update_model_by_id(
-    request: Request, id: str, form_data: ModelForm, user=Depends(get_admin_user)
+    request: Request,
+    id: str,
+    form_data: ModelForm,
+    user=Depends(get_admin_user),
+    db=Depends(get_db),
 ):
-    model = Models.get_model_by_id(id)
+    model = Models.get_model_by_id(db, id)
     if model:
-        model = Models.update_model_by_id(id, form_data)
+        model = Models.update_model_by_id(db, id, form_data)
         return model
     else:
         if form_data.id in request.app.state.MODELS:
-            model = Models.insert_new_model(form_data, user.id)
+            model = Models.insert_new_model(db, form_data, user.id)
             if model:
                 return model
             else:
@@ -102,6 +111,6 @@ async def update_model_by_id(
 
 
 @router.delete("/delete", response_model=bool)
-async def delete_model_by_id(id: str, user=Depends(get_admin_user)):
-    result = Models.delete_model_by_id(id)
+async def delete_model_by_id(id: str, user=Depends(get_admin_user), db=Depends(get_db)):
+    result = Models.delete_model_by_id(db, id)
     return result