app.py

import gradio as gr
import sqlite3
import json
import pandas as pd
from openai import OpenAI
import traceback
from typing import Dict, List, Tuple, Any
import re
from datetime import datetime
import threading
import queue
import html
import sys
import os

# Force stdout to use UTF-8 encoding to handle Unicode characters
if sys.stdout.encoding != 'utf-8':
    sys.stdout = open(sys.stdout.fileno(), mode='w', encoding='utf-8', buffering=1)

class DatabaseQueryAgent:
    def __init__(self, db_path: str = "innovativeskills.db"):
        self.db_path = db_path
        self.client = None
        
        # Available models
        self.models = {
            "llama": "meta-llama/llama-3.3-70b-instruct:free",
            "mistral": "mistralai/mistral-7b-instruct:free",
            "gemma": "google/gemma-2-9b-it:free"  # Verification model
        }
        
        # Initialize database connection
        self.init_db_connection()
        
    def init_db_connection(self):
        """Initialize database connection with UTF-8 encoding"""
        try:
            conn = sqlite3.connect(self.db_path, check_same_thread=False)
            conn.execute("PRAGMA encoding = 'UTF-8';")
            cursor = conn.cursor()
            
            # Load table metadata
            self.table_metadata = self.get_table_metadata(conn, cursor)
            self.column_metadata = self.get_column_metadata(conn, cursor)
            self.actual_schema = self.get_actual_schema(conn, cursor)
            
            conn.close()
            
        except Exception as e:
            print(f"Database initialization error: {e}")
            self.table_metadata = {}
            self.column_metadata = {}
            self.actual_schema = {}
    
    def get_db_connection(self):
        """Get a new database connection with UTF-8 encoding"""
        conn = sqlite3.connect(self.db_path, check_same_thread=False)
        conn.execute("PRAGMA encoding = 'UTF-8';")
        return conn
    
    def get_actual_schema(self, conn, cursor) -> Dict:
        """Get actual database schema"""
        try:
            cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name NOT LIKE 'sqlite_%'")
            tables = [row[0] for row in cursor.fetchall()]
            schema = {}
            for table in tables:
                cursor.execute(f"PRAGMA table_info({table})")
                columns = cursor.fetchall()
                try:
                    cursor.execute(f"SELECT * FROM {table} LIMIT 3")
                    sample_data = cursor.fetchall()
                except Exception:
                    sample_data = []
                try:
                    cursor.execute(f"SELECT COUNT(*) FROM {table}")
                    row_count = cursor.fetchone()[0]
                except Exception:
                    row_count = 0
                schema[table] = {
                    'columns': [{'name': col[1], 'type': col[2], 'notnull': col[3], 'pk': col[5]} for col in columns],
                    'sample_data': sample_data,
                    'row_count': row_count
                }
            return schema
        except Exception as e:
            print(f"Error getting actual schema: {e}")
            return {}
    
    def get_table_metadata(self, conn, cursor) -> Dict:
        """Get table metadata"""
        try:
            query = """
            SELECT table_name, domain, description, row_count 
            FROM table_catalog 
            WHERE table_name NOT IN ('table_catalog', 'column_catalog')
            """
            results = cursor.execute(query).fetchall()
            metadata = {}
            for table_name, domain, description, row_count in results:
                metadata[table_name] = {
                    'domain': domain,
                    'description': description,
                    'row_count': row_count
                }
            return metadata
        except Exception as e:
            print(f"Error loading table metadata: {e}")
            return {}
    
    def get_column_metadata(self, conn, cursor) -> Dict:
        """Get column metadata"""
        try:
            query = """
            SELECT table_name, column_name, data_type, is_foreign_key, references_table, description
            FROM column_catalog
            """
            results = cursor.execute(query).fetchall()
            metadata = {}
            for table_name, column_name, data_type, is_fk, ref_table, description in results:
                if table_name not in metadata:
                    metadata[table_name] = []
                metadata[table_name].append({
                    'name': column_name,
                    'type': data_type,
                    'is_foreign_key': bool(is_fk),
                    'references': ref_table,
                    'description': description
                })
            return metadata
        except Exception as e:
            print(f"Error loading column metadata: {e}")
            return {}
    
    def setup_client(self, api_key: str):
        """Setup OpenRouter client"""
        self.client = OpenAI(
            base_url="https://openrouter.ai/api/v1",
            api_key=api_key,
        )
    
    def get_relevant_tables_for_query(self, query: str) -> str:
        """Analyze query and return relevant table info"""
        query_lower = query.lower()
        relevant_tables = []
        keywords = {
            'customer': ['customer', 'client', 'buyer', 'user'],
            'order': ['order', 'purchase', 'transaction', 'sale'],
            'product': ['product', 'item', 'inventory', 'stock'],
            'employee': ['employee', 'staff', 'worker', 'personnel'],
            'patient': ['patient', 'medical', 'health'],
            'student': ['student', 'enrollment', 'grade', 'course'],
            'supplier': ['supplier', 'vendor', 'provider'],
            'shipping': ['shipping', 'delivery', 'logistics'],
            'payment': ['payment', 'invoice', 'billing'],
            'account': ['account', 'financial', 'balance']
        }
        for concept, search_terms in keywords.items():
            if any(term in query_lower for term in search_terms):
                for table_name in self.actual_schema.keys():
                    table_lower = table_name.lower()
                    if any(term in table_lower for term in search_terms):
                        if table_name not in relevant_tables:
                            relevant_tables.append(table_name)
        if not relevant_tables:
            relevant_tables = [name for name, info in self.actual_schema.items() 
                             if info['row_count'] > 10][:10]
        schema_info = ""
        for table in relevant_tables[:15]:
            if table in self.actual_schema:
                info = self.actual_schema[table]
                columns_str = ", ".join([f"{col['name']}({col['type']})" for col in info['columns']])
                schema_info += f"\nTable: {table}\n"
                schema_info += f"  Columns: {columns_str}\n"
                schema_info += f"  Rows: {info['row_count']}\n"
                if table in self.table_metadata:
                    meta = self.table_metadata[table]
                    schema_info += f"  Domain: {meta['domain']}\n"
                    schema_info += f"  Description: {meta['description']}\n"
                if info['sample_data']:
                    schema_info += f"  Sample: {info['sample_data'][0] if info['sample_data'] else 'No data'}\n"
        return schema_info
    
    def get_system_prompt(self, user_query: str) -> str:
        """Generate system prompt with actual schema"""
        relevant_schema = self.get_relevant_tables_for_query(user_query)
        return f"""You are an intelligent database query agent that specializes in identifying relevant tables and generating accurate SQL queries.

DATABASE SCHEMA INFORMATION:
{relevant_schema}

CRITICAL SQL RULES:
1. NEVER use reserved words as table aliases (like 'to', 'from', 'where', 'select', etc.)
2. Use descriptive aliases like 'cust', 'ord', 'prod' instead
3. Only JOIN tables if you can identify a logical relationship between them
4. If no clear JOIN relationship exists, use separate SELECT statements or UNION
5. Always use the EXACT column names shown in the schema
6. Do not assume foreign key relationships unless explicitly shown

CRITICAL: You MUST respond with ONLY a valid JSON object. No markdown, no explanations outside the JSON. 

Your response must be exactly in this JSON format:
{{
  "analysis": "Brief analysis of the query and table selection reasoning",
  "identified_tables": ["table1", "table2", "table3"],
  "domains_involved": ["domain1", "domain2"],
  "sql_query": "SELECT ... FROM ... WHERE ...",
  "explanation": "Step-by-step explanation of the query logic",
  "confidence": 0.95,
  "alternative_queries": ["Alternative SQL if applicable"]
}}

IMPORTANT RULES:
1. Respond with ONLY valid JSON - no markdown formatting
2. Use ONLY the actual table names shown in the schema above
3. Use ONLY the actual column names shown in the schema above
4. Generate syntactically correct SQL queries with proper aliases
5. Focus on tables that actually exist and have relevant data
6. Include confidence scores between 0.0 and 1.0
7. Provide clear explanations
8. Ensure table names in 'identified_tables' match those used in 'sql_query'
9. Check that columns referenced in SQL actually exist in the tables
10. If no perfect match exists, choose the closest relevant tables and explain the compromise
11. Avoid reserved word aliases like 'to', 'from', 'order', 'select'

QUERY ANALYSIS GUIDELINES:
- For customer/order queries: Look for tables with customer-related or order-related names and columns
- For employee queries: Look for tables with employee, staff, or HR-related names
- For product queries: Look for tables with product, inventory, or item-related names
- Always verify column names exist before using them in SQL
- Use proper JOIN syntax when combining tables, but only if logical relationships exist
- Include appropriate WHERE clauses when filtering is implied
- If unsure about relationships, prefer simpler queries or multiple separate queries"""

    def extract_json_from_response(self, response_text: str) -> Dict:
        """Extract JSON from response text"""
        try:
            return json.loads(response_text)
        except json.JSONDecodeError:
            json_pattern = r'```json\s*(.*?)\s*```'
            json_match = re.search(json_pattern, response_text, re.DOTALL)
            if json_match:
                try:
                    return json.loads(json_match.group(1))
                except json.JSONDecodeError:
                    pass
            json_pattern = r'\{.*\}'
            json_match = re.search(json_pattern, response_text, re.DOTALL)
            if json_match:
                try:
                    return json.loads(json_match.group(0))
                except json.JSONDecodeError:
                    pass
            return self.create_fallback_response(response_text)
    
    def create_fallback_response(self, response_text: str) -> Dict:
        """Create a fallback response when JSON parsing fails"""
        sql_pattern = r'SELECT.*?(?:;|$)'
        sql_match = re.search(sql_pattern, response_text, re.IGNORECASE | re.DOTALL)
        sql_query = sql_match.group(0).strip(';') if sql_match else ""
        identified_tables = [table_name for table_name in self.actual_schema.keys() 
                           if table_name.lower() in response_text.lower()]
        domains_involved = [self.table_metadata[table]['domain'] for table in identified_tables 
                          if table in self.table_metadata and self.table_metadata[table]['domain'] not in domains_involved]
        return {
            "analysis": "Fallback analysis from unparseable response",
            "identified_tables": identified_tables[:5],
            "domains_involved": domains_involved[:3],
            "sql_query": sql_query,
            "explanation": "Response could not be parsed as JSON, extracted information where possible",
            "confidence": 0.5,
            "alternative_queries": []
        }
    
    def validate_sql_query(self, sql_query: str, identified_tables: List[str]) -> Tuple[bool, str]:
        """Validate SQL query against schema"""
        try:
            if not sql_query.strip():
                return False, "Empty SQL query"
            for table in identified_tables:
                if table not in self.actual_schema:
                    return False, f"Table '{table}' does not exist in database"
            sql_upper = sql_query.upper()
            if not sql_upper.strip().startswith('SELECT'):
                return False, "Only SELECT queries are allowed"
            reserved_words = ['TO', 'FROM', 'WHERE', 'SELECT', 'ORDER', 'GROUP', 'HAVING', 'UNION', 'JOIN', 'ON']
            alias_pattern = r'(?:FROM|JOIN)\s+(\w+)\s+(\w+)'
            aliases = re.findall(alias_pattern, sql_query, re.IGNORECASE)
            for table, alias in aliases:
                if alias.upper() in reserved_words:
                    return False, f"Cannot use reserved word '{alias}' as table alias"
            for table in identified_tables:
                if table in sql_query:
                    table_info = self.actual_schema[table]
                    available_columns = [col['name'] for col in table_info['columns']]
                    column_patterns = [
                        rf'{re.escape(table)}\.(\w+)', 
                        rf'\b(\w+)\.(\w+)', 
                        rf'SELECT\s+([^FROM]+)'
                    ]
                    for pattern in column_patterns:
                        matches = re.findall(pattern, sql_query, re.IGNORECASE)
                        for match in matches:
                            if isinstance(match, tuple):
                                column = match[1] if len(match) == 2 else match[0] if match else ''
                            else:
                                column = match
                            if column.upper() in ['*', 'COUNT', 'SUM', 'AVG', 'MAX', 'MIN', 'DISTINCT']:
                                continue
                            if column and column not in available_columns and f'{table}.{column}' in sql_query:
                                return False, f"Column '{column}' does not exist in table '{table}'"
            return True, "Query validation passed"
        except Exception as e:
            return False, f"Validation error: {str(e)}"
    
    def call_model(self, model_key: str, prompt: str, user_query: str) -> Dict:
        """Call specific model with prompt"""
        try:
            messages = [
                {"role": "system", "content": prompt},
                {"role": "user", "content": f"Query: {user_query}\n\nRespond with ONLY a valid JSON object following the exact format specified in the system prompt."}
            ]
            completion = self.client.chat.completions.create(
                model=self.models[model_key],
                messages=messages,
                temperature=0.1,
                max_tokens=2000
            )
            response = completion.choices[0].message.content.strip()
            parsed_response = self.extract_json_from_response(response)
            sql_query = parsed_response.get('sql_query', '')
            identified_tables = parsed_response.get('identified_tables', [])
            if sql_query:
                is_valid, validation_message = self.validate_sql_query(sql_query, identified_tables)
                parsed_response['sql_validation'] = {
                    'is_valid': is_valid,
                    'message': validation_message
                }
            return {
                "success": True,
                "response": parsed_response,
                "raw_response": response,
                "model": model_key
            }
        except Exception as e:
            return {
                "success": False,
                "error": str(e),
                "model": model_key
            }
    
    def verify_response(self, api_key: str, original_query: str, llama_response: Dict, mistral_response: Dict) -> Dict:
        """Use Gemma to verify responses"""
        self.setup_client(api_key)
        relevant_schema = self.get_relevant_tables_for_query(original_query)
        verification_prompt = f"""You are a database query verification expert. You have access to the actual database schema and must verify responses against it.

ACTUAL DATABASE SCHEMA:
{relevant_schema}

ORIGINAL QUERY: {original_query}

LLAMA RESPONSE: {json.dumps(llama_response.get('response', {}), indent=2)}

MISTRAL RESPONSE: {json.dumps(mistral_response.get('response', {}), indent=2)}

Verify these responses against the ACTUAL schema above. Check:
1. Do the table names actually exist in the schema?
2. Do the column names actually exist in those tables?
3. Are the table selections appropriate for the query?
4. Is the SQL syntax correct?
5. Are table aliases proper (not reserved words)?

Respond with ONLY a valid JSON object:
{{
  "verification_summary": "Overall assessment based on actual schema",
  "table_selection_accuracy": "Assessment of table choices against actual schema",
  "sql_correctness": "SQL syntax and schema validation",
  "consistency_check": "Comparison between responses",
  "recommended_response": "llama, mistral, or neither",
  "confidence_score": 0.85,
  "suggested_improvements": ["improvement1", "improvement2"],
  "potential_issues": ["issue1", "issue2"],
  "schema_compliance": "Assessment of how well responses match actual schema"
}}"""
        return self.call_model("gemma", verification_prompt, "Verify the above responses against the actual database schema.")
    
    def execute_query_in_thread(self, sql_query: str, result_queue: queue.Queue):
        """Execute SQL query in a thread"""
        try:
            if not sql_query.strip().upper().startswith('SELECT'):
                result_queue.put((False, "Only SELECT queries are allowed"))
                return
            sql_query = sql_query.strip().rstrip(';')
            conn = self.get_db_connection()
            try:
                df = pd.read_sql_query(sql_query, conn)
                result_queue.put((True, df))
            except Exception as e:
                result_queue.put((False, str(e)))
            finally:
                conn.close()
        except Exception as e:
            result_queue.put((False, f"Query execution error: {str(e)}"))
    
    def execute_query(self, sql_query: str) -> Tuple[bool, Any]:
        """Execute SQL query using thread-safe approach"""
        try:
            result_queue = queue.Queue()
            thread = threading.Thread(
                target=self.execute_query_in_thread, 
                args=(sql_query, result_queue)
            )
            thread.start()
            thread.join(timeout=30)
            if thread.is_alive():
                return False, "Query execution timed out"
            if not result_queue.empty():
                return result_queue.get()
            else:
                return False, "No result returned from query execution"
        except Exception as e:
            return False, f"Execution error: {str(e)}"
    
    def process_query(self, api_key: str, user_query: str) -> Dict:
        """Process user query"""
        if not api_key:
            return {"error": "Please provide OpenRouter API key"}
        try:
            self.setup_client(api_key)
            system_prompt = self.get_system_prompt(user_query)
            llama_result = self.call_model("llama", system_prompt, user_query)
            mistral_result = self.call_model("mistral", system_prompt, user_query)
            verification_result = self.verify_response(api_key, user_query, llama_result, mistral_result)
            execution_results = {}
            for model_name, result in [("llama", llama_result), ("mistral", mistral_result)]:
                if result.get("success") and result.get("response", {}).get("sql_query"):
                    sql_query = result["response"]["sql_query"]
                    validation_info = result["response"].get("sql_validation", {})
                    if sql_query.strip():
                        if validation_info.get("is_valid", True):
                            success, data = self.execute_query(sql_query)
                            execution_results[model_name] = {
                                "success": success,
                                "data": data.to_dict('records') if success and isinstance(data, pd.DataFrame) else str(data),
                                "row_count": len(data) if success and isinstance(data, pd.DataFrame) else 0,
                                "sql_query": sql_query,
                                "validation": validation_info
                            }
                        else:
                            execution_results[model_name] = {
                                "success": False,
                                "data": f"Query validation failed: {validation_info.get('message', 'Unknown error')}",
                                "row_count": 0,
                                "sql_query": sql_query,
                                "validation": validation_info
                            }
                    else:
                        execution_results[model_name] = {
                            "success": False,
                            "data": "No SQL query generated",
                            "row_count": 0,
                            "sql_query": "",
                            "validation": {"is_valid": False, "message": "Empty query"}
                        }
                else:
                    execution_results[model_name] = {
                        "success": False,
                        "data": "Model failed to generate response",
                        "row_count": 0,
                        "sql_query": "",
                        "validation": {"is_valid": False, "message": "Model error"}
                    }
            return {
                "llama_response": llama_result,
                "mistral_response": mistral_result,
                "verification": verification_result,
                "execution_results": execution_results,
                "timestamp": datetime.now().isoformat(),
                "schema_info": self.get_relevant_tables_for_query(user_query)
            }
        except Exception as e:
            return {"error": f"Processing error: {str(e)}", "traceback": traceback.format_exc()}

