app.py

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