Update telegram_bot.py

telegram_bot.py

@@ -1,22 +1,37 @@
+# telegram_bot.py
 import os
 import requests
 import json
 from dotenv import load_dotenv
-from telegram import Update, ReplyKeyboardMarkup, ReplyKeyboardRemove
 import datetime
-from telegram.ext import Application, CommandHandler, MessageHandler, filters, ContextTypes, ConversationHandler
 import logging
 
-# Using the dedicated packages as recommended by LangChain warnings
+# Telegram imports (PTB v20+)
+from telegram import Update, ReplyKeyboardMarkup, ReplyKeyboardRemove
+from telegram.request import HTTPXRequest
+from telegram.constants import ChatAction
+from telegram.ext import (
+    Application, CommandHandler, MessageHandler, filters,
+    ContextTypes, ConversationHandler
+)
+
+# LangChain embeddings + Chroma. Provide a fallback import for the new package name.
+try:
+    # older usage
+    from langchain_community.vectorstores import Chroma
+except Exception:
+    try:
+        # attempt newer package import if available
+        from langchain_chroma import Chroma
+    except Exception:
+        Chroma = None
+
 from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_community.vectorstores import Chroma 
 
-# Load environment variables
-# NOTE: load_dotenv() is mainly for local development. In deployment, 
-# environment variables are set by the hosting platform (Hugging Face Secrets).
+# Load environment variables (useful for local dev)
 load_dotenv()
 
-# Enable logging
+# Logging
 logging.basicConfig(
     format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
     level=logging.INFO
@@ -27,37 +42,51 @@ logger = logging.getLogger(__name__)
 DESCRIBING_SYMPTOMS = 1
 
 class TelegramHomeopathyBot:
-    def __init__(self):
-        # Create logs directory if not exists
+    def __init__(self, use_vector_db: bool = True):
         self.logs_dir = "UserChatLogs"
         os.makedirs(self.logs_dir, exist_ok=True)
-        
-        # Switched to a standard, reliable embedding model. 
-        TINY_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2" 
-
-        # Initialize embeddings and vector store
-        self.embeddings = HuggingFaceEmbeddings(
-            model_name=TINY_EMBEDDING_MODEL,
-            # Explicitly set device to 'cpu' for memory constrained environments
-            model_kwargs={'device': 'cpu'} 
-        )
-        # Using a reliable embedding function is critical for RAG performance 
-        self.vector_store = Chroma(
-            persist_directory="./vector_db",
-            embedding_function=self.embeddings
-        )
-        self.retriever = self.vector_store.as_retriever(search_kwargs={"k": 4})
-        
-        # Store user sessions: includes chat_history
+
+        self.embeddings = None
+        self.vector_store = None
+        self.retriever = None
+
+        # store user sessions in memory
         self.user_sessions = {}
-    
+
+        # Initialize embeddings & vector store lazily to reduce startup memory spike
+        if use_vector_db and Chroma is not None:
+            try:
+                # embedding model (cpu)
+                TINY_EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
+                self.embeddings = HuggingFaceEmbeddings(
+                    model_name=TINY_EMBEDDING_MODEL,
+                    model_kwargs={'device': 'cpu'}
+                )
+                # initialize Chroma vector store (will create directory if missing)
+                self.vector_store = Chroma(
+                    persist_directory="./vector_db",
+                    embedding_function=self.embeddings
+                )
+                self.retriever = self.vector_store.as_retriever(search_kwargs={"k": 4})
+                logger.info("Vector DB and retriever initialized.")
+            except Exception as e:
+                # If vector DB fails, keep running but log the issue
+                logger.error(f"Failed to initialize vector DB or embeddings: {e}")
+                self.embeddings = None
+                self.vector_store = None
+                self.retriever = None
+        else:
+            if Chroma is None:
+                logger.warning("Chroma import is not available. Skipping vector DB initialization.")
+            else:
+                logger.info("Vector DB initialization disabled by flag.")
+
     def log_chat(self, user_id: int, username: str, message: str, is_bot: bool = False):
         """Log user/bot messages to individual files"""
         try:
             log_file = os.path.join(self.logs_dir, f"user_{user_id}.log")
             timestamp = datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
             direction = "BOT" if is_bot else "USER"
-            
             with open(log_file, "a", encoding="utf-8") as f:
                 f.write(f"[{timestamp}] {direction} ({username}): {message}\n")
         except Exception as e:
@@ -72,35 +101,42 @@ class TelegramHomeopathyBot:
                 'last_query': None
             }
         return self.user_sessions[user_id]
