Update app.py

app.py

@@ -9,6 +9,13 @@ import os
 import sys
 from datetime import datetime
 import time
+import logging
+
+# Set up logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s',
+)
 
 # Enable GPU if available
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
@@ -31,6 +38,7 @@ GLOBAL_TOKENIZER = None
 def initialize_model():
     """Initialize model and tokenizer globally"""
     global GLOBAL_MODEL, GLOBAL_TOKENIZER
+    logging.info("Initializing model and tokenizer...")
     st.write("Initializing model and tokenizer...")
     start_time = time.time()
     
@@ -44,11 +52,12 @@ def initialize_model():
     # Set model to evaluation mode
     GLOBAL_MODEL.eval()
     
-    st.write(f"Model initialization took {time.time() - start_time:.2f} seconds")
+    logging.info(f"Model initialization took {time.time() - start_time:.2f} seconds")
 
 def test_db_connection():
     """Test database connection with timeout"""
     try:
+        logging.info("Testing database connection...")
         connection = mysql.connector.connect(
             **DB_CONFIG,
             connect_timeout=10
@@ -60,33 +69,41 @@ def test_db_connection():
             db_name = cursor.fetchone()[0]
             cursor.close()
             connection.close()
+            logging.info(f"Successfully connected to MySQL Server version {db_info} - Database: {db_name}")
             return True, f"Successfully connected to MySQL Server version {db_info}\nDatabase: {db_name}"
     except Error as e:
+        logging.error(f"Error connecting to MySQL database: {e}")
         return False, f"Error connecting to MySQL database: {e}"
     return False, "Unable to establish database connection"
 
 def get_db_connection():
     """Get database connection from pool"""
+    logging.info("Getting database connection from pool...")
     return mysql.connector.connect(**DB_CONFIG)
 
 def execute_query(query):
     """Execute SQL query with timeout and connection pooling"""
+    logging.info(f"Executing query: {query}")
     connection = None
     try:
         connection = get_db_connection()
         cursor = connection.cursor(dictionary=True, buffered=True)
         cursor.execute(query)
         results = cursor.fetchall()
+        logging.info(f"Query executed successfully, retrieved {len(results)} records.")
         return results
     except Error as e:
+        logging.error(f"Error executing query: {e}")
         return f"Error executing query: {e}"
     finally:
         if connection and connection.is_connected():
             cursor.close()
             connection.close()
+            logging.info("Database connection closed.")
 
 def generate_sql(natural_language_query):
     """Generate SQL query with performance optimizations"""
+    logging.info(f"Generating SQL for query: {natural_language_query}")
     try:
         start_time = time.time()
         
@@ -138,18 +155,21 @@ def generate_sql(natural_language_query):
         generated_query = GLOBAL_TOKENIZER.decode(outputs[0], skip_special_tokens=True)
         sql_query = generated_query.split("### SQL Query:")[-1].strip()
         
-        st.write(f"SQL generation took {time.time() - start_time:.2f} seconds")
+        logging.info(f"SQL generation took {time.time() - start_time:.2f} seconds")
         return sql_query
 
     except Exception as e:
+        logging.error(f"Error generating SQL query: {str(e)}")
         return f"Error generating SQL query: {str(e)}"
 
 def format_result(query_result):
     """Format query results efficiently"""
     if isinstance(query_result, str) and "Error" in query_result:
+        logging.warning(f"Query result contains an error: {query_result}")
         return query_result
     
     if not query_result:
+        logging.info("No results found.")
         return "No results found."
     
     # Use list comprehension for better performance
@@ -177,6 +197,7 @@ def main():
     st.write(db_message)
     
     if not db_success:
+        logging.error("Could not connect to the database. Exiting.")
         st.write("Could not connect to the database. Exiting.")
         return
     
@@ -202,6 +223,7 @@ def main():
             st.write("Human-Readable Response:")
             st.text(formatted_result)
         else:
+            logging.warning("User did not enter a query.")
             st.write("Please enter a query.")
 
 if __name__ == "__main__":