rag_utils.py

"""
TableMind AI — Phase 6: RAG + tool layer.

Handles: parsing arbitrary uploaded documents, chunking (with header-repeat for tables),
embedding + FAISS retrieval, a numeric/aggregate query router that generates and safely
executes a single line of pandas against the parsed data, and lightweight chart generation.

All functions here are CPU-only by design so they never touch the ZeroGPU quota in app.py.
"""

import io
import os
import re
import tempfile
import textwrap
import uuid

import numpy as np
import pandas as pd

# ---------------------------------------------------------------------------
# Document parsing
# ---------------------------------------------------------------------------

def parse_document(file_path: str):
    """
    Returns a dict:
      { "dataframes": {name: pd.DataFrame, ...}, "text_chunks": [str, ...] }
    Supports .xlsx/.xls/.csv/.pdf/.docx/.txt/.md
    """
    ext = os.path.splitext(file_path)[1].lower()
    dataframes = {}
    text_chunks = []

    if ext in (".xlsx", ".xls"):
        sheets = pd.read_excel(file_path, sheet_name=None)
        for name, df in sheets.items():
            dataframes[name] = df.dropna(how="all")

    elif ext == ".csv":
        try:
            df = pd.read_csv(file_path)
        except Exception:
            df = pd.read_csv(file_path, sep=None, engine="python")
        dataframes["sheet1"] = df

    elif ext == ".pdf":
        import pdfplumber
        with pdfplumber.open(file_path) as pdf:
            for page_num, page in enumerate(pdf.pages):
                tables = page.extract_tables()
                for t_idx, table in enumerate(tables):
                    if not table or len(table) < 2:
                        continue
                    header, *rows = table
                    header = [str(h) if h else f"col_{i}" for i, h in enumerate(header)]
                    df = pd.DataFrame(rows, columns=header)
                    dataframes[f"page{page_num+1}_table{t_idx+1}"] = df
                page_text = page.extract_text() or ""
                if page_text.strip():
                    text_chunks.extend(_chunk_text(page_text, source=f"page {page_num+1}"))

    elif ext == ".docx":
        import docx
        d = docx.Document(file_path)
        for t_idx, table in enumerate(d.tables):
            rows = [[cell.text for cell in row.cells] for row in table.rows]
            if len(rows) < 2:
                continue
            header, *body = rows
            dataframes[f"table{t_idx+1}"] = pd.DataFrame(body, columns=header)
        full_text = "\n".join(p.text for p in d.paragraphs if p.text.strip())
        if full_text.strip():
            text_chunks.extend(_chunk_text(full_text, source="document body"))

    elif ext in (".txt", ".md"):
        with open(file_path, "r", encoding="utf-8", errors="ignore") as f:
            text_chunks.extend(_chunk_text(f.read(), source="text file"))

    else:
        raise ValueError(f"Unsupported file type: {ext}")

    # Turn every dataframe into header-repeated row-window chunks for retrieval too,
    # so lookup-style questions can be answered via RAG even when the router doesn't
    # trigger the numeric path.
    for name, df in dataframes.items():
        text_chunks.extend(_chunk_dataframe(df, name))

    return {"dataframes": dataframes, "text_chunks": text_chunks}


def _chunk_text(text: str, source: str, words_per_chunk: int = 300, overlap: int = 50):
    words = text.split()
    chunks = []
    i = 0
    while i < len(words):
        window = words[i:i + words_per_chunk]
        chunks.append(f"[Source: {source}]\n" + " ".join(window))
        i += words_per_chunk - overlap
    return chunks


def _chunk_dataframe(df: pd.DataFrame, name: str, rows_per_chunk: int = 20):
    """Chunk a dataframe into row windows, repeating the header/column names in every
    chunk so each chunk is independently understandable to the embedding model and LLM."""
    chunks = []
    header = " | ".join(str(c) for c in df.columns)
    for start in range(0, len(df), rows_per_chunk):
        window = df.iloc[start:start + rows_per_chunk]
        lines = [f"[Table: {name}]", f"Columns: {header}"]
        for _, row in window.iterrows():
            lines.append(" | ".join(str(v) for v in row.values))
        chunks.append("\n".join(lines))
    return chunks


# ---------------------------------------------------------------------------
# Embedding + retrieval
# ---------------------------------------------------------------------------

_EMBEDDER = None

