Merge pull request #297 from ParisNeo/main

fixes

api/__init__.py

@@ -111,7 +111,7 @@ class ModelProcess:
 
     def reset_config_result(self):
         self._set_config_result = {
-            'status': 'succeeded',
+            'status': True,
             'binding_status':True,
             'model_status':True,
             'personalities_status':True,
@@ -253,10 +253,12 @@ class ModelProcess:
     def rebuild_personalities(self):
         mounted_personalities=[]
         ASCIIColors.success(f" ******************* Building mounted Personalities from main Process *************************")
-        for personality in self.config['personalities']:
+        for i,personality in enumerate(self.config['personalities']):
             try:
-                print(f" {personality}")
                 personality_path = self.lollms_paths.personalities_zoo_path/f"{personality}"
+                if i==self.config["active_personality_id"]:
+                    ASCIIColors.red("*", end="")
+                print(f" {personality}")
                 print(f"Loading from {personality_path}")
                 personality = AIPersonality(self.lollms_paths, personality_path, run_scripts=False)
                 mounted_personalities.append(personality)
@@ -265,7 +267,7 @@ class ModelProcess:
                 if self.config["debug"]:
                     print(ex)
                 personality = AIPersonality(self.lollms_paths)
-            
+        print(f'selected : {self.config["active_personality_id"]}')
         ASCIIColors.success(f" ************ Personalities mounted (Main process) ***************************")
             
         return mounted_personalities
@@ -275,8 +277,10 @@ class ModelProcess:
         failed_personalities=[]
         self.reset_config_result()
         ASCIIColors.success(f" ******************* Building mounted Personalities from generation Process *************************")
-        for personality in self.config['personalities']:
+        for i,personality in enumerate(self.config['personalities']):
             try:
+                if i==self.config["active_personality_id"]:
+                    ASCIIColors.red("*", end="")
                 print(f" {personality}")
                 personality_path = self.lollms_paths.personalities_zoo_path/f"{personality}"
                 personality = AIPersonality(self.lollms_paths, personality_path, run_scripts=True, model=self.model)
@@ -288,6 +292,7 @@ class ModelProcess:
                 failed_personalities.append(personality_path)
                 self._set_config_result['errors'].append(f"couldn't build personalities:{ex}")
                 
+        print(f'selected : {self.config["active_personality_id"]}')            
         ASCIIColors.success(f" ************ Personalities mounted (Generation process) ***************************")
         if len(failed_personalities)==len(self.config['personalities']):
             self._set_config_result['status'] ='failed'
@@ -458,7 +463,8 @@ class ModelProcess:
     def _set_config(self, config):
         bk_cfg = self.config
         self.config = config
-        print("Changing configuration")
+        ASCIIColors.info("Upgrading configuration...",end="")
+        
         # verify that the binding is the same
         if self.config["binding_name"]!=bk_cfg["binding_name"] or self.config["model_name"]!=bk_cfg["model_name"]:
             self._rebuild_model()
@@ -508,7 +514,10 @@ class LoLLMsAPPI():
             self.db = DiscussionsDB(self.lollms_paths.personal_path/"databases"/self.db_path)
 
         # If the database is empty, populate it with tables
-        self.db.populate()
+        ASCIIColors.info("Checking discussions database... ",end="")
+        self.db.create_tables()
+        self.db.add_missing_columns()
+        ASCIIColors.success("ok")
 
         # This is used to keep track of messages 
         self.full_message_list = []
@@ -539,7 +548,7 @@ class LoLLMsAPPI():
 
                 if installation_path.exists():
                     print("Error: Model already exists")
-                    socketio.emit('install_progress',{'status': 'failed', 'error': 'model already exists'}, room=room_id)
+                    socketio.emit('install_progress',{'status': False, 'error': 'model already exists'}, room=room_id)
                 
                 socketio.emit('install_progress',{'status': 'progress', 'progress': progress}, room=room_id)
                 
@@ -550,7 +559,7 @@ class LoLLMsAPPI():
                     self.binding.download_model(model_path, installation_path, callback)
                 else:
                     self.download_file(model_path, installation_path, callback)
-                socketio.emit('install_progress',{'status': 'succeeded', 'error': ''}, room=room_id)
+                socketio.emit('install_progress',{'status': True, 'error': ''}, room=room_id)
             tpe = threading.Thread(target=install_model_, args=())
             tpe.start()
             
@@ -563,10 +572,10 @@ class LoLLMsAPPI():
             installation_path = installation_dir / filename
 
             if not installation_path.exists():
-                socketio.emit('install_progress',{'status': 'failed', 'error': 'The model does not exist'}, room=request.sid)
+                socketio.emit('install_progress',{'status': False, 'error': 'The model does not exist'}, room=request.sid)
 
             installation_path.unlink()
-            socketio.emit('install_progress',{'status': 'succeeded', 'error': ''}, room=request.sid)
+            socketio.emit('install_progress',{'status': True, 'error': ''}, room=request.sid)
             
 
         
@@ -587,7 +596,7 @@ class LoLLMsAPPI():
                 )
 
                 self.current_user_message_id = message_id
