Experiments/visualizer.py

import os
import sys
import re
from io import StringIO
from typing import Optional
from dotenv import load_dotenv
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from pydantic import BaseModel, Field
from langchain_core.tools import tool
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain.agents import AgentExecutor, create_tool_calling_agent

# --- 1. SETUP & CONFIGURATION ---

# Load environment variables

load_dotenv()

api_key = os.getenv("GOOGLE_API_KEY")

if not api_key:
    raise ValueError("GOOGLE_API_KEY not found. Make sure it's set in the .env file.")

# Define paths
CSV_PATH = r"D:\Code Assistant\avocado.csv"
VISUALIZATIONS_DIR = "visualizations"

# Create the output directory if it doesn't exist
if not os.path.exists(VISUALIZATIONS_DIR):
    os.makedirs(VISUALIZATIONS_DIR)
    print(f"Created directory: {VISUALIZATIONS_DIR}")

# Initialize Gemini LLM
llm_model = ChatGoogleGenerativeAI(
    model="gemini-2.5-flash", # Using 1.5 Flash for better instruction following
    google_api_key=api_key
)

# --- 2. THE NEW, VISUALIZATION-FOCUSED PROMPT ---

VISUALIZATION_PROMPT = """You are an expert Data Analyst and Storyteller. Your sole mission is to analyze the provided CSV dataset and generate a series of 10 insightful, professional-quality visualizations that tell a compelling story about the data.

**Your Mandated Workflow:**

1.  **Initial Reconnaissance:** You MUST first call the `eda_fact_sheet` tool on the provided `df_path` to understand the data's structure, columns, and content.
2.  **Strategic Visualization Plan:** After analyzing the fact sheet, you must decide on 10 key insights you want to visualize. Your goal is to cover different aspects of the data, such as trends over time, comparisons between categories, distributions, and relationships between variables.
3.  **Execution & Generation:** For each of the 10 insights, you MUST perform the following steps:
    a. Use the `python_repl_ast` tool to write and execute Python code using `pandas`, `matplotlib`, and `seaborn`.
    b. Your code MUST generate a high-quality plot (e.g., line chart, bar chart, histogram, scatter plot, box plot).
    c. **CRITICAL SAVE COMMAND:** You MUST save each plot as a unique `.png` file directly into the `visualizations/` directory. Use a clear, descriptive filename (e.g., `visualizations/avg_price_by_region.png`).
    d. **CRITICAL REPORTING COMMAND:** After generating all 10 plots, your final output must be a well-structured text report.
4.  **Final Report Synthesis:** Your final output will be a markdown-formatted report that lists each visualization you created. For each visualization, you MUST provide:
    a. A clear title for the visualization.
    b. The exact file path in the format `(File: visualizations/your_filename.png)`.
    c. A brief, 1-2 sentence explanation of what the visualization shows and the business insight it reveals.

**CRITICAL DIRECTIVES:**
- You MUST generate exactly 10 visualizations.
- Every visualization MUST be saved to the `visualizations/` directory.
- Your final output MUST be the text report containing the list of all 10 generated files and their explanations. Do not output anything else.
- Ensure your Python code for plotting is robust: create a figure, generate the plot, add a title and labels, save the figure, and then close the plot using `plt.close()` to avoid memory issues.
"""

# --- 3. TOOL DEFINITIONS (Your existing tools are perfect) ---

@tool
def eda_fact_sheet(path:str,sample_size: int = 3, freq_top: int = 3):
    """
    Generates a batchwise fact sheet for a dataset.

    Args:
        path (str): Path to CSV or Excel dataset.
        sample_size (int): Number of example values per column.
        freq_top (int): Number of top frequent values per column.

    Returns:
        dict: Fact sheet containing dataset-level info and column-level stats in batches.
    """
    path = path.strip().strip('"').strip("'")
    path = os.path.normpath(path)

    #--- Load dataset ---
    if path.endswith(".csv"):
        df = pd.read_csv(path)
    elif path.endswith(".xlsx") or path.endswith(".xls"):
        df = pd.read_excel(path)
    else:
        raise ValueError("Unsupported file format. Use CSV or Excel.")

    df = pd.read_csv(path)

    n_rows, n_columns = df.shape

    fact_sheet = {
        "n_rows": n_rows,
        "n_columns": n_columns,
        "batches": []
    }

    # --- Determine batch splits ---
    batch_size = 15
    batch_splits = [df.columns[i:i+batch_size].tolist() for i in range(0, n_columns, batch_size)]

    # --- Process each batch ---
    for batch_cols in batch_splits:
        batch_profile = {"columns": {}}
        for col in batch_cols:
            series = df[col]
            col_profile = {}

            total = len(series)

            # Core stats
            col_profile["dtype"] = str(series.dtype)
            col_profile["null_percent"] = round(float(series.isna().sum() / total * 100), 2)
            col_profile["unique_percent"] = round(float(series.nunique(dropna=True) / total * 100), 2)

            # Example values
            try:
                col_profile["examples"] = series.dropna().sample(
                    min(sample_size, series.dropna().shape[0]),
                    random_state=42
                ).tolist()
            except ValueError:
                col_profile["examples"] = []

            # Top frequent values
            if not series.isna().all():
                top_freq = series.value_counts(dropna=True).head(freq_top)
                col_profile["top_values"] = {str(k): int(v) for k, v in top_freq.to_dict().items()}

            # Numeric columns
            if pd.api.types.is_numeric_dtype(series):
                col_profile.update({
                    "min": float(series.min(skipna=True)),
                    "max": float(series.max(skipna=True)),
                    "mean": float(series.mean(skipna=True)),
                    "std": float(series.std(skipna=True))
                })
                if series.std(skipna=True) > 0:
                    z_scores = ((series - series.mean(skipna=True)) / series.std(skipna=True)).abs()
                    col_profile["has_outliers"] = bool((z_scores > 3).any())
                else:
                    col_profile["has_outliers"] = False

            # Datetime columns
            elif pd.api.types.is_datetime64_any_dtype(series):
                if not series.dropna().empty:
                    col_profile["min_date"] = str(series.min())
                    col_profile["max_date"] = str(series.max())

            # Text/categorical columns
            elif pd.api.types.is_object_dtype(series):
                lengths = series.dropna().astype(str).map(len)
                if not lengths.empty:
                    col_profile["avg_length"] = float(lengths.mean())
                    col_profile["max_length"] = int(lengths.max())
                unusual = series.dropna().astype(str).str.contains(r"[^a-zA-Z0-9\s]").sum()
                col_profile["unusual_char_percent"] = round(float(unusual / total * 100), 2)

            # Flags
            if series.nunique(dropna=True) == total:
                col_profile["is_identifier"] = True
            elif series.nunique(dropna=True) <= 1:
                col_profile["is_constant"] = True

            batch_profile["columns"][col] = col_profile

        fact_sheet["batches"].append(batch_profile)

    return fact_sheet


