refac/enh: db session sharing

backend/open_webui/models/models.py

@@ -2,7 +2,8 @@ import logging
 import time
 from typing import Optional
 
-from open_webui.internal.db import Base, JSONField, get_db
+from sqlalchemy.orm import Session
+from open_webui.internal.db import Base, JSONField, get_db, get_db_context
 
 from open_webui.models.groups import Groups
 from open_webui.models.users import User, UserModel, Users, UserResponse
@@ -150,7 +151,7 @@ class ModelForm(BaseModel):
 
 class ModelsTable:
     def insert_new_model(
-        self, form_data: ModelForm, user_id: str
+        self, form_data: ModelForm, user_id: str, db: Optional[Session] = None
     ) -> Optional[ModelModel]:
         model = ModelModel(
             **{
@@ -161,7 +162,7 @@ class ModelsTable:
             }
         )
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 result = Model(**model.model_dump())
                 db.add(result)
                 db.commit()
@@ -175,17 +176,17 @@ class ModelsTable:
             log.exception(f"Failed to insert a new model: {e}")
             return None
 
-    def get_all_models(self) -> list[ModelModel]:
-        with get_db() as db:
+    def get_all_models(self, db: Optional[Session] = None) -> list[ModelModel]:
+        with get_db_context(db) as db:
             return [ModelModel.model_validate(model) for model in db.query(Model).all()]
 
-    def get_models(self) -> list[ModelUserResponse]:
-        with get_db() as db:
+    def get_models(self, db: Optional[Session] = None) -> list[ModelUserResponse]:
+        with get_db_context(db) as db:
             all_models = db.query(Model).filter(Model.base_model_id != None).all()
 
             user_ids = list(set(model.user_id for model in all_models))
 
-            users = Users.get_users_by_user_ids(user_ids) if user_ids else []
+            users = Users.get_users_by_user_ids(user_ids, db=db) if user_ids else []
             users_dict = {user.id: user for user in users}
 
             models = []
@@ -201,18 +202,18 @@ class ModelsTable:
                 )
             return models
 
-    def get_base_models(self) -> list[ModelModel]:
-        with get_db() as db:
+    def get_base_models(self, db: Optional[Session] = None) -> list[ModelModel]:
+        with get_db_context(db) as db:
             return [
                 ModelModel.model_validate(model)
                 for model in db.query(Model).filter(Model.base_model_id == None).all()
             ]
 
     def get_models_by_user_id(
-        self, user_id: str, permission: str = "write"
+        self, user_id: str, permission: str = "write", db: Optional[Session] = None
     ) -> list[ModelUserResponse]:
-        models = self.get_models()
-        user_group_ids = {group.id for group in Groups.get_groups_by_member_id(user_id)}
+        models = self.get_models(db=db)
+        user_group_ids = {group.id for group in Groups.get_groups_by_member_id(user_id, db=db)}
         return [
             model
             for model in models
@@ -263,9 +264,9 @@ class ModelsTable:
         return query
 
     def search_models(
-        self, user_id: str, filter: dict = {}, skip: int = 0, limit: int = 30
+        self, user_id: str, filter: dict = {}, skip: int = 0, limit: int = 30, db: Optional[Session] = None
     ) -> ModelListResponse:
-        with get_db() as db:
+        with get_db_context(db) as db:
             # Join GroupMember so we can order by group_id when requested
             query = db.query(Model, User).outerjoin(User, User.id == Model.user_id)
             query = query.filter(Model.base_model_id != None)
@@ -349,24 +350,24 @@ class ModelsTable:
 
             return ModelListResponse(items=models, total=total)
 
-    def get_model_by_id(self, id: str) -> Optional[ModelModel]:
+    def get_model_by_id(self, id: str, db: Optional[Session] = None) -> Optional[ModelModel]:
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 model = db.get(Model, id)
                 return ModelModel.model_validate(model)
         except Exception:
             return None
 
-    def get_models_by_ids(self, ids: list[str]) -> list[ModelModel]:
+    def get_models_by_ids(self, ids: list[str], db: Optional[Session] = None) -> list[ModelModel]:
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 models = db.query(Model).filter(Model.id.in_(ids)).all()
                 return [ModelModel.model_validate(model) for model in models]
         except Exception:
             return []
 
-    def toggle_model_by_id(self, id: str) -> Optional[ModelModel]:
-        with get_db() as db:
+    def toggle_model_by_id(self, id: str, db: Optional[Session] = None) -> Optional[ModelModel]:
+        with get_db_context(db) as db:
             try:
                 is_active = db.query(Model).filter_by(id=id).first().is_active
 
@@ -378,13 +379,13 @@ class ModelsTable:
                 )
                 db.commit()
 
-                return self.get_model_by_id(id)
+                return self.get_model_by_id(id, db=db)
             except Exception:
                 return None
 
-    def update_model_by_id(self, id: str, model: ModelForm) -> Optional[ModelModel]:
+    def update_model_by_id(self, id: str, model: ModelForm, db: Optional[Session] = None) -> Optional[ModelModel]:
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 # update only the fields that are present in the model
                 data = model.model_dump(exclude={"id"})
                 result = db.query(Model).filter_by(id=id).update(data)
@@ -398,9 +399,9 @@ class ModelsTable:
             log.exception(f"Failed to update the model by id {id}: {e}")
             return None
 
-    def delete_model_by_id(self, id: str) -> bool:
+    def delete_model_by_id(self, id: str, db: Optional[Session] = None) -> bool:
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 db.query(Model).filter_by(id=id).delete()
                 db.commit()
 
@@ -408,9 +409,9 @@ class ModelsTable:
         except Exception:
             return False
 
-    def delete_all_models(self) -> bool:
+    def delete_all_models(self, db: Optional[Session] = None) -> bool:
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 db.query(Model).delete()
                 db.commit()
 
@@ -418,9 +419,9 @@ class ModelsTable:
         except Exception:
             return False
 
-    def sync_models(self, user_id: str, models: list[ModelModel]) -> list[ModelModel]:
+    def sync_models(self, user_id: str, models: list[ModelModel], db: Optional[Session] = None) -> list[ModelModel]:
         try:
-            with get_db() as db:
+            with get_db_context(db) as db:
                 # Get existing models
                 existing_models = db.query(Model).all()
                 existing_ids = {model.id for model in existing_models}