app.py

"""
app_hf.py -- Hugging Face Spaces version of the Streamlit UI
Calls Groq directly (no Flask layer needed on HF Spaces).

Deploy:
    1. Create Space at huggingface.co -> New Space -> Streamlit -> Public
    2. Add GROQ_API_KEY to Space Secrets (Settings -> Repository Secrets)
    3. Copy this file as app.py to your Space repo
    4. Copy requirements_hf.txt as requirements.txt
"""

import os
import sys
import json
import time
import re
import numpy as np
import pandas as pd
import streamlit as st
from pathlib import Path
from datetime import datetime
from groq import Groq
from dotenv import load_dotenv

load_dotenv()

# -------------------------------------------------------
# Config
# -------------------------------------------------------

MODEL       = "llama-3.3-70b-versatile"
TEMPERATURE = 0
MAX_TOKENS  = 300
TIMEOUT_SEC = 30

STRATEGY_LABELS = {
    "zero_shot"       : "Zero-Shot",
    "few_shot"        : "Few-Shot",
    "chain_of_thought": "Chain-of-Thought",
}

SQL_KEYWORDS = [
    "SELECT","FROM","WHERE","AND","OR","NOT","IN","EXISTS",
    "JOIN","LEFT","RIGHT","INNER","OUTER","ON","AS","DISTINCT",
    "ORDER","BY","GROUP","HAVING","LIMIT","COUNT","SUM","AVG",
    "MAX","MIN","ASC","DESC","UNION","INTERSECT","EXCEPT",
]

SYSTEM_PROMPTS = {
    "zero_shot": (
        "You are an expert SQL assistant. Given a database schema and a natural "
        "language question, return ONLY the SQL query. Do not explain. Do not use "
        "markdown. Do not add any text before or after the SQL. "
        "Use DISTINCT when the question implies unique or different values. "
        "Use foreign key relationships between tables when writing JOINs."
    ),
    "few_shot": (
        "You are an expert SQL assistant. Study the examples below, then generate "
        "SQL for the new question. Return ONLY the SQL query. Do not explain. "
        "Do not use markdown. Do not add any text before or after the SQL. "
        "Use DISTINCT when the question implies unique or different values. "
        "Use foreign key relationships between tables when writing JOINs."
    ),
    "chain_of_thought": (
        "You are a SQL expert. Think through the query step by step, then write "
        "the final SQL. Rules: Use DISTINCT when the question implies uniqueness. "
        "Use foreign key relationships between tables when writing JOINs. "
        "Always include all relevant columns. "
        "Your final answer MUST start with SELECT and contain only SQL -- "
        "no explanation after the query."
    ),
}

# -------------------------------------------------------
# Page config
# -------------------------------------------------------

st.set_page_config(
    page_title = "AI SQL Assistant",
    page_icon  = "🛢",
    layout     = "wide",
    initial_sidebar_state = "expanded",
)

# -------------------------------------------------------
# Session state
# -------------------------------------------------------

if "history" not in st.session_state:
    st.session_state.history = []

# -------------------------------------------------------
# Groq client (cached)
# -------------------------------------------------------

@st.cache_resource
def get_groq_client():
    api_key = os.getenv("GROQ_API_KEY")
    if not api_key:
        return None
    return Groq(api_key=api_key)


# -------------------------------------------------------
# SQL helpers
# -------------------------------------------------------

def clean_sql(raw: str) -> str:
    if not raw or not raw.strip():
        return ""
    text = raw.strip()
    # Strip markdown fences
    match = re.search(r'```(?:sql|SQL)?\s*\n?(.*?)\n?```', text, re.DOTALL)
    if match:
        text = match.group(1).strip()
    # Strip prefix explanation
    upper = text.upper()
    idx = upper.find('SELECT')
    if idx > 0:
        prefix = text[:idx].strip()
        if prefix:
            text = text[idx:].strip()
    # Uppercase keywords
    for kw in SQL_KEYWORDS:
        text = re.sub(rf'\b{kw}\b', kw, text, flags=re.IGNORECASE)
    # Collapse whitespace
    text = re.sub(r'\s+', ' ', text).strip()
    # Remove trailing semicolon
    text = text.rstrip(';').rstrip()
    return text


def is_valid_sql(sql: str) -> bool:
    try:
        import sqlglot
        sqlglot.parse_one(sql)
        return True
    except Exception:
        return 'SELECT' in sql.upper() and 'FROM' in sql.upper()


def build_prompt(nl_query: str, schema: dict, strategy: str) -> str:
    tables  = ", ".join(schema.get("tables", []))
    columns = ", ".join(schema.get("columns", [])[:15])
    schema_text = ""
    if tables:
        schema_text += f"Tables: {tables}\n"
    if columns:
        schema_text += f"Columns: {columns}"

    if strategy == "chain_of_thought":
        return (
            f"Database schema:\n{schema_text}\n\n"
            f"Question: {nl_query}\n\n"
            f"Think step by step:\n"
            f"1. Which tables do I need?\n"
            f"2. Which columns are relevant?\n"
            f"3. What SQL pattern fits?\n"
            f"4. Does the question imply DISTINCT?\n\n"
            f"Now write ONLY the SQL query. Start with SELECT:"
        )
    else:
        return (
            f"Database schema:\n{schema_text}\n\n"
            f"Question: {nl_query}\n\n"
            f"SQL:"
        )


