agent_graph.py

"""
LangGraph analytical agent for SheetOracle.

Pipeline:
  1. Resolve chart intent (keyword match, no LLM)
  2. Generate and execute SQL against the session SQLite DB
  3. Optionally render a matplotlib chart
  4. Summarize results in natural language
"""

import ast
import os
import re
import sqlite3
import time
from typing import Any, Dict, List, Literal, Tuple, TypedDict, Union

import matplotlib
import matplotlib.pyplot as plt
from langchain_core.messages import HumanMessage, SystemMessage
from langchain_openai import ChatOpenAI
from langgraph.graph import END, START, StateGraph

matplotlib.use("Agg")

# ---------------------------------------------------------------------------
# State
# ---------------------------------------------------------------------------


class AgentState(TypedDict):
    """Shared state passed between graph nodes."""

    user_question: str
    db_path: str
    generated_sql: str
    query_result: Union[str, List[tuple]]
    error_message: str
    final_response: str
    retry_count: int
    chart_requested: bool
    chart_intent_resolved: bool
    chart_file_path: str
    chart_error: str
    chart_warning: str


# ---------------------------------------------------------------------------
# Configuration
# ---------------------------------------------------------------------------

llm = ChatOpenAI(model="gpt-4o-mini", temperature=0)

MAX_CHART_ROWS = 50
MAX_PIE_SLICES = 10
MAX_PROMPT_ROWS = 20
MAX_SQL_RETRIES = 3
MAX_RESULT_ROWS = 10_000
MAX_QUERY_SECONDS = 20
PROGRESS_HANDLER_INSTRUCTIONS = 10_000
CHART_OUTPUT_DIR = "./tmp_charts"

CHART_KEYWORDS = (
    "plot",
    "chart",
    "graph",
    "visualize",
    "visualise",
    "visual",
    "bar chart",
    "line chart",
    "pie chart",
    "scatter plot",
    "scatter",
    "histogram",
    "trend",
    "over time",
    "timeline",
    "draw",
)

GRAPH_NODE_PROGRESS = {
    "generate_sql": "Querying data...",
    "execute_sql": "Running query...",
    "generate_chart": "Building chart...",
    "formulate_response": "Summarizing results...",
}


def get_agent_task_order(chart_requested: bool) -> List[str]:
    """Return the thinking-step task order for a given analysis run."""
    order = ["generate_sql", "execute_sql"]
    if chart_requested:
        order.append("generate_chart")
    order.append("formulate_response")
    return order

CHART_COLORS = ("#36C5F0", "#ECB22E", "#2EB67D", "#E01E5A")


# ---------------------------------------------------------------------------
# Public helpers (used by app.py and tests)
# ---------------------------------------------------------------------------


def is_chart_requested(question: str) -> bool:
    """Return True when the question asks for a visual/chart output."""
    normalized = question.lower()
    return any(keyword in normalized for keyword in CHART_KEYWORDS)


def build_initial_state(user_question: str, db_path: str) -> Dict[str, Any]:
    """Create the initial graph state with chart routing already resolved."""
    return {
        "user_question": user_question,
        "db_path": db_path,
        "generated_sql": "",
        "query_result": "",
        "error_message": "",
        "final_response": "",
        "retry_count": 0,
        "chart_requested": is_chart_requested(user_question),
        "chart_intent_resolved": True,
        "chart_file_path": "",
        "chart_error": "",
        "chart_warning": "",
    }


# ---------------------------------------------------------------------------
# Query-result helpers
# ---------------------------------------------------------------------------


def _parse_query_result(query_result: Union[str, List[tuple], Any]) -> List[tuple]:
    """Normalize SQL output into a list of row tuples."""
    if isinstance(query_result, list):
        return query_result
    if isinstance(query_result, str):
        if not query_result:
            return []
        return ast.literal_eval(query_result)
    raise ValueError(f"Unexpected query_result type: {type(query_result)}")


def _is_numeric(value: Any) -> bool:
    try:
        float(value)
        return True
    except (TypeError, ValueError):
        return False


def _coerce_numeric(value: Any, label: str) -> float:
    if value is None:
        raise ValueError(f"{label} contains null values that cannot be plotted.")
    try:
        return float(value)
    except (TypeError, ValueError) as exc:
        raise ValueError(f"{label} must be numeric; got {value!r}.") from exc