-                print("Starting message generation")
+                ASCIIColors.green("Starting message generation by"+self.personality.name)
                 tpe = threading.Thread(target=self.start_message_generation, args=(message, message_id))
                 tpe.start()
             else:
@@ -642,8 +651,10 @@ class LoLLMsAPPI():
             print("update_settings : New model selected")
             if hasattr(self,"process"):
                 result = self.process.set_config(self.config)
-                print("Set config results:")
+                if result["status"]:
-                print(result)
+                    ASCIIColors.success("OK")
+                else:
+                    ASCIIColors.error("NOK")
         except Exception as ex:
             ASCIIColors.error(f"Couldn't load model.Process returned exception : {ex}")
             ASCIIColors.error(f"Binding returned this exception : {ex}")
@@ -657,8 +668,10 @@ class LoLLMsAPPI():
             print("update_settings : New personality selected")
             if hasattr(self,"process"):
                 result = self.process.set_config(self.config)
-                print("Set config results:")
+                if result["status"]:
-                print(result)
+                    ASCIIColors.success("OK")
+                else:
+                    ASCIIColors.error("NOK")
         except Exception as ex:
             ASCIIColors.error(f"Couldn't load personality. Please verify your configuration file at {self.configuration_path} or use the next menu to select a valid personality")
             ASCIIColors.error(f"Binding returned this exception : {ex}")

api/db.py

@@ -2,6 +2,8 @@
 import sqlite3
 from pathlib import Path
 from datetime import datetime
+from lollms.helpers import ASCIIColors
+
 __author__ = "parisneo"
 __github__ = "https://github.com/ParisNeo/lollms-webui"
 __copyright__ = "Copyright 2023, "
@@ -17,157 +19,98 @@ class DiscussionsDB:
         self.db_path = Path(db_path)
         self.db_path .parent.mkdir(exist_ok=True, parents= True)
 
-    def populate(self):
-        """
-        create database schema
-        """
-        db_version = 5
-        # Verify encoding and change it if it is not complient
-        with sqlite3.connect(self.db_path) as conn:
-            # Execute a PRAGMA statement to get the current encoding of the database
-            cur = conn.execute('PRAGMA encoding')
-            current_encoding = cur.fetchone()[0]
 
-            if current_encoding != 'UTF-8':
+    def create_tables(self):
-                # The current encoding is not UTF-8, so we need to change it
+        db_version = 6
-                print(f"The current encoding is {current_encoding}, changing to UTF-8...")
-                conn.execute('PRAGMA encoding = "UTF-8"')
-                conn.commit()        
-
-        print("Checking discussions database...")
         with sqlite3.connect(self.db_path) as conn:
             cursor = conn.cursor()
 
