""" 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")