def generate_sql_direct(nl_query: str, schema: dict, strategy: str) -> dict:
    client = get_groq_client()
    if not client:
        return {"error": "GROQ_API_KEY not set in Space Secrets"}

    system = SYSTEM_PROMPTS.get(strategy, SYSTEM_PROMPTS["zero_shot"])
    prompt = build_prompt(nl_query, schema, strategy)

    t0 = time.time()
    try:
        resp = client.chat.completions.create(
            model       = MODEL,
            temperature = TEMPERATURE,
            max_tokens  = MAX_TOKENS,
            messages    = [
                {"role": "system", "content": system},
                {"role": "user",   "content": prompt},
            ],
        )
        raw = resp.choices[0].message.content.strip()
    except Exception as e:
        return {"error": str(e)}

    latency = round(time.time() - t0, 3)
    sql     = clean_sql(raw)
    valid   = is_valid_sql(sql) if sql else False

    return {
        "sql"      : sql,
        "raw"      : raw,
        "valid"    : valid,
        "strategy" : strategy,
        "model"    : MODEL,
        "latency_s": latency,
        "rag_used" : False,
        "retrieved_pairs": [],
    }


# -------------------------------------------------------
# Sidebar
# -------------------------------------------------------

with st.sidebar:
    st.title("🛢 AI SQL Assistant")
    st.caption("RAG-powered NL to SQL generation")
    st.divider()

    client = get_groq_client()
    if client:
        st.success("Groq API connected")
        st.caption(f"Model: {MODEL}")
    else:
        st.error("GROQ_API_KEY not set")
        st.caption("Add it to Space Secrets")

    st.divider()

    strategy_key = st.selectbox(
        "Prompt Strategy",
        options     = list(STRATEGY_LABELS.keys()),
        format_func = lambda k: STRATEGY_LABELS[k],
    )

    st.divider()
    st.caption(f"Session queries: {len(st.session_state.history)}")
    st.caption("Built with Groq + LLaMA-3.3-70B")


# -------------------------------------------------------
# Tabs
# -------------------------------------------------------

tab_generate, tab_history = st.tabs(["Generate", "History"])


# ===============================================================
# TAB 1 -- GENERATE
# ===============================================================

with tab_generate:
    st.header("Generate SQL from Natural Language")

    with st.expander("Optional: provide database schema hint", expanded=False):
        col_t, col_c = st.columns(2)
        with col_t:
            tables_input = st.text_input("Tables (comma-separated)",
                                          placeholder="orders, customers, products")
        with col_c:
            columns_input = st.text_input("Columns (comma-separated)",
                                           placeholder="id, name, amount, region")

    schema = {}
    if tables_input or columns_input:
        schema = {
            "tables" : [t.strip() for t in tables_input.split(",") if t.strip()],
            "columns": [c.strip() for c in columns_input.split(",") if c.strip()],
        }

    nl_query = st.text_area(
        "Natural language question",
        placeholder="e.g. Find the top 5 customers by total spend in 2024",
        height=100,
    )

    generate_clicked = st.button("Generate SQL", type="primary")

    if generate_clicked:
        if not nl_query.strip():
            st.warning("Please enter a question.")
        elif not client:
            st.error("GROQ_API_KEY not set in Space Secrets.")
        else:
            with st.spinner("Generating SQL..."):
                result = generate_sql_direct(
                    nl_query = nl_query.strip(),
                    schema   = schema,
                    strategy = strategy_key,
                )

            if result and "error" not in result:
                st.divider()
                st.subheader("Result")

                m1, m2, m3 = st.columns(3)
                m1.metric("Valid SQL", "Yes" if result.get("valid") else "No")
                m2.metric("Strategy", STRATEGY_LABELS.get(result.get("strategy",""), ""))
                m3.metric("Latency", f"{result.get('latency_s', 0):.2f}s")

                st.subheader("Generated SQL")
                st.code(result.get("sql", ""), language="sql")

                st.session_state.history.append({
                    "timestamp": datetime.now().strftime("%H:%M:%S"),
                    "nl_query" : nl_query.strip(),
                    "sql"      : result.get("sql", ""),
                    "strategy" : result.get("strategy", ""),
                    "valid"    : result.get("valid", False),
                    "latency_s": result.get("latency_s", 0),
                    "rag_used" : False,
                })
            elif result and "error" in result:
                st.error(f"Error: {result['error']}")


# ===============================================================
# TAB 2 -- HISTORY
# ===============================================================

with tab_history:
    st.header("Session History")

    if not st.session_state.history:
        st.info("No queries yet. Use the Generate tab to start.")
    else:
        col_exp, col_dl = st.columns([1, 1])
        with col_exp:
            if st.button("Clear history"):
                st.session_state.history = []
                st.rerun()
        with col_dl:
            history_json = json.dumps(st.session_state.history, indent=2)
            st.download_button(
                label     = "Export history (JSON)",
                data      = history_json,
                file_name = f"sql_history_{datetime.now().strftime('%Y%m%d_%H%M%S')}.json",
                mime      = "application/json",
            )

        st.divider()
        for i, entry in enumerate(reversed(st.session_state.history), 1):
            valid_tag = "Valid" if entry.get("valid") else "Invalid"
            with st.expander(
                f"[{entry['timestamp']}]  {entry['nl_query'][:70]}  |  {valid_tag}",
                expanded = i == 1,
            ):
                st.write(f"**Question:** {entry['nl_query']}")
                st.code(entry["sql"], language="sql")
                c1, c2, c3 = st.columns(3)
                c1.caption(f"Strategy: {STRATEGY_LABELS.get(entry['strategy'], entry['strategy'])}")
                c2.caption(f"Valid: {'Yes' if entry.get('valid') else 'No'}")
                c3.caption(f"Latency: {entry.get('latency_s', 0):.2f}s")