-            discussion_table_exist=False
+            cursor.execute("""
-            message_table_exist=False
+                CREATE TABLE IF NOT EXISTS schema_version (
-            schema_table_exist=False
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    version INTEGER NOT NULL
+                )
+            """)
 
-            # Check if the 'schema_version' table exists
+            cursor.execute("""
-            try:
+                CREATE TABLE IF NOT EXISTS discussion (
-                cursor.execute("""
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
-                    CREATE TABLE schema_version (
+                    title TEXT,
-                        id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-                        version INTEGER NOT NULL
+                )
-                    )
+            """)
-                """)
-            except:
-                schema_table_exist = True
-            try:
-                cursor.execute("""
-                    CREATE TABLE discussion (
-                        id INTEGER PRIMARY KEY AUTOINCREMENT,
-                        title TEXT,
-                        created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
-                    )
-                    """)
-            except Exception:
-                discussion_table_exist=True        
-            try:
-                cursor.execute("""
-                        CREATE TABLE message (
-                            id INTEGER PRIMARY KEY AUTOINCREMENT,
-                            binding TEXT,
-                            model TEXT,
-                            personality TEXT,
-                            sender TEXT NOT NULL,
-                            content TEXT NOT NULL,
-                            type INT NOT NULL,
-                            rank INT NOT NULL,
-                            parent INT,
-                            created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
-                            discussion_id INTEGER NOT NULL,
-                            FOREIGN KEY (discussion_id) REFERENCES discussion(id),
-                            FOREIGN KEY (parent) REFERENCES message(id)
-                        )
-                    """
-                    )
-            except :
-                message_table_exist=True
 
-            # Get the current version from the schema_version table
+            cursor.execute("""
-            cursor.execute("SELECT version FROM schema_version WHERE id = 1")
+                CREATE TABLE IF NOT EXISTS message (
+                    id INTEGER PRIMARY KEY AUTOINCREMENT,
+                    binding TEXT,
+                    model TEXT,
+                    personality TEXT,
+                    sender TEXT NOT NULL,
+                    content TEXT NOT NULL,
+                    type INT NOT NULL,
+                    rank INT NOT NULL,
+                    parent INT,
+                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    finished_generating_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
+                    discussion_id INTEGER NOT NULL,
+                    FOREIGN KEY (discussion_id) REFERENCES discussion(id),
+                    FOREIGN KEY (parent) REFERENCES message(id)
+                )
+            """)
+
+            cursor.execute("SELECT * FROM schema_version")
             row = cursor.fetchone()
+
             if row is None:
-                # If the table is empty, assume version 0
+                cursor.execute("INSERT INTO schema_version (version) VALUES (?)", (db_version,))
-                version = 0
             else:
