feat: langchain buffermemory try

langchain_mcp_client.py

@@ -12,12 +12,16 @@ from utils import parse_mcp_output, classify_intent
 import logging
 from dotenv import load_dotenv
 from langgraph.checkpoint.memory import MemorySaver
+from langchain.memory import ConversationBufferMemory
+from langchain.agents import AgentExecutor
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_core.prompts import PromptTemplate
 
 
 
 logger = logging.getLogger(__name__)
 load_dotenv()
-async def lc_mcp_exec(request: str, history) -> tuple[Any, Any]:
+async def lc_mcp_exec(request: str, history) -> str:
     """
     Execute the full PostgreSQL MCP pipeline: load summary, connect session,
     load memory and tools, build prompt, run agent, update memory.
@@ -34,9 +38,14 @@ async def lc_mcp_exec(request: str, history) -> tuple[Any, Any]:
     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="gemini-2.0-flash",
+                model_provider="google_genai",
+                api_key=api_key,
+                temperature=0.5,
+            )
+    print(type(llm))
     async with stdio_client(server_params) as (read, write):
         async with ClientSession(read, write) as session:
             await session.initialize()
@@ -45,45 +54,19 @@ async def lc_mcp_exec(request: str, history) -> tuple[Any, Any]:
 
             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))
+            memory = ConversationBufferMemory(return_messages=True)
+            agent = create_react_agent(llm, tools)
 
-            prompt = await build_prompt(session, intent, request, tools, table_summary, messages)
+            prompt = await build_prompt(session, intent, request, tools, table_summary)
+            agent_executor = AgentExecutor(agent=agent, tools=tools, memory=memory, verbose=True)
             config = {"configurable": {"thread_id": "conversation_123"}}
-            agent_response = await agent.ainvoke(
+            result = await agent_executor.ainvoke(
                 {"messages": prompt},
                 config
             )
+            response = result.get("output", "No response generated")
 
-            if "messages" in agent_response:
-                response = agent_response["messages"][-1].content
-            else:
-                response = "No response generated"
-
-            messages.append(AIMessage(content=response))
-
-
-            return  response, messages
+            return response
 
 
 # ---------------- Helper Functions ---------------- #
@@ -103,7 +86,7 @@ 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):
+async def build_prompt(session, intent, request, tools, 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
@@ -114,11 +97,8 @@ async def build_prompt(session, intent, request, tools, summary, messages):
         )
     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
-        )
+        #tools_str = "\n".join([f"- {tool.name}: {tool.description}" for tool in tools])
+        #return template.format(  new_request=request,  tools=tools_str, descriptions=summary  )
+        prompt = PromptTemplate.from_template(template)
+        return prompt
 

postgre_mcp_server.py

@@ -134,6 +134,7 @@ async def base_prompt_query() -> str:
                     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}
+                    {tool_names}
 
                     ---
 
@@ -149,37 +150,38 @@ async def base_prompt_query() -> str:
                       > SELECT * FROM customers WHERE country = 'Germany';
                       > ```
 
-==========================
-# 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.:
-
-```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.}}
+                    ==========================
+                    # 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.:
+                    
+                    ```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.}}
 
                     **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}
+                    {input}
+                    
+                    Thought:{agent_scratchpad}
                     """
 
     return base_prompt