refac

backend/apps/webui/models/chats.py

@@ -7,7 +7,7 @@ import time
 
 from sqlalchemy import Column, String, BigInteger, Boolean, Text
 
-from apps.webui.internal.db import Base, Session
+from apps.webui.internal.db import Base, get_db
 
 
 ####################
@@ -79,87 +79,99 @@ class ChatTitleIdResponse(BaseModel):
 class ChatTable:
 
     def insert_new_chat(self, user_id: str, form_data: ChatForm) -> Optional[ChatModel]:
-        id = str(uuid.uuid4())
-        chat = ChatModel(
-            **{
-                "id": id,
-                "user_id": user_id,
-                "title": (
-                    form_data.chat["title"] if "title" in form_data.chat else "New Chat"
-                ),
-                "chat": json.dumps(form_data.chat),
-                "created_at": int(time.time()),
-                "updated_at": int(time.time()),
-            }
-        )
+        with get_db() as db:
 
-        result = Chat(**chat.model_dump())
-        Session.add(result)
-        Session.commit()
-        Session.refresh(result)
-        return ChatModel.model_validate(result) if result else None
+            id = str(uuid.uuid4())
+            chat = ChatModel(
+                **{
+                    "id": id,
+                    "user_id": user_id,
+                    "title": (
+                        form_data.chat["title"]
+                        if "title" in form_data.chat
+                        else "New Chat"
+                    ),
+                    "chat": json.dumps(form_data.chat),
+                    "created_at": int(time.time()),
+                    "updated_at": int(time.time()),
+                }
+            )
+
+            result = Chat(**chat.model_dump())
+            db.add(result)
+            db.commit()
+            db.refresh(result)
+            return ChatModel.model_validate(result) if result else None
 
     def update_chat_by_id(self, id: str, chat: dict) -> Optional[ChatModel]:
         try:
-            chat_obj = Session.get(Chat, id)
-            chat_obj.chat = json.dumps(chat)
-            chat_obj.title = chat["title"] if "title" in chat else "New Chat"
-            chat_obj.updated_at = int(time.time())
-            Session.commit()
-            Session.refresh(chat_obj)
+            with get_db() as db:
 
-            return ChatModel.model_validate(chat_obj)
+                chat_obj = db.get(Chat, id)
+                chat_obj.chat = json.dumps(chat)
+                chat_obj.title = chat["title"] if "title" in chat else "New Chat"
+                chat_obj.updated_at = int(time.time())
+                db.commit()
+                db.refresh(chat_obj)
+
+                return ChatModel.model_validate(chat_obj)
         except Exception as e:
             return None
 
     def insert_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
-        # Get the existing chat to share
-        chat = Session.get(Chat, chat_id)
-        # Check if the chat is already shared
-        if chat.share_id:
-            return self.get_chat_by_id_and_user_id(chat.share_id, "shared")
-        # Create a new chat with the same data, but with a new ID
-        shared_chat = ChatModel(
-            **{
-                "id": str(uuid.uuid4()),
-                "user_id": f"shared-{chat_id}",
-                "title": chat.title,
-                "chat": chat.chat,
-                "created_at": chat.created_at,
-                "updated_at": int(time.time()),
-            }
-        )
-        shared_result = Chat(**shared_chat.model_dump())
-        Session.add(shared_result)
-        Session.commit()
-        Session.refresh(shared_result)
-        # Update the original chat with the share_id
-        result = (
-            Session.query(Chat)
-            .filter_by(id=chat_id)
-            .update({"share_id": shared_chat.id})
-        )
+        with get_db() as db:
 
-        return shared_chat if (shared_result and result) else None
+            # Get the existing chat to share
+            chat = db.get(Chat, chat_id)
+            # Check if the chat is already shared
+            if chat.share_id:
+                return self.get_chat_by_id_and_user_id(chat.share_id, "shared")
+            # Create a new chat with the same data, but with a new ID
+            shared_chat = ChatModel(
+                **{
+                    "id": str(uuid.uuid4()),
+                    "user_id": f"shared-{chat_id}",
+                    "title": chat.title,
+                    "chat": chat.chat,
+                    "created_at": chat.created_at,
+                    "updated_at": int(time.time()),
+                }
+            )
+            shared_result = Chat(**shared_chat.model_dump())
+            db.add(shared_result)
+            db.commit()
+            db.refresh(shared_result)
+            # Update the original chat with the share_id
+            result = (
+                db.query(Chat)
+                .filter_by(id=chat_id)
+                .update({"share_id": shared_chat.id})
+            )
+
+            return shared_chat if (shared_result and result) else None
 
     def update_shared_chat_by_chat_id(self, chat_id: str) -> Optional[ChatModel]:
         try:
-            print("update_shared_chat_by_id")
-            chat = Session.get(Chat, chat_id)
-            print(chat)
-            chat.title = chat.title
-            chat.chat = chat.chat
-            Session.commit()
-            Session.refresh(chat)
+            with get_db() as db:
 