-                # Otherwise, use the version from the table
+                cursor.execute("UPDATE schema_version SET version = ?", (db_version,))            
-                version = row[0]
-
-            # Upgrade the schema to version 1
-            if version < 1:
-                print(f"Upgrading schema to version {db_version}...")
-                # Add the 'created_at' column to the 'message' table
-                if message_table_exist:
-                    cursor.execute("ALTER TABLE message ADD COLUMN type INT DEFAULT 0") # Added in V1
-                    cursor.execute("ALTER TABLE message ADD COLUMN rank INT DEFAULT 0") # Added in V1
-                    cursor.execute("ALTER TABLE message ADD COLUMN parent INT DEFAULT 0") # Added in V2
-                    cursor.execute("ALTER TABLE message ADD COLUMN created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP") # Added in V3
-                    cursor.execute("ALTER TABLE discussion ADD COLUMN created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP") # Added in V3
-
-                    cursor.execute("ALTER TABLE message ADD COLUMN model TEXT DEFAULT ''") # Added in V4
-                    cursor.execute("ALTER TABLE message ADD COLUMN personality INT DEFAULT ''") # Added in V4
-                    
-            # Upgrade the schema to version 1
-            elif version < 2:
-                print(f"Upgrading schema to version {db_version}...")
-                # Add the 'created_at' column to the 'message' table
-                if message_table_exist:
-                    try:
-                        cursor.execute("ALTER TABLE message ADD COLUMN parent INT DEFAULT 0") # Added in V2
-                        cursor.execute("ALTER TABLE message ADD COLUMN created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP") # Added in V3
-                        cursor.execute("ALTER TABLE discussion ADD COLUMN created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP") # Added in V3
-
-                        cursor.execute("ALTER TABLE message ADD COLUMN model TEXT DEFAULT ''") # Added in V4
-                        cursor.execute("ALTER TABLE message ADD COLUMN personality INT DEFAULT ''") # Added in V4
-                        
-                        cursor.execute("ALTER TABLE message ADD COLUMN binding TEXT DEFAULT ''") # Added in V5
-                        
-                        
-                    except :
-                        pass
-            # Upgrade the schema to version 1
-            elif version < 3:
-                print(f"Upgrading schema to version {db_version}...")
-                # Add the 'created_at' column to the 'message' table
-                if message_table_exist:
-                    try:
-                        cursor.execute("ALTER TABLE message ADD COLUMN created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP") # Added in V3
-                        cursor.execute("ALTER TABLE discussion ADD COLUMN created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP") # Added in V3
-                        cursor.execute("ALTER TABLE message ADD COLUMN model TEXT DEFAULT ''") # Added in V4
-                        cursor.execute("ALTER TABLE message ADD COLUMN personality INT DEFAULT ''") # Added in V4
-
-                        cursor.execute("ALTER TABLE message ADD COLUMN binding TEXT DEFAULT ''") # Added in V5
-                    except :
-                        pass
-            # Upgrade the schema to version 1
-            elif version < 4:
-                print(f"Upgrading schema to version {db_version}...")
-                # Add the 'created_at' column to the 'message' table
-                if message_table_exist:
-                    try:
-                        cursor.execute("ALTER TABLE message ADD COLUMN model TEXT DEFAULT ''") # Added in V4
-                        cursor.execute("ALTER TABLE message ADD COLUMN personality INT DEFAULT ''") # Added in V4
-
-                        cursor.execute("ALTER TABLE message ADD COLUMN binding TEXT DEFAULT ''") # Added in V5
-                    except :
-                        pass                    
-            elif version < 5:
-                print(f"Upgrading schema to version {db_version}...")
-                # Add the 'created_at' column to the 'message' table
-                if message_table_exist:
-                    try:
-                        cursor.execute("ALTER TABLE message ADD COLUMN binding TEXT DEFAULT ''") # Added in V5
-                    except :
-                        pass                    
-            # Update the schema version
-            if not schema_table_exist:
-                cursor.execute(f"INSERT INTO schema_version (id, version) VALUES (1, {db_version})")
-            else:
-                cursor.execute(f"UPDATE schema_version SET version=? WHERE id=?",(db_version,1))
 
             conn.commit()
 
