DuckDB + HNSW vectors
1956919 3 months ago
31.2 kB
r"""
NHTSA complaints  →  DuckDB  →  natural-language Q&A (Streamlit).
Flow:
  1. Pick make / model / year (cascading dropdowns) and click "Fetch & store".
  2. Complaints are fetched from the NHTSA API and stored in a DuckDB table.
  3. Ask a question in plain English. The LLM writes a DuckDB SQL query,
     we run it, and the LLM turns the rows into a natural-language list answer.
DuckDB is an embedded database (a single file) — no server needed. The docker
compose here just containerizes this app and keeps the .duckdb file on a volume.
Local run (interpreter that has streamlit + duckdb + langchain):
  d:\2026\July\Agentic-RAG-with-LangGraph-and-Ollama\.venv\Scripts\python.exe -m streamlit run duckdb_app.py
Docker run:
  docker compose up --build
  # then open http://localhost:8501
"""

import os
import time
import logging
from functools import reduce
from typing import Any, TypedDict

import duckdb
import requests
import pandas as pd
import streamlit as st
from dotenv import load_dotenv
from langchain_openai import ChatOpenAI
from langchain_huggingface import HuggingFaceEndpoint, ChatHuggingFace
from langgraph.graph import StateGraph, START, END

# Load config from .env (fail loudly if a required key is missing) --------
load_dotenv()


def env(key):
    val = os.getenv(key)
    if val is None or val == "":
        raise RuntimeError(f"Missing required config '{key}' in .env")
    return val


# Data source APIs
VPIC = env("VPIC_URL")
COMPLAINTS_URL = env("COMPLAINTS_URL")

# DuckDB
DUCKDB_PATH = env("DUCKDB_PATH")
TABLE = env("TABLE")
HNSW_INDEX = f"{TABLE}_embedding_hnsw"   # vector (ANN) index for semantic search

# Embedding backend: EMBED_PROVIDER selects hf (local) | google (Gemini API);
# EMBED_MODEL and EMBED_DIM must match the chosen provider (see .env).
EMBED_PROVIDER = env("EMBED_PROVIDER")
EMBED_MODEL = env("EMBED_MODEL")
EMBED_DIM = int(env("EMBED_DIM"))

# Chat LLMs
OPENAI_MODEL = env("OPENAI_MODEL")
HF_MODEL = env("HF_MODEL")
HF_PROVIDER = env("HF_PROVIDER")
GOOGLE_CHAT_MODEL = env("GOOGLE_CHAT_MODEL")

logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s | %(levelname)-7s | %(message)s",
    datefmt="%H:%M:%S",
)
log = logging.getLogger("nhtsa-duckdb")


