From ede71740d293c28fbd350b099cc9dd037bd4f761 Mon Sep 17 00:00:00 2001
From: "Timothy J. Baek" <timothyjrbeck@gmail.com>
Date: Wed, 16 Oct 2024 15:13:38 -0700
Subject: [PATCH] feat: folders table

---
 .../open_webui/apps/webui/models/folders.py   | 132 ++++++++++++++++++
 1 file changed, 132 insertions(+)
 create mode 100644 backend/open_webui/apps/webui/models/folders.py

diff --git a/backend/open_webui/apps/webui/models/folders.py b/backend/open_webui/apps/webui/models/folders.py
new file mode 100644
index 000000000..e8e6ffef1
--- /dev/null
+++ b/backend/open_webui/apps/webui/models/folders.py
@@ -0,0 +1,132 @@
+import logging
+import time
+import uuid
+from typing import Optional
+
+from open_webui.apps.webui.internal.db import Base, get_db
+
+
+from open_webui.env import SRC_LOG_LEVELS
+from pydantic import BaseModel, ConfigDict
+from sqlalchemy import BigInteger, Column, Text, JSON, PrimaryKeyConstraint
+
+log = logging.getLogger(__name__)
+log.setLevel(SRC_LOG_LEVELS["MODELS"])
+
+
+####################
+# Folder DB Schema
+####################
+
+
+class FolderItems(BaseModel):
+    chat_ids: Optional[list[str]] = None
+    file_ids: Optional[list[str]] = None
+    folder_ids: Optional[list[str]] = None
+
+    model_config = ConfigDict(extra="allow")
+
+
+class Folder(Base):
+    __tablename__ = "folder"
+    id = Column(Text, primary_key=True)
+    parent_id = Column(Text, nullable=True)
+    user_id = Column(Text)
+    name = Column(Text)
+    items = Column(JSON, nullable=True)
+    meta = Column(JSON, nullable=True)
+    created_at = Column(BigInteger)
+    updated_at = Column(BigInteger)
+
+
+class FolderModel(BaseModel):
+    id: str
+    parent_id: Optional[str] = None
+    user_id: str
+    name: str
+    items: Optional[FolderItems] = None
+    meta: Optional[dict] = None
+    created_at: int
+    updated_at: int
+
+    model_config = ConfigDict(from_attributes=True)
+
+
+class FolderTable:
+    def insert_new_folder(self, name: str, user_id: str) -> Optional[FolderModel]:
+        with get_db() as db:
+            id = name.lower()
+            folder = FolderModel(
+                **{
+                    "id": id,
+                    "user_id": user_id,
+                    "name": name,
+                    "created_at": int(time.time()),
+                    "updated_at": int(time.time()),
+                }
+            )
+            try:
+                result = Folder(**folder.model_dump())
+                db.add(result)
+                db.commit()
+                db.refresh(result)
+                if result:
+                    return FolderModel.model_validate(result)
+                else:
+                    return None
+            except Exception as e:
+                print(e)
+                return None
+
+    def get_folder_by_name_and_user_id(
+        self, name: str, user_id: str
+    ) -> Optional[FolderModel]:
+        try:
+            id = name.lower()
+            with get_db() as db:
+                folder = db.query(Folder).filter_by(id=id, user_id=user_id).first()
+                return FolderModel.model_validate(folder)
+        except Exception:
+            return None
+
+    def get_folders_by_user_id(self, user_id: str) -> list[FolderModel]:
+        with get_db() as db:
+            return [
+                FolderModel.model_validate(folder)
+                for folder in db.query(Folder).filter_by(user_id=user_id).all()
+            ]
+
+    def update_folder_by_name_and_user_id(
+        self, name: str, user_id: str, items: FolderItems
+    ) -> Optional[FolderModel]:
+        try:
+            id = name.lower()
+            with get_db() as db:
+                folder = db.query(Folder).filter_by(id=id, user_id=user_id).first()
+
+                folder.items = items.model_dump()
+                folder.updated_at = int(time.time())
+
+                db.commit()
+
+                return FolderModel.model_validate(folder)
+        except Exception as e:
+            log.error(f"update_folder: {e}")
+            return
+
+    def delete_folder_by_name_and_user_id(self, name: str, user_id: str) -> bool:
+        try:
+            with get_db() as db:
+                id = name.lower()
+
+                folder = db.query(Folder).filter_by(id=id, user_id=user_id).first()
+                db.delete(folder)
+
+                db.commit()
+                return True
+        except Exception as e:
+            log.error(f"delete_folder: {e}")
+            return False
+
+
+Folders = FolderTable()