def response_to_markdown(response_dict: Dict) -> str:
    """Convert model response to Markdown"""
    if not response_dict.get("success", False):
        return f"**Error**: {response_dict.get('error', 'Unknown error')}"
    response = response_dict.get("response", {})
    markdown = "**Query Analysis Results**\n\n"
    markdown += f"- **Analysis**: {response.get('analysis', 'N/A')}\n\n"
    identified_tables = response.get('identified_tables', [])
    markdown += f"- **Identified Tables**: {', '.join(identified_tables) if identified_tables else 'None'}\n\n"
    domains_involved = response.get('domains_involved', [])
    markdown += f"- **Domains Involved**: {', '.join(domains_involved) if domains_involved else 'None'}\n\n"
    sql_query = response.get('sql_query', '')
    if sql_query:
        markdown += "- **SQL Query**:\n\n```sql\n" + sql_query + "\n```\n\n"
    else:
        markdown += "- **SQL Query**: None\n\n"
    markdown += f"- **Explanation**: {response.get('explanation', 'N/A')}\n\n"
    markdown += f"- **Confidence**: {response.get('confidence', 'N/A')}\n\n"
    alternative_queries = response.get('alternative_queries', [])
    if alternative_queries:
        markdown += "- **Alternative Queries**:\n"
        for query in alternative_queries:
            markdown += f"  - {query}\n"
    else:
        markdown += "- **Alternative Queries**: None\n"
    validation = response.get('sql_validation', {})
    if validation:
        is_valid = validation.get('is_valid', False)
        message = validation.get('message', 'N/A')
        markdown += f"\n- **SQL Validation**: {'Passed' if is_valid else 'Failed'} - {message}\n"
    return markdown