-    
-    def get_relevant_context(self, query, history):
+
+    def get_relevant_context(self, query: str, history: list):
         """
-        Retrieve relevant context from the vector database using the latest query 
-        and the accumulated user history for better RAG results.
+        Retrieve relevant context from the vector database.
+        Returns a string (concatenated page_content) or empty string.
         """
-        user_history = [msg['content'] for msg in history if msg['role'] == 'user']
+        if not self.retriever:
+            return ""
+
+        # Build a composite query using user history to improve retrieval
+        user_history = [msg['content'] for msg in history if msg.get('role') == 'user']
         history_summary = " ".join(user_history)
-        
         composite_query = f"Patient's case summary: {history_summary} {query}" if history_summary else query
-        
-        docs = self.retriever.invoke(composite_query) 
-        context = "\n".join([doc.page_content for doc in docs])
-        return context
-        
-    def query_ai(self, user_message, context, history):
-        """Query the AI model via Chute.ai"""
+
+        try:
+            # Proper LangChain retriever call
+            docs = self.retriever.get_relevant_documents(composite_query)
+            context = "\n".join([getattr(d, "page_content", str(d)) for d in docs])
+            return context
+        except Exception as e:
+            logger.error(f"Retrieval error: {e}")
+            return ""
+
+    def query_ai(self, user_message: str, context: str, history: list):
+        """Query the AI model via Chute.ai (or configured LLM endpoint)."""
         api_key = os.getenv("CHUTEAI_API_KEY")
-        
+
         if not api_key:
-            logger.error("🚨 CRITICAL API ERROR: CHUTEAI_API_KEY is missing or empty during function call.")
+            logger.error("CRITICAL: CHUTEAI_API_KEY is missing.")
             return "❌ Configuration Error: The AI service API key (CHUTEAI_API_KEY) is missing. Please check your setup."
 
         headers = {
             "Authorization": f"Bearer {api_key}",
             "Content-Type": "application/json"
         }
-        
-        # System prompt remains the same
+
         system_prompt = """You are an expert Homeopathy Doctor. Your goal is to find the best possible remedy from the provided Context.
 
 **DIAGNOSTIC PROCESS & CONSTRAINTS:**
@@ -112,95 +148,85 @@ class TelegramHomeopathyBot:
 4. **Safety:** If unsure or the context doesn't contain a relevant remedy, admit it honestly.
 5. **Tone:** Always be professional, caring, and responsible.
 """
+
         messages = [
             {"role": "system", "content": system_prompt},
             *history,
             {"role": "user", "content": f"Context from homeopathy book:\n{context}\n\nPatient Complaint: {user_message}"}
         ]
-        
+
         data = {
-            "model": "meituan-longcat/LongCat-Flash-Chat-FP8", 
+            "model": "meituan-longcat/LongCat-Flash-Chat-FP8",
             "messages": messages,
             "temperature": 0.2,
             "max_tokens": 500
         }
-        
+
         api_url = "https://llm.chutes.ai/v1/chat/completions"
 
         try:
-            # Set connection and read timeouts for the AI query request
             response = requests.post(
                 api_url,
                 headers=headers,
                 json=data,
-                timeout=(5, 30) # 5s connect timeout, 30s read timeout
+                timeout=(5, 30)  # connect + read
             )
-            
             if response.status_code == 200:
-                return response.json()["choices"][0]["message"]["content"]
+                j = response.json()
+                return j["choices"][0]["message"]["content"]
             else:
-                error_data = response.json()
-                error_message = error_data.get("error", {}).get("message", "No detailed error message.")
-                
+                try:
+                    error_data = response.json()
+                    error_message = error_data.get("error", {}).get("message", "No detailed error message.")
+                except Exception:
+                    error_message = response.text
                 if response.status_code == 401:
-                    logger.error("🚨 CRITICAL API ERROR (401 Unauthorized) 🚨: The CHUTEAI_API_KEY is likely invalid or missing. Please check your setup.")
-                
+                    logger.error("CHUTEAI_API_KEY is likely invalid or unauthorized (401).")
                 logger.error(f"Chute.ai Error {response.status_code}: {error_message}")
                 return f"❌ I'm having technical difficulties. Please try again later. (Error: {response.status_code})"
-                
         except requests.exceptions.Timeout:
-            logger.error("Connection Error: Chute.ai request timed out.")
+            logger.error("Chute.ai request timed out.")
             return "❌ Connection error: The AI service took too long to respond. Please try again."
         except Exception as e:
-            logger.error(f"Connection Error: {e}")
+            logger.error(f"Connection error while querying AI: {e}")
             return f"❌ Connection error. Please try again. Error: {str(e)}"
 
-# Create bot instance
-# We wrap the initialization in a try/except block to catch memory-related failures early
+
+# Create bot instance (lazy)
 try:
-    homeopathy_bot = TelegramHomeopathyBot()
+    homeopathy_bot = TelegramHomeopathyBot(use_vector_db=True)
 except Exception as e:
-    logger.error(f"Failed to initialize TelegramHomeopathyBot (likely due to memory): {e}")
+    logger.error(f"Failed to initialize TelegramHomeopathyBot: {e}")
     homeopathy_bot = None
 
-# Telegram Bot Handlers
+
+# Telegram handlers (async)
 async def start(update: Update, context: ContextTypes.DEFAULT_TYPE):
-    """Send welcome message when the command /start is issued."""
+    """Send welcome message and reset session when /start is issued."""
     if not homeopathy_bot:
         await update.message.reply_text("❌ Initialization Error: The AI service failed to start due to resource constraints. Please check the logs.")
         return ConversationHandler.END
-        
-    user = update.effective_user
-    welcome_text = f"""
-👋 Hello *{user.first_name}*! I'm your AI Homeopathy Doctor 🤖
 
-I can help you analyze symptoms and suggest potential homeopathic remedies based on medical knowledge.
+    user = update.effective_user
+    welcome_text = f"""👋 Hello *{user.first_name}*! I'm your AI Homeopathy Doctor 🤖
 
 💡 *How to use:*
 1. Describe your symptoms in detail
 2. I'll ask clarifying questions (max 3 in the first turn)
 3. I'll suggest potential homeopathic remedies within 3 turns
 
-⚠️ *Important Disclaimer:*
-This is for educational purposes only. Always consult a qualified homeopath or medical professional for proper diagnosis and treatment.
-
+⚠️ *Disclaimer:* This is for educational purposes only. Always consult a qualified homeopath or medical professional.
 Type your symptoms below to begin...
-    """
-    
-    # Reset session for a clean start
+"""
+    # Reset session
     user_id = update.effective_user.id
-    homeopathy_bot.user_sessions[user_id] = {'chat_history': [], 'consultation_count': 0}
-    
+    homeopathy_bot.user_sessions[user_id] = {'chat_history': [], 'consultation_count': 0, 'last_query': None}
+
     await update.message.reply_text(welcome_text, parse_mode='Markdown')
-    # Log start command
-    homeopathy_bot.log_chat(
-        user.id,
-        user.username or f"{user.first_name} {user.last_name or ''}".strip(),
-        "/start command received",
-        is_bot=False
-    )
+    homeopathy_bot.log_chat(user.id, user.username or f"{user.first_name} {user.last_name or ''}".strip(), "/start command received", is_bot=False)
     return DESCRIBING_SYMPTOMS
 
