agents/sql_agent/nodes.py

import sqlite3
import pandas as pd
import re
from agents.llms import LLM
from agents.tools import PlotSQLTool
from .states import SQLAgentState
from utils.consts import DB_PATH


def choose_visualization(state: SQLAgentState) -> SQLAgentState:
    """Use LLM to suggest a suitable chart type for the SQL result."""
    question = state['question']
    sql_query = state['sql_query']
    sql_result = state['sql_result']
    # Convert sql_result DataFrame to markdown or string preview (or sample rows)
    if sql_result is not None:
        if hasattr(sql_result, 'head'):
            preview = sql_result.head(5).to_markdown(index=False)
        else:
            preview = str(sql_result)
    else:
        preview = "No results"

    prompt = f'''
You are an AI assistant that recommends appropriate data visualizations. Based on the user's question, SQL query, and query results, suggest the most suitable type of graph or chart to visualize the data. If no visualization is appropriate, indicate that.

Available chart types and their use cases:
- Bar Graphs: Best for comparing categorical data or showing changes over time when categories are discrete and the number of categories is more than 2.
- Horizontal Bar Graphs: Best for comparing categorical data or showing changes over time when the number of categories is small or the disparity between categories is large.
- Scatter Plots: Useful for identifying relationships or correlations between two numerical variables or plotting distributions of data. Best used when both x axis and y axis are continuous.
- Pie Charts: Ideal for showing proportions or percentages within a whole.
- Line Graphs: Best for showing trends and distributions over time. Best used when both x axis and y axis are continuous or time-based.

Provide your response in the following format:
Recommended Visualization: [Chart type or "None"]. ONLY use the following names: bar, horizontal_bar, line, pie, scatter, none
Reason: [Brief explanation for your recommendation]

User question: {question}
SQL query: {sql_query}
Query results: {preview}

Recommend a visualization:
'''
    llm = LLM()
    response = llm.generate(prompt)
    lines = response.split('\n')
    visualization = 'none'
    reason = ''
    for line in lines:
        if line.lower().startswith('recommended visualization:'):
            visualization = line.split(':', 1)[1].strip()
        elif line.lower().startswith('reason:'):
            reason = line.split(':', 1)[1].strip()
    state['visualization'] = visualization
    state['visualization_reason'] = reason
    state['step'] = 'choose_visualization'
    return state


