""" Text-to-SQL Trace RAG — Gradio Prototype (HuggingFace Spaces, no Docker) -------------------------------------------------------------------------- Same trace RAG logic as before, wrapped in a Gradio UI so it deploys on HF Spaces' free Gradio SDK (Docker Spaces are no longer free). Set GROQ_API_KEY as a Space secret (Settings -> Repository secrets), or export it locally before running. """ import os import re import duckdb import faiss import numpy as np import gradio as gr from sentence_transformers import SentenceTransformer from groq import Groq # --------------------------------------------------------------------------- # 1. Toy retail dataset (swap this for a real BinDawood extract later) # --------------------------------------------------------------------------- con = duckdb.connect(":memory:") con.execute(""" CREATE TABLE sales ( sale_id INTEGER, region VARCHAR, store VARCHAR, category VARCHAR, product VARCHAR, quantity INTEGER, revenue DOUBLE, sale_date DATE ) """) con.execute(""" INSERT INTO sales VALUES (1,'Riyadh','Store A','Dairy','Milk 1L',120,600.0,'2026-07-05'), (2,'Riyadh','Store A','Bakery','Bread',80,240.0,'2026-07-05'), (3,'Jeddah','Store B','Dairy','Milk 1L',95,475.0,'2026-07-06'), (4,'Jeddah','Store B','Produce','Tomatoes',200,300.0,'2026-07-06'), (5,'Dammam','Store C','Dairy','Milk 1L',60,300.0,'2026-07-07'), (6,'Riyadh','Store A','Produce','Tomatoes',150,225.0,'2026-08-01'), (7,'Jeddah','Store B','Bakery','Bread',110,330.0,'2026-08-01') """) SCHEMA_DESC = """ Table: sales Columns: sale_id (INTEGER), region (VARCHAR), store (VARCHAR), category (VARCHAR), product (VARCHAR), quantity (INTEGER), revenue (DOUBLE), sale_date (DATE) """ # --------------------------------------------------------------------------- # 2. Trace library -- VERIFIED past question -> SQL -> reasoning. # Grows over time; new verified entries get appended here at runtime. # --------------------------------------------------------------------------- TRACE_LIBRARY = [ { "question": "What is total revenue by region for July 2026?", "reasoning": ( "Step 1: Need SUM(revenue) grouped by region.\n" "Step 2: Filter sale_date within July 2026.\n" "Step 3: GROUP BY region." ), "sql": ( "SELECT region, SUM(revenue) AS total_revenue " "FROM sales WHERE sale_date >= '2026-07-01' " "AND sale_date < '2026-08-01' GROUP BY region;" ), }, { "question": "Show total quantity sold by category for August 2026.", "reasoning": ( "Step 1: Need SUM(quantity) grouped by category.\n" "Step 2: Filter sale_date within August 2026.\n" "Step 3: GROUP BY category." ), "sql": ( "SELECT category, SUM(quantity) AS total_qty " "FROM sales WHERE sale_date >= '2026-08-01' " "AND sale_date < '2026-09-01' GROUP BY category;" ), }, { "question": "Which store had the highest revenue overall?", "reasoning": ( "Step 1: Need SUM(revenue) grouped by store.\n" "Step 2: ORDER BY total revenue descending.\n" "Step 3: LIMIT 1 to get the top store." ), "sql": ( "SELECT store, SUM(revenue) AS total_revenue FROM sales " "GROUP BY store ORDER BY total_revenue DESC LIMIT 1;" ), }, ] # --------------------------------------------------------------------------- # 3. Embedder + FAISS index (rebuilt whenever the library grows) # --------------------------------------------------------------------------- embedder = SentenceTransformer("all-MiniLM-L6-v2") def build_index(): texts = [t["question"] for t in TRACE_LIBRARY] embs = embedder.encode(texts, normalize_embeddings=True) idx = faiss.IndexFlatIP(embs.shape[1]) idx.add(np.array(embs, dtype="float32")) return idx index = build_index() def retrieve_similar(query: str, k: int = 2): q_emb = embedder.encode([query], normalize_embeddings=True) scores, idxs = index.search(np.array(q_emb, dtype="float32"), k) return [{**TRACE_LIBRARY[i], "score": float(s)} for s, i in zip(scores[0], idxs[0])] # --------------------------------------------------------------------------- # 4. Generation + execution # --------------------------------------------------------------------------- def generate_sql(new_question: str, k: int = 2): examples = retrieve_similar(new_question, k=k) exemplar_block = "" for i, ex in enumerate(examples, 1): exemplar_block += ( f"\nExample {i} (similarity {ex['score']:.2f}):\n" f"Question: {ex['question']}\n" f"Reasoning:\n{ex['reasoning']}\n" f"SQL: {ex['sql']}\n" ) prompt = f"""You write DuckDB SQL for a retail sales table. {SCHEMA_DESC} Here are verified past questions with their reasoning and correct SQL. Use the SAME reasoning pattern for the new question. {exemplar_block} Now answer this new question the same way. Give brief reasoning, then output ONLY the final SQL on a line starting with "SQL:". Question: {new_question} Reasoning:""" client = Groq(api_key=os.environ["GROQ_API_KEY"]) response = client.chat.completions.create( model="llama-3.3-70b-versatile", messages=[{"role": "user", "content": prompt}], temperature=0.1, ) text = response.choices[0].message.content match = re.search(r"SQL:\s*(.+)", text, re.DOTALL) sql = match.group(1).strip() if match else None return examples, text, sql def run_pipeline(question: str): if not question or not question.strip(): return "Enter a question first.", "", None, "" if not os.environ.get("GROQ_API_KEY"): return "GROQ_API_KEY is not set (add it as a Space secret).", "", None, "" examples, full_response, sql = generate_sql(question) retrieved_md = "\n".join( f"- **[{ex['score']:.2f}]** {ex['question']}" for ex in examples ) if not sql: return retrieved_md, full_response, None, "No SQL extracted from model output." try: df = con.execute(sql).fetchdf() return retrieved_md, f"```sql\n{sql}\n```", df, "Query ran successfully." except Exception as e: return retrieved_md, f"```sql\n{sql}\n```", None, f"SQL failed: {e}" def add_verified_example(question: str, reasoning: str, sql: str): """Manually confirm a generated query was correct, then grow the trace library with it. This is the feedback loop that makes the library compound in usefulness over time -- only call this after a human has checked the result.""" global index if not question or not sql: return "Need both a question and SQL to add.", library_preview() TRACE_LIBRARY.append({ "question": question.strip(), "reasoning": reasoning.strip() or "Verified example.", "sql": sql.strip(), }) index = build_index() return f"Added. Library now has {len(TRACE_LIBRARY)} verified examples.", library_preview() def library_preview(): return "\n".join(f"{i+1}. {t['question']}" for i, t in enumerate(TRACE_LIBRARY)) # --------------------------------------------------------------------------- # 5. Gradio UI # --------------------------------------------------------------------------- with gr.Blocks(title="Text-to-SQL Trace RAG Prototype") as demo: gr.Markdown( "# Text-to-SQL Trace RAG — Prototype\n" "Retrieves **verified past question -> SQL -> reasoning** examples " "(not documents) and uses them as few-shot patterns to generate new SQL. " "Toy `sales` table included -- swap for a real schema for production use." ) with gr.Tab("Ask a question"): question_in = gr.Textbox( label="Ask a BI question", placeholder="e.g. What is total revenue by category for August 2026?", ) ask_btn = gr.Button("Generate & Run", variant="primary") retrieved_out = gr.Markdown(label="Retrieved examples used") sql_out = gr.Markdown(label="Generated SQL") result_out = gr.Dataframe(label="Query result") status_out = gr.Textbox(label="Status", interactive=False) ask_btn.click( run_pipeline, inputs=[question_in], outputs=[retrieved_out, sql_out, result_out, status_out], ) with gr.Tab("Add verified example"): gr.Markdown( "Once you've manually confirmed a generated query was correct, " "add it here so future similar questions retrieve it too." ) v_question = gr.Textbox(label="Question") v_reasoning = gr.Textbox(label="Reasoning (steps taken)", lines=3) v_sql = gr.Textbox(label="Verified correct SQL", lines=3) add_btn = gr.Button("Add to trace library") add_status = gr.Textbox(label="Status", interactive=False) lib_view = gr.Textbox( label="Current library", value=library_preview(), lines=8, interactive=False ) add_btn.click( add_verified_example, inputs=[v_question, v_reasoning, v_sql], outputs=[add_status, lib_view], ) if __name__ == "__main__": demo.launch()