class PythonInputs(BaseModel):
    query: str = Field(description="A valid python command to run.")

@tool(args_schema=PythonInputs)
def python_repl_ast(query: str, path: Optional[str] = None) -> str:
    """
    Runs a Python command and returns the result.
    If `path` is provided, the dataframe at that path will be loaded as `df`.
    This tool has access to pandas, matplotlib, and seaborn.
    All plots MUST be saved to the 'visualizations/' directory.
    """
    # Use a global df to persist changes between calls if needed, but it's cleaner to reload
    df = None
    if path:
        path = path.strip().strip('"').strip("'")
        path = os.path.normpath(path)
        if os.path.exists(path):
            df = pd.read_csv(path)
        else:
            return f"Error: Path does not exist - {path}"

    # Restricted namespaces
    # Important: Pass the visualization libraries to the exec environment
    local_namespace = {
        "pd": pd,
        "plt": plt,
        "sns": sns,
        "df": df
    }
    global_namespace = {}

    # Capture stdout for print()
    old_stdout = sys.stdout
    sys.stdout = captured_output = StringIO()

    try:
        exec(query, global_namespace, local_namespace)
        output = captured_output.getvalue()
        if not output:
            output = "Code executed successfully. Plot saved."
        return output
    except Exception as e:
        return f"Execution failed with error: {e}"
    finally:
        sys.stdout = old_stdout

# --- 4. AGENT AND EXECUTOR SETUP ---

# Note: The system prompt is simplified. The main instructions are now in the user-facing prompt.
system_prompt_template = ChatPromptTemplate.from_messages(
    [
        (
           "system",
            "You are a helpful data analysis agent. Your primary goal is to help the user by generating visualizations based on their instructions. You have access to a CSV file at the path: {df_path}."
        ),
        ("placeholder", "{chat_history}"),
        ("human", "{input}"),
        MessagesPlaceholder(variable_name="agent_scratchpad")
    ]
)

# Create the agent
visualizer_agent = create_tool_calling_agent(
    llm=llm_model,
    tools=[eda_fact_sheet, python_repl_ast],
    prompt=system_prompt_template
)

# Create the Agent Executor
agent_executor = AgentExecutor(
    agent=visualizer_agent,
    tools=[eda_fact_sheet, python_repl_ast],
    verbose=True,
    handle_parsing_errors=True,
)

# --- 5. EXECUTION AND POST-PROCESSING ---

def main():
    print("--- Starting Visualization Agent ---")

    # The task_prompt now directly uses our new visualization-focused prompt
    task_prompt = VISUALIZATION_PROMPT

    # Invoke the agent
    result = agent_executor.invoke({
        "input": task_prompt,
        "df_path": CSV_PATH,
        "chat_history": []
    })

    # --- Post-processing and presenting the final output ---
    final_report_text = result.get("output", "Agent did not produce a final report.")

    # Use regular expressions to find all visualization paths in the report
    visualization_paths = re.findall(r"\(File: (visualizations/[-_a-zA-Z0-9]+\.png)\)", final_report_text)

    print("\n" + "="*50)
    print("✅ AGENT RUN COMPLETE")
    print("="*50 + "\n")

    # Print the text of the report
    print("--- AGENT'S FINAL REPORT ---")
    print(final_report_text)

    # Print the list of generated assets and verify they exist
    print("\n" + "="*50)
    print("🖼️  VERIFIED VISUAL ASSETS")
    print("="*50)

    if visualization_paths:
        print(f"Agent reported creating {len(visualization_paths)} visualizations:\n")
        for path in visualization_paths:
            if os.path.exists(path):
                print(f"  - ✅ {path} (File found)")
            else:
                print(f"  - ❌ {path} (ERROR: File not found!)")
    else:
        print("Agent did not report creating any visualization files.")

    # Count actual files in the directory for a final check
    actual_files = os.listdir(VISUALIZATIONS_DIR)
    print(f"\nFound {len(actual_files)} files in the '{VISUALIZATIONS_DIR}' directory.")
    print("="*50)

if __name__ == "__main__":
    main()