fix: memory and system messages

.gitignore

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

chat_history.json

@@ -1 +0,0 @@
-[{"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}}, {"type": "human", "data": {"content": "can you find the foreign keys and relationships ?", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "I am unable to retrieve the foreign keys using the `get_foreign_keys` tool. It seems there's an issue with the tool's execution. However, based on the table schemas provided, I can infer the relationships.\n\n# Result\n\n*   **dim\\_agreement**:  Related to `dim_customer` via `customer_id`.\n*   **dim\\_customer**:  Likely related to `dim_product`, `dim_product_order_item`, and potentially other tables through their respective `customer_id` columns.\n*   **dim\\_product**: Related to `dim_customer` via `customer_id`, `dim_agreement` via `agreement_id`, and likely other tables through foreign key relationships.\n*   **dim\\_product\\_order\\_item**: Related to `dim_customer` via `customer_id`, `dim_agreement` via `agreement_id`, and likely related to a product order table via a `product_order_id`.\n\n# Explanation\n\nThe `get_foreign_keys` tool did not return the foreign key relationships. Therefore, I am relying on the schema descriptions to determine the relationships between the tables. The tables are connected through foreign keys, as described in the table descriptions.\n\n# Query\n\nN/A", "additional_kwargs": {}, "response_metadata": {}, "type": "ai", "name": null, "id": null, "example": false, "tool_calls": [], "invalid_tool_calls": [], "usage_metadata": null}}, {"type": "human", "data": {"content": "get all the foreign keys and summarize them", "additional_kwargs": {}, "response_metadata": {}, "type": "human", "name": null, "id": null, "example": false}}, {"type": "ai", "data": {"content": "I am sorry, I cannot directly retrieve and summarize all foreign keys using the available tools. The `get_foreign_keys` tool did not work in a previous turn. However, I can describe the relationships between the tables based on the schema descriptions.\n\n# Result\n\n*   **dim\\_agreement**: Related to `dim_customer` via `customer_id`.\n*   **dim\\_customer**: Likely related to `dim_product` and `dim_product_order_item` through their respective `customer_id` columns.\n*   **dim\\_product**: Related to `dim_customer` via `customer_id`, `dim_agreement` via `agreement_id`, and likely other tables through foreign key relationships.\n*   **dim\\_product\\_order\\_item**: Related to `dim_customer` via `customer_id`, `dim_agreement` via `agreement_id`, and likely related to a product order table via a `product_order_id`.\n\n# Explanation\n\nThe `get_foreign_keys` tool did not return the foreign key relationships. Therefore, I am relying on the schema descriptions to determine the relationships between the tables. The tables are connected through foreign keys, as described in the table descriptions.\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}}]

gradio_app.py

@@ -56,6 +56,20 @@ 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])
+
+    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

@@ -5,57 +5,35 @@ 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_core.prompts import PromptTemplate
-from langchain_core.messages import AIMessage, HumanMessage
-from langchain.memory import ChatMessageHistory
+from langchain_core.messages import AIMessage, HumanMessage, ToolMessage, SystemMessage
 from langchain_community.chat_message_histories import FileChatMessageHistory
 from langchain.chat_models import init_chat_model
-from utils import parse_mcp_output, classify_intent
 import logging
 from dotenv import load_dotenv
 
-
 # Set up logging
 logger = logging.getLogger(__name__)
 
-# Load environment variables from .env file
+# Load environment variables
 load_dotenv()
 
-async def lc_mcp_exec(request: str, history) -> tuple[Any, Any]:
+async def lc_mcp_exec(request: str, history=None) -> Tuple[str, list]:
     """
-    Execute the full PostgreSQL MCP pipeline with persistent memory.
+    Execute the PostgreSQL MCP pipeline with persistent memory in a JSON file.
+    Returns the response and the updated message history.
     """
     try:
+        # Define the path for the chat history JSON file
         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)
-
+        # Initialize FileChatMessageHistory for persistent storage
         message_history = FileChatMessageHistory(file_path=history_file)
 
-        try:
-            # Load existing messages or handle bootstrap scenario
-            existing_messages = message_history.messages
-        except json.JSONDecodeError:
-            # If JSON is corrupted, reinitialize the file
-            logger.warning("Chat history file corrupted, reinitializing...")
-            with open(history_file, 'w') as f:
-                json.dump({"messages": []}, f)
-            existing_messages = []
-
-        # Format existing messages properly
-        formatted_history = []
-        for msg in existing_messages:
-            if isinstance(msg, HumanMessage):
-                formatted_history.append(HumanMessage(content=msg.content))
-            elif isinstance(msg, AIMessage):
-                formatted_history.append(AIMessage(content=msg.content))
-
+        # Load table summary and server parameters
         table_summary = load_table_summary(os.environ["TABLE_SUMMARY_PATH"])
         server_params = get_server_params()
 
