Enhance visualization support and update project documentation

Added example visualizations and integrated PandasAI-based chart generation into the application workflow. Updated README with detailed project structure, installation, usage, and visualization support. Included minor code refactoring and logging enhancements, along with configuration adjustments for tool chaining and debugging.

.idea/misc.xml

@@ -3,5 +3,5 @@
   <component name="Black">
     <option name="sdkName" value="pyfapi_update" />
   </component>
-  <component name="ProjectRootManager" version="2" project-jdk-name="pyfapi_update" project-jdk-type="Python SDK" />
+  <component name="ProjectRootManager" version="2" project-jdk-name="mcp-vis" project-jdk-type="Python SDK" />
 </project>

.idea/query_mcp_server.iml

@@ -4,7 +4,7 @@
     <content url="file://$MODULE_DIR$">
       <excludeFolder url="file://$MODULE_DIR$/.idea/dataSources" />
     </content>
-    <orderEntry type="jdk" jdkName="mcptry" jdkType="Python SDK" />
+    <orderEntry type="jdk" jdkName="mcp-vis" jdkType="Python SDK" />
     <orderEntry type="sourceFolder" forTests="false" />
   </component>
 </module>

README.md

@@ -1,108 +1,190 @@
-# PostgreSQL NLP Query Bot
+# Natural Language SQL Query Agent with Visualization
 
-A PostgreSQL MCP server and client system that accepts natural language queries from the user, translates them into SQL using an LLM agent (powered by LangChain), executes the query on the database, and returns the results in a human-readable format.
+A PostgreSQL-based query system that converts natural language requests into SQL queries, executes them, and provides visualizations using PandasAI. Built with LangChain, FastMCP, and Gradio.
 
-## Features
+![Architecture](resources/visualization_demo.png)
 
-* Natural Language to SQL Translation
-* Live Interaction with PostgreSQL via MCP Protocol
-* Tool Integration with Database Metadata
-* In-Session Conversational Memory
-* Supports multiple LLM clients (Langchain and Smolagents)
-* Gradio web interface for easy interaction
 
+## Description
 
-## Tech Stack
+This project combines several components to create a powerful natural language interface to PostgreSQL databases with visualization capabilities:
 
