feat: example script for pandas AI

.gitignore

@@ -1,6 +1,8 @@
 .idea
 .env
 .vscode
+.xml
+.iml
 
 # Byte-compiled / optimized / DLL files
 __pycache__/

pandasai_visualization.py

@@ -0,0 +1,103 @@
+#!/usr/bin/env python3
+"""
+Visualization script using PandasAI.
+
+This script creates a sample dataframe and uses PandasAI to generate
+and save visualizations based on user queries.
+
+Usage:
+    python visualize.py "Create a bar chart of sales by region"
+
+Requirements:
+    - pandas
+    - pandasai
+    - matplotlib
+"""
+
+import os
+import sys
+import pandas as pd
+import matplotlib.pyplot as plt
+import pandasai as pai
+from dotenv import load_dotenv
+
+
+def create_sample_dataframe():
+    """Create a sample dataframe with sales data."""
+    data = {
+        'Region': ['North', 'South', 'East', 'West', 'North', 'South', 'East', 'West'],
+        'Product': ['Widget', 'Widget', 'Widget', 'Widget', 'Gadget', 'Gadget', 'Gadget', 'Gadget'],
+        'Sales': [150, 200, 120, 180, 90, 110, 95, 130],
+        'Quarter': ['Q1', 'Q1', 'Q1', 'Q1', 'Q2', 'Q2', 'Q2', 'Q2'],
+        'Year': [2023, 2023, 2023, 2023, 2023, 2023, 2023, 2023]
+    }
+    return pai.DataFrame(data)
+
+
+def visualize_data(df, query):
+    """
+    Generate visualization based on user query using PandasAI.
+
+    Args:
+        df: Pandas DataFrame containing the data
+        query: User query string describing the desired visualization
+
+    Returns:
+        Path to the saved visualization file
+    """
+    # Initialize PandasAI with an LLM
+    # Note: In a real application, you would need to set up your OpenAI API key
+    # Either set OPENAI_API_KEY environment variable or pass it directly
+    try:
+
+        # llm = OpenAI(api_token=api_key)
+        # pandas_ai = PandasAI(llm)
+
+        load_dotenv()
+        pai.api_key.set(os.environ["PANDAS_KEY"])
+
+        df.chat(query)
+
+        # Generate the visualization
+        print(f"Generating visualization for query: '{query}'")
+
+        # Save the current figure
+        output_file = "visualization_output.png"
+        plt.savefig(output_file)
+        plt.close()
+
+        print(f"Visualization saved to {output_file}")
+        return output_file
+
+    except Exception as e:
+        print(f"Error generating visualization: {str(e)}")
+        return None
+
+
+def main():
+    """Main function to run the visualization script."""
+    # Get query from command line argument
+    # if len(sys.argv) < 2:
+    #     print("Usage: python visualize.py \"Your visualization query here\"")
+    #     print("Example: python visualize.py \"Create a bar chart of sales by region\"")
+    #     return
+
+    # query = sys.argv[1]
+    query = "Plot a bar chart of sales by region"
+
+    # Create sample dataframe
+    df = create_sample_dataframe()
+    print("Sample DataFrame created:")
+    print(df.head())
+
+    # Generate and save visualization
+    output_file = visualize_data(df, query)
+
+    if output_file:
+        print(f"Visualization process completed. Output saved to: {output_file}")
+    else:
+        print("Visualization process failed.")
+
+
+if __name__ == "__main__":
+    main()

postgre_mcp_server.py

@@ -2,11 +2,13 @@ import os
 from contextlib import asynccontextmanager
 from dataclasses import dataclass
 from typing import Optional, AsyncIterator
-
 import asyncpg
-from flask.cli import load_dotenv
+from dotenv import load_dotenv
 from mcp.server.fastmcp import FastMCP, Context
 from pydantic import Field
+import pandasai as pai
+import matplotlib as plt
+import pandas as pd
 
 # Constants
 DEFAULT_QUERY_LIMIT = 100
@@ -578,5 +580,67 @@ def find_relationships(table_name: str, schema: str = 'public') -> str:
         return f"Error finding relationships: {str(e)}"
 
 
+@mcp.tool(description="Visualizes query results using a prompt and JSON data.")
+async def visualize_results(json_data: dict, vis_prompt: str) -> str:
+    """
+    Generates a visualization based on query results using PandasAI.
+
+    Args:
+        json_data (dict): A dictionary containing the query results.
+            It should have two keys:
+                - 'columns': A list of column names (strings).
+                - 'data': A list of lists, where each inner list represents a row of data.
+                  Each element in the inner list corresponds to a column in 'columns'.
+            Example:
+            {
+                'columns': ['Region', 'Product', 'Sales'],
+                'data': [
+                    ['North', 'Widget', 150],
+                    ['South', 'Widget', 200]
+                ]
+            }
+        vis_prompt (str): A natural language prompt describing the desired visualization
+            (e.g., "Create a bar chart showing sales by region").
+
+    Returns:
+        str: The path to the saved visualization file (e.g., 'visualization_output.png')
+             or an error message if the visualization fails.
+    """
+    try:
+        # Debug prints to see what's being received
+        print("\nVisualization Tool Debug:")
+        print(f"Received json_data: {json_data}")
+        print(f"Received vis_prompt: {vis_prompt}")
+
+        # Convert JSON to DataFrame
+        df = pd.DataFrame(json_data["data"], columns=json_data["columns"])
+        print(f"Created DataFrame:\n{df.head()}")
+
+        # Initialize PandasAI
+        df_ai = pai.DataFrame(df)
+        print("Initialized PandasAI DataFrame")
+
+        load_dotenv()
+        api_key = os.environ.get("PANDAS_KEY")
+        print(f"Using PandasAI API key: {api_key[:5]}...")
+        pai.api_key.set(api_key)
+
+        # Generate visualization
+        print(f"Attempting to generate visualization with prompt: '{vis_prompt}'")
+        df_ai.chat(vis_prompt)
+
+        # Save plot
+        output_file = "visualization_output.png"
+        plt.savefig(output_file)
+        plt.close()
+        print(f"Saved visualization to {output_file}")
+
+        return f"Visualization saved as {output_file}"
+    except Exception as e:
+        print(f"Visualization error: {str(e)}")
+        print(f"Error type: {type(e)}")
+        return f"Visualization error: {str(e)}"
+
+
 if __name__ == "__main__":
     mcp.run()