initial commit

agents.py

@@ -0,0 +1,744 @@
+import os
+import re
+import json
+import logging
+import httpx
+import pandas as pd
+import numpy as np
+from typing import Dict, List, TypedDict, Annotated, Literal
+from fastapi import HTTPException
+from langgraph.graph import StateGraph, START, END
+from langgraph.prebuilt import ToolNode
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage
+from langchain_core.tools import tool
+import matplotlib.pyplot as plt
+import seaborn as sns
+from scipy import stats
+import warnings
+import io
+import base64
+import tempfile
+from dotenv import load_dotenv
+
+# Load environment variables from .env file
+load_dotenv()
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+# Gemini API configuration
+GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")
+GEMINI_MODEL = os.getenv("GEMINI_MODEL", "gemini-2.5-flash-preview-05-20")
+GEMINI_BASE_URL = "https://generativelanguage.googleapis.com/v1beta/models"
+
+if not GEMINI_API_KEY:
+    raise ValueError("GEMINI_API_KEY environment variable is required")
+
+# Define the agent state
+class AgentState(TypedDict):
+    messages: Annotated[List[BaseMessage], "The conversation messages"]
+    prompt: str
+    dataframe: pd.DataFrame
+    columns: List[str]
+    intent: Dict
+    chart_config: Dict
+    code: str
+    result: Dict
+    error: str
+    next_action: str
+    plot_path: str
+
+async def generate(prompt, temperature=0.2, model="gemma3:12b-it-qat"):
+    """Generate response using your deployed Ollama API."""
+    url = "https://sumansuriya7010--ollama-server3-ollamaserver-serve.modal.run/v1/chat/completions"
+    
+    headers = {
+        "Content-Type": "application/json",
+    }
+    
+    payload = {
+        "model": model, 
+        "messages": [
+            {
+                "role": "user",
+                "content": prompt
+            }
+        ],
+        "temperature": temperature,
+        "max_tokens": 8192,
+        "stream": False
+    }
+    
+    try:
+        async with httpx.AsyncClient(timeout=120.0) as client:
+            response = await client.post(
+                url,
+                json=payload,
+                headers=headers
+            )
+            response.raise_for_status()
+            result = response.json()
+            
+            # Extract text from Ollama/OpenAI compatible response
+            if "choices" in result and len(result["choices"]) > 0:
+                choice = result["choices"][0]
+                if "message" in choice and "content" in choice["message"]:
+                    return choice["message"]["content"]
+            
+            return ""
+            
+    except httpx.HTTPStatusError as e:
+        logger.error(f"HTTP error from Ollama API: {e.response.status_code} - {e.response.text}")
+        raise HTTPException(status_code=e.response.status_code, detail=f"Ollama API error: {e.response.text}")
+    except Exception as e:
+        logger.error(f"Error generating response with Ollama: {str(e)}")
+        raise HTTPException(status_code=500, detail=f"Error generating response with Ollama: {str(e)}")
+
+def create_chart(df: pd.DataFrame, chart_config: Dict) -> str:
+    """Create a matplotlib chart and return the base64 encoded image."""
+    try:
+        plt.style.use('seaborn-v0_8')
+        fig, ax = plt.subplots(figsize=(12, 8))
+        
+        chart_type = chart_config.get("chart_type", "bar")
+        x_axis = chart_config.get("x_axis")
+        y_axis = chart_config.get("y_axis")
+        title = chart_config.get("title", "Chart")
+        aggregation = chart_config.get("aggregation", "none")
+        
+        # Handle data aggregation if needed
+        plot_df = df.copy()
+        if aggregation != "none" and x_axis and y_axis:
+            if aggregation == "sum":
+                plot_df = df.groupby(x_axis)[y_axis].sum().reset_index()
+            elif aggregation == "mean":
+                plot_df = df.groupby(x_axis)[y_axis].mean().reset_index()
+            elif aggregation == "count":
+                plot_df = df.groupby(x_axis)[y_axis].count().reset_index()
+        
+        # Create the chart based on type
+        if chart_type == "bar":
+            if aggregation != "none":
+                ax.bar(plot_df[x_axis], plot_df[y_axis])
+            else:
+                sns.barplot(data=plot_df, x=x_axis, y=y_axis, ax=ax)
+                
+        elif chart_type == "line":
+            if aggregation != "none":
+                ax.plot(plot_df[x_axis], plot_df[y_axis], marker='o')
+            else:
+                sns.lineplot(data=plot_df, x=x_axis, y=y_axis, ax=ax)
+                
+        elif chart_type == "scatter":
+            sns.scatterplot(data=plot_df, x=x_axis, y=y_axis, ax=ax)
+            
+        elif chart_type == "histogram":
+            if x_axis in df.columns:
+                ax.hist(df[x_axis].dropna(), bins=30, alpha=0.7)
+            
+        elif chart_type == "boxplot":
+            if y_axis and x_axis:
+                sns.boxplot(data=plot_df, x=x_axis, y=y_axis, ax=ax)
+            else:
+                ax.boxplot(df.select_dtypes(include=[np.number]).dropna())
+                
+        elif chart_type == "pie":
+            if x_axis:
+                value_counts = df[x_axis].value_counts()
+                ax.pie(value_counts.values, labels=value_counts.index, autopct='%1.1f%%')
+                
+        elif chart_type == "area":
+            if x_axis and y_axis:
+                ax.fill_between(plot_df[x_axis], plot_df[y_axis], alpha=0.7)
+        
+        # Customize the chart
+        ax.set_title(title, fontsize=16, fontweight='bold')
+        if x_axis and chart_type != "pie":
+            ax.set_xlabel(x_axis.replace('_', '').title(), fontsize=12)
+        if y_axis and chart_type not in ["pie", "histogram"]:
+            ax.set_ylabel(y_axis.replace('_', ' ').title(), fontsize=12)
+        
+        # Rotate x-axis labels if they're long
+        if chart_type not in ["pie", "histogram"]:
+            plt.xticks(rotation=45, ha='right')
+        
+        plt.tight_layout()
+        
+        # Save to base64
+        buffer = io.BytesIO()
+        plt.savefig(buffer, format='png', dpi=300, bbox_inches='tight')
+        buffer.seek(0)
+        image_base64 = base64.b64encode(buffer.read()).decode()
+        plt.close(fig)
+        
+        return image_base64
+        
+    except Exception as e:
+        logger.error(f"Error creating chart: {str(e)}")
+        plt.close('all')  # Clean up any open figures
+        return None
+
+# Agent nodes
+async def analyze_intent_node(state: AgentState) -> AgentState:
+    """Analyze the user's prompt to determine intent."""
+    prompt = state["prompt"]
+    columns = state["columns"]
+    
+    response_format = {
+        "intent": "statistical",
+        "reason": "Prompt requests statistical analysis",
+        "visualization_type": None,
+        "transformation_type": None,
+        "statistical_type": "correlation"
+    }
+
+    input_text = f"""Analyze the following prompt and determine if it's requesting data transformation, visualization, or statistical analysis:
+
+Prompt: {prompt}
+Available columns: {', '.join(columns)}
+
+Provide a JSON response with:
+1. intent: Either 'visualization', 'transformation', or 'statistical'
+2. reason: Brief explanation of why this classification was chosen
+3. visualization_type: If intent is 'visualization', specify the chart type ('bar', 'line', 'pie', 'scatter', 'area', 'histogram', 'boxplot')
+4. transformation_type: If intent is 'transformation', specify the operation type ('aggregate', 'filter', 'join', 'compute', 'sort', 'group')
+5. statistical_type: If intent is 'statistical', specify the test type ('correlation', 'ttest', 'regression', 'descriptive'), 
+
+Example response format:
+{json.dumps(response_format)}"""
+
+    try:
+        json_text = await generate(input_text, temperature=0.4)
+        
+        # Try to extract JSON from markdown code blocks if present
+        json_match = re.search(r"```(?:json)?\n(.*?)\n```", json_text, re.DOTALL)
+        if json_match:
+            json_text = json_match.group(1)
+        
+        json_text = json_text.strip()
+        
+        try:
+            intent = json.loads(json_text)
+        except json.JSONDecodeError:
+            # If direct parsing fails, try to extract just the JSON object
+            json_obj_match = re.search(r"(\{.*\})", json_text, re.DOTALL)
+            if json_obj_match:
+                intent = json.loads(json_obj_match.group(1))
+            else:
+                # Fallback classification based on keywords
+                prompt_lower = prompt.lower()
+                if any(word in prompt_lower for word in ['chart', 'plot', 'graph', 'visualiz', 'show']):
+                    intent = {"intent": "visualization", "reason": "Keywords suggest visualization"}
+                elif any(word in prompt_lower for word in ['filter', 'transform', 'add', 'modify', 'create column']):
+                    intent = {"intent": "transformation", "reason": "Keywords suggest transformation"}
+                else:
+                    intent = {"intent": "statistical", "reason": "Default to statistical analysis"}
+        
+        state["intent"] = intent
+        state["next_action"] = intent["intent"]
+        logger.info(f"Intent analysis result: {intent}")
+        
+    except Exception as e:
+        state["error"] = f"Error analyzing prompt intent: {str(e)}"
+        state["next_action"] = "error"
+        logger.error(f"Error in analyze_intent_node: {str(e)}")
+    
+    return state
+
+async def generate_visualization_node(state: AgentState) -> AgentState:
+    """Generate visualization configuration and create the chart."""
+    prompt = state["prompt"]
+    columns = state["columns"]
+    df = state["dataframe"]
+    
+    response_format = {
+        "chart_type": "bar",
+        "x_axis": "date",
+        "y_axis": "sales",
+        "aggregation": "sum",
+        "title": "Total Sales by Date"
+    }
+    
+    input_text = f"""Based on the following prompt, determine the appropriate chart configuration:
+
+Prompt: {prompt}
+Available columns: {', '.join(columns)}
+
+Generate a JSON configuration with:
+1. chart_type: 'bar', 'line', 'pie', 'scatter', 'area', 'histogram', 'boxplot'
+2. x_axis: column name for x-axis (choose from available columns)
+3. y_axis: column name for y-axis (can be None for histograms, choose from available columns)
+4. aggregation: 'sum', 'mean', 'count', 'none'
+5. title: descriptive chart title
+
+Example response format:
+{json.dumps(response_format)}
+
+Provide only the JSON configuration, no explanations."""
+
+    try:
+        json_text = await generate(input_text, temperature=0.5)
+        
+        json_match = re.search(r"```(?:json)?\n(.*?)\n```", json_text, re.DOTALL)
+        if json_match:
+            json_text = json_match.group(1)
+        
+        json_text = json_text.strip()
+        
+        try:
+            chart_config = json.loads(json_text)
+        except json.JSONDecodeError:
+            json_obj_match = re.search(r"(\{.*\})", json_text, re.DOTALL)
+            if json_obj_match:
+                chart_config = json.loads(json_obj_match.group(1))
+            else:
+                # Fallback configuration
+                numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
+                categorical_cols = df.select_dtypes(include=['object']).columns.tolist()
+                
+                chart_config = {
+                    "chart_type": "bar",
+                    "x_axis": categorical_cols[0] if categorical_cols else columns[0],
+                    "y_axis": numeric_cols[0] if numeric_cols else columns[1] if len(columns) > 1 else None,
+                    "aggregation": "mean" if numeric_cols else "count",
+                    "title": "Data Visualization"
+                }
+        
+        # Validate column names exist
+        if chart_config.get("x_axis") not in columns:
+            chart_config["x_axis"] = columns[0]
+        if chart_config.get("y_axis") and chart_config["y_axis"] not in columns:
+            numeric_cols = df.select_dtypes(include=[np.number]).columns.tolist()
+            chart_config["y_axis"] = numeric_cols[0] if numeric_cols else None
+                
+        state["chart_config"] = chart_config
+        
+        # Create the chart immediately
+        image_base64 = create_chart(df, chart_config)
+        if image_base64:
+            state["result"] = {
+                "type": "visualization",
+                "chart_type": chart_config["chart_type"],
+                "config": chart_config,
+                "image": image_base64,
+                "message": "Visualization created successfully"
+            }
+            state["next_action"] = "complete"
+        else:
+            state["error"] = "Failed to create visualization"
+            state["next_action"] = "error"
+            
+        logger.info(f"Generated chart config: {chart_config}")
+        
+    except Exception as e:
+        state["error"] = f"Error generating chart configuration: {str(e)}"
+        state["next_action"] = "error"
+        logger.error(f"Error in generate_visualization_node: {str(e)}")
+    
+    return state
+
+async def generate_transformation_node(state: AgentState) -> AgentState:
+    """Generate pandas transformation code."""
+    prompt = state["prompt"]
+    columns = state["columns"]
+    
+    input_text = f"""Write Python code to perform the following pandas DataFrame transformation:
+
+{prompt}
+
+Available columns: {', '.join(columns)}
+
+Pandas Knowledge Base:
+1. DataFrame Operations:
+   - select columns: df[['col1', 'col2']]
+   - filter rows: df[df['column'] > value]
+   - group data: df.groupby('column')
+   - sort data: df.sort_values('column')
+   - add/modify columns: df['new_col'] = df['col1'] * 2
+   - drop columns: df.drop(['col1'], axis=1)
+   - remove duplicates: df.drop_duplicates()
+   - merge dataframes: pd.merge(df1, df2)
+
+2. Common Functions:
+   - df.apply(): Apply function to columns/rows
+   - df.fillna(): Fill missing values
+   - df.dropna(): Drop missing values
+   - df.replace(): Replace values
+   - pd.to_datetime(): Convert to datetime
+   - df.astype(): Convert data types
+   - df.round(): Round numbers
+   - df.sum(), df.mean(), df.count(): Aggregations
+
+3. String Operations:
+   - df['col'].str.contains(): String contains
+   - df['col'].str.split(): Split strings
+   - df['col'].str.replace(): Replace in strings
+   - df['col'].str.upper(): Convert to uppercase
+
+4. Window Operations:
+   - df.rolling(): Rolling window operations
+   - df.shift(): Shift values
+   - df.expanding(): Expanding window
+
+Requirements:
+1. Use pandas DataFrame operations
+2. Handle missing values appropriately
+3. Store result in 'transformed_df'
+4. DO NOT define functions
+5. Return a pandas DataFrame
+6. Use proper type conversions if needed
+
+Available variables:
+- df: pandas DataFrame
+- pd: pandas module
+- np: numpy module
+
+Example format:
+```python
+transformed_df = df.copy()
+transformed_df['new_column'] = df['column1'] * df['column2']
+transformed_df = transformed_df.fillna(0)  # Handle nulls
+```
+
+Provide only the code, no explanations. DO NOT DEFINE functions, directly perform the operations on the df."""
+
+    try:
+        code = await generate(input_text, temperature=0.4)
+        
+        code_match = re.search(r"```python\n(.*?)\n```", code, re.DOTALL)
+        code = code_match.group(1) if code_match else code
+        
+        state["code"] = code
+        state["next_action"] = "execute"
+        logger.info(f"Generated transformation code: {code}")
+        
+    except Exception as e:
+        state["error"] = f"Error generating transformation code: {str(e)}"
+        state["next_action"] = "error"
+        logger.error(f"Error in generate_transformation_node: {str(e)}")
+    
+    return state
+
+async def generate_statistical_node(state: AgentState) -> AgentState:
+    """Generate robust pandas/numpy code for statistical analysis with fallbacks."""
+    prompt = state.get("prompt", "")
+    print(prompt+" - Prompt received in generate_statistical_node")
+
+    columns = state.get("columns", [])
+    # Use predefined templates based on prompt keywords
+    operations = []
+    if any(x in prompt.lower() for x in ["describe", "summary"]):
+        operations.append("describe")
+    if any(x in prompt.lower() for x in ["correlation", "corr"]):
+        operations.append("correlation")
+    if any(x in prompt.lower() for x in ["ttest", "hypothesis"]):
+        operations.append("ttest")
+    if not operations:
+        operations = ["describe"]  # default
+
+    code_blocks = []
+    # Build code blocks robustly
+    if "describe" in operations:
+        code_blocks.append(
+            "# Descriptive statistics\n"
+            "desc = df.describe(include='all')\n"
+        )
+    if "correlation" in operations:
+        code_blocks.append(
+            "# Correlation for numeric columns\n"
+            "num_cols = df.select_dtypes(include=[np.number]).columns.tolist()\n"
+            "corr = df[num_cols].corr() if len(num_cols) > 1 else pd.DataFrame()\n"
+        )
+    if "ttest" in operations and 'category' in columns:
+        # safe t-test only if category and value exist
+        code_blocks.append(
+            "# Independent T-test between two groups in 'category' on 'value' column\n"
+            "groups = df['category'].dropna().unique().tolist()[:2]\n"
+            "if len(groups) == 2:\n"
+            "    g1 = df[df['category'] == groups[0]]['value'].dropna()\n"
+            "    g2 = df[df['category'] == groups[1]]['value'].dropna()\n"
+            "    t_stat, p_val = stats.ttest_ind(g1, g2, nan_policy='omit')\n"
+            "else:\n"
+            "    t_stat, p_val = None, None\n"
+        )
+    # Assemble result dict
+    code_blocks.append(
+        "# Assemble results\n"
+        "results = {}\n"
+        "if 'desc' in locals(): results['descriptive'] = desc\n"
+        "if 'corr' in locals(): results['correlation'] = corr\n"
+        "if 't_stat' in locals(): results['ttest'] = {'t_statistic': t_stat, 'p_value': p_val}\n"
+        "# Final assignment\n"
+        "stat_result = results\n"
+    )
+
+    state['code'] = '\n'.join(code_blocks)
+    state['next_action'] = 'execute'
+    logger.info(f"Generated statistical code with operations {operations}")
+    return state
+
+
+
+async def execute_code_node(state: AgentState) -> AgentState:
+    """Execute the generated code safely."""
+    code = state["code"]
+    df = state["dataframe"]
+    
+    if not code:
+        state["error"] = "No code to execute"
+        state["next_action"] = "error"
+        return state
+    
+    try:
+        # Create safe execution environment
+        safe_globals = {
+            'df': df,
+            'pd': pd,
+            'np': np,
+            'stats': stats,
+            'plt': plt,
+            'sns': sns
+        }
+        
+        # Execute the code
+        exec(code, safe_globals)
+        
+        # Extract results based on intent
+        intent = state["intent"]["intent"]
+        
+        if intent == "transformation":
+            if 'transformed_df' in safe_globals:
+                result_df = safe_globals['transformed_df']
+                state["result"] = {
+                    "type": "transformation",
+                    "shape": result_df.shape,
+                    "columns": result_df.columns.tolist(),
+                    "preview": result_df.head(10).to_html(classes='table table-striped'),
+                    "dataframe": result_df,
+                    "message": f"Data transformed successfully. New shape: {result_df.shape}"
+                }
+            else:
+                state["error"] = "No 'transformed_df' found in execution result"
+                
+        elif intent == "statistical":
+            exec(code, safe_globals)
+            stat_result = safe_globals.get('stat_result')
+            if stat_result is None:
+                raise ValueError("'stat_result' not found after execution")
+            if not isinstance(stat_result, dict):
+                stat_result = {'result': stat_result}
+            formatted = format_statistical_result(stat_result)
+            state['result'] = {
+                'type': 'statistical',
+                'data': formatted,
+                'message': 'Statistical analysis completed successfully'
+            }
+        
+        state["next_action"] = "complete"
+        logger.info("Code executed successfully")
+        
+    except Exception as e:
+        state["error"] = f"Error executing code: {str(e)}"
+        state["next_action"] = "error"
+        logger.error(f"Error in execute_code_node: {str(e)}")
+    
+    return state
+
+def format_statistical_result(stat_result) -> str:
+    """Format statistical results for display in Gradio."""
+    try:
+        if isinstance(stat_result, pd.DataFrame):
+            return stat_result.to_html(classes='table table-striped')
+        elif isinstance(stat_result, dict):
+            html_parts = []
+            for key, value in stat_result.items():
+                html_parts.append(f"<h4>{key.replace('_', ' ').title()}</h4>")
+                if isinstance(value, pd.DataFrame):
+                    html_parts.append(value.to_html(classes='table table-striped'))
+                elif isinstance(value, (int, float)):
+                    html_parts.append(f"<p><strong>{value:.6f}</strong></p>")
+                else:
+                    html_parts.append(f"<p>{str(value)}</p>")
+            return ''.join(html_parts)
+        else:
+            return f"<p><strong>Result:</strong> {str(stat_result)}</p>"
+    except Exception as e:
+        return f"<p><strong>Error formatting result:</strong> {str(e)}</p>"
+
+async def error_handler_node(state: AgentState) -> AgentState:
+    """Handle errors and provide feedback."""
+    error = state.get("error", "Unknown error occurred")
+    logger.error(f"Error in agent workflow: {error}")
+    
+    state["result"] = {
+        "type": "error",
+        "message": error,
+        "suggestions": [
+            "Check if the column names are correct",
+            "Verify that the data types are appropriate",
+            "Ensure the prompt is clear and specific"
+        ]
+    }
+    state["next_action"] = "complete"
+    return state
+
+def route_based_on_intent(state: AgentState) -> Literal["visualization", "transformation", "statistical", "error"]:
+    """Route to appropriate node based on intent analysis."""
+    if state.get("error"):
+        return "error"
+    
+    intent = state.get("intent", {}).get("intent", "error")
+    return intent
+
+def route_to_execution(state: AgentState) -> Literal["execute", "error", "complete"]:
+    """Route to execution or error handling."""
+    if state.get("error"):
+        return "error"
+    
+    next_action = state.get("next_action", "error")
+    if next_action == "execute":
+        return "execute"
+    elif next_action == "complete":
+        return "complete"
+    else:
+        return "error"
+
+# Build the LangGraph workflow
+def create_data_analysis_agent():
+    """Create the data analysis agent using LangGraph."""
+    
+    # Create the state graph
+    workflow = StateGraph(AgentState)
+    
+    # Add nodes
+    workflow.add_node("analyze_intent", analyze_intent_node)
+    workflow.add_node("visualization", generate_visualization_node)
+    workflow.add_node("transformation", generate_transformation_node)
+    workflow.add_node("statistical", generate_statistical_node)
+    workflow.add_node("execute", execute_code_node)
+    workflow.add_node("error_handler", error_handler_node)
+    
+    # Add edges
+    workflow.add_edge(START, "analyze_intent")
+    
+    # Conditional edges based on intent
+    workflow.add_conditional_edges(
+        "analyze_intent",
+        route_based_on_intent,
+        {
+            "visualization": "visualization",
+            "transformation": "transformation", 
+            "statistical": "statistical",
+            "error": "error_handler"
+        }
+    )
+    
+    # Route from generation nodes to execution
+    workflow.add_conditional_edges(
+        "visualization",
+        route_to_execution,
+        {
+            "execute": "execute",
+            "complete": END,
+            "error": "error_handler"
+        }
+    )
+    workflow.add_conditional_edges(
+        "transformation",
+        route_to_execution,
+        {
+            "execute": "execute",
+            "complete": END,
+            "error": "error_handler"
+        }
+    )
+    workflow.add_conditional_edges(
+        "statistical",
+        route_to_execution,
+        {
+            "execute": "execute",
+            "complete": END,
+            "error": "error_handler"
+        }
+    )
+    
+    # Final edges
+    workflow.add_edge("execute", END)
+    workflow.add_edge("error_handler", END)
+    
+    # Compile the graph
+    app = workflow.compile()
+    return app
+
+# Main execution function
+async def analyze_data_with_agent(prompt: str, dataframe: pd.DataFrame) -> Dict:
+    """
+    Analyze data using the LangGraph agent.
+    
+    Args:
+        prompt: Natural language prompt describing the analysis
+        dataframe: Pandas DataFrame to analyze
+        
+    Returns:
+        Dictionary containing the analysis results
+    """
+    # Create the agent
+    agent = create_data_analysis_agent()
+    
+    # Initialize state
+    initial_state = {
+        "messages": [HumanMessage(content=prompt)],
+        "prompt": prompt,
+        "dataframe": dataframe,
+        "columns": dataframe.columns.tolist(),
+        "intent": {},
+        "chart_config": {},
+        "code": "",
+        "result": {},
+        "error": "",
+        "next_action": "",
+        "plot_path": ""
+    }
+    
+    # Run the agent
+    try:
+        final_state = await agent.ainvoke(initial_state)
+        return final_state["result"]
+    except Exception as e:
+        logger.error(f"Error running agent: {str(e)}")
+        return {
+            "type": "error",
+            "message": f"Agent execution failed: {str(e)}"
+        }
+
+# Test function
+async def test_agent():
+    """Test the data analysis agent."""
+    # Create sample data
+    data = {
+        'date': pd.date_range('2024-01-01', periods=100),
+        'sales': np.random.normal(1000, 200, 100),
+        'category': np.random.choice(['A', 'B', 'C'], 100),
+        'region': np.random.choice(['North', 'South', 'East', 'West'], 100)
+    }
+    df = pd.DataFrame(data)
+    
+    # Test different types of prompts
+    test_prompts = [
+        "Create a bar chart showing average sales by category",
+        "Calculate correlation between date and sales",
+        "Filter the data to show only category A and add a profit column that is 20% of sales"
+    ]
+    
+    for prompt in test_prompts:
+        print(f"\n--- Testing: {prompt} ---")
+        result = await analyze_data_with_agent(prompt, df)
+        print(f"Result: {result}")
+
+if __name__ == "__main__":
+    import asyncio
+    asyncio.run(test_agent())