# ---- DuckDB ----------------------------------------------------------------
@st.cache_resource
def get_db():
    """One DuckDB connection for the app; creates the file + table if needed."""
    os.makedirs(os.path.dirname(DUCKDB_PATH) or ".", exist_ok=True)
    log.info("DB: connecting to DuckDB at %s", DUCKDB_PATH)
    con = duckdb.connect(DUCKDB_PATH)

    # vss extension enables the HNSW vector index used for semantic search.
    con.execute("INSTALL vss; LOAD vss;")
    con.execute("SET hnsw_enable_experimental_persistence = true;")

    con.execute(f"""
        CREATE TABLE IF NOT EXISTS {TABLE} (
            odiNumber          BIGINT PRIMARY KEY,
            make               VARCHAR,
            model              VARCHAR,
            modelYear          INTEGER,
            manufacturer       VARCHAR,
            component          VARCHAR,
            summary            VARCHAR,
            vin                VARCHAR,
            crash              BOOLEAN,
            fire               BOOLEAN,
            numberOfInjuries   INTEGER,
            numberOfDeaths     INTEGER,
            dateComplaintFiled VARCHAR,
            embed_text         VARCHAR,
            embedding          FLOAT[{EMBED_DIM}]
        )
    """)
    # Migrate older DB files created before these columns existed.
    con.execute(f"ALTER TABLE {TABLE} ADD COLUMN IF NOT EXISTS vin VARCHAR")
    con.execute(f"ALTER TABLE {TABLE} ADD COLUMN IF NOT EXISTS embed_text VARCHAR")
    con.execute(f"ALTER TABLE {TABLE} ADD COLUMN IF NOT EXISTS embedding FLOAT[{EMBED_DIM}]")

    # If the embedding dimension changed (e.g. switching embedding models),
    # drop and recreate the column so the type matches. Existing embeddings are
    # cleared and must be re-computed (Fetch & store, or backfill script).
    row = con.execute(
        "SELECT data_type FROM information_schema.columns "
        "WHERE table_name = ? AND column_name = 'embedding'", [TABLE]).fetchone()
    want = f"FLOAT[{EMBED_DIM}]"
    if row and row[0] != want:
        log.info("DB: embedding dim changed %s → %s; recreating column", row[0], want)
        con.execute(f"DROP INDEX IF EXISTS {HNSW_INDEX}")   # index must go before column
        con.execute(f"ALTER TABLE {TABLE} DROP COLUMN embedding")
        con.execute(f"ALTER TABLE {TABLE} ADD COLUMN embedding {want}")

    # Create the HNSW (approximate nearest-neighbor) index for cosine similarity.
    # Inserts/updates maintain it automatically.
    con.execute(f"""
        CREATE INDEX IF NOT EXISTS {HNSW_INDEX}
        ON {TABLE} USING HNSW (embedding) WITH (metric = 'cosine')
    """)
    return con


@st.cache_resource
def get_embedder():
    """The embedding model (cached across reruns), per EMBED_PROVIDER."""
    log.info("EMBED: loading %s model %s", EMBED_PROVIDER, EMBED_MODEL)
    if EMBED_PROVIDER == "google":
        from langchain_google_genai import GoogleGenerativeAIEmbeddings
        return GoogleGenerativeAIEmbeddings(model=EMBED_MODEL,
                                            output_dimensionality=EMBED_DIM)
    from sentence_transformers import SentenceTransformer
    return SentenceTransformer(EMBED_MODEL)


def embed_texts(texts):
    """Return a list of EMBED_DIM-length float vectors (one per text)."""
    embedder = get_embedder()
    if EMBED_PROVIDER == "google":
        return embedder.embed_documents(list(texts))
    # local sentence-transformers
    vecs = embedder.encode(list(texts), normalize_embeddings=True)
    return [v.tolist() for v in vecs]


SCHEMA_DESC = f"""
Table "{TABLE}" columns:
  odiNumber (BIGINT)         - unique complaint id
  make (VARCHAR)             - vehicle make, e.g. 'Tesla'
  model (VARCHAR)            - vehicle model, e.g. 'Model 3'
  modelYear (INTEGER)        - model year, e.g. 2020
  manufacturer (VARCHAR)     - reporting manufacturer
  component (VARCHAR)        - affected component(s), comma-separated
  summary (VARCHAR)          - free-text complaint description
  vin (VARCHAR)              - vehicle VIN (often a partial VIN; casing varies)
  crash (BOOLEAN)            - did it involve a crash
  fire (BOOLEAN)             - did it involve a fire
  numberOfInjuries (INTEGER) - injuries reported
  numberOfDeaths (INTEGER)   - deaths reported
  dateComplaintFiled (VARCHAR) - date filed, MM/DD/YYYY
""".strip()


# ---- Cascading dropdown data ----------------------------------------------
@st.cache_data(show_spinner=False)
def get_makes():
    log.info("STEP 1: Fetching car makes")
    r = requests.get(f"{VPIC}/GetMakesForVehicleType/car?format=json", timeout=30)
    r.raise_for_status()
    return sorted(m["MakeName"].title() for m in r.json()["Results"])