def verification_to_markdown(verification_dict: Dict) -> str:
    """Convert verification response to Markdown"""
    if not verification_dict.get("success", False):
        return f"**Error**: {verification_dict.get('error', 'Unknown error')}"
    response = verification_dict.get("response", {})
    markdown = "**Verification Results**\n\n"
    markdown += f"- **Verification Summary**: {response.get('verification_summary', 'N/A')}\n\n"
    markdown += f"- **Table Selection Accuracy**: {response.get('table_selection_accuracy', 'N/A')}\n\n"
    markdown += f"- **SQL Correctness**: {response.get('sql_correctness', 'N/A')}\n\n"
    markdown += f"- **Consistency Check**: {response.get('consistency_check', 'N/A')}\n\n"
    markdown += f"- **Recommended Response**: {response.get('recommended_response', 'N/A')}\n\n"
    markdown += f"- **Confidence Score**: {response.get('confidence_score', 'N/A')}\n\n"
    suggested_improvements = response.get('suggested_improvements', [])
    if suggested_improvements:
        markdown += "- **Suggested Improvements**:\n"
        for improvement in suggested_improvements:
            markdown += f"  - {improvement}\n"
    else:
        markdown += "- **Suggested Improvements**: None\n"
    potential_issues = response.get('potential_issues', [])
    if potential_issues:
        markdown += "- **Potential Issues**:\n"
        for issue in potential_issues:
            markdown += f"  - {issue}\n"
    else:
        markdown += "- **Potential Issues**: None\n"
    markdown += f"- **Schema Compliance**: {response.get('schema_compliance', 'N/A')}\n"
    return markdown