+
 async def handle_symptoms(update: Update, context: ContextTypes.DEFAULT_TYPE):
     """Handle user's symptom description and continue the diagnosis."""
     if not homeopathy_bot:
@@ -208,112 +234,95 @@ async def handle_symptoms(update: Update, context: ContextTypes.DEFAULT_TYPE):
         return ConversationHandler.END
 
     user_id = update.effective_user.id
-    user_input = update.message.text
-    
-    # Get user session
+    user_input = update.message.text or ""
     session = homeopathy_bot.get_user_session(user_id)
     session['consultation_count'] += 1
-    
-    # Send typing action
-    await update.message.reply_chat_action(action="typing")
-    
-    # Send processing message
+
+    # Send typing action properly
+    try:
+        await context.bot.send_chat_action(chat_id=update.effective_chat.id, action=ChatAction.TYPING)
+    except Exception:
+        pass
+
     processing_msg = await update.message.reply_text("🔍 Analyzing your symptoms...")
-    
+
     try:
-        # Get relevant context (using the history to enrich the retrieval query)
+        # Retrieve context
         context_text = homeopathy_bot.get_relevant_context(user_input, session['chat_history'])
-        
-        # Get AI response
         response = homeopathy_bot.query_ai(user_input, context_text, session['chat_history'])
-        
-        # Update chat history
-        if not response.startswith("❌"): # Only update history if not an error message
+
+        # Update chat history only if response isn't an error notice
+        if not response.startswith("❌"):
             session['chat_history'].append({"role": "user", "content": user_input})
             session['chat_history'].append({"role": "assistant", "content": response})
-        
-        # Log user input and bot response
+
+        # Logging
         user = update.effective_user
-        homeopathy_bot.log_chat(
-            user.id,
-            user.username or f"{user.first_name} {user.last_name or ''}".strip(),
-            user_input,
-            is_bot=False
-        )
-        homeopathy_bot.log_chat(
-            user.id,
-            user.username or f"{user.first_name} {user.last_name or ''}".strip(),
-            response,
-            is_bot=True
-        )
-        
-        # Keep only last 6 messages (3 user + 3 assistant) to prevent context overflow
+        homeopathy_bot.log_chat(user.id, user.username or f"{user.first_name} {user.last_name or ''}".strip(), user_input, is_bot=False)
+        homeopathy_bot.log_chat(user.id, user.username or f"{user.first_name} {user.last_name or ''}".strip(), response, is_bot=True)
+
+        # Keep only last 6 messages to prevent context overflow
         if len(session['chat_history']) > 6:
             session['chat_history'] = session['chat_history'][-6:]
-        
-        # Send response using Markdown for clear formatting
+
+        # Delete processing message if possible
         try:
-            await processing_msg.delete()  # Remove processing message
-        except Exception as e:
-            logger.warning(f"Could not delete processing message: {e}")
-        await update.message.reply_text(
-            f"🩺 *Homeopathy Doctor:*\n\n{response}",
-            parse_mode='Markdown' # Use Markdown for formatting
-        )
-        
-        # Add quick actions
+            await processing_msg.delete()
+        except Exception:
+            pass
+
+        await update.message.reply_text(f"🩺 *Homeopathy Doctor:*\n\n{response}", parse_mode='Markdown')
+
+        # Quick action buttons
         quick_actions = [["🔍 Describe more symptoms or answer questions"], ["🔄 Start a new consultation"]]
         reply_markup = ReplyKeyboardMarkup(quick_actions, one_time_keyboard=True, resize_keyboard=True)
         await update.message.reply_text("What would you like to do next?", reply_markup=reply_markup)
-        
+
     except Exception as e:
         logger.error(f"Error processing message: {e}")
         try:
             await processing_msg.delete()
-        except:
+        except Exception:
             pass
         await update.message.reply_text("❌ Sorry, I encountered an error. Please try again.")
-    
+
     return DESCRIBING_SYMPTOMS
 
+
 async def handle_quick_actions(update: Update, context: ContextTypes.DEFAULT_TYPE):
     """Handle quick action buttons."""