@st.cache_data(show_spinner=False)
def get_models(make):
    log.info("STEP 2: Fetching models for make=%r", make)
    r = requests.get(f"{VPIC}/GetModelsForMake/{make}?format=json", timeout=30)
    r.raise_for_status()
    return sorted(m["Model_Name"] for m in r.json()["Results"])


# ---- Fetch complaints + store into DuckDB ---------------------------------
@st.cache_data(show_spinner=False)
def fetch_complaints(make, model, model_year):
    params = {"make": make, "model": model, "modelYear": model_year}
    log.info("STEP 3: Fetching complaints  params=%s", params)
    t0 = time.perf_counter()
    r = requests.get(COMPLAINTS_URL, params=params, timeout=30)
    r.raise_for_status()
    results = r.json().get("results", [])
    log.info("        -> %d complaints in %.2fs", len(results), time.perf_counter() - t0)
    return results


def to_record(c, make, model, model_year):
    products = c.get("products") or []
    component = ", ".join(
        p.get("productType", "") for p in products if isinstance(p, dict)
    ) or c.get("components", "")
    rec = {
        "odiNumber": c.get("odiNumber"),
        "make": make,
        "model": model,
        "modelYear": int(model_year),
        "manufacturer": c.get("manufacturer", ""),
        "component": component,
        "summary": c.get("summary", ""),
        "vin": (c.get("vin") or ""),
        "crash": bool(c.get("crash")),
        "fire": bool(c.get("fire")),
        "numberOfInjuries": c.get("numberOfInjuries", 0) or 0,
        "numberOfDeaths": c.get("numberOfDeaths", 0) or 0,
        "dateComplaintFiled": c.get("dateComplaintFiled", ""),
    }
    # All details combined into one text column, which we then embed.
    rec["embed_text"] = (
        f"{rec['make']} {rec['model']} {rec['modelYear']} | "
        f"manufacturer: {rec['manufacturer']} | component: {rec['component']} | "
        f"vin: {rec['vin']} | crash: {rec['crash']} | fire: {rec['fire']} | "
        f"injuries: {rec['numberOfInjuries']} | deaths: {rec['numberOfDeaths']} | "
        f"filed: {rec['dateComplaintFiled']} | summary: {rec['summary']}"
    )
    return rec


def store_dataset(con, make, model, model_year):
    raw = fetch_complaints(make, model, model_year)
    records = [to_record(c, make, model, model_year) for c in raw]
    if not records:
        return 0
    df = pd.DataFrame(records)

    # Embed the combined text of every record into the `embedding` column.
    log.info("STEP 4a: Embedding %d records with %s", len(df), EMBED_MODEL)
    df["embedding"] = embed_texts(df["embed_text"].tolist())

    log.info("STEP 4b: Upserting %d rows into DuckDB.%s", len(df), TABLE)
    con.register("incoming", df)
    # INSERT OR REPLACE upserts by primary key (odiNumber): new rows are added and
    # existing rows are refreshed (this backfills vin/embed_text/embedding). BY NAME
    # makes it robust to column order / added columns.
    con.execute(f"""
        INSERT OR REPLACE INTO {TABLE} BY NAME
        SELECT * FROM incoming
    """)
    con.unregister("incoming")
    return len(df)


def semantic_search(con, query, k=10):
    """Embed the query and return the k most similar complaints by cosine distance."""
    log.info("SEMANTIC: search %r (k=%d)", query, k)
    qvec = embed_texts([query])[0]
    # Pure `ORDER BY array_cosine_distance(...) LIMIT k` lets DuckDB use the HNSW
    # index (an ANN scan) instead of a full brute-force scan. Similarity is
    # derived in pandas afterwards so it doesn't disrupt the index optimization.
    df = con.execute(
        f"""
        SELECT odiNumber, make, model, modelYear, component, vin,
               numberOfInjuries, numberOfDeaths, summary,
               array_cosine_distance(embedding, ?::FLOAT[{EMBED_DIM}]) AS dist
        FROM {TABLE}
        ORDER BY dist
        LIMIT {int(k)}
        """,
        [qvec],
    ).fetchdf()
    df = df[df["dist"].notna()].copy()          # drop rows without an embedding
    df.insert(0, "similarity", (1 - df["dist"]).round(3))
    return df.drop(columns=["dist"])


# ---- LLM: text -> SQL -> natural language ---------------------------------
def get_llm(provider):
    if provider == "OpenAI":
        # stream_usage=True so token counts arrive in the final streamed chunk.
        return ChatOpenAI(model=OPENAI_MODEL, temperature=0, stream_usage=True)
    if provider == "Google":
        from langchain_google_genai import ChatGoogleGenerativeAI
        return ChatGoogleGenerativeAI(model=GOOGLE_CHAT_MODEL, temperature=0)
    endpoint = HuggingFaceEndpoint(
        repo_id=HF_MODEL,
        task="conversational",
        provider=HF_PROVIDER,
        huggingfacehub_api_token=os.getenv("HUGGINGFACEHUB_API_TOKEN"),
    )
    return ChatHuggingFace(llm=endpoint)


def timed_invoke(llm, prompt, label):
    """Call the LLM, logging wall-clock time and token usage for this call."""
    t0 = time.perf_counter()
    resp = llm.invoke(prompt)
    dt = time.perf_counter() - t0

    # LangChain normalizes usage into usage_metadata; fall back to raw metadata.
    usage = getattr(resp, "usage_metadata", None) or {}
    in_tok = usage.get("input_tokens")
    out_tok = usage.get("output_tokens")
    tot_tok = usage.get("total_tokens")
    if in_tok is None:
        md = getattr(resp, "response_metadata", {}) or {}
        tu = md.get("token_usage") or md.get("usage") or {}
        in_tok = tu.get("prompt_tokens")
        out_tok = tu.get("completion_tokens")
        tot_tok = tu.get("total_tokens")

    log.info("LLM[%s]: %.2fs | tokens in=%s out=%s total=%s",
             label, dt, in_tok, out_tok, tot_tok)
    meta = {"label": label, "seconds": dt,
            "in": in_tok, "out": out_tok, "total": tot_tok}
    return resp, meta


def clean_sql(text):
    """Strip markdown fences / stray prose, keep the SQL statement."""
    t = text.strip()
    if "```" in t:
        # take the content of the first fenced block
        t = t.split("```")[1]
        if t.lower().startswith("sql"):
            t = t[3:]
    return t.strip().rstrip(";").strip()


def history_text(history, limit=6):
    """Render the last few turns so follow-up questions have context."""
    if not history:
        return ""
    recent = history[-limit:]
    lines = ["Conversation so far (for resolving follow-up questions):"]
    for m in recent:
        who = "User" if m["role"] == "user" else "Assistant"
        lines.append(f"{who}: {m['content']}")
    return "\n".join(lines) + "\n\n"


def plan_answer(llm, question, history):
    """Router: decide whether to answer with SQL or with semantic vector search.
    Returns (mode, payload, meta) where mode is "sql" or "semantic".
    """
    prompt = (
        "You are the planner for a vehicle-complaints assistant backed by DuckDB. "
        "Decide how to answer the latest question and output exactly ONE line:\n\n"
        "1. If it needs exact filtering, counting, grouping, sorting, dates, VIN "
        "lookup, or precise conditions → output a DuckDB SQL SELECT query.\n"
        "2. If it asks to FIND complaints similar in meaning/topic/theme "
        "(e.g. 'complaints about sudden acceleration', 'anything like phantom "
        "braking') → output: SEMANTIC: <short search phrase>\n\n"
        "SQL rules: SELECT only. Quote all string values in single quotes. Use "
        "ILIKE '%...%' for text/VIN partial matches (VIN casing varies). Limit "
        "free-text result sets to 50 rows.\n"
        "Resolve follow-ups (e.g. 'those', 'what about fires') using the "
        "conversation. Output ONLY the SQL or the SEMANTIC line — no markdown, no "
        "explanation.\n\n"
        f"{SCHEMA_DESC}\n\n{history_text(history)}Latest question: {question}\nPlan:"
    )
    resp, meta = timed_invoke(llm, prompt, "router")
    out = resp.content.strip()
    log.info("STEP 5: ROUTER raw output → %r", out)

    # Strip any code fences first.
    if out.startswith("```"):
        out = clean_sql(out)

    if out.upper().lstrip().startswith("SEMANTIC"):
        phrase = out.split(":", 1)[1].strip() if ":" in out else question
        log.info("STEP 5: ROUTER decision → mode=semantic | phrase=%r", phrase)
        return ("semantic", phrase or question, meta)

    sql = clean_sql(out)
    log.info("STEP 5: ROUTER decision → mode=sql | query=%s", sql)
    if not sql.lower().lstrip().startswith("select"):
        raise ValueError(f"Refusing to run non-SELECT query: {sql!r}")
    return ("sql", sql, meta)


