ui/copilot.py

# Crewlyze
# Copyright (c) 2025 Sowmiyan S
# Licensed under the MIT License

"""
Interactive AI Data Copilot module.

Accepts a natural language query from the user, generates Python code via LLM,
executes it securely in an isolated subprocess, and returns text results and
new dynamically generated visualizations.

Key improvements:
- Column names, dtypes, and per-column statistics are injected into the prompt
  so the LLM generates accurate, runnable code without guessing column names.
- Supports /column slash command prefix for column-aware queries.
"""

import os
import re
import sys
import textwrap
import uuid
from pathlib import Path

import pandas as pd
from crewai import LLM
from config.llm_config import get_llm_params
from tools.dataset_tools import _run_in_subprocess, _strip_markdown_fences, read_csv_robust


# ---------------------------------------------------------------------------
# Column schema builder
# ---------------------------------------------------------------------------

def _build_column_context(csv_path: str, max_rows: int = 500) -> str:
    """
    Load the CSV and build a compact column schema string for injection into
    the LLM prompt.  Includes dtypes, missing%, and key statistics so the LLM
    can write correct, runnable pandas code without hallucinating column names.
    """
    try:
        df = read_csv_robust(csv_path, nrows=max_rows)
    except Exception as exc:
        return f"[Could not load dataset: {exc}]"

    lines = [
        f"Dataset: {max_rows if len(df) == max_rows else len(df)} rows × {len(df.columns)} columns",
        "",
        "Columns (name | dtype | missing% | stats):",
    ]
    for col in df.columns:
        dtype    = df[col].dtype
        miss_pct = round(df[col].isnull().sum() / max(len(df), 1) * 100, 1)
        if pd.api.types.is_numeric_dtype(dtype):
            stats = (
                f"min={df[col].min():.4g}, "
                f"mean={df[col].mean():.4g}, "
                f"max={df[col].max():.4g}, "
                f"std={df[col].std():.4g}"
            )
        else:
            top3 = df[col].dropna().value_counts().head(3).index.tolist()
            stats = "top: " + ", ".join(str(v) for v in top3) if top3 else "—"
        lines.append(f"  - {col!r}: {dtype} | missing={miss_pct}% | {stats}")

    lines.append("")
    lines.append("Sample rows (first 3):")
    for _, row in df.head(3).iterrows():
        lines.append("  " + str(dict(row)))

    return "\n".join(lines)


# ---------------------------------------------------------------------------
# Main copilot entry point
# ---------------------------------------------------------------------------

def run_copilot_query(query: str, csv_path: str, output_dir_str: str) -> dict:
    """
    Accepts a user query, generates Python code using the current LLM,
    runs the code in a sandbox subprocess, and returns {text, plot_path}.

    The column schema (names, dtypes, stats) is injected into the prompt to
    prevent NameError / KeyError in LLM-generated code.
    """
    # 1. Initialise LLM from current session env vars
    try:
        llm_params = get_llm_params()
        llm = LLM(**llm_params)
    except Exception as exc:
        return {
            "success": False,
            "text": f"LLM not configured: {exc}\nSet your API key in the sidebar.",
            "plot_path": None,
        }

    # 2. Build column context (prevents wrong-column NameErrors)
    column_context = _build_column_context(csv_path)

    # 3. Prepare plot output path
    output_dir = Path(output_dir_str)
    output_dir.mkdir(parents=True, exist_ok=True)
    plot_name = f"copilot_plot_{_uuid_short()}.png"
    plot_path = output_dir / plot_name

    # Clean up previous copilot plots
    for prev in output_dir.glob("copilot_plot_*.png"):
        try:
            prev.unlink(missing_ok=True)
        except OSError:
            pass

    # 4. Build LLM prompt with full column context
    prompt = textwrap.dedent(f"""
    You are an expert AI Data Analyst. You have access to a CSV dataset at:
      FILE_PATH = '{Path(csv_path).as_posix()}'

    === DATASET SCHEMA ===
    {column_context}
    =====================

    USER QUERY: "{query}"

    INSTRUCTIONS:
    1. Read the dataset: df = pd.read_csv(FILE_PATH)
    2. Use ONLY the column names listed in the dataset schema (exact spelling, case-sensitive).
    3. Perform any required analysis, aggregation, computation, or modifications.
    4. Print a clear, detailed, and nicely formatted answer to stdout detailing the results or actions taken.
       - Use rich Markdown formatting (e.g. Markdown tables, bulleted lists, bold text, headers) to structure the output like a professional report.
       - If the user asks for a table or for N values, print a Markdown table.
    5. If the query asks to modify, clean, fix, rename, delete columns, drop rows, replace missing values, or update values in the dataset:
       - Perform the operation on the DataFrame `df`.
       - Save the modified DataFrame back to the CSV file at the end of the script: `df.to_csv(FILE_PATH, index=False)`.
       - Print a confirmation message to stdout using Markdown (e.g., bulleted list) explaining exactly what dataset modifications were made.
    6. If the query asks for a chart/plot/graph:
       - You can use either Matplotlib/Seaborn OR Plotly.
       - If using Matplotlib/Seaborn: Call `import matplotlib; matplotlib.use('Agg')` BEFORE importing pyplot. Save with `plt.savefig('{plot_path.as_posix()}')`.
       - If using Plotly: Do NOT use `fig.write_image()`. Instead, you MUST export the figure using the Kaleido API directly:
         ```python
         import kaleido
         kaleido.write_fig_sync(fig, '{plot_path.as_posix()}')
         ```
       - Generate a professional chart. Apply any specific styles, colors, layouts, grids, or palettes requested by the user.

    Return ONLY valid Python code inside a ```python ... ``` block.
    Do NOT include explanations or text outside the code block.
    """).strip()

    try:
        # 5. Generate code
        response  = llm.call([{"role": "user", "content": prompt}])
        raw_code  = response if isinstance(response, str) else str(response)
        code      = _strip_markdown_fences(raw_code)

        if not code.strip():
            return {
                "success": False,
                "text": "The model returned empty code. Try rephrasing your query.",
                "plot_path": None,
            }

        # 6. Execute in sandboxed subprocess
        success, exec_output = _run_in_subprocess(code)

        plot_saved      = plot_path.exists() and plot_path.stat().st_size > 0
        final_plot_path = str(plot_path) if plot_saved else None

        if success:
            answer_text = exec_output.strip() if exec_output.strip() not in ("", "(no output)") \
                          else "✅ Query executed successfully (no text output)."
            return {"success": True, "text": answer_text, "plot_path": final_plot_path}
        else:
            return {
                "success": False,
                "text": f"⚠️ Execution error:\n```\n{exec_output}\n```",
                "plot_path": None,
            }

    except Exception as exc:
        return {
            "success": False,
            "text": f"Copilot error: {exc}",
            "plot_path": None,
        }


def _uuid_short() -> str:
    return uuid.uuid4().hex[:6]


# ---------------------------------------------------------------------------
# Column list helper (used by the /column slash picker in app.py)
# ---------------------------------------------------------------------------

def get_column_names(csv_path: str) -> list[str]:
    """Return column names from the CSV, or empty list on error."""
    try:
        return list(read_csv_robust(csv_path, nrows=0).columns)
    except Exception:
        return []