PraisonAI 3.0.0__py3-none-any.whl

praisonai/ui/db.py

@@ -0,0 +1,294 @@
+import os
+import sqlite3
+import asyncio
+import shutil
+import logging
+from sqlalchemy import text
+from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession
+from sqlalchemy.orm import sessionmaker
+from sql_alchemy import SQLAlchemyDataLayer
+import chainlit.data as cl_data
+from chainlit.types import ThreadDict
+from database_config import get_database_url_with_sqlite_override
+from praisonai.ui.chainlit_compat import create_local_storage_client
+
+def ensure_directories():
+    """Ensure required directories exist"""
+    if "CHAINLIT_APP_ROOT" not in os.environ:
+        chainlit_root = os.path.join(os.path.expanduser("~"), ".praison")
+        os.environ["CHAINLIT_APP_ROOT"] = chainlit_root
+    else:
+        chainlit_root = os.environ["CHAINLIT_APP_ROOT"]
+
+    os.makedirs(chainlit_root, exist_ok=True)
+    os.makedirs(os.path.join(chainlit_root, ".files"), exist_ok=True)
+    
+    # Copy public folder and chainlit.md if they don't exist
+    public_folder = os.path.join(os.path.dirname(__file__), "public")
+    config_folder = os.path.join(os.path.dirname(__file__), "config")
+    
+    # Copy public folder
+    if not os.path.exists(os.path.join(chainlit_root, "public")):
+        if os.path.exists(public_folder):
+            shutil.copytree(public_folder, os.path.join(chainlit_root, "public"), dirs_exist_ok=True)
+            logging.info("Public folder copied successfully!")
+        else:
+            logging.info("Public folder not found in the package.")
+    
+    # Copy all files from config folder to root if translations doesn't exist
+    if not os.path.exists(os.path.join(chainlit_root, "translations")):
+        os.makedirs(os.path.join(chainlit_root, "translations"), exist_ok=True)
+        
+        if os.path.exists(config_folder):
+            for item in os.listdir(config_folder):
+                src_path = os.path.join(config_folder, item)
+                dst_path = os.path.join(chainlit_root, item)
+                if os.path.isfile(src_path):
+                    shutil.copy2(src_path, dst_path)
+                    logging.info(f"File {item} copied to root successfully!")
+                elif os.path.isdir(src_path):
+                    if os.path.exists(dst_path):
+                        shutil.rmtree(dst_path)
+                    shutil.copytree(src_path, dst_path)
+                    logging.info(f"Directory {item} copied to root successfully!")
+        else:
+            logging.info("Config folder not found in the package.")
+
+# Create directories at module import time
+ensure_directories()
+
+class DatabaseManager(SQLAlchemyDataLayer):
+    def __init__(self):
+        # Check FORCE_SQLITE flag to bypass external database detection
+        self.database_url = get_database_url_with_sqlite_override()
+        
+        if self.database_url:
+            self.conninfo = self.database_url
+        else:
+            chainlit_root = os.environ["CHAINLIT_APP_ROOT"]  # Now using CHAINLIT_APP_ROOT
+            self.db_path = os.path.join(chainlit_root, "database.sqlite")
+            self.conninfo = f"sqlite+aiosqlite:///{self.db_path}"
+        
+        # Create local file storage client for element persistence
+        storage_client = create_local_storage_client()
+        
+        # Initialize SQLAlchemyDataLayer with the connection info and storage client
+        super().__init__(conninfo=self.conninfo, storage_provider=storage_client)
+
+    async def create_schema_async(self):
+        """Create the database schema for PostgreSQL"""
+        if not self.database_url:
+            return
+        engine = create_async_engine(self.database_url, echo=False)
+        async with engine.begin() as conn:
+            await conn.execute(text('''
+                CREATE TABLE IF NOT EXISTS users (
+                    "id" TEXT PRIMARY KEY,
+                    "identifier" TEXT NOT NULL UNIQUE,
+                    "meta" TEXT NOT NULL DEFAULT '{}',
+                    "createdAt" TEXT
+                );
+            '''))
+            await conn.execute(text('''
+                CREATE TABLE IF NOT EXISTS threads (
+                    "id" TEXT PRIMARY KEY,
+                    "createdAt" TEXT,
+                    "name" TEXT,
+                    "userId" TEXT,
+                    "userIdentifier" TEXT,
+                    "tags" TEXT DEFAULT '[]',
+                    "meta" TEXT NOT NULL DEFAULT '{}',
+                    FOREIGN KEY ("userId") REFERENCES users("id") ON DELETE CASCADE
+                );
+            '''))
+            await conn.execute(text('''
+                CREATE TABLE IF NOT EXISTS steps (
+                    "id" TEXT PRIMARY KEY,
+                    "name" TEXT NOT NULL,
+                    "type" TEXT NOT NULL,
+                    "threadId" TEXT NOT NULL,
+                    "parentId" TEXT,
+                    "disableFeedback" BOOLEAN NOT NULL DEFAULT FALSE,
+                    "streaming" BOOLEAN NOT NULL DEFAULT FALSE,
+                    "waitForAnswer" BOOLEAN DEFAULT FALSE,
+                    "isError" BOOLEAN NOT NULL DEFAULT FALSE,
+                    "meta" TEXT DEFAULT '{}',
+                    "tags" TEXT DEFAULT '[]',
+                    "input" TEXT,
+                    "output" TEXT,
+                    "createdAt" TEXT,
+                    "startTime" TEXT,
+                    "endTime" TEXT,
+                    "generation" TEXT,
+                    "showInput" TEXT,
+                    "language" TEXT,
+                    "indent" INT,
+                    FOREIGN KEY ("threadId") REFERENCES threads("id") ON DELETE CASCADE
+                );
+            '''))
+            await conn.execute(text('''
+                CREATE TABLE IF NOT EXISTS elements (
+                    "id" TEXT PRIMARY KEY,
+                    "threadId" TEXT,
+                    "type" TEXT,
+                    "url" TEXT,
+                    "chainlitKey" TEXT,
+                    "name" TEXT NOT NULL,
+                    "display" TEXT,
+                    "objectKey" TEXT,
+                    "size" TEXT,
+                    "page" INT,
+                    "language" TEXT,
+                    "forId" TEXT,
+                    "mime" TEXT,
+                    FOREIGN KEY ("threadId") REFERENCES threads("id") ON DELETE CASCADE
+                );
+            '''))
+            await conn.execute(text('''
+                CREATE TABLE IF NOT EXISTS feedbacks (
+                    "id" TEXT PRIMARY KEY,
+                    "forId" TEXT NOT NULL,
+                    "value" INT NOT NULL,
+                    "threadId" TEXT,
+                    "comment" TEXT
+                );
+            '''))
+            await conn.execute(text('''
+                CREATE TABLE IF NOT EXISTS settings (
+                    "id" SERIAL PRIMARY KEY,
+                    "key" TEXT UNIQUE,
+                    "value" TEXT
+                );
+            '''))
+        await engine.dispose()
+
+    def create_schema_sqlite(self):
+        """Create the database schema for SQLite"""
+        chainlit_root = os.environ["CHAINLIT_APP_ROOT"]  # Now using CHAINLIT_APP_ROOT
+        self.db_path = os.path.join(chainlit_root, "database.sqlite")
+        conn = sqlite3.connect(self.db_path)
+        cursor = conn.cursor()
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS users (
+                id TEXT PRIMARY KEY,
+                identifier TEXT NOT NULL UNIQUE,
+                meta TEXT NOT NULL DEFAULT '{}',
+                createdAt TEXT
+            );
+        ''')
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS threads (
+                id TEXT PRIMARY KEY,
+                createdAt TEXT,
+                name TEXT,
+                userId TEXT,
+                userIdentifier TEXT,
+                tags TEXT DEFAULT '[]',
+                meta TEXT NOT NULL DEFAULT '{}',
+                FOREIGN KEY (userId) REFERENCES users(id) ON DELETE CASCADE
+            );
+        ''')
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS steps (
+                id TEXT PRIMARY KEY,
+                name TEXT NOT NULL,
+                type TEXT NOT NULL,
+                threadId TEXT NOT NULL,
+                parentId TEXT,
+                disableFeedback BOOLEAN NOT NULL DEFAULT 0,
+                streaming BOOLEAN NOT NULL DEFAULT 0,
+                waitForAnswer BOOLEAN DEFAULT 0,
+                isError BOOLEAN NOT NULL DEFAULT 0,
+                meta TEXT DEFAULT '{}',
+                tags TEXT DEFAULT '[]',
+                input TEXT,
+                output TEXT,
+                createdAt TEXT,
+                startTime TEXT,
+                endTime TEXT,
+                generation TEXT,
+                showInput TEXT,
+                language TEXT,
+                indent INT,
+                FOREIGN KEY (threadId) REFERENCES threads(id) ON DELETE CASCADE
+            );
+        ''')
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS elements (
+                id TEXT PRIMARY KEY,
+                threadId TEXT,
+                type TEXT,
+                url TEXT,
+                chainlitKey TEXT,
+                name TEXT NOT NULL,
+                display TEXT,
+                objectKey TEXT,
+                size TEXT,
+                page INT,
+                language TEXT,
+                forId TEXT,
+                mime TEXT,
+                FOREIGN KEY (threadId) REFERENCES threads(id) ON DELETE CASCADE
+            );
+        ''')
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS feedbacks (
+                id TEXT PRIMARY KEY,
+                forId TEXT NOT NULL,
+                value INT NOT NULL,
+                threadId TEXT,
+                comment TEXT
+            );
+        ''')
+        cursor.execute('''
+            CREATE TABLE IF NOT EXISTS settings (
+                id INTEGER PRIMARY KEY AUTOINCREMENT,
+                key TEXT UNIQUE,
+                value TEXT
+            );
+        ''')
+        conn.commit()
+        conn.close()
+
+    def initialize(self):
+        """Initialize the database with schema based on the configuration"""
+        if self.database_url:
+            asyncio.run(self.create_schema_async())
+        else:
+            self.create_schema_sqlite()
+
+    async def save_setting(self, key: str, value: str):
+        """Save a setting to the database"""
+        if self.database_url:
+            async with self.engine.begin() as conn:
+                await conn.execute(text("""
+                    INSERT INTO settings ("key", "value") VALUES (:key, :value)
+                    ON CONFLICT ("key") DO UPDATE SET "value" = EXCLUDED."value"
+                """), {"key": key, "value": value})
+        else:
+            conn = sqlite3.connect(self.db_path)
+            cursor = conn.cursor()
+            cursor.execute(
+                """
+                INSERT OR REPLACE INTO settings (id, key, value)
+                VALUES ((SELECT id FROM settings WHERE key = ?), ?, ?)
+                """,
+                (key, key, value),
+            )
+            conn.commit()
+            conn.close()
+
+    async def load_setting(self, key: str) -> str:
+        """Load a setting from the database"""
+        if self.database_url:
+            async with self.engine.connect() as conn:
+                result = await conn.execute(text('SELECT "value" FROM settings WHERE "key" = :key'), {"key": key})
+                row = result.fetchone()
+                return row[0] if row else None
+        else:
+            conn = sqlite3.connect(self.db_path)
+            cursor = conn.cursor()
+            cursor.execute('SELECT value FROM settings WHERE key = ?', (key,))
+            result = cursor.fetchone()
+            conn.close()
+            return result[0] if result else None