HEAVY_COLS = ["embedding", "embed_text"]   # never send these to the LLM / UI


def drop_heavy(df):
    """Remove the big vector/text columns (embedding, embed_text) that would
    otherwise bloat the LLM prompt (SELECT * pulls a 384/768-float array per row)
    and clutter the displayed table."""
    cols = [c for c in HEAVY_COLS if c in df.columns]
    return df.drop(columns=cols) if cols else df


def is_tabular(result_df):
    """A table makes sense for multi-row or multi-column results, but not for a
    single scalar like COUNT(*)."""
    return not (result_df.shape[0] <= 1 and result_df.shape[1] <= 1)


def build_answer_prompt(question, query_used, result_df, history):
    # Keep the context bounded — and never include the huge embedding/embed_text
    # columns (a SELECT * would otherwise blow past the model's token limit).
    preview = drop_heavy(result_df).head(50).to_markdown(index=False)

    if is_tabular(result_df):
        # The table itself shows the detail, so only ask for a short summary line.
        instruction = (
            "Reply with ONE short sentence summarizing the result (e.g. how many "
            "rows matched). Do NOT list the rows — a table will be shown below."
        )
    else:
        instruction = (
            "Answer in ONE short natural-language sentence. If the result is "
            "empty, say no matching complaints were found."
        )

    return (
        "You are a helpful data analyst chatbot. Using ONLY the query result below, "
        f"answer the latest question. {instruction}\n\n"
        f"{history_text(history)}Latest question: {question}\nQuery used: {query_used}\n\n"
        f"Result:\n{preview}\n\nAnswer:"
    )


def _usage_from(msg):
    usage = getattr(msg, "usage_metadata", None) or {}
    if usage:
        return usage.get("input_tokens"), usage.get("output_tokens"), usage.get("total_tokens")
    md = getattr(msg, "response_metadata", {}) or {}
    tu = md.get("token_usage") or md.get("usage") or {}
    return tu.get("prompt_tokens"), tu.get("completion_tokens"), tu.get("total_tokens")


def stream_answer(llm, provider, question, query_used, result_df, history, holder):
    """Generator yielding answer text chunks (for st.write_stream).
    Stores timing/token metrics in holder['meta'] when done.
    OpenAI streams natively. HuggingFace's streaming path is buggy with some
    providers (novita sends a final chunk with choices=null), so we invoke once
    and simulate word-by-word streaming for the same typing effect.
    """
    prompt = build_answer_prompt(question, query_used, result_df, history)
    t0 = time.perf_counter()
    in_tok = out_tok = tot_tok = None

    if provider == "OpenAI":
        chunks = []
        for chunk in llm.stream(prompt):
            chunks.append(chunk)
            if chunk.content:
                yield chunk.content
        if chunks:
            in_tok, out_tok, tot_tok = _usage_from(reduce(lambda a, b: a + b, chunks))
    else:
        # HuggingFace (and any non-OpenAI): single call, then fake streaming.
        resp = llm.invoke(prompt)
        in_tok, out_tok, tot_tok = _usage_from(resp)
        for word in (resp.content or "").split(" "):
            yield word + " "

    dt = time.perf_counter() - t0
    log.info("LLM[answer]: %.2fs | tokens in=%s out=%s total=%s (%s)",
             dt, in_tok, out_tok, tot_tok,
             "streamed" if provider == "OpenAI" else "invoke")
    holder["meta"] = {"label": "answer", "seconds": dt,
                      "in": in_tok, "out": out_tok, "total": tot_tok}


