feat: langchain memory

.env.sample

@@ -0,0 +1 @@
+API_KEY = ""

.gitignore

@@ -0,0 +1,2 @@
+.idea
+.env

gradio_app.py

@@ -4,6 +4,7 @@ import gradio as gr
 import asyncio
 from postgre_mcp_client import pg_mcp_exec
 from postgre_smolagent_clinet import pg_mcp_smolagent_exec
+from langchain_mcp_client import lc_mcp_exec
 
 def load_db_configs():
     """Load database configurations from configs.yaml"""
@@ -18,12 +19,12 @@ def load_db_configs():
     return configs["db_configs"]
 
 # Async-compatible wrapper
-async def run_agent(request):
+async def run_agent(request, history):
     # configs = load_db_configs()
     # final_answer, last_tool_answer, = await pg_mcp_exec(request)
     # return final_answer, last_tool_answer
 
-    result = await pg_mcp_smolagent_exec(request)
+    result = await lc_mcp_exec(request, history)
     return result
 
 # Gradio UI

langchain_mcp_client.py

@@ -0,0 +1,124 @@
+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_core.prompts import PromptTemplate
+from langchain_core.messages import AIMessage, HumanMessage
+from langchain.chat_models import init_chat_model
+from utils import parse_mcp_output, classify_intent
+import logging
+from dotenv import load_dotenv
+from langgraph.checkpoint.memory import MemorySaver
+
+
+
+logger = logging.getLogger(__name__)
+load_dotenv()
+async def lc_mcp_exec(request: str, history) -> 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()
+
+    api_key = os.getenv("API_KEY")
+    llm = init_chat_model(model="gemini-2.0-flash", model_provider="google_genai",
+                          api_key=api_key)
+
+    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)
+
+            messages = []
+            memory = MemorySaver()
+            agent = create_react_agent(llm, tools, checkpointer=memory)
+
+            messages.append(HumanMessage(content=request))
+            if history:
+                # Clear existing messages and rebuild from history
+                messages = []
+
+                # Process Gradio chat history format
+                for msg in history:
+                    # Gradio format: {'role': 'user/assistant', 'metadata': None, 'content': 'message', 'options': None}
+                    role = msg.get('role', '')
+                    content = msg.get('content', '')
+
+                    if role == 'user' and content:
+                        messages.append(HumanMessage(content=content))
+                    elif role == 'assistant' and content:
+                        messages.append(AIMessage(content=content))
+
+                # Add the current query
+                messages.append(HumanMessage(content=request))
+
+            prompt = await build_prompt(session, intent, request, tools, table_summary, messages)
+            config = {"configurable": {"thread_id": "conversation_123"}}
+            agent_response = await agent.ainvoke(
+                {"messages": prompt},
+                config
+            )
+
+            if "messages" in agent_response:
+                response = agent_response["messages"][-1].content
+            else:
+                response = "No response generated"
+
+            messages.append(AIMessage(content=response))
+
+
+            return  response, messages
+
+
+# ---------------- 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"C:\Users\yukcus\Desktop\query_mcp_server\postgre_mcp_server.py"],
+    )
+
+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, messages):
+    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(
+            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(
+            new_request=request,
+            tools=tools_str,
+            descriptions=summary,
+            chat_history = messages
+        )
+

postgre_mcp_server.py

@@ -71,10 +71,6 @@ async def base_prompt_query() -> str:
 
 
                     Each tool may also return previews or summaries of table contents to help you better understand the data structure.
-
-                    You also have access to **short-term memory**, which stores relevant context from earlier queries. If memory contains the needed information, you **must use it** instead of repeating tool calls with the same input. Avoid redundant tool usage unless:
-                    - The memory is empty, or
-                    - A tool's output is outdated or missing
                     
                     ---
 
@@ -121,20 +117,6 @@ async def base_prompt_query() -> str:
                     - 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.
 
-
-                    ---
-
-                    ==========================
-                    # Short-Term Memory
-                    ==========================
-
-                    You have access to the following memory from this conversation. Use it if applicable for the current request.
-
-                    - Previous user requests: {user_requests}
-                    - Tools used so far: {past_tools}
-                    - Last SQL queries: {last_queries}
-                    - Last result preview: {last_results}
-
                     ---
                     
                     ==========================
@@ -189,6 +171,8 @@ Present your final answer using the following structure **exactly**. When necess
                     **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
                     =========================
@@ -218,8 +202,6 @@ async def last_prompt() -> str:
 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.
 
-You have access to a **short-term memory**, which stores relevant context from earlier interactions in the current conversation.
-
 ---
 
 ==========================
@@ -228,7 +210,7 @@ You have access to a **short-term memory**, which stores relevant context from e
 
 When a user submits a request (e.g., *"send me that table"*, *"send the last query"*, etc.), follow these steps:
 
-1. Identify which previous result the user is referring to, using your short-term memory.
+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.
@@ -241,35 +223,12 @@ When a user submits a request (e.g., *"send me that table"*, *"send the last que
 
 - Do **not** ask questions or request clarification.
 - Do **not** explain anything to the user.
-- Only use the **memory** to determine which query is relevant.
 - Respond with the **exact ANSI SQL query only**, formatted cleanly.
-- Do **not** guess — only retrieve queries that actually exist in memory.
+- Do **not** guess — only retrieve queries.
 - If no query fits, respond with: "Query not found."
 
 ---
 
-==========================
-# Short-Term Memory
-==========================
-
-You have access to the following memory from this conversation:
-
-- **Previous user requests**:  
-  `{user_requests}`
-
-- **Tools used so far**:  
-  `{past_tools}`
-
-- **Recent SQL queries**:  
-  `{last_queries}`
-
-- **Result preview from last query**:  
-  `{last_results}`
-
-Use this memory to resolve any references in the user's latest request.
-
----
-
 ==========================
 # Reference Conversion Rules for PostgreSQL to ANSI SQL
 ==========================

run.sh

@@ -0,0 +1,6 @@
+#!/bin/bash
+
+# Replace 'myenv' with the name of your conda environment
+conda activate myenv
+
+python gradio_app.py