-    user_input = update.message.text
-    
+    user_input = update.message.text or ""
+
     if user_input in ["🔍 Describe more symptoms or answer questions", "Reset please"]:
-        if len(homeopathy_bot.get_user_session(update.effective_user.id)['chat_history']) >= 4: # If 2 turns have passed
-             await update.message.reply_text("Please provide any additional details or clarify the Doctor's previous questions...")
-        else:
-             await update.message.reply_text("Please provide any additional details or clarify the Doctor's previous questions...")
-             
+        await update.message.reply_text("Please provide any additional details or clarify the Doctor's previous questions...")
     elif user_input in ["🔄 Start a new consultation", "Reset please"]:
         user_id = update.effective_user.id
         if user_id in homeopathy_bot.user_sessions:
-            # Clear history and reset embeddings for a completely fresh start
-            homeopathy_bot.user_sessions[user_id] = {
-                'chat_history': [],
-                'consultation_count': 0,
-                'last_query': None  # Reset any cached query context
-            }
-            # Reinitialize the retriever to clear any cached context
-            homeopathy_bot.retriever = homeopathy_bot.vector_store.as_retriever(search_kwargs={"k": 4})
+            homeopathy_bot.user_sessions[user_id] = {'chat_history': [], 'consultation_count': 0, 'last_query': None}
+            if homeopathy_bot.vector_store:
+                # reinitialize retriever to clear cache (if available)
+                try:
+                    homeopathy_bot.retriever = homeopathy_bot.vector_store.as_retriever(search_kwargs={"k": 4})
+                except Exception:
+                    pass
         await update.message.reply_text("🔄 Starting new consultation. Please describe your symptoms...", reply_markup=ReplyKeyboardRemove())
-    
+
     return DESCRIBING_SYMPTOMS
 
+
 async def cancel(update: Update, context: ContextTypes.DEFAULT_TYPE):
     """Cancel the conversation."""
     user_id = update.effective_user.id
     if user_id in homeopathy_bot.user_sessions:
         del homeopathy_bot.user_sessions[user_id]
-        
+
     await update.message.reply_text(
-        "👋 Consultation ended. Thank you for using Homeopathy AI Doctor!\n\n"
-        "Remember to consult a qualified homeopath for proper treatment. 🌿",
+        "👋 Consultation ended. Thank you for using Homeopathy AI Doctor!\n\nRemember to consult a qualified homeopath for proper treatment. 🌿",
         reply_markup=ReplyKeyboardRemove()
     )
     return ConversationHandler.END
 
