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)