refactor: separate helper functions and main file from client function. refactor output prints

main.py

@@ -0,0 +1,32 @@
+import yaml
+from pathlib import Path
+from typing import List
+import asyncio
+from postgre_mcp_client import pg_mcp_exec
+import logging
+
+#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 = "Show me the table of join posts and users tables."
+    await pg_mcp_exec(request)
+
+
+if __name__ == "__main__":
+
+    asyncio.run(main())

postgre_mcp_client.py

@@ -1,118 +1,126 @@
-import asyncio
 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
 
-llm = init_chat_model(model="gemini-2.0-flash-lite", model_provider="google_genai",api_key ="AIzaSyAuxYmci0DVU5l5L_YcxLlxHzR5MLn70js")
+logger = logging.getLogger(__name__)
 
-server_params = StdioServerParameters(
-    command="python",
-    args=[r"C:\Users\yukcus\Desktop\MCPTest\postgre_mcp_server.py"],
-)
+async def pg_mcp_exec(request: str) -> str:
+    """
+    Execute the full PostgreSQL MCP pipeline: load summary, connect session,
+    load memory and tools, build prompt, run agent, update memory.
 
-table_summary = """
-The users table stores information about the individuals who use the application. Each user is assigned a unique, auto-incrementing id that serves as the primary key. The username field holds the user's chosen display name and cannot be null, while the email field stores the user’s unique email address, also required and constrained to be unique to avoid duplicates. To track when a user was added to the system, the created_at column records the timestamp of their creation, with a default value set to the current time.
+    Args:
+        request (str): User's request input.
+        llm (Any): Language model for reasoning agent.
 
