feat: smolagent option

postgre_smolagent_clinet.py

@@ -0,0 +1,128 @@
+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
+
+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-lite", 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()
+
+            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)
+                agent_response = agent.run(task=prompt, stream=False)
+
+
+                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
+    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.")