Update modules/pmbl.py

modules/pmbl.py

@@ -1,56 +1,54 @@
-import sqlite3
+import aiosqlite
 from datetime import datetime
 from ctransformers import AutoModelForCausalLM
-from concurrent.futures import ThreadPoolExecutor
+import asyncio
 
 class PMBL:
     def __init__(self, model_path, gpu_layers=50):
         self.model_path = model_path
         self.gpu_layers = gpu_layers
-        self.init_db()
-        self.executor = ThreadPoolExecutor(max_workers=6)
-
-    def init_db(self):
-        conn = sqlite3.connect('chat_history.db')
-        c = conn.cursor()
-        c.execute('''CREATE TABLE IF NOT EXISTS chats
-                     (id INTEGER PRIMARY KEY AUTOINCREMENT,
-                     timestamp TEXT,
-                     prompt TEXT,
-                     response TEXT,
-                     topic TEXT)''')
-        conn.commit()
-        conn.close()
-
-    def get_chat_history(self, mode="full", user_message=""):
-        conn = sqlite3.connect('chat_history.db')
-        c = conn.cursor()
-
-        if mode == "full":
-            c.execute("SELECT prompt, response FROM chats ORDER BY id")
-            history = []
-            for row in c.fetchall():
-                history.append({"role": "user", "content": row[0]})
-                history.append({"role": "PMB", "content": row[1]})
-        else:
-            c.execute("SELECT id, prompt, response FROM chats WHERE topic != 'Untitled'")
-            chats = c.fetchall()
-            relevant_chat_id = self.find_relevant_chat(chats, user_message)
-
-            if relevant_chat_id:
-                c.execute("SELECT prompt, response FROM chats WHERE id = ?", (relevant_chat_id,))
-                row = c.fetchone()
-                history = [
-                    {"role": "user", "content": row[0]},
-                    {"role": "PMB", "content": row[1]}
-                ]
+        self.db_name = 'chat_history.db'
+        self.init_db_lock = asyncio.Lock()
+
+    async def init_db(self):
+        async with self.init_db_lock:
+            async with aiosqlite.connect(self.db_name) as db:
+                await db.execute('''CREATE TABLE IF NOT EXISTS chats
+                                 (id INTEGER PRIMARY KEY AUTOINCREMENT,
+                                 timestamp TEXT,
+                                 prompt TEXT,
+                                 response TEXT,
+                                 topic TEXT)''')
+                await db.commit()
+
+    async def get_chat_history(self, mode="full", user_message=""):
+        await self.init_db()
+        async with aiosqlite.connect(self.db_name) as db:
+            if mode == "full":
+                async with db.execute("SELECT prompt, response FROM chats ORDER BY id") as cursor:
+                    rows = await cursor.fetchall()
+                    history = []
+                    for row in rows:
+                        history.append({"role": "user", "content": row[0]})
+                        history.append({"role": "PMB", "content": row[1]})
             else:
-                history = []
+                async with db.execute("SELECT id, prompt, response FROM chats WHERE topic != 'Untitled'") as cursor:
+                    chats = await cursor.fetchall()
+                    relevant_chat_id = await self.find_relevant_chat(chats, user_message)
+
+                    if relevant_chat_id:
+                        async with db.execute("SELECT prompt, response FROM chats WHERE id = ?", (relevant_chat_id,)) as cursor:
+                            row = await cursor.fetchone()
+                            history = [
+                                {"role": "user", "content": row[0]},
+                                {"role": "PMB", "content": row[1]}
+                            ]
+                    else:
+                        history = []
 
-        conn.close()
         return history
 
-    def find_relevant_chat(self, chats, user_message):
+    async def find_relevant_chat(self, chats, user_message):
         max_score = 0
         relevant_chat_id = None
 
@@ -75,15 +73,14 @@ class PMBL:
 
         return score
 
-    def save_chat_history(self, prompt, response):
-        conn = sqlite3.connect('chat_history.db')
-        c = conn.cursor()
-        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-        c.execute("INSERT INTO chats (timestamp, prompt, response, topic) VALUES (?, ?, ?, 'Untitled')", (timestamp, prompt, response))
-        conn.commit()
-        conn.close()
+    async def save_chat_history(self, prompt, response):
+        await self.init_db()
+        async with aiosqlite.connect(self.db_name) as db:
+            timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
+            await db.execute("INSERT INTO chats (timestamp, prompt, response, topic) VALUES (?, ?, ?, 'Untitled')", (timestamp, prompt, response))
+            await db.commit()
 
-    def generate_response(self, prompt, history, mode):
+    async def generate_response(self, prompt, history, mode):
         history.append({"role": "user", "content": prompt})
 
         formatted_history = ""
@@ -117,7 +114,7 @@ class PMBL:
             response_text += chunk
             yield chunk
 
-        self.save_chat_history(prompt, response_text)
+        await self.save_chat_history(prompt, response_text)
 
     def calculate_context(self, system_prompt, formatted_history):
         system_prompt_tokens = len(system_prompt) // 4
@@ -131,21 +128,19 @@ class PMBL:
         else:
             return context_ceiling
 
-    def sleep_mode(self):
-        conn = sqlite3.connect('chat_history.db')
-        c = conn.cursor()
-        c.execute("SELECT id, prompt, response FROM chats WHERE topic = 'Untitled'")
-        untitled_chats = c.fetchall()
+    async def sleep_mode(self):
+        await self.init_db()
+        async with aiosqlite.connect(self.db_name) as db:
+            async with db.execute("SELECT id, prompt, response FROM chats WHERE topic = 'Untitled'") as cursor:
+                untitled_chats = await cursor.fetchall()
 
-        for chat in untitled_chats:
-            chat_id, prompt, response = chat
-            topic = self.generate_topic(prompt, response)
-            c.execute("UPDATE chats SET topic = ? WHERE id = ?", (topic, chat_id))
-            conn.commit()
-
-        conn.close()
+            for chat in untitled_chats:
+                chat_id, prompt, response = chat
+                topic = await self.generate_topic(prompt, response)
+                await db.execute("UPDATE chats SET topic = ? WHERE id = ?", (topic, chat_id))
+            await db.commit()
 
-    def generate_topic(self, prompt, response):
+    async def generate_topic(self, prompt, response):
         llm = AutoModelForCausalLM.from_pretrained(
             self.model_path,
             model_type="llama",
@@ -162,4 +157,8 @@ class PMBL:
             stop=["\n"]
         )
 
-        return topic.strip()
+        return topic.strip()
+
+    async def close(self):
+        # Implement any cleanup operations here if needed
+        pass