-The posts table represents content created by users, such as blog posts or messages. Like the users table, each entry has a unique, auto-incrementing id as the primary key. The user_id field links each post to its author by referencing the id field in the users table, establishing a one-to-many relationship between users and posts. The title column holds a brief summary or headline of the post, while the content field contains the full text. A created_at timestamp is also included to record when each post was created, with a default value of the current time.
-"""
+    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-lite", model_provider="google_genai",
+                          api_key="AIzaSyAuxYmci0DVU5l5L_YcxLlxHzR5MLn70js")
 
-#request = "can you show me the result of the join of posts and users tables?"
-#request = "But you did not execute the query. Can you tell me why?"
-#request = "May ı see the table?"
-#request = "stop"
-#request = "how many columns are there in this joined table?"
-request = "send the table"
-async def main():
     async with stdio_client(server_params) as (read, write):
         async with ClientSession(read, write) as session:
-            # Initialize the connection
             await session.initialize()
+            memory = await load_or_create_memory()
 
-            memory = ConversationMemory()
-
-            # Get tools
-            tools = await load_mcp_tools(session)
-            for tool in tools:
-                tool.description += f" {table_summary}"
-
-            if os.path.exists("memory.json"):
-                memory = memory.load_memory()
-                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()
-
-            else:
-                past_tools = "No tools found"
-                past_queries ="No queries found"
-                past_results = "No results found"
-                past_requests = "No requests found"
-
+            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)
 
-            if intent == "superset_request":
-                uri = f"resource://last_prompt"
-                resource = await session.read_resource(uri)
-                base_prompt = resource.contents[0].text
-
-                prompt = base_prompt.format(
-                    user_requests=past_requests,
-                    past_tools=past_tools,
-                    last_queries=past_queries,
-                    last_results=past_results,
-                    new_request=request
-
-                )
-
-            else:
-                uri = f"resource://base_prompt"
-                resource = await session.read_resource(uri)
-                base_prompt = resource.contents[0].text
-
-                # Create a formatted string of tools
-                tools_str = "\n".join([f"- {tool.name}: {tool.description}" for tool in tools])
-
-
-                prompt = base_prompt.format(
-                    user_requests=past_requests,
-                    past_tools=past_tools,
-                    last_queries=past_queries,
-                    last_results=past_results,
-                    new_request = request,
-                    tools = tools_str
-                )
-
-
-
-            # Create and run the agent
             agent = create_react_agent(llm, tools)
             agent_response = await agent.ainvoke({"messages": prompt})
 
-
-            parsed_steps, query_store = parse_mcp_output(agent_response)
-            print("************")
-            print(parsed_steps)
+            parsed_steps, _ = parse_mcp_output(agent_response)
             memory.update_from_parsed(parsed_steps, request)
 
-            if request.strip().lower() == "stop":
-                memory.reset()
-                print("Conversation memory reset.")
-
-            else:
-                memory.save_memory()
-            #pprint(parsed_steps)
-
-
-
-asyncio.run(main())
-
-# open up a new cmd shell and run awith python mcp_client.py
+            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
+    return StdioServerParameters(
+        command="python",
+        args=[r"C:\Users\yukcus\Desktop\MCPTest\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.")

table_summary.txt

@@ -0,0 +1,3 @@
+The users table stores information about the individuals who use the application. Each user is assigned a unique, auto-incrementing id that serves as the primary key. The username field holds the user's chosen display name and cannot be null, while the email field stores the user’s unique email address, also required and constrained to be unique to avoid duplicates. To track when a user was added to the system, the created_at column records the timestamp of their creation, with a default value set to the current time.
+
+The posts table represents content created by users, such as blog posts or messages. Like the users table, each entry has a unique, auto-incrementing id as the primary key. The user_id field links each post to its author by referencing the id field in the users table, establishing a one-to-many relationship between users and posts. The title column holds a brief summary or headline of the post, while the content field contains the full text. A created_at timestamp is also included to record when each post was created, with a default value of the current time.

utils.py

@@ -1,10 +1,12 @@
 import re
+import os
+from conversation_memory import ConversationMemory
+import logging
 
-
+logger = logging.getLogger(__name__)
 def parse_mcp_output(output_dict):
     result = []
     messages = output_dict.get("messages", [])
-
     query_store = []
 
     for msg in messages:
@@ -30,9 +32,20 @@ def parse_mcp_output(output_dict):
 
                 # Check for presence of "query" key
                 if "query" in arguments_dict:
-                    print("query detected!!!")
-                    print(f"ai said:{content[0]}")
-                    print(arguments_dict["query"])
+                    #print("query detected!!!")
+                    print(f"=============== AI Reasoning Step ===============")
+                    print(content[0])
+                    print()
+                    print("=============== AI used the following tools ===============")
+                    print(tool_name)
+                    print()
+                    print("=============== AI generated the following query ===============")
+                    print(arguments_dict['query'])
+
+                    logger.info(f"ai said:{content[0]}")
+                    logger.info(f"ai used:{tool_name}")
+                    logger.info(f"generated query:{arguments_dict['query']}")
+                    #print(arguments_dict["query"])
                     query_store.append(arguments_dict["query"])
 
                     result.append({
@@ -42,7 +55,16 @@ def parse_mcp_output(output_dict):
                         "args": arguments
                     })
                 else:
-                    print(f"ai said:{content}")
+                    #print(f"ai said:{content}")
+                    logger.info(f"ai said:{content}")
+                    logger.info(f"ai used:{tool_name}")
+                    print(f"=============== AI Reasoning Step ===============")
+                    print(content)
+                    print()
+                    print("=============== AI used the following tools ===============")
+                    print(tool_name)
+                    print()
+
                     result.append({
                         "type": "ai_function_call",
                         "ai_said": content,
@@ -51,7 +73,10 @@ def parse_mcp_output(output_dict):
                     })
 
             else:
-                print(f"ai final answer:{content}")
+                #print(f"ai final answer:{content}")
+                logger.info(f"ai final answer:{content}")
+                print("=============== AI's final answer ===============")
+                print(content)
                 result.append({
                     "type": "ai_final_answer",
                     "ai_said": content
@@ -60,7 +85,9 @@ def parse_mcp_output(output_dict):
         # ToolMessage
         elif role_name == "ToolMessage":
             tool_name = getattr(msg, "name", None)
-            print(f"tool response:{content}")
+            print("=============== The tool returned the following response ===============")
+            print(content)
+            logger.info(f"tool response:{content}")
             result.append({
                 "type": "tool_response",
                 "tool": tool_name,
@@ -88,3 +115,4 @@ def classify_intent(user_input: str) -> str:
 
     # Fallback
     return "sql_request"
+