def _summarize_query_result_for_prompt(query_result: Union[str, List[tuple], Any]) -> str:
    """Build a compact dataset summary for the response LLM prompt."""
    try:
        rows = _parse_query_result(query_result)
    except (ValueError, SyntaxError):
        text = str(query_result)
        return text[:2000] + ("..." if len(text) > 2000 else "")

    if not rows:
        return "No rows returned."

    row_count = len(rows)
    if row_count == 1:
        return f"Query returned 1 answer row: {rows[0]}"

    if row_count <= MAX_PROMPT_ROWS:
        return f"Query returned {row_count} rows: {rows}"

    total = len(rows)
    col_count = len(rows[0]) if isinstance(rows[0], (tuple, list)) else 1
    lines = [
        f"Total rows: {total} (truncated for summarization).",
        f"First rows: {rows[:5]}",
    ]
    if total > 8:
        lines.append(f"Last rows: {rows[-3:]}")

    value_index = 1 if col_count >= 2 else 0
    numeric_rows = [
        row for row in rows
        if len(row) > value_index and _is_numeric(row[value_index])
    ]
    if numeric_rows:
        values = [float(row[value_index]) for row in numeric_rows]
        lines.append(
            f"Numeric stats (column {value_index + 1}): "
            f"min={min(values):.4g}, max={max(values):.4g}, "
            f"avg={sum(values) / len(values):.4g}, sum={sum(values):.4g}"
        )
        top_rows = sorted(
            numeric_rows,
            key=lambda row: float(row[value_index]),
            reverse=True,
        )[:5]
        lines.append(f"Top 5 rows by value: {top_rows}")

    return "\n".join(lines)


def _connect_readonly(db_path: str) -> sqlite3.Connection:
    """Open the session database in read-only mode.

    Prevent SQL injection by activate the read-only mode at first when modifying.
    """
    return sqlite3.connect(f"file:{os.path.abspath(db_path)}?mode=ro", uri=True)


def _load_table_schema(db_path: str) -> str:
    conn = _connect_readonly(db_path)
    cursor = conn.cursor()
    cursor.execute("PRAGMA table_info(uploaded_data)")
    columns = cursor.fetchall()
    conn.close()
    return "\n".join(f"- {name}: {col_type}" for _, name, col_type, *_ in columns)


# ---------------------------------------------------------------------------
# SQL safety
# ---------------------------------------------------------------------------

FORBIDDEN_SQL_KEYWORDS = (
    "attach", "detach", "pragma", "vacuum", "insert", "update", "delete",
    "drop", "alter", "create", "reindex", "begin",
    "commit", "rollback", "savepoint", "analyze", "load_extension",
)


def _strip_sql_literals_and_comments(sql: str) -> str:
    """Blank out string/identifier literals and comments for safe keyword scanning."""
    out: List[str] = []
    index = 0
    length = len(sql)
    while index < length:
        char = sql[index]
        if char in ("'", '"', "`"):
            quote = char
            index += 1
            while index < length:
                if sql[index] == quote:
                    # Doubled quote is an escaped quote inside the literal.
                    if index + 1 < length and sql[index + 1] == quote:
                        index += 2
                        continue
                    index += 1
                    break
                index += 1
            out.append(" ")
            continue
        if char == "[":
            while index < length and sql[index] != "]":
                index += 1
            index += 1
            out.append(" ")
            continue
        if sql.startswith("--", index):
            while index < length and sql[index] != "\n":
                index += 1
            continue
        if sql.startswith("/*", index):
            end = sql.find("*/", index + 2)
            index = length if end == -1 else end + 2
            out.append(" ")
            continue
        out.append(char)
        index += 1
    return "".join(out)


def _strip_sql_fences(sql: str) -> str:
    """Remove markdown code fences the LLM sometimes emits despite instructions."""
    cleaned = sql.strip()
    if cleaned.startswith("```"):
        cleaned = re.sub(r"^```[a-zA-Z]*\s*", "", cleaned)
        if cleaned.endswith("```"):
            cleaned = cleaned[:-3]
    return cleaned.strip()


def _validate_sql(sql: str) -> str:
    """Return a single read-only SELECT statement, or raise ValueError."""
    cleaned = _strip_sql_fences(sql)
    if not cleaned:
        raise ValueError("No SQL query was generated.")

    scannable = _strip_sql_literals_and_comments(cleaned)

    statements = [part for part in scannable.split(";") if part.strip()]
    if len(statements) > 1:
        raise ValueError("Only a single SQL statement may be executed.")

    normalized = scannable.strip().lstrip("(").lstrip()
    first_word = normalized.split(None, 1)[0].lower() if normalized.split() else ""
    if first_word not in ("select", "with"):
        raise ValueError("Only SELECT queries are allowed.")

    lowered = scannable.lower()
    for keyword in FORBIDDEN_SQL_KEYWORDS:
        if re.search(rf"\b{keyword}\b", lowered):
            raise ValueError(f"Query contains a disallowed keyword: {keyword.upper()}.")

    return cleaned.rstrip().rstrip(";")