def format_data_for_visualization(state: SQLAgentState) -> SQLAgentState:
    """
    Format the data for the chosen visualization type.
    Hỗ trợ line, bar, scatter, grouped bar, fallback LLM cho các visualization khác.
    """
    import json
    import pandas as pd
    llm = LLM()

    visualization = state.get('visualization', 'none')
    sql_result = state.get('sql_result')
    question = state.get('question')
    sql_query = state.get('sql_query')

    # Convert DataFrame to list of lists for processing
    if sql_result is not None and hasattr(sql_result, 'values'):
        data = sql_result.values.tolist()
        columns = list(sql_result.columns)
    elif isinstance(sql_result, list):
        data = sql_result
        columns = []
    else:
        state['formatted_data_for_visualization'] = None
        return state

    def _format_line_data(data, question):
        if len(data[0]) == 2:
            x_values = [str(row[0]) for row in data]
            y_values = [float(row[1]) for row in data]
            prompt = f"""
You are a data labeling expert. Given a question and some data, provide a concise and relevant label for the data series.
Question: {question}
Data (first few rows): {data[:2]}
Provide a concise label for this y axis.
"""
            label = llm.generate(prompt).strip()
            formatted_data = {
                "xValues": x_values,
                "yValues": [
                    {
                        "data": y_values,
                        "label": label
                    }
                ]
            }
            return formatted_data
        elif len(data[0]) == 3:
            data_by_label = {}
            x_values = []
            labels = list(set(item2 for item1, item2, item3 in data if isinstance(item2, str) and not item2.replace(".", "").isdigit() and "/" not in item2))
            if not labels:
                labels = list(set(item1 for item1, item2, item3 in data if isinstance(item1, str) and not item1.replace(".", "").isdigit() and "/" not in item1))
            for item1, item2, item3 in data:
                if isinstance(item1, str) and not item1.replace(".", "").isdigit() and "/" not in item1:
                    label, x, y = item1, item2, item3
                else:
                    x, label, y = item1, item2, item3
                if str(x) not in x_values:
                    x_values.append(str(x))
                if label not in data_by_label:
                    data_by_label[label] = []
                data_by_label[label].append(float(y))
                for other_label in labels:
                    if other_label != label:
                        if other_label not in data_by_label:
                            data_by_label[other_label] = []
                        data_by_label[other_label].append(None)
            y_values = [
                {
                    "data": data,
                    "label": label
                }
                for label, data in data_by_label.items()
            ]
            formatted_data = {
                "xValues": x_values,
                "yValues": y_values,
                "yAxisLabel": ""
            }
            prompt = f"""
You are a data labeling expert. Given a question and some data, provide a concise and relevant label for the y-axis.
Question: {question}
Data (first few rows): {data[:2]}
Provide a concise label for the y-axis.
"""
            y_axis_label = llm.generate(prompt).strip()
            formatted_data["yAxisLabel"] = y_axis_label
            return formatted_data
        return None

    def _format_scatter_data(data):
        formatted_data = {"series": []}
        if len(data[0]) == 2:
            formatted_data["series"].append({
                "data": [
                    {"x": float(x), "y": float(y), "id": i+1}
                    for i, (x, y) in enumerate(data)
                ],
                "label": "Data Points"
            })
        elif len(data[0]) == 3:
            entities = {}
            for item1, item2, item3 in data:
                if isinstance(item1, str) and not item1.replace(".", "").isdigit() and "/" not in item1:
                    label, x, y = item1, item2, item3
                else:
                    x, label, y = item1, item2, item3
                if label not in entities:
                    entities[label] = []
                entities[label].append({"x": float(x), "y": float(y), "id": len(entities[label])+1})
            for label, d in entities.items():
                formatted_data["series"].append({
                    "data": d,
                    "label": label
                })
        else:
            raise ValueError("Unexpected data format in results")
        return formatted_data

    def _format_bar_data(data, question):
        if len(data[0]) == 2:
            labels = [str(row[0]) for row in data]
            values = [float(row[1]) for row in data]
            prompt = f"""
You are a data labeling expert. Given a question and some data, provide a concise and relevant label for the data series.
Question: {question}
Data (first few rows): {data[:2]}
Provide a concise label for this y axis.
"""
            label = llm.generate(prompt).strip()
            y_values = [{"data": values, "label": label}]
        elif len(data[0]) == 3:
            categories = set(row[1] for row in data)
            labels = list(categories)
            entities = set(row[0] for row in data)
            y_values = []
            for entity in entities:
                entity_data = [float(row[2]) for row in data if row[0] == entity]
                y_values.append({"data": entity_data, "label": str(entity)})
        else:
            raise ValueError("Unexpected data format in results")
        formatted_data = {
            "labels": labels,
            "values": y_values
        }
        return formatted_data

    def _format_other_visualizations(visualization, question, sql_query, data):
        # Fallback: use LLM to format data
        prompt = f"""
You are a Data expert who formats data according to the required needs. You are given the question asked by the user, its sql query, the result of the query and the format you need to format it in.
For the given question: {question}\n\nSQL query: {sql_query}\n\nResult: {data}\n\nFormat this data for visualization type: {visualization}. Just give the json string. Do not format it.
"""
        response = llm.generate(prompt)
        try:
            formatted_data_for_visualization = json.loads(response)
            return formatted_data_for_visualization
        except json.JSONDecodeError:
            return {"error": "Failed to format data for visualization", "raw_response": response}

    visualization_map = {
        "none": lambda data: None,
        "scatter": lambda data: _format_scatter_data(data),
        "bar": lambda data, question: _format_bar_data(data, question),
        "horizontal_bar": lambda data, question: _format_bar_data(data, question),
        "line": lambda data, question: _format_line_data(data, question)
    }
    try:
        state["formatted_data_for_visualization"] = visualization_map[visualization](data, question)
    except (KeyError, Exception):
        state["formatted_data_for_visualization"] = _format_other_visualizations(visualization, question, sql_query, data)
    state['step'] = 'format_data_for_visualization'
    return state