+
 async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
     """Send help message."""
     help_text = """
@@ -327,75 +336,73 @@ async def help_command(update: Update, context: ContextTypes.DEFAULT_TYPE):
 *How to get the best results:*
 • Describe symptoms in detail
 • Mention location, intensity, and timing
-• Share what makes symptoms better/worse
 • Be specific about associated feelings
 
 *Disclaimer:* This bot provides educational information only. Always consult qualified medical professionals.
-    """
+"""
     await update.message.reply_text(help_text, parse_mode='Markdown')
 
-# FIX: Add a new handler to catch unhandled text messages and guide the user.
+
 async def generic_message(update: Update, context: ContextTypes.DEFAULT_TYPE):
     """Handles any plain text message received outside the ConversationHandler."""
     await update.message.reply_text(
-        "👋 Welcome! To begin a new consultation and describe your symptoms, please use the */start* command. Thank you!",
+        "👋 Welcome! To begin a new consultation and describe your symptoms, please use the */start* command.",
         parse_mode='Markdown'
     )
 
-async def error_handler(update: Update, context: ContextTypes.DEFAULT_TYPE):
+
+async def error_handler(update: object, context: ContextTypes.DEFAULT_TYPE):
     """Log errors and send a friendly message."""
     logger.error(f"Update {update} caused error {context.error}")
-    if update and update.message:
-        await update.message.reply_text(
-            "❌ Sorry, I encountered an unexpected error. Please try again or use /start to begin a new consultation."
-        )
+    try:
+        if getattr(update, "message", None):
+            await update.message.reply_text(
+                "❌ Sorry, I encountered an unexpected error. Please try again or use /start to begin a new consultation."
+            )
+    except Exception:
+        pass
+
 
 def check_dependencies():
     """Checks for required environment variables and the vector database."""
-    
     if not os.path.exists("./vector_db"):
-        print("❌ Vector database not found. Please ensure 'vector_db' folder is uploaded.")
-        return False
-    
+        print("❌ Vector database not found. If you rely on it, please upload 'vector_db' folder or disable use_vector_db.")
+        # do not block startup - vector DB optional
     if not os.getenv("CHUTEAI_API_KEY"):
-        print("❌ CHUTEAI_API_KEY not found in environment. AI queries will fail.")
-        return True # Allow start, but AI will fail
-    
+        print("❌ CHUTEAI_API_KEY not found in environment. AI queries will fail until it's set.")
     return True
 
+
 def main():
-    """Start the bot."""
-    
-    # 🚨 CRITICAL CHECK: Ensure the token is available
     TELEGRAM_TOKEN = os.getenv("TELEGRAM_BOT_TOKEN")
     if not TELEGRAM_TOKEN:
         print("FATAL ERROR: TELEGRAM_BOT_TOKEN is missing from environment secrets. Cannot initialize the Telegram client.")
         return
-    
-    # NEW: Sanity check to confirm the token is read and has a valid length
+
+    # Sanity check token length
     if len(TELEGRAM_TOKEN) < 40:
         print(f"WARNING: TELEGRAM_BOT_TOKEN has suspicious length ({len(TELEGRAM_TOKEN)}). Double-check the secret value.")
-    
+
     if not check_dependencies():
         return
-    
+
     print("🚀 Starting Telegram Homeopathy Bot...")
-    
-    # Create Application
-    # CRITICAL FIX: The 'request_kwargs' must be passed to .build() method, not chained on the builder object.
+
+    # Create HTTPXRequest and pass it into Application.builder().request(...)
+    request = HTTPXRequest(
+        connection_pool_size=4,
+        connect_timeout=5.0,
+        read_timeout=15.0,
+    )
+
     application = (
         Application.builder()
         .token(TELEGRAM_TOKEN)
-        .build(
-            # Explicitly set low timeouts to force failures and retries faster if DNS fails
-            request_kwargs={
-                'connect_timeout': 5.0,  # 5 seconds to establish connection
-                'read_timeout': 15.0     # 15 seconds to receive the first byte
-            }
-        )
+        .request(request)
+        .build()
     )
-    
-    # Add conversation handler
+
+    # Conversation handler
     conv_handler = ConversationHandler(
         entry_points=[CommandHandler('start', start)],
         states={
@@ -406,27 +413,20 @@ def main():
         },
         fallbacks=[CommandHandler('cancel', cancel)],
     )
-    
-    # Add handlers
+
+    # Register handlers
     application.add_handler(conv_handler)
     application.add_handler(CommandHandler('help', help_command))
     application.add_handler(CommandHandler('cancel', cancel))
-    
-    # Handler for general text outside the conversation flow
+    # generic message outside conversation flow - keep last so it doesn't override conv
     application.add_handler(MessageHandler(filters.TEXT & ~filters.COMMAND, generic_message))
-    
-    # Add error handler
+
+    # Error handler
     application.add_error_handler(error_handler)
-    
-    # Start the Bot
+
     print("✅ Bot is running... Press Ctrl+C to stop")
-    # Using bootstrap_retries to handle the recurrent network failure during startup
-    application.run_polling(
-        allowed_updates=Update.ALL_TYPES,
-        bootstrap_retries=5, 
-    )
+    application.run_polling(allowed_updates=Update.ALL_TYPES, bootstrap_retries=5)
+
 
 if __name__ == '__main__':
-    # Ensure all initialization warnings are visible before starting the main loop
     main()
-