+        # Initialize the LLM
         llm = init_chat_model(
             model=os.getenv("GEMINI_MODEL"),
             model_provider=os.getenv("GEMINI_MODEL_PROVIDER"),
@@ -67,60 +45,62 @@ async def lc_mcp_exec(request: str, history) -> tuple[Any, Any]:
             async with ClientSession(read, write) as session:
                 await session.initialize()
 
+                # Load tools
                 tools = await load_and_enrich_tools(session)
 
-                # Add new user message before processing
+                # Create the ReAct agent with tools
+                agent = create_react_agent(llm, tools)
+
+                # Add the new user message to history as HumanMessage
                 message_history.add_user_message(request)
 
-                # Create agent and prepare system message
-                agent = create_react_agent(llm, tools)
+                # Prepare the system prompt for the agent
+                system_prompt = await build_prompt(session, tools, table_summary)
 
-                # 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))
-
-                # Build prompt with conversation context
-                prompt = await build_prompt(
-                    session=session,
-                    request=request,
-                    tools=tools,
-                    table_summary=table_summary,
-                    chat_history=formatted_history
+                # Init the system message
+                system_message = SystemMessage(
+                    content=system_prompt
                 )
 
-                # Invoke agent with proper message structure
+                # Add the system message to the history
+                input_messages = [system_message] + message_history.messages
+
+                # Invoke the agent with the message list
                 agent_response = await agent.ainvoke(
-                    {
-                        "messages": prompt,
-                        "chat_history": [msg.content for msg in formatted_history]
-                    },
+                    {"messages": input_messages},
                     config={"configurable": {"thread_id": "conversation_123"}}
                 )
 
-                if "messages" in agent_response:
-                    response = agent_response["messages"][-1].content
-                    # Save assistant response
-                    message_history.add_ai_message(response)
+                # Extract the latest response and save all new messages
+                response_content = "No response generated"
+                if "messages" in agent_response and agent_response["messages"]:
+                    # Identify new messages (those not in input_messages)
+                    new_messages = agent_response["messages"][len(input_messages):]
+
+                    # Save all new messages to history
+                    for msg in new_messages:
+                        if isinstance(msg, AIMessage):
+                            # Save AIMessage, including tool_calls if present
+                            message_history.add_message(msg)
+                        elif isinstance(msg, ToolMessage):
+                            # Save ToolMessage with content and tool_call_id
+                            message_history.add_message(msg)
+                        else:
+                            # Log unexpected message types for debugging
+                            logger.debug(f"Skipping unexpected message type: {type(msg)}")
+
+                    # Use the last message’s content as the response
+                    response_content = agent_response["messages"][-1].content
+
                 else:
-                    response = "No response generated"
-                    message_history.add_ai_message(response)
+                    message_history.add_ai_message(response_content)
 
-                # Return current response and up-to-date messages
-                return response, message_history.messages
+                # Return the response and the updated history
+                return response_content, message_history.messages
 
     except Exception as e:
-        logger.error(f"Error in chat history handling: {str(e)}", exc_info=True)
-        # Fallback to stateless response if history fails
-        return f"Error in conversation: {str(e)}", []
-
+        logger.error(f"Error in execution: {str(e)}", exc_info=True)
+        return f"Error: {str(e)}", message_history.messages if 'message_history' in locals() else []
 
 # ---------------- Helper Functions ---------------- #
 
@@ -128,10 +108,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()
     return StdioServerParameters(
         command="python",
         args=[os.environ["MCP_SERVER_PATH"]],
@@ -143,30 +120,13 @@ async def load_and_enrich_tools(session: ClientSession):
     return tools
 
 
-async def build_prompt(session, request, tools, table_summary, chat_history):
+async def build_prompt(session, 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])
 
-    # 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_server.py

@@ -155,7 +155,7 @@ You can use the following FastMCP tools. These allow you to create **read-only**
 # Output Format
 ==========================
 
-Present your final answer using the following structure **exactly**. When necessary, bold the important parts of your answer or use `` inline code blocks.:
+Present your final answer using the following structure **exactly** in markdown language. When necessary, bold the important parts of your answer or use `` for inline code blocks:
 
 ```markdown
 # Result
@@ -167,21 +167,15 @@ Present your final answer using the following structure **exactly**. When necess
 # 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
-=========================
-
-Please fulfill the following request based on the above context:
 
-{new_request}
 """
 
     return base_prompt

utils.py

@@ -102,22 +102,3 @@ def parse_mcp_output(output_dict):
     return result, final_answer, last_tool_answer, query_store
 
 
-
-
-
-def classify_intent(user_input: str) -> str:
-    user_input = user_input.lower().strip()
-
-    superset_keywords = [
-        "send to superset", "chart", "visualize", "visualise",
-        "plot", "graph", "send this", "send that", "create a chart",
-        "push to superset", "make a chart", "show chart", "dashboard", "send"
-    ]
-
-    # Check for superset intent
-    if any(kw in user_input for kw in superset_keywords):
-        return "superset_request"
-
-    # Fallback
-    return "sql_request"
-