def get_embedder():
    global _EMBEDDER
    if _EMBEDDER is None:
        from sentence_transformers import SentenceTransformer
        # device="cpu" is REQUIRED here, not optional: SentenceTransformer auto-detects
        # CUDA by default, but this function runs outside any @spaces.GPU-decorated
        # function (it's called from file-upload handling). On a ZeroGPU Space, any
        # CUDA touch outside the decorated function is blocked with a protective error -
        # this is also exactly what we want anyway, since embeddings should stay on CPU
        # and never consume ZeroGPU quota in the first place.
        _EMBEDDER = SentenceTransformer("BAAI/bge-small-en-v1.5", device="cpu")
    return _EMBEDDER


def build_index(text_chunks):
    """Build an in-memory FAISS index for one session's document. Returns (index, chunks)."""
    import faiss
    if not text_chunks:
        return None, []
    embedder = get_embedder()
    vectors = embedder.encode(text_chunks, normalize_embeddings=True, show_progress_bar=False)
    vectors = np.asarray(vectors, dtype="float32")
    index = faiss.IndexFlatIP(vectors.shape[1])
    index.add(vectors)
    return index, text_chunks


def retrieve(question: str, index, chunks, k: int = 6):
    if index is None or not chunks:
        return []
    embedder = get_embedder()
    q_vec = embedder.encode([question], normalize_embeddings=True)
    q_vec = np.asarray(q_vec, dtype="float32")
    scores, idxs = index.search(q_vec, min(k, len(chunks)))
    return [chunks[i] for i in idxs[0] if i != -1]


# ---------------------------------------------------------------------------
# Numeric / aggregate router
# ---------------------------------------------------------------------------

NUMERIC_KEYWORDS = re.compile(
    r"\b(sum|total|average|avg|mean|count|how many|maximum|max|minimum|min|top \d+|"
    r"highest|lowest|percent|percentage|compare|median|std|standard deviation|ratio)\b",
    re.IGNORECASE,
)


def is_numeric_question(question: str) -> bool:
    return bool(NUMERIC_KEYWORDS.search(question))


def describe_dataframes(dataframes: dict, max_sample_rows: int = 3) -> str:
    """Compact schema description fed to the LLM so it can write a pandas query
    without ever seeing the full table (keeps prompts cheap and scalable to any size)."""
    parts = []
    for name, df in dataframes.items():
        dtypes = ", ".join(f"{c} ({df[c].dtype})" for c in df.columns)
        sample = df.head(max_sample_rows).to_string(index=False)
        parts.append(f"DataFrame `{name}` — columns: {dtypes}\nSample rows:\n{sample}")
    return "\n\n".join(parts)


def run_pandas_query(code: str, dataframes: dict, timeout_seconds: int = 5):
    """
    Execute a single expression/line of pandas code in a restricted namespace.
    `code` must reference dataframes by the names given in describe_dataframes(),
    available in the namespace as `dfs["name"]`.

    NOTE: this restricted-exec sandbox is adequate for a personal portfolio demo with
    trusted/low-volume traffic. It is NOT a hardened multi-tenant sandbox — for a real
    production product, run this in an isolated subprocess/container with a real timeout.
    """
    safe_globals = {"__builtins__": {}}
    safe_locals = {"pd": pd, "np": np, "dfs": dataframes, "result": None}
    guarded_code = f"result = {code.strip()}"
    try:
        exec(guarded_code, safe_globals, safe_locals)
        return safe_locals["result"], None
    except Exception as e:
        return None, str(e)


# ---------------------------------------------------------------------------
# Chart generation
# ---------------------------------------------------------------------------

CHART_KEYWORDS = re.compile(
    r"\b(chart|graph|plot|trend|visuali[sz]e|distribution|over time)\b", re.IGNORECASE
)


def wants_chart(question: str) -> bool:
    return bool(CHART_KEYWORDS.search(question))


def make_chart(data, chart_type: str = "bar", title: str = "TableMind chart"):
    """
    data: a pandas Series (index=labels, values=numbers) or a small DataFrame with
    exactly two columns (label, value). Returns a path to a saved PNG (unique tempfile,
    safe for concurrent ZeroGPU requests).
    """
    import matplotlib
    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    if isinstance(data, pd.DataFrame):
        if data.shape[1] < 2:
            return None
        labels = data.iloc[:, 0].astype(str)
        values = data.iloc[:, 1]
    elif isinstance(data, pd.Series):
        labels = data.index.astype(str)
        values = data.values
    else:
        return None

    fig, ax = plt.subplots(figsize=(7, 4))
    if chart_type == "line":
        ax.plot(labels, values, marker="o")
    else:
        ax.bar(labels, values)
    ax.set_title(title)
    plt.xticks(rotation=45, ha="right")
    plt.tight_layout()

    out_path = os.path.join(tempfile.gettempdir(), f"tablemind_chart_{uuid.uuid4().hex}.png")
    fig.savefig(out_path, dpi=150)
    plt.close(fig)
    return out_path