Upload app.py

app.py

@@ -0,0 +1,2187 @@
+# %%
+import os
+import uuid
+import sqlite3
+import datetime
+import json
+import re
+import time
+from typing import Dict, Any, List, Tuple, Optional
+import hashlib
+
+import pandas as pd
+from groq import Groq
+from langchain_community.vectorstores import Chroma
+from langchain_community.embeddings import HuggingFaceEmbeddings
+
+# %%
+# 0. Global config
+DB_PATH = "olist.db"
+DATA_DIR = "data"
+
+# Groq Model
+GROQ_MODEL_NAME = "llama-3.3-70b-versatile"
+
+GROQ_API_KEY = os.getenv("GROQ_API_KEY")
+if not GROQ_API_KEY:
+    raise ValueError("GROQ_API_KEY not found.")
+
+groq_client = Groq(api_key=GROQ_API_KEY)
+
+# Embedding model
+EMBED_MODEL_NAME = "sentence-transformers/all-mpnet-base-v2"
+
+# Logging paths
+PROMPT_LOG_PATH = "prompt_logs.txt"
+RUN_LOG_PATH = "run_logs.txt"
+
+# %%
+# SINGLE, CORRECT, PERSISTENT CHROMA CLIENT
+
+import os
+import chromadb
+from chromadb.config import Settings
+
+CHROMA_PERSIST_DIR = os.path.abspath("./chroma_data")
+os.makedirs(CHROMA_PERSIST_DIR, exist_ok=True)
+
+_CHROMA_CLIENT = None
+
+def get_chroma_client():
+    global _CHROMA_CLIENT
+    if _CHROMA_CLIENT is None:
+        _CHROMA_CLIENT = chromadb.PersistentClient(
+            path=CHROMA_PERSIST_DIR,
+            settings=Settings(
+                anonymized_telemetry=False,
+            ),
+        )
+    return _CHROMA_CLIENT
+
+
+# %%
+
+
+# %%
+# 1. Logging helpers
+
+
+def log_prompt(tag: str, prompt: str) -> None:
+    """
+    Append the full prompt to a log file, with a tag and timestamp.
+    """
+    timestamp = datetime.datetime.now().isoformat(timespec="seconds")
+    header = f"\n\n================ {tag} @ {timestamp} ================\n"
+    with open(PROMPT_LOG_PATH, "a", encoding="utf-8") as f:
+        f.write(header)
+        f.write(prompt)
+        f.write("\n================ END PROMPT ================\n")
+
+
+def log_run_event(tag: str, content: str) -> None:
+    """
+    Append model response, final SQL, and error info into a run log.
+    """
+    timestamp = datetime.datetime.now().isoformat(timespec="seconds")
+    header = f"\n\n================ {tag} @ {timestamp} ================\n"
+    with open(RUN_LOG_PATH, "a", encoding="utf-8") as f:
+        f.write(header)
+        f.write(content)
+        f.write("\n================ END EVENT ================\n")
+
+
+
+# %%
+# 2. Feedback table + helpers
+
+def init_feedback_table(conn: sqlite3.Connection) -> None:
+    """
+    Create (or upgrade) a table to capture user feedback on model answers.
+    """
+    conn.execute("""
+    CREATE TABLE IF NOT EXISTS user_feedback (
+        id INTEGER PRIMARY KEY AUTOINCREMENT,
+        created_at TEXT NOT NULL,
+        question TEXT NOT NULL,
+        generated_sql TEXT,
+        model_answer TEXT,
+        rating TEXT CHECK(rating IN ('good','bad')) NOT NULL,
+        comment TEXT,
+        corrected_sql TEXT
+    )
+    """)
+    conn.commit()
+
+
+def record_feedback(
+    conn: sqlite3.Connection,
+    question: str,
+    generated_sql: str,
+    model_answer: str,
+    rating: str,                     # "good" or "bad"
+    comment: Optional[str] = None,
+    corrected_sql: Optional[str] = None,
+) -> None:
+    """
+    Store user feedback about a particular model answer / SQL query.
+    If corrected_sql is provided, it is treated as an external correction.
+    """
+    rating = rating.lower()
+    if rating not in ("good", "bad"):
+        raise ValueError("rating must be 'good' or 'bad'")
+
+    ts = datetime.datetime.now().isoformat(timespec="seconds")
+    conn.execute(
+        """
+        INSERT INTO user_feedback (
+            created_at, question, generated_sql, model_answer,
+            rating, comment, corrected_sql
+        )
+        VALUES (?, ?, ?, ?, ?, ?, ?)
+        """,
+        (ts, question, generated_sql, model_answer, rating, comment, corrected_sql),
+    )
+    conn.commit()
+
+
+def get_last_feedback_for_question(
+    conn: sqlite3.Connection,
+    question: str,
+) -> Optional[Dict[str, Any]]:
+    """
+    Return the most recent feedback row for this question (if any).
+    """
+    cur = conn.cursor()
+    cur.execute(
+        """
+        SELECT created_at, generated_sql, model_answer,
+               rating, comment, corrected_sql
+        FROM user_feedback
+        WHERE question = ?
+        ORDER BY created_at DESC
+        LIMIT 1
+        """,
+        (question,),
+    )
+    row = cur.fetchone()
+    if not row:
+        return None
+
+    return {
+        "created_at": row[0],
+        "generated_sql": row[1],
+        "model_answer": row[2],
+        "rating": row[3],
+        "comment": row[4],
+        "corrected_sql": row[5],
+    }
+
+
+
+# %%
+# 3. Database setup (from CSVs)   
+
+def init_db() -> sqlite3.Connection:
+    """
+    Load all CSVs from the data/ folder into a local SQLite DB.
+    Table names are derived from file names (without .csv).
+    """
+    conn = sqlite3.connect(DB_PATH, check_same_thread=False)
+
+    csv_files = [
+        "olist_customers_dataset.csv",
+        "olist_orders_dataset.csv",
+        "olist_order_items_dataset.csv",
+        "olist_products_dataset.csv",
+        "olist_order_reviews_dataset.csv",
+        "olist_order_payments_dataset.csv",
+        "product_category_name_translation.csv",
+        "olist_sellers_dataset.csv",
+        "olist_geolocation_dataset.csv",
+    ]
+
+    for fname in csv_files:
+        path = os.path.join(DATA_DIR, fname)
+        print(path)
+        if not os.path.exists(path):
+            print(f"CSV not found: {path} - skipping")
+            continue
+
+        table_name = os.path.splitext(fname)[0]
+        print(f"Loading {path} into table {table_name}...")
+        df = pd.read_csv(path)
+        df.to_sql(table_name, conn, if_exists="replace", index=False)
+
+    
+    init_feedback_table(conn)
+
+    return conn
+
+
+conn = init_db()
+
+
+
+# %%
+# 4. Manual docs for Olist tables
+
+OLIST_DOCS: Dict[str, Dict[str, Any]] = {
+    "olist_customers_dataset": {
+        "description": "Customer master data, one row per customer_id (which can change over time for the same end-user).",
+        "columns": {
+            "customer_id": "Primary key for this table. Unique technical identifier for a customer at a point in time. Used to join with olist_orders_dataset.customer_id.",
+            "customer_unique_id": "Stable unique identifier for the end-user. A single customer_unique_id can map to multiple customer_id records over time.",
+            "customer_zip_code_prefix": "Customer ZIP/postal code prefix. Used to join with olist_geolocation_dataset.geolocation_zip_code_prefix.",
+            "customer_city": "Customer's city as captured at the time of the order or registration.",
+            "customer_state": "Customer's state (two-letter Brazilian state code, e.g. SP, RJ)."
+        }
+    },
+    "olist_orders_dataset": {
+        "description": "Customer orders placed on the Olist marketplace, one row per order.",
+        "columns": {
+            "order_id": "Primary key. Unique identifier for each order. Used to join with items, payments, and reviews.",
+            "customer_id": "Foreign key to olist_customers_dataset.customer_id indicating who placed the order.",
+            "order_status": "Current lifecycle status of the order (e.g. created, shipped, delivered, canceled, unavailable).",
+            "order_purchase_timestamp": "Timestamp when the customer completed the purchase (event time for order placement).",
+            "order_approved_at": "Timestamp when the payment was approved by the system or financial gateway.",
+            "order_delivered_carrier_date": "Timestamp when the order was handed over by the seller to the carrier/logistics provider.",
+            "order_delivered_customer_date": "Timestamp when the carrier reported the order as delivered to the final customer.",
+            "order_estimated_delivery_date": "Estimated delivery date promised to the customer at checkout."
+        }
+    },
+    "olist_order_items_dataset": {
+        "description": "Order line items, one row per product per order.",
+        "columns": {
+            "order_id": "Foreign key to olist_orders_dataset.order_id. Multiple order_items can belong to the same order.",
+            "order_item_id": "Sequential item number within an order (1, 2, 3, ...). Uniquely identifies a line inside an order.",
+            "product_id": "Foreign key to olist_products_dataset.product_id representing the purchased product.",
+            "seller_id": "Foreign key to olist_sellers_dataset.seller_id representing the seller that fulfilled this item.",
+            "shipping_limit_date": "Deadline for the seller to hand the item over to the carrier for shipping.",
+            "price": "Item price paid by the customer for this line (in BRL, not including freight).",
+            "freight_value": "Freight (shipping) cost attributed to this line item (in BRL)."
+        }
+    },
+    "olist_products_dataset": {
+        "description": "Product catalog with physical and category attributes, one row per product.",
+        "columns": {
+            "product_id": "Primary key. Unique identifier for each product. Used to join with olist_order_items_dataset.product_id.",
+            "product_category_name": "Product category name in Portuguese. Join to product_category_name_translation.product_category_name for English.",
+            "product_name_lenght": "Number of characters in the product name (field name misspelled as 'lenght' in the original dataset).",
+            "product_description_lenght": "Number of characters in the product description (also misspelled as 'lenght').",
+            "product_photos_qty": "Number of product images associated with the listing.",
+            "product_weight_g": "Product weight in grams.",
+            "product_length_cm": "Product length in centimeters (package dimension).",
+            "product_height_cm": "Product height in centimeters (package dimension).",
+            "product_width_cm": "Product width in centimeters (package dimension)."
+        }
+    },
+    "olist_order_reviews_dataset": {
+        "description": "Post-purchase customer reviews and satisfaction scores, one row per review.",
+        "columns": {
+            "review_id": "Primary key. Unique identifier for each review record.",
+            "order_id": "Foreign key to olist_orders_dataset.order_id for the reviewed order.",
+            "review_score": "Star rating given by the customer on a 1–5 scale (5 = very satisfied, 1 = very dissatisfied).",
+            "review_comment_title": "Optional short text title or summary of the review.",
+            "review_comment_message": "Optional detailed free-text comment describing the customer experience.",
+            "review_creation_date": "Date when the customer created the review.",
+            "review_answer_timestamp": "Timestamp when Olist or the seller responded to the review (if applicable)."
+        }
+    },
+    "olist_order_payments_dataset": {
+        "description": "Payments associated with orders, one row per payment record (order can have multiple payments).",
+        "columns": {
+            "order_id": "Foreign key to olist_orders_dataset.order_id.",
+            "payment_sequential": "Sequence number for multiple payments of the same order (1 for first payment, 2 for second, etc.).",
+            "payment_type": "Payment method used (e.g. credit_card, boleto, voucher, debit_card).",
+            "payment_installments": "Number of installments chosen by the customer for this payment.",
+            "payment_value": "Monetary amount paid in this payment record (in BRL)."
+        }
+    },
+    "product_category_name_translation": {
+        "description": "Lookup table mapping Portuguese product category names to English equivalents.",
+        "columns": {
+            "product_category_name": "Product category name in Portuguese as used in olist_products_dataset.",
+            "product_category_name_english": "Translated product category name in English."
+        }
+    },
+    "olist_sellers_dataset": {
+        "description": "Seller master data, one row per seller operating on the Olist marketplace.",
+        "columns": {
+            "seller_id": "Primary key. Unique identifier for each seller. Used to join with olist_order_items_dataset.seller_id.",
+            "seller_zip_code_prefix": "Seller ZIP/postal code prefix, used to join with olist_geolocation_dataset.geolocation_zip_code_prefix.",
+            "seller_city": "City where the seller is located.",
+            "seller_state": "State where the seller is located (two-letter Brazilian state code)."
+        }
+    },
+    "olist_geolocation_dataset": {
+        "description": "Geolocation reference data for ZIP code prefixes in Brazil, not unique per prefix.",
+        "columns": {
+            "geolocation_zip_code_prefix": "ZIP/postal code prefix, used to link customers and sellers via zip code.",
+            "geolocation_lat": "Latitude coordinate of the location.",
+            "geolocation_lng": "Longitude coordinate of the location.",
+            "geolocation_city": "City name for the location.",
+            "geolocation_state": "State code for the location (two-letter Brazilian state code)."
+        }
+    }
+}
+
+
+
+# %%
+# 5. Schema extractor + metadata YAML 
+
+def extract_schema_metadata(connection: sqlite3.Connection) -> Dict[str, Any]:
+    """
+    Introspect SQLite tables, columns and foreign key relationships.
+    """
+    cursor = connection.cursor()
+    cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
+    tables = [row[0] for row in cursor.fetchall()]
+
+    metadata: Dict[str, Any] = {"tables": {}}
+
+    for table in tables:
+        cursor.execute(f"PRAGMA table_info('{table}')")
+        cols = cursor.fetchall()
+
+        # Manual docs for the table
+        table_docs = OLIST_DOCS.get(table, {})
+        table_desc: str = table_docs.get(
+            "description",
+            f"Table '{table}' from Olist dataset."
+        )
+        column_docs: Dict[str, str] = table_docs.get("columns", {})
+
+        columns_meta: Dict[str, str] = {}
+        for c in cols:
+            col_name = c[1]
+            col_type = c[2] or "TEXT"
+
+            if col_name in column_docs:
+                columns_meta[col_name] = column_docs[col_name]
+            else:
+                columns_meta[col_name] = f"Column '{col_name}' of type {col_type}"
+
+        cursor.execute(f"PRAGMA foreign_key_list('{table}')")
+        fk_rows = cursor.fetchall()
+        relationships: List[str] = []
+        for fk in fk_rows:
+            ref_table = fk[2]
+            from_col = fk[3]
+            to_col = fk[4]
+            relationships.append(
+                f"{table}.{from_col} → {ref_table}.{to_col} (foreign key)"
+            )
+
+        metadata["tables"][table] = {
+            "description": table_desc,
+            "columns": columns_meta,
+            "relationships": relationships,
+        }
+
+    return metadata
+
+
+schema_metadata = extract_schema_metadata(conn)
+
+
+def build_schema_yaml(metadata: Dict[str, Any]) -> str:
+    """
+    Render metadata dict into a YAML-style string.
+    """
+    lines: List[str] = ["tables:"]
+    for tname, tinfo in metadata["tables"].items():
+        lines.append(f"  {tname}:")
+        desc = tinfo.get("description", "").replace('"', "'")
+        lines.append(f'    description: "{desc}"')
+        lines.append("    columns:")
+        for col_name, col_desc in tinfo.get("columns", {}).items():
+            col_desc_clean = col_desc.replace('"', "'")
+            lines.append(f'      {col_name}: "{col_desc_clean}"')
+        rels = tinfo.get("relationships", [])
+        if rels:
+            lines.append("    relationships:")
+            for rel in rels:
+                rel_clean = rel.replace('"', "'")
+                lines.append(f'      - "{rel_clean}"')
+    return "\n".join(lines)
+
+
+schema_yaml = build_schema_yaml(schema_metadata)
+
+
+
+# %%
+# 6. Build schema documents for RAG (taking samples from the table)
+
+def build_schema_documents(
+    connection: sqlite3.Connection,
+    schema_metadata: Dict[str, Any],
+    sample_rows: int = 5,
+) -> Tuple[List[str], List[Dict[str, Any]]]:
+    """
+    Build one rich RAG document per table, using schema_metadata.
+
+    Each document includes:
+    - Table name
+    - Table description
+    - Columns with type + description
+    - Relationships (FKs)
+    - A few sample rows
+    """
+    cursor = connection.cursor()
+    cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
+    tables = [row[0] for row in cursor.fetchall()]
+
+    docs: List[str] = []
+    metadatas: List[Dict[str, Any]] = []
+
+    for table in tables:
+        tmeta = schema_metadata["tables"][table]
+        table_desc = tmeta.get("description", "")
+        columns_meta = tmeta.get("columns", {})
+        relationships = tmeta.get("relationships", [])
+
+        # Use PRAGMA to get types, then enrich with descriptions
+        cursor.execute(f"PRAGMA table_info('{table}')")
+        cols = cursor.fetchall()
+
+        col_lines = []
+        for c in cols:
+            col_name = c[1]
+            col_type = c[2] or "TEXT"
+            col_desc = columns_meta.get(col_name, f"Column '{col_name}' of type {col_type}")
+            col_lines.append(f"- {col_name} ({col_type}): {col_desc}")
+
+        # Sample rows
+        try:
+            sample_df = pd.read_sql_query(
+                f"SELECT * FROM '{table}' LIMIT {sample_rows}",
+                connection,
+            )
+            sample_text = sample_df.to_markdown(index=False)
+        except Exception:
+            sample_text = "(could not fetch sample rows)"
+
+        # Relationships block
+        rel_block = ""
+        if relationships:
+            rel_block = "Relationships:\n" + "\n".join(
+                f"- {rel}" for rel in relationships
+            ) + "\n"
+
+        doc_text = (
+            f"Table: {table}\n"
+            f"Description: {table_desc}\n\n"
+            f"Columns:\n" + "\n".join(col_lines) + "\n\n"
+            f"{rel_block}\n"
+            f"Example rows:\n{sample_text}\n"
+        )
+
+        docs.append(doc_text)
+        metadatas.append({
+            "doc_type": "table_schema",
+            "table_name": table,
+        })
+
+    return docs, metadatas
+
+
+# Build RAG texts + metadata
+schema_docs, schema_doc_metas = build_schema_documents(conn, schema_metadata)
+
+RAG_TEXTS: List[str] = []
+RAG_METADATAS: List[Dict[str, Any]] = []
+
+# 1) Per-table docs
+RAG_TEXTS.extend(schema_docs)
+RAG_METADATAS.extend(schema_doc_metas)
+
+# 2) Global YAML as a separate doc
+RAG_TEXTS.append("SCHEMA_METADATA_YAML:\n" + schema_yaml)
+RAG_METADATAS.append({"doc_type": "global_schema"})
+
+
+# %%
+def build_store_final():
+    embedding_model = HuggingFaceEmbeddings(
+        model_name=EMBED_MODEL_NAME,
+        encode_kwargs={"normalize_embeddings": True},
+    )
+
+    rag_store = Chroma(
+        client=get_chroma_client(),
+        collection_name="rag_schema_store",
+        embedding_function=embedding_model,
+    )
+
+    if rag_store._collection.count() == 0:
+        rag_store.add_texts(
+            texts=RAG_TEXTS,
+            metadatas=RAG_METADATAS,
+        )
+
+    rag_retriever = rag_store.as_retriever(
+        search_kwargs={"k": 3, "filter": {"doc_type": "table_schema"}}
+    )
+
+    return rag_store, rag_retriever
+
+# Initialize once
+rag_store, rag_retriever = build_store_final()
+
+
+# %%
+from langchain.vectorstores import Chroma
+from langchain.embeddings import HuggingFaceEmbeddings
+
+_sql_cache_store = None
+_sql_embedding_fn = None
+
+SQL_CACHE_COLLECTION = "sql_cache_mpnet"
+
+def get_sql_cache_store():
+    global _sql_cache_store, _sql_embedding_fn
+
+    if _sql_cache_store is not None:
+        return _sql_cache_store
+
+    if _sql_embedding_fn is None:
+        _sql_embedding_fn = HuggingFaceEmbeddings(
+            model_name=EMBED_MODEL_NAME,
+            encode_kwargs={"normalize_embeddings": True},
+        )
+
+    _sql_cache_store = Chroma(
+        client=get_chroma_client(),
+        collection_name=SQL_CACHE_COLLECTION,
+        embedding_function=_sql_embedding_fn,
+    )
+
+    return _sql_cache_store
+
+
+# %%
+def sanitize_metadata_for_chroma(metadata: dict) -> dict:
+    safe = {}
+    for k, v in (metadata or {}).items():
+        if v is None:
+            continue
+        if isinstance(v, (int, float, bool)):
+            safe[k] = v
+        elif isinstance(v, str):
+            safe[k] = v
+        else:
+            safe[k] = str(v)
+    return safe
+
+
+# %%
+def normalize_question_strict(q: str) -> str:
+    """
+    Deterministic normalization for exact cache hits.
+    """
+    if not q:
+        return ""
+    q = q.lower().strip()
+    q = re.sub(r"[^\w\s]", "", q)
+    q = re.sub(r"\s+", " ", q)
+    return q
+
+
+# %%
+# Helper: normalize question
+
+def normalize_question_text(q: str) -> str:
+    if not q:
+        return ""
+    q = q.strip().lower()
+    q = re.sub(r"[^\w\s]", " ", q)
+    q = re.sub(r"\s+", " ", q).strip()
+    return q
+
+# Helper: compute success rate
+def compute_success_rate(md: dict) -> float:
+    sc = md.get("success_count", 0) or 0
+    tf = md.get("total_feedbacks", 0) or 0
+    if tf <= 0:
+        return 0.0
+    return float(sc) / float(tf)
+
+# Insert initial cache entry (no feedback yet)
+def cache_sql_answer_initial(question: str, sql: str, answer_md: str, store=None, extra_metadata: dict = None):
+    """
+    Insert a cached entry when you run a query and want to cache it regardless of feedback.
+    initial metrics: views=1, success_count=0, total_feedbacks=0, success_rate=1.0
+    """
+    if store is None:
+        store = get_sql_cache_store()
+
+    ident = uuid.uuid4().hex
+    norm = normalize_question_text(question)
+    md = {
+        "id": ident,
+        "normalized_question": norm,
+        "sql": sql,
+        "answer_md": answer_md,
+        "saved_at": time.time(),
+        "views": 1,
+        "success_count": 0,
+        "total_feedbacks": 0,
+        "success_rate": 1.0,
+    }
+    if extra_metadata:
+        md.update(extra_metadata)
+
+    # Use store's API; 
+    store.add_texts([question], metadatas=[md])
+    return ident
+
+
+
+# %%
+import time
+import logging
+from typing import Optional, List, Dict, Any
+from difflib import SequenceMatcher
+
+_logger = logging.getLogger(__name__)
+
+# -------------------------------------------------------------------------
+# Utility helpers
+# -------------------------------------------------------------------------
+
+def _now_ts() -> float:
+    return time.time()
+
+def similarity_score(a: str, b: str) -> float:
+    return SequenceMatcher(None, (a or ""), (b or "")).ratio()
+
+# -------------------------------------------------------------------------
+# Persist helper 
+# -------------------------------------------------------------------------
+def langchain_upsert(
+    store,
+    text: str,
+    metadata: dict,
+    cache_id: str,
+):
+    safe_md = sanitize_metadata_for_chroma(metadata)
+
+    try:
+        store.add_texts(
+            texts=[text],
+            metadatas=[safe_md],
+            ids=[cache_id],
+        )
+    except TypeError:
+        store.add_texts(
+            texts=[text],
+            metadatas=[safe_md],
+        )
+
+
+# -------------------------------------------------------------------------
+# Cache insert / update
+# -------------------------------------------------------------------------
+def cache_sql_answer_dedup(
+    question: str,
+    sql: str,
+    answer_md: str,
+    metadata: dict,
+    store,
+):
+    norm_q_semantic = normalize_text(question)
+    norm_q_exact = normalize_question_strict(question)
+
+    cache_id = generate_cache_id(question, sql)
+    now = _now_ts()
+
+    md = {
+        # identity
+        "cache_id": cache_id,
+        "sql": sql,
+        "answer_md": answer_md,
+
+        # exact match key
+        "normalized_question": norm_q_exact,
+
+        # timestamps
+        "saved_at": metadata.get("saved_at", now),
+        "last_updated_at": now,
+        "last_viewed_at": metadata.get("last_viewed_at", 0),
+
+        # metrics
+        "good_count": metadata.get("good_count", 0),
+        "bad_count": metadata.get("bad_count", 0),
+        "total_feedbacks": metadata.get("total_feedbacks", 0),
+        "success_rate": metadata.get("success_rate", 0.5),
+        "views": metadata.get("views", 0),
+    }
+
+    langchain_upsert(
+        store=store,
+        text=norm_q_semantic,   # semantic vector
+        metadata=md,
+        cache_id=cache_id,
+    )
+
+    return {
+        "question": question,
+        "sql": sql,
+        "answer_md": answer_md,
+        "metadata": md,
+    }
+
+# -------------------------------------------------------------------------
+# Find exact cached entry (question + SQL)
+# -------------------------------------------------------------------------
+
+def find_cached_doc_by_sql(question: str, sql: str, store):
+    cache_id = generate_cache_id(question, sql)
+    coll = getattr(store, "_collection", None)
+
+    if coll and hasattr(coll, "get"):
+        try:
+            res = coll.get(ids=[cache_id])
+            if res and res.get("metadatas"):
+                md = res["metadatas"][0]
+                return {
+                    "id": cache_id,
+                    "question": question,
+                    "sql": md.get("sql"),
+                    "answer_md": md.get("answer_md"),
+                    "metadata": md,
+                }
+        except Exception:
+            pass
+
+    return None
+
+# -------------------------------------------------------------------------
+# Retrieve cached answers ranked primarily by success rate
+# -------------------------------------------------------------------------
+
+import re
+import unicodedata
+
+
+def normalize_text(text: str) -> str:
+    if not text:
+        return ""
+    text = unicodedata.normalize("NFKD", text).encode("ascii", "ignore").decode("ascii")
+    text = text.lower()
+    text = re.sub(r"[^\w\s]", " ", text)
+    text = re.sub(r"\s+", " ", text).strip()
+    return text
+
+import hashlib
+
+def generate_cache_id(question: str, sql: str) -> str:
+    q = normalize_text(question)
+    s = (sql or "").strip()
+    key = f"{q}||{s}".encode("utf-8")
+    return hashlib.sha1(key).hexdigest()
+
+
+def rank_cached_candidates(candidates: list[dict]) -> dict:
+    """
+    Deterministic quality-first ranking.
+    """
+    candidates.sort(
+        key=lambda c: (
+            -float(c["metadata"].get("success_rate", 0.5)),
+            -int(c["metadata"].get("good_count", 0)),
+            int(c["metadata"].get("bad_count", 0)),
+            -float(c["metadata"].get("last_updated_at", 0)),
+        )
+    )
+    return candidates[0]
+
+
+
+def retrieve_exact_cached_sql(question: str, store):
+    """
+    Exact question match, but still quality-ranked.
+    """
+    norm_q = normalize_question_strict(question)
+    coll = store._collection
+
+    try:
+        res = coll.get(where={"normalized_question": norm_q})
+        if not res or not res.get("metadatas"):
+            return None
+
+        candidates = []
+        for md in res["metadatas"]:
+            candidates.append({
+                "matched_question": question,
+                "sql": md["sql"],
+                "answer_md": md.get("answer_md", ""),
+                "distance": 0.0,
+                "metadata": md,
+            })
+
+        return rank_cached_candidates(candidates)
+
+    except Exception:
+        return None
+
+
+def retrieve_best_cached_sql(
+    question: str,
+    store,
+    max_distance: float = 0.25,
+):
+    norm_q = normalize_text(question)
+    results = store.similarity_search_with_score(norm_q, k=20)
+
+    candidates = []
+
+    for doc, score in results:
+        distance = float(score)
+        if distance > max_distance:
+            continue
+
+        md = doc.metadata or {}
+        if "sql" not in md:
+            continue
+
+        candidates.append({
+            "matched_question": doc.page_content,
+            "sql": md["sql"],
+            "answer_md": md.get("answer_md", ""),
+            "distance": distance,
+            "metadata": md,
+
+            # ranking signals (FORCED numeric)
+            "success_rate": float(md.get("success_rate", 0.5)),
+            "good": int(md.get("good_count", 0)),
+            "bad": int(md.get("bad_count", 0)),
+            "views": int(md.get("views", 0)),
+            "last_updated": float(md.get("last_updated_at", 0)),
+        })
+
+    if not candidates:
+        return None
+
+    # QUALITY-FIRST RANKING
+    candidates.sort(
+    key=lambda c: (
+        -c["success_rate"],
+        -c["good"],
+        c["bad"],
+        c["distance"],
+        -c["last_updated"],
+    )
+)
+
+
+    return candidates[0]
+
+# -------------------------------------------------------------------------
+# Increment views
+# -------------------------------------------------------------------------
+
+def increment_cache_views(metadata: dict, store):
+    if not metadata:
+        return False
+
+    cache_id = metadata.get("cache_id")
+    if not cache_id:
+        return False
+
+    md = dict(metadata)
+    md["views"] = int(md.get("views", 0)) + 1
+    md["last_viewed_at"] = _now_ts()
+    md["last_updated_at"] = _now_ts()
+    md["saved_at"] = md.get("saved_at", _now_ts())
+
+    try:
+        langchain_upsert(
+            store=store,
+            text=normalize_text(md.get("normalized_question", "")),
+            metadata=md,
+            cache_id=cache_id,
+        )
+        return True
+    except Exception:
+        _logger.exception("increment_cache_views failed")
+        return False
+
+
+# Update metrics on feedback
+def update_cache_on_feedback(
+    question: str,
+    original_doc_md: dict,
+    user_marked_good: bool,
+    llm_corrected_sql: str | None,
+    llm_corrected_answer_md: str | None,
+    store,
+):
+    if not original_doc_md:
+        return
+
+    md = dict(original_doc_md["metadata"])
+    cache_id = md["cache_id"]
+
+    # ---- feedback counts ----
+    if user_marked_good:
+        md["good_count"] = md.get("good_count", 0) + 1
+    else:
+        md["bad_count"] = md.get("bad_count", 0) + 1
+
+    md["total_feedbacks"] = md.get("total_feedbacks", 0) + 1
+    md["success_rate"] = (
+        md["good_count"] / md["total_feedbacks"]
+        if md["total_feedbacks"] > 0 else 0.5
+    )
+
+    # ---- timestamps ----
+    md["saved_at"] = md.get("saved_at", _now_ts())   # preserve
+    md["last_updated_at"] = _now_ts()
+
+    langchain_upsert(
+        store=store,
+        text=normalize_text(question),
+        metadata=md,
+        cache_id=cache_id,
+    )
+
+    # -------------------------
+    # Corrected SQL --> NEW ENTRY
+    # -------------------------
+    if llm_corrected_sql and llm_corrected_answer_md:
+        cache_sql_answer_dedup(
+            question=question,
+            sql=llm_corrected_sql,
+            answer_md=llm_corrected_answer_md,
+            metadata={
+                "good_count": 1,
+                "bad_count": 0,
+                "total_feedbacks": 1,
+                "success_rate": 1.0,
+                "views": 0,
+                "saved_at": _now_ts(),
+            },
+            store=store,
+        )
+
+
+# %%
+# ### 8. Groq LLM via LangChain
+
+from langchain_groq import ChatGroq
+import re
+import gradio as gr
+
+llm = ChatGroq(model=GROQ_MODEL_NAME, groq_api_key=GROQ_API_KEY)
+
+# %%
+def get_rag_context(question: str) -> str:
+    """
+    Retrieve the most relevant schema documents for the question.
+    """
+    docs = rag_retriever.invoke(question)
+    return "\n\n---\n\n".join(d.page_content for d in docs)
+
+def clean_sql(sql: str) -> str:
+    sql = sql.strip()
+    if "```" in sql:
+        sql = sql.replace("```sql", "").replace("```", "").strip()
+    return sql
+
+def extract_sql_from_markdown(text: str) -> str:
+    """
+    Extract the first ```sql ... ``` block from LLM output.
+    If not found, return the whole text.
+    """
+    match = re.search(r"```sql(.*?)```", text, flags=re.DOTALL | re.IGNORECASE)
+    if match:
+        return match.group(1).strip()
+    return text.strip()
+
+def extract_explanation_after_marker(text: str, marker: str = "EXPLANATION:") -> str:
+    """
+    After the given marker, return the rest of the text as explanation.
+    """
+    idx = text.upper().find(marker.upper())
+    if idx == -1:
+        return text.strip()
+    return text[idx + len(marker):].strip()
+
+# 6b. General / descriptive questions
+
+GENERAL_DESC_KEYWORDS = [
+    "what is this dataset about",
+    "what is this data about",
+    "describe this dataset",
+    "describe the dataset",
+    "dataset overview",
+    "data overview",
+    "summary of the dataset",
+    "explain this dataset",
+]
+
+
+def is_general_question(question: str) -> bool:
+    """
+    Detect high-level descriptive questions where we should answer
+    directly from schema context instead of generating SQL.
+    """
+    q = question.lower().strip()
+    return any(key in q for key in GENERAL_DESC_KEYWORDS)
+
+
+def answer_general_question(question: str) -> str:
+    """
+    Use the RAG schema docs to generate a rich, high-level description
+    of the Olist dataset for conceptual questions.
+    """
+    rag_context = get_rag_context(question)
+
+    system_instructions = """
+You are a data documentation expert.
+
+You will be given:
+- Schema documentation for the Olist dataset (tables, descriptions, columns, relationships).
+- A high-level user question like "what is this dataset about?".
+
+Your job:
+- Write a clear, structured overview of the dataset.
+- Explain the main entities (customers, orders, items, products, sellers, payments, reviews, geolocation).
+- Mention typical analysis use-cases (delivery performance, customer behavior, seller performance, product/category analysis, etc.).
+- Target a non-technical person.
+
+Do NOT write SQL. Answer in Markdown.
+"""
+
+    prompt = (
+        system_instructions
+        + "\n\n=== SCHEMA CONTEXT ===\n"
+        + rag_context
+        + "\n\n=== USER QUESTION ===\n"
+        + question
+        + "\n\nDetailed dataset overview:"
+    )
+
+    log_prompt("GENERAL_DATASET_QUESTION", prompt)
+
+    response = llm.invoke(prompt)
+    log_run_event("RAW_MODEL_RESPONSE_GENERAL_DATASET", response.content)
+
+    return response.content.strip()
+
+# %%
+# ### 9. SQL execution / validation
+
+def execute_sql(sql: str, connection: sqlite3.Connection) -> Tuple[Optional[pd.DataFrame], Optional[str]]:
+    """
+    Execute SQL on SQLite and return a DataFrame, else return an error.
+    """
+    try:
+        df = pd.read_sql_query(sql, connection)
+        return df, None
+    except Exception as e:
+        return None, str(e)
+
+def validate_sql(sql: str, connection: sqlite3.Connection) -> Tuple[bool, Optional[str]]:
+    """
+    Basic SQL validator:
+    - Uses EXPLAIN QUERY PLAN to detect syntax or schema issues.
+    """
+    try:
+        cursor = connection.cursor()
+        cursor.execute(f"EXPLAIN QUERY PLAN {sql}")
+        return True, None
+    except Exception as e:
+        return False, str(e)
+
+# %%
+# ### 10. SQL Generation + Repair with LLM (feedback-aware)
+
+def build_sql_review_prompt(
+    question: str,
+    generated_sql: str,
+    user_feedback_comment: str,
+    rag_context: str,
+) -> str:
+    """
+    Prompt to let the LLM compare its SQL with the user's feedback,
+    decide if the join/logic is wrong, and produce a corrected SQL + explanation.
+
+    We explicitly allow the model to say:
+    - "query was already correct, user mistaken" OR
+    - "query was wrong, here is the fix".
+    """
+    prompt = f"""
+You previously generated the following SQL for a SQLite database:
+
+```sql
+{generated_sql}
+
+The user now says this query is WRONG and provided this feedback:
+"{user_feedback_comment}"
+
+TASKS:
+
+    Compare the SQL with the database schema (given in the context) and the user's feedback.
+
+    Decide whether the query is actually correct or incorrect.
+    Make sure that the user clearly specifies why they think it is incorrect. 
+    It can be if they are unsatisfied with the numbers, or the logic is incorrect, or SQL is invalid etc.
+    If any reason is not clearly specified, return the previous result as it is.
+
+        If it is already correct, keep it unchanged and explain why the user might be mistaken.
+
+        If it is incorrect (e.g., wrong joins, missing filters, wrong aggregation), fix it.
+
+        Produce a corrected SQL query that better answers the question.
+
+        If the original query is already correct, just repeat the same SQL.
+
+        Explain in a few sentences WHO is correct (you or the user) and WHY.
+
+DATABASE SCHEMA (partial, from RAG):
+{rag_context}
+
+User question:
+{question}
+
+Return your answer in this format:
+
+CORRECTED_SQL:
+
+-- your (possibly unchanged) SQL here
+SELECT ...
+
+EXPLANATION:
+Your explanation here, clearly stating whether:
+
+    the original query was correct or not, and
+
+    how your corrected SQL addresses the issue (or why it didn't need changes).
+    """
+    return prompt.strip()
+
+
+# %%
+# Review and Correct SQL based on Feedback
+
+def review_and_correct_sql_with_llm(
+    question: str,
+    generated_sql: str,
+    user_feedback_comment: str,
+    rag_context: str,
+) -> Tuple[str, str]:
+    """
+    Ask the LLM to compare its SQL with user's feedback, decide what is wrong (or not),
+    and propose a corrected SQL (possibly unchanged) + explanation.
+
+    Returns:
+        corrected_sql, explanation
+    """
+    prompt = build_sql_review_prompt(
+        question=question,
+        generated_sql=generated_sql,
+        user_feedback_comment=user_feedback_comment,
+        rag_context=rag_context,
+    )
+
+    log_prompt("SQL_REVIEW_FEEDBACK", prompt)
+    response = llm.invoke(prompt)
+    log_run_event("RAW_MODEL_RESPONSE_SQL_REVIEW", response.content)
+
+    # Try to extract SQL; if none, fall back to original
+    extracted_sql = extract_sql_from_markdown(response.content)
+    corrected_sql = clean_sql(extracted_sql) if extracted_sql else generated_sql
+    if not corrected_sql.strip():
+        corrected_sql = generated_sql
+
+    explanation = extract_explanation_after_marker(
+        response.content,
+        marker="EXPLANATION:",
+    )
+    return corrected_sql, explanation
+
+# %%
+# Generate SQL
+
+def generate_sql(question: str, rag_context: str) -> str:
+    """
+    Generate SQL using LLM, but also pull in any past user feedback
+    (corrected SQL) for this question as external guidance.
+    """
+    # External correction from past feedback, if any
+    last_fb = get_last_feedback_for_question(conn, question)
+    if last_fb and last_fb.get("corrected_sql"):
+        previous_feedback_block = f"""
+EXTERNAL USER FEEDBACK FROM PAST RUNS:
+
+Previous generated SQL:
+{last_fb['generated_sql']}
+
+Corrected SQL (preferred reference):
+{last_fb['corrected_sql']}
+
+User comment & prior explanation:
+{last_fb.get('comment') or '(none)'}
+
+You must avoid repeating the same mistake and should follow the logic of
+the corrected SQL when appropriate, while still reasoning from the schema.
+"""
+    else:
+        previous_feedback_block = ""
+
+    system_instructions = f"""
+You are a senior data analyst writing SQL for a SQLite database.
+
+You will be given:
+
+- A description of available tables, columns, and relationships (schema + YAML metadata).
+- A natural language question from the user.
+
+Your job:
+
+- Write ONE valid SQLite SQL query that answers the question.
+- ONLY use tables and columns that exist in the schema_context.
+- Use correct JOINS (Left, Right, Inner, Outer, Full Outer etc.) with ON conditions.
+- Do not use DROP, INSERT, UPDATE, DELETE or other destructive operations.
+- Always use floating-point division for percentage calculations using 1.0 * numerator / denominator,
+  and round to 2 decimals when appropriate.
+
+{previous_feedback_block}
+"""
+
+    prompt = (
+        system_instructions
+        + "\n\n=== RAG CONTEXT ===\n"
+        + rag_context
+        + "\n\n=== USER QUESTION ===\n"
+        + question
+        + "\n\nSQL query:"
+    )
+
+    log_prompt("SQL_GENERATION", prompt)
+    response = llm.invoke(prompt)
+    log_run_event("RAW_MODEL_RESPONSE_SQL_GENERATION", response.content)
+
+    sql = clean_sql(response.content)
+    return sql
+
+
+# %%
+# Repair SQL
+
+def repair_sql(
+    question: str,
+    rag_context: str,
+    bad_sql: str,
+    error_message: str,
+) -> str:
+    """
+    Ask the LLM to correct a failing SQL query.
+    """
+    system_instructions = """
+You are a senior data analyst fixing an existing SQL query for a SQLite database.
+
+You will be given:
+
+- Schema context (tables, columns, relationships).
+- The user's question.
+- A previously generated SQL query that failed.
+- The SQLite error message.
+
+Your job:
+
+- Diagnose why the query failed.
+- Rewrite ONE valid SQLite SQL query that answers the question.
+- ONLY use tables and columns that exist in the schema_context.
+- Use correct JOINS (Left, Right, Inner, Outer, Full Outer etc.) with ON conditions.
+- Do not use DROP, INSERT, UPDATE, DELETE or other destructive operations.
+- Return ONLY the corrected SQL query, no explanation or markdown.
+"""
+
+    prompt = (
+        system_instructions
+        + "\n\n=== RAG CONTEXT ===\n"
+        + rag_context
+        + "\n\n=== USER QUESTION ===\n"
+        + question
+        + "\n\n=== PREVIOUS (FAILING) SQL ===\n"
+        + bad_sql
+        + "\n\n=== SQLITE ERROR ===\n"
+        + error_message
+        + "\n\nCorrected SQL query:"
+    )
+
+    log_prompt("SQL_REPAIR", prompt)
+    response = llm.invoke(prompt)
+    log_run_event("RAW_MODEL_RESPONSE_SQL_REPAIR", response.content)
+
+    sql = clean_sql(response.content)
+    return sql
+
+
+# %%
+### 11. Result summarization
+
+def summarize_results(
+    question: str,
+    sql: str,
+    df: Optional[pd.DataFrame],
+    rag_context: str,
+    error: Optional[str] = None,
+) -> str:
+    """
+    Ask the LLM to produce a concise, human-readable answer.
+    """
+    system_instructions = """
+    You are a senior data analyst.
+
+    You will be given:
+
+    The user's question.
+    The final SQL that was executed.
+    A small preview of the query result (as a Markdown table, if available).
+    Optional error information if the query failed.
+
+    Your job:
+
+    Provide a clear, concise answer in Markdown.
+    If the result is numeric / aggregated, explain what it means in business terms.
+    If there was an error, explain it simply and suggest how the user could rephrase.
+    Do NOT show raw SQL unless it is helpful to the user.
+    """
+
+    # Build a markdown table preview if we have data
+    if df is not None and not df.empty:
+        preview_rows = min(len(df), 50)
+        df_preview_md = df.head(preview_rows).to_markdown(index=False)
+    else:
+        df_preview_md = "(no rows returned)"
+
+    prompt = (
+        system_instructions
+        + "\n\n=== USER QUESTION ===\n"
+        + question
+        + "\n\n=== EXECUTED SQL ===\n"
+        + sql
+        + "\n\n=== QUERY RESULT PREVIEW ===\n"
+        + df_preview_md
+        + "\n\n=== RAG CONTEXT (schema) ===\n"
+        + rag_context
+    )
+
+    if error:
+        prompt += "\n\n=== ERROR ===\n" + error
+
+    # Logging helpers assumed to exist
+    log_prompt("RESULT_SUMMARY", prompt)
+
+    response = llm.invoke(prompt)
+    log_run_event("RAW_MODEL_RESPONSE_RESULT_SUMMARY", response.content)
+
+    return response.content.strip()
+
+
+# %%
+def backend_pipeline(question: str):
+    """
+    STRICT cache-first backend.
+
+    Priority:
+      1. Exact question cache hit
+      2. Best semantic cache hit
+      3. LLM fallback
+    """
+
+    # ------------------------------------------------------------------
+    # Guards
+    # ------------------------------------------------------------------
+    if not question or not question.strip():
+        return (
+            "Please type a question.",
+            pd.DataFrame(),
+            "",
+            "",
+            "",
+            "",
+            pd.DataFrame(),
+            [],
+            [],
+            False,
+            4,
+            "**Feedback attempts remaining: 4**",
+            gr.update(value="", visible=False),
+        )
+
+    attempts_left = 4
+    attempts_text = f"**Feedback attempts remaining: {attempts_left}**"
+
+    # ------------------------------------------------------------------
+    # General / descriptive questions
+    # ------------------------------------------------------------------
+    if is_general_question(question):
+        overview_md = answer_general_question(question)
+        return (
+            overview_md,
+            pd.DataFrame(),
+            "",
+            "",
+            question,
+            overview_md,
+            pd.DataFrame(),
+            [],
+            [],
+            False,
+            attempts_left,
+            attempts_text,
+            gr.update(value="", visible=False),
+        )
+
+    store = get_sql_cache_store()
+
+    # ------------------------------------------------------------------
+    # STEP 0A: EXACT CACHE LOOKUP (deterministic)
+    # ------------------------------------------------------------------
+    cached = retrieve_exact_cached_sql(question, store)
+
+    # ------------------------------------------------------------------
+    # STEP 0B: SEMANTIC CACHE LOOKUP (ranked)
+    # ------------------------------------------------------------------
+    if cached is None:
+        cached = retrieve_best_cached_sql(
+            question=question,
+            store=store,
+            max_distance=0.25,
+        )
+
+    # ------------------------------------------------------------------
+    # CACHE HIT PATH
+    # ------------------------------------------------------------------
+    if cached:
+        try:
+            increment_cache_views(cached["metadata"], store=store)
+        except Exception:
+            pass
+
+        rag_context = get_rag_context(question)
+
+        header = (
+            "### Cache Hit\n"
+            f"- **Matched question:** \"{cached['matched_question']}\"\n"
+            f"- **Success rate:** {cached['metadata'].get('success_rate', 0.5):.2f}\n"
+            f"- **Similarity distance:** {cached['distance']:.4f}\n\n"
+            "---\n\n"
+        )
+
+        answer_md = header + (cached.get("answer_md") or "")
+
+        try:
+            df, exec_error = execute_sql(cached["sql"], conn)
+            if exec_error:
+                df = pd.DataFrame()
+                answer_md += f"\n\n Error re-running cached SQL: `{exec_error}`"
+        except Exception as e:
+            df = pd.DataFrame()
+            answer_md += f"\n\n Exception re-running cached SQL: `{e}`"
+
+        md = cached["metadata"]
+
+        stats_md = (
+            f"**Cached entry stats**\n\n"
+            f"- **Success rate:** {md.get('success_rate', 0.5):.2f}  \n"
+            f"- **Total feedbacks:** {md.get('total_feedbacks', 0)}  \n"
+            f"- **Good / Bad:** {md.get('good_count', 0)} / {md.get('bad_count', 0)}  \n"
+            f"- **Views:** {md.get('views', 0)}  \n"
+            f"- **Saved at:** "
+            f"{datetime.datetime.fromtimestamp(md.get('saved_at')).strftime('%Y-%m-%d %H:%M') if md.get('saved_at') else 'unknown'}  \n"
+            f"- **Last updated:** "
+            f"{datetime.datetime.fromtimestamp(md.get('last_updated_at')).strftime('%Y-%m-%d %H:%M') if md.get('last_updated_at') else 'unknown'}\n\n"
+            f"**SQL preview:**\n\n```sql\n{cached['sql']}\n```\n"
+        )
+
+        return (
+            answer_md,
+            df,
+            cached["sql"],
+            rag_context,
+            question,
+            answer_md,
+            df,
+            [],
+            [],
+            False,
+            attempts_left,
+            attempts_text,
+            gr.update(value=stats_md, visible=True),
+        )
+
+    # ------------------------------------------------------------------
+    # STEP 1: LLM FLOW (NO CACHE HIT)
+    # ------------------------------------------------------------------
+    rag_context = get_rag_context(question)
+
+    sql = generate_sql(question, rag_context)
+    original_sql = sql
+
+    is_valid, validation_error = validate_sql(sql, conn)
+    repaired = False
+
+    if not is_valid and validation_error:
+        repaired_sql = repair_sql(question, rag_context, sql, validation_error)
+        repaired_valid, repaired_error = validate_sql(repaired_sql, conn)
+        if repaired_valid:
+            sql = repaired_sql
+            repaired = True
+            validation_error = None
+        else:
+            validation_error = repaired_error or validation_error
+
+    df, exec_error = (None, None)
+    if not validation_error:
+        df, exec_error = execute_sql(sql, conn)
+    else:
+        exec_error = validation_error
+
+    summary_text = summarize_results(
+        question=question,
+        sql=sql,
+        df=df,
+        rag_context=rag_context,
+        error=exec_error,
+    )
+
+    sql_status = []
+    if exec_error:
+        sql_status.append(f"**Error:** `{exec_error}`")
+    else:
+        sql_status.append("Query ran successfully.")
+    if repaired:
+        sql_status.append("_Note: SQL was auto-repaired._")
+
+    sql_status.append("\n**Final SQL used:**\n")
+    sql_status.append(f"```sql\n{sql}\n```")
+
+    answer_md = summary_text + "\n\n---\n\n" + "\n".join(sql_status)
+    df_preview = df if df is not None and exec_error is None else pd.DataFrame()
+
+    # ------------------------------------------------------------------
+    # STEP 2: CACHE LLM RESULT (ALWAYS)
+    # ------------------------------------------------------------------
+    try:
+        cache_sql_answer_dedup(
+            question=question,
+            sql=sql,
+            answer_md=answer_md,
+            metadata={
+                "good_count": 0,
+                "bad_count": 0,
+                "total_feedbacks": 0,
+                "success_rate": 0.5,
+                "views": 1,
+                "saved_at": _now_ts(),
+            },
+            store=store,
+        )
+    except Exception:
+        _logger.exception("backend_pipeline: failed to cache LLM result")
+
+    return (
+        answer_md,
+        df_preview,
+        sql,
+        rag_context,
+        question,
+        answer_md,
+        df_preview,
+        [],
+        [],
+        False,
+        attempts_left,
+        attempts_text,
+        gr.update(value="", visible=False),
+    )
+
+
+# %%
+def _looks_like_sql(text: str) -> bool:
+    """Quick heuristic: does text contain SQL keywords / SELECT ?"""
+    if not text:
+        return False
+    return bool(re.search(r"\bselect\b|\bfrom\b|\bwhere\b|\bjoin\b|\bgroup by\b|\border by\b", text, flags=re.I))
+
+
+def is_feedback_sufficient(feedback_text: str) -> bool:
+    """
+    Heuristic to decide whether the user's free-text feedback is actionable.
+
+    Returns True if:
+      - length >= 20 characters AND contains a signal word (e.g., 'filter', 'year', 'should', 'instead', 'missing', 'wrong', digits),
+      OR
+      - it looks like SQL (user pasted corrected SQL),
+      OR
+      - length >= 60 characters (long feedback).
+    """
+    if not feedback_text:
+        return False
+
+    text = feedback_text.strip()
+    if len(text) >= 60:
+        return True
+
+    if _looks_like_sql(text):
+        return True
+
+    # look for signal words that indicate specificity
+    signal_words = [
+        "filter", "where", "year", "month", "should", "instead", "expected",
+        "wrong", "missing", "aggregate", "sum", "avg", "count", "distinct",
+        "join", "left join", "inner join", "group by", "order by", "date",
+        "range", "exclude", "include", "only"
+    ]
+    lower = text.lower()
+    signals = sum(1 for w in signal_words if w in lower)
+    if signals >= 1 and len(text) >= 20:
+        return True
+
+    # short hits like "numbers look off" are insufficient
+    return False
+
+
+def build_followup_prompt_for_user(sample_feedback: str = "") -> str:
+    """
+    Deterministic follow-up question to ask the user when feedback is vague.
+    Returns a friendly prompt that the UI can display to the user.
+    """
+    base = (
+        "Thanks — I need a bit more detail to act on this feedback.\n\n"
+        "Please tell me one (or more) of the following so I can check and correct the result:\n\n"
+        "1. Which part looks wrong — the **numbers**, the **aggregation** (sum/count/avg),\n"
+        "   the **time range** (year/month), or the **filters** applied?\n"
+        "2. If you expected a different number, what was the expected number (and how was it computed)?\n"
+        "3. If you have a corrected SQL snippet, paste it (I can run and compare it).\n\n"
+        "Examples you can copy-paste:\n"
+    )
+    examples = (
+        "- \"I think the query should count DISTINCT customer_unique_id, not customer_id.\"\n"
+        "- \"This looks off for year 2018 — I expected the count for 2018 to be ~40k.\"\n"
+        "- \"Please exclude canceled orders (order_status = 'canceled').\"\n"
+        "- \"SELECT COUNT(DISTINCT customer_unique_id) FROM olist_customers_dataset;\"\n"
+    )
+    hint = "\nIf you prefer, just paste a corrected SQL snippet and I'll run it and compare."
+    prompt = base + examples + hint
+    if sample_feedback:
+        prompt = f"I saw your feedback: \"{sample_feedback}\"\n\n" + prompt
+    return prompt
+
+
+
+# %%
+def feedback_pipeline_interactive(
+    feedback_rating: str,
+    feedback_comment: str,
+    last_sql: str,
+    last_rag_context: str,
+    last_question: str,
+    last_answer_md: str,
+    last_df: pd.DataFrame,
+    feedback_sql: str,
+    attempts_left: int,
+):
+    rating = (feedback_rating or "").strip().lower()
+    comment = (feedback_comment or "").strip()
+    attempts_left = int(attempts_left or 0)
+
+    # ---------------- Guard ----------------
+    if not last_question or not last_sql:
+        return (
+            last_answer_md,
+            last_df,
+            last_sql,
+            last_rag_context,
+            last_question,
+            last_answer_md,
+            last_df,
+            False,
+            "",
+            attempts_left,
+        )
+
+    if rating not in ("correct", "wrong"):
+        return (
+            last_answer_md + "\n\n Please select **Correct** or **Wrong**.",
+            last_df,
+            last_sql,
+            last_rag_context,
+            last_question,
+            last_answer_md,
+            last_df,
+            False,
+            "",
+            attempts_left,
+        )
+
+    # ============================================================
+    # CORRECT -> no attempt decrement
+    # ============================================================
+    if rating == "correct":
+        original_doc = find_cached_doc_by_sql(
+            last_question, last_sql, store=get_sql_cache_store()
+        )
+
+        update_cache_on_feedback(
+            question=last_question,
+            original_doc_md=original_doc,
+            user_marked_good=True,
+            llm_corrected_sql=None,
+            llm_corrected_answer_md=None,
+            store=get_sql_cache_store(),
+        )
+
+        record_feedback(
+            conn=conn,
+            question=last_question,
+            generated_sql=last_sql,
+            model_answer=last_answer_md,
+            rating="good",
+            comment=comment or None,
+            corrected_sql=None,
+        )
+
+        return (
+            last_answer_md + "\n\n **Feedback recorded as GOOD.**",
+            last_df,
+            last_sql,
+            last_rag_context,
+            last_question,
+            last_answer_md,
+            last_df,
+            False,
+            "",
+            attempts_left,
+        )
+
+    # ============================================================
+    #  WRONG -> decrement immediately
+    # ============================================================
+    attempts_left = max(0, attempts_left - 1)
+
+    # ============================================================
+    #  Attempts exhausted → FORCE LLM
+    # ============================================================
+    if attempts_left == 0:
+        comment = comment or "User marked result as wrong."
+
+    # ============================================================
+    # Insufficient feedback -> FOLLOW-UP (only if attempts remain)
+    # ============================================================
+    if attempts_left > 0 and not is_feedback_sufficient(comment):
+        return (
+            last_answer_md,
+            last_df,
+            last_sql,
+            last_rag_context,
+            last_question,
+            last_answer_md,
+            last_df,
+            True,   # awaiting follow-up
+            build_followup_prompt_for_user(comment),
+            attempts_left,
+        )
+
+    # ============================================================
+    # Run LLM review
+    # ============================================================
+    original_doc = find_cached_doc_by_sql(
+        last_question, last_sql, store=get_sql_cache_store()
+    )
+
+    corrected_sql, explanation = review_and_correct_sql_with_llm(
+        question=last_question,
+        generated_sql=last_sql,
+        user_feedback_comment=comment,
+        rag_context=last_rag_context,
+    )
+
+    corrected_sql = corrected_sql or last_sql
+    df_new, exec_error = execute_sql(corrected_sql, conn)
+
+    if exec_error:
+        answer_core = summarize_results(
+            question=last_question,
+            sql=corrected_sql,
+            df=None,
+            rag_context=last_rag_context,
+            error=exec_error,
+        )
+        df_new = pd.DataFrame()
+    else:
+        answer_core = summarize_results(
+            question=last_question,
+            sql=corrected_sql,
+            df=df_new,
+            rag_context=last_rag_context,
+            error=None,
+        )
+
+    update_cache_on_feedback(
+        question=last_question,
+        original_doc_md=original_doc,
+        user_marked_good=False,
+        llm_corrected_sql=(
+            corrected_sql if corrected_sql.strip() != last_sql.strip() else None
+        ),
+        llm_corrected_answer_md=(
+            answer_core if corrected_sql.strip() != last_sql.strip() else None
+        ),
+        store=get_sql_cache_store(),
+    )
+
+    record_feedback(
+        conn=conn,
+        question=last_question,
+        generated_sql=last_sql,
+        model_answer=last_answer_md,
+        rating="bad",
+        comment=comment + "\n\nLLM explanation:\n" + (explanation or ""),
+        corrected_sql=corrected_sql,
+    )
+
+    final_md = (
+        answer_core
+        + "\n\n---\n\n"
+        + f"**Final corrected SQL:**\n```sql\n{corrected_sql}\n```\n\n"
+        + "### LLM Review Explanation\n"
+        + (explanation or "")
+    )
+
+    return (
+        final_md,
+        df_new,
+        corrected_sql,
+        last_rag_context,
+        last_question,
+        final_md,
+        df_new,
+        False,
+        "",
+        attempts_left,
+    )
+
+
+# %%
+import gradio as gr
+import pandas as pd
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Olist Analytics Assistant (RAG + SQL + Feedback)")
+
+    # ==================== STATE ====================
+    last_sql_state = gr.State("")
+    last_rag_state = gr.State("")
+    last_question_state = gr.State("")
+    last_answer_state = gr.State("")
+    last_df_state = gr.State(pd.DataFrame())
+
+    attempts_state = gr.State(4)
+    feedback_sql_state = gr.State("") 
+
+    # ==================== MAIN UI ====================
+    with gr.Row():
+        with gr.Column(scale=1):
+            question_in = gr.Textbox(
+                label="Your question",
+                placeholder="e.g. Total number of customers",
+                lines=4,
+            )
+            submit_btn = gr.Button("Run")
+        with gr.Column(scale=2):
+            answer_out = gr.Markdown()
+            table_out = gr.Dataframe()
+
+    attempts_display = gr.Markdown("**Feedback attempts remaining: 4**")
+    cached_stats_md = gr.Markdown(visible=False)
+
+    # ==================== FEEDBACK ====================
+    gr.Markdown("### Feedback")
+
+    feedback_rating = gr.Radio(
+        ["Correct", "Wrong"],
+        label="Is the answer correct?",
+        value=None,
+    )
+    feedback_comment = gr.Textbox(
+        label="Explain (required if Wrong)",
+        lines=3,
+    )
+    feedback_btn = gr.Button("Submit feedback")
+
+    # ==================== FOLLOW-UP ====================
+    followup_prompt_md = gr.Markdown(visible=False)
+    followup_input = gr.Textbox(
+        label="Please clarify",
+        visible=False,
+        lines=4,
+    )
+    followup_submit_btn = gr.Button(
+        "Submit follow-up",
+        visible=False,
+    )
+
+    exhausted_md = gr.Markdown(
+        "**You have exhausted your feedback attempts. Please ask a new question to continue.**",
+        visible=False,
+    )
+
+    # ==================== UI HELPERS ====================
+    def reset_feedback_ui():
+        return (
+            gr.update(value=None, visible=True),   # rating
+            gr.update(value="", visible=True),     # comment
+            gr.update(visible=True),               # submit
+            gr.update(visible=False),              # followup input
+            gr.update(visible=False),              # followup btn
+            gr.update(visible=False),              # followup prompt
+            gr.update(visible=False),              # exhausted
+        )
+
+    def show_followup_ui(prompt: str):
+        return (
+            gr.update(visible=False),              # rating
+            gr.update(visible=False),              # comment
+            gr.update(visible=False),              # submit
+            gr.update(value="", visible=True),     # followup input
+            gr.update(visible=True),               # followup btn
+            gr.update(value=prompt, visible=True), # followup prompt
+            gr.update(visible=False),              # exhausted
+        )
+
+    def show_exhausted_ui():
+        return (
+            gr.update(visible=False),  # rating
+            gr.update(visible=False),  # comment
+            gr.update(visible=False),  # submit
+            gr.update(visible=False),  # followup input
+            gr.update(visible=False),  # followup btn
+            gr.update(visible=False),  # followup prompt
+            gr.update(visible=True),   # exhausted
+        )
+
+    # ==================== RUN PIPELINE ====================
+    def run_and_render(question):
+        (
+            answer_md,
+            df,
+            sql,
+            rag,
+            q,
+            answer_state,
+            df_state,
+            _cached_matches,
+            _dropdown_choices,
+            _dropdown_visible,
+            attempts,
+            attempts_text,
+            cached_stats_update,
+        ) = backend_pipeline(question)
+
+        return (
+            answer_md,
+            df,
+            sql,
+            rag,
+            q,
+            answer_state,
+            df_state,
+            attempts,
+            attempts_text,
+            cached_stats_update,
+            *reset_feedback_ui(),
+        )
+
+    submit_btn.click(
+        run_and_render,
+        inputs=[question_in],
+        outputs=[
+            answer_out,
+            table_out,
+            last_sql_state,
+            last_rag_state,
+            last_question_state,
+            last_answer_state,
+            last_df_state,
+            attempts_state,
+            attempts_display,
+            cached_stats_md,
+            feedback_rating,
+            feedback_comment,
+            feedback_btn,
+            followup_input,
+            followup_submit_btn,
+            followup_prompt_md,
+            exhausted_md,
+        ],
+    )
+
+    # ==================== FEEDBACK HANDLER ====================
+    def handle_feedback(
+        rating,
+        comment,
+        last_sql,
+        last_rag,
+        last_question,
+        last_answer,
+        last_df,
+        feedback_sql,
+        attempts_left,
+    ):
+        (
+            answer_md,
+            df_new,
+            sql_new,
+            rag_new,
+            q_new,
+            ans_state,
+            df_state,
+            awaiting_followup,
+            followup_prompt,
+            attempts_new,
+        ) = feedback_pipeline_interactive(
+            rating,
+            comment,
+            last_sql,
+            last_rag,
+            last_question,
+            last_answer,
+            last_df,
+            feedback_sql,
+            attempts_left,
+        )
+
+        attempts_md = f"**Feedback attempts remaining: {attempts_new}**"
+
+        # Exhausted
+        if attempts_new <= 0:
+            ui = show_exhausted_ui()
+            return (
+                answer_md,
+                df_new,
+                sql_new,
+                rag_new,
+                q_new,
+                ans_state,
+                df_state,
+                attempts_new,
+                attempts_md,
+                cached_stats_md,
+                *ui,
+            )
+
+        # Follow-up only
+        if awaiting_followup:
+            ui = show_followup_ui(followup_prompt)
+            return (
+                answer_md,
+                df_new,
+                sql_new,
+                rag_new,
+                q_new,
+                ans_state,
+                df_state,
+                attempts_new,
+                attempts_md,
+                cached_stats_md,
+                *ui,
+            )
+
+        # Normal reset
+        ui = reset_feedback_ui()
+        return (
+            answer_md,
+            df_new,
+            sql_new,
+            rag_new,
+            q_new,
+            ans_state,
+            df_state,
+            attempts_new,
+            attempts_md,
+            cached_stats_md,
+            *ui,
+        )
+
+    feedback_btn.click(
+        handle_feedback,
+        inputs=[
+            feedback_rating,
+            feedback_comment,
+            last_sql_state,
+            last_rag_state,
+            last_question_state,
+            last_answer_state,
+            last_df_state,
+            feedback_sql_state,
+            attempts_state,
+        ],
+        outputs=[
+            answer_out,
+            table_out,
+            last_sql_state,
+            last_rag_state,
+            last_question_state,
+            last_answer_state,
+            last_df_state,
+            attempts_state,
+            attempts_display,
+            cached_stats_md,
+            feedback_rating,
+            feedback_comment,
+            feedback_btn,
+            followup_input,
+            followup_submit_btn,
+            followup_prompt_md,
+            exhausted_md,
+        ],
+    )
+
+    followup_submit_btn.click(
+        handle_feedback,
+        inputs=[
+            feedback_rating,
+            followup_input,
+            last_sql_state,
+            last_rag_state,
+            last_question_state,
+            last_answer_state,
+            last_df_state,
+            feedback_sql_state,
+            attempts_state,
+        ],
+        outputs=[
+            answer_out,
+            table_out,
+            last_sql_state,
+            last_rag_state,
+            last_question_state,
+            last_answer_state,
+            last_df_state,
+            attempts_state,
+            attempts_display,
+            cached_stats_md,
+            feedback_rating,
+            feedback_comment,
+            feedback_btn,
+            followup_input,
+            followup_submit_btn,
+            followup_prompt_md,
+            exhausted_md,
+        ],
+    )
+
+# %%
+if __name__ == "__main__":
+    demo.launch()
+
+