update space

app.py

@@ -6,8 +6,9 @@ import logging
 import dotenv
 import gradio as gr
 import numpy as np
-from pymongo import MongoClient
+from pymongo import MongoClient, errors
 from datetime import datetime
+from werkzeug.security import generate_password_hash, check_password_hash
 from utils import PoultryFarmBot, llama3_response
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
@@ -20,26 +21,41 @@ logger = logging.getLogger(__name__)
 
 # MongoDB Setup for logging and audit
 MONGO_URI = os.getenv("MONGO_URI")
-logger.info("Connecting to MongoDB.")
-client = MongoClient(MONGO_URI)
-db = client.poultry_farm  # Connect to the 'poultry_farm' database
-enquiries_collection = db.enquiries  # Collection to store farmer enquiries
-users_collection = db.users  # Collection to store user credentials
-logs_collection = db.logs  # Collection to store application logs
+if not MONGO_URI:
+    logger.error("MONGO_URI is not set in the environment variables.")
+    raise ValueError("MONGO_URI environment variable is required but not set.")
+
+try:
+    logger.info("Connecting to MongoDB.")
+    client = MongoClient(MONGO_URI, serverSelectionTimeoutMS=5000)  # Timeout after 5 seconds
+    client.server_info()  # Trigger exception if cannot connect
+    db = client.poultry_farm  # Connect to the 'poultry_farm' database
+    enquiries_collection = db.enquiries  # Collection to store farmer enquiries
+    users_collection = db.users  # Collection to store user credentials
+    logs_collection = db.logs  # Collection to store application logs
+except errors.ServerSelectionTimeoutError as e:
+    logger.error(f"Failed to connect to MongoDB: {e}")
+    raise ConnectionError("Could not connect to MongoDB. Please check the MONGO_URI and ensure the database is running.")
 
 def log_to_db(level, message):
-    log_entry = {
-        "level": level,
-        "message": message,
-        "timestamp": datetime.utcnow()
-    }
-    logs_collection.insert_one(log_entry)
+    try:
+        log_entry = {
+            "level": level,
+            "message": message,
+            "timestamp": datetime.utcnow()
+        }
+        logs_collection.insert_one(log_entry)
+    except Exception as e:
+        logger.error(f"Failed to log to database: {e}")
 
 # Override logger methods to also log to MongoDB
 class MongoHandler(logging.Handler):
     def emit(self, record):
         log_entry = self.format(record)
-        log_to_db(record.levelname, log_entry)
+        try:
+            log_to_db(record.levelname, log_entry)
+        except Exception as e:
+            logger.error(f"Failed to emit log to MongoDB: {e}")
 
 mongo_handler = MongoHandler()
 mongo_handler.setLevel(logging.INFO)
@@ -59,17 +75,21 @@ else:
 logger.info("Initializing PoultryFarmBot instance.")
 bot = PoultryFarmBot(db)
 
-# Load Llama 3.1 model and tokenizer for text generation
-logger.info("Loading Llama 3.1 model and tokenizer.")
-model_name = "meta-llama/Llama-2-7b-hf"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
-
-# Set the padding token to EOS token or add a new padding token
-if tokenizer.pad_token is None:
-    logger.info("Adding padding token to tokenizer.")
-    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
-    model.resize_token_embeddings(len(tokenizer))
+# Load Llama 3.2 model and tokenizer for text generation
+try:
+    logger.info("Loading Llama 3.2 model and tokenizer.")
+    model_name = "meta-llama/Llama-3.2-3B"
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    model = AutoModelForCausalLM.from_pretrained(model_name)
+
+    # Set the padding token to EOS token or add a new padding token
+    if tokenizer.pad_token is None:
+        logger.info("Adding padding token to tokenizer.")
+        tokenizer.add_special_tokens({'pad_token': '[PAD]'})
+        model.resize_token_embeddings(len(tokenizer))
+except Exception as e:
+    logger.error(f"Failed to load Llama 3.2 model or tokenizer: {e}")
+    raise RuntimeError("Could not load the Llama 3.2 model or tokenizer. Please check the configuration.")
 
 def chatbot_response(image, text, username, password):
     """
@@ -84,30 +104,34 @@ def chatbot_response(image, text, username, password):
     Returns:
         str: Response generated by the chatbot.
     """
-    user = bot.authenticate_user(username, password)
-    if not user:
-        return "Authentication failed. Please check your username and password."
-
-    user_id = user['_id']
-
-    # If an image is provided, diagnose the disease
-    if image is not None:
-        logger.info("Image input detected. Proceeding with disease diagnosis.")
-        diagnosis, name, status, recom = bot.diagnose_disease(image)
-        if name and status and recom:
-            logger.info("Diagnosis complete.")
-            bot.log_enquiry("image", "Image Enquiry", diagnosis, user_id)
-            return diagnosis
+    try:
+        user = bot.authenticate_user(username, password)
+        if not user:
+            return "Authentication failed. Please check your username and password."
+
+        user_id = user['_id']
+
+        # If an image is provided, diagnose the disease
+        if image is not None:
+            logger.info("Image input detected. Proceeding with disease diagnosis.")
+            diagnosis, name, status, recom = bot.diagnose_disease(image)
+            if name and status and recom:
+                logger.info("Diagnosis complete.")
+                bot.log_enquiry("image", "Image Enquiry", diagnosis, user_id)
+                return diagnosis
+            else:
+                logger.warning("Diagnosis incomplete.")
+                bot.log_enquiry("image", "Image Enquiry", diagnosis, user_id)
+                return diagnosis
         else:
-            logger.warning("Diagnosis incomplete.")
-            bot.log_enquiry("image", "Image Enquiry", diagnosis, user_id)
-            return diagnosis
-    else:
-        # Generate a response using Llama 3.2 for general text input
-        logger.info("Text input detected. Generating response.")
-        response = llama3_response(text, tokenizer, model)
-        bot.log_enquiry("text", text, response, user_id)
-        return response
+            # Generate a response using Llama 3.2 for general text input
+            logger.info("Text input detected. Generating response.")
+            response = llama3_response(text, tokenizer, model)
+            bot.log_enquiry("text", text, response, user_id)
+            return response
+    except Exception as e:
+        logger.error(f"Error during chatbot response generation: {e}")
+        return "An error occurred while processing your request. Please try again later."
 
 # Gradio interface
 def build_gradio_interface():
@@ -117,66 +141,74 @@ def build_gradio_interface():
     Returns:
         gr.Blocks: Gradio Blocks object representing the chatbot interface.
     """
-    logger.info("Building Gradio interface.")
-    with gr.Blocks(theme=gr.themes.Base()):
-        gr.Markdown("# 🐔 Poultry Management Chatbot")
-        gr.Markdown("Welcome! This chatbot helps you manage your poultry with ease. You can upload an image for disease diagnosis or ask any questions about poultry management.")
-
-        chat_history = gr.Chatbot()
-        with gr.Row():
-            with gr.Column(scale=1):
-                fecal_image = gr.Image(
-                    label="Upload Image of Poultry Feces (Optional)",
-                    type="numpy",
-                    elem_id="image-upload",
-                    show_label=True,
-                )
-            with gr.Column(scale=2):
-                user_input = gr.Textbox(
-                    label="Ask a question",
-                    placeholder="Ask about poultry management...",
-                    lines=3,
-                    elem_id="user-input",
-                )
-                username = gr.Textbox(
-                    label="Username",
-                    placeholder="Enter your username",
-                    lines=1,
-                    elem_id="username-input",
-                )
-                password = gr.Textbox(
-                    label="Password",
-                    placeholder="Enter your password",
-                    type="password",
-                    lines=1,
-                    elem_id="password-input",
-                )
-
-        output_box = gr.Textbox(
-            label="Response",
-            placeholder="Response will appear here...",
-            interactive=False,
-            lines=10,
-            elem_id="output-box",
-        )
-
-        submit_button = gr.Button(
-            "Submit",
-            variant="primary",
-            elem_id="submit-button"
-        )
-        # Connect the submit button to the chatbot response function
-        submit_button.click(
-            fn=chatbot_response,
-            inputs=[fecal_image, user_input, username, password],
-            outputs=[output_box]
-        )
-    logger.info("Gradio interface built successfully.")
-    return chatbot_interface
+    try:
+        logger.info("Building Gradio interface.")
+        with gr.Blocks(theme=gr.themes.Base()):
+            gr.Markdown("# 🐔 Poultry Management Chatbot")
+            gr.Markdown("Welcome! This chatbot helps you manage your poultry with ease. You can upload an image for disease diagnosis or ask any questions about poultry management.")
+
+            chat_history = gr.Chatbot()
+            with gr.Row():
+                with gr.Column(scale=1):
+                    fecal_image = gr.Image(
+                        label="Upload Image of Poultry Feces (Optional)",
+                        type="numpy",
+                        elem_id="image-upload",
+                        show_label=True,
+                    )
+                with gr.Column(scale=2):
+                    user_input = gr.Textbox(
+                        label="Ask a question",
+                        placeholder="Ask about poultry management...",
+                        lines=3,
+                        elem_id="user-input",
+                    )
+                    username = gr.Textbox(
+                        label="Username",
+                        placeholder="Enter your username",
+                        lines=1,
+                        elem_id="username-input",
+                    )
+                    password = gr.Textbox(
+                        label="Password",
+                        placeholder="Enter your password",
+                        type="password",
+                        lines=1,
+                        elem_id="password-input",
+                    )
+
+            output_box = gr.Textbox(
+                label="Response",
+                placeholder="Response will appear here...",
+                interactive=False,
+                lines=10,
+                elem_id="output-box",
+            )
+
+            submit_button = gr.Button(
+                "Submit",
+                variant="primary",
+                elem_id="submit-button"
+            )
+            # Connect the submit button to the chatbot response function
+            submit_button.click(
+                fn=chatbot_response,
+                inputs=[fecal_image, user_input, username, password],
+                outputs=[output_box]
+            )
+        logger.info("Gradio interface built successfully.")
+        return chatbot_interface
+    except Exception as e:
+        logger.error(f"Error building Gradio interface: {e}")
+        raise RuntimeError("Could not build the Gradio interface. Please check the configuration.")
 
 # Launch the Gradio interface
 if __name__ == "__main__":
-    logger.info("Launching Gradio interface.")
-    interface = build_gradio_interface()
-    # Launch the interface with queuing enabled for concurrent requests
-    interface.queue().launch(debug=True, share=True)
+    try:
+        logger.info("Launching Gradio interface.")
+        interface = build_gradio_interface()
+        # Launch the interface with queuing enabled for concurrent requests
+        interface.queue().launch(debug=True, share=True)
+    except Exception as e:
+        logger.error(f"Failed to launch Gradio interface: {e}")
+        raise RuntimeError("Could not launch the Gradio interface. Please check the application setup.")