# Warning: This project may not work due to Hugging Face restrictions. Please check out the GitHub repo for the latest updates. import gradio as gr import pandas as pd import os import json import re from typing import Optional, Tuple, Dict, Any, List import traceback from datetime import datetime import time # Database imports import mysql.connector from sqlalchemy import create_engine, inspect, text from sqlalchemy.exc import SQLAlchemyError # LangChain imports from langchain_community.agent_toolkits.sql.base import create_sql_agent from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit from langchain_community.utilities import SQLDatabase from langchain.agents.agent_types import AgentType from langchain_community.callbacks.manager import get_openai_callback from langchain_google_genai import ChatGoogleGenerativeAI # Environment setup from dotenv import load_dotenv load_dotenv() class DatabaseManager: def __init__(self): self.db_connection = None self.db_context = None self.sql_agent = None self.connection_status = "Not Connected" self.db_type = None self.query_history = [] # Store query history self.max_history_items = 20 # Maximum number of history items to keep self.user_api_key = None # Store user-provided API key def set_api_key(self, api_key: str) -> str: """Set user-provided API key""" if not api_key or not api_key.strip(): self.user_api_key = None return "❌ API key cleared. Using environment variable if available." # Store the API key self.user_api_key = api_key.strip() return "✅ API key set successfully!" def get_api_key(self) -> str: """Get API key with priority to user-provided key""" if self.user_api_key: return self.user_api_key return os.getenv("GOOGLE_API_KEY", "") def connect_mysql(self, host: str, port: str, username: str, password: str, database: str) -> Tuple[str, str]: """Connect to MySQL database""" try: # Clean and validate inputs host = host.strip() if host else "localhost" port_num = int(port.strip()) if port and port.strip() else 3306 username = username.strip() if username else "" password = str(password) if password else "" # Ensure password is treated as string database = database.strip() if database else "" if not username or not database: return "❌ Missing required fields", "Please provide username and database name." # Test connection first with mysql.connector # Using raw credentials without URL encoding for direct connection conn = mysql.connector.connect( host=host, port=port_num, user=username, password=password, database=database, autocommit=True ) conn.close() # Create SQLAlchemy engine with proper URL encoding from urllib.parse import quote_plus # Make sure to properly encode all special characters in password encoded_password = quote_plus(str(password)) encoded_username = quote_plus(username) encoded_database = quote_plus(database) # Add binary_prefix=true to handle binary data warnings connection_string = f"mysql+pymysql://{encoded_username}:{encoded_password}@{host}:{port_num}/{encoded_database}?binary_prefix=true" engine = create_engine(connection_string, echo=False) # Test SQLAlchemy connection with engine.connect() as conn: conn.execute(text("SELECT 1")) # Create LangChain SQLDatabase self.db_connection = SQLDatabase(engine) self.db_type = "MySQL" self.connection_status = f"✅ Connected to MySQL: {host}:{port_num}/{database}" return self.connection_status, "Connection successful! You can now analyze the database." except Exception as e: error_msg = f"❌ MySQL Connection Failed: {str(e)}" self.connection_status = "Not Connected" return error_msg, f"Connection failed. Please check your credentials.\nError details: {str(e)}" def validate_sql_query(self, sql_query: str) -> Tuple[bool, str]: """ Validate SQL query for common errors and security issues Args: sql_query: SQL query string to validate Returns: Tuple of (is_valid, message) """ if not sql_query or not isinstance(sql_query, str): return False, "Invalid or empty SQL query" sql_query = sql_query.strip() # Check for basic SQL injection patterns dangerous_patterns = [ "DROP TABLE", "DROP DATABASE", "DELETE FROM", "TRUNCATE TABLE", "ALTER TABLE", "UPDATE", "INSERT INTO", "CREATE TABLE", "GRANT", "REVOKE", "--", ";--", ";", "/*", "*/" ] for pattern in dangerous_patterns: if pattern.upper() in sql_query.upper(): return False, f"Potentially harmful SQL detected: {pattern}" # Check for common SQL errors common_errors = [ # NOT IN with NULL values (r"NOT\s+IN.*NULL", "Using NOT IN with NULL values can lead to unexpected results"), # BETWEEN for exclusive ranges (r"BETWEEN.*AND", "Check BETWEEN usage for correct inclusive/exclusive ranges"), # Potential data type mismatches (r"CAST\(|CONVERT\(", "Verify data type casting is correct"), # Potential quoting issues (r"[^']'[^']|[^']'$", "Check for proper quoting of identifiers") ] import re for pattern, message in common_errors: if re.search(pattern, sql_query, re.IGNORECASE): # This is just a warning, not an error return True, f"Warning: {message}" # Check for SELECT statement if not sql_query.upper().startswith("SELECT"): return False, "Only SELECT queries are allowed" return True, "Query validation passed" def fix_sql_query(self, sql_query: str, error_message: str, db_schema: Optional[dict] = None) -> str: """ Use LLM to fix an invalid SQL query Args: sql_query: The original invalid SQL query error_message: The error message from validation or execution db_schema: Optional database schema information to help with correction Returns: Corrected SQL query """ api_key = self.get_api_key() if not api_key: raise ValueError("No API key available. Please set a Google API key.") # Initialize LLM llm = ChatGoogleGenerativeAI( model="gemini-2.5-flash-preview-05-20", temperature=0, google_api_key=api_key ) # Prepare schema information if available schema_info = "" if db_schema and isinstance(db_schema, dict): schema_info = "Database schema information:\n" for table, info in db_schema.items(): schema_info += f"Table: {table}\n" if "columns" in info: schema_info += "Columns:\n" for col in info["columns"]: schema_info += f"- {col['name']} ({col['type']})\n" schema_info += "\n" # Build prompt for the LLM prompt = f""" Fix the following SQL query that has errors: ```sql {sql_query} ``` Error message: {error_message} {schema_info} Please provide ONLY the corrected SQL query with no additional text or explanation. The query should be a valid SELECT statement. """ # Get the corrected query try: response = llm.invoke(prompt) corrected_query = response.content # Extract SQL from response if needed if "```sql" in corrected_query: corrected_query = corrected_query.split("```sql")[1].split("```")[0].strip() elif "```" in corrected_query: corrected_query = corrected_query.split("```")[1].strip() return corrected_query except Exception as e: # If correction fails, return the original query return sql_query def analyze_database(self) -> Tuple[str, str]: """Analyze database structure and create context""" if not self.db_connection: return "❌ No database connection", "Please connect to a database first." try: # Get database schema information inspector = inspect(self.db_connection._engine) tables = inspector.get_table_names() context_info = { "database_type": self.db_type, "total_tables": len(tables), "tables": {}, "analysis_timestamp": datetime.now().isoformat() } # Analyze each table for table in tables[:10]: # Limit to first 10 tables for performance try: columns = inspector.get_columns(table) primary_keys = inspector.get_pk_constraint(table) foreign_keys = inspector.get_foreign_keys(table) # Get sample data count with self.db_connection._engine.connect() as conn: result = conn.execute(text(f"SELECT COUNT(*) FROM {table}")) row_count = result.scalar() context_info["tables"][table] = { "columns": [{"name": col["name"], "type": str(col["type"])} for col in columns], "primary_keys": primary_keys["constrained_columns"] if primary_keys else [], "foreign_keys": [{"columns": fk["constrained_columns"], "refers_to": f"{fk['referred_table']}.{fk['referred_columns']}"} for fk in foreign_keys], "row_count": row_count } except Exception as table_error: context_info["tables"][table] = {"error": str(table_error)} self.db_context = context_info # Initialize Gemini LLM api_key = self.get_api_key() if not api_key: return "❌ Analysis Failed", "Please set a Google API key in the settings or environment variables" llm = ChatGoogleGenerativeAI( model="gemini-2.5-flash-preview-05-20", temperature=0, google_api_key=api_key ) # Create SQL agent toolkit = SQLDatabaseToolkit(db=self.db_connection, llm=llm) self.sql_agent = create_sql_agent( llm=llm, toolkit=toolkit, agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION, verbose=True, handle_parsing_errors=True ) summary = f""" ✅ Database Analysis Complete! 📊 Database: {self.db_type} 📋 Tables Found: {len(tables)} 🔍 Analyzed Tables: {min(len(tables), 10)} Ready for natural language queries! """ detailed_info = json.dumps(context_info, indent=2) return summary, f"Database context saved. You can now ask questions!\n\nDetailed Analysis:\n{detailed_info}" except Exception as e: error_msg = f"❌ Analysis Failed: {str(e)}" return error_msg, f"Error during analysis: {traceback.format_exc()}" def query_database(self, question: str) -> Tuple[str, str]: """Process natural language query and return results""" if not self.sql_agent: # Convert error to table format with clean RDBMS style df_error = pd.DataFrame({"Message": ["Please connect and analyze database first."]}) table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0) return "❌ Not Ready", table_html if not question.strip(): # Convert error to table format with clean RDBMS style df_error = pd.DataFrame({"Message": ["Please enter a question about your database."]}) table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0) return "❌ Empty Query", table_html try: # Track query start time for overall performance start_time = time.time() # Process the query with the agent result = self.sql_agent.run(question) # Try to extract and execute the SQL query for tabular display try: # Look for SQL in the result if "SELECT" in result.upper(): # Extract SQL query (this is a simple extraction, could be improved) lines = result.split('\n') sql_lines = [line for line in lines if 'SELECT' in line.upper()] if sql_lines: sql_query = sql_lines[0].strip() # Clean up the SQL query sql_query = sql_query.replace('sql', '').replace('```', '').strip() # Validate the SQL query before execution is_valid, validation_message = self.validate_sql_query(sql_query) # If query is invalid, try to fix it correction_applied = False if not is_valid: # Get schema information for the correction agent schema_info = self.db_context["tables"] if self.db_context else None # Try to fix the query corrected_query = self.fix_sql_query(sql_query, validation_message, schema_info) # Validate the corrected query is_valid_corrected, validation_message_corrected = self.validate_sql_query(corrected_query) if is_valid_corrected: sql_query = corrected_query validation_message = validation_message_corrected correction_applied = True is_valid = True else: # If correction also failed, return both errors in table format error_msg = f"The generated SQL query failed validation: {validation_message}\n\nAttempted correction also failed: {validation_message_corrected}\n\nOriginal result:\n{result}" df_error = pd.DataFrame({"Error": [error_msg]}) table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0) return "❌ Query Validation Failed", table_html # If there's a warning but query is valid, add it to the result warning_message = "" if validation_message.startswith("Warning:"): warning_message = f"\n\n⚠️ {validation_message}" # Add correction notice if applicable if correction_applied: warning_message += f"\n\n🔧 Query was automatically corrected. Original query had issues: {validation_message}" # Execute the query to get structured data try: # Measure query performance performance_metrics = self.measure_query_performance(sql_query) if performance_metrics.get("success", False): # Get the data from the metrics with self.db_connection._engine.connect() as conn: df = pd.read_sql(sql_query, conn) # Calculate overall processing time total_time_ms = round((time.time() - start_time) * 1000, 2) # Add query to history history_item = { "question": question, "sql_query": sql_query, "execution_time_ms": performance_metrics["execution_time_ms"], "total_time_ms": total_time_ms, "row_count": performance_metrics["row_count"], "complexity": performance_metrics["complexity"]["level"], "timestamp": datetime.now().isoformat() } self.add_to_query_history(history_item) # Generate performance and complexity insights complexity = performance_metrics["complexity"] perf_insights = f"\n\n📊 Query Metrics:\n" perf_insights += f"• Execution time: {performance_metrics['execution_time_ms']}ms\n" perf_insights += f"• Total processing time: {total_time_ms}ms\n" perf_insights += f"• Rows returned: {performance_metrics['row_count']}\n" perf_insights += f"• Complexity: {complexity['level']}\n" if complexity["insights"]: perf_insights += "\n🔍 Insights:\n" for insight in complexity["insights"]: perf_insights += f"• {insight}\n" if not df.empty: # Format table in RDBMS style table_html = df.to_html(index=False, classes="table table-bordered table-striped", border=0) # Add custom styling to make it look more like RDBMS output table_html = f""" {table_html} """ return f"✅ Query Successful\n\n{result}{warning_message}{perf_insights}", table_html else: # If performance measurement failed, continue with normal execution with self.db_connection._engine.connect() as conn: df = pd.read_sql(sql_query, conn) if not df.empty: # Format table in RDBMS style table_html = df.to_html(index=False, classes="table table-bordered table-striped", border=0) # Add custom styling table_html = f""" {table_html} """ return f"✅ Query Successful\n\n{result}{warning_message}", table_html except SQLAlchemyError as exec_error: # If execution fails, try to fix the query again with the specific error if not correction_applied: schema_info = self.db_context["tables"] if self.db_context else None corrected_query = self.fix_sql_query(sql_query, str(exec_error), schema_info) # Try executing the corrected query try: # Measure performance of corrected query performance_metrics = self.measure_query_performance(corrected_query) if performance_metrics.get("success", False): # Get the data from the metrics with self.db_connection._engine.connect() as conn: df = pd.read_sql(corrected_query, conn) # Calculate overall processing time total_time_ms = round((time.time() - start_time) * 1000, 2) # Add query to history history_item = { "question": question, "sql_query": corrected_query, "execution_time_ms": performance_metrics["execution_time_ms"], "total_time_ms": total_time_ms, "row_count": performance_metrics["row_count"], "complexity": performance_metrics["complexity"]["level"], "timestamp": datetime.now().isoformat(), "corrected": True, "original_query": sql_query } self.add_to_query_history(history_item) # Generate performance and complexity insights complexity = performance_metrics["complexity"] perf_insights = f"\n\n📊 Query Metrics:\n" perf_insights += f"• Execution time: {performance_metrics['execution_time_ms']}ms\n" perf_insights += f"• Total processing time: {total_time_ms}ms\n" perf_insights += f"• Rows returned: {performance_metrics['row_count']}\n" perf_insights += f"• Complexity: {complexity['level']}\n" if complexity["insights"]: perf_insights += "\n🔍 Insights:\n" for insight in complexity["insights"]: perf_insights += f"• {insight}\n" if not df.empty: # Format table in RDBMS style table_html = df.to_html(index=False, classes="table table-bordered table-striped", border=0) # Add custom styling table_html = f""" {table_html} """ return f"✅ Query Successful (after correction)\n\n{result}\n\n🔧 Query was automatically corrected due to execution error: {str(exec_error)}{perf_insights}", table_html else: # If performance measurement failed, continue with normal execution with self.db_connection._engine.connect() as conn: df = pd.read_sql(corrected_query, conn) if not df.empty: # Format table in RDBMS style table_html = df.to_html(index=False, classes="table table-bordered table-striped", border=0) # Add custom styling table_html = f""" {table_html} """ return f"✅ Query Successful (after correction)\n\n{result}\n\n🔧 Query was automatically corrected due to execution error: {str(exec_error)}", table_html except Exception: # If correction fails, return the original error pass # Return the execution error in table format error_msg = f"The query failed to execute:\n\n{str(exec_error)}\n\nOriginal result:\n{result}" df_error = pd.DataFrame({"Error": [error_msg]}) table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0) # Add custom styling table_html = f""" {table_html} """ return "❌ SQL Execution Error", table_html except SQLAlchemyError as sql_error: # Handle SQL execution errors error_details = str(sql_error) error_msg = f"❌ SQL Execution Error" details = f"The query failed to execute:\n\n{error_details}\n\nOriginal result:\n{result}" df_error = pd.DataFrame({"Error": [details]}) table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0) # Add custom styling table_html = f""" {table_html} """ return error_msg, table_html except Exception as table_error: # If table extraction fails, just return the text result pass # If we got here, we just have the text result without structured data # Convert to table format with RDBMS style df_text = pd.DataFrame({"Result": [result]}) table_html = df_text.to_html(index=False, classes="table table-bordered table-striped", border=0) # Add custom styling table_html = f""" {table_html} """ # Add to history history_item = { "question": question, "result": result, "timestamp": datetime.now().isoformat() } self.add_to_query_history(history_item) return f"✅ Query Successful", table_html except Exception as e: # Convert exception to table format with RDBMS style error_msg = f"❌ Query Failed: {str(e)}" details = f"Error processing query: {traceback.format_exc()}" df_error = pd.DataFrame({"Error": [details]}) table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0) # Add custom styling table_html = f""" {table_html} """ return error_msg, table_html def analyze_query_complexity(self, sql_query: str) -> Dict[str, Any]: """ Analyze SQL query complexity and provide insights Args: sql_query: SQL query to analyze Returns: Dictionary with complexity metrics and insights """ if not sql_query or not isinstance(sql_query, str): return {"error": "Invalid query provided"} sql_query = sql_query.strip().upper() # Initialize complexity metrics complexity = { "level": "Simple", "score": 0, "joins": 0, "tables": [], "aggregations": False, "grouping": False, "ordering": False, "limiting": False, "subqueries": 0, "complex_functions": [], "insights": [] } # Count number of JOINs join_count = len(re.findall(r'\bJOIN\b', sql_query)) complexity["joins"] = join_count if join_count > 0: complexity["score"] += join_count * 2 if join_count >= 3: complexity["insights"].append(f"Query uses {join_count} joins, which may impact performance") # Detect tables used from_clause = re.search(r'\bFROM\b\s+(.*?)(?:\bWHERE\b|\bGROUP\b|\bHAVING\b|\bORDER\b|\bLIMIT\b|$)', sql_query) if from_clause: # Extract table names from FROM clause tables_text = from_clause.group(1).strip() # Handle JOIN syntax in FROM clause tables = re.findall(r'([a-zA-Z0-9_]+)(?:\s+(?:AS\s+)?[a-zA-Z0-9_]+)?', tables_text) complexity["tables"] = list(set(tables)) # Remove duplicates # Check for aggregations agg_functions = ["COUNT", "SUM", "AVG", "MIN", "MAX"] for func in agg_functions: if re.search(rf'\b{func}\s*\(', sql_query): complexity["aggregations"] = True complexity["score"] += 1 break # Check for GROUP BY if re.search(r'\bGROUP\s+BY\b', sql_query): complexity["grouping"] = True complexity["score"] += 2 # Check for ORDER BY if re.search(r'\bORDER\s+BY\b', sql_query): complexity["ordering"] = True complexity["score"] += 1 # Check for LIMIT if re.search(r'\bLIMIT\b', sql_query): complexity["limiting"] = True complexity["score"] += 0.5 # Check for subqueries subquery_count = len(re.findall(r'\(\s*SELECT', sql_query)) complexity["subqueries"] = subquery_count if subquery_count > 0: complexity["score"] += subquery_count * 3 complexity["insights"].append(f"Query contains {subquery_count} subqueries, which may affect performance") # Check for complex functions complex_funcs = ["CASE", "COALESCE", "NULLIF", "CAST", "CONVERT", "SUBSTRING", "CONCAT", "DATE_FORMAT", "EXTRACT"] for func in complex_funcs: if re.search(rf'\b{func}\b', sql_query): complexity["complex_functions"].append(func) complexity["score"] += 1 # Determine complexity level if complexity["score"] <= 2: complexity["level"] = "Simple" elif complexity["score"] <= 5: complexity["level"] = "Moderate" elif complexity["score"] <= 10: complexity["level"] = "Complex" else: complexity["level"] = "Very Complex" complexity["insights"].append("This is a highly complex query that may benefit from optimization") # Add insights based on complexity if complexity["level"] in ["Complex", "Very Complex"] and not complexity["limiting"]: complexity["insights"].append("Consider adding a LIMIT clause to prevent large result sets") if complexity["joins"] >= 2 and not any(idx for idx in complexity["insights"] if "index" in idx.lower()): complexity["insights"].append("Ensure proper indexes exist on join columns") return complexity def add_to_query_history(self, query_data: Dict[str, Any]) -> None: """ Add a query to the history Args: query_data: Dictionary containing query information """ # Add timestamp if not present if "timestamp" not in query_data: query_data["timestamp"] = datetime.now().isoformat() # Add to history (at the beginning for most recent first) self.query_history.insert(0, query_data) # Trim history if needed if len(self.query_history) > self.max_history_items: self.query_history = self.query_history[:self.max_history_items] def get_query_history(self) -> List[Dict[str, Any]]: """ Get the query history Returns: List of query history items """ return self.query_history def clear_query_history(self) -> None: """Clear the query history""" self.query_history = [] def measure_query_performance(self, sql_query: str) -> Dict[str, Any]: """ Measure the performance of a SQL query Args: sql_query: SQL query to execute and measure Returns: Dictionary with performance metrics """ if not self.db_connection: return {"error": "No database connection"} metrics = { "query": sql_query, "execution_time_ms": 0, "row_count": 0, "success": False, "error": None } try: # Measure execution time start_time = time.time() with self.db_connection._engine.connect() as conn: result = conn.execute(text(sql_query)) # Convert to DataFrame to get row count df = pd.DataFrame(result.fetchall(), columns=result.keys()) end_time = time.time() # Calculate metrics metrics["execution_time_ms"] = round((end_time - start_time) * 1000, 2) metrics["row_count"] = len(df) metrics["success"] = True # Add complexity analysis metrics["complexity"] = self.analyze_query_complexity(sql_query) return metrics except Exception as e: metrics["error"] = str(e) return metrics def generate_schema_diagram(self, include_all_tables: bool = False) -> str: """ Generate a Mermaid ER diagram for the database schema Args: include_all_tables: Whether to include all tables or just a subset Returns: Mermaid diagram code """ # Return a message that this functionality is not available return "This functionality has been removed" # Initialize the database manager db_manager = DatabaseManager() def create_interface(): """Create the Gradio interface""" with gr.Blocks(title="AI Database Query Assistant", theme=gr.themes.Soft()) as demo: # Warning banner at the top gr.Markdown("""
⚠️ WARNING: This project may not work due to Hugging Face restrictions. Please check out the GitHub repo for the latest updates. https://github.com/yash-8923/gradio.git
""") gr.Markdown(""" # 🤖 AI Database Query Assistant Connect to your MySQL database and query it using natural language! ### Steps: 1. **Connect** to your database 2. **Analyze** your database structure 3. **Ask questions** in natural language """) # Connection Status connection_status = gr.Textbox( label="Connection Status", value="Not Connected", interactive=False ) with gr.Tabs(): # MySQL Connection Tab with gr.TabItem("MySQL Connection"): gr.Markdown(""" **MySQL Connection Details:** - Enter your MySQL server connection details - Password will be securely handled (not stored) - Default port is 3306 if not specified - Special characters in passwords are supported """) with gr.Row(): mysql_host = gr.Textbox( label="Host", value="localhost", placeholder="localhost or IP address" ) mysql_port = gr.Textbox( label="Port", value="3306", placeholder="3306" ) with gr.Row(): mysql_username = gr.Textbox( label="Username", placeholder="root or your username" ) mysql_password = gr.Textbox( label="DB Password(optional)", type="password", placeholder="Your MySQL password" ) mysql_database = gr.Textbox( label="Database Name", placeholder="my_database" ) mysql_connect_btn = gr.Button("Connect to MySQL", variant="primary") mysql_message = gr.Textbox(label="Connection Message", interactive=False) # API Key Section gr.Markdown(""" ### 🔑 Google API Key Enter your Google API key for Gemini model. If not provided, will use environment variable. """) with gr.Row(): api_key_input = gr.Textbox( label="Google API Key", type="password", placeholder="Enter your Gemini API key here", info="Get your API key from: https://makersuite.google.com/app/apikey" ) api_key_btn = gr.Button("Set API Key", variant="secondary") api_key_message = gr.Textbox(label="API Key Status", interactive=False) # Database Analysis Section with gr.Tabs(): with gr.TabItem("Database Analysis"): gr.Markdown("## 🔍 Database Analysis") analyze_btn = gr.Button("Analyze Database", variant="secondary", size="lg") with gr.Row(): analysis_status = gr.Textbox(label="Analysis Status", interactive=False) analysis_details = gr.Textbox(label="Analysis Details", lines=10, interactive=False) # Schema Visualization gr.Markdown("### 📊 Database Schema Visualization") with gr.Row(): schema_table_select = gr.Dropdown(label="Select Table", choices=[], interactive=True) visualize_schema_btn = gr.Button("Visualize Schema", variant="secondary") schema_output = gr.HTML(label="Schema Visualization") # Removed ER Diagram Visualization section # Query Section with gr.TabItem("Query Database"): gr.Markdown("## 💬 Ask Questions") question_input = gr.Textbox( label="Your Question", placeholder="Example: Show me all customers from New York, What are the top 5 selling products?", lines=2 ) query_btn = gr.Button("Ask Question", variant="primary", size="lg") with gr.Row(): query_status = gr.Textbox(label="Query Result", lines=5, interactive=False) query_output = gr.HTML(label="Data Output") # Example questions gr.Markdown(""" ### 💡 Example Questions: - "Show me all users registered in the last month" - "What are the top 5 products by sales?" - "How many orders were placed yesterday?" - "Show me customers with more than 10 orders" - "What's the average order value?" """) # Query History Tab with gr.TabItem("Query History"): gr.Markdown("## 📜 Query History") with gr.Row(): refresh_history_btn = gr.Button("Refresh History", variant="secondary") clear_history_btn = gr.Button("Clear History", variant="secondary") history_output = gr.HTML(label="Query History") # Reuse Query Section gr.Markdown("### 🔄 Reuse Previous Query") with gr.Row(): history_question_select = gr.Dropdown(label="Select Previous Question", choices=[], interactive=True) reuse_query_btn = gr.Button("Use Selected Query", variant="primary") # Event handlers mysql_connect_btn.click( fn=lambda h, p, u, pw, d: db_manager.connect_mysql(h, p, u, pw, d) + (db_manager.connection_status,), inputs=[mysql_host, mysql_port, mysql_username, mysql_password, mysql_database], outputs=[mysql_message, connection_status] ) # API Key event handler api_key_btn.click( fn=db_manager.set_api_key, inputs=[api_key_input], outputs=[api_key_message] ) # Database analysis event handler def on_analyze_database(): status, details = db_manager.analyze_database() # Update schema table dropdown if analysis was successful table_choices = [] if "✅" in status and db_manager.db_context: table_choices = list(db_manager.db_context.get("tables", {}).keys()) return status, details, gr.Dropdown(choices=table_choices) analyze_btn.click( fn=on_analyze_database, outputs=[analysis_status, analysis_details, schema_table_select] ) # Schema visualization event handler def visualize_table_schema(table_name): if not table_name or not db_manager.db_context or table_name not in db_manager.db_context.get("tables", {}): return "

Please select a valid table

" table_info = db_manager.db_context["tables"][table_name] # Create HTML visualization html = f"

Table: {table_name}

" html += f"

Row count: {table_info.get('row_count', 'Unknown')}

" # Create table for columns html += "" html += "" html += "" # Add columns primary_keys = table_info.get("primary_keys", []) foreign_keys_flat = [] # Flatten foreign key references for fk in table_info.get("foreign_keys", []): for col in fk.get("columns", []): foreign_keys_flat.append(col) for col in table_info.get("columns", []): col_name = col.get("name", "") col_type = col.get("type", "") # Determine key type key_type = "" if col_name in primary_keys: key_type = "🔑 Primary" elif col_name in foreign_keys_flat: key_type = "🔗 Foreign" html += f"" html += "
ColumnTypeKey
{col_name}{col_type}{key_type}
" # Add foreign key relationships if table_info.get("foreign_keys"): html += "

Foreign Key Relationships

" html += "" return html visualize_schema_btn.click( fn=visualize_table_schema, inputs=[schema_table_select], outputs=[schema_output] ) # Removed ER Diagram event handler # Query event handler query_btn.click( fn=db_manager.query_database, inputs=[question_input], outputs=[query_status, query_output] ) # Query history event handlers def format_query_history(): history = db_manager.get_query_history() if not history: return "