gradio_app.py

@@ -0,0 +1,639 @@
+import gradio as gr
+import pandas as pd
+import json
+from agents import analyze_data_with_agent 
+import io
+import asyncio
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+async def process_data_and_prompt(file, prompt):
+    """Process uploaded file and prompt using the data analysis agent."""
+    try:
+        if not file:
+            return "Please upload a data file.", None, None
+            
+        if not prompt or prompt.strip() == "":
+            return "Please enter an analysis prompt.", None, None
+
+        # Read the uploaded file
+        if file.name.endswith('.csv'):
+            df = pd.read_csv(file.name)
+        elif file.name.endswith(('.xlsx', '.xls')):
+            df = pd.read_excel(file.name)
+        elif file.name.endswith('.json'):
+            df = pd.read_json(file.name)
+        else:
+            return "Error: Unsupported file format. Please upload CSV, Excel, or JSON files.", None, None
+
+        # Clean column names
+        df.columns = [str(col).strip().lower().replace(' ', '_').replace('-', '_') for col in df.columns]
+        
+        # Show data preview
+        # data_preview = f"""
+        # <div class="data-section">
+        #     <h3>Data Preview</h3>
+        #     <p><strong>Shape:</strong> {df.shape[0]} rows × {df.shape[1]} columns</p>
+        #     <p><strong>Columns:</strong> {', '.join(df.columns.tolist())}</p>
+        #     {df.head().to_html(classes='table data-table', table_id='data-preview')}
+        # </div>
+        # """
+        data_preview = f"""
+<div></div>"""
+
+        # Process with agent
+        logger.info(f"Processing prompt: {prompt}")
+        result = await analyze_data_with_agent(prompt, df)
+        logger.info(f"Agent result type: {result.get('type')}")
+
+        # Handle different result types
+        if result["type"] == "error":
+            error_html = f"""
+            <div class="error-box">
+                <h3>Error</h3>
+                <p><strong>Message:</strong> {result['message']}</p>
+                {f"<p><strong>Suggestions:</strong></p><ul>{''.join([f'<li>{s}</li>' for s in result.get('suggestions', [])])}</ul>" if result.get('suggestions') else ""}
+            </div>
+            """
+            return data_preview + error_html, None, None
+        
+        elif result["type"] == "visualization":
+            # Display the chart
+            image_base64 = result.get("image")
+            if image_base64:
+                chart_html = f"""
+                <div class="analysis-result">
+                    <h3>Visualization Result</h3>
+                    <p><strong>Chart Type:</strong> {result.get('chart_type', 'Unknown').title()}</p>
+                    <div class="chart-container">
+                        <img src="data:image/png;base64,{image_base64}" class="chart-image">
+                    </div>
+                    <p><em>{result.get('message', 'Visualization created successfully')}</em></p>
+                </div>
+                """
+                return data_preview + chart_html, None, None
+            else:
+                return data_preview + "<p>Error: Could not generate visualization</p>", None, None
+        
+        elif result["type"] == "statistical":
+            # Format statistical results
+            stat_html = f"""
+            <div class="analysis-result">
+                <h3>Statistical Analysis Results</h3>
+                <div class="stat-output-box">
+                    {result.get('data', 'No statistical results available')}
+                </div>
+                <p><em>{result.get('message', 'Statistical analysis completed')}</em></p>
+            </div>
+            """
+            return data_preview + stat_html, None, None
+        
+        elif result["type"] == "transformation":
+            # Return transformed data
+            transformed_df = result.get("dataframe")
+            if transformed_df is not None:
+                # Create CSV for download
+                csv_buffer = io.StringIO()
+                transformed_df.to_csv(csv_buffer, index=False)
+                csv_data = csv_buffer.getvalue()
+                
+                # Create temporary file for download (Gradio handles temporary files for downloads)
+                temp_file_name = "transformed_data.csv"
+                with open(temp_file_name, 'w', encoding='utf-8') as f:
+                    f.write(csv_data)
+                
+                transform_html = f"""
+                <div class="analysis-result">
+                    <h3>Data Transformation Results</h3>
+                    <p><strong>Original Shape:</strong> {df.shape[0]} rows × {df.shape[1]} columns</p>
+                    <p><strong>New Shape:</strong> {result.get('shape', 'Unknown')}</p>
+                    <p><strong>New Columns:</strong> {', '.join(result.get('columns', []))}</p>
+                    <div class="transformed-data-preview">
+                        <h4>Preview of Transformed Data:</h4>
+                        {result.get('preview', 'No preview available')}
+                    </div>
+                    <p><em>{result.get('message', 'Data transformation completed')}</em></p>
+                    <p><strong>Download the transformed data using the button below.</strong></p>
+                </div>
+                """
+                return data_preview + transform_html, temp_file_name, None
+            else:
+                return data_preview + "<p>Error: Could not retrieve transformed data</p>", None, None
+        
+        else:
+            return data_preview + f"<p>Unknown result type: {result.get('type')}</p>", None, None
+
+    except Exception as e:
+        logger.error(f"Error processing data: {str(e)}")
+        error_html = f"""
+        <div class="error-box">
+            <h3>Processing Error</h3>
+            <p><strong>Error:</strong> {str(e)}</p>
+            <p><strong>Please check:</strong></p>
+            <ul>
+                <li>File format is supported (CSV, Excel)</li>
+                <li>File is not corrupted</li>
+                <li>Prompt is clear and specific</li>
+                <li>Ollama server is running</li>
+            </ul>
+        </div>
+        """
+        return error_html, None, None
+
+def process_sync(file, prompt):
+    """Synchronous wrapper for the async processing function."""
+    try:
+        # Check if an event loop is already running
+        try:
+            loop = asyncio.get_running_loop()
+        except RuntimeError:
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+        return loop.run_until_complete(process_data_and_prompt(file, prompt))
+    except Exception as e:
+        logger.error(f"Error in sync wrapper: {str(e)}")
+        return f"Error: {str(e)}", None, None
+
+def generate_preview(file):
+    """Generate a preview of the uploaded file."""
+    try:
+        if not file:
+            return "Please upload a data file to see preview."
+            
+        # Read the uploaded file
+        if file.name.endswith('.csv'):
+            df = pd.read_csv(file.name)
+        elif file.name.endswith(('.xlsx', '.xls')):
+            df = pd.read_excel(file.name)
+        elif file.name.endswith('.json'):
+            df = pd.read_json(file.name)
+        else:
+            return "Error: Unsupported file format. Please upload CSV, Excel, or JSON files."
+
+        # Clean column names
+        df.columns = [str(col).strip().lower().replace(' ', '_').replace('-', '_') for col in df.columns]
+        
+        # Show data preview
+        data_preview = f"""
+        <div class="data-section">
+            <h3>📊 Data Preview</h3>
+            <div class="data-stats">
+                <span class="stat-badge">📏 {df.shape[0]} rows</span>
+                <span class="stat-badge">📋 {df.shape[1]} columns</span>
+            </div>
+            <div class="columns-info">
+                <strong>Columns:</strong> {', '.join(df.columns.tolist())}
+            </div>
+            <div class="table-container">
+                {df.head(4).to_html(classes='table data-table', table_id='data-preview')}
+            </div>
+        </div>
+        """
+        return data_preview
+    except Exception as e:
+        logger.error(f"Error generating preview: {str(e)}")
+        return f"<div class='error-box'>Error generating preview: {str(e)}</div>"
+
+# Sample prompts for different analysis types
+sample_prompts = {
+    "Data Transformation": [
+        "Filter data where [column] > 1000 ",
+        "Group by [column] and calculate average [values]",
+        "Create new columns based on existing ones",
+        "Remove duplicates and sort by date",
+    ],
+    "Visualization": [
+        "Create a bar chart showing the distribution of [categories]",
+        "Generate a line plot of sales over time",
+        "Make a scatter plot of [column1] vs [column2]",
+        "Show a histogram of [column2]",
+        "Create a pie chart of market share by region"
+    ],
+    "Statistical Analysis": [
+        "Calculate correlation matrix for all numeric columns",
+        "Perform descriptive statistics analysis",
+    ]
+}
+
+# Create the Gradio interface
+with gr.Blocks(
+    title="Data Analysis Agent",
+    theme=gr.themes.Soft(),
+    css="""
+    /* Main container */
+    .gradio-container {
+        max-width: 900px;
+        margin: auto;
+        padding: 20px;
+    }
+
+    /* Header styling */
+    .main-header {
+        text-align: center;
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        color: white;
+        padding: 30px;
+        border-radius: 15px;
+        margin-bottom: 30px;
+        box-shadow: 0 8px 32px rgba(0,0,0,0.1);
+    }
+    
+    .main-header h1 {
+        margin: 0;
+        font-size: 2.5em;
+        font-weight: 600;
+    }
+    
+    .main-header p {
+        margin: 10px 0 0 0;
+        font-size: 1.1em;
+        opacity: 0.9;
+    }
+
+    /* Accordion styling */
+    .gr-accordion {
+        margin-bottom: 20px !important;
+        border-radius: 12px !important;
+        border: 1px solid var(--border-color-primary) !important;
+        box-shadow: 0 2px 8px rgba(0,0,0,0.05) !important;
+        overflow: hidden !important;
+    }
+    
+    .gr-accordion-header {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
+        color: white !important;
+        padding: 15px 20px !important;
+        font-weight: 600 !important;
+        font-size: 1.1em !important;
+        border: none !important;
+        cursor: pointer !important;
+        transition: all 0.3s ease !important;
+    }
+    
+    .gr-accordion-header:hover {
+        background: linear-gradient(135deg, #5a6fd8 0%, #6a4190 100%) !important;
+        transform: translateY(-1px) !important;
+    }
+    
+    .gr-accordion-content {
+        background: var(--background-fill-secondary) !important;
+        padding: 25px !important;
+        border-top: 1px solid var(--border-color-primary) !important;
+    }
+    
+    /* Special styling for example prompt accordions */
+    .gr-accordion .gr-accordion {
+        margin-bottom: 15px !important;
+        border-radius: 8px !important;
+        box-shadow: 0 1px 4px rgba(0,0,0,0.1) !important;
+    }
+    
+    .gr-accordion .gr-accordion .gr-accordion-header {
+        background: var(--color-accent-soft) !important;
+        color: var(--text-color-body) !important;
+        padding: 12px 16px !important;
+        font-size: 1em !important;
+        font-weight: 500 !important;
+    }
+    
+    .gr-accordion .gr-accordion .gr-accordion-header:hover {
+        background: var(--color-accent) !important;
+        color: white !important;
+        transform: none !important;
+    }
+    
+    .gr-accordion .gr-accordion .gr-accordion-content {
+        background: var(--background-fill-primary) !important;
+        padding: 15px !important;
+    }
+
+    /* Section styling (keeping for compatibility) */
+    .section {
+        background: var(--background-fill-secondary);
+        border-radius: 12px;
+        padding: 25px;
+        margin-bottom: 25px;
+        border: 1px solid var(--border-color-primary);
+        box-shadow: 0 2px 8px rgba(0,0,0,0.05);
+    }
+    
+    .section h2 {
+        margin: 0 0 20px 0;
+        color: var(--text-color-body);
+        font-size: 1.4em;
+        font-weight: 600;
+        display: flex;
+        align-items: center;
+        gap: 10px;
+    }
+
+    /* File upload styling */
+    .upload-area {
+        border: 2px dashed var(--border-color-accent);
+        border-radius: 10px;
+        padding: 20px;
+        text-align: center;
+        background: var(--background-fill-primary);
+        transition: all 0.3s ease;
+    }
+    
+    .upload-area:hover {
+        border-color: var(--color-accent);
+        background: var(--background-fill-hover);
+    }
+
+    /* Data preview styling */
+    .data-section {
+        background: var(--background-fill-primary);
+        border-radius: 10px;
+        padding: 20px;
+        border: 1px solid var(--border-color-primary);
+        margin: 15px 0;
+    }
+    
+    .data-section h3 {
+        margin: 0 0 15px 0;
+        color: var(--text-color-body);
+        font-size: 1.2em;
+    }
+    
+    .data-stats {
+        display: flex;
+        gap: 10px;
+        margin-bottom: 15px;
+        flex-wrap: wrap;
+    }
+    
+    .stat-badge {
+        background: var(--color-accent-soft);
+        color: var(--text-color-body);
+        padding: 6px 12px;
+        border-radius: 20px;
+        font-size: 0.9em;
+        font-weight: 500;
+    }
+    
+    .columns-info {
+        margin-bottom: 15px;
+        padding: 10px;
+        background: var(--background-fill-secondary);
+        border-radius: 8px;
+        font-size: 0.9em;
+    }
+    
+    .table-container {
+        overflow-x: auto;
+        border-radius: 8px;
+    }
+
+    /* Table styling */
+    .table {
+        width: 100%;
+        border-collapse: collapse;
+        font-size: 0.85em;
+        background: var(--background-fill-primary);
+    }
+    
+    .table th {
+        background: var(--background-fill-secondary);
+        color: var(--text-color-body);
+        font-weight: 600;
+        padding: 12px 8px;
+        border: 1px solid var(--border-color-primary);
+        text-align: left;
+    }
+    
+    .table td {
+        padding: 10px 8px;
+        border: 1px solid var(--border-color-primary);
+        color: var(--text-color-body);
+    }
+    
+    .table tr:nth-child(even) {
+        background: var(--background-fill-hover);
+    }
+
+    /* Prompt examples styling */
+    .prompt-examples {
+        display: grid;
+        gap: 15px;
+        margin-top: 15px;
+    }
+    
+    .prompt-category {
+        background: var(--background-fill-primary);
+        border-radius: 8px;
+        padding: 15px;
+        border: 1px solid var(--border-color-primary);
+    }
+    
+    .prompt-category h4 {
+        margin: 0 0 10px 0;
+        color: var(--text-color-body);
+        font-size: 1em;
+    }
+    
+    .prompt-buttons {
+        display: flex;
+        flex-wrap: wrap;
+        gap: 8px;
+    }
+    
+    .prompt-btn {
+        font-size: 0.8em !important;
+        padding: 6px 12px !important;
+        border-radius: 15px !important;
+        background: var(--color-accent-soft) !important;
+        color: var(--text-color-body) !important;
+        border: 1px solid var(--border-color-accent) !important;
+        cursor: pointer;
+        transition: all 0.2s ease;
+    }
+    
+    .prompt-btn:hover {
+        background: var(--color-accent) !important;
+        color: white !important;
+    }
+
+    /* Analysis results styling */
+    .analysis-result {
+        background: var(--background-fill-primary);
+        border-radius: 10px;
+        padding: 20px;
+        margin: 15px 0;
+        border: 1px solid var(--border-color-primary);
+    }
+    
+    .analysis-result h3 {
+        margin: 0 0 15px 0;
+        color: var(--text-color-body);
+    }
+
+    /* Chart styling */
+    .chart-container {
+        text-align: center;
+        margin: 20px 0;
+        background: var(--background-fill-primary);
+        padding: 15px;
+        border-radius: 8px;
+        border: 1px solid var(--border-color-primary);
+    }
+    
+    .chart-image {
+        max-width: 100%;
+        height: auto;
+        border-radius: 8px;
+        box-shadow: 0 4px 12px rgba(0,0,0,0.1);
+    }
+
+    /* Error styling */
+    .error-box {
+        background: #fee;
+        border: 1px solid #fcc;
+        color: #c33;
+        padding: 15px;
+        border-radius: 8px;
+        margin: 15px 0;
+    }
+    
+    .error-box h3 {
+        margin: 0 0 10px 0;
+        color: #c33;
+    }
+
+    /* Button styling */
+    .analyze-btn {
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
+        color: white !important;
+        border: none !important;
+        border-radius: 25px !important;
+        padding: 15px 30px !important;
+        font-size: 1.1em !important;
+        font-weight: 600 !important;
+        box-shadow: 0 4px 15px rgba(102, 126, 234, 0.4) !important;
+        transition: all 0.3s ease !important;
+    }
+    
+    .analyze-btn:hover {
+        transform: translateY(-2px) !important;
+        box-shadow: 0 6px 20px rgba(102, 126, 234, 0.6) !important;
+    }
+
+    /* Responsive design */
+    @media (max-width: 768px) {
+        .gradio-container {
+            padding: 10px;
+        }
+        
+        .main-header h1 {
+            font-size: 2em;
+        }
+        
+        .section {
+            padding: 15px;
+        }
+        
+        .data-stats {
+            flex-direction: column;
+        }
+        
+        .prompt-buttons {
+            flex-direction: column;
+        }
+    }
+    """
+) as demo:
+    
+    # Header
+    gr.Markdown("""
+    # 🤖 Data Analysis Agent
+    
+    Upload your data file and describe what analysis you want to perform. The AI agent will:
+    - 📊 Create visualizations (charts, plots, graphs)
+    - 🔢 Perform statistical analysis (correlations, tests, summaries)
+    - 🔧 Transform your data (filter, aggregate, compute new columns)
+    
+    **Supported formats:** CSV, Excel (.xlsx, .xls)
+    """)
+    
+    # Step 1: File Upload
+    with gr.Accordion("📁 Step 1: Upload Your Data", open=True):
+        file_input = gr.File(
+            label="Choose your data file (CSV, Excel)",
+            file_types=[".csv", ".xlsx", ".xls"],
+            type="filepath"
+        )
+    
+    # Step 2: Data Preview
+    with gr.Accordion("👀 Step 2: Data Preview", open=True):
+        preview_output = gr.HTML(value="<p style='text-align: center; color: #888; padding: 40px;'>Upload a file to see data preview</p>")
+    
+    # Step 3: Analysis Prompt
+    with gr.Accordion("💬 Step 3: Describe Your Analysis", open=True):
+        prompt_input = gr.Textbox(
+            label="What would you like to analyze?",
+            placeholder="e.g., 'Create a bar chart showing sales by category' or 'Calculate correlation between price and quantity'",
+            lines=3
+        )
+        
+        # Example prompts in separate collapsible sections
+        gr.HTML('<h4 style="margin: 20px 0 10px 0;">💡 Need inspiration? Try these examples:</h4>')
+        
+
+        with gr.Accordion("🔧 Data Transformation Examples", open=False):
+            for prompt in sample_prompts["Data Transformation"]:
+                gr.Button(prompt, size="sm", elem_classes=["prompt-btn"]).click(
+                    lambda p=prompt: p, inputs=[], outputs=prompt_input, queue=False
+                )
+
+        with gr.Accordion("📊 Visualization Examples", open=False):
+            for prompt in sample_prompts["Visualization"]:
+                gr.Button(prompt, size="sm", elem_classes=["prompt-btn"]).click(
+                    lambda p=prompt: p, inputs=[], outputs=prompt_input, queue=False
+                )
+        
+        with gr.Accordion("📈 Statistical Analysis Examples", open=False):
+            for prompt in sample_prompts["Statistical Analysis"]:
+                gr.Button(prompt, size="sm", elem_classes=["prompt-btn"]).click(
+                    lambda p=prompt: p, inputs=[], outputs=prompt_input, queue=False
+                )
+        
+
+    
+    # Step 4: Analysis Button
+    with gr.Accordion("🚀 Step 4: Run Analysis", open=True):
+        submit_btn = gr.Button("🚀 Analyze Data", variant="primary", size="lg", elem_classes=["analyze-btn"])
+    
+    # Step 5: Results
+    with gr.Accordion("📊 Step 5: Analysis Results", open=True):
+        output = gr.HTML(value="<p style='text-align: center; color: #888; padding: 40px;'>Click 'Analyze Data' to see results here</p>")
+    
+    # Step 6: Downloads
+    with gr.Accordion("📥 Step 6: Downloads", open=True):
+        download_output = gr.File(label="Transformed Data (if applicable)", visible=True)
+        gr.HTML("<p style='color: #666; font-size: 0.9em;'>Download will appear here for data transformation results</p>")
+    
+    # Event handlers
+    file_input.change(
+        fn=generate_preview,
+        inputs=[file_input],
+        outputs=[preview_output]
+    )
+    
+    submit_btn.click(
+        fn=process_sync,
+        inputs=[file_input, prompt_input],
+        outputs=[output, download_output],
+        show_progress=True
+    )
+
+if __name__ == "__main__":
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        share=False,
+        debug=True
+    )