def render_visualization(state: SQLAgentState) -> SQLAgentState:
    """
    Render the visualization from formatted data.
    Output: path to saved image file.
    """
    import matplotlib.pyplot as plt
    import os
    from io import BytesIO
    import uuid

    data = state.get("formatted_data_for_visualization")
    visualization = state.get("visualization", "none")

    if not data:
        state["visualization_output"] = None
        return state

    output_dir = "output/plots"
    os.makedirs(output_dir, exist_ok=True)

    def save_fig(fig):
        file_path = os.path.join(output_dir, f"visualization_{uuid.uuid4().hex[:8]}.png")
        fig.savefig(file_path, format="png", bbox_inches="tight")
        plt.close(fig)
        return file_path

    def render_line(data):
        fig, ax = plt.subplots()
        x = data["xValues"]
        for series in data["yValues"]:
            ax.plot(x, series["data"], label=series["label"])
        ax.set_xlabel("X")
        ax.set_ylabel(data.get("yAxisLabel", "Y"))
        ax.legend()
        return save_fig(fig)

    def render_bar(data, horizontal=False):
        fig, ax = plt.subplots()
        labels = data["labels"]
        n_series = len(data["values"])
        width = 0.8 / n_series
        x_indexes = list(range(len(labels)))
        for i, series in enumerate(data["values"]):
            offset = (i - n_series / 2) * width + width / 2
            if horizontal:
                ax.barh(
                    [x + offset for x in x_indexes], 
                    series["data"], 
                    height=width, 
                    label=series["label"]
                )
                ax.set_yticks(x_indexes)
                ax.set_yticklabels(labels)
            else:
                ax.bar(
                    [x + offset for x in x_indexes], 
                    series["data"], 
                    width=width, 
                    label=series["label"]
                )
                ax.set_xticks(x_indexes)
                ax.set_xticklabels(labels, rotation=45, ha='right')
        ax.legend()
        return save_fig(fig)

    def render_scatter(data):
        fig, ax = plt.subplots()
        for series in data["series"]:
            xs = [point["x"] for point in series["data"]]
            ys = [point["y"] for point in series["data"]]
            ax.scatter(xs, ys, label=series["label"])
        ax.set_xlabel("X")
        ax.set_ylabel("Y")
        ax.legend()
        return save_fig(fig)

    try:
        if visualization == "line":
            image_path = render_line(data)
        elif visualization == "bar":
            image_path = render_bar(data, horizontal=False)
        elif visualization == "horizontal_bar":
            image_path = render_bar(data, horizontal=True)
        elif visualization == "scatter":
            image_path = render_scatter(data)
        else:
            state["visualization_output"] = None
            return state

        state["visualization_output"] = image_path
    except Exception as e:
        state["visualization_output"] = None
        state["error"] = f"Failed to render visualization: {str(e)}"

    state["step"] = "render_visualization"
    return state