# ---------------------------------------------------------------------------
# Chart planning and rendering
# ---------------------------------------------------------------------------


def _looks_like_time_series(x_values: List[str]) -> bool:
    patterns = (
        r"^\d{4}-\d{2}",
        r"^\d{4}$",
        r"\b(jan|feb|mar|apr|may|jun|jul|aug|sep|oct|nov|dec)\b",
    )
    sample = x_values[: min(len(x_values), 5)]
    return any(
        any(re.search(pattern, value.lower()) for pattern in patterns)
        for value in sample
    )


def _infer_chart_spec(question: str, data_points: List[tuple]) -> Dict[str, str]:
    """Pick chart type and labels using simple rules (no LLM)."""
    question_lower = question.lower()
    x_values = [str(row[0]) for row in data_points]

    pie_keywords = ("pie", "donut", "share", "percentage", "percent", "proportion", "composition")
    line_keywords = ("over time", "trend", "timeline", "time series", "monthly", "weekly", "daily")
    scatter_keywords = ("scatter", "correlation", " vs ", " versus ")

    if any(keyword in question_lower for keyword in pie_keywords) and len(data_points) <= MAX_PIE_SLICES:
        chart_type = "pie"
    elif (
        any(keyword in question_lower for keyword in scatter_keywords)
        and data_points
        and _is_numeric(data_points[0][0])
    ):
        chart_type = "scatter"
    elif any(keyword in question_lower for keyword in line_keywords) or _looks_like_time_series(x_values):
        chart_type = "line"
    elif "line chart" in question_lower:
        chart_type = "line"
    else:
        chart_type = "bar"

    title = question.strip().rstrip("?")
    if not title or len(title) > 80:
        title = "Data Visualization"
    else:
        title = title[0].upper() + title[1:]

    if chart_type == "line" and _looks_like_time_series(x_values):
        xlabel = "Time Period"
    elif chart_type == "pie":
        xlabel = ""
    else:
        xlabel = "Category"

    ylabel = "Value"
    if any(word in question_lower for word in ("expense", "cost", "spend", "amount", "revenue", "sales")):
        ylabel = "Amount (USD)" if "usd" in question_lower or "expense" in question_lower else "Amount"

    return {
        "chart_type": chart_type,
        "title": title,
        "xlabel": xlabel,
        "ylabel": ylabel,
    }


def _validate_and_prepare_chart_data(
    data_points: List[tuple],
    chart_type: str,
) -> Tuple[List[tuple], str]:
    """Ensure rows are plottable and cap how many points are rendered."""
    if not data_points:
        raise ValueError("Query returned no rows to plot.")

    normalized: List[tuple] = []
    for index, row in enumerate(data_points):
        if not isinstance(row, (tuple, list)):
            raise ValueError(f"Row {index + 1} is not a plottable record.")
        if len(row) < 2:
            raise ValueError("Each row must have at least two columns (label and value).")
        normalized.append(tuple(row))

    for index, row in enumerate(normalized):
        _coerce_numeric(row[1], f"Value in row {index + 1}")

    if chart_type == "scatter":
        for index, row in enumerate(normalized):
            if not _is_numeric(row[0]):
                raise ValueError(
                    f"Scatter plots need numeric X values; got {row[0]!r} in row {index + 1}."
                )

    warnings: List[str] = []
    total_rows = len(normalized)

    if chart_type == "pie" and total_rows > MAX_PIE_SLICES:
        normalized = normalized[:MAX_PIE_SLICES]
        warnings.append(f"Pie chart limited to {MAX_PIE_SLICES} slices.")

    if len(normalized) > MAX_CHART_ROWS:
        normalized = normalized[:MAX_CHART_ROWS]
        warnings.append(f"Chart shows the first {MAX_CHART_ROWS} of {total_rows} rows.")

    return normalized, " ".join(warnings)


def _session_id_from_db_path(db_path: str) -> str:
    return db_path.split("/")[-1].replace(".db", "")


