app.py

import streamlit as st
import sqlite3
import pandas as pd
import sqlparse
import google.generativeai as genai

st.set_page_config(page_title="NL2SQL", layout="wide")

api_key = st.secrets.get("GEMINI_API_KEY", "")
if not api_key:
    st.error("Please set GEMINI_API_KEY in .streamlit/secrets.toml")
    st.stop()

genai.configure(api_key=api_key)

def get_model():
    """Return the best available Gemini model with fallback."""
    preferred_models = [
        "models/gemini-2.5-flash",
        "models/gemini-flash-latest",
        "models/gemini-2.0-flash",
        "models/gemini-pro-latest",
    ]
    for name in preferred_models:
        try:
            model = genai.GenerativeModel(name)
            #st.caption(f"Using model: **{name}**")
            return model
        except Exception:
            continue
    st.error("⚠️ No suitable Gemini model found. Please check your API quota.")
    st.stop()

# Initialize model
model = get_model()

def get_schema(conn):
    """Return schema as dict: {table: [columns]}"""
    schema = {}
    cursor = conn.cursor()
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table';")
    tables = [row[0] for row in cursor.fetchall()]
    for table in tables:
        cursor.execute(f"PRAGMA table_info({table});")
        columns = [row[1] for row in cursor.fetchall()]
        schema[table] = columns
    return schema

def format_schema(schema):
    """Nicely format schema for display & prompt"""
    return "\n".join([f"{table}: {', '.join(cols)}" for table, cols in schema.items()])

def get_raw_schema(conn):
    """Return raw CREATE TABLE DDLs"""
    cursor = conn.cursor()
    cursor.execute("SELECT name, sql FROM sqlite_master WHERE type='table';")
    return cursor.fetchall()

def format_raw_schema(raw_schema):
    """Format DDLs into text block"""
    return "\n".join([f"{tbl}: {ddl}" for tbl, ddl in raw_schema if ddl])

def clean_sql(sql_text: str) -> str:
    """Clean LLM output to get valid SQL for SQLite execution."""
    sql_text = sql_text.strip()
    if sql_text.startswith("```"):
        parts = sql_text.split("```")
        if len(parts) >= 2:
            sql_text = parts[1]
    sql_text = sql_text.replace("sql\n", "").replace("sqlite", "").replace("ite", "")
    sql_text = sql_text.replace("sql", "")
    sql_text = sql_text.strip()
    if sql_text.endswith(";"):
        sql_text = sql_text[:-1]
    return sql_text

def execute_sql(conn, query):
    """Run SQL query (SELECT or non-SELECT) and return results or status"""
    cursor = conn.cursor()
    try:
        cursor.execute(query)
        if query.strip().lower().startswith("select"):
            rows = cursor.fetchall()
            cols = [desc[0] for desc in cursor.description]
            return pd.DataFrame(rows, columns=cols), None
        else:
            conn.commit()
            return None, f"Query executed successfully: {query}"
    except Exception as e:
        return None, f"⚠️ Error executing SQL: {e}"

st.title("💬 Natural Language → SQL Query Generator")

uploaded_file = st.file_uploader("Upload your SQLite database (.db)", type=["db"])

if uploaded_file:
    db_path = "uploaded.db"
    with open(db_path, "wb") as f:
        f.write(uploaded_file.read())
    st.success("Database uploaded successfully!")

    conn = sqlite3.connect(db_path)
    schema = get_schema(conn)

    st.subheader("Database Schema")
    for table, cols in schema.items():
        st.markdown(f"**{table}**: {', '.join(cols)}")

    show_ddl = st.checkbox("Show raw schema (DDL statements)")
    raw_schema = None
    if show_ddl:
        raw_schema = get_raw_schema(conn)
        st.subheader("Raw Schema DDL")
        for table, ddl in raw_schema:
            if ddl:
                st.code(ddl, language="sql")

    nl_query = st.text_area(
        "Write your query in English",
        placeholder="e.g. Show me top 5 students by GPA in 2023",
    )

    if st.button("Generate SQL"):
        if nl_query.strip() == "":
            st.warning("Please enter a query.")
        else:
            ddl_context = format_raw_schema(raw_schema) if raw_schema else ""
            prompt = f"""
            You are a SQL assistant. Convert the following natural language query 
            into a valid SQLite SQL query. Only use the provided tables and columns.

            Database schema:
            {format_schema(schema)}

            DDL:
            {ddl_context}

            Query: {nl_query}
            """
            try:
                response = model.generate_content(prompt)
                raw_sql = response.text.strip()
                sql_query = clean_sql(raw_sql)
                sql_query = sqlparse.format(sql_query, reindent=True)

                st.subheader("🧠 Generated SQL")
                st.code(sql_query, language="sql")

                df, msg = execute_sql(conn, sql_query)
                if df is not None:
                    st.dataframe(df, use_container_width=True)
                    csv = df.to_csv(index=False).encode("utf-8")
                    st.download_button("Download Results", csv, "results.csv", "text/csv")
                elif msg:
                    st.success(msg)
            except Exception as e:
                st.error(f"⚠️ Error generating SQL: {e}")

    conn.close()