def render_visualization(state: SQLAgentState) -> SQLAgentState:
    """
    Render the visualization from formatted data.
    Output: path to saved image file.
    """
    import matplotlib.pyplot as plt
    import os
    import uuid
    from typing import Dict, Any, Optional

    def save_fig(fig: plt.Figure) -> str:
        """Save figure to file and return the file path."""
        try:
            output_dir = "output/plots"
            os.makedirs(output_dir, exist_ok=True)
            file_path = os.path.join(output_dir, f"visualization_{uuid.uuid4().hex[:8]}.png")
            fig.savefig(file_path, format="png", bbox_inches="tight", dpi=100)
            plt.close(fig)
            return file_path
        except Exception as e:
            print(f"Error saving figure: {e}")
            return ""

    def validate_data(data: Dict[str, Any], required_keys: list) -> bool:
        """Validate that data contains all required keys and has valid values."""
        if not all(key in data for key in required_keys):
            return False
        # Check if there's actual data to plot
        if "values" in data and not data["values"]:
            return False
        if "yValues" in data and not data["yValues"]:
            return False
        return True

    def render_line(data: Dict[str, Any]) -> Optional[str]:
        """Render line chart."""
        required_keys = ["xValues", "yValues"]
        if not validate_data(data, required_keys):
            return None
            
        try:
            fig, ax = plt.subplots(figsize=(10, 6))
            x = data["xValues"]
            for series in data["yValues"]:
                if len(x) == len(series["data"]):
                    ax.plot(x, series["data"], label=series.get("label", ""), marker='o')
            
            ax.set_xlabel(data.get("xAxisLabel", "X"))
            ax.set_ylabel(data.get("yAxisLabel", "Y"))
            ax.set_title(data.get("title", ""))
            if any(series.get("label") for series in data["yValues"]):
                ax.legend()
            plt.tight_layout()
            return save_fig(fig)
        except Exception as e:
            print(f"Error rendering line chart: {e}")
            return None

    def render_bar(data: Dict[str, Any], horizontal: bool = False) -> Optional[str]:
        """Render bar chart (vertical or horizontal)."""
        required_keys = ["labels", "values"]
        if not validate_data(data, required_keys) or not data["values"]:
            return None

        try:
            fig, ax = plt.subplots(figsize=(10, 6))
            labels = data["labels"]
            n_series = len(data["values"])
            width = 0.8 / max(1, n_series)  # Prevent division by zero
            x_indexes = list(range(len(labels)))
            
            for i, series in enumerate(data["values"]):
                if not series["data"]:  # Skip empty series
                    continue
                    
                offset = (i - n_series / 2) * width + width / 2
                if horizontal:
                    ax.barh(
                        [x + offset for x in x_indexes],
                        series["data"],
                        height=width,
                        label=series.get("label", f"Series {i+1}")
                    )
                    ax.set_yticks(x_indexes)
                    ax.set_yticklabels(labels)
                    ax.set_xlabel(data.get("xAxisLabel", "Value"))
                    ax.set_ylabel(data.get("yAxisLabel", "Category"))
                else:
                    ax.bar(
                        [x + offset for x in x_indexes],
                        series["data"],
                        width=width,
                        label=series.get("label", f"Series {i+1}")
                    )
                    ax.set_xticks(x_indexes)
                    ax.set_xticklabels(labels, rotation=45, ha='right')
                    ax.set_xlabel(data.get("xAxisLabel", "Category"))
                    ax.set_ylabel(data.get("yAxisLabel", "Value"))
            
            if any(series.get("label") for series in data["values"]):
                ax.legend()
            ax.set_title(data.get("title", ""))
            plt.tight_layout()
            return save_fig(fig)
        except Exception as e:
            print(f"Error rendering {'horizontal ' if horizontal else ''}bar chart: {e}")
            return None

    def render_scatter(data: Dict[str, Any]) -> Optional[str]:
        """Render scatter plot."""
        required_keys = ["series"]
        if not validate_data(data, required_keys):
            return None

        try:
            fig, ax = plt.subplots(figsize=(10, 6))
            for series in data["series"]:
                if not series.get("data"):
                    continue
                xs = [point.get("x", 0) for point in series["data"]]
                ys = [point.get("y", 0) for point in series["data"]]
                if len(xs) == len(ys):
                    ax.scatter(
                        xs, 
                        ys, 
                        label=series.get("label"),
                        alpha=0.6,
                        edgecolors='w'
                    )
            
            ax.set_xlabel(data.get("xAxisLabel", "X"))
            ax.set_ylabel(data.get("yAxisLabel", "Y"))
            ax.set_title(data.get("title", ""))
            if any(series.get("label") for series in data["series"]):
                ax.legend()
            plt.tight_layout()
            return save_fig(fig)
        except Exception as e:
            print(f"Error rendering scatter plot: {e}")
            return None

    # Main function logic
    data = state.get("formatted_data_for_visualization")
    visualization = state.get("visualization", "none")
    state["visualization_output"] = None

    if not data or visualization == "none":
        return state

    try:
        renderers = {
            "line": lambda: render_line(data),
            "bar": lambda: render_bar(data, horizontal=False),
            "horizontal_bar": lambda: render_bar(data, horizontal=True),
            "scatter": lambda: render_scatter(data)
        }

        if visualization in renderers:
            image_path = renderers[visualization]()
            if image_path and os.path.exists(image_path):
                state["visualization_output"] = image_path
            else:
                state["error"] = "Failed to generate visualization: No valid data to display"
        else:
            state["error"] = f"Unsupported visualization type: {visualization}"

    except Exception as e:
        state["error"] = f"Error in visualization: {str(e)}"
        print(f"Visualization error: {e}")

    state["step"] = "render_visualization"
    return state


def finalize_output(state: SQLAgentState) -> SQLAgentState:
    """
    Node hợp nhất kết quả cuối cùng (answer, visualization_output, error, ...).
    Hiện tại chỉ trả về state, có thể mở rộng xử lý sau.
    """
    state['step'] = 'finalize_output'
    return state

# def ingest(state: SQLAgentState) -> SQLAgentState:
#     """Populate state.tables with list of tables in the DB."""
#     db_info = state['db_info']
#     conn = sqlite3.connect(DB_PATH)
#     try:
#         db_info['tables'] = [row[0] for row in conn.execute(
#             "SELECT name FROM sqlite_master WHERE type='table';"
#         )]
#         # Populate columns for each table
#         columns = {}
#         for table in db_info['tables']:
#             col_rows = conn.execute(f'PRAGMA table_info("{table}")').fetchall()
#             columns[table] = [r[1] for r in col_rows]
#         db_info['columns'] = columns
#         state.db_info = db_info
#     finally:
#         conn.close()
#     return state


from agents.safe_guardrails import OffTopicValidator
from guardrails import Guard

