Add main app.py for text-to-sql agent

app.py

@@ -0,0 +1,438 @@
+"""
+Multi-Turn Text-to-SQL Agent with Clarification Capabilities
+=============================================================
+
+An intelligent SQL assistant that:
+- Answers clear database questions with accurate SQL
+- Detects ambiguous questions and asks targeted clarifications
+- Explains when questions can't be answered with available data
+- Self-corrects SQL errors via ReAct reasoning loop
+- Maintains multi-turn conversation context
+
+Architecture based on:
+- MMSQL (arXiv:2412.17867) — 4-type question classification
+- PRACTIQ (arXiv:2410.11076) — clarification dialogue patterns
+- SQLFixAgent (arXiv:2406.13408) — self-correcting SQL generation
+
+Built with smolagents CodeAgent + Gradio UI.
+"""
+
+import os
+import sqlite3
+from textwrap import dedent
+
+from smolagents import (
+    tool,
+    CodeAgent,
+    InferenceClientModel,
+    GradioUI,
+)
+
+# ─────────────────────────────────────────────
+# 1. Database Setup — Sample multi-table DB
+# ─────────────────────────────────────────────
+
+DB_PATH = "demo_company.db"
+
+
+def create_demo_database(db_path: str = DB_PATH):
+    """Creates a rich demo company database with realistic data and some ambiguous schema elements."""
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+
+    for table in ["order_items", "orders", "products", "customers", "employees", "departments"]:
+        cursor.execute(f"DROP TABLE IF EXISTS {table}")
+
+    cursor.execute("""
+        CREATE TABLE departments (
+            dept_id INTEGER PRIMARY KEY,
+            name TEXT NOT NULL,
+            location TEXT,
+            budget REAL
+        )
+    """)
+    cursor.executemany("INSERT INTO departments VALUES (?, ?, ?, ?)", [
+        (1, "Engineering", "San Francisco", 2500000.00),
+        (2, "Sales", "New York", 1800000.00),
+        (3, "Marketing", "New York", 1200000.00),
+        (4, "HR", "Chicago", 800000.00),
+        (5, "Finance", "Chicago", 950000.00),
+    ])
+
+    cursor.execute("""
+        CREATE TABLE employees (
+            emp_id INTEGER PRIMARY KEY,
+            name TEXT NOT NULL,
+            email TEXT,
+            dept_id INTEGER REFERENCES departments(dept_id),
+            salary REAL,
+            hire_date TEXT,
+            manager_id INTEGER REFERENCES employees(emp_id),
+            status TEXT DEFAULT 'active'
+        )
+    """)
+    cursor.executemany("INSERT INTO employees VALUES (?, ?, ?, ?, ?, ?, ?, ?)", [
+        (1, "Alice Chen", "alice@company.com", 1, 145000, "2019-03-15", None, "active"),
+        (2, "Bob Martinez", "bob@company.com", 1, 128000, "2020-06-01", 1, "active"),
+        (3, "Carol Smith", "carol@company.com", 2, 95000, "2021-01-10", None, "active"),
+        (4, "David Lee", "david@company.com", 2, 88000, "2021-08-20", 3, "active"),
+        (5, "Eva Johnson", "eva@company.com", 3, 102000, "2020-11-05", None, "active"),
+        (6, "Frank Wilson", "frank@company.com", 1, 135000, "2019-07-22", 1, "active"),
+        (7, "Grace Kim", "grace@company.com", 4, 78000, "2022-02-14", None, "active"),
+        (8, "Henry Brown", "henry@company.com", 5, 115000, "2020-04-30", None, "active"),
+        (9, "Iris Davis", "iris@company.com", 2, 92000, "2022-09-01", 3, "active"),
+        (10, "Jack Taylor", "jack@company.com", 1, 140000, "2019-11-18", 1, "inactive"),
+        (11, "Karen White", "karen@company.com", 3, 98000, "2021-05-12", 5, "active"),
+        (12, "Leo Garcia", "leo@company.com", 5, 105000, "2021-03-28", 8, "active"),
+    ])
+
+    cursor.execute("""
+        CREATE TABLE customers (
+            customer_id INTEGER PRIMARY KEY,
+            name TEXT NOT NULL,
+            email TEXT,
+            city TEXT,
+            state TEXT,
+            signup_date TEXT,
+            tier TEXT DEFAULT 'standard'
+        )
+    """)
+    cursor.executemany("INSERT INTO customers VALUES (?, ?, ?, ?, ?, ?, ?)", [
+        (1, "Acme Corp", "contact@acme.com", "San Francisco", "CA", "2020-01-15", "premium"),
+        (2, "Beta Industries", "info@beta.com", "New York", "NY", "2020-03-22", "standard"),
+        (3, "Gamma Solutions", "hello@gamma.com", "Chicago", "IL", "2020-06-10", "premium"),
+        (4, "Delta Systems", "sales@delta.com", "Austin", "TX", "2021-02-05", "enterprise"),
+        (5, "Epsilon LLC", "team@epsilon.com", "Seattle", "WA", "2021-08-18", "standard"),
+        (6, "Zeta Partners", "info@zeta.com", "Boston", "MA", "2022-01-30", "premium"),
+        (7, "Eta Global", "contact@eta.com", "Denver", "CO", "2022-07-14", "standard"),
+        (8, "Theta Inc", "hello@theta.com", "Portland", "OR", "2023-03-01", "enterprise"),
+    ])
+
+    cursor.execute("""
+        CREATE TABLE products (
+            product_id INTEGER PRIMARY KEY,
+            name TEXT NOT NULL,
+            category TEXT,
+            price REAL,
+            cost REAL,
+            stock_quantity INTEGER,
+            status TEXT DEFAULT 'active'
+        )
+    """)
+    cursor.executemany("INSERT INTO products VALUES (?, ?, ?, ?, ?, ?, ?)", [
+        (1, "Widget Pro", "Hardware", 299.99, 150.00, 500, "active"),
+        (2, "Widget Basic", "Hardware", 149.99, 75.00, 1200, "active"),
+        (3, "DataSync Cloud", "Software", 49.99, 10.00, None, "active"),
+        (4, "DataSync Enterprise", "Software", 199.99, 40.00, None, "active"),
+        (5, "SecureVault", "Software", 89.99, 20.00, None, "active"),
+        (6, "PowerAdapter X", "Hardware", 39.99, 18.00, 3000, "active"),
+        (7, "Legacy Suite", "Software", 299.99, 60.00, None, "discontinued"),
+        (8, "SmartHub", "Hardware", 449.99, 220.00, 200, "active"),
+    ])
+
+    cursor.execute("""
+        CREATE TABLE orders (
+            order_id INTEGER PRIMARY KEY,
+            customer_id INTEGER REFERENCES customers(customer_id),
+            employee_id INTEGER REFERENCES employees(emp_id),
+            order_date TEXT,
+            status TEXT,
+            total_amount REAL
+        )
+    """)
+    cursor.executemany("INSERT INTO orders VALUES (?, ?, ?, ?, ?, ?)", [
+        (1001, 1, 3, "2024-01-15", "completed", 1499.95),
+        (1002, 2, 4, "2024-01-22", "completed", 599.96),
+        (1003, 3, 3, "2024-02-10", "completed", 899.97),
+        (1004, 1, 9, "2024-02-28", "shipped", 449.99),
+        (1005, 4, 4, "2024-03-05", "completed", 2499.90),
+        (1006, 5, 3, "2024-03-18", "pending", 149.99),
+        (1007, 6, 9, "2024-04-02", "completed", 749.97),
+        (1008, 3, 3, "2024-04-15", "completed", 339.98),
+        (1009, 7, 4, "2024-05-01", "cancelled", 299.99),
+        (1010, 8, 9, "2024-05-20", "shipped", 1349.97),
+        (1011, 1, 3, "2024-06-01", "completed", 199.98),
+        (1012, 4, 4, "2024-06-15", "completed", 3599.88),
+    ])
+
+    cursor.execute("""
+        CREATE TABLE order_items (
+            item_id INTEGER PRIMARY KEY,
+            order_id INTEGER REFERENCES orders(order_id),
+            product_id INTEGER REFERENCES products(product_id),
+            quantity INTEGER,
+            unit_price REAL,
+            discount REAL DEFAULT 0.0
+        )
+    """)
+    cursor.executemany("INSERT INTO order_items VALUES (?, ?, ?, ?, ?, ?)", [
+        (1, 1001, 1, 5, 299.99, 0.0),
+        (2, 1002, 2, 4, 149.99, 0.0),
+        (3, 1003, 3, 6, 49.99, 0.0),
+        (4, 1003, 5, 3, 89.99, 10.0),
+        (5, 1004, 8, 1, 449.99, 0.0),
+        (6, 1005, 1, 5, 299.99, 0.0),
+        (7, 1005, 4, 5, 199.99, 0.0),
+        (8, 1006, 2, 1, 149.99, 0.0),
+        (9, 1007, 5, 3, 89.99, 0.0),
+        (10, 1007, 3, 9, 49.99, 10.0),
+        (11, 1008, 6, 5, 39.99, 0.0),
+        (12, 1008, 3, 3, 49.99, 10.0),
+        (13, 1009, 1, 1, 299.99, 0.0),
+        (14, 1010, 8, 3, 449.99, 0.0),
+        (15, 1011, 3, 4, 49.99, 0.0),
+        (16, 1012, 4, 12, 199.99, 15.0),
+        (17, 1012, 8, 2, 449.99, 10.0),
+    ])
+
+    conn.commit()
+    conn.close()
+    return db_path
+
+
+# ─────────────────────────────────────────────
+# 2. Build Dynamic Schema Description
+# ─────────────────────────────────────────────
+
+def get_schema_description(db_path: str = DB_PATH) -> str:
+    conn = sqlite3.connect(db_path)
+    cursor = conn.cursor()
+    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name")
+    tables = [row[0] for row in cursor.fetchall()]
+
+    schema_parts = []
+    for table in tables:
+        cursor.execute(f"PRAGMA table_info({table})")
+        columns = cursor.fetchall()
+        cursor.execute(f"PRAGMA foreign_key_list({table})")
+        fks = cursor.fetchall()
+        fk_map = {fk[3]: f"→ {fk[2]}({fk[4]})" for fk in fks}
+        cursor.execute(f"SELECT COUNT(*) FROM {table}")
+        row_count = cursor.fetchone()[0]
+
+        table_desc = f"Table '{table}' ({row_count} rows):\n  Columns:\n"
+        for col in columns:
+            col_id, col_name, col_type, not_null, default, pk = col
+            parts = [f"    - {col_name}: {col_type or 'TEXT'}"]
+            if pk: parts.append("PRIMARY KEY")
+            if not_null and not pk: parts.append("NOT NULL")
+            if default is not None: parts.append(f"DEFAULT {default}")
+            if col_name in fk_map: parts.append(f"FK {fk_map[col_name]}")
+            table_desc += " ".join(parts) + "\n"
+
+        for col in columns:
+            col_name, col_type = col[1], col[2]
+            if col_type in ("TEXT", None) and col_name not in ("email",):
+                try:
+                    cursor.execute(f"SELECT DISTINCT {col_name} FROM {table} WHERE {col_name} IS NOT NULL LIMIT 8")
+                    vals = [str(r[0]) for r in cursor.fetchall()]
+                    if vals:
+                        table_desc += f"  Sample '{col_name}' values: {', '.join(vals)}\n"
+                except:
+                    pass
+        schema_parts.append(table_desc)
+
+    conn.close()
+    return "\n".join(schema_parts)
+
+
+# ─────────────────────────────────────────────
+# 3. Define Agent Tools
+# ─────────────────────────────────────────────
+
+SCHEMA_DESCRIPTION = ""
+
+
+@tool
+def execute_sql(query: str) -> str:
+    """
+    Executes a SQL query against the company database and returns the results.
+    Use this tool to run SELECT queries to answer user questions about the data.
+    
+    IMPORTANT RULES:
+    - Only use SELECT statements (no INSERT, UPDATE, DELETE, DROP)
+    - Always use table and column names exactly as shown in the schema
+    - Use JOINs when data spans multiple tables
+    - Use LIMIT to avoid overwhelming output (max 50 rows)
+    
+    DATABASE SCHEMA:
+    {schema}
+
+    Args:
+        query: A valid SQL SELECT query to execute against the database.
+    """
+    cleaned = query.strip().upper()
+    if not cleaned.startswith("SELECT") and not cleaned.startswith("WITH"):
+        return "ERROR: Only SELECT queries are allowed."
+
+    try:
+        conn = sqlite3.connect(DB_PATH)
+        cursor = conn.cursor()
+        cursor.execute(query)
+        columns = [desc[0] for desc in cursor.description] if cursor.description else []
+        rows = cursor.fetchall()
+        conn.close()
+
+        if not rows:
+            return f"Query executed successfully.\nColumns: {', '.join(columns)}\nResult: No rows returned."
+
+        result = f"Query executed successfully. {len(rows)} row(s) returned.\n"
+        result += "Columns: " + " | ".join(columns) + "\n"
+        result += "-" * 60 + "\n"
+        for row in rows[:50]:
+            result += " | ".join(str(v) for v in row) + "\n"
+        if len(rows) > 50:
+            result += f"... ({len(rows) - 50} more rows truncated)\n"
+        return result
+
+    except Exception as e:
+        return f"SQL ERROR: {str(e)}\n\nPlease check your query syntax and column/table names against the schema."
+
+
+@tool
+def inspect_schema(table_name: str = "") -> str:
+    """
+    Inspect the database schema. If a table_name is provided, shows detailed info
+    about that specific table including column types, foreign keys, and sample data.
+    If no table_name is given, shows an overview of all tables.
+    
+    Use this tool BEFORE writing SQL to understand the database structure,
+    especially when the user's question is ambiguous about which tables or columns to use.
+    
+    Args:
+        table_name: Name of a specific table to inspect. Leave empty for full schema overview.
+    """
+    conn = sqlite3.connect(DB_PATH)
+    cursor = conn.cursor()
+
+    if not table_name:
+        return f"DATABASE SCHEMA OVERVIEW:\n\n{SCHEMA_DESCRIPTION}"
+
+    try:
+        cursor.execute(f"PRAGMA table_info({table_name})")
+        columns = cursor.fetchall()
+        if not columns:
+            conn.close()
+            return f"Table '{table_name}' not found. Use inspect_schema() with no arguments to see all tables."
+
+        result = f"DETAILED INSPECTION OF TABLE '{table_name}':\n\n"
+        result += "Columns:\n"
+        for col in columns:
+            result += f"  {col[1]} ({col[2] or 'TEXT'})"
+            if col[5]: result += " [PRIMARY KEY]"
+            if col[3]: result += " [NOT NULL]"
+            result += "\n"
+
+        cursor.execute(f"PRAGMA foreign_key_list({table_name})")
+        fks = cursor.fetchall()
+        if fks:
+            result += "\nForeign Keys:\n"
+            for fk in fks:
+                result += f"  {fk[3]} → {fk[2]}({fk[4]})\n"
+
+        cursor.execute(f"SELECT COUNT(*) FROM {table_name}")
+        count = cursor.fetchone()[0]
+        result += f"\nTotal rows: {count}\n"
+
+        cursor.execute(f"SELECT * FROM {table_name} LIMIT 3")
+        sample_rows = cursor.fetchall()
+        col_names = [c[1] for c in columns]
+        result += f"\nSample rows (first 3):\n"
+        result += " | ".join(col_names) + "\n"
+        result += "-" * 60 + "\n"
+        for row in sample_rows:
+            result += " | ".join(str(v) for v in row) + "\n"
+
+        conn.close()
+        return result
+
+    except Exception as e:
+        conn.close()
+        return f"Error inspecting table: {str(e)}"
+
+
+# ─────────────────────────────────────────────
+# 4. Agent System Prompt
+# ─────────────────────────────────────────────
+
+SYSTEM_INSTRUCTIONS = dedent("""\
+You are an expert SQL assistant that helps users query a company database. You follow a structured multi-turn approach:
+
+## YOUR DECISION PROCESS
+
+For EVERY user question, follow these steps:
+
+### Step 1: Classify the Question
+Determine if the question is:
+- **ANSWERABLE**: The question is clear and maps directly to the database schema
+- **AMBIGUOUS**: The question could have multiple valid SQL interpretations (e.g., "show me the top employees" — top by salary? by sales? by tenure?)
+- **UNANSWERABLE**: The question asks for data that doesn't exist in the database
+
+### Step 2: Handle Based on Classification
+
+**If AMBIGUOUS:**
+- Identify ALL possible interpretations
+- Use `final_answer()` to return a targeted clarification question listing the specific options
+- Example: call `final_answer("Your question could mean several things:\\n1. Employees with highest salary\\n2. Employees who handled the most orders\\n3. Employees with the longest tenure\\n\\nWhich interpretation do you mean?")`
+- Do NOT generate SQL — return the clarification question immediately using `final_answer()`
+- The user will respond in the next turn with their clarification
+
+**If UNANSWERABLE:**
+- Use `final_answer()` to explain clearly what data is missing and why the question can't be answered
+- Include a suggestion for a related question that CAN be answered with the available data
+
+**If ANSWERABLE:**
+- First inspect the schema to confirm the right tables/columns
+- Generate and execute the SQL query
+- Present results clearly with a natural language summary
+
+### Step 3: Self-Correct
+- If your SQL returns an error, analyze the error and fix the query
+- If the result seems wrong or empty, verify your joins and filters
+- Always sanity-check: does the result make sense given what was asked?
+
+## COMMON AMBIGUITY PATTERNS TO WATCH FOR
+
+1. **Column ambiguity**: "Show employee names" — the 'name' column appears in employees, departments, customers, and products tables
+2. **Metric ambiguity**: "Top customers" — by total spending? by number of orders? by most recent activity?
+3. **Filter ambiguity**: "Recent orders" — last week? last month? last quarter?
+4. **Scope ambiguity**: "Total sales" — all time? this year? by product? by employee?
+5. **Status ambiguity**: "List products" — all products? only active ones? including discontinued?
+6. **Value ambiguity**: "Expensive products" — what price threshold?
+
+## FORMATTING RULES
+
+- When presenting query results, format them as a clear table
+- Always explain what the query does in plain language
+- If you make assumptions (e.g., "I'm assuming you mean active employees only"), state them explicitly
+- For numerical results, include relevant aggregations (count, sum, average) when helpful
+""")
+
+
+# ─────────────────────────────────────────────
+# 5. Main
+# ─────────────────────────────────────────────
+
+def create_agent(model_id: str = "Qwen/Qwen2.5-Coder-32B-Instruct"):
+    create_demo_database()
+
+    global SCHEMA_DESCRIPTION
+    SCHEMA_DESCRIPTION = get_schema_description()
+    execute_sql.description = execute_sql.description.replace("{schema}", SCHEMA_DESCRIPTION)
+
+    model = InferenceClientModel(model_id=model_id)
+
+    agent = CodeAgent(
+        tools=[execute_sql, inspect_schema],
+        model=model,
+        instructions=SYSTEM_INSTRUCTIONS,
+        max_steps=15,
+        additional_authorized_imports=["json", "re"],
+    )
+    return agent
+
+
+if __name__ == "__main__":
+    agent = create_agent()
+    ui = GradioUI(agent, reset_agent_memory=False)
+    ui.launch()