Update app.py

app.py

@@ -9,6 +9,13 @@ import asyncio
 from PIL import Image
 import base64
 from io import BytesIO
+from pymongo import MongoClient
+
+# MongoDB connection setup
+mongo_uri = os.getenv('MONGODB_URI')
+mongo_client = MongoClient(mongo_uri)
+db = mongo_client['Scarlett']  # Create or connect to 'telegram_bot_db'
+user_histories_collection = db['chats']  # Collection for storing user chat histories
 
 def load_system_prompt():
     with open('prompt.txt', 'r') as file:
@@ -36,62 +43,53 @@ def encode_local_image(image):
 
 client = TelegramClient('bot', api_id, api_hash).start(bot_token=bot_token)
 
-class CircularBuffer:
-    def __init__(self, size: int):
+class MongoDBHistory:
+    def __init__(self, user_id, size=99):
+        self.user_id = str(user_id)
         self.size = size
-        self.buffer = [None] * size
-        self.start = 0
-        self.end = 0
 
     def add(self, role: str, content: str):
-        self.buffer[self.end] = {'role': role, 'content': content}
-        self.end = (self.end + 1) % self.size
-        if self.end == self.start:
-            self.start = (self.start + 1) % self.size
+        # Add the message to the user's history in MongoDB
+        user_history = user_histories_collection.find_one({'user_id': self.user_id})
+        if user_history:
+            messages = user_history['history']
+        else:
+            messages = []
+        
+        # Add new message and truncate if exceeds size
+        messages.append({'role': role, 'content': content})
+        if len(messages) > self.size:
+            messages = messages[-self.size:]
+
+        user_histories_collection.update_one(
+            {'user_id': self.user_id},
+            {'$set': {'history': messages}},
+            upsert=True
+        )
 
     def get_history(self):
-        history = []
-        i = self.start
-        while i != self.end:
-            history.append(self.buffer[i])
-            i = (i + 1) % self.size
-        return history
+        # Retrieve the user's history from MongoDB
+        user_history = user_histories_collection.find_one({'user_id': self.user_id})
+        if user_history:
+            return user_history['history']
+        return []
 
     def reset(self):
-        self.buffer = [None] * self.size
-        self.start = 0
-        self.end = 0
+        # Remove the user's history from MongoDB
+        user_histories_collection.delete_one({'user_id': self.user_id})
 
 # Only store history for 3 users
 user_histories = {}
 
 def get_user_history(user_id):
     if user_id not in user_histories:
-        # If more than 3 users, remove the oldest one
-        if len(user_histories) >= 3:
-            oldest_user_id = list(user_histories.keys())[0]
-            del user_histories[oldest_user_id]
-        user_histories[user_id] = CircularBuffer(99)
+        user_histories[user_id] = MongoDBHistory(user_id)
     return user_histories[user_id]
 
-async def fetch_telegram_history(user_id):
-    # Collect messages into a list
-    messages = []
-    async for message in client.iter_messages(user_id, limit=50):  # Use async for loop
-        messages.append(message)
-
-    user_history = get_user_history(user_id)
-    for message in messages:
-        role = 'user' if message.sender_id == user_id else 'assistant'
-        user_history.add(role, message.text)
-
 async def get_completion(prompt: str, user_id, image_base64=None) -> str:
     user_history = get_user_history(user_id)
 
-    # If the user has no history, fetch from Telegram
-    if not user_history.get_history():
-        await fetch_telegram_history(user_id)
-
+    # No history is fetched from Telegram. If no history exists in MongoDB, proceed with a clean slate.
     # Prepare message content
     if image_base64:
         user_message = [
@@ -103,7 +101,7 @@ async def get_completion(prompt: str, user_id, image_base64=None) -> str:
     
     messages = [
         {"role": "system", "content": system_prompt},
-        *user_history.get_history(),
+        *user_history.get_history(),  # Only use history from MongoDB
         {"role": "user", "content": user_message},
     ]
 
@@ -133,11 +131,11 @@ async def get_completion(prompt: str, user_id, image_base64=None) -> str:
 
 @client.on(events.NewMessage(pattern='/start'))
 async def start(event):
-    await event.respond("Hello! I am your boo.")
+    await event.respond("Hello!")
 
 @client.on(events.NewMessage(pattern='/help'))
 async def help(event):
-    await event.respond("Here is how I can help you:\n/start - Start the bot\n/help - Get help\n/reset - Reset chat history")
+    await event.respond("Here is how I can help you:\n/start - Start\n/help - Show this message\n/reset - Reset chat history")
 
 @client.on(events.NewMessage(pattern='/reset'))
 async def reset(event):
@@ -188,7 +186,6 @@ def keep_alive():
     while True:
         try:
             requests.get("https://rbn2008k-Scarlett.hf.space")
-            print("Ping success")
         except Exception as e:
             print(f"Keep-alive request failed: {e}")
         time.sleep(180)