def detect_off_topic(state: SQLAgentState) -> SQLAgentState:
    """Check if the input question is off-topic."""
    question = state['question']
    validator = Guard().use(
        OffTopicValidator,
        on_fail="fix"
    )
    metadata = {
        "topic": "Database Queries",
        "additional_context": "Only accept queries related to the data on Database/CSV"
        # "additional_context": "The database is about ecommerce products with tables: products, laptops, phones, tablets, promotions, category"
    }

    validation_result = validator.validate(question, metadata=metadata)
    if validation_result.validated_output == "OFF_TOPIC":
        state['error'] = True
    else:
        state['error'] = False
    state['step'] = 'detect_off_topic'
    state['off_topic'] = validation_result.validated_output

    print(state)
    return state


def get_db_info(state: SQLAgentState) -> SQLAgentState:
    """Get database information."""
    db_info = state['db_info']
    conn = sqlite3.connect(DB_PATH)
    try:
        db_info['tables'] = [row[0] for row in conn.execute(
            "SELECT name FROM sqlite_master WHERE type='table';"
        )]
        # Populate columns for each table
        columns = {}
        for table in db_info['tables']:
            col_rows = conn.execute(f'PRAGMA table_info("{table}")').fetchall()
            columns[table] = [r[1] for r in col_rows]
        db_info['columns'] = columns
        schema = "; ".join(f"{t}({', '.join(db_info['columns'][t])})" for t in db_info['tables'])
        db_info['schema'] = schema
    finally:
        conn.close()
    state['step'] = 'get_db_info'
    return state


def generate_sql(state: SQLAgentState) -> SQLAgentState:
    """Use LLM to translate user_query into SQL."""
    llm = LLM()
    # Include detailed schema with columns
    schema = state['db_info']['schema']
    prompt = (
        f"Given this database schema: {schema}, "
        f"write an SQL query to: {state['question']}. "
        "Respond with only the SQL enclosed in triple backticks."
    )
    raw = llm.generate(prompt)
    # print('raw', raw)
    lines = raw.splitlines()
    if lines and lines[0].strip().startswith("```"):
        lines = lines[1:]
    if lines and lines[-1].strip().startswith("```"):
        lines = lines[:-1]
    state['sql_query'] = "\n".join(lines).strip()
    state['step'] = 'generate_sql'
    return state



def execute_sql(state: SQLAgentState) -> SQLAgentState:
    """Run the SQL in state.sql and store result DataFrame."""
    sql_query = state['sql_query']
    conn = sqlite3.connect(DB_PATH)
    try:
        state['sql_result'] = pd.read_sql_query(sql_query, conn)
    except Exception as e:
        state['error'] = str(e)
    finally:
        conn.close()
    state['step'] = 'execute_sql'
    return state

def generate_answer(state: SQLAgentState) -> SQLAgentState:
    """Generate answer using LLM based on SQL result."""
    llm = LLM()
    if state['sql_result'] is not None and not state['sql_result'].empty:
        result_str = state['sql_result'].to_string(index=False)
        prompt = (
            f"Given the question: {state['question']},\n"
            f"SQL Query: {state['sql_query']},\n"
            f"and the following SQL query result: {result_str},\n"
            "provide a concise answer:"
        )
        state['answer'] = llm.generate(prompt)
    else:
        state['error'] = state['error'] or "No results found."
        if state["off_topic"] == "OFF_TOPIC":
            state['error'] = "The question is off-topic." 
            # state["answer"] = "Sorry, I can't assist you with that request."
            state["answer"] = "Sorry, I can only help you with questions about the data! What information would you like to explore from the data?"
    state['step'] = 'generate_answer'
    return state


def optional_plot(state: SQLAgentState) -> SQLAgentState:
    """If user_query requests plotting, generate plot and set state.plot_path."""
    if any(k in state['question'].lower() for k in ['plot', 'vẽ', 'biểu đồ']):
        tool = PlotSQLTool()
        md = tool._run(state['sql_query'])
        m = re.search(r'!\[.*\]\((.*?)\)', md)
        if m:
            state['plot_path'] = m.group(1)
        else:
            state['error'] = state['error'] or 'Plot generation failed'
    return state


def format_response(state: SQLAgentState) -> SQLAgentState:
    """Build markdown response including SQL, table preview, and plot."""
    parts = []
    if state['sql_query']:
        parts.append(f"```sql\n{state['sql_query']}\n```")
    if state['sql_result'] is not None:
        parts.append(state['sql_result'].to_markdown(index=False))
    if state['plot_path']:
        parts.append(f"![Plot]({state['plot_path']})")
    if state['error']:
        parts.append(f"**Error**: {state['error']}")
    state['response_md'] = "\n\n".join(parts)
    return state