Update bot_telegram.py

bot_telegram.py

@@ -1,18 +1,33 @@
+import os
 import logging
+import asyncio
 import requests
 from telegram import Update
 from telegram.ext import Application, CommandHandler, MessageHandler, filters, CallbackContext
+import aiohttp
 import io
 
 # Configure logging
-logging.basicConfig(format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", level=logging.INFO)
-logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+
+# Load environment variables from Hugging Face Secrets
+BOT_TOKEN = os.getenv("BOT_TOKEN")
+BASE_URL = os.getenv("BASE_URL")
+API_USERNAME = os.getenv("API_USERNAME")
+API_PASSWORD = os.getenv("API_PASSWORD")
+
+# Set the secret password for authentication
+SECRET_PASSWORD = "secure123"  # Change this to your desired password
+
+# Dictionary to store authenticated users
+AUTHENTICATED_USERS = set()
+AWAITING_PASSWORD = set()
 
 class TelegramBot:
-    """A Telegram bot that communicates with an AI model via API."""
+    """A Telegram bot with password-based authentication."""
 
     def __init__(self, bot_token, base_url, username, password):
-        """Initialize the bot with Telegram API token, API credentials, and endpoints."""
+        """Initialize the bot with Telegram API token, API credentials, and authentication."""
         self.bot_token = bot_token
         self.base_url = base_url
         self.username = username
@@ -28,98 +43,86 @@ class TelegramBot:
         self.app = Application.builder().token(self.bot_token).build()
         self.setup_handlers()
 
-        # Authenticate on startup
+        # Authenticate with API
+        logging.info("Authenticating with API...")
         self.authenticate()
 
     def authenticate(self):
         """Authenticate with the API and retrieve an access token."""
-        payload = {
-            "username": self.username,
-            "password": self.password,
-            "grant_type": "password"
-        }
-
-        headers = {
-            "Content-Type": "application/json",
-            "accept": "application/json"
-        }
+        payload = {"username": self.username, "password": self.password}
+        headers = {"Content-Type": "application/json", "accept": "application/json"}
 
         try:
             response = requests.post(self.login_url, headers=headers, json=payload)
 
             if response.status_code == 200:
                 self.auth_token = response.json().get("access_token")
-                logger.info("Successfully authenticated with API!")
+                logging.info("Successfully authenticated with API")
             else:
-                logger.error(f"Authentication failed: {response.status_code} - {response.text}")
+                logging.error(f"Authentication failed: {response.status_code} - {response.text}")
 
         except Exception as e:
-            logger.error(f"Authentication Error: {e}")
+            logging.error(f"Authentication Error: {e}")
 
     async def start_command(self, update: Update, context: CallbackContext):
-        """Handles the /start command."""
-        await update.message.reply_text(
-            "Hello! You can:\n"
-            "1. Send me any question as text\n"
-            "2. Send me an Excel file with questions (must have a 'question' column in 'rfp' sheet)"
-        )
-
-    async def handle_document(self, update: Update, context: CallbackContext):
-        """Handles Excel file uploads."""
-        if not self.auth_token:
-            self.authenticate()
-
-        if not self.auth_token:
-            await update.message.reply_text("Authentication failed. Please try again later.")
+        """Handles the /start command and asks for a password if the user is not authenticated."""
+        user_id = update.message.from_user.id
+
+        if user_id in AUTHENTICATED_USERS:
+            await update.message.reply_text(
+                "✅ You are already authenticated!\n\n"
+                "You can:\n"
+                "1. Send me any question as text\n"
+                "2. Send me an Excel file with questions (must have a 'question' column in 'rfp' sheet)\n\n"
+                "Note: Excel files must contain no more than 50 questions."
+            )
+        else:
+            AWAITING_PASSWORD.add(user_id)
+            await update.message.reply_text("🔑 Please enter the secret password to access the bot.")
+
+    async def handle_message(self, update: Update, context: CallbackContext):
+        """Handles all incoming messages."""
+        user_id = update.message.from_user.id
+        user_message = update.message.text.strip()
+
+        # If user is waiting to enter a password, validate it
+        if user_id in AWAITING_PASSWORD:
+            await self.check_password(update, context)
             return
 
-        document = update.message.document
-        if not document.file_name.endswith(('.xls', '.xlsx')):
-            await update.message.reply_text("Please send only Excel files (.xls or .xlsx)")
-            return
-
-        try:
-            # Download the file from Telegram
-            file = await context.bot.get_file(document.file_id)
-            file_content = await file.download_as_bytearray()
-
-            # Prepare the file for upload to our API
-            files = {
-                'file': (document.file_name, io.BytesIO(file_content), 'application/vnd.openxmlformats-officedocument.spreadsheetml.sheet')
-            }
-            headers = {
-                "Authorization": f"Bearer {self.auth_token}",
-                "accept": "application/json"
-            }
-
-            # Send processing message
-            status_message = await update.message.reply_text("Processing your Excel file... This may take a few minutes.")
-
-            # Send to our API
-            response = requests.post(self.excel_url, headers=headers, files=files)
+        # If user is authenticated, process AI request
+        if user_id in AUTHENTICATED_USERS:
+            await self.chat_with_ai(update, context)
+        else:
+            await update.message.reply_text("❌ You are not authenticated. Please enter the password first.")
+
+    async def check_password(self, update: Update, context: CallbackContext):
+        """Checks if the password is correct and authenticates the user."""
+        user_id = update.message.from_user.id
+        user_message = update.message.text.strip()
+
+        if user_message == SECRET_PASSWORD:
+            AUTHENTICATED_USERS.add(user_id)
+            AWAITING_PASSWORD.discard(user_id)
+            logging.info(f"User {user_id} authenticated successfully.")
+            await update.message.reply_text(
+                "✅ Authentication successful!\n\n"
+                "You can:\n"
+                "1. Send me any question as text\n"
+                "2. Send me an Excel file with questions (must have a 'question' column in 'rfp' sheet)\n\n"
+                "Note: Excel files must contain no more than 50 questions."
+            )
+        else:
+            await update.message.reply_text("❌ Wrong password. Try again.")
 
-            if response.status_code == 200:
-                # Get the processed Excel file from the response
-                result = response.json()
-                file_content = result.get("file_content")
-                filename = result.get("filename")
-
-                # Send the processed Excel file back to Telegram
-                await context.bot.send_document(
-                    chat_id=update.effective_chat.id,
-                    document=io.BytesIO(bytes(file_content)),
-                    filename=filename
-                )
-                await status_message.delete()
-            else:
-                error_message = f"Error processing file: {response.text}"
-                await status_message.edit_text(error_message)
+    async def chat_with_ai(self, update: Update, context: CallbackContext):
+        """Handles messages and sends them to the AI API."""
+        user_id = update.message.from_user.id
 
-        except Exception as e:
-            await update.message.reply_text(f"Error processing file: {str(e)}")
+        if user_id not in AUTHENTICATED_USERS:
+            await update.message.reply_text("❌ You are not authenticated. Please enter the password first.")
+            return
 
-    async def chat_with_ai(self, update: Update, context: CallbackContext):
-        """Handles text messages and sends them to the AI API."""
         if not self.auth_token:
             self.authenticate()
 
@@ -131,19 +134,24 @@ class TelegramBot:
 
         headers = {
             "Authorization": f"Bearer {self.auth_token}",
-            "Content-Type": "application/x-www-form-urlencoded",
             "accept": "application/json"
         }
 
-        payload = {"question": user_message}
+        json_payload = {"question": user_message}
+        form_payload = {"question": user_message}
 
         try:
-            response = requests.post(self.ai_url, headers=headers, data=payload)
+            logging.info(f"Sending payload as JSON: {json_payload}")
+            response = requests.post(self.ai_url, headers={**headers, "Content-Type": "application/json"}, json=json_payload)
+
+            if response.status_code == 422:
+                logging.warning("JSON format rejected. Retrying with form-data...")
+                response = requests.post(self.ai_url, headers={**headers, "Content-Type": "application/x-www-form-urlencoded"}, data=form_payload)
 
             if response.status_code == 200:
                 bot_reply = response.json().get("answer", "I didn't understand that.")
             elif response.status_code == 401:
-                logger.warning("Authorization expired. Re-authenticating...")
+                logging.warning("Authorization expired. Re-authenticating...")
                 self.authenticate()
                 await self.chat_with_ai(update, context)
                 return
@@ -155,23 +163,97 @@ class TelegramBot:
 
         await update.message.reply_text(bot_reply)
 
+    async def handle_excel(self, update: Update, context: CallbackContext):
+        """Handles Excel file uploads."""
+        user_id = update.message.from_user.id
+
+        if user_id not in AUTHENTICATED_USERS:
+            await update.message.reply_text("❌ You are not authenticated. Please enter the password first.")
+            return
+
+        if not self.auth_token:
+            self.authenticate()
+
+        if not self.auth_token:
+            await update.message.reply_text("Authentication failed. Please try again later.")
+            return
+
+        document = update.message.document
+        if not document.file_name.endswith(('.xls', '.xlsx')):
+            await update.message.reply_text("Please send only Excel files (.xls or .xlsx)")
+            return
+
+        try:
+            # Send initial processing message
+            processing_msg = await update.message.reply_text("📊 Processing your Excel file... This may take a few minutes.")
+
+            # Get file from Telegram
+            file = await context.bot.get_file(document.file_id)
+            file_bytes = await file.download_as_bytearray()
+
+            headers = {
+                "Authorization": f"Bearer {self.auth_token}",
+                "accept": "application/json"
+            }
+
+            async with aiohttp.ClientSession() as session:
+                # Create form data with the file
+                form_data = aiohttp.FormData()
+                form_data.add_field('file',
+                                  file_bytes,
+                                  filename=document.file_name,
+                                  content_type='application/vnd.openxmlformats-officedocument.spreadsheetml.sheet')
+
+                # Send file to API
+                async with session.post(self.excel_url, headers=headers, data=form_data) as response:
+                    if response.status == 200:
+                        response_json = await response.json()
+                        file_content = response_json.get('file_content')
+                        filename = response_json.get('filename', 'rfp_responses.xlsx')
+
+                        if not file_content:
+                            await processing_msg.edit_text("❌ Error: No file content received from server")
+                            return
+
+                        # Convert base64 string to bytes and send back to Telegram
+                        file_bytes = bytes(file_content)
+                        await context.bot.send_document(
+                            chat_id=update.effective_chat.id,
+                            document=io.BytesIO(file_bytes),
+                            filename=filename,
+                            caption="✅ Here's your processed Excel file with answers!"
+                        )
+                        await processing_msg.delete()
+
+                    elif response.status == 401:
+                        logging.warning("Authorization expired. Re-authenticating...")
+                        self.authenticate()
+                        await self.handle_excel(update, context)
+                    else:
+                        error_text = await response.text()
+                        await processing_msg.edit_text(f"❌ Error processing Excel file: {error_text}")
+
+        except Exception as e:
+            logging.error(f"Error handling Excel file: {e}")
+            await update.message.reply_text(f"❌ Error processing Excel file: {str(e)}")
+
     def setup_handlers(self):
         """Set up Telegram command and message handlers."""
         self.app.add_handler(CommandHandler("start", self.start_command))
-        self.app.add_handler(MessageHandler(filters.DOCUMENT, self.handle_document))
-        self.app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, self.chat_with_ai))
+        self.app.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, self.handle_message))
+        self.app.add_handler(MessageHandler(filters.Document.FileExtension("xlsx") | filters.Document.FileExtension("xls"), self.handle_excel))
 
     def run(self):
         """Start the bot and listen for messages."""
-        logger.info("Bot is now listening for messages...")
+        logging.info("Starting Telegram bot...")
         self.app.run_polling()
 
-# --- Initialize and Run Bot ---
 if __name__ == "__main__":
     bot = TelegramBot(
-        bot_token="7218319671:AAF6XyrAJrsent3fMOjQQus9XnZJxRH99BQ",
-        base_url="https://22ab-86-19-236-226.ngrok-free.app",
-        username="admin",
-        password="admin"
+        bot_token=BOT_TOKEN,
+        base_url=BASE_URL,
+        username=API_USERNAME,
+        password=API_PASSWORD
     )
     bot.run()
+