requirements.txt

@@ -0,0 +1,10 @@
+gradio>=4.19.2
+pandas>=2.2.0
+numpy>=1.26.0
+matplotlib>=3.8.0
+seaborn>=0.13.0
+scipy>=1.12.0
+httpx>=0.26.0
+langgraph>=0.0.20
+langchain-core>=0.1.27
+python-dotenv>=1.0.0 

visualize_agent.py

@@ -0,0 +1,156 @@
+import os
+from langgraph.graph import StateGraph, START, END
+from IPython.display import Image, display
+
+# Assuming 'agents' module and its contents (AgentState, nodes, routes) are available.
+# For a runnable example, you'd need to define these or mock them.
+# Example placeholders if 'agents.py' isn't provided:
+class AgentState:
+    """A placeholder for AgentState."""
+    pass
+
+def analyze_intent_node(state):
+    """Placeholder for analyze_intent_node."""
+    print("Analyzing intent...")
+    # In a real scenario, this would determine the next step
+    # For demonstration, let's simulate routing to visualization
+    return {"next_step": "visualization"}
+
+def generate_visualization_node(state):
+    """Placeholder for generate_visualization_node."""
+    print("Generating visualization code...")
+    # Simulate success
+    return {"code_generated": True}
+
+def generate_transformation_node(state):
+    """Placeholder for generate_transformation_node."""
+    print("Generating transformation code...")
+    return {"code_generated": True}
+
+def generate_statistical_node(state):
+    """Placeholder for generate_statistical_node."""
+    print("Generating statistical code...")
+    return {"code_generated": True}
+
+def execute_code_node(state):
+    """Placeholder for execute_code_node."""
+    print("Executing code...")
+    return {"execution_successful": True}
+
+def error_handler_node(state):
+    """Placeholder for error_handler_node."""
+    print("Handling error...")
+    return {}
+
+def route_based_on_intent(state):
+    """Placeholder for route_based_on_intent."""
+    # In a real app, this would use state to determine the route
+    if state.get("next_step") == "visualization":
+        return "visualization"
+    elif state.get("next_step") == "transformation":
+        return "transformation"
+    elif state.get("next_step") == "statistical":
+        return "statistical"
+    return "error"
+
+def route_to_execution(state):
+    """Placeholder for route_to_execution."""
+    # In a real app, this would check if code generation was successful
+    if state.get("code_generated"):
+        return "execute"
+    return "error"
+
+
+def create_visualization():
+    """Create and save a visualization of the agent workflow."""
+    # Create the state graph
+    workflow = StateGraph(AgentState)
+    
+    # Add nodes
+    workflow.add_node("analyze_intent", analyze_intent_node)
+    workflow.add_node("visualization", generate_visualization_node)
+    workflow.add_node("transformation", generate_transformation_node)
+    workflow.add_node("statistical", generate_statistical_node)
+    workflow.add_node("execute", execute_code_node)
+    workflow.add_node("error_handler", error_handler_node)
+    
+    # Add edges
+    workflow.add_edge(START, "analyze_intent")
+    
+    # Conditional edges based on intent
+    workflow.add_conditional_edges(
+        "analyze_intent",
+        route_based_on_intent,
+        {
+            "visualization": "visualization",
+            "transformation": "transformation", 
+            "statistical": "statistical",
+            "error": "error_handler"
+        }
+    )
+    
+    # Route from generation nodes to execution
+    workflow.add_conditional_edges(
+        "visualization",
+        route_to_execution,
+        {
+            "execute": "execute",
+            "complete": END, # Added 'complete' to allow direct END from visualization if needed
+            "error": "error_handler"
+        }
+    )
+    workflow.add_conditional_edges(
+        "transformation",
+        route_to_execution,
+        {
+            "execute": "execute",
+            "complete": END,
+            "error": "error_handler"
+        }
+    )
+    workflow.add_conditional_edges(
+        "statistical",
+        route_to_execution,
+        {
+            "execute": "execute",
+            "complete": END,
+            "error": "error_handler"
+        }
+    )
+    
+    # Final edges
+    workflow.add_edge("execute", END)
+    workflow.add_edge("error_handler", END)
+    
+    # Create visualization directory if it doesn't exist
+    os.makedirs("visualizations", exist_ok=True)
+    
+    # Generate and save the visualization
+    graph = workflow.compile()
+    
+    try:
+        # Get the graph as a Mermaid diagram and draw it to PNG
+        # This requires 'mermaid-py' and potentially 'puppeteer' (for playwright backend)
+        png_data = graph.get_graph().draw_mermaid_png()
+
+        # Define the filename
+        filename = os.path.join("visualizations", "mermaid_graph.png")
+
+        # Save the PNG data to a file
+        with open(filename, "wb") as f:
+            f.write(png_data)
+
+        print(f"Image successfully saved as '{filename}'")
+
+        # Optionally, display the image after saving
+        display(Image(png_data))
+    except ImportError:
+        print("Please install 'mermaid-py' to generate PNG visualizations.")
+        print("You might also need to install a browser automation tool like 'playwright' for mermaid-py.")
+    except Exception as e:
+        print(f"An error occurred during visualization: {e}")
+        # This requires some extra dependencies and is optional
+        pass
+
+if __name__ == "__main__":
+    create_visualization()