Agents-MCP-Hackathon/text2sql
0
1# Warning: This project may not work due to Hugging Face restrictions. Please check out the GitHub repo for the latest updates.2 3import gradio as gr4import pandas as pd5import os6import json7import re8from typing import Optional, Tuple, Dict, Any, List9import traceback10from datetime import datetime11import time12 13# Database imports14import mysql.connector15from sqlalchemy import create_engine, inspect, text16from sqlalchemy.exc import SQLAlchemyError17 18# LangChain imports19from langchain_community.agent_toolkits.sql.base import create_sql_agent20from langchain_community.agent_toolkits.sql.toolkit import SQLDatabaseToolkit21from langchain_community.utilities import SQLDatabase22from langchain.agents.agent_types import AgentType23from langchain_community.callbacks.manager import get_openai_callback24from langchain_google_genai import ChatGoogleGenerativeAI25 26# Environment setup27from dotenv import load_dotenv28load_dotenv()29 30class DatabaseManager:31 def __init__(self):32 self.db_connection = None33 self.db_context = None34 self.sql_agent = None35 self.connection_status = "Not Connected"36 self.db_type = None37 self.query_history = [] # Store query history38 self.max_history_items = 20 # Maximum number of history items to keep39 self.user_api_key = None # Store user-provided API key40 41 def set_api_key(self, api_key: str) -> str:42 """Set user-provided API key"""43 if not api_key or not api_key.strip():44 self.user_api_key = None45 return "❌ API key cleared. Using environment variable if available."46 47 # Store the API key48 self.user_api_key = api_key.strip()49 return "✅ API key set successfully!"50 51 def get_api_key(self) -> str:52 """Get API key with priority to user-provided key"""53 if self.user_api_key:54 return self.user_api_key55 return os.getenv("GOOGLE_API_KEY", "")56 57 def connect_mysql(self, host: str, port: str, username: str, password: str, database: str) -> Tuple[str, str]:58 """Connect to MySQL database"""59 try:60 # Clean and validate inputs61 host = host.strip() if host else "localhost"62 port_num = int(port.strip()) if port and port.strip() else 330663 username = username.strip() if username else ""64 password = str(password) if password else "" # Ensure password is treated as string65 database = database.strip() if database else ""66 67 if not username or not database:68 return "❌ Missing required fields", "Please provide username and database name."69 70 # Test connection first with mysql.connector71 # Using raw credentials without URL encoding for direct connection72 conn = mysql.connector.connect(73 host=host,74 port=port_num,75 user=username,76 password=password,77 database=database,78 autocommit=True79 )80 conn.close()81 82 # Create SQLAlchemy engine with proper URL encoding83 from urllib.parse import quote_plus84 # Make sure to properly encode all special characters in password85 encoded_password = quote_plus(str(password))86 encoded_username = quote_plus(username)87 encoded_database = quote_plus(database)88 89 # Add binary_prefix=true to handle binary data warnings90 connection_string = f"mysql+pymysql://{encoded_username}:{encoded_password}@{host}:{port_num}/{encoded_database}?binary_prefix=true"91 engine = create_engine(connection_string, echo=False)92 93 # Test SQLAlchemy connection94 with engine.connect() as conn:95 conn.execute(text("SELECT 1"))96 97 # Create LangChain SQLDatabase98 self.db_connection = SQLDatabase(engine)99 self.db_type = "MySQL"100 self.connection_status = f"✅ Connected to MySQL: {host}:{port_num}/{database}"101 102 return self.connection_status, "Connection successful! You can now analyze the database."103 104 except Exception as e:105 error_msg = f"❌ MySQL Connection Failed: {str(e)}"106 self.connection_status = "Not Connected"107 return error_msg, f"Connection failed. Please check your credentials.\nError details: {str(e)}"108 109 def validate_sql_query(self, sql_query: str) -> Tuple[bool, str]:110 """111 Validate SQL query for common errors and security issues112 113 Args:114 sql_query: SQL query string to validate115 116 Returns:117 Tuple of (is_valid, message)118 """119 if not sql_query or not isinstance(sql_query, str):120 return False, "Invalid or empty SQL query"121 122 sql_query = sql_query.strip()123 124 # Check for basic SQL injection patterns125 dangerous_patterns = [126 "DROP TABLE", "DROP DATABASE", "DELETE FROM", "TRUNCATE TABLE",127 "ALTER TABLE", "UPDATE", "INSERT INTO", "CREATE TABLE", "GRANT",128 "REVOKE", "--", ";--", ";", "/*", "*/"129 ]130 131 for pattern in dangerous_patterns:132 if pattern.upper() in sql_query.upper():133 return False, f"Potentially harmful SQL detected: {pattern}"134 135 # Check for common SQL errors136 common_errors = [137 # NOT IN with NULL values138 (r"NOT\s+IN.*NULL", "Using NOT IN with NULL values can lead to unexpected results"),139 # BETWEEN for exclusive ranges140 (r"BETWEEN.*AND", "Check BETWEEN usage for correct inclusive/exclusive ranges"),141 # Potential data type mismatches142 (r"CAST\(|CONVERT\(", "Verify data type casting is correct"),143 # Potential quoting issues144 (r"[^']'[^']|[^']'$", "Check for proper quoting of identifiers")145 ]146 147 import re148 for pattern, message in common_errors:149 if re.search(pattern, sql_query, re.IGNORECASE):150 # This is just a warning, not an error151 return True, f"Warning: {message}"152 153 # Check for SELECT statement154 if not sql_query.upper().startswith("SELECT"):155 return False, "Only SELECT queries are allowed"156 157 return True, "Query validation passed"158 159 def fix_sql_query(self, sql_query: str, error_message: str, db_schema: Optional[dict] = None) -> str:160 """161 Use LLM to fix an invalid SQL query162 163 Args:164 sql_query: The original invalid SQL query165 error_message: The error message from validation or execution166 db_schema: Optional database schema information to help with correction167 168 Returns:169 Corrected SQL query170 """171 api_key = self.get_api_key()172 if not api_key:173 raise ValueError("No API key available. Please set a Google API key.")174 175 # Initialize LLM176 llm = ChatGoogleGenerativeAI(177 model="gemini-2.5-flash-preview-05-20",178 temperature=0,179 google_api_key=api_key180 )181 182 # Prepare schema information if available183 schema_info = ""184 if db_schema and isinstance(db_schema, dict):185 schema_info = "Database schema information:\n"186 for table, info in db_schema.items():187 schema_info += f"Table: {table}\n"188 if "columns" in info:189 schema_info += "Columns:\n"190 for col in info["columns"]:191 schema_info += f"- {col['name']} ({col['type']})\n"192 schema_info += "\n"193 194 # Build prompt for the LLM195 prompt = f"""196 Fix the following SQL query that has errors:197 198 ```sql199 {sql_query}200 ```201 202 Error message:203 {error_message}204 205 {schema_info}206 207 Please provide ONLY the corrected SQL query with no additional text or explanation.208 The query should be a valid SELECT statement.209 """210 211 # Get the corrected query212 try:213 response = llm.invoke(prompt)214 corrected_query = response.content215 216 # Extract SQL from response if needed217 if "```sql" in corrected_query:218 corrected_query = corrected_query.split("```sql")[1].split("```")[0].strip()219 elif "```" in corrected_query:220 corrected_query = corrected_query.split("```")[1].strip()221 222 return corrected_query223 except Exception as e:224 # If correction fails, return the original query225 return sql_query226 227 def analyze_database(self) -> Tuple[str, str]:228 """Analyze database structure and create context"""229 if not self.db_connection:230 return "❌ No database connection", "Please connect to a database first."231 232 try:233 # Get database schema information234 inspector = inspect(self.db_connection._engine)235 tables = inspector.get_table_names()236 237 context_info = {238 "database_type": self.db_type,239 "total_tables": len(tables),240 "tables": {},241 "analysis_timestamp": datetime.now().isoformat()242 }243 244 # Analyze each table245 for table in tables[:10]: # Limit to first 10 tables for performance246 try:247 columns = inspector.get_columns(table)248 primary_keys = inspector.get_pk_constraint(table)249 foreign_keys = inspector.get_foreign_keys(table)250 251 # Get sample data count252 with self.db_connection._engine.connect() as conn:253 result = conn.execute(text(f"SELECT COUNT(*) FROM {table}"))254 row_count = result.scalar()255 256 context_info["tables"][table] = {257 "columns": [{"name": col["name"], "type": str(col["type"])} for col in columns],258 "primary_keys": primary_keys["constrained_columns"] if primary_keys else [],259 "foreign_keys": [{"columns": fk["constrained_columns"], "refers_to": f"{fk['referred_table']}.{fk['referred_columns']}"} for fk in foreign_keys],260 "row_count": row_count261 }262 except Exception as table_error:263 context_info["tables"][table] = {"error": str(table_error)}264 265 self.db_context = context_info266 267 # Initialize Gemini LLM268 api_key = self.get_api_key()269 if not api_key:270 return "❌ Analysis Failed", "Please set a Google API key in the settings or environment variables"271 272 llm = ChatGoogleGenerativeAI(273 model="gemini-2.5-flash-preview-05-20",274 temperature=0,275 google_api_key=api_key276 )277 278 # Create SQL agent279 toolkit = SQLDatabaseToolkit(db=self.db_connection, llm=llm)280 self.sql_agent = create_sql_agent(281 llm=llm,282 toolkit=toolkit,283 agent_type=AgentType.ZERO_SHOT_REACT_DESCRIPTION,284 verbose=True,285 handle_parsing_errors=True286 )287 288 summary = f"""289 ✅ Database Analysis Complete!290 291 📊 Database: {self.db_type}292 📋 Tables Found: {len(tables)}293 🔍 Analyzed Tables: {min(len(tables), 10)}294 295 Ready for natural language queries!296 """297 298 detailed_info = json.dumps(context_info, indent=2)299 return summary, f"Database context saved. You can now ask questions!\n\nDetailed Analysis:\n{detailed_info}"300 301 except Exception as e:302 error_msg = f"❌ Analysis Failed: {str(e)}"303 return error_msg, f"Error during analysis: {traceback.format_exc()}"304 305 def query_database(self, question: str) -> Tuple[str, str]:306 """Process natural language query and return results"""307 if not self.sql_agent:308 # Convert error to table format with clean RDBMS style309 df_error = pd.DataFrame({"Message": ["Please connect and analyze database first."]})310 table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0)311 return "❌ Not Ready", table_html312 313 if not question.strip():314 # Convert error to table format with clean RDBMS style315 df_error = pd.DataFrame({"Message": ["Please enter a question about your database."]})316 table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0)317 return "❌ Empty Query", table_html318 319 try:320 # Track query start time for overall performance321 start_time = time.time()322 323 # Process the query with the agent324 result = self.sql_agent.run(question)325 326 # Try to extract and execute the SQL query for tabular display327 try:328 # Look for SQL in the result329 if "SELECT" in result.upper():330 # Extract SQL query (this is a simple extraction, could be improved)331 lines = result.split('\n')332 sql_lines = [line for line in lines if 'SELECT' in line.upper()]333 334 if sql_lines:335 sql_query = sql_lines[0].strip()336 # Clean up the SQL query337 sql_query = sql_query.replace('sql', '').replace('```', '').strip()338 339 # Validate the SQL query before execution340 is_valid, validation_message = self.validate_sql_query(sql_query)341 342 # If query is invalid, try to fix it343 correction_applied = False344 if not is_valid:345 # Get schema information for the correction agent346 schema_info = self.db_context["tables"] if self.db_context else None347 348 # Try to fix the query349 corrected_query = self.fix_sql_query(sql_query, validation_message, schema_info)350 351 # Validate the corrected query352 is_valid_corrected, validation_message_corrected = self.validate_sql_query(corrected_query)353 354 if is_valid_corrected:355 sql_query = corrected_query356 validation_message = validation_message_corrected357 correction_applied = True358 is_valid = True359 else:360 # If correction also failed, return both errors in table format361 error_msg = f"The generated SQL query failed validation: {validation_message}\n\nAttempted correction also failed: {validation_message_corrected}\n\nOriginal result:\n{result}"362 df_error = pd.DataFrame({"Error": [error_msg]})363 table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0)364 return "❌ Query Validation Failed", table_html365 366 # If there's a warning but query is valid, add it to the result367 warning_message = ""368 if validation_message.startswith("Warning:"):369 warning_message = f"\n\n⚠️ {validation_message}"370 371 # Add correction notice if applicable372 if correction_applied:373 warning_message += f"\n\n🔧 Query was automatically corrected. Original query had issues: {validation_message}"374 375 # Execute the query to get structured data376 try:377 # Measure query performance378 performance_metrics = self.measure_query_performance(sql_query)379 380 if performance_metrics.get("success", False):381 # Get the data from the metrics382 with self.db_connection._engine.connect() as conn:383 df = pd.read_sql(sql_query, conn)384 385 # Calculate overall processing time386 total_time_ms = round((time.time() - start_time) * 1000, 2)387 388 # Add query to history389 history_item = {390 "question": question,391 "sql_query": sql_query,392 "execution_time_ms": performance_metrics["execution_time_ms"],393 "total_time_ms": total_time_ms,394 "row_count": performance_metrics["row_count"],395 "complexity": performance_metrics["complexity"]["level"],396 "timestamp": datetime.now().isoformat()397 }398 self.add_to_query_history(history_item)399 400 # Generate performance and complexity insights401 complexity = performance_metrics["complexity"]402 perf_insights = f"\n\n📊 Query Metrics:\n"403 perf_insights += f"• Execution time: {performance_metrics['execution_time_ms']}ms\n"404 perf_insights += f"• Total processing time: {total_time_ms}ms\n"405 perf_insights += f"• Rows returned: {performance_metrics['row_count']}\n"406 perf_insights += f"• Complexity: {complexity['level']}\n"407 408 if complexity["insights"]:409 perf_insights += "\n🔍 Insights:\n"410 for insight in complexity["insights"]:411 perf_insights += f"• {insight}\n"412 413 if not df.empty:414 # Format table in RDBMS style415 table_html = df.to_html(index=False, classes="table table-bordered table-striped", border=0)416 # Add custom styling to make it look more like RDBMS output417 table_html = f"""418 <style>419 .table-bordered {{420 border-collapse: collapse;421 width: 100%;422 font-family: 'Courier New', Courier, monospace;423 }}424 .table-bordered th {{425 background-color: #f2f2f2;426 color: #333;427 font-weight: bold;428 text-align: left;429 padding: 8px;430 border: 1px solid #ddd;431 }}432 .table-bordered td {{433 padding: 8px;434 border: 1px solid #ddd;435 }}436 .table-striped tbody tr:nth-of-type(odd) {{437 background-color: rgba(0,0,0,.05);438 }}439 </style>440 {table_html}441 """442 return f"✅ Query Successful\n\n{result}{warning_message}{perf_insights}", table_html443 else:444 # If performance measurement failed, continue with normal execution445 with self.db_connection._engine.connect() as conn:446 df = pd.read_sql(sql_query, conn)447 448 if not df.empty:449 # Format table in RDBMS style450 table_html = df.to_html(index=False, classes="table table-bordered table-striped", border=0)451 # Add custom styling452 table_html = f"""453 <style>454 .table-bordered {{455 border-collapse: collapse;456 width: 100%;457 font-family: 'Courier New', Courier, monospace;458 }}459 .table-bordered th {{460 background-color: #f2f2f2;461 color: #333;462 font-weight: bold;463 text-align: left;464 padding: 8px;465 border: 1px solid #ddd;466 }}467 .table-bordered td {{468 padding: 8px;469 border: 1px solid #ddd;470 }}471 .table-striped tbody tr:nth-of-type(odd) {{472 background-color: rgba(0,0,0,.05);473 }}474 </style>475 {table_html}476 """477 return f"✅ Query Successful\n\n{result}{warning_message}", table_html478 479 except SQLAlchemyError as exec_error:480 # If execution fails, try to fix the query again with the specific error481 if not correction_applied:482 schema_info = self.db_context["tables"] if self.db_context else None483 corrected_query = self.fix_sql_query(sql_query, str(exec_error), schema_info)484 485 # Try executing the corrected query486 try:487 # Measure performance of corrected query488 performance_metrics = self.measure_query_performance(corrected_query)489 490 if performance_metrics.get("success", False):491 # Get the data from the metrics492 with self.db_connection._engine.connect() as conn:493 df = pd.read_sql(corrected_query, conn)494 495 # Calculate overall processing time496 total_time_ms = round((time.time() - start_time) * 1000, 2)497 498 # Add query to history499 history_item = {500 "question": question,501 "sql_query": corrected_query,502 "execution_time_ms": performance_metrics["execution_time_ms"],503 "total_time_ms": total_time_ms,504 "row_count": performance_metrics["row_count"],505 "complexity": performance_metrics["complexity"]["level"],506 "timestamp": datetime.now().isoformat(),507 "corrected": True,508 "original_query": sql_query509 }510 self.add_to_query_history(history_item)511 512 # Generate performance and complexity insights513 complexity = performance_metrics["complexity"]514 perf_insights = f"\n\n📊 Query Metrics:\n"515 perf_insights += f"• Execution time: {performance_metrics['execution_time_ms']}ms\n"516 perf_insights += f"• Total processing time: {total_time_ms}ms\n"517 perf_insights += f"• Rows returned: {performance_metrics['row_count']}\n"518 perf_insights += f"• Complexity: {complexity['level']}\n"519 520 if complexity["insights"]:521 perf_insights += "\n🔍 Insights:\n"522 for insight in complexity["insights"]:523 perf_insights += f"• {insight}\n"524 525 if not df.empty:526 # Format table in RDBMS style527 table_html = df.to_html(index=False, classes="table table-bordered table-striped", border=0)528 # Add custom styling529 table_html = f"""530 <style>531 .table-bordered {{532 border-collapse: collapse;533 width: 100%;534 font-family: 'Courier New', Courier, monospace;535 }}536 .table-bordered th {{537 background-color: #f2f2f2;538 color: #333;539 font-weight: bold;540 text-align: left;541 padding: 8px;542 border: 1px solid #ddd;543 }}544 .table-bordered td {{545 padding: 8px;546 border: 1px solid #ddd;547 }}548 .table-striped tbody tr:nth-of-type(odd) {{549 background-color: rgba(0,0,0,.05);550 }}551 </style>552 {table_html}553 """554 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_html555 else:556 # If performance measurement failed, continue with normal execution557 with self.db_connection._engine.connect() as conn:558 df = pd.read_sql(corrected_query, conn)559 560 if not df.empty:561 # Format table in RDBMS style562 table_html = df.to_html(index=False, classes="table table-bordered table-striped", border=0)563 # Add custom styling564 table_html = f"""565 <style>566 .table-bordered {{567 border-collapse: collapse;568 width: 100%;569 font-family: 'Courier New', Courier, monospace;570 }}571 .table-bordered th {{572 background-color: #f2f2f2;573 color: #333;574 font-weight: bold;575 text-align: left;576 padding: 8px;577 border: 1px solid #ddd;578 }}579 .table-bordered td {{580 padding: 8px;581 border: 1px solid #ddd;582 }}583 .table-striped tbody tr:nth-of-type(odd) {{584 background-color: rgba(0,0,0,.05);585 }}586 </style>587 {table_html}588 """589 return f"✅ Query Successful (after correction)\n\n{result}\n\n🔧 Query was automatically corrected due to execution error: {str(exec_error)}", table_html590 except Exception:591 # If correction fails, return the original error592 pass593 594 # Return the execution error in table format595 error_msg = f"The query failed to execute:\n\n{str(exec_error)}\n\nOriginal result:\n{result}"596 df_error = pd.DataFrame({"Error": [error_msg]})597 table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0)598 # Add custom styling599 table_html = f"""600 <style>601 .table-bordered {{602 border-collapse: collapse;603 width: 100%;604 font-family: 'Courier New', Courier, monospace;605 }}606 .table-bordered th {{607 background-color: #f2f2f2;608 color: #333;609 font-weight: bold;610 text-align: left;611 padding: 8px;612 border: 1px solid #ddd;613 }}614 .table-bordered td {{615 padding: 8px;616 border: 1px solid #ddd;617 }}618 .table-striped tbody tr:nth-of-type(odd) {{619 background-color: rgba(0,0,0,.05);620 }}621 </style>622 {table_html}623 """624 return "❌ SQL Execution Error", table_html625 626 except SQLAlchemyError as sql_error:627 # Handle SQL execution errors628 error_details = str(sql_error)629 error_msg = f"❌ SQL Execution Error"630 details = f"The query failed to execute:\n\n{error_details}\n\nOriginal result:\n{result}"631 df_error = pd.DataFrame({"Error": [details]})632 table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0)633 # Add custom styling634 table_html = f"""635 <style>636 .table-bordered {{637 border-collapse: collapse;638 width: 100%;639 font-family: 'Courier New', Courier, monospace;640 }}641 .table-bordered th {{642 background-color: #f2f2f2;643 color: #333;644 font-weight: bold;645 text-align: left;646 padding: 8px;647 border: 1px solid #ddd;648 }}649 .table-bordered td {{650 padding: 8px;651 border: 1px solid #ddd;652 }}653 .table-striped tbody tr:nth-of-type(odd) {{654 background-color: rgba(0,0,0,.05);655 }}656 </style>657 {table_html}658 """659 return error_msg, table_html660 except Exception as table_error:661 # If table extraction fails, just return the text result662 pass663 664 # If we got here, we just have the text result without structured data665 # Convert to table format with RDBMS style666 df_text = pd.DataFrame({"Result": [result]})667 table_html = df_text.to_html(index=False, classes="table table-bordered table-striped", border=0)668 # Add custom styling669 table_html = f"""670 <style>671 .table-bordered {{672 border-collapse: collapse;673 width: 100%;674 font-family: 'Courier New', Courier, monospace;675 }}676 .table-bordered th {{677 background-color: #f2f2f2;678 color: #333;679 font-weight: bold;680 text-align: left;681 padding: 8px;682 border: 1px solid #ddd;683 }}684 .table-bordered td {{685 padding: 8px;686 border: 1px solid #ddd;687 }}688 .table-striped tbody tr:nth-of-type(odd) {{689 background-color: rgba(0,0,0,.05);690 }}691 </style>692 {table_html}693 """694 695 # Add to history696 history_item = {697 "question": question,698 "result": result,699 "timestamp": datetime.now().isoformat()700 }701 self.add_to_query_history(history_item)702 703 return f"✅ Query Successful", table_html704 705 except Exception as e:706 # Convert exception to table format with RDBMS style707 error_msg = f"❌ Query Failed: {str(e)}"708 details = f"Error processing query: {traceback.format_exc()}"709 df_error = pd.DataFrame({"Error": [details]})710 table_html = df_error.to_html(index=False, classes="table table-bordered table-striped", border=0)711 # Add custom styling712 table_html = f"""713 <style>714 .table-bordered {{715 border-collapse: collapse;716 width: 100%;717 font-family: 'Courier New', Courier, monospace;718 }}719 .table-bordered th {{720 background-color: #f2f2f2;721 color: #333;722 font-weight: bold;723 text-align: left;724 padding: 8px;725 border: 1px solid #ddd;726 }}727 .table-bordered td {{728 padding: 8px;729 border: 1px solid #ddd;730 }}731 .table-striped tbody tr:nth-of-type(odd) {{732 background-color: rgba(0,0,0,.05);733 }}734 </style>735 {table_html}736 """737 return error_msg, table_html738 739 def analyze_query_complexity(self, sql_query: str) -> Dict[str, Any]:740 """741 Analyze SQL query complexity and provide insights742 743 Args:744 sql_query: SQL query to analyze745 746 Returns:747 Dictionary with complexity metrics and insights748 """749 if not sql_query or not isinstance(sql_query, str):750 return {"error": "Invalid query provided"}751 752 sql_query = sql_query.strip().upper()753 754 # Initialize complexity metrics755 complexity = {756 "level": "Simple",757 "score": 0,758 "joins": 0,759 "tables": [],760 "aggregations": False,761 "grouping": False,762 "ordering": False,763 "limiting": False,764 "subqueries": 0,765 "complex_functions": [],766 "insights": []767 }768 769 # Count number of JOINs770 join_count = len(re.findall(r'\bJOIN\b', sql_query))771 complexity["joins"] = join_count772 if join_count > 0:773 complexity["score"] += join_count * 2774 if join_count >= 3:775 complexity["insights"].append(f"Query uses {join_count} joins, which may impact performance")776 777 # Detect tables used778 from_clause = re.search(r'\bFROM\b\s+(.*?)(?:\bWHERE\b|\bGROUP\b|\bHAVING\b|\bORDER\b|\bLIMIT\b|$)', sql_query)779 if from_clause:780 # Extract table names from FROM clause781 tables_text = from_clause.group(1).strip()782 # Handle JOIN syntax in FROM clause783 tables = re.findall(r'([a-zA-Z0-9_]+)(?:\s+(?:AS\s+)?[a-zA-Z0-9_]+)?', tables_text)784 complexity["tables"] = list(set(tables)) # Remove duplicates785 786 # Check for aggregations787 agg_functions = ["COUNT", "SUM", "AVG", "MIN", "MAX"]788 for func in agg_functions:789 if re.search(rf'\b{func}\s*\(', sql_query):790 complexity["aggregations"] = True791 complexity["score"] += 1792 break793 794 # Check for GROUP BY795 if re.search(r'\bGROUP\s+BY\b', sql_query):796 complexity["grouping"] = True797 complexity["score"] += 2798 799 # Check for ORDER BY800 if re.search(r'\bORDER\s+BY\b', sql_query):801 complexity["ordering"] = True802 complexity["score"] += 1803 804 # Check for LIMIT805 if re.search(r'\bLIMIT\b', sql_query):806 complexity["limiting"] = True807 complexity["score"] += 0.5808 809 # Check for subqueries810 subquery_count = len(re.findall(r'\(\s*SELECT', sql_query))811 complexity["subqueries"] = subquery_count812 if subquery_count > 0:813 complexity["score"] += subquery_count * 3814 complexity["insights"].append(f"Query contains {subquery_count} subqueries, which may affect performance")815 816 # Check for complex functions817 complex_funcs = ["CASE", "COALESCE", "NULLIF", "CAST", "CONVERT", "SUBSTRING", "CONCAT", "DATE_FORMAT", "EXTRACT"]818 for func in complex_funcs:819 if re.search(rf'\b{func}\b', sql_query):820 complexity["complex_functions"].append(func)821 complexity["score"] += 1822 823 # Determine complexity level824 if complexity["score"] <= 2:825 complexity["level"] = "Simple"826 elif complexity["score"] <= 5:827 complexity["level"] = "Moderate"828 elif complexity["score"] <= 10:829 complexity["level"] = "Complex"830 else:831 complexity["level"] = "Very Complex"832 complexity["insights"].append("This is a highly complex query that may benefit from optimization")833 834 # Add insights based on complexity835 if complexity["level"] in ["Complex", "Very Complex"] and not complexity["limiting"]:836 complexity["insights"].append("Consider adding a LIMIT clause to prevent large result sets")837 838 if complexity["joins"] >= 2 and not any(idx for idx in complexity["insights"] if "index" in idx.lower()):839 complexity["insights"].append("Ensure proper indexes exist on join columns")840 841 return complexity842 843 def add_to_query_history(self, query_data: Dict[str, Any]) -> None:844 """845 Add a query to the history846 847 Args:848 query_data: Dictionary containing query information849 """850 # Add timestamp if not present851 if "timestamp" not in query_data:852 query_data["timestamp"] = datetime.now().isoformat()853 854 # Add to history (at the beginning for most recent first)855 self.query_history.insert(0, query_data)856 857 # Trim history if needed858 if len(self.query_history) > self.max_history_items:859 self.query_history = self.query_history[:self.max_history_items]860 861 def get_query_history(self) -> List[Dict[str, Any]]:862 """863 Get the query history864 865 Returns:866 List of query history items867 """868 return self.query_history869 870 def clear_query_history(self) -> None:871 """Clear the query history"""872 self.query_history = []873 874 def measure_query_performance(self, sql_query: str) -> Dict[str, Any]:875 """876 Measure the performance of a SQL query877 878 Args:879 sql_query: SQL query to execute and measure880 881 Returns:882 Dictionary with performance metrics883 """884 if not self.db_connection:885 return {"error": "No database connection"}886 887 metrics = {888 "query": sql_query,889 "execution_time_ms": 0,890 "row_count": 0,891 "success": False,892 "error": None893 }894 895 try:896 # Measure execution time897 start_time = time.time()898 899 with self.db_connection._engine.connect() as conn:900 result = conn.execute(text(sql_query))901 # Convert to DataFrame to get row count902 df = pd.DataFrame(result.fetchall(), columns=result.keys())903 904 end_time = time.time()905 906 # Calculate metrics907 metrics["execution_time_ms"] = round((end_time - start_time) * 1000, 2)908 metrics["row_count"] = len(df)909 metrics["success"] = True910 911 # Add complexity analysis912 metrics["complexity"] = self.analyze_query_complexity(sql_query)913 914 return metrics915 916 except Exception as e:917 metrics["error"] = str(e)918 return metrics919 920 def generate_schema_diagram(self, include_all_tables: bool = False) -> str:921 """922 Generate a Mermaid ER diagram for the database schema923 924 Args:925 include_all_tables: Whether to include all tables or just a subset926 927 Returns:928 Mermaid diagram code929 """930 # Return a message that this functionality is not available931 return "This functionality has been removed"932 933# Initialize the database manager934db_manager = DatabaseManager()935 936def create_interface():937 """Create the Gradio interface"""938 939 with gr.Blocks(title="AI Database Query Assistant", theme=gr.themes.Soft()) as demo:940 # Warning banner at the top941 gr.Markdown("""942 <div style="background-color: #FFF3CD; color: #856404; padding: 15px; border-radius: 5px; border: 1px solid #FFEEBA; margin-bottom: 20px; font-weight: bold; text-align: center;">943 ⚠️ 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.git944 </div>945 """)946 947 gr.Markdown("""948 # 🤖 AI Database Query Assistant949 950 Connect to your MySQL database and query it using natural language!951 952 ### Steps:953 1. **Connect** to your database954 2. **Analyze** your database structure 955 3. **Ask questions** in natural language956 """)957 958 # Connection Status959 connection_status = gr.Textbox(960 label="Connection Status", 961 value="Not Connected", 962 interactive=False963 )964 965 with gr.Tabs():966 # MySQL Connection Tab967 with gr.TabItem("MySQL Connection"):968 gr.Markdown("""969 **MySQL Connection Details:**970 - Enter your MySQL server connection details971 - Password will be securely handled (not stored)972 - Default port is 3306 if not specified973 - Special characters in passwords are supported974 """)975 976 with gr.Row():977 mysql_host = gr.Textbox(978 label="Host", 979 value="localhost", 980 placeholder="localhost or IP address"981 )982 mysql_port = gr.Textbox(983 label="Port", 984 value="3306", 985 placeholder="3306"986 )987 988 with gr.Row():989 mysql_username = gr.Textbox(990 label="Username", 991 placeholder="root or your username"992 )993 mysql_password = gr.Textbox(994 label="DB Password(optional)", 995 type="password",996 placeholder="Your MySQL password"997 )998 999 mysql_database = gr.Textbox(1000 label="Database Name", 1001 placeholder="my_database"1002 )1003 1004 mysql_connect_btn = gr.Button("Connect to MySQL", variant="primary")1005 mysql_message = gr.Textbox(label="Connection Message", interactive=False)1006 1007 # API Key Section1008 gr.Markdown("""1009 ### 🔑 Google API Key1010 Enter your Google API key for Gemini model. If not provided, will use environment variable.1011 """)1012 1013 with gr.Row():1014 api_key_input = gr.Textbox(1015 label="Google API Key", 1016 type="password",1017 placeholder="Enter your Gemini API key here",1018 info="Get your API key from: https://makersuite.google.com/app/apikey"1019 )1020 api_key_btn = gr.Button("Set API Key", variant="secondary")1021 1022 api_key_message = gr.Textbox(label="API Key Status", interactive=False)1023 1024 # Database Analysis Section1025 with gr.Tabs():1026 with gr.TabItem("Database Analysis"):1027 gr.Markdown("## 🔍 Database Analysis")1028 analyze_btn = gr.Button("Analyze Database", variant="secondary", size="lg")1029 1030 with gr.Row():1031 analysis_status = gr.Textbox(label="Analysis Status", interactive=False)1032 analysis_details = gr.Textbox(label="Analysis Details", lines=10, interactive=False)1033 1034 # Schema Visualization1035 gr.Markdown("### 📊 Database Schema Visualization")1036 with gr.Row():1037 schema_table_select = gr.Dropdown(label="Select Table", choices=[], interactive=True)1038 visualize_schema_btn = gr.Button("Visualize Schema", variant="secondary")1039 1040 schema_output = gr.HTML(label="Schema Visualization")1041 1042 # Removed ER Diagram Visualization section1043 1044 # Query Section1045 with gr.TabItem("Query Database"):1046 gr.Markdown("## 💬 Ask Questions")1047 1048 question_input = gr.Textbox(1049 label="Your Question",1050 placeholder="Example: Show me all customers from New York, What are the top 5 selling products?",1051 lines=21052 )1053 1054 query_btn = gr.Button("Ask Question", variant="primary", size="lg")1055 1056 with gr.Row():1057 query_status = gr.Textbox(label="Query Result", lines=5, interactive=False)1058 query_output = gr.HTML(label="Data Output")1059 1060 # Example questions1061 gr.Markdown("""1062 ### 💡 Example Questions:1063 - "Show me all users registered in the last month"1064 - "What are the top 5 products by sales?"1065 - "How many orders were placed yesterday?"1066 - "Show me customers with more than 10 orders"1067 - "What's the average order value?"1068 """)1069 1070 # Query History Tab1071 with gr.TabItem("Query History"):1072 gr.Markdown("## 📜 Query History")1073 1074 with gr.Row():1075 refresh_history_btn = gr.Button("Refresh History", variant="secondary")1076 clear_history_btn = gr.Button("Clear History", variant="secondary")1077 1078 history_output = gr.HTML(label="Query History")1079 1080 # Reuse Query Section1081 gr.Markdown("### 🔄 Reuse Previous Query")1082 with gr.Row():1083 history_question_select = gr.Dropdown(label="Select Previous Question", choices=[], interactive=True)1084 reuse_query_btn = gr.Button("Use Selected Query", variant="primary")1085 1086 # Event handlers1087 mysql_connect_btn.click(1088 fn=lambda h, p, u, pw, d: db_manager.connect_mysql(h, p, u, pw, d) + (db_manager.connection_status,),1089 inputs=[mysql_host, mysql_port, mysql_username, mysql_password, mysql_database],1090 outputs=[mysql_message, connection_status]1091 )1092 1093 # API Key event handler1094 api_key_btn.click(1095 fn=db_manager.set_api_key,1096 inputs=[api_key_input],1097 outputs=[api_key_message]1098 )1099 1100 # Database analysis event handler1101 def on_analyze_database():1102 status, details = db_manager.analyze_database()1103 1104 # Update schema table dropdown if analysis was successful1105 table_choices = []1106 if "✅" in status and db_manager.db_context:1107 table_choices = list(db_manager.db_context.get("tables", {}).keys())1108 1109 return status, details, gr.Dropdown(choices=table_choices)1110 1111 analyze_btn.click(1112 fn=on_analyze_database,1113 outputs=[analysis_status, analysis_details, schema_table_select]1114 )1115 1116 # Schema visualization event handler1117 def visualize_table_schema(table_name):1118 if not table_name or not db_manager.db_context or table_name not in db_manager.db_context.get("tables", {}):1119 return "<p>Please select a valid table</p>"1120 1121 table_info = db_manager.db_context["tables"][table_name]1122 1123 # Create HTML visualization1124 html = f"<h3>Table: {table_name}</h3>"1125 html += f"<p>Row count: {table_info.get('row_count', 'Unknown')}</p>"1126 1127 # Create table for columns1128 html += "<table class='table table-bordered table-striped'>"1129 html += "<thead><tr><th>Column</th><th>Type</th><th>Key</th></tr></thead>"1130 html += "<tbody>"1131 1132 # Add columns1133 primary_keys = table_info.get("primary_keys", [])1134 foreign_keys_flat = []1135 1136 # Flatten foreign key references1137 for fk in table_info.get("foreign_keys", []):1138 for col in fk.get("columns", []):1139 foreign_keys_flat.append(col)1140 1141 for col in table_info.get("columns", []):1142 col_name = col.get("name", "")1143 col_type = col.get("type", "")1144 1145 # Determine key type1146 key_type = ""1147 if col_name in primary_keys:1148 key_type = "🔑 Primary"1149 elif col_name in foreign_keys_flat:1150 key_type = "🔗 Foreign"1151 1152 html += f"<tr><td>{col_name}</td><td>{col_type}</td><td>{key_type}</td></tr>"1153 1154 html += "</tbody></table>"1155 1156 # Add foreign key relationships1157 if table_info.get("foreign_keys"):1158 html += "<h4>Foreign Key Relationships</h4>"1159 html += "<ul>"1160 for fk in table_info.get("foreign_keys", []):1161 cols = ", ".join(fk.get("columns", []))1162 refs = fk.get("refers_to", "")1163 html += f"<li>{cols} → {refs}</li>"1164 html += "</ul>"1165 1166 return html1167 1168 visualize_schema_btn.click(1169 fn=visualize_table_schema,1170 inputs=[schema_table_select],1171 outputs=[schema_output]1172 )1173 1174 # Removed ER Diagram event handler1175 1176 # Query event handler1177 query_btn.click(1178 fn=db_manager.query_database,1179 inputs=[question_input],1180 outputs=[query_status, query_output]1181 )1182 1183 # Query history event handlers1184 def format_query_history():1185 history = db_manager.get_query_history()1186 if not history:1187 return "<p>No queries in history</p>", gr.Dropdown(choices=[])1188 1189 # Format history as HTML table1190 html = "<table class='table table-bordered table-striped'>"1191 html += "<thead><tr><th>Time</th><th>Question</th><th>SQL Query</th><th>Execution Time</th><th>Rows</th><th>Complexity</th></tr></thead>"1192 html += "<tbody>"1193 1194 # Collect questions for dropdown1195 questions = []1196 1197 for i, item in enumerate(history):1198 # Format timestamp1199 timestamp = item.get("timestamp", "")1200 if timestamp: