app.py

from flask import Flask, request, render_template
import os
import json
import re
import logging
import importlib.util
import tkinter as tk
from tkinter import simpledialog, messagebox
import uuid
import webbrowser
from groq import Groq

app = Flask(__name__)

# Configure logging
logging.basicConfig(filename='AI_SQL_Assistant.log', level=logging.INFO)
logger = logging.getLogger(__name__)

# Prompt for Groq API key using Tkinter
def get_groq_api_key():
    groq_api_key = os.getenv('GROQ_API_KEY')
    if not groq_api_key:
        root = tk.Tk()
        root.withdraw()
        groq_api_key = simpledialog.askstring("Groq API Key", "Enter your Groq API key (get from https://console.groq.com):", show='*')
        root.destroy()
    if groq_api_key:
        os.environ['GROQ_API_KEY'] = groq_api_key
        with open('groq_config.json', 'w') as f:
            json.dump({'GROQ_API_KEY': groq_api_key}, f)
    else:
        messagebox.showerror("Error", "Groq API key is required to run the app.")
        logger.error("No Groq API key provided")
        exit(1)
    return groq_api_key

# Load Groq API key
if os.path.exists('groq_config.json'):
    try:
        with open('groq_config.json', 'r') as f:
            config = json.load(f)
            os.environ['GROQ_API_KEY'] = config.get('GROQ_API_KEY', '')
    except Exception as e:
        logger.error("Failed to load groq_config.json: %s", str(e))

try:
    groq_client = Groq(api_key=get_groq_api_key())
    logger.info("Groq client initialized successfully")
except Exception as e:
    groq_client = None
    logger.error("Failed to initialize Groq client: %s", str(e))
    messagebox.showerror("Error", f"Failed to initialize Groq client: {str(e)}")
    exit(1)

# Storage for current schema, summary, and mock data
current_schema = {}
current_summary = {}
current_db_name = None
results = None
generated_query = None
mock_data = {}

def parse_sql_file(file_content):
    """Parse SQL file to extract schema and mock data."""
    global current_db_name, current_schema, mock_data
    file_content = file_content.decode('utf-8') if isinstance(file_content, bytes) else file_content
    statements = []
    current_statement = ""
    in_comment = False

    db_name_match = re.search(r"CREATE\s+DATABASE\s+[`']?(\w+)[`']?", file_content, re.IGNORECASE)
    current_db_name = db_name_match.group(1) if db_name_match else f"temp_db_{uuid.uuid4().hex[:8]}"
    logger.info("Parsed SQL file: database name=%s", current_db_name)

    for line in file_content.splitlines():
        line = line.strip()
        if not line or line.startswith('--'):
            continue
        if line.startswith('/*'):
            in_comment = True
            continue
        if line.endswith('*/'):
            in_comment = False
            continue
        if not in_comment:
            current_statement += line + ' '
            if line.endswith(';'):
                statements.append(current_statement.strip())
                current_statement = ""

    current_schema = {}
    mock_data = {}
    for statement in statements:
        if statement.startswith("CREATE TABLE"):
            table_match = re.search(r"CREATE TABLE\s+[`']?(\w+)[`']?\s*\(", statement, re.IGNORECASE)
            if table_match:
                table_name = table_match.group(1)
                columns = re.findall(r"\b(\w+)\s+\w+(?:\(.*?\))?(?:,\s|\))", statement)
                current_schema[table_name] = columns
                # Mock some data (e.g., 3 rows with sample values)
                mock_data[table_name] = [[f"val{i}_{col}" for col in columns] for i in range(3)]
    logger.info("Extracted schema: %s", current_schema)
    return True, current_schema, None

