app.py

import streamlit as st
from transformers import AutoModelForCausalLM, AutoTokenizer
import torch
from functools import lru_cache
import json
import mysql.connector
from mysql.connector import Error
import os
import sys
from datetime import datetime
import time
import logging
import threading

# 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")

# Database configuration
DB_CONFIG = {
    'host': 'sql12.freemysqlhosting.net',
    'database': 'sql12740625',
    'user': 'sql12740625',
    'password': 'QGG9kdrE4g',
    'port': 3306,
    'pool_size': 5,
    'pool_reset_session': True
}

# Global variables for model and tokenizer
GLOBAL_MODEL = None
GLOBAL_TOKENIZER = None
db_connection_status = False

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()
    
    model_name_sql = "premai-io/prem-1B-SQL"
    GLOBAL_TOKENIZER = AutoTokenizer.from_pretrained(model_name_sql)
    GLOBAL_MODEL = AutoModelForCausalLM.from_pretrained(
        model_name_sql,
        torch_dtype=torch.float32,  # Use float32 for CPU
    ).to(device)
    
    # Set model to evaluation mode
    GLOBAL_MODEL.eval()
    
    logging.info(f"Model initialization took {time.time() - start_time:.2f} seconds")

def test_db_connection():
    """Test database connection with timeout"""
    global db_connection_status
    try:
        logging.info("Testing database connection...")
        connection = mysql.connector.connect(
            **DB_CONFIG,
            connect_timeout=10
        )
        if connection.is_connected():
            db_info = connection.get_server_info()
            cursor = connection.cursor()
            cursor.execute("SELECT DATABASE();")
            db_name = cursor.fetchone()[0]
            cursor.close()
            connection.close()
            db_connection_status = True
            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:
        db_connection_status = False
        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()
        
        schema_info = """
        CREATE TABLE sales (
          pizza_id DECIMAL(8,2) PRIMARY KEY,
          order_id DECIMAL(8,2),
          pizza_name_id VARCHAR(14),
          quantity DECIMAL(4,2),
          order_date DATE,
          order_time VARCHAR(8),
          unit_price DECIMAL(5,2),
          total_price DECIMAL(5,2),
          pizza_size VARCHAR(3),
          pizza_category VARCHAR(7),
          pizza_ingredients VARCHAR(97),
          pizza_name VARCHAR(42)
        );
        """
        
        prompt = f"""### Task: Generate a SQL query to answer the following question.
        ### Database Schema:
        {schema_info}
        ### Question: {natural_language_query}
        ### SQL Query:"""

        inputs = GLOBAL_TOKENIZER(
            prompt, 
            return_tensors="pt", 
            padding=True, 
            truncation=True, 
            max_length=512,
            return_attention_mask=True
        )
        inputs = {k: v.to(device) for k, v in inputs.items()}
        
        with torch.no_grad():
            outputs = GLOBAL_MODEL.generate(
                input_ids=inputs["input_ids"],
                attention_mask=inputs["attention_mask"],
                max_length=256,
                temperature=0.1,
                do_sample=True,
                top_p=0.95,
                num_return_sequences=1,
                pad_token_id=GLOBAL_TOKENIZER.eos_token_id,
            )
            
        generated_query = GLOBAL_TOKENIZER.decode(outputs[0], skip_special_tokens=True)
        sql_query = generated_query.split("### SQL Query:")[-1].strip()
        
        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
    if len(query_result) == 1:
        return "\n".join(f"{k}: {v}" for k, v in query_result[0].items())
    
    results = [f"Found {len(query_result)} results:\n"]
    for i, row in enumerate(query_result[:5], 1):
        results.append(f"Result {i}:")
        results.extend(f"{k}: {v}" for k, v in row.items())
        results.append("")
    
    if len(query_result) > 5:
        results.append(f"(Showing first 5 of {len(query_result)} results)")
    
    return "\n".join(results)

def check_live_connection():
    """Check the database connection status periodically."""
    while True:
        test_db_connection()
        time.sleep(10)  # Check every 10 seconds

def main():
    """Main function with Streamlit UI components"""
    st.title("Natural Language to SQL Query")
    st.write("Ask questions about pizza sales data in plain English.")
    
    # Start the live connection check in a separate thread
    threading.Thread(target=check_live_connection, daemon=True).start()
    
    # Test and display database connection status
    if db_connection_status:
        st.success("Database connection is live.")
    else:
        st.error("Database connection is not live.")
    
    # Initialize model
    initialize_model()
    
    # Input field for natural language query
    natural_language_query = st.text_input("Enter your question", placeholder="e.g., What were the total sales for each pizza category?")
    
    if st.button("Generate and Execute Query"):
        if natural_language_query:
            # Generate SQL query
            sql_query = generate_sql(natural_language_query)
            st.write("Generated SQL Query:", sql_query)

            # Execute the generated query
            query_result = execute_query(sql_query)
            formatted_result = format_result(query_result)
            
            st.write("Query Result:")
            st.code(json.dumps(query_result, indent=2))
            
            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__":
    main()