| """ |
| 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 |
|
|
| |
| |
| |
| 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) |
| """ |
|
|
| |
| |
| |
| |
| 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;" |
| ), |
| }, |
| ] |
|
|
| |
| |
| |
| 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])] |
|
|
|
|
| |
| |
| |
| 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)) |
|
|
|
|
| |
| |
| |
| 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() |