def generate_sql_query(question, schema):
    """Generate SQL query using Groq API based on schema."""
    if not groq_client:
        logger.error("Groq client not initialized")
        return "ERROR: Groq client not initialized. Check API key and try again."
    schema_text = "\n".join([f"Table: {table}\nColumns: {', '.join(columns)}" for table, columns in schema.items()])
    prompt = f"""
You are a SQL expert. Based on the following database schema from an uploaded .sql file, generate a valid MySQL query for the user's question. Only use tables and columns that exist in the schema. Use user-friendly aliases for column names (e.g., 'cust_id' becomes 'Customer ID', 'admission_date' becomes 'Admission Date'). Return ONLY the SQL query, without explanations, markdown, or code block formatting (e.g., no ```). If the question references non-existent tables or columns, return an error message starting with 'ERROR:'. Do not use GROUP BY or aggregation functions (e.g., SUM, COUNT, AVG) unless the question explicitly requests aggregation (e.g., 'sum of all bills', 'average cost', 'count of patients'). Treat 'total bill amount' as the individual bill amount (e.g., bill.amount) unless aggregation is clearly specified. For names, concatenate first_name and last_name if applicable (e.g., CONCAT(first_name, ' ', last_name) AS 'Full Name'). Use direct JOINs with correct foreign key relationships if implied (e.g., table_id columns). Avoid subqueries unless absolutely necessary. Place filtering conditions (e.g., department name, status) in the WHERE clause, not JOIN clauses. Handle case sensitivity in string comparisons by using LOWER() for status fields (e.g., LOWER(status) = 'unpaid'). Verify table relationships before joining.

Schema:
{schema_text}

User Question: {question}
"""
    try:
        response = groq_client.chat.completions.create(
            messages=[{"role": "user", "content": prompt}],
            model="llama3-70b-8192"
        )
        query = response.choices[0].message.content.strip()
        query = re.sub(r'```(?:sql)?\n?', '', query)
        query = query.strip()
        logger.info("Generated SQL query: %s", query[:100])
        return query
    except Exception as e:
        logger.error("Failed to generate SQL query: %s", str(e))
        return f"ERROR: Failed to generate SQL query: {str(e)}"

def execute_mock_query(query):
    """Simulate query execution with mock data."""
    if not current_schema or not mock_data:
        logger.error("No schema or mock data available")
        return False, "No data loaded. Please upload an .sql file first.", None
    try:
        # Simple mock execution: return data from the first table if query matches
        table_match = re.search(r"FROM\s+[`']?(\w+)[`']?", query, re.IGNORECASE)
        if table_match and table_match.group(1) in mock_data:
            return True, mock_data[table_match.group(1)], None
        return False, "Mock query execution failed: Table not found or query not supported.", None
    except Exception as e:
        logger.error("Mock query execution failed: %s", str(e))
        return False, f"Mock query execution failed: {str(e)}", None

@app.route('/', methods=['GET', 'POST'])
def index():
    global current_schema, current_summary, results, generated_query
    error = None

    if not groq_client:
        error = "Groq client not initialized. Please restart the app and enter a valid Groq API key."
        logger.error(error)

    if request.method == 'POST':
        logger.info("Received POST request")
        if 'sql_file' in request.files:
            file = request.files['sql_file']
            logger.info("SQL file upload detected: %s", file.filename if file else "No file")
            if file and file.filename.endswith('.sql'):
                success, schema, _ = parse_sql_file(file.read())
                if success:
                    current_schema = schema
                    logger.info("SQL file parsed successfully")
                else:
                    error = "Failed to parse SQL file."
                    logger.error(error)
            else:
                error = "Please upload a valid .sql file."
                logger.error(error)
        elif 'question' in request.form:
            question = request.form['question']
            logger.info("Received question: %s", question)
            if not current_schema:
                error = "No schema loaded. Please upload an .sql file first."
                logger.error(error)
            else:
                generated_query = generate_sql_query(question, current_schema)
                if not generated_query.startswith('ERROR:'):
                    success, result, _ = execute_mock_query(generated_query)
                    if success:
                        results = result
                        logger.info("Mock query executed successfully, results: %d rows", len(result))
                    else:
                        error = result
                        logger.error(error)
                else:
                    error = generated_query
                    logger.error(error)

    logger.info("Rendering index.html: error=%s, schema=%s, summary=%s, results=%s",
                error, bool(current_schema), bool(current_summary), bool(results))
    return render_template('index.html', error=error, schema=current_schema, summary=current_summary, results=results, query=generated_query)

if __name__ == '__main__':
    try:
        spec = importlib.util.find_spec("webbrowser")
        if spec is None:
            logger.error("Standard library webbrowser module not found")
            raise ImportError("Could not find webbrowser module")
        webbrowser = importlib.util.module_from_spec(spec)
        spec.loader.exec_module(webbrowser)
    except ImportError as e:
        logger.error("Failed to import webbrowser: %s", str(e))
        raise

    url = 'http://localhost:7860'
    try:
        webbrowser.open(url)
    except AttributeError:
        logger.error("webbrowser.open() failed, possibly due to environment issue")
        print(f"Warning: Could not open browser automatically. Please navigate to {url} manually.")
    app.run(host='0.0.0.0', port=7860, debug=False)