# ---- LangGraph pipeline: router -> (sql | semantic) -> answer --------------
class QAState(TypedDict, total=False):
    question: str
    history: list
    llm: Any
    con: Any
    mode: str          # "sql" | "semantic"
    payload: str       # the SQL string or the semantic phrase
    query_used: str
    result_df: Any
    answer: str
    metrics: list      # per-LLM-call timing/token dicts


def router_node(state: QAState):
    """LLM decides the route (and produces the SQL or semantic phrase)."""
    mode, payload, meta = plan_answer(state["llm"], state["question"], state["history"])
    return {"mode": mode, "payload": payload, "metrics": [meta]}


def sql_node(state: QAState):
    df = state["con"].execute(state["payload"]).fetchdf()
    log.info("        -> %d rows returned (mode=sql)", len(df))
    return {"result_df": df, "query_used": state["payload"]}


def semantic_node(state: QAState):
    df = semantic_search(state["con"], state["payload"], k=10)
    log.info("        -> %d rows returned (mode=semantic)", len(df))
    return {"result_df": df, "query_used": f"semantic search for: {state['payload']!r}"}


@st.cache_resource
def get_qa_graph():
    """Compile the routing+retrieval graph once (cached across reruns).
    The final NL answer is streamed at the UI layer, so the graph ends after
    the sql/semantic retrieval nodes.
    """
    log.info("GRAPH: compiling LangGraph QA pipeline")
    g = StateGraph(QAState)
    g.add_node("router", router_node)
    g.add_node("sql", sql_node)
    g.add_node("semantic", semantic_node)

    g.add_edge(START, "router")
    # Conditional edge: branch on the mode the router chose.
    g.add_conditional_edges("router", lambda s: s["mode"],
                            {"sql": "sql", "semantic": "semantic"})
    g.add_edge("sql", END)
    g.add_edge("semantic", END)
    return g.compile()


def format_metrics(metrics):
    """Turn per-call metrics into a compact caption like
    'router 2.6s · answer 1.8s · 248 tokens'."""
    if not metrics:
        return ""
    parts = [f"{m['label']} {m['seconds']:.1f}s" for m in metrics]
    total_tokens = sum(m["total"] or 0 for m in metrics)
    return " · ".join(parts) + f" · {total_tokens} tokens"


def retrieve(provider, question, con, history):
    """Run the graph (router + sql/semantic retrieval). The answer is streamed
    separately by the UI. Returns the final graph state (dict)."""
    return get_qa_graph().invoke({
        "question": question,
        "history": history,
        "llm": get_llm(provider),
        "con": con,
    })


# ---- UI --------------------------------------------------------------------
st.set_page_config(page_title="NHTSA · DuckDB Q&A", page_icon=":duck:")
st.header("NHTSA Complaints · DuckDB + NLP Q&A")

con = get_db()