-            return self.get_chat_by_id(chat.share_id)
+                print("update_shared_chat_by_id")
+                chat = db.get(Chat, chat_id)
+                print(chat)
+                chat.title = chat.title
+                chat.chat = chat.chat
+                db.commit()
+                db.refresh(chat)
+
+                return self.get_chat_by_id(chat.share_id)
         except:
             return None
 
     def delete_shared_chat_by_chat_id(self, chat_id: str) -> bool:
         try:
-            Session.query(Chat).filter_by(user_id=f"shared-{chat_id}").delete()
-            return True
+            with get_db() as db:
+
+                db.query(Chat).filter_by(user_id=f"shared-{chat_id}").delete()
+                return True
         except:
             return False
 
@@ -167,42 +179,50 @@ class ChatTable:
         self, id: str, share_id: Optional[str]
     ) -> Optional[ChatModel]:
         try:
-            chat = Session.get(Chat, id)
-            chat.share_id = share_id
-            Session.commit()
-            Session.refresh(chat)
-            return ChatModel.model_validate(chat)
+            with get_db() as db:
+
+                chat = db.get(Chat, id)
+                chat.share_id = share_id
+                db.commit()
+                db.refresh(chat)
+                return ChatModel.model_validate(chat)
         except:
             return None
 
     def toggle_chat_archive_by_id(self, id: str) -> Optional[ChatModel]:
         try:
-            chat = Session.get(Chat, id)
-            chat.archived = not chat.archived
-            Session.commit()
-            Session.refresh(chat)
-            return ChatModel.model_validate(chat)
+            with get_db() as db:
+
+                chat = db.get(Chat, id)
+                chat.archived = not chat.archived
+                db.commit()
+                db.refresh(chat)
+                return ChatModel.model_validate(chat)
         except:
             return None
 
     def archive_all_chats_by_user_id(self, user_id: str) -> bool:
         try:
-            Session.query(Chat).filter_by(user_id=user_id).update({"archived": True})
-            return True
+            with get_db() as db:
+
+                db.query(Chat).filter_by(user_id=user_id).update({"archived": True})
+                return True
         except:
             return False
 
     def get_archived_chat_list_by_user_id(
         self, user_id: str, skip: int = 0, limit: int = 50
     ) -> List[ChatModel]:
-        all_chats = (
-            Session.query(Chat)
-            .filter_by(user_id=user_id, archived=True)
-            .order_by(Chat.updated_at.desc())
-            # .limit(limit).offset(skip)
-            .all()
-        )
-        return [ChatModel.model_validate(chat) for chat in all_chats]
+        with get_db() as db:
+
+            all_chats = (
+                db.query(Chat)
+                .filter_by(user_id=user_id, archived=True)
+                .order_by(Chat.updated_at.desc())
+                # .limit(limit).offset(skip)
+                .all()
+            )
+            return [ChatModel.model_validate(chat) for chat in all_chats]
 
     def get_chat_list_by_user_id(
         self,
@@ -211,110 +231,136 @@ class ChatTable:
         skip: int = 0,
         limit: int = 50,
     ) -> List[ChatModel]:
-        query = Session.query(Chat).filter_by(user_id=user_id)
-        if not include_archived:
-            query = query.filter_by(archived=False)
-        all_chats = (
-            query.order_by(Chat.updated_at.desc())
-            # .limit(limit).offset(skip)
-            .all()
-        )
-        return [ChatModel.model_validate(chat) for chat in all_chats]
+        with get_db() as db:
+            query = db.query(Chat).filter_by(user_id=user_id)
+            if not include_archived:
+                query = query.filter_by(archived=False)
+            all_chats = (
+                query.order_by(Chat.updated_at.desc())
+                # .limit(limit).offset(skip)
+                .all()
+            )
+            return [ChatModel.model_validate(chat) for chat in all_chats]
 
     def get_chat_list_by_chat_ids(
         self, chat_ids: List[str], skip: int = 0, limit: int = 50
     ) -> List[ChatModel]:
-        all_chats = (
-            Session.query(Chat)
-            .filter(Chat.id.in_(chat_ids))
-            .filter_by(archived=False)
-            .order_by(Chat.updated_at.desc())
-            .all()
-        )
-        return [ChatModel.model_validate(chat) for chat in all_chats]
+
+        with get_db() as db:
+
+            all_chats = (
+                db.query(Chat)
+                .filter(Chat.id.in_(chat_ids))
+                .filter_by(archived=False)
+                .order_by(Chat.updated_at.desc())
+                .all()
+            )
+            return [ChatModel.model_validate(chat) for chat in all_chats]
 
     def get_chat_by_id(self, id: str) -> Optional[ChatModel]:
         try:
-            chat = Session.get(Chat, id)
-            return ChatModel.model_validate(chat)
+            with get_db() as db:
+
+                chat = db.get(Chat, id)
+                return ChatModel.model_validate(chat)
         except:
             return None
 
     def get_chat_by_share_id(self, id: str) -> Optional[ChatModel]:
         try:
-            chat = Session.query(Chat).filter_by(share_id=id).first()
+            with get_db() as db:
 
-            if chat:
-                return self.get_chat_by_id(id)
-            else:
-                return None
+                chat = db.query(Chat).filter_by(share_id=id).first()
+
+                if chat:
+                    return self.get_chat_by_id(id)
+                else:
+                    return None
         except Exception as e:
             return None
 
     def get_chat_by_id_and_user_id(self, id: str, user_id: str) -> Optional[ChatModel]:
         try:
-            chat = Session.query(Chat).filter_by(id=id, user_id=user_id).first()
-            return ChatModel.model_validate(chat)
+            with get_db() as db:
+
+                chat = db.query(Chat).filter_by(id=id, user_id=user_id).first()
+                return ChatModel.model_validate(chat)
         except:
             return None
 
     def get_chats(self, skip: int = 0, limit: int = 50) -> List[ChatModel]:
-        all_chats = (
-            Session.query(Chat)
-            # .limit(limit).offset(skip)
-            .order_by(Chat.updated_at.desc())
-        )
-        return [ChatModel.model_validate(chat) for chat in all_chats]
+        with get_db() as db:
+
+            all_chats = (
+                db.query(Chat)
+                # .limit(limit).offset(skip)
+                .order_by(Chat.updated_at.desc())
+            )
+            return [ChatModel.model_validate(chat) for chat in all_chats]
 
     def get_chats_by_user_id(self, user_id: str) -> List[ChatModel]:
-        all_chats = (
-            Session.query(Chat)
-            .filter_by(user_id=user_id)
-            .order_by(Chat.updated_at.desc())
-        )
-        return [ChatModel.model_validate(chat) for chat in all_chats]
+        with get_db() as db:
+
+            all_chats = (
+                db.query(Chat)
+                .filter_by(user_id=user_id)
+                .order_by(Chat.updated_at.desc())
+            )
+            return [ChatModel.model_validate(chat) for chat in all_chats]
 
     def get_archived_chats_by_user_id(self, user_id: str) -> List[ChatModel]:
-        all_chats = (
-            Session.query(Chat)
-            .filter_by(user_id=user_id, archived=True)
-            .order_by(Chat.updated_at.desc())
-        )
-        return [ChatModel.model_validate(chat) for chat in all_chats]
+        with get_db() as db:
+
+            all_chats = (
+                db.query(Chat)
+                .filter_by(user_id=user_id, archived=True)
+                .order_by(Chat.updated_at.desc())
+            )
+            return [ChatModel.model_validate(chat) for chat in all_chats]
 
     def delete_chat_by_id(self, id: str) -> bool:
         try:
-            Session.query(Chat).filter_by(id=id).delete()
+            with get_db() as db:
 
-            return True and self.delete_shared_chat_by_chat_id(id)
+                db.query(Chat).filter_by(id=id).delete()
+
+                return True and self.delete_shared_chat_by_chat_id(id)
         except:
             return False
 
     def delete_chat_by_id_and_user_id(self, id: str, user_id: str) -> bool:
         try:
-            Session.query(Chat).filter_by(id=id, user_id=user_id).delete()
+            with get_db() as db:
 
-            return True and self.delete_shared_chat_by_chat_id(id)
+                db.query(Chat).filter_by(id=id, user_id=user_id).delete()
+
+                return True and self.delete_shared_chat_by_chat_id(id)
         except:
             return False
 
     def delete_chats_by_user_id(self, user_id: str) -> bool:
         try:
-            self.delete_shared_chats_by_user_id(user_id)
 
-            Session.query(Chat).filter_by(user_id=user_id).delete()
-            return True
+            with get_db() as db:
+
+                self.delete_shared_chats_by_user_id(user_id)
+
+                db.query(Chat).filter_by(user_id=user_id).delete()
+                return True
         except:
             return False
 
     def delete_shared_chats_by_user_id(self, user_id: str) -> bool:
         try:
-            chats_by_user = Session.query(Chat).filter_by(user_id=user_id).all()
-            shared_chat_ids = [f"shared-{chat.id}" for chat in chats_by_user]
 
-            Session.query(Chat).filter(Chat.user_id.in_(shared_chat_ids)).delete()
+            with get_db() as db:
 
-            return True
+                chats_by_user = db.query(Chat).filter_by(user_id=user_id).all()
+                shared_chat_ids = [f"shared-{chat.id}" for chat in chats_by_user]
+
+                db.query(Chat).filter(Chat.user_id.in_(shared_chat_ids)).delete()
+
+                return True
         except:
             return False