def create_gradio_interface():
    """Create Gradio interface"""
    agent = DatabaseQueryAgent()
    sample_queries = [
        "Find all customers from customer tables",
        "Show me employee information from HR tables",
        "Get patient data from healthcare tables",
        "List all products with their details",
        "Find students enrolled in courses",
        "Show financial transaction records",
        "Get shipping information for deliveries",
        "Find all suppliers and their information",
        "Show retail store data",
        "Get manufacturing production records"
    ]
    
    def process_user_query(api_key, query):
        """Process query and return formatted results"""
        if not query.strip():
            return "Please enter a query", "", "", "", "", ""
        results = agent.process_query(api_key, query)
        if "error" in results:
            return f"**Error**: {results['error']}", "", "", "", "", ""
        
        # Format responses as Markdown
        llama_markdown = response_to_markdown(results.get("llama_response", {}))
        mistral_markdown = response_to_markdown(results.get("mistral_response", {}))
        verification_markdown = verification_to_markdown(results.get("verification", {}))
        
        # Format execution results
        exec_results = results.get("execution_results", {})
        execution_formatted = ""
        for model, result in exec_results.items():
            execution_formatted += f"\n=== {model.upper()} EXECUTION ===\n"
            execution_formatted += f"SQL Query: {result.get('sql_query', 'N/A')}\n"
            validation = result.get('validation', {})
            if validation.get('is_valid'):
                execution_formatted += f"✅ Query Validation: PASSED\n"
            else:
                execution_formatted += f"❌ Query Validation: FAILED - {validation.get('message', 'Unknown error')}\n"
            if result["success"]:
                execution_formatted += f"✅ Execution: Success! Retrieved {result['row_count']} rows\n"
                if result["row_count"] > 0:
                    sample_data = result['data'][:3] if isinstance(result['data'], list) else []
                    execution_formatted += f"Sample data:\n{json.dumps(sample_data, indent=2)}\n"
                else:
                    execution_formatted += "No data returned (empty result set)\n"
            else:
                execution_formatted += f"❌ Execution Error: {result['data']}\n"
            execution_formatted += "\n"
        if not execution_formatted:
            execution_formatted = "No queries were executed. Check if valid SQL was generated."
        
        schema_info = results.get('schema_info', 'No schema information available')
        
        # Format summary as Markdown
        verification_resp = results.get('verification', {}).get('response', {})
        summary = f"""
**🔍 QUERY ANALYSIS COMPLETE**

━━━━━━━━━━━━━━━━━━━━━━━━

**📊 Models Used**: Llama 3.1 8B, Mistral 7B, Gemma 2 9B (verification)

**⏰ Processed**: {results.get('timestamp', 'N/A')}

**🎯 Verification Summary**:

{verification_resp.get('verification_summary', 'N/A')}

**💡 Recommended Model**: {verification_resp.get('recommended_response', 'N/A')}

**📈 Confidence**: {verification_resp.get('confidence_score', 'N/A')}

**🗄️ Schema Compliance**: {verification_resp.get('schema_compliance', 'N/A')}

**🗄️ Query Execution Status**:

{len(exec_results)} queries attempted
        """
        
        return summary, llama_markdown, mistral_markdown, verification_markdown, execution_formatted, schema_info
    
    with gr.Blocks(
        title="Fixed Intelligent Database Query Agent",
        theme=gr.themes.Soft(),
        css="""
        .gradio-container {
            max-width: 1200px !important;
            margin: 0 auto !important;
        }
        .result-box {
            background-color: #f8f9fa;
            border: 1px solid #dee2e6;
            border-radius: 8px;
            padding: 15px;
        }
        """
    ) as interface:
        gr.HTML("""
        <div style="text-align: center; padding: 20px; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; border-radius: 10px; margin-bottom: 20px;">
            <h1>🤖 Fixed Intelligent Database Query Agent</h1>
            <p>AI-powered agent that intelligently selects relevant tables from 100+ tables and generates optimized SQL queries</p>
            <p><strong>Database:</strong> 100 tables across 10 business domains | <strong>Models:</strong> Llama 3.1 8B + Mistral 7B + Gemma 2 9B</p>
            <p><strong>✅ FIXED:</strong> Reserved Word Aliases | Enhanced Column Validation | Better SQL Syntax Checking</p>
        </div>
        """)
        
        with gr.Row():
            with gr.Column(scale=1):
                api_key_input = gr.Textbox(
                    label="🔑 OpenRouter API Key",
                    type="password",
                    placeholder="Enter your OpenRouter API key...",
                    info="Get your free API key from openrouter.ai"
                )
                query_input = gr.Textbox(
                    label="💬 Database Query",
                    placeholder="Enter your natural language query...",
                    lines=3,
                    info="Example: 'Find all customers who placed orders in the last month'"
                )
                with gr.Row():
                    submit_btn = gr.Button("🚀 Process Query", variant="primary", size="lg")
                    clear_btn = gr.Button("🗑️ Clear", variant="secondary")
                gr.HTML("<h3>📝 Sample Test Queries</h3>")
                sample_dropdown = gr.Dropdown(
                    choices=sample_queries,
                    label="Quick Test Examples",
                    info="Select a sample query to test the agent"
                )
        
            with gr.Column(scale=2):
                summary_output = gr.Markdown(label="📊 Analysis Summary")
                with gr.Tabs():
                    with gr.Tab("🦙 Llama 3.1 8B Response"):
                        llama_output = gr.Markdown(label="Llama Response")
                    with gr.Tab("🌟 Mistral 7B Response"):
                        mistral_output = gr.Markdown(label="Mistral Response")
                    with gr.Tab("✅ Verification (Gemma 2 9B)"):
                        verification_output = gr.Markdown(label="Verification Analysis")
                    with gr.Tab("🗄️ Query Execution Results"):
                        execution_output = gr.Textbox(
                            label="Database Execution Results",
                            lines=15,
                            max_lines=20,
                            elem_classes=["result-box"]
                        )
                    with gr.Tab("📋 Database Schema"):
                        schema_output = gr.Textbox(
                            label="Relevant Database Schema",
                            lines=15,
                            max_lines=20,
                            elem_classes=["result-box"]
                        )
        
        submit_btn.click(
            fn=process_user_query,
            inputs=[api_key_input, query_input],
            outputs=[summary_output, llama_output, mistral_output, verification_output, execution_output, schema_output]
        )
        clear_btn.click(
            fn=lambda: ("", "", "", "", "", "", ""),
            outputs=[query_input, summary_output, llama_output, mistral_output, verification_output, execution_output, schema_output]
        )
        sample_dropdown.change(
            fn=lambda x: x,
            inputs=[sample_dropdown],
            outputs=[query_input]
        )
        gr.HTML("""
        <div style="margin-top: 20px; padding: 15px; background-color: #f8f9fa; border-radius: 8px;">
            <h3>🎯 How to Use</h3>
            <ol>
                <li><strong>API Key:</strong> Get a free API key from <a href="https://openrouter.ai" target="_blank">openrouter.ai</a></li>
                <li><strong>Query:</strong> Enter your natural language database query</li>
                <li><strong>Process:</strong> The agent will analyze your query across 100+ tables and generate optimized SQL</li>
                <li><strong>Results:</strong> View responses from multiple AI models, verification analysis, and actual query execution results</li>
            </ol>
            <p><strong>Features:</strong></p>
            <ul>
                <li>🧠 Multi-model AI analysis (Llama, Mistral, Gemma)</li>
                <li>🔍 Intelligent table selection from 100+ tables</li>
                <li>✅ SQL validation and syntax checking</li>
                <li>🗄️ Real database query execution with results</li>
                <li>📊 Cross-model verification and comparison</li>
            </ul>
        </div>
        """)
    
    return interface

def main():
    """Main function to launch the application"""
    print("🚀 Starting Intelligent Database Query Agent...")
    print("📊 Loading database schema and metadata...")
    interface = create_gradio_interface()
    print("✅ Database Query Agent Ready!")
    print("🌐 Access the interface at: http://localhost:7860")
    print("🔑 Don't forget to add your OpenRouter API key!")
    interface.launch(share=True)

if __name__ == "__main__":
    main()