Update app.py

app.py

@@ -21,6 +21,7 @@ def load_system_prompt():
 
 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')
@@ -28,10 +29,10 @@ 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')
 
+# OpenAI and MongoDB clients
 openai_client = OpenAI(api_key=openai_api_key, base_url=api_url)
 mongo_client = MongoClient(mongoURI)
 db = mongo_client['Scarlett']
@@ -40,6 +41,7 @@ chat_collection = db['chats']
 local_chat_history = OrderedDict()
 MAX_LOCAL_USERS = 5
 
+# Functions for MongoDB-based chat history storage and retrieval
 def get_history_from_mongo(user_id):
     result = chat_collection.find_one({"user_id": user_id})
     return result.get("messages", []) if result else []
@@ -77,28 +79,45 @@ def update_chat_history(user_id, role, content):
         local_chat_history.popitem(last=False)
     store_message_in_mongo(user_id, role, content)
 
-def encode_local_image(image_path):
-    im = Image.open(image_path)
-    buffered = BytesIO()
-    im.save(buffered, format="PNG")
-    image_bytes = buffered.getvalue()
-    image_base64 = base64.b64encode(image_bytes).decode('ascii')
-    return image_base64
+# Fixing image encoding
+def encode_local_image(image_file):
+    try:
+        im = Image.open(image_file)
+        buffered = BytesIO()
+        im.save(buffered, format="PNG")
+        image_bytes = buffered.getvalue()
+        image_base64 = base64.b64encode(image_bytes).decode('ascii')
+        return image_base64
+    except Exception as e:
+        print(f"Error encoding image: {e}")
+        return None
 
+# Image description function, calling external inference model
 def inference_calling_idefics(image_path, question=""):
     system_prompt = os.getenv('USER_PROMPT')
     model_id = model2
     client = InferenceClient(model=model_id)
-    image_base64 = describe_image(image_path)
+    
+    # Use the fixed `encode_local_image` to encode the image
+    image_base64 = encode_local_image(image_path)
+    
+    if not image_base64:
+        return "Error: Invalid image or unable to encode image."
+
     image_info = f"data:image/png;base64,{image_base64}"
     prompt = f"{system_prompt}\n![]({image_info})\n{question}\n\n"
-    response = client.text_generation(
-        prompt, 
-        max_new_tokens=512, 
-        do_sample=True,
-        temperature=0.2
-    )
-    return response
+    
+    try:
+        response = client.text_generation(
+            prompt, 
+            max_new_tokens=512, 
+            do_sample=True,
+            temperature=0.2
+        )
+        return response
+    except Exception as e:
+        print(f"Error in inference call: {e}")
+        return "Error while processing the image."
 
 def describe_image(image_path, question=""):
     try:
@@ -108,12 +127,14 @@ def describe_image(image_path, question=""):
         print(e)
         return "Error while seeing the image."
 
+# Telegram bot setup
 client = TelegramClient('bot', api_id, api_hash).start(bot_token=bot_token)
 
 async def get_bot_id():
     me = await client.get_me()
     return me.id
 
+# OpenAI completion handler
 async def get_completion(event, user_id, prompt):
     async with client.action(event.chat_id, 'typing'):
         await asyncio.sleep(3)
@@ -139,12 +160,13 @@ async def get_completion(event, user_id, prompt):
                 if chunk.choices[0].delta.content is not None:
                     message += chunk.choices[0].delta.content
         except Exception as e:
-            message = f"Whoops!"
+            message = "Whoops!"
             print(e)
         update_chat_history(user_id, "user", prompt)
         update_chat_history(user_id, "assistant", message)
         return message
 
+# Telegram bot commands
 @client.on(events.NewMessage(pattern='/start'))
 async def start(event):
     await event.respond("Hello!")
@@ -173,8 +195,9 @@ async def handle_message(event):
         user_message = event.raw_text
         if event.photo:
             photo = await event.download_media()
-            image_description = describe_image(photo, user_message)
-            user_message += f"\n\nI sent you an image. Content of the image: {image_description}"
+            if photo:
+                image_description = describe_image(photo, user_message)
+                user_message += f"\n\nI sent you an image. Content of the image: {image_description}"
         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)
@@ -183,6 +206,7 @@ async def handle_message(event):
         print(f"An error occurred: {e}")
         await event.respond("Whoopsie!")
 
+# Gradio interface
 def launch_gradio():
     welcome_message = """
     # Meet Scarlett!
@@ -197,12 +221,13 @@ def launch_gradio():
         """)
     demo.launch(show_api=False)
 
+# Keep-alive functionality for the bot
 def keep_alive():
     ping_client = OpenAI(api_key=ping_key, base_url=api_url)
     while True:
         try:
             messages = [
-                {"role": "system", "content": "Repeat what i say."},
+                {"role": "system", "content": "Repeat what I say."},
                 {"role": "user", "content": "Repeat: 'Ping success'"}
             ]
             request = ping_client.chat.completions.create(
@@ -217,6 +242,7 @@ def keep_alive():
             print(f"Keep-alive request failed: {e}")
         time.sleep(1800)
 
+# Main execution
 if __name__ == "__main__":
     threading.Thread(target=keep_alive).start()
     threading.Thread(target=launch_gradio).start()