def _draw_chart(
    ax: plt.Axes,
    chart_type: str,
    x_data: List[str],
    y_data: List[float],
    data_points: List[tuple],
) -> None:
    if chart_type == "line":
        ax.plot(x_data, y_data, color=CHART_COLORS[0], marker="o", linewidth=2.5, markersize=8)
    elif chart_type == "bar":
        ax.bar(x_data, y_data, color=CHART_COLORS[1], edgecolor="none", alpha=0.9, width=0.6)
    elif chart_type == "scatter":
        x_numeric = [_coerce_numeric(row[0], "X value") for row in data_points]
        ax.scatter(x_numeric, y_data, color=CHART_COLORS[2], s=100, alpha=0.8)
    elif chart_type == "pie":
        ax.pie(
            y_data,
            labels=x_data,
            autopct="%1.1f%%",
            startangle=140,
            colors=CHART_COLORS,
        )
    else:
        ax.bar(x_data, y_data, color=CHART_COLORS[1], edgecolor="none", alpha=0.9, width=0.6)


def _apply_chart_style(
    fig: plt.Figure,
    ax: plt.Axes,
    chart_type: str,
    title: str,
    xlabel: str,
    ylabel: str,
) -> None:
    if chart_type == "pie":
        fig.suptitle(title, fontsize=13, fontweight="bold", y=0.98)
        return

    ax.set_title(title, fontsize=13, fontweight="bold", pad=15)
    ax.set_xlabel(xlabel, fontsize=10, labelpad=10)
    ax.set_ylabel(ylabel, fontsize=10, labelpad=10)
    ax.grid(axis="y", linestyle="--", alpha=0.3)
    plt.setp(ax.get_xticklabels(), rotation=25, ha="right")
    for spine in ("top", "right", "left"):
        ax.spines[spine].set_visible(False)


def _save_chart(fig: plt.Figure, session_id: str) -> str:
    os.makedirs(CHART_OUTPUT_DIR, exist_ok=True)
    chart_path = f"{CHART_OUTPUT_DIR}/{session_id}.png"
    fig.savefig(chart_path, dpi=150)
    return chart_path


def _chart_node_result(
    *,
    chart_file_path: str = "",
    chart_error: str = "",
    chart_warning: str = "",
) -> Dict[str, str]:
    return {
        "chart_file_path": chart_file_path,
        "chart_error": chart_error,
        "chart_warning": chart_warning,
    }


# ---------------------------------------------------------------------------
# Graph nodes
# ---------------------------------------------------------------------------


def classify_intent_node(state: AgentState) -> Dict[str, Any]:
    """Fallback routing when chart intent was not set in build_initial_state()."""
    chart_requested = is_chart_requested(state["user_question"])
    print(f"[Router] chart_requested={chart_requested}")
    return {"chart_requested": chart_requested}


def route_from_start(state: AgentState) -> Literal["generate_sql", "classify_intent"]:
    if state.get("chart_intent_resolved"):
        return "generate_sql"
    return "classify_intent"


def generate_sql_node(state: AgentState) -> Dict[str, Any]:
    question = state["user_question"]
    error = state.get("error_message", "")
    schema_context = _load_table_schema(state["db_path"])

    system_prompt = (
        "You are an elite SQL expert. Write a raw SQL query targeting a table named "
        "'uploaded_data' to answer the user's question.\n"
        f"Available columns:\n{schema_context}\n\n"
        "When aggregating or ranking, include the computed values in SELECT "
        "(e.g. totals, counts, averages), not just category labels.\n"
        "CRITICAL: Return ONLY the executable SQL query string. "
        "Do not use markdown backticks or explanations."
    )
    if error:
        system_prompt += f"\n\nPrevious query failed with: {error}. Fix the query and try again."

    response = llm.invoke([
        SystemMessage(content=system_prompt),
        HumanMessage(content=question),
    ])
    return {"generated_sql": response.content.strip(), "error_message": ""}


def execute_sql_node(state: AgentState) -> Dict[str, Any]:
    conn = None
    try:
        safe_sql = _validate_sql(state["generated_sql"])
        conn = _connect_readonly(state["db_path"])

        # Abort runaway queries (e.g. accidental cross joins) instead of hanging
        # the Slack thread forever.
        deadline = time.monotonic() + MAX_QUERY_SECONDS
        conn.set_progress_handler(
            lambda: 1 if time.monotonic() > deadline else 0,
            PROGRESS_HANDLER_INSTRUCTIONS,
        )

        cursor = conn.cursor()
        cursor.execute(safe_sql)
        rows = cursor.fetchmany(MAX_RESULT_ROWS)
        return {"query_result": rows, "generated_sql": safe_sql, "error_message": ""}
    except Exception as exc:
        return {
            "error_message": str(exc),
            "retry_count": state.get("retry_count", 0) + 1,
        }
    finally:
        if conn is not None:
            conn.close()