+    def add_missing_columns(self):
+        with sqlite3.connect(self.db_path) as conn:
+            cursor = conn.cursor()
+
+            table_columns = {
+                'discussion': [
+                    'id',
+                    'title',
+                    'created_at'
+                ],
+                'message': [
+                    'id',
+                    'binding',
+                    'model',
+                    'personality',
+                    'sender',
+                    'content',
+                    'type',
+                    'rank',
+                    'parent',
+                    'created_at',
+                    'finished_generating_at',
+                    'discussion_id'
+                ]
+            }
+
+            for table, columns in table_columns.items():
+                cursor.execute(f"PRAGMA table_info({table})")
+                existing_columns = [column[1] for column in cursor.fetchall()]
+
+                for column in columns:
+                    if column not in existing_columns:
+                        if column == 'id':
+                            cursor.execute(f"ALTER TABLE {table} ADD COLUMN {column} INTEGER PRIMARY KEY AUTOINCREMENT")
+                        elif column.endswith('_at'):
+                            cursor.execute(f"ALTER TABLE {table} ADD COLUMN {column} TIMESTAMP")
+                        else:
+                            cursor.execute(f"ALTER TABLE {table} ADD COLUMN {column} TEXT")
+
+            conn.commit()
+
+
     def select(self, query, params=None, fetch_all=True):
         """
         Execute the specified SQL select query on the database,
@@ -275,7 +218,7 @@ class DiscussionsDB:
             discussion_id = row[0]
             discussion_title = row[1]
             discussion = {"id": discussion_id, "title":discussion_title, "messages": []}
-            rows = self.select(f"SELECT sender, content, message_type, rank, parent, binding, model, personality, created_at FROM message WHERE discussion_id=?",(discussion_id,))
+            rows = self.select(f"SELECT sender, content, message_type, rank, parent, binding, model, personality, created_at, finished_generating_at FROM message WHERE discussion_id=?",(discussion_id,))
             for message_row in rows:
                 sender = message_row[1]
                 content = message_row[2]
@@ -286,9 +229,10 @@ class DiscussionsDB:
                 model = message_row[7]
                 personality = message_row[8]
                 created_at = message_row[9]
+                finished_generating_at = message_row[10]
                 
                 discussion["messages"].append(
-                    {"sender": sender, "content": content, "type": content_type, "rank": rank, "parent": parent, "binding": binding, "model":model, "personality":personality, "created_at":created_at}
+                    {"sender": sender, "content": content, "type": content_type, "rank": rank, "parent": parent, "binding": binding, "model":model, "personality":personality, "created_at":created_at, "finished_generating_at":finished_generating_at}
                 )
             discussions.append(discussion)
         return discussions
@@ -306,7 +250,7 @@ class DiscussionsDB:
             discussion_id = row[0]
             discussion_title = row[1]
             discussion = {"id": discussion_id, "title":discussion_title, "messages": []}
-            rows = self.select(f"SELECT sender, content, message_type, rank, parent, binding, model, personality, created_at FROM message WHERE discussion_id=?",(discussion_id,))
+            rows = self.select(f"SELECT sender, content, message_type, rank, parent, binding, model, personality, created_at, finished_generating_at FROM message WHERE discussion_id=?",(discussion_id,))
             for message_row in rows:
                 sender = message_row[1]
                 content = message_row[2]
@@ -317,9 +261,10 @@ class DiscussionsDB:
                 model = message_row[7]
                 personality = message_row[8]
                 created_at = message_row[9]
+                finished_generating_at = message_row[10]
                 
                 discussion["messages"].append(
-                    {"sender": sender, "content": content, "type": content_type, "rank": rank, "parent": parent, "binding": binding, "model":model, "personality":personality, "created_at":created_at}
+                    {"sender": sender, "content": content, "type": content_type, "rank": rank, "parent": parent, "binding": binding, "model":model, "personality":personality, "created_at":created_at, "finished_generating_at": finished_generating_at}
                 )
             discussions.append(discussion)
         return discussions
@@ -346,13 +291,14 @@ class DiscussionsDB:
                 model = message_data.get("model","")
                 personality = message_data.get("personality","")
                 created_at = message_data.get("created_at",datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
+                finished_generating_at = message_data.get("finished_generating_at",datetime.now().strftime('%Y-%m-%d %H:%M:%S'))
                 discussion["messages"].append(
-                    {"sender": sender, "content": content, "type": content_type, "rank": rank, "model": model, "binding": binding, "personality": personality, "created_at": created_at}
+                    {"sender": sender, "content": content, "type": content_type, "rank": rank, "binding": binding, "model": model, "personality": personality, "created_at": created_at, "finished_generating_at": finished_generating_at}
                 )
 
                 # Insert message into the database
-                self.insert("INSERT INTO message (sender, content, type, rank, parent, discussion_id) VALUES (?, ?, ?, ?, ?, ?)",
+                self.insert("INSERT INTO message (sender, content, type, rank, parent, binding, model, personality, created_at, finished_generating_at, discussion_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)",
-                            (sender, content, content_type, rank, parent, discussion_id))
+                            (sender, content, content_type, rank, parent, model, personality, created_at, finished_generating_at, discussion_id))
 
             discussions.append(discussion)
 
@@ -363,7 +309,7 @@ class Discussion:
         self.discussion_id = discussion_id
         self.discussions_db = discussions_db
 
-    def add_message(self, sender, content, message_type=0, rank=0, parent=0, binding="", model ="", personality="", created_at=datetime.now().strftime('%Y-%m-%d %H:%M:%S')):
+    def add_message(self, sender, content, message_type=0, rank=0, parent=0, binding="", model ="", personality="", created_at=datetime.now().strftime('%Y-%m-%d %H:%M:%S'), finished_generating_at=datetime.now().strftime('%Y-%m-%d %H:%M:%S')):
         """Adds a new message to the discussion
 
         Args:
@@ -374,8 +320,8 @@ class Discussion:
             int: The added message id
         """
         message_id = self.discussions_db.insert(
-            "INSERT INTO message (sender, content, type, rank, parent, binding, model, personality, created_at, discussion_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", 
+            "INSERT INTO message (sender, content, type, rank, parent, binding, model, personality, created_at, finished_generating_at, discussion_id) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)", 
-            (sender, content, message_type, rank, parent, binding, model, personality, created_at, self.discussion_id)
+            (sender, content, message_type, rank, parent, binding, model, personality, created_at, finished_generating_at, self.discussion_id)
         )
         return message_id
 
@@ -406,10 +352,10 @@ class Discussion:
             list: List of entries in the format {"id":message id, "sender":sender name, "content":message content, "type":message type, "rank": message rank}
         """
         rows = self.discussions_db.select(
-            "SELECT id, sender, content, type, rank, parent, binding, model, personality, created_at FROM message WHERE discussion_id=?", (self.discussion_id,)
+            "SELECT id, sender, content, type, rank, parent, binding, model, personality, created_at, finished_generating_at FROM message WHERE discussion_id=?", (self.discussion_id,)
         )
 
-        return [{"id": row[0], "sender": row[1], "content": row[2], "type": row[3], "rank": row[4], "parent": row[5], "binding":row[6], "model": row[7], "personality": row[8], "created_at": row[9]} for row in rows]
+        return [{"id": row[0], "sender": row[1], "content": row[2], "type": row[3], "rank": row[4], "parent": row[5], "binding":row[6], "model": row[7], "personality": row[8], "created_at": row[9], "finished_generating_at": row[10]} for row in rows]
 
     def update_message(self, message_id, new_content):
         """Updates the content of a message

app.py

@@ -488,8 +488,11 @@ class LoLLMsWebUI(LoLLMsAPPI):
 
     def apply_settings(self):
         result = self.process.set_config(self.config)
-        print("Set config results:")
+        if result["status"]:
-        print(result)
+            ASCIIColors.success("OK")
+        else:
+            ASCIIColors.error("NOK")
+            
         return jsonify(result)
     
 
@@ -606,9 +609,9 @@ class LoLLMsWebUI(LoLLMsAPPI):
         path = Path("personalities")/lang/category/name
         try:
             shutil.rmtree(path)
-            return jsonify({'status':'success'})
+            return jsonify({'status':True})
         except Exception as ex:
-            return jsonify({'status':'failure','error':str(ex)})
+            return jsonify({'status':False,'error':str(ex)})
 
     def add_endpoint(
         self,