init: clean up unnecessary components for demo

.env.sample

@@ -1,2 +1,12 @@
-API_KEY=
-MCP_SERVER_PATH=/home/andre/andre/postgre_mcp_server.py
+MCP_SERVER_PATH=
+TABLE_SUMMARY_PATH=
+
+DB_URL=
+
+GEMINI_API_KEY=
+GEMINI_MODEL=
+
+PANDAS_KEY=
+
+OPENAI_MODEL=
+OPENAI_API_KEY=

.gitignore

@@ -122,4 +122,5 @@ dmypy.json
 .pytype/
 
 # Cython debug symbols
-cython_debug/
+cython_debug/
+/.gradio/

base_chat_history.json

@@ -0,0 +1 @@
+[{"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "# Result\nThe tables in the database are:\n*   dim\\_agreement\n*   dim\\_customer\n*   dim\\_product\n*   dim\\_product\\_order\\_item\n\n# Explanation\nThe `list_tables` tool was called to retrieve a list of all available tables in the database schema. The result shows the names of these tables.\n\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}]

chat_history.json

@@ -1 +1 @@
-[{"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "# Result\nThe tables in the database are:\n*   dim\\_agreement\n*   dim\\_customer\n*   dim\\_product\n*   dim\\_product\\_order\\_item\n\n# Explanation\nThe `list_tables` tool was called to retrieve a list of all available tables in the database schema. The result shows the names of these tables.\n\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "what was my previous request ?", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "# Result\nlist all tables\n# Explanation\nThe user's previous request was to list all the tables in the database. This was determined by reviewing the previous turn in the conversation history.\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "count tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\n4\n# Explanation\nThe user requested the number of tables. The previous response listed the tables, and counting them gives the answer.\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "what was my last message ?", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\ncount tables\n# Explanation\nThe user's last message was \"count tables\". This was determined by reviewing the previous turn in the conversation history.\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "what was my very first message in this conversation?", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\nlist all tables\n# Explanation\nThe user's very first message in this conversation was \"list all tables\". This was determined by reviewing the conversation history.\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "total number of sales in 2024", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\nThe tables in the database are:\n*   dim\\_agreement\n*   dim\\_customer\n*   dim\\_product\n*   dim\\_product\\_order\\_item\n\n# Explanation\nThe `list_tables` tool was called to retrieve a list of all available tables in the database schema. The result shows the names of these tables.\n\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\nThe tables in the database are:\n*   dim_agreement\n*   dim_customer\n*   dim_product\n*   dim_product_order_item\n\n# Explanation\nThe `list_tables` tool was called to retrieve a list of all available tables in the database schema. The result shows the names of these tables.\n\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}]
+[{"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "# Result\nThe tables in the database are:\n*   dim\\_agreement\n*   dim\\_customer\n*   dim\\_product\n*   dim\\_product\\_order\\_item\n\n# Explanation\nThe `list_tables` tool was called to retrieve a list of all available tables in the database schema. The result shows the names of these tables.\n\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "what was my previous request ?", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "# Result\nlist all tables\n# Explanation\nThe user's previous request was to list all the tables in the database. This was determined by reviewing the previous turn in the conversation history.\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "count tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\n4\n# Explanation\nThe user requested the number of tables. The previous response listed the tables, and counting them gives the answer.\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "what was my last message ?", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\ncount tables\n# Explanation\nThe user's last message was \"count tables\". This was determined by reviewing the previous turn in the conversation history.\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "what was my very first message in this conversation?", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\nlist all tables\n# Explanation\nThe user's very first message in this conversation was \"list all tables\". This was determined by reviewing the conversation history.\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "total number of sales in 2024", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\nThe tables in the database are:\n*   dim\\_agreement\n*   dim\\_customer\n*   dim\\_product\n*   dim\\_product\\_order\\_item\n\n# Explanation\nThe `list_tables` tool was called to retrieve a list of all available tables in the database schema. The result shows the names of these tables.\n\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "list all tables", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "```markdown\n# Result\nThe tables in the database are:\n*   dim_agreement\n*   dim_customer\n*   dim_product\n*   dim_product_order_item\n\n# Explanation\nThe `list_tables` tool was called to retrieve a list of all available tables in the database schema. The result shows the names of these tables.\n\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "hi there. show me the largest table name", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "I am sorry, I cannot directly determine the \"largest table name\" based on the available tools. I can list the tables and show their schemas, but I do not have a tool to determine table sizes. I would need a tool to get table sizes to fulfill this request. However, I can list the tables for you.\n\n```markdown\n# Result\nThe tables in the database are:\n*   dim\\_agreement\n*   dim\\_customer\n*   dim\\_product\n*   dim\\_product\\_order\\_item\n\n# Explanation\nThe `list_tables` tool was called to retrieve a list of all available tables in the database schema. The result shows the names of these tables.\n\n# Query\n```sql\nN/A\n```", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "you are free to execute any sql query and get the size of each table and return me the count of largest one", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "I am sorry, I cannot directly determine the \"largest table\" based on the available tools. I can list the tables and show their schemas, but I do not have a tool to determine table sizes. I would need a tool to get table sizes to fulfill this request. I can list the tables for you, but I cannot provide the count of the largest one.", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}]

gradio_app.py

@@ -2,10 +2,9 @@ import yaml
 from pathlib import Path
 import gradio as gr
 import asyncio
-from postgre_mcp_client import pg_mcp_exec
-from postgre_smolagent_client import pg_mcp_smolagent_exec
 from langchain_mcp_client import lc_mcp_exec
 
+# ======================================= Load DB configs
 def load_db_configs():
     """Load database configurations from configs.yaml"""
     configs_path = Path("configs.yaml")
@@ -18,7 +17,8 @@ def load_db_configs():
 
     return configs["db_configs"]
 
-# Async-compatible wrapper
+
+# ====================================== Async-compatible wrapper
 async def run_agent(request, history):
     # configs = load_db_configs()
     # final_answer, last_tool_answer, = await pg_mcp_exec(request)
@@ -27,6 +27,8 @@ async def run_agent(request, history):
     response, message = await lc_mcp_exec(request, history)
     return response
 
+
+# ====================================== Gradio UI with history
 demo = gr.ChatInterface(
     run_agent,
     type="messages",
@@ -35,6 +37,7 @@ demo = gr.ChatInterface(
     save_history=True,
 )
 
+
 # Gradio UI
 # demo = gr.Interface(
 #     fn=run_agent,
@@ -53,5 +56,6 @@ demo = gr.ChatInterface(
 #     flagging_mode="never"
 # )
 
+
 if __name__ == "__main__":
     demo.launch()

langchain_mcp_client.py

@@ -1,7 +1,6 @@
 import os.path
 import json
 from typing import Tuple, Any
-
 from mcp import ClientSession, StdioServerParameters
 from mcp.client.stdio import stdio_client
 from langchain_mcp_adapters.tools import load_mcp_tools
@@ -16,9 +15,12 @@ import logging
 from dotenv import load_dotenv
 
 
-
+# Set up logging
 logger = logging.getLogger(__name__)
+
+# Load environment variables from .env file
 load_dotenv()
+
 async def lc_mcp_exec(request: str, history) -> tuple[Any, Any]:
     """
     Execute the full PostgreSQL MCP pipeline with persistent memory.
@@ -51,20 +53,21 @@ async def lc_mcp_exec(request: str, history) -> tuple[Any, Any]:
             elif isinstance(msg, AIMessage):
                 formatted_history.append(AIMessage(content=msg.content))
 
-        # TODO: give summary file path from config
-        table_summary = load_table_summary("table_summary.txt")
+        table_summary = load_table_summary(os.environ["TABLE_SUMMARY_PATH"])
         server_params = get_server_params()
 
-        api_key = os.getenv("API_KEY")
-        llm = init_chat_model(model="gemini-2.0-flash", model_provider="google_genai",
-                              api_key=api_key)
+        llm = init_chat_model(
+            model=os.getenv("GEMINI_MODEL"),
+            model_provider=os.getenv("GEMINI_MODEL_PROVIDER"),
+            api_key=os.getenv("GEMINI_API_KEY"),
+        )
 
+        # Initialize the MCP client
         async with stdio_client(server_params) as (read, write):
             async with ClientSession(read, write) as session:
                 await session.initialize()
 
                 tools = await load_and_enrich_tools(session)
-                intent = classify_intent(request)
 
                 # Add new user message before processing
                 message_history.add_user_message(request)
@@ -85,7 +88,13 @@ async def lc_mcp_exec(request: str, history) -> tuple[Any, Any]:
                 messages.append(HumanMessage(content=request))
 
                 # Build prompt with conversation context
-                prompt = await build_prompt(session, intent, request, tools, table_summary, formatted_history)
+                prompt = await build_prompt(
+                    session=session,
+                    request=request,
+                    tools=tools,
+                    table_summary=table_summary,
+                    chat_history=formatted_history
+                )
 
                 # Invoke agent with proper message structure
                 agent_response = await agent.ainvoke(
@@ -119,6 +128,7 @@ def load_table_summary(path: str) -> str:
     with open(path, 'r') as file:
         return file.read()
 
+
 def get_server_params() -> StdioServerParameters:
     # TODO: give server params from config
     load_dotenv()
@@ -127,41 +137,36 @@ def get_server_params() -> StdioServerParameters:
         args=[os.environ["MCP_SERVER_PATH"]],
     )
 
+
 async def load_and_enrich_tools(session: ClientSession):
     tools = await load_mcp_tools(session)
     return tools
 
-async def build_prompt(session, intent, request, tools, summary, chat_history):
-    superset_prompt = await session.read_resource("resource://last_prompt")
+
+async def build_prompt(session, request, tools, table_summary, chat_history):
     conversation_prompt = await session.read_resource("resource://base_prompt")
 
-    if intent == "superset_request":
-        template = superset_prompt.contents[0].text
-        return template.format(
-            new_request=request
-        )
-    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:
-            history_sections = []
-            for msg in chat_history:
-                if isinstance(msg, HumanMessage):
-                    history_sections.append(f"Previous Human Question:\n{msg.content}\n")
-                elif isinstance(msg, AIMessage):
-                    history_sections.append(f"Previous Assistant Response:\n{msg.content}\n")
-            history_str = "\n".join(history_sections)
-
-        return template.format(
-            new_request=request,
-            tools=tools_str,
-            descriptions=summary,
-            chat_history=f"\nPrevious Conversation History:\n{history_str}" if history_str else "\nThis is a new conversation.",
-            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."""
-        )
+    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:
+        history_sections = []
+        for msg in chat_history:
+            if isinstance(msg, HumanMessage):
+                history_sections.append(f"Previous Human Question:\n{msg.content}\n")
+            elif isinstance(msg, AIMessage):
+                history_sections.append(f"Previous Assistant Response:\n{msg.content}\n")
+        history_str = "\n".join(history_sections)
+
+    return template.format(
+        new_request=request,
+        tools=tools_str,
+        descriptions=table_summary,
+        chat_history=f"\nPrevious Conversation History:\n{history_str}" if history_str else "\nThis is a new conversation.",
+        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."""
+    )
 

postgre_mcp_client.py

@@ -1,128 +0,0 @@
-import os.path
-from typing import Tuple, Any
-
-from mcp import ClientSession, StdioServerParameters
-from mcp.client.stdio import stdio_client
-from langchain_mcp_adapters.tools import load_mcp_tools
-from langgraph.prebuilt import create_react_agent
-from langchain.chat_models import init_chat_model
-from conversation_memory import ConversationMemory
-from utils import parse_mcp_output, classify_intent
-import logging
-
-logger = logging.getLogger(__name__)
-
-async def pg_mcp_exec(request: str) -> tuple[Any, Any]:
-    """
-    Execute the full PostgreSQL MCP pipeline: load summary, connect session,
-    load memory and tools, build prompt, run agent, update memory.
-
-    Args:
-        request (str): User's request input.
-        llm (Any): Language model for reasoning agent.
-
-    Returns:
-        str: Agent response message.
-    """
-    # TODO: give summary file path from config
-    table_summary = load_table_summary("table_summary.txt")
-    server_params = get_server_params()
-
-    # TODO: give key from env
-    llm = init_chat_model(model="gemini-2.0-flash", model_provider="google_genai",
-                          api_key="AIzaSyAuxYmci0DVU5l5L_YcxLlxHzR5MLn70js")
-
-    async with stdio_client(server_params) as (read, write):
-        async with ClientSession(read, write) as session:
-            await session.initialize()
-            memory = await load_or_create_memory()
-
-            tools = await load_and_enrich_tools(session, table_summary)
-            past_data = get_memory_snapshot(memory)
-
-            intent = classify_intent(request)
-            prompt = await build_prompt(session, intent, request, tools, past_data)
-
-            agent = create_react_agent(llm, tools)
-            agent_response = await agent.ainvoke({"messages": prompt})
-
-            parsed_steps, final_answer, last_tool_answer, _ = parse_mcp_output(agent_response)
-            memory.update_from_parsed(parsed_steps, request)
-
-            await handle_memory_save_or_reset(memory, request)
-
-            return  final_answer, last_tool_answer
-
-
-# ---------------- Helper Functions ---------------- #
-
-def load_table_summary(path: str) -> str:
-    with open(path, 'r') as file:
-        return file.read()
-
-def get_server_params() -> StdioServerParameters:
-    # TODO: give server params from config
-    return StdioServerParameters(
-        command="python",
-        args=[r"/home/amirkia/Desktop/query_mcp_server/postgre_mcp_server.py"],
-    )
-
-async def load_or_create_memory() -> ConversationMemory:
-    memory = ConversationMemory()
-    if os.path.exists("memory.json"):
-        return memory.load_memory()
-    return memory
-
-async def load_and_enrich_tools(session: ClientSession, summary: str):
-    tools = await load_mcp_tools(session)
-    for tool in tools:
-        tool.description += f" {summary}"
-    return tools
-
-def get_memory_snapshot(memory: ConversationMemory) -> dict:
-    if os.path.exists("memory.json"):
-        return {
-            "past_tools": memory.get_all_tools_used(),
-            "past_queries": memory.get_last_n_queries(),
-            "past_results": memory.get_last_n_results(),
-            "past_requests": memory.get_all_user_messages()
-        }
-    return {
-        "past_tools": "No tools found",
-        "past_queries": "No queries found",
-        "past_results": "No results found",
-        "past_requests": "No requests found"
-    }
-
-async def build_prompt(session, intent, request, tools, past_data):
-    superset_prompt = await session.read_resource("resource://last_prompt")
-    conversation_prompt = await session.read_resource("resource://base_prompt")
-    # TODO: add uri's from config
-    if intent == "superset_request":
-        template = superset_prompt.contents[0].text
-        return template.format(
-            user_requests=past_data["past_requests"],
-            past_tools=past_data["past_tools"],
-            last_queries=past_data["past_queries"],
-            last_results=past_data["past_results"],
-            new_request=request
-        )
-    else:
-        template = conversation_prompt.contents[0].text
-        tools_str = "\n".join([f"- {tool.name}: {tool.description}" for tool in tools])
-        return template.format(
-            user_requests=past_data["past_requests"],
-            past_tools=past_data["past_tools"],
-            last_queries=past_data["past_queries"],
-            last_results=past_data["past_results"],
-            new_request=request,
-            tools=tools_str
-        )
-
-async def handle_memory_save_or_reset(memory: ConversationMemory, request: str):
-    if request.strip().lower() == "stop":
-        memory.reset()
-        logger.info("Conversation memory reset.")
-    else:
-        memory.save_memory()
-        logger.info("Conversation memory saved.")

postgre_mcp_server.py

@@ -4,12 +4,12 @@ from dataclasses import dataclass
 from typing import Optional, AsyncIterator
 
 import asyncpg
+from flask.cli import load_dotenv
 from mcp.server.fastmcp import FastMCP, Context
 from pydantic import Field
 
 # Constants
 DEFAULT_QUERY_LIMIT = 100
-DEFAULT_SCHEMA = "public"
 
 
 # Define our own PromptMessage class if the MCP one isn't available
@@ -31,8 +31,9 @@ class DbContext:
 async def db_lifespan(server: FastMCP) -> AsyncIterator[DbContext]:
     """Manage database connection lifecycle"""
     # Initialize DB connection from environment variables
-    dsn = os.environ.get("DSN", "postgresql://postgres:postgres@localhost:5432/DNextReporting")
-    schema = os.environ.get("SCHEMA", DEFAULT_SCHEMA)
+    load_dotenv()
+    dsn = os.getenv("DB_URL")
+    schema = os.getenv("DB_SCHEMA")
     pool = await asyncpg.create_pool(dsn)
     try:
         yield DbContext(pool=pool, schema=schema)
@@ -58,149 +59,33 @@ async def base_prompt_query() -> str:
     """Returns a base prompt to generate sql queries and answer questions"""
 
     base_prompt = """
-
-                    ==========================
-                    # 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:
-
-                    - List available tables
-                    - Retrieve schema details
-                    - Execute SQL queries
-
-
-                    Each tool may also return previews or summaries of table contents to help you better understand the data structure.
-                    
-                    ---
-
-                    ==========================
-                    # 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.
-
-                    ---
-
-                    ==========================
-                    # 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. 
-
-                    ---
-
-                    ==========================
-                    # Critical Rules
-                    ==========================
-
-                    - Only use **read-only** SQL queries such as **SELECT**, **COUNT**, or queries with **GROUP BY**, **ORDER BY**, etc.
-                    - **Never** use destructive operations like **DELETE**, **UPDATE**, **INSERT**, or **DROP**.
-                    - Always show the SQL query you generate along with the execution result.
-                    - 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.
-
-                    ---
-                    
-                    ==========================
-                    # Database Description
-                    ==========================
-                    
-                    {descriptions}
-                    
-                    ---
-
-                    ==========================
-                    # 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.
-
-                    {tools}
-
-                    ---
-
-
-                    ### Invalid Example — DELETE Operation (Not Allowed):
-                    **User Request:** "Delete all customers from Germany."
-
-                    **Response Guidance:**
-                    - **Do not generate or execute** destructive queries such as `DELETE`.
-                    - Instead, respond with a message like:
-                      > Destructive operations such as `DELETE` are not permitted. I can help you retrieve the customers from Germany using a `SELECT` query instead:
-                      > ```sql
-                      > SELECT * FROM customers WHERE country = 'Germany';
-                      > ```
-
 ==========================
-# Output Format
+# Your Role
 ==========================
 
-Present your final answer using the following structure **exactly**. When necessary, bold the important parts of your answer or use `` inline code blocks.:
-
-```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.}}
-
-# 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.}}
-
-# Query
-```sql
-{{Display the exact SQL query you generated and executed here to answer the user's request.}}
+You are an expert in generating and executing SQL queries and interacting with a PostgreSQL database using **FastMCP tools**. These tools allow you to:
 
-                    **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.** 
-                    ---
-                    {chat_history}
-                    ---
-                    =========================
-                    # New User Request
-                    =========================
+- List available tables
+- Retrieve schema details
+- Execute SQL queries
 
-                    Please fulfill the following request based on the above context:
 
-                    {new_request}
-                    """
+Each tool may also return previews or summaries of table contents to help you better understand the data structure.
 
-    return base_prompt
-
-
-@mcp.resource(
-    uri="resource://last_prompt",
-    name="last_prompt",
-    description="A prompt that identifies the most recent SQL query related to the user's request and reformats it into ANSI SQL syntax for use in Superset."
-)
-async def last_prompt() -> str:
-    """A prompt that identifies the most recent SQL query related to the user's request and reformats it into ANSI SQL syntax for use in Superset."""
-
-    base_prompt = """
+---
 
 ==========================
-# Your Role
+# Your Objective
 ==========================
 
-You are an expert at reading and understanding SQL queries.  
-Your task is to retrieve the **exact SQL query** that produced a previously seen result, convert the query to the **ANSI SQL query** and return **only the ANSI SQL query** — no explanation, reasoning, or commentary.
+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.
 
 ---
 
@@ -208,12 +93,15 @@ Your task is to retrieve the **exact SQL query** that produced a previously seen
 # Your Objective
 ==========================
 
-When a user submits a request (e.g., *"send me that table"*, *"send the last query"*, etc.), follow these steps:
+When a user submits a request, follow these steps:
 
-1. Identify which previous result the user is referring to.
-2. Retrieve the corresponding SQL query that produced that result.
-3. Convert the SQL query to the ANSI SQL query
-3. Return **only** that ANSI SQL query.
+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. 
 
 ---
 
@@ -221,220 +109,78 @@ When a user submits a request (e.g., *"send me that table"*, *"send the last que
 # Critical Rules
 ==========================
 
-- Do **not** ask questions or request clarification.
-- Do **not** explain anything to the user.
-- Respond with the **exact ANSI SQL query only**, formatted cleanly.
-- Do **not** guess — only retrieve queries.
-- If no query fits, respond with: "Query not found."
+- Only use **read-only** SQL queries such as **SELECT**, **COUNT**, or queries with **GROUP BY**, **ORDER BY**, etc.
+- **Never** use destructive operations like **DELETE**, **UPDATE**, **INSERT**, or **DROP**.
+- Always show the SQL query you generate along with the execution result.
+- 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.
 
 ---
 
 ==========================
-# Reference Conversion Rules for PostgreSQL to ANSI SQL
+# Database Description
 ==========================
-## Use the following rules as examples and guidance. You are not limited to these rules — apply your own knowledge of SQL standards and PostgreSQL to perform accurate conversions, even in cases not explicitly covered below.
-
-### 1. String Concatenation
-- **PostgreSQL:** || for string concatenation
-- **ANSI SQL:** CONCAT() function
-
-**Example:**
-```sql
--- PostgreSQL
-SELECT first_name || ' ' || last_name AS full_name FROM users;
-
--- ANSI SQL
-SELECT CONCAT(first_name, ' ', last_name) AS full_name FROM users;
-
-### 2. LIMIT Clause
-
-- **PostgreSQL:** Uses `LIMIT`  
-- **ANSI SQL:** Uses `FETCH FIRST n ROWS ONLY`
-
-**Example:**
-```sql
--- PostgreSQL
-SELECT * FROM posts LIMIT 10;
-
--- ANSI SQL
-SELECT * FROM posts FETCH FIRST 10 ROWS ONLY;
-
-### 3. ILIKE (Case-insensitive LIKE)
-
-- **PostgreSQL:** `ILIKE`  
-- **ANSI SQL:** Use `LOWER(column) LIKE` or `UPPER(column) LIKE`
 
-**Example:**
-```sql
--- PostgreSQL
-SELECT * FROM users WHERE username ILIKE 'john%';
-
--- ANSI SQL
-SELECT * FROM users WHERE LOWER(username) LIKE 'john%';
-
-### 4. Type Casting
-
-- **PostgreSQL:** `'123'::INTEGER`  
-- **ANSI SQL:** `CAST('123' AS INTEGER)`
-
-**Example:**
-```sql
--- PostgreSQL
-SELECT '123'::INTEGER;
-
--- ANSI SQL
-SELECT CAST('123' AS INTEGER);
-
-### 5. INSERT ... RETURNING
-
-- **PostgreSQL:** Supports `RETURNING` to get values immediately  
-- **ANSI SQL:** No equivalent; use a separate `SELECT` query
-
-**Example:**
-```sql
--- PostgreSQL
-INSERT INTO users (username) VALUES ('john') RETURNING id;
-
--- ANSI SQL
--- No direct equivalent; you must SELECT after INSERT
-INSERT INTO users (username) VALUES ('john');
-SELECT id FROM users WHERE username = 'john';
-
-
-### 6. UPSERT (INSERT ON CONFLICT)
-
-- **PostgreSQL:** `INSERT ... ON CONFLICT`  
-- **ANSI SQL:** Use `MERGE` (available in ANSI SQL-2003+)
-
-**Example:**
-```sql
--- PostgreSQL
-INSERT INTO users (id, username)
-VALUES (1, 'john')
-ON CONFLICT (id)
-DO UPDATE SET username = EXCLUDED.username;
-
--- ANSI SQL
-MERGE INTO users AS target
-USING (SELECT 1 AS id, 'john' AS username) AS source
-ON (target.id = source.id)
-WHEN MATCHED THEN
-  UPDATE SET username = source.username
-WHEN NOT MATCHED THEN
-  INSERT (id, username) VALUES (source.id, source.username);
-
-### 7. SERIAL (Auto-increment IDs)
-
-- **PostgreSQL:** `SERIAL` keyword  
-- **ANSI SQL:** Use `IDENTITY` in table creation
-
-**Example:**
-```sql
--- PostgreSQL
-CREATE TABLE users (
-  id SERIAL PRIMARY KEY,
-  username TEXT
-);
-
--- ANSI SQL
-CREATE TABLE users (
-  id INTEGER GENERATED ALWAYS AS IDENTITY PRIMARY KEY,
-  username VARCHAR(255)
-);
-
-
-### 8. Arrays and Accessing Arrays
-
-- **PostgreSQL:** Supports arrays and uses `[index]`  
-- **ANSI SQL:** No array type; use normalized relational design
-
-**Example:**
-```sql
--- PostgreSQL
-SELECT tags[1] FROM posts;
-
--- ANSI SQL
--- No equivalent. Use a separate `post_tags` table instead.
-
-### 9. PostgreSQL Extensions and Functions
-
-- Replace PostgreSQL-specific functions like `NOW()`, `CURRENT_DATE`, etc., with ANSI SQL equivalents when available.
-
-**Example:**
-```sql
--- PostgreSQL
-SELECT NOW(), CURRENT_DATE;
-
--- ANSI SQL
-SELECT CURRENT_TIMESTAMP, CURRENT_DATE;
+{descriptions}
 
-### 10. Table Aliases with 'AS'
-
-- **PostgreSQL:** `FROM users AS u`  
-- **ANSI SQL:** `FROM users u` (without `AS` is also valid)
-
-**Note:** Both are valid ANSI SQL — no change needed.
-
-### 11. Boolean Type
-
-- **PostgreSQL:** Uses `BOOLEAN`, and accepts `TRUE` / `FALSE`  
-- **ANSI SQL:** Also supports `BOOLEAN`, but in some systems may need conversion to `CHAR(1)` with `'Y'` / `'N'`
-
-**Example (if needed):**
-```sql
--- PostgreSQL
-SELECT * FROM flags WHERE is_enabled = TRUE;
+---
 
--- ANSI SQL (fallback style)
-SELECT * FROM flags WHERE is_enabled = 'Y';
+==========================
+# 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.
 
-### 12. Text vs. VARCHAR
+{tools}
 
-- **PostgreSQL:** Allows unbounded `TEXT`  
-- **ANSI SQL:** Use `VARCHAR(n)`
+---
 
-**Example:**
-```sql
--- PostgreSQL
-CREATE TABLE posts (content TEXT);
 
--- ANSI SQL
-CREATE TABLE posts (content VARCHAR(65535));
+### Invalid Example — DELETE Operation (Not Allowed):
+**User Request:** "Delete all customers from Germany."
 
----
+**Response Guidance:**
+- **Do not generate or execute** destructive queries such as `DELETE`.
+- Instead, respond with a message like:
+  > Destructive operations such as `DELETE` are not permitted. I can help you retrieve the customers from Germany using a `SELECT` query instead:
+  > ```sql
+  > SELECT * FROM customers WHERE country = 'Germany';
+  > ```
 
 ==========================
-# Examples
+# Output Format
 ==========================
 
-### Example 1 — Referring last query (Check the memory and find the most recent query that generates a table) :
-**User Request:** "send the last table"
-**You return:**
-```sql
-SELECT * FROM posts INNER JOIN users ON posts.user_id = users.id;
+Present your final answer using the following structure **exactly**. When necessary, bold the important parts of your answer or use `` inline code blocks.:
 
-### Example 2 — Referring to a specific query (check the memory and find the query that returns the count of users):
-**User Request:** "send the query that gave us the count of users"
-**You return:**
-```sql
-SELECT COUNT(*) FROM users;
+```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.}}
 
-### Example 3 — Referring latest known query (Check the memory and find the most recent query.)
-**User Request:** "send"
-**You return:**
-(latest known query):
-```sql
-SELECT * FROM posts WHERE user_id = 1;
+# 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.}}
 
-Remember: Only respond with valid SQL from memory converted to the ANSI SQL. No assumptions. No explanations.
+# 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.** 
+---
+{chat_history}
+---
 =========================
 # New User Request
 =========================
-{new_request}
 
-                        """
+Please fulfill the following request based on the above context:
+
+{new_request}
+"""
 
     return base_prompt
 
@@ -565,7 +311,7 @@ async def list_tables() -> str:
 async def get_table_schema(table_name: str) -> str:
     """Get schema information for a specific table"""
     try:
-        schema = os.environ.get("SCHEMA", DEFAULT_SCHEMA)
+        schema = os.getenv("DB_SCHEMA")
 
         async with db_lifespan(mcp) as db_ctx:
             columns = await get_table_schema_info(db_ctx.pool, schema, table_name)
@@ -588,7 +334,7 @@ def get_foreign_keys(table_name: str) -> str:
         table_name: The name of the table to get foreign keys from
         schema: The schema name (defaults to 'public')
     """
-    schema = os.environ.get("SCHEMA", DEFAULT_SCHEMA)
+    schema = os.getenv("DB_SCHEMA")
 
     sql = """
     SELECT 
@@ -617,7 +363,7 @@ def get_foreign_keys(table_name: str) -> str:
 async def get_all_schemas() -> str:
     """Get schema information for all tables in the database"""
     try:
-        schema = os.environ.get("SCHEMA", DEFAULT_SCHEMA)
+        schema = os.getenv("DB_SCHEMA")
 
         async with db_lifespan(mcp) as db_ctx:
             tables = await get_all_tables(db_ctx.pool, db_ctx.schema)

postgre_smolagent_client.py

@@ -1,133 +0,0 @@
-import os.path
-from mcp import ClientSession, StdioServerParameters
-from mcp.client.stdio import stdio_client
-from langchain_mcp_adapters.tools import load_mcp_tools
-from langgraph.prebuilt import create_react_agent
-from langchain.chat_models import init_chat_model
-from conversation_memory import ConversationMemory
-from utils import parse_mcp_output, classify_intent
-import logging
-from smolagents import LiteLLMModel, ToolCollection, CodeAgent
-from dotenv import load_dotenv
-
-
-logger = logging.getLogger(__name__)
-
-async def pg_mcp_smolagent_exec(request: str) -> str:
-    """
-    Execute the full PostgreSQL MCP pipeline: load summary, connect session,
-    load memory and tools, build prompt, run agent, update memory.
-
-    Args:
-        request (str): User's request input.
-        llm (Any): Language model for reasoning agent.
-
-    Returns:
-        str: Agent response message.
-    """
-    # TODO: give summary file path from config
-    table_summary = load_table_summary("table_summary.txt")
-    server_params = get_server_params()
-
-    # TODO: give key from env
-    llm = LiteLLMModel(model_id="gemini/gemini-2.0-flash", api_key="AIzaSyDRoQ8yG0Rmpt3NtfHS7fdm5hgfL9su5aY")
-
-    async with stdio_client(server_params) as (read, write):
-        async with ClientSession(read, write) as session:
-            await session.initialize()
-            memory = await load_or_create_memory()
-
-            intent = classify_intent(request)
-            with ToolCollection.from_mcp(server_params, trust_remote_code=True) as tool_collection:
-                # Create a Code agent using the LLM and tools
-                agent = CodeAgent(model=llm, tools=[*tool_collection.tools], add_base_tools=True)
-                tools = await load_and_enrich_tools(session, table_summary)
-                past_data = get_memory_snapshot(memory)
-
-                prompt = await build_prompt(session, intent, request, tools, past_data, table_summary)
-                agent_response = agent.run(task=prompt, stream=False)
-
-                # TODO: add a smolagent output parser
-
-                #parsed_steps, _ = parse_mcp_output(agent_response)
-                #memory.update_from_parsed(parsed_steps, request)
-
-                #await handle_memory_save_or_reset(memory, request)
-
-            return agent_response
-
-
-# ---------------- Helper Functions ---------------- #
-
-def load_table_summary(path: str) -> str:
-    with open(path, 'r') as file:
-        return file.read()
-
-def get_server_params() -> StdioServerParameters:
-    # TODO: give server params from config
-    load_dotenv()
-    return StdioServerParameters(
-        command="python",
-        args=[os.environ["MCP_SERVER_PATH"]],
-    )
-
-async def load_or_create_memory() -> ConversationMemory:
-    memory = ConversationMemory()
-    if os.path.exists("memory.json"):
-        return memory.load_memory()
-    return memory
-
-async def load_and_enrich_tools(session: ClientSession, summary: str):
-    tools = await load_mcp_tools(session)
-    #for tool in tools:
-    #    tool.description += f" {summary}"
-    return tools
-
-def get_memory_snapshot(memory: ConversationMemory) -> dict:
-    if os.path.exists("memory.json"):
-        return {
-            "past_tools": memory.get_all_tools_used(),
-            "past_queries": memory.get_last_n_queries(),
-            "past_results": memory.get_last_n_results(),
-            "past_requests": memory.get_all_user_messages()
-        }
-    return {
-        "past_tools": "No tools found",
-        "past_queries": "No queries found",
-        "past_results": "No results found",
-        "past_requests": "No requests found"
-    }
-
-async def build_prompt(session, intent, request, tools, past_data, summary):
-    superset_prompt = await session.read_resource("resource://last_prompt")
-    conversation_prompt = await session.read_resource("resource://base_prompt")
-    # TODO: add uri's from config
-    if intent == "superset_request":
-        template = superset_prompt.contents[0].text
-        return template.format(
-            user_requests=past_data["past_requests"],
-            past_tools=past_data["past_tools"],
-            last_queries=past_data["past_queries"],
-            last_results=past_data["past_results"],
-            new_request=request
-        )
-    else:
-        template = conversation_prompt.contents[0].text
-        tools_str = "\n".join([f"- {tool.name}: {tool.description}" for tool in tools])
-        return template.format(
-            user_requests=past_data["past_requests"],
-            past_tools=past_data["past_tools"],
-            last_queries=past_data["past_queries"],
-            last_results=past_data["past_results"],
-            new_request=request,
-            tools=tools_str,
-            descriptions=summary
-        )
-
-async def handle_memory_save_or_reset(memory: ConversationMemory, request: str):
-    if request.strip().lower() == "stop":
-        memory.reset()
-        logger.info("Conversation memory reset.")
-    else:
-        memory.save_memory()
-        logger.info("Conversation memory saved.")