with st.sidebar:
    st.subheader("Add data to DuckDB")
    make = st.selectbox("Make", get_makes(), index=None, placeholder="Select a make…")
    if make:
        model = st.selectbox("Model", get_models(make), index=None,
                             placeholder="Select a model…")
    else:
        model = st.selectbox("Model", [], index=None,
                             placeholder="Pick a make first…", disabled=True)
    model_year = st.selectbox("Model year", list(range(2026, 2014, -1)))
    fetch = st.button("Fetch & store", disabled=not (make and model))

    # One-click sample datasets (verified to return complaints).
    st.caption("⚡ Or load quick test data:")
    QUICK_PRESETS = [
        ("Tesla", "Model 3", 2020),
        ("Ford", "Mustang", 2020),
        ("Suzuki", "SX4", 2010),
        ("Toyota", "Camry", 2019),
        ("Honda", "Civic", 2019),
    ]
    for mk, md, yr in QUICK_PRESETS:
        if st.button(f"{mk} · {md} · {yr}", key=f"preset_{mk}_{md}_{yr}",
                     width='stretch'):
            st.session_state["fetch_target"] = (mk, md, yr)
            st.rerun()

    st.divider()
    st.subheader("Ask (LLM) settings")
    # Primary: HuggingFace. Secondary fallbacks: OpenAI, Google.
    provider = st.selectbox("Provider", ["HuggingFace", "OpenAI", "Google"])

# ---- Fetch & store action --------------------------------------------------
# Target comes from the manual "Fetch & store" button or a quick-test preset.
target = None
if fetch and make and model:
    target = (make, model, model_year)
else:
    target = st.session_state.pop("fetch_target", None)

if target:
    t_make, t_model, t_year = target
    log.info("STEP 0: Fetch → %s %s %s", t_make, t_model, t_year)
    try:
        with st.spinner(f"Fetching + storing {t_make} {t_model} {t_year}…"):
            n = store_dataset(con, t_make, t_model, t_year)
        if n == 0:
            st.info(f"No complaints found for {t_make} {t_model} {t_year}.")
        else:
            st.success(f"Stored {n} complaints for {t_make} {t_model} {t_year} in DuckDB.")
    except Exception as e:
        log.exception("ERROR while storing")
        st.error(f"{type(e).__name__}: {e}")

# ---- What's in the DB ------------------------------------------------------
total = con.execute(f"SELECT COUNT(*) FROM {TABLE}").fetchone()[0]
embedded = con.execute(
    f"SELECT COUNT(*) FROM {TABLE} WHERE embedding IS NOT NULL").fetchone()[0]
has_index = con.execute(
    "SELECT COUNT(*) FROM duckdb_indexes WHERE table_name = ? AND index_name = ?",
    [TABLE, HNSW_INDEX]).fetchone()[0] > 0
coverage = (embedded / total * 100) if total else 0
index_badge = "HNSW ✓" if has_index else "HNSW ✗"

st.caption(
    f"DuckDB `{DUCKDB_PATH}` · table `{TABLE}` · **{total}** rows · "
    f"embedded **{embedded}** ({coverage:.0f}%) · "
    f"vector index **{index_badge}** ({EMBED_PROVIDER}, {EMBED_DIM}d)"
)
if total and embedded < total:
    st.caption(f"⚠️ {total - embedded} rows have no embedding — run the backfill "
               f"to make them searchable.")

if total:
    inv = con.execute(f"""
        SELECT make, model, modelYear, COUNT(*) AS complaints
        FROM {TABLE} GROUP BY 1,2,3 ORDER BY complaints DESC
    """).fetchdf()
    with st.expander("Datasets currently in DuckDB"):
        st.dataframe(inv, width='stretch')

# ---- Sample-question templates (shown on the right) ------------------------
TEMPLATES = {
    "🧠 Semantic (meaning)": [
        "Find complaints about the car suddenly accelerating on its own",
        "Anything similar to phantom braking on the highway?",
        "Complaints describing the battery catching fire",
        "Show me issues that feel like the steering locking up",
    ],
    "Counts & filters": [
        "How many complaints involved a fire?",
        "How many complaints reported at least one injury?",
        "How many complaints had zero crashes, fires, injuries, and deaths?",
        "What percentage of all complaints involved a crash?",
    ],
    "Group-by / rankings": [
        "What are the top 10 most common components?",
        "Which component has the most crash-related complaints?",
        "How many complaints are there per model year?",
        "Break down the complaint count by make and model.",
    ],
    "Text search": [
        'How many complaints mention "brake"?',
        'List complaints that mention "autopilot" or "self driving".',
        'Show me complaints where the summary mentions "battery".',
    ],
    "Aggregates": [
        "What is the total number of injuries across all complaints?",
        "Which single complaint reported the most injuries? Show its summary.",
    ],
    "Dates": [
        "How many complaints were filed in 2020?",
        "What is the most recent complaint filed, and what component was it about?",
    ],
    "Trickier": [
        "Are there any complaints involving both a fire and an injury?",
        "List the odiNumbers of the 5 complaints with the longest summaries.",
    ],
}

