Merge pull request #6 from pia-team/selcan_test

.idea/query_mcp_server.iml

@@ -1,8 +1,10 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <module type="PYTHON_MODULE" version="4">
   <component name="NewModuleRootManager">
-    <content url="file://$MODULE_DIR$" />
-    <orderEntry type="inheritedJdk" />
+    <content url="file://$MODULE_DIR$">
+      <excludeFolder url="file://$MODULE_DIR$/.idea/dataSources" />
+    </content>
+    <orderEntry type="jdk" jdkName="mcptry" jdkType="Python SDK" />
     <orderEntry type="sourceFolder" forTests="false" />
   </component>
 </module>

README.md

@@ -0,0 +1,108 @@
+# PostgreSQL NLP Query Bot
+
+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.
+
+## Features
+
+* 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
+
+
+## Tech Stack
+
+| 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               |
+
+
+## Installation
+
+
+1. **Clone the repository**
+
+   ```bash
+   git clone https://github.com/yourusername/query_mcp_server.git
+   cd query_mcp_server
+
+2. **Create and activate a conda environment**
+   
+   ```bash
+   conda create -n mcp-agent python=3.10
+   conda activate mcp-agent
+
+3. **Install dependencies**
+   ```bash
+   pip install -r requirements.txt
+
+
+Make sure the server path is correctly set in get_server_params() in your code.
+
+## Configuration
+
+
+Before running the application, make sure to set your configuration values in a `.env` file.
+
+You can start by copying the example file provided:
+
+    ```bash
+    cp .env.sample .env
+
+### .env variables
+* API_KEY: Your Google Gemini API key
+* MCP_SERVER_PATH: Path to the MCP server script (e.g. ./postgres_mcp_server.py)
+
+## Running the Project
+
+
+Once you've installed the dependencies and configured the `.env` file, you're ready to run the app.
+
+### Run the Agent
+
+   ```bash
+   python gradio_app.py
+   ```
+
+
+## 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
+```
+
+### 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

@@ -3,7 +3,7 @@ from pathlib import Path
 import gradio as gr
 import asyncio
 from postgre_mcp_client import pg_mcp_exec
-from postgre_smolagent_clinet import pg_mcp_smolagent_exec
+from postgre_smolagent_client import pg_mcp_smolagent_exec
 from langchain_mcp_client import lc_mcp_exec
 
 def load_db_configs():

langchain_mcp_client.py

@@ -25,14 +25,14 @@ async def lc_mcp_exec(request: str, history) -> tuple[Any, Any]:
     """
     try:
         history_file = os.path.join(os.path.dirname(__file__), "chat_history.json")
-        
+
         # Initialize chat history file if it doesn't exist or is empty
         if not os.path.exists(history_file) or os.path.getsize(history_file) == 0:
             with open(history_file, 'w') as f:
                 json.dump({"messages": []}, f)
 
         message_history = FileChatMessageHistory(file_path=history_file)
-        
+
         try:
             # Load existing messages or handle bootstrap scenario
             existing_messages = message_history.messages
@@ -71,16 +71,16 @@ async def lc_mcp_exec(request: str, history) -> tuple[Any, Any]:
 
                 # Create agent and prepare system message
                 agent = create_react_agent(llm, tools)
-                
+
                 # Create base messages list with system message
                 base_message = HumanMessage(content="""You are a PostgreSQL database expert assistant. 
                     Use the conversation history for context when available.""")
                 messages = [base_message]
-                
+
                 # Add history if exists
                 if formatted_history:
                     messages.extend(formatted_history)
-                
+
                 # Add current request
                 messages.append(HumanMessage(content=request))
 
@@ -134,7 +134,7 @@ async def load_and_enrich_tools(session: ClientSession):
 async def build_prompt(session, intent, request, tools, summary, chat_history):
     superset_prompt = await session.read_resource("resource://last_prompt")
     conversation_prompt = await session.read_resource("resource://base_prompt")
-    
+
     if intent == "superset_request":
         template = superset_prompt.contents[0].text
         return template.format(
@@ -143,7 +143,7 @@ async def build_prompt(session, intent, request, tools, summary, chat_history):
     else:
         template = conversation_prompt.contents[0].text
         tools_str = "\n".join([f"- {tool.name}: {tool.description}" for tool in tools])
-        
+
         # Handle history formatting with proper message access
         history_str = ""
         if chat_history:

main.py

@@ -1,38 +0,0 @@
-import yaml
-from pathlib import Path
-from typing import List
-import asyncio
-from postgre_mcp_client import pg_mcp_exec
-import logging
-
-from postgre_smolagent_clinet import pg_mcp_smolagent_exec
-
-
-#logger = logging.getLogger(__name__)
-# TODO add config
-def load_db_configs():
-    """Load database configurations from databases.yaml"""
-    configs_path = Path("configs.yaml")
-
-    if not configs_path.exists():
-        raise FileNotFoundError("configs.yaml not found")
-
-    with open(configs_path) as f:
-        configs = yaml.safe_load(f)
-
-    return configs["db_configs"]
-
-
-async def main():
-    #configs = load_db_configs()
-
-    request = "send to superset"
-    # request = "Show me the table of join posts and users tables."
-
-    # await pg_mcp_exec(request)
-    # await pg_mcp_smolagent_exec(request)
-
-
-if __name__ == "__main__":
-
-    asyncio.run(main())

prompt_temp.txt

@@ -1,58 +0,0 @@
-=========================
-# Your Role
-=========================
-
-You are an expert in generating SQL queries and interacting with a PostgreSQL database using FastMCP tools. These tools allow you to:
-
-- List available tables
-- Retrieve schema details
-- Execute SQL queries
-
-Each tool may also provide summaries of table contents to help you understand the data structure. You have access to **short-term memory**, which stores relevant information from earlier steps or previous queries. If the memory is not empty, you **must** use it when processing the current request. Avoid repeating the same tool with identical input unless the result is **not already present in memory**.
-
-=========================
-# Your Objective
-=========================
-
-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).
-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.
-
-=========================
-# Critical Rules
-=========================
-
-- Only use **SELECT** queries. Never use destructive operations (e.g., DELETE, DROP, UPDATE).
-- Always display any SQL query you generate along with the result of its execution.
-- Validate SQL syntax before execution to ensure correctness and safety.
-- Never make assumptions about the database structure — always use tools to confirm table and column names.
-- Be cautious, deliberate, and transparent in your actions.
-
-=========================
-# Short-Term Memory
-=========================
-
-You have access to the following memory from this conversation. Use it if applicable for the current request.
-
-### Conversation Context:
-- **Previous user requests**:
-  {user_requests}
-- **Tools used so far**:
-  {past_tools}
-- **Last SQL queries (if any)**:
-  {last_queries}
-- **Last result preview**:
-  {last_results}
-
-=========================
-# New User Request
-=========================
-
-Please fulfill the following request based on the above context:
-
-{new_request}

streamlit_app.py

@@ -1,38 +0,0 @@
-import yaml
-from pathlib import Path
-import streamlit as st
-from postgre_mcp_client import pg_mcp_exec
-import asyncio
-import nest_asyncio
-nest_asyncio.apply()
-
-def load_db_configs():
-    """Load database configurations from configs.yaml"""
-    configs_path = Path("configs.yaml")
-
-    if not configs_path.exists():
-        st.error("configs.yaml not found")
-        return None
-
-    with open(configs_path) as f:
-        configs = yaml.safe_load(f)
-
-    return configs.get("db_configs", {})
-
-def run_agent(message):
-    response = asyncio.run(pg_mcp_exec(message))
-    # Return in message format
-    return {"role": "assistant", "content": response}
-
-# Streamlit UI
-st.title("PostgreSQL Query Agent")
-st.write("Ask your database in natural language and get results using the smolagent executor.")
-
-user_input = st.text_input("Natural Language Request", placeholder="e.g., Show me the table of join posts and users tables.")
-
-if st.button("Run Query"):
-    if user_input.strip():
-        result = run_agent(user_input)
-        st.text_area("SQL Query / Result", value=str(result), height=300)
-    else:
-        st.warning("Please enter a natural language request.")