melikakheirieh/nl2sql-copilot-prototype

1from config import (2    LLM_MODEL,3    LLM_TEMPERATURE,4    FORBIDDEN_KEYWORDS,5    FORBIDDEN_TABLES6)7import os8import sqlite39import json10import re11from typing import Optional, Tuple, List12 13import gradio as gr14import sqlglot15from sqlglot import exp16 17from langchain_openai import ChatOpenAI18from langchain_community.utilities import SQLDatabase19from langchain.chains import create_sql_query_chain20from langchain.prompts import ChatPromptTemplate21 22 23def get_readonly_sqlite_url(db_path: str) -> str:24    return f"file:{db_path}?mode=ro&uri=true"25 26def get_schema_preview(db_path: str, limit_per_table: int = 0) -> str:27    uri = get_readonly_sqlite_url(db_path)28    with sqlite3.connect(uri, uri=True, timeout=3) as conn:29        conn.row_factory = sqlite3.Row30        cur = conn.cursor()31        cur.execute("SELECT name FROM sqlite_master WHERE type='table' ORDER BY name;")32        tables = [r["name"] for r in cur.fetchall()]33        lines = []34        for t in tables:35            # skip SQLite internals36            if t in FORBIDDEN_TABLES:37                continue38            cur.execute(f"PRAGMA table_info({t});")39            cols = cur.fetchall()40            col_line = ", ".join([f"{c['name']}:{c['type']}" for c in cols])41            lines.append(f"- {t} ({col_line})")42            if limit_per_table > 0:43                try:44                    cur.execute(f"SELECT * FROM {t} LIMIT {limit_per_table};")45                    sample = cur.fetchall()46                    if sample:47                        lines.append(f"  sample rows: {len(sample)}")48                except Exception:49                    pass50        if not lines:51            return "(no user tables found)"52        return "\n".join(lines)53 54 55def validate_sql_safe(sql: str) -> Tuple[bool, str]:56    if sql.count(";") > 0:57        if sql.strip().endswith(";"):58            if sql.strip()[:-1].count(";") > 0:59                return False, "Multiple statements are not allowed."60        else:61            return False, "Multiple statements are not allowed."62 63    upper = re.sub(r"\s+", " ", sql).strip()64    for kw in FORBIDDEN_KEYWORDS:65        if re.search(rf"\b{kw}\b", upper):66            return False, f"Keyword '{kw}' is not allowed."67 68    try:69        parsed = sqlglot.parse(sql, read='sqlite')70    except Exception as e:71        return False, f"SQL parse error: {e}"72 73    if not parsed or len(parsed) != 1:74        return False, "Exactly one SQL statement is allowed."75 76    stmt = parsed[0]77    if not isinstance(stmt, exp.Select):78        return False, "Only SELECT statements are allowed."79 80    for table in stmt.find_all(exp.Table):81        table_name = table.name.lower() if table.name else ""82        if table_name in FORBIDDEN_TABLES:83            return False, f"Access to {table_name} is not allowed."84 85    return True, "OK"86 87def execute_select(db_path: str, sql: str, max_rows: int = 1000, timeout: float = 5.0) -> Tuple[list[str], List[List]]:88    uri = get_readonly_sqlite_url(db_path)89    if not re.search(r"\bLIMIT\b", sql, re.IGNORECASE):90        sql = f"{sql.rstrip(';')} LIMIT {max_rows}"91 92    with sqlite3.connect(uri, uri=True, timeout=timeout) as conn:93        conn.row_factory = sqlite3.Row94        cur = conn.cursor()95        cur.execute(sql)96        rows = cur.fetchall()97        if rows:98            cols = rows[0].keys()99            data = [list(r) for r in rows]100            return list(cols), data101        else:102            return [], []103 104 105 106custom_prompt = ChatPromptTemplate.from_template("""107Given the following question, return ONLY a valid SQL query in JSON form.108 109Question: {input}110Database schema: {table_info}111 112You may sample/preview at most {top_k} rows if you need examples.113 114Respond in this exact JSON format:115{{116  "sql": "<SQL_QUERY_HERE>"117}}118""")119 120 121def make_sql_chain(sql_db: SQLDatabase):122    llm = ChatOpenAI(model=LLM_MODEL, temperature=LLM_TEMPERATURE)123    chain = create_sql_query_chain(llm, sql_db, prompt=custom_prompt, k=20)124    return chain125 126 127def on_upload_database(db_file, state):128    if db_file is None:129        return state, "No file provided.", "(no schema)"130    path = db_file.name131 132    sql_db = SQLDatabase.from_uri(f"sqlite:///{path}")133 134    schema_text = get_schema_preview(path, limit_per_table=0)135 136    chain = make_sql_chain(sql_db)137 138    new_state = {139        "db_path": path,140        "sql_db": sql_db,141        "schema_text": schema_text,142        "chain": chain,143    }144    return new_state, f"Database '{os.path.basename(path)}' uploaded successfully.", schema_text145 146def extract_sql_safe(output_text: str) -> str:147    try:148        obj = json.loads(output_text)149        if isinstance(obj, dict) and "sql" in obj:150            return obj["sql"].strip()151    except Exception:152        pass153    m = re.search(r"```sql\s*(.*?)\s*```", output_text, re.DOTALL | re.IGNORECASE)154    if m:155        return m.group(1).strip()156    return output_text.strip()157 158def on_generate_query(question , max_rows, state):159    if not state or not state.get("db_path") or not state.get("chain"):160        return "Please upload a database first.", "", ""161    if not question or not question.strip():162        return "Please enter a question.", "", ""163 164    try:165        generated_sql = state["chain"].invoke({"question": question})166 167        sql = extract_sql_safe(str(generated_sql))168 169        ok, msg = validate_sql_safe(sql)170        if not ok:171            return f"Blocked SQL: {msg}", sql, ""172 173        cols, rows = execute_select(state["db_path"], sql, max_rows=max_rows)174        if not cols:175            return f"No rows returned.", sql, "[]"176 177        sample = [dict(zip(cols, r)) for r in rows[:50]]178        return f"Returned {len(rows)} row(s). Showing up to 50.", sql, json.dumps(sample, indent=2)179 180    except Exception as e:181        return f"Error: {e}", "", ""182 183 184with gr.Blocks(title="nl2sql-copilot-prototype (safe)") as demo:185    gr.Markdown("# nl2sql-copilot-prototype (Sqlite, safe)")186    gr.Markdown(187        "Upload a **SQLite** file, ask a question in natural language, "188        "and I will: (1) generate SQL, (2) validate it (SELECT-only), (3) execute read-only, "189        "and (4) show you the results."190    )191 192    state = gr.State({"db_path": None, "sql_db": None, "schema_text": "", "chain": None})193 194    with gr.Row():195        db_file = gr.File(label="Upload SQlite Database", file_types=[".sqlite", ".db"])196        upload_status = gr.Textbox(label="upload Status", interactive=False)197 198    schema_box = gr.Accordion("Database schema (preview)", open=False)199    with schema_box:200        schema_md = gr.Markdown("(no schema)")201 202    gr.Markdown("---")203 204    with gr.Row():205        question = gr.Textbox(label="Your question", placeholder="e.g., Top 10 tracks by total sales")206    with gr.Row():207        max_row= gr.Slider(10, 5000, value=1000, step=10, label="Max rows")208 209    with gr.Row():210        run_btn = gr.Button("Generate & Run SQL", variant="primary")211 212    with gr.Row():213        status_out = gr.Textbox(label="Status")214    with gr.Row():215        sql_out = gr.Code(label="Generated SQL (validated)")216    with gr.Row():217        result_out = gr.Code(label="Result (JSON sample)")218 219    db_file.change(220        fn=on_upload_database,221        inputs=[db_file, state],222        outputs=[state, upload_status, schema_md],223    )224 225    run_btn.click(226        fn=on_generate_query,227        inputs=[question, max_row, state],228        outputs=[status_out, sql_out, result_out],229    )230 231 232 233if __name__ == "__main__":234    demo.launch()