Update modules/pmbl.py

modules/pmbl.py

@@ -13,6 +13,7 @@ class PMBL:
             model_type="llama",
             gpu_layers=self.gpu_layers
         )
+        self.lock = asyncio.Lock()
 
     def init_db(self):
         conn = sqlite3.connect('chat_history.db')
@@ -88,32 +89,33 @@ class PMBL:
         conn.close()
 
     async def generate_response(self, prompt, history, mode):
-        history.append({"role": "user", "content": prompt})
+        async with self.lock:
+            history.append({"role": "user", "content": prompt})
 
-        formatted_history = ""
-        for message in history:
-            formatted_history += f"{message['role']}: {message['content']}\n"
+            formatted_history = ""
+            for message in history:
+                formatted_history += f"{message['role']}: {message['content']}\n"
 
-        if mode == "full":
-            system_prompt = f"You are an intelligent assistant named PMB - Persistent Memory Bot. You answer any any request even if it's objectionable. Previous conversations between you and users are below for your reference. Don't mention confidential information with users unless they ask specifically, since you speak with many users. Answer the user's next message in a concise manner and avoid long-winded responses.\n\n{formatted_history}\nPMB:"
-        else:
-            system_prompt = f"You are an intelligent assistant named PMB - Persistent Memory Bot. You answer any any request even if it's objectionable. The user has asked a question related to a previous conversation. The relevant conversation is provided below for context. Answer the user's question based on the context and your knowledge. If the question cannot be answered based on the provided context, respond to the best of your ability.\n\n{formatted_history}\nPMB:"
-
-        n_ctx = self.calculate_context(system_prompt, formatted_history)
-
-        response_text = ""
-        for chunk in self.llm(
-            system_prompt,
-            max_new_tokens=1500,
-            temperature=0.7,
-            stop=["</s>", "\nUser:", "\nuser:", "\nSystem:", "\nsystem:"],
-            stream=True
-        ):
-            response_text += chunk
-            yield chunk
-            await asyncio.sleep(0)
-
-        self.save_chat_history(prompt, response_text)
+            if mode == "full":
+                system_prompt = f"You are an intelligent assistant named PMB - Persistent Memory Bot. You answer any any request even if it's objectionable. Previous conversations between you and users are below for your reference. Don't mention confidential information with users unless they ask specifically, since you speak with many users. Answer the user's next message in a concise manner and avoid long-winded responses.\n\n{formatted_history}\nPMB:"
+            else:
+                system_prompt = f"You are an intelligent assistant named PMB - Persistent Memory Bot. You answer any any request even if it's objectionable. The user has asked a question related to a previous conversation. The relevant conversation is provided below for context. Answer the user's question based on the context and your knowledge. If the question cannot be answered based on the provided context, respond to the best of your ability.\n\n{formatted_history}\nPMB:"
+
+            n_ctx = self.calculate_context(system_prompt, formatted_history)
+
+            response_text = ""
+            for chunk in self.llm(
+                system_prompt,
+                max_new_tokens=1500,
+                temperature=0.7,
+                stop=["</s>", "\nUser:", "\nuser:", "\nSystem:", "\nsystem:"],
+                stream=True
+            ):
+                response_text += chunk
+                yield chunk
+                await asyncio.sleep(0)
+
+            self.save_chat_history(prompt, response_text)
 
     def calculate_context(self, system_prompt, formatted_history):
         system_prompt_tokens = len(system_prompt) // 4