# ---- Conversational chatbot ------------------------------------------------
st.subheader("Chat with your data")

# Persistent chat history
if "messages" not in st.session_state:
    st.session_state["messages"] = []

chat_col, tmpl_col = st.columns([3, 1])

# ---- Right side: clickable templates --------------------------------------
with tmpl_col:
    st.markdown("**📋 Sample questions**")
    st.caption("Click one to run it")
    for category, questions in TEMPLATES.items():
        with st.expander(category):
            for i, q in enumerate(questions):
                # A clicked template becomes the pending question, then rerun.
                if st.button(q, key=f"tmpl_{category}_{i}", width='stretch'):
                    st.session_state["pending"] = q
                    st.rerun()

# ---- Left side: the conversation ------------------------------------------
with chat_col:
    if st.session_state["messages"]:
        if st.button("🗑️ Clear chat"):
            st.session_state["messages"] = []
            st.rerun()

    # Replay the whole conversation on every rerun
    for m in st.session_state["messages"]:
        with st.chat_message(m["role"]):
            st.markdown(m["content"])
            if m.get("table") is not None:
                st.dataframe(m["table"], width='stretch', hide_index=True)

    # A pending question comes from either the chat box or a template click.
    prompt = st.session_state.pop("pending", None)

    if prompt:
        st.session_state["messages"].append({"role": "user", "content": prompt})
        with st.chat_message("user"):
            st.markdown(prompt)

        with st.chat_message("assistant"):
            if total == 0:
                msg = "The database is empty — fetch & store a dataset first (sidebar)."
                st.warning(msg)
                st.session_state["messages"].append({"role": "assistant", "content": msg})
            else:
                try:
                    history = st.session_state["messages"][:-1]

                    # Step 1: router + retrieval (LangGraph). Show a spinner here.
                    with st.spinner(f"Searching for an answer to: {prompt}"):
                        state = retrieve(provider, prompt, con, history)
                    mode = state["mode"]
                    result_df = state["result_df"]
                    metrics = list(state.get("metrics", []))

                    badge = ("🧠 semantic search" if mode == "semantic"
                             else "🗄️ SQL query")
                    st.caption(badge)

                    # Step 2: stream the natural-language answer token by token.
                    holder = {}
                    answer = st.write_stream(stream_answer(
                        state["llm"], provider, prompt, state["query_used"],
                        result_df, history, holder))
                    if holder.get("meta"):
                        metrics.append(holder["meta"])

                    # Metrics caption (timing/tokens) after streaming completes.
                    stats = format_metrics(metrics)
                    if stats:
                        st.caption(stats)

                    # Show a table for multi-row/column results (without the
                    # heavy embedding/embed_text columns).
                    table = drop_heavy(result_df) if is_tabular(result_df) else None
                    if table is not None:
                        st.dataframe(table, width='stretch', hide_index=True)

                    caption = f"{badge}  ·  {stats}" if stats else badge
                    st.session_state["messages"].append({
                        "role": "assistant",
                        "content": f"{caption}\n\n{answer}",
                        "table": table,
                    })
                except Exception as e:
                    log.exception("ERROR during answer")
                    err = f"{type(e).__name__}: {e}"
                    st.error(err)
                    st.session_state["messages"].append(
                        {"role": "assistant", "content": err})

# ---- Chat input (must be top-level, not inside a column) -------------------
typed = st.chat_input("Ask about the complaints… (e.g. how many involved a crash?)")
if typed:
    st.session_state["pending"] = typed
    st.rerun()