Spaces:
Sleeping
Sleeping
| import os | |
| import re | |
| import json | |
| import duckdb | |
| import faiss | |
| import numpy as np | |
| from fastapi import FastAPI | |
| from fastapi.responses import HTMLResponse | |
| from pydantic import BaseModel | |
| from sentence_transformers import SentenceTransformer | |
| from groq import Groq | |
| # ============================================================ | |
| # 1. DATABASE | |
| # ============================================================ | |
| 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 CORPUS | |
| # ============================================================ | |
| TRACE_LIBRARY = [ | |
| { | |
| "question": "What is total revenue by region for July 2026?", | |
| "trace": { | |
| "intent": "aggregation", | |
| "metric": "revenue", | |
| "aggregation": "SUM", | |
| "dimension": "region", | |
| "time_filter": "July 2026", | |
| "grouping": "region", | |
| "ordering": None, | |
| "limit": None | |
| }, | |
| "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.", | |
| "trace": { | |
| "intent": "aggregation", | |
| "metric": "quantity", | |
| "aggregation": "SUM", | |
| "dimension": "category", | |
| "time_filter": "August 2026", | |
| "grouping": "category", | |
| "ordering": None, | |
| "limit": None | |
| }, | |
| "sql": """ | |
| SELECT | |
| category, | |
| SUM(quantity) AS total_quantity | |
| 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?", | |
| "trace": { | |
| "intent": "ranking", | |
| "metric": "revenue", | |
| "aggregation": "SUM", | |
| "dimension": "store", | |
| "time_filter": None, | |
| "grouping": "store", | |
| "ordering": "total revenue DESC", | |
| "limit": 1 | |
| }, | |
| "sql": """ | |
| SELECT | |
| store, | |
| SUM(revenue) AS total_revenue | |
| FROM sales | |
| GROUP BY store | |
| ORDER BY total_revenue DESC | |
| LIMIT 1; | |
| """ | |
| } | |
| ] | |
| # ============================================================ | |
| # 3. EMBEDDING MODEL | |
| # ============================================================ | |
| embedder = SentenceTransformer( | |
| "sentence-transformers/all-MiniLM-L6-v2" | |
| ) | |
| # ============================================================ | |
| # 4. FAISS TRACE INDEX | |
| # ============================================================ | |
| def build_index(): | |
| questions = [item["question"] for item in TRACE_LIBRARY] | |
| embeddings = embedder.encode(questions, normalize_embeddings=True) | |
| embeddings = np.asarray(embeddings, dtype="float32") | |
| index = faiss.IndexFlatIP(embeddings.shape[1]) | |
| index.add(embeddings) | |
| return index | |
| index = build_index() | |
| # Minimum cosine similarity (embeddings are normalized, so FAISS IP | |
| # score == cosine similarity) for a retrieved trace to be considered | |
| # "relevant" enough to inject into the prompt as a verified example. | |
| # | |
| # 0.55 is a reasonable starting cutoff for all-MiniLM-L6-v2 on short | |
| # question-style text: near-duplicate/paraphrased questions usually | |
| # score 0.7+, loosely related questions land around 0.4-0.6, and | |
| # unrelated questions fall below 0.3. Tune against your own | |
| # TRACE_LIBRARY as it grows -- it's a heuristic, not a guarantee. | |
| MIN_SIMILARITY_SCORE = 0.55 | |
| def retrieve_traces(question, k=3, min_score=MIN_SIMILARITY_SCORE): | |
| if not TRACE_LIBRARY: | |
| return [] | |
| k = min(k, len(TRACE_LIBRARY)) | |
| query_embedding = embedder.encode([question], normalize_embeddings=True) | |
| query_embedding = np.asarray(query_embedding, dtype="float32") | |
| scores, indices = index.search(query_embedding, k) | |
| results = [] | |
| for score, idx in zip(scores[0], indices[0]): | |
| if idx < 0: | |
| continue | |
| score = float(score) | |
| # Below-threshold matches are dropped rather than passed to | |
| # the LLM as "verified" -- a weak match is worse than no | |
| # match, since it can anchor the model to the wrong pattern. | |
| if score < min_score: | |
| continue | |
| item = TRACE_LIBRARY[idx].copy() | |
| item["score"] = score | |
| results.append(item) | |
| return results | |
| # ============================================================ | |
| # 5. GROQ CLIENT | |
| # ============================================================ | |
| def get_client(): | |
| api_key = os.environ.get("GROQ_API_KEY") | |
| if not api_key: | |
| raise RuntimeError("GROQ_API_KEY is not set.") | |
| return Groq(api_key=api_key) | |
| # ============================================================ | |
| # 6. FORMAT RETRIEVED TRACES | |
| # ============================================================ | |
| def format_retrieved_traces(examples): | |
| output = "" | |
| for i, example in enumerate(examples, start=1): | |
| output += f""" | |
| ============================== | |
| VERIFIED EXAMPLE {i} | |
| Similarity: {example["score"]:.3f} | |
| ============================== | |
| Question: | |
| {example["question"]} | |
| Structured Solution Trace: | |
| {json.dumps(example["trace"], indent=2)} | |
| Verified SQL: | |
| {example["sql"]} | |
| """ | |
| return output | |
| # ============================================================ | |
| # 7. GENERATE STRUCTURED TRACE | |
| # ============================================================ | |
| def generate_trace(question, retrieved_examples): | |
| client = get_client() | |
| examples = format_retrieved_traces(retrieved_examples) | |
| prompt = f""" | |
| You are an expert retail Text-to-SQL planner. | |
| Database schema: | |
| {SCHEMA_DESC} | |
| You have access to VERIFIED previous SQL solutions. | |
| {examples} | |
| New user question: | |
| {question} | |
| Create a structured SQL solution trace. | |
| IMPORTANT: | |
| Do NOT provide hidden chain-of-thought. | |
| Only provide the following explicit fields: | |
| intent | |
| metric | |
| aggregation | |
| dimension | |
| time_filter | |
| grouping | |
| ordering | |
| limit | |
| Use null when a field does not apply. | |
| Return ONLY valid JSON. | |
| Example format: | |
| {{ | |
| "intent": "aggregation", | |
| "metric": "revenue", | |
| "aggregation": "SUM", | |
| "dimension": "category", | |
| "time_filter": "August 2026", | |
| "grouping": "category", | |
| "ordering": null, | |
| "limit": null | |
| }} | |
| """ | |
| response = client.chat.completions.create( | |
| model="llama-3.3-70b-versatile", | |
| messages=[{"role": "user", "content": prompt}], | |
| temperature=0 | |
| ) | |
| text = response.choices[0].message.content.strip() | |
| text = re.sub(r"```json\s*", "", text) | |
| text = re.sub(r"```\s*", "", text) | |
| try: | |
| return json.loads(text) | |
| except Exception: | |
| return {"raw_trace": text} | |
| # ============================================================ | |
| # 8. GENERATE SQL FROM TRACE | |
| # ============================================================ | |
| def generate_sql(question, trace, retrieved_examples): | |
| client = get_client() | |
| examples = format_retrieved_traces(retrieved_examples) | |
| prompt = f""" | |
| You are an expert DuckDB Text-to-SQL generator. | |
| Database schema: | |
| {SCHEMA_DESC} | |
| VERIFIED PREVIOUS EXAMPLES: | |
| {examples} | |
| NEW USER QUESTION: | |
| {question} | |
| NEW STRUCTURED SOLUTION TRACE: | |
| {json.dumps(trace, indent=2)} | |
| Generate SQL that correctly answers the question. | |
| Rules: | |
| 1. Use only the provided schema. | |
| 2. Use DuckDB SQL. | |
| 3. Follow the structured trace. | |
| 4. Do not invent columns. | |
| 5. Handle date ranges correctly. | |
| 6. Return ONLY SQL. | |
| """ | |
| response = client.chat.completions.create( | |
| model="llama-3.3-70b-versatile", | |
| messages=[{"role": "user", "content": prompt}], | |
| temperature=0 | |
| ) | |
| sql = response.choices[0].message.content.strip() | |
| sql = re.sub(r"```sql\s*", "", sql, flags=re.IGNORECASE) | |
| sql = re.sub(r"```\s*", "", sql) | |
| return sql.strip() | |
| # ============================================================ | |
| # 9. TRACE LIBRARY PREVIEW | |
| # ============================================================ | |
| def library_preview(): | |
| return [item["question"] for item in TRACE_LIBRARY] | |
| # ============================================================ | |
| # 10. FASTAPI APP | |
| # ============================================================ | |
| app = FastAPI(title="Text-to-SQL Trace RAG") | |
| class AskRequest(BaseModel): | |
| question: str | |
| class AddTraceRequest(BaseModel): | |
| question: str | |
| trace: dict | |
| sql: str | |
| def ask(req: AskRequest): | |
| question = req.question.strip() | |
| if not question: | |
| return { | |
| "retrieved": [], | |
| "trace": None, | |
| "sql": "", | |
| "columns": [], | |
| "rows": [], | |
| "status": "Please enter a question." | |
| } | |
| try: | |
| examples = retrieve_traces(question, k=3) | |
| retrieved = [ | |
| {"question": e["question"], "score": e["score"]} | |
| for e in examples | |
| ] | |
| trace = generate_trace(question, examples) | |
| sql = generate_sql(question, trace, examples) | |
| try: | |
| df = con.execute(sql).fetchdf() | |
| columns = list(df.columns) | |
| rows = df.to_dict(orient="records") | |
| status = "SQL executed successfully." | |
| except Exception as e: | |
| columns = [] | |
| rows = [] | |
| status = f"SQL execution failed: {str(e)}" | |
| return { | |
| "retrieved": retrieved, | |
| "trace": trace, | |
| "sql": sql, | |
| "columns": columns, | |
| "rows": rows, | |
| "status": status | |
| } | |
| except Exception as e: | |
| return { | |
| "retrieved": [], | |
| "trace": None, | |
| "sql": "", | |
| "columns": [], | |
| "rows": [], | |
| "status": f"Pipeline error: {str(e)}" | |
| } | |
| def add_trace(req: AddTraceRequest): | |
| global index | |
| question = req.question.strip() | |
| sql = req.sql.strip() | |
| if not question: | |
| return {"status": "Question is required.", "library": library_preview()} | |
| if not sql: | |
| return {"status": "SQL is required.", "library": library_preview()} | |
| TRACE_LIBRARY.append({ | |
| "question": question, | |
| "trace": req.trace, | |
| "sql": sql | |
| }) | |
| # Rebuild FAISS index. Note: this mutates process-global state | |
| # with no locking -- fine for a single-worker prototype, not | |
| # safe under concurrent requests or multiple uvicorn workers. | |
| index = build_index() | |
| return { | |
| "status": f"Verified trace added. Trace corpus now contains {len(TRACE_LIBRARY)} examples.", | |
| "library": library_preview() | |
| } | |
| def get_library(): | |
| return {"library": library_preview()} | |
| # ============================================================ | |
| # 11. FRONTEND (single inline HTML page, vanilla JS, no build step) | |
| # ============================================================ | |
| INDEX_HTML = """ | |
| <!DOCTYPE html> | |
| <html lang="en"> | |
| <head> | |
| <meta charset="UTF-8"> | |
| <title>Text-to-SQL Trace RAG</title> | |
| <style> | |
| body { font-family: -apple-system, Segoe UI, Roboto, sans-serif; max-width: 900px; margin: 40px auto; padding: 0 20px; color: #1a1a1a; } | |
| h1 { font-size: 1.4rem; } | |
| .tabs { display: flex; gap: 8px; margin-bottom: 20px; } | |
| .tab-btn { padding: 8px 16px; border: 1px solid #ccc; background: #f5f5f5; cursor: pointer; border-radius: 6px; } | |
| .tab-btn.active { background: #2563eb; color: white; border-color: #2563eb; } | |
| .tab { display: none; } | |
| .tab.active { display: block; } | |
| textarea, input[type=text] { width: 100%; box-sizing: border-box; padding: 8px; font-family: inherit; } | |
| textarea { min-height: 70px; } | |
| button.run { margin-top: 10px; padding: 10px 18px; background: #2563eb; color: white; border: none; border-radius: 6px; cursor: pointer; } | |
| pre { background: #f5f5f5; padding: 12px; border-radius: 6px; overflow-x: auto; white-space: pre-wrap; } | |
| table { border-collapse: collapse; width: 100%; margin-top: 8px; } | |
| th, td { border: 1px solid #ddd; padding: 6px 10px; text-align: left; font-size: 0.9rem; } | |
| th { background: #f0f0f0; } | |
| .status { margin-top: 8px; font-size: 0.9rem; color: #444; } | |
| .section-title { font-weight: 600; margin-top: 18px; margin-bottom: 6px; } | |
| </style> | |
| </head> | |
| <body> | |
| <h1>Text-to-SQL Trace RAG</h1> | |
| <p>Retrieves verified SQL solution traces (FAISS, cosine ≥ 0.55) and uses them to ground new Text-to-SQL generation.</p> | |
| <div class="tabs"> | |
| <button class="tab-btn active" onclick="showTab('ask', event)">Ask Question</button> | |
| <button class="tab-btn" onclick="showTab('add', event)">Add Verified Trace</button> | |
| </div> | |
| <div id="ask" class="tab active"> | |
| <textarea id="question" placeholder="e.g. What is total revenue by category for August 2026?"></textarea> | |
| <button class="run" onclick="runAsk()">Generate & Run</button> | |
| <div class="section-title">Retrieved Similar Traces</div> | |
| <pre id="retrieved">-</pre> | |
| <div class="section-title">Structured Solution Trace</div> | |
| <pre id="trace">-</pre> | |
| <div class="section-title">Generated SQL</div> | |
| <pre id="sql">-</pre> | |
| <div class="section-title">Query Result</div> | |
| <div id="result">-</div> | |
| <div class="status" id="askStatus"></div> | |
| </div> | |
| <div id="add" class="tab"> | |
| <p>Only add an example after a human has verified the trace and SQL are correct and produce the expected result.</p> | |
| <label>Question</label> | |
| <input type="text" id="vQuestion"> | |
| <label>Structured Trace (JSON)</label> | |
| <textarea id="vTrace">{}</textarea> | |
| <label>Verified SQL</label> | |
| <textarea id="vSql"></textarea> | |
| <button class="run" onclick="runAddTrace()">Add Verified Trace</button> | |
| <div class="status" id="addStatus"></div> | |
| <div class="section-title">Current Trace Corpus</div> | |
| <pre id="library">-</pre> | |
| </div> | |
| <script> | |
| function showTab(name, evt) { | |
| document.querySelectorAll('.tab').forEach(el => el.classList.remove('active')); | |
| document.querySelectorAll('.tab-btn').forEach(el => el.classList.remove('active')); | |
| document.getElementById(name).classList.add('active'); | |
| evt.target.classList.add('active'); | |
| if (name === 'add') refreshLibrary(); | |
| } | |
| async function runAsk() { | |
| const question = document.getElementById('question').value; | |
| document.getElementById('askStatus').textContent = 'Running...'; | |
| const res = await fetch('/api/ask', { | |
| method: 'POST', | |
| headers: {'Content-Type': 'application/json'}, | |
| body: JSON.stringify({question}) | |
| }); | |
| const data = await res.json(); | |
| document.getElementById('retrieved').textContent = | |
| data.retrieved.length | |
| ? data.retrieved.map(r => `${r.score.toFixed(3)} ${r.question}`).join('\\n') | |
| : 'No verified trace scored above the similarity threshold -- generated zero-shot (no examples injected into the prompt).'; | |
| document.getElementById('trace').textContent = JSON.stringify(data.trace, null, 2); | |
| document.getElementById('sql').textContent = data.sql; | |
| const resultDiv = document.getElementById('result'); | |
| if (data.columns.length) { | |
| let html = '<table><tr>' + data.columns.map(c => `<th>${c}</th>`).join('') + '</tr>'; | |
| data.rows.forEach(row => { | |
| html += '<tr>' + data.columns.map(c => `<td>${row[c]}</td>`).join('') + '</tr>'; | |
| }); | |
| html += '</table>'; | |
| resultDiv.innerHTML = html; | |
| } else { | |
| resultDiv.textContent = 'No rows returned.'; | |
| } | |
| document.getElementById('askStatus').textContent = data.status; | |
| } | |
| async function runAddTrace() { | |
| const question = document.getElementById('vQuestion').value; | |
| const sql = document.getElementById('vSql').value; | |
| let trace; | |
| try { | |
| trace = JSON.parse(document.getElementById('vTrace').value); | |
| } catch (e) { | |
| document.getElementById('addStatus').textContent = 'Trace must be valid JSON.'; | |
| return; | |
| } | |
| const res = await fetch('/api/add_trace', { | |
| method: 'POST', | |
| headers: {'Content-Type': 'application/json'}, | |
| body: JSON.stringify({question, trace, sql}) | |
| }); | |
| const data = await res.json(); | |
| document.getElementById('addStatus').textContent = data.status; | |
| document.getElementById('library').textContent = data.library.map((q,i) => `${i+1}. ${q}`).join('\\n'); | |
| } | |
| async function refreshLibrary() { | |
| const res = await fetch('/api/library'); | |
| const data = await res.json(); | |
| document.getElementById('library').textContent = data.library.map((q,i) => `${i+1}. ${q}`).join('\\n'); | |
| } | |
| </script> | |
| </body> | |
| </html> | |
| """ | |
| def home(): | |
| return INDEX_HTML | |
| # ============================================================ | |
| # 12. START | |
| # ============================================================ | |
| if __name__ == "__main__": | |
| import uvicorn | |
| uvicorn.run(app, host="0.0.0.0", port=7860) |