def generate_chart_node(state: AgentState) -> Dict[str, str]:
    question = state["user_question"]
    session_id = _session_id_from_db_path(state["db_path"])

    try:
        raw_rows = _parse_query_result(state["query_result"])
        chart_spec = _infer_chart_spec(question, raw_rows)
        chart_type = chart_spec["chart_type"]
        rows, chart_warning = _validate_and_prepare_chart_data(raw_rows, chart_type)

        x_data = [str(row[0]) for row in rows]
        y_data = [_coerce_numeric(row[1], "Value") for row in rows]

        fig, ax = plt.subplots(figsize=(10, 5))
        _draw_chart(ax, chart_type, x_data, y_data, rows)
        _apply_chart_style(
            fig,
            ax,
            chart_type,
            chart_spec["title"],
            chart_spec["xlabel"],
            chart_spec["ylabel"],
        )
        fig.tight_layout()

        chart_path = _save_chart(fig, session_id)
        plt.close(fig)

        print(f"[Chart] saved {chart_type} chart to {chart_path}")
        if chart_warning:
            print(f"[Chart] {chart_warning}")

        return _chart_node_result(chart_file_path=chart_path, chart_warning=chart_warning)

    except ValueError as exc:
        print(f"[Chart] failed: {exc}")
        return _chart_node_result(chart_error=str(exc))
    except Exception as exc:
        print(f"[Chart] unexpected error: {exc}")
        return _chart_node_result(chart_error=f"Unexpected chart rendering error: {exc}")


def _build_summary_prompt(
    question: str,
    data_summary: str,
    generated_sql: str,
    chart_file_path: str,
    chart_error: str,
    chart_requested: bool,
) -> str:
    prompt = (
        "You are an elite corporate financial data analyst answering a Slack user.\n"
        f"User question: '{question}'\n\n"
        f"SQL executed on their uploaded spreadsheet:\n{generated_sql}\n\n"
        f"Query result (this is the computed answer — not a profile of the full dataset):\n"
        f"{data_summary}\n\n"
        "Answer the question directly using the query result. "
        "If the result has one or a few rows, each row is the answer — do not claim "
        "the spreadsheet lacks other categories or data.\n"
        "Lead with the direct answer, then add brief context if helpful."
    )

    if chart_file_path:
        prompt += (
            "\n\nThe system already attached a chart to the user's message. "
            "Do NOT provide code or instructions for building a chart. "
            "Focus on business insights from the data."
        )
    elif chart_requested and chart_error:
        prompt += (
            f"\n\nChart generation failed ({chart_error}). "
            "Do not claim a chart was shown. Summarize the data in text instead."
        )

    return prompt


def formulate_response_node(state: AgentState) -> Dict[str, Any]:
    prompt = _build_summary_prompt(
        question=state["user_question"],
        data_summary=_summarize_query_result_for_prompt(state["query_result"]),
        generated_sql=state.get("generated_sql", ""),
        chart_file_path=state.get("chart_file_path", ""),
        chart_error=state.get("chart_error", ""),
        chart_requested=state.get("chart_requested", False),
    )
    response = llm.invoke([HumanMessage(content=prompt)])
    return {"final_response": str(response.content).strip()}


def route_after_sql(state: AgentState) -> Literal["generate_sql", "formulate_response", "generate_chart"]:
    if state.get("error_message") and state.get("retry_count", 0) < MAX_SQL_RETRIES:
        print(f"[SQL] retrying after error: {state['error_message']}")
        return "generate_sql"
    if state.get("chart_requested"):
        return "generate_chart"
    return "formulate_response"


# ---------------------------------------------------------------------------
# Graph assembly
# ---------------------------------------------------------------------------

workflow = StateGraph(AgentState)

workflow.add_node("classify_intent", classify_intent_node)
workflow.add_node("generate_sql", generate_sql_node)
workflow.add_node("execute_sql", execute_sql_node)
workflow.add_node("generate_chart", generate_chart_node)
workflow.add_node("formulate_response", formulate_response_node)

workflow.add_conditional_edges(
    START,
    route_from_start,
    {"generate_sql": "generate_sql", "classify_intent": "classify_intent"},
)
workflow.add_edge("classify_intent", "generate_sql")
workflow.add_edge("generate_sql", "execute_sql")
workflow.add_conditional_edges(
    "execute_sql",
    route_after_sql,
    {
        "generate_sql": "generate_sql",
        "generate_chart": "generate_chart",
        "formulate_response": "formulate_response",
    },
)
workflow.add_edge("generate_chart", "formulate_response")
workflow.add_edge("formulate_response", END)

analytical_agent = workflow.compile()