No queries in history

", gr.Dropdown(choices=[]) # Format history as HTML table html = "" html += "" html += "" # Collect questions for dropdown questions = [] for i, item in enumerate(history): # Format timestamp timestamp = item.get("timestamp", "") if timestamp: try: dt = datetime.fromisoformat(timestamp) timestamp = dt.strftime("%Y-%m-%d %H:%M:%S") except: pass question = item.get("question", "") sql_query = item.get("sql_query", "") exec_time = f"{item.get('execution_time_ms', 0)}ms" if "execution_time_ms" in item else "-" row_count = item.get("row_count", "-") complexity = item.get("complexity", "-") # Add question to dropdown options if question: questions.append(question) # Format row with corrected query highlight row_class = " class='table-warning'" if item.get("corrected", False) else "" html += f"" html += f"" html += f"" html += f"" html += f"" html += f"" html += f"" html += "" html += "
TimeQuestionSQL QueryExecution TimeRowsComplexity
{timestamp}{question}{sql_query}{exec_time}{row_count}{complexity}
" return html, gr.Dropdown(choices=questions) refresh_history_btn.click( fn=format_query_history, outputs=[history_output, history_question_select] ) clear_history_btn.click( fn=lambda: (db_manager.clear_query_history(), "

History cleared

", gr.Dropdown(choices=[])), outputs=[history_output, history_question_select] ) # Reuse query event handler def reuse_question(selected_question): if not selected_question: return gr.Textbox(value="") return gr.Textbox(value=selected_question) reuse_query_btn.click( fn=reuse_question, inputs=[history_question_select], outputs=[question_input] ) return demo if __name__ == "__main__": # Check for required environment variables if not os.getenv("GOOGLE_API_KEY"): print("⚠️ Warning: GOOGLE_API_KEY not found in environment variables") print("You will need to provide an API key in the interface or set the environment variable.") print("Get your API key from: https://makersuite.google.com/app/apikey") # Create and launch the interface demo = create_interface() demo.launch( #server_name="0.0.0.0", server_port=7860, share=False, debug=False )