-| Layer       | Technology / Library               | Purpose                                                      |
-|-------------|------------------------------------|--------------------------------------------------------------|
-| LLM Framework | [LangChain](https://www.langchain.com/) | Manages agent behavior, memory, and tool interaction         |
-| LLM Provider | Google Gemini (via LangChain)      | Converts natural language into SQL queries                   |
-| Agent Runtime | LangGraph’s `create_react_agent`   | Runs the ReAct-style LLM agent                               |
-| Communication | Modular Command Protocol (MCP)     | Handles communication between client and PostgreSQL backend|
-| Tooling     | `langchain_mcp_adapters`           | Loads and manages tools for the agent                        |
-| Database    | PostgreSQL                         | Serves queried data                                          |
-| UI (optional) | [Gradio](https://www.gradio.app/)  | Provides a lightweight web interface for querying            |
-| Memory      | `FileChatMessageHistory`         | Stores short-term in-session memory for conversation context |
-| Language    | Python 3.10+                        | Core programming language used for development               |
+### Core Components:
 
+1. **PostgreSQL MCP Server** (`postgre_mcp_server.py`):
+   - Handles database connections and query execution
+   - Provides tools for table listing, schema retrieval
+   - Implements visualization using PandasAI
+   - Manages database lifecycle and connection pooling
 
-## Installation
+2. **LangChain Client** (`langchain_mcp_client.py`):
+   - Converts natural language to SQL using LLM
+   - Manages conversation history
+   - Integrates with MCP tools
+   - Handles agent execution flow
+
+3. **Gradio Interface** (`gradio_app.py`):
+   - Provides web-based chat interface
+   - Handles user interactions
+   - Displays query results and visualizations
 
+4. **Memory Management** (`conversation_memory.py`):
+   - Implements conversation history persistence
+   - Tracks tool usage and queries
+   - Manages session state
 
-1. **Clone the repository**
+5. **Utilities** (`utils.py`):
+   - Provides helper functions for output parsing
+   - Handles MCP response formatting
+   - Manages logging
 
+6. **Visualization** (`pandasai_visualization.py` and MCP Tools):
+   - Implements PandasAI integration for intelligent chart generation
+   - Custom MCP tool `visualize_results` that:
+     * Takes query results as JSON and a visualization prompt
+     * Uses PandasAI to automatically generate appropriate visualizations
+     * Saves charts in the `exports/charts/` directory
+   - Supports various chart types (bar charts, line plots, pie charts, etc.)
+   - Intelligent prompt-based visualization selection
+   - Includes standalone testing script:
+     * Located at `pandasai_visualization.py`
+     * Can be run directly to test PandasAI functionality
+     * Creates sample data and generates test visualizations
+     * Usage: `python pandasai_visualization.py`
+     * Helps verify PandasAI setup and API key configuration
+
+## Installation
+
+1. **Clone the Repository:**
    ```bash
-   git clone https://github.com/yourusername/query_mcp_server.git
+   git clone <repository-url>
    cd query_mcp_server
+   ```
 
-2. **Create and activate a conda environment**
-   
+2. **Create and Activate Virtual Environment:**
    ```bash
-   conda create -n mcp-agent python=3.10
-   conda activate mcp-agent
+   python -m venv venv
+   source venv/bin/activate  # On Linux/Mac
+   # or
+   .\venv\Scripts\activate  # On Windows
+   ```
 
-3. **Install dependencies**
+3. **Install Dependencies:**
    ```bash
    pip install -r requirements.txt
+   ```
 
+4. **Configure Environment Variables:**
+   Create a `.env` file in the project root with the following variables:
+   ```
+   # Database Configuration
+   DB_URL=postgresql://username:password@localhost:5432/your_database
+   DB_SCHEMA=public
+
+   # Test the PandasAI Setup (Optional)
+   # Before running the main application, you can test the visualization component:
+   python pandasai_visualization.py
+   # This will create a sample visualization using PandasAI
+
+   # API Keys
+   PANDAS_KEY=your-pandasai-key          # Required for PandasAI visualization
+   GEMINI_API_KEY=your-gemini-api-key    # For LLM query understanding
+   GEMINI_MODEL=gemini-2.0-flash-lite    # LLM model selection
+   GEMINI_MODEL_PROVIDER=google_genai    # LLM provider
+
+   # Path Configuration
+   MCP_SERVER_PATH=/absolute/path/to/postgre_mcp_server.py
+   TABLE_SUMMARY_PATH=table_summary.txt
+   ```
 
-Make sure the server path is correctly set in get_server_params() in your code.
-
-## Configuration
-
+## Running the Application
 
-Before running the application, make sure to set your configuration values in a `.env` file.
+1. **Ensure PostgreSQL Database is Running:**
+   - Make sure your PostgreSQL instance is up and accessible
+   - Verify the connection details in `.env` are correct
 
-You can start by copying the example file provided:
+2. **Start the Application:**
+   ```bash
+   # Using the run script
+   chmod +x run.sh
+   ./run.sh
 
-    ```bash
-    cp .env.sample .env
+   # Or directly with Python
+   python gradio_app.py
+   ```
 
-### .env variables
-* API_KEY: Your Google Gemini API key
-* MCP_SERVER_PATH: Path to the MCP server script (e.g. ./postgres_mcp_server.py)
+3. **Access the Interface:**
+   - Open your browser to `http://localhost:7860`
+   - The chat interface will be ready for queries
 
-## Running the Project
+## Usage
 
+### Input Examples
 
-Once you've installed the dependencies and configured the `.env` file, you're ready to run the app.
+Example prompts:
+```
+1. Simple queries:
+   "List all tables in the database"
+   "Show me the schema of table X"
 
-### Run the Agent
+2. Analysis queries:
+   "Count the number of active customers by region"
+   "Show me total sales by product category for the last month"
 
-   ```bash
-   python gradio_app.py
-   ```
+3. Visualization requests:
+   "Plot a bar chart showing sales distribution by region"
+   "Create a pie chart of customer segments"
+```
 
+### Output Structure
+
+1. **Query Results:**
+   - Text results are displayed directly in the chat interface
+   - Tabular data is formatted as markdown tables
+
+2. **Visualizations:**
+   - Generated charts are saved as PNG files in `./exports/charts/`
+   - Files are named with unique IDs: `temp_chart_{uuid}.png`
+   - Visualization is handled by the `visualize_results` MCP tool which:
+     * Automatically converts SQL results to pandas DataFrames
+     * Uses PandasAI to interpret visualization requests
+     * Generates appropriate chart types based on data and prompt
+   - Supports a wide range of visualization types:
+     * Bar charts for categorical comparisons
+     * Line plots for time series
+     * Pie charts for proportions
+     * Scatter plots for correlations
+     * And more based on data characteristics
+
+### Response Format
+
+The system provides responses in a structured format:
+```markdown
+# Result
+[Query results in table or list format]
+
+# Visualization (if requested)
+[Path to generated visualization file]
+
+# Explanation
+[Brief interpretation of results]
+
+# Query
+```sql
+[The executed SQL query]
+```
+```
 
 ## Project Structure
 
 ```
-.
-├── gradio_app.py                # Gradio UI to interact with any client (LangChain or SmolAgent)
-├── langchain_mcp_client.py      # Uses LangChain ReAct agent + ConversationBufferMemory (in-memory, also saves to memory.json)
-├── postgre_mcp_client.py        # Uses ReAct agent + custom memory class (in-session only), **WIP**: needs prompt fix
-├── postgre_smolagent_client.py  # Uses SmolAgent (lightweight), no memory support yet, **WIP**: output parsing + prompt building
-├── postgre_mcp_server.py        # MCP server that handles actual PostgreSQL communication
-├── conversation_memory.py       # Built-in memory class used by some clients
-├── chat_history.json                  # Keeps session memory. Needs to be deleted after each session
-├── table_summary.txt            # Database table descriptions used to enrich tools
-├── utils.py                     # Helper functions: intent classification, output parsing, etc.
-├── .env.sample                  # Template for your environment config
-├── requirements.txt             # Python dependencies
-└── README.md
+├── gradio_app.py              # Main application and UI
+├── postgre_mcp_server.py      # Database server and tools
+├── langchain_mcp_client.py    # LangChain integration
+├── conversation_memory.py     # Memory management
+├── utils.py                   # Helper utilities
+├── pandasai_visualization.py  # Visualization handling
+├── requirements.txt          # Project dependencies
+├── run.sh                    # Run script
+└── .env                      # Environment configuration
 ```
-
-### Work In Progress (WIP) Notes
-* postgre_mcp_client.py: 
-   - needs improvement in build_prompt() function
-
-* postgre_smolagent_client.py:
-
-   - Doesn’t support memory yet
-   
-   - Uses a different output format — not yet compatible with parse_mcp_output()
-   
-   - Also needs build_prompt() logic finalized
-* gradio_app.py: 
-   - Doesn't support conversation chat ui 

gradio_app.py

@@ -34,6 +34,17 @@ demo = gr.ChatInterface(
     type="messages",
     flagging_mode="manual",
     flagging_options=["Like", "Spam", "Inappropriate", "Other"],
+    # additional_outputs=[gr.Image(label="Visualization Output")],
+    examples=[
+        "List all tables in the database",
+        "Show me the schema of the dim_customer table",
+        "Count the number of active customers in 2024",
+        "Create a pie chart showing the distribution of customer statuses in the dim_customer table.",
+        "Plot a line chart showing the trend of order quantities over order dates from the dim_product_order_item table.",
+        "Generate a bar chart displaying the count of products in each product class from the dim_product table. Use beautiful and vivid colors fro visualization.",
+        "Visualize the relationship between tax_rate and tax_included_amount from the dim_product table using a scatter plot.",
+        "Visualize the number of product orders over time from the dim_product_order_item table, using the order_date. Show the trend monthly.",
+    ],
     save_history=True,
 )
 
@@ -56,20 +67,9 @@ demo = gr.ChatInterface(
 #     flagging_mode="never"
 # )
 
-async def build_prompt(session, request, tools, table_summary):
-    conversation_prompt = await session.read_resource("resource://base_prompt")
 
-    template = conversation_prompt.contents[0].text
-    tools_str = "\n".join([f"- {tool.name}: {tool.description}" for tool in tools])
+# TODO: maybe we can add a mcp tool to validate the results (those converted to DataFrame) to make sure the valid type is passed to the visualization tool by ReAct agent
 
-    return template.format(
-        new_request=request,
-        tools=tools_str,
-        descriptions=table_summary,
-        system_instructions="""You are a PostgreSQL database expert assistant.
-            Use the conversation history when available to maintain context.
-            For new conversations, focus on understanding the initial request."""
-    )
 
 if __name__ == "__main__":
     demo.launch()

langchain_mcp_client.py

@@ -10,6 +10,10 @@ from langchain_community.chat_message_histories import FileChatMessageHistory
 from langchain.chat_models import init_chat_model
 import logging
 from dotenv import load_dotenv
+from langchain.globals import set_debug
+
+# set_debug(True)
+
 
 # Set up logging
 logger = logging.getLogger(__name__)

pandasai_visualization.py

@@ -24,11 +24,11 @@ from dotenv import load_dotenv
 
 def create_sample_dataframe():
     """Create a sample dataframe with sales 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'],
     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)

postgre_mcp_server.py

@@ -9,10 +9,13 @@ from pydantic import Field
 import pandasai as pai
 import matplotlib as plt
 import pandas as pd
+import logging
 
 # Constants
 DEFAULT_QUERY_LIMIT = 100
 
+# logging info
+logging.basicConfig(level=logging.INFO)
 
 # Define our own PromptMessage class if the MCP one isn't available
 @dataclass
@@ -65,14 +68,14 @@ async def base_prompt_query() -> str:
 # Your Role
 ==========================
 
-You are an expert in generating and executing SQL queries and interacting with a PostgreSQL database using **FastMCP tools**. These tools allow you to:
+You are an expert in generating and executing SQL queries, interacting with a PostgreSQL database using **FastMCP tools**, and visualizing results when requested. These tools allow you to:
 
 - List available tables
 - Retrieve schema details
 - Execute SQL queries
+- Visualize query results using PandasAI
 
-
-Each tool may also return previews or summaries of table contents to help you better understand the data structure.
+Each tool may return previews or summaries of table contents to help you understand the data structure.
 
 ---
 
@@ -82,28 +85,15 @@ Each tool may also return previews or summaries of table contents to help you be
 
 When a user submits a request, you must:
 
-1. **Analyze the request** to determine the required data or action.
-2. **Use FastMCP tools** to gather any necessary information (e.g., list tables or retrieve schema).
+1. **Analyze the request** to determine the required data, action, or visualization.
+2. **Use FastMCP tools** to gather necessary information (e.g., list tables, retrieve schema).
 3. **Generate a valid SQL SELECT query**, if needed, and clearly show the full query.
-4. **Execute the SQL query** and return the results.
-5. **Chain tools logically**, such as: List Tables → Get Schema → Write and Run Query.
-6. **Explain your reasoning and each step taken** to ensure clarity and transparency.
-
----
-
-==========================
-# Your Objective
-==========================
-
-When a user submits a request, follow these steps:
-
-1. Analyze the request to determine the desired data or action.
-2. Use tools to gather any necessary information (e.g., list tables, get schema).
-3. Generate a valid SQL query (such as **SELECT**, **COUNT**, or other read-only operations) and clearly display the full query.
-4. Execute the query and **return the execution result of the query**.
-5. **Chain tools logically to build toward the answer.**
-6. Explain your reasoning at every step for clarity and transparency.
-7. Show the result of the **execute_query** in your final answer. 
+4. **Execute the SQL query** using the `execute_query` tool and return the results.
+5. **Visualize results** using the `visualize_results` tool if the user explicitly requests a visualization (e.g., "create a chart", "visualize", "plot"). For visualizations:
+   - Craft a visualization prompt starting with "plot" (e.g., "plot a bar chart showing sales by region").
+   - Send JSON data in the format `{{'columns': [list of column names], 'data': [list of rows, each row a list of values]}}`.
+6. **Chain tools logically**, such as: List Tables → Get Schema → Write and Run Query → Visualize Results (if requested).
+7. **Explain your reasoning and each step taken** to ensure clarity and transparency.
 
 ---
 
@@ -117,7 +107,11 @@ When a user submits a request, follow these steps:
 - Validate SQL syntax before execution.
 - Never assume table or column names. Use tools to confirm structure.
 - Use memory efficiently. Don’t rerun a tool unless necessary.
-- If you generate a SQL query, immediately call the **execute_query** tool using that query. Do not delay or wait for user confirmation.
+- If you generate a SQL query, immediately call the **execute_query** tool.
+- If the user requests a visualization, call the **visualize_results** tool with:
+  - A visualization prompt starting with "plot" (e.g., "plot a bar chart showing sales by region").
+  - JSON data in the format `{{'columns': [list of column names], 'data': [list of rows, each row a list of values]}}`.
+- For non-query or non-visualization requests (e.g., history questions), respond appropriately without forcing a query or visualization.
 
 ---
 
@@ -133,13 +127,12 @@ When a user submits a request, follow these steps:
 # Tools
 ==========================
 
-You can use the following FastMCP tools. These allow you to create **read-only** queries, such as `SELECT`, `COUNT`, or queries with `GROUP BY`, `ORDER BY`, and similar clauses. You may chain tools together to gather the necessary information before generating your SQL query.
+You can use the following FastMCP tools to create **read-only** queries (e.g., `SELECT`, `COUNT`, `GROUP BY`, `ORDER BY`) and visualize results when requested. Chain tools logically to gather information, execute queries, or generate visualizations.
 
 {tools}
 
 ---
 
-
 ### Invalid Example — DELETE Operation (Not Allowed):
 **User Request:** "Delete all customers from Germany."
 
@@ -161,21 +154,26 @@ Present your final answer using the following structure **exactly** in markdown
 # Result
 {{Take the result from the execute_query tool and format it nicely using Markdown. Use a Markdown table for tabular data (rows and columns) including headers. Use bullet points or items in markdown for answers that include lists of names or descriptions. Use plain text for single values or simple messages. Ensure data alignment and clarity.}}
 
+# Visualization (if requested)
+{{If the user requested a visualization, include the result from the visualize_results tool, e.g., "Visualization saved as visualization_output.png". Otherwise, omit this section.}}
+
 # Explanation
-{{Provide a concise explanation or interpretation of the results in 1-3 sentences. Explain what the data in the 'Result' section represents in the context of the user's request.}}
+{{Provide a concise explanation or interpretation of the results (and visualization, if applicable) in 1-3 sentences. Explain what the data and visualization (if any) represent in the context of the user's request.}}
 
 # Query
 ```sql
 {{Display the exact SQL query you generated and executed here to answer the user's request.}}
 ```
 
-```
-
-**Reminder:** 
-**Every time you generate a SQL query, call **execute_query** right after and include the result in your final response.**
-**If you do not execute the generated SQL query, this will be the violation of the instructions**
-**Your final answer cannot be only a SQL query, you will have to call **execute_query** and give the result of the call with the SQL query.** 
+**Reminder:**
+- **Every time you generate a SQL query, call `execute_query` immediately and include the result.**
+- **If the user requests a visualization (e.g., "create a chart", "visualize", "plot"), call `visualize_results` with:**
+  - A visualization prompt starting with "plot" (e.g., "plot a bar chart showing sales by region").
+  - JSON data in the format `{{'columns': [list of column names], 'data': [list of rows, each row a list of values]}}`.
+- **For non-query or non-visualization requests, respond appropriately without forcing a query or visualization.**
 
+**Conversation History:**
+Use the conversation history for context when available to maintain continuity.
 """
 
     return base_prompt
@@ -201,6 +199,7 @@ async def execute_query(
 ) -> str:
     """Execute a read-only SQL query against the database"""
     # Validate query - simple check for read-only
+
     query = query.strip()
     if not query.lower().startswith("select"):
         return "Error: Only SELECT queries are allowed for security reasons."
@@ -222,6 +221,16 @@ async def execute_query(
             rows = [" | ".join(str(val) for val in row.values())
                     for row in result[:limit if limit else DEFAULT_QUERY_LIMIT]]
 
+            # print(f"{header}\n{separator}\n" + "\n".join(rows))
+
+            # print(f"===== Header Type: ======\n {type(header)}")
+            # print(f"===== Row Type: ======\n {type(rows)}")
+            #
+            # # print the  data itself
+            # print(f"===== Header Data: ======\n {header}")
+            # print(f"===== Row Data: ======\n {rows}")
+
+
             return f"{header}\n{separator}\n" + "\n".join(rows)
     except asyncpg.exceptions.PostgresError as e:
         return f"SQL Error: {str(e)}"