Update app.py

app.py

@@ -1,189 +1,224 @@
 import os
 import time
 import threading
+import base64
+from io import BytesIO
+from collections import OrderedDict
+from datetime import datetime
+
 import requests
+from openai import OpenAI
 from telethon import TelegramClient, events
-import gradio as gr
-import asyncio
-from io import BytesIO
+from PIL import Image
+from huggingface_hub import InferenceClient
+from transformers import AutoProcessor, AutoTokenizer
+import pymongo
 from pymongo import MongoClient
 
-# MongoDB connection setup
-mongo_uri = os.getenv('MONGO_URI')
-mongo_client = MongoClient(mongo_uri)
-db = mongo_client['Scarlett']
-user_histories_collection = db['chats']
-
+# Load system prompt from a file
 def load_system_prompt():
     with open('prompt.txt', 'r') as file:
         return file.read()
 
 system_prompt = load_system_prompt()
 
+# Environment variables
 api_id = os.getenv('api_id')
 api_hash = os.getenv('api_hash')
 bot_token = os.getenv('bot_token')
-cloudflare_api_key = os.getenv('glhf')
-cloudflare_base_url = os.getenv('yolo')
+openai_api_key = os.getenv('glhf')
+ping_key = os.getenv('bolo')
+api_url = os.getenv('yolo')
+model = os.getenv('model')
+model1 = os.getenv('model1')
+model2 = os.getenv('model2')
+mongoURI = os.getenv('MONGO_URI')
+
+# Initialize OpenAI and MongoDB clients
+openai_client = OpenAI(api_key=openai_api_key, base_url=api_url)
+mongo_client = MongoClient(mongoURI)
+db = mongo_client['Scarlett']
+chat_collection = db['chats']
+
+# Initialize Hugging Face models for image processing
+idefics_processor = AutoProcessor.from_pretrained(model1)
+idefics_client = InferenceClient(model2)
+tokenizer = AutoTokenizer.from_pretrained(model1)
+
+# Local cache for up to 5 users
+local_chat_history = OrderedDict()
+MAX_LOCAL_USERS = 5
+
+# Retrieve chat history from MongoDB
+def get_history_from_mongo(user_id):
+    result = chat_collection.find_one({"user_id": user_id})
+    return result.get("messages", []) if result else []
+
+# Store message in MongoDB (limit to last 99 messages)
+def store_message_in_mongo(user_id, role, content):
+    chat_collection.update_one(
+        {"user_id": user_id},
+        {
+            "$push": {
+                "messages": {
+                    "$each": [{"role": role, "content": content}],
+                    "$slice": -99
+                }
+            }
+        },
+        upsert=True
+    )
+
+# Get chat history from local cache or MongoDB
+def get_chat_history(user_id):
+    if user_id in local_chat_history:
+        local_chat_history.move_to_end(user_id)  # Mark as most recently used
+        return local_chat_history[user_id]
+
+    # Load from MongoDB if not in local cache
+    history = get_history_from_mongo(user_id)
+    local_chat_history[user_id] = history
+
+    if len(local_chat_history) > MAX_LOCAL_USERS:
+        local_chat_history.popitem(last=False)  # Remove LRU user
+    
+    return history
 
-client = TelegramClient('bot', api_id, api_hash).start(bot_token=bot_token)
+# Update chat history (both local and MongoDB)
+def update_chat_history(user_id, role, content):
+    if user_id not in local_chat_history:
+        local_chat_history[user_id] = get_history_from_mongo(user_id)
 
-class MongoDBHistory:
-    def __init__(self, user_id, size=99):
-        self.user_id = str(user_id)
-        self.size = size
+    local_chat_history[user_id].append({"role": role, "content": content})
+    local_chat_history.move_to_end(user_id)
+
+    if len(local_chat_history) > MAX_LOCAL_USERS:
+        local_chat_history.popitem(last=False)
+
+    store_message_in_mongo(user_id, role, content)
+
+# Encode image to base64
+def encode_local_image(image):
+    pil_image = Image.open(image)
+    buffer = BytesIO()
+    pil_image.save(buffer, format="JPEG")
+    return f"data:image/jpeg;base64,{base64.b64encode(buffer.getvalue()).decode('utf-8')}"
+
+# Describe image using the model
+def describe_image(image_path):
+    image_string = encode_local_image(image_path)
 
-    def add(self, role: str, content: str):
-        user_history = user_histories_collection.find_one({'user_id': self.user_id})
-        if user_history:
-            messages = user_history['history']
-        else:
-            messages = []
-        
-        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):
-        user_history = user_histories_collection.find_one({'user_id': self.user_id})
-        if user_history:
-            return user_history['history']
-        return []
-
-    def reset(self):
-        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:
-        user_histories[user_id] = MongoDBHistory(user_id)
-    return user_histories[user_id]
-
-def fetch_image_as_u8(image_binary):
-    return list(image_binary)
-
-async def get_completion(prompt: str, user_id, image_u8=None) -> str:
-    user_history = get_user_history(user_id)
-
-    # Prepare the message content
     messages = [
-        {"role": "system", "content": system_prompt},
-        *user_history.get_history(),
-        {"role": "user", "content": prompt}
+        {"role": "user", "content": [{"type": "image"}, {"type": "text", "text": os.getenv('USER_PROMPT')}]}
     ]
 
-    try:
-        # Prepare the payload with additional parameters
-        data = {
-            "messages": messages,
-            "temperature": 0.5,  # Adjust creativity
-            "top_p": 0.9,  # Nucleus sampling parameter
-            "max_tokens": 512,  # Maximum token limit
-            "frequency_penalty": 1.6,  # Penalize repeated tokens
-            "presence_penalty": 1.9  # Encourage new topics
-        }
-
-        # If there is an image, include it as an array of unsigned 8-bit integers (u8)
-        if image_u8:
-            print(image_u8)
-            data["image"] = image_u8  # Send as array of integers
-
-        headers = {
-            "Authorization": f"Bearer {cloudflare_api_key}"
-        }
-
-        # Send request to Cloudflare API
-        request = requests.post(
-            url=cloudflare_base_url,
-            json=data,
-            headers=headers
-        )
-        response = request.json()
-        print(response)
-
-        if response['success'] == True:
-            message = response["result"]["response"]
-        else:
-            message = f"Error: {response['errors'][0]['message']}"
+    prompt_with_template = idefics_processor.apply_chat_template(
+        messages, add_generation_prompt=True, chat_template=os.getenv('CHAT_TEMPLATE')
+    )
 
-    except Exception as e:
-        message = f"Error: {str(e)}"
-        print(e)
-    
-    # Add the user prompt and assistant response to the history
-    user_history.add("user", prompt)
-    user_history.add("assistant", message)
-    return message
+    prompt_with_images = prompt_with_template.replace("<image>", f"![]({image_string})")
+
+    payload = {
+        "inputs": prompt_with_images,
+        "parameters": {"return_full_text": False, "max_new_tokens": 2048},
+    }
+
+    response = idefics_client.post(json=payload).decode()
+    return response
+
+# Telegram bot client
+client = TelegramClient('bot', api_id, api_hash).start(bot_token=bot_token)
+
+# Async function to get OpenAI completion
+async def get_completion(event, user_id, prompt):
+    async with client.action(event.chat_id, 'typing'):
+        history = get_chat_history(user_id)
+        messages = [
+            {"role": "system", "content": system_prompt},
+            *history,
+            {"role": "user", "content": prompt},
+        ]
 
+        try:
+            response = openai_client.chat.completions.create(
+                model=model,
+                messages=messages,
+                max_tokens=512,
+                temperature=0.5,
+                top_p=1.0,
+                frequency_penalty=0.9,
+                presence_penalty=0.9,
+            )
+            message = response.choices[0].message.content
+        except Exception as e:
+            message = f"Error: {str(e)}"
+            print(e)
+
+        update_chat_history(user_id, "user", prompt)  # Update history
+        update_chat_history(user_id, "assistant", message)  # Update assistant's response
+        return message
+
+# Telegram bot events
 @client.on(events.NewMessage(pattern='/start'))
 async def start(event):
-    await event.respond("Hello!")
+    await event.respond("Hello! I am your chat assistant.")
 
 @client.on(events.NewMessage(pattern='/help'))
 async def help(event):
-    await event.respond("Here is how I can help you:\n/start - Start\n/help - Show this message\n/reset - Reset chat history")
+    await event.respond("Here is how I can help you:\n/start - Start the bot\n/help - Get help\n/reset - Reset chat history")
 
 @client.on(events.NewMessage(pattern='/reset'))
 async def reset(event):
-    user_history = get_user_history(event.sender_id)
-    user_history.reset()
+    user_id = event.chat_id
+    chat_collection.delete_one({"user_id": user_id})  # Reset MongoDB chat history for the user
+    if user_id in local_chat_history:
+        del local_chat_history[user_id]  # Remove from local cache if present
     await event.respond("History has been reset.")
 
 @client.on(events.NewMessage)
 async def handle_message(event):
-    if event.raw_text.startswith('/start') or event.raw_text.startswith('/help') or event.raw_text.startswith('/reset'):
-        return
     try:
-        async with client.action(event.chat_id, 'typing'):
-            await asyncio.sleep(3)
+        user_id = event.chat_id  # Use chat_id to distinguish between users
 
-        user_id = event.sender_id
-        user_message = event.raw_text
-
-        # Check if the message contains an image
-        image_u8 = None
         if event.photo:
-            image = await event.download_media(file=BytesIO())
-            image_u8 = fetch_image_as_u8(image.getvalue())  # Convert to u8
+            photo = await event.download_media()
+            image_description = describe_image(photo)
+            user_message = f"{event.raw_text}\n\nContent of the image: {image_description}"
+        else:
+            user_message = event.raw_text
 
-        response = await get_completion(user_message, user_id, image_u8)
-        await event.respond(response)
+        if user_message.startswith('/start') or user_message.startswith('/help') or user_message.startswith('/reset'):
+            return
 
+        response = await get_completion(event, user_id, user_message)
+        await event.respond(response)
     except Exception as e:
         print(f"An error occurred: {e}")
         await event.respond("Whoopsie 🤭")
 
-def launch_gradio():
-    welcome_message = """
-    # Meet Scarlett!
-    - Scarlett is a girl residing in Telegram
-    - She can chat with you without any censorship.
-    - Head over to https://t.me/Scarlett_Rogue_isNotBot
-    """
-    with gr.Blocks() as demo:
-        gr.Markdown(f"""
-        # Follow me: https://example.com
-        {welcome_message}
-        """)
-    demo.launch(show_api=False)
-
+# Keep-alive function to keep the bot running
 def keep_alive():
+    ping_client = OpenAI(api_key=ping_key, base_url=api_url)
     while True:
         try:
-            requests.get("https://huggingface.co/spaces/rbn2008k/Scarlett/")
+            messages = [
+                {"role": "system", "content": "Be a helpful assistant."},
+                {"role": "user", "content": "Hello"}
+            ]
+            request = ping_client.chat.completions.create(
+                model=model,
+                messages=messages,
+                max_tokens=10,
+                temperature=0.6,
+                top_p=0.9,
+                frequency_penalty=0.2,
+                presence_penalty=0.6,
+            )
+            print(request)
         except Exception as e:
             print(f"Keep-alive request failed: {e}")
-        time.sleep(180)
+        time.sleep(1800)  # Ping every 30 minutes
 
 if __name__ == "__main__":
-    threading.Thread(target=launch_gradio).start()
     threading.Thread(target=keep_alive).start()
     client.run_until_disconnected()