feat: gradio and streamlit app options, formatting output

gradio_app.py

@@ -0,0 +1,38 @@
+import yaml
+from pathlib import Path
+import gradio as gr
+import asyncio
+from postgre_mcp_client import pg_mcp_exec
+
+def load_db_configs():
+    """Load database configurations from configs.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-compatible wrapper
+async def run_agent(request):
+    # configs = load_db_configs()  # Optional if needed
+    final_answer, last_tool_answer, = await pg_mcp_exec(request)
+    return final_answer, last_tool_answer
+
+# Gradio UI
+demo = gr.Interface(
+    fn=run_agent,
+    inputs=gr.Textbox(
+        label="Natural Language Request",
+        placeholder="e.g., Show me the table of join posts and users tables."
+    ),
+    outputs=gr.Textbox(label="SQL Query / Result"),
+    title="PostgreSQL Query Agent",
+    description="Ask your database in natural language and get results using the smolagent executor."
+)
+
+if __name__ == "__main__":
+    demo.launch()

postgre_mcp_client.py

@@ -44,12 +44,12 @@ async def pg_mcp_exec(request: str) -> str:
             agent = create_react_agent(llm, tools)
             agent_response = await agent.ainvoke({"messages": prompt})
 
-            parsed_steps, _ = parse_mcp_output(agent_response)
+            parsed_steps, final_answer, last_tool_answer, _ = parse_mcp_output(agent_response)
             memory.update_from_parsed(parsed_steps, request)
 
             await handle_memory_save_or_reset(memory, request)
 
-            return agent_response
+            return  final_answer, last_tool_answer
 
 
 # ---------------- Helper Functions ---------------- #

postgre_smolagent_clinet.py

@@ -46,10 +46,10 @@ async def pg_mcp_smolagent_exec(request: str) -> str:
                 agent_response = agent.run(task=prompt, stream=False)
 
 
-                parsed_steps, _ = parse_mcp_output(agent_response)
-                memory.update_from_parsed(parsed_steps, request)
+                #parsed_steps, _ = parse_mcp_output(agent_response)
+                #memory.update_from_parsed(parsed_steps, request)
 
-                await handle_memory_save_or_reset(memory, request)
+                #await handle_memory_save_or_reset(memory, request)
 
             return agent_response
 

streamlit_app.py

@@ -0,0 +1,38 @@
+import yaml
+from pathlib import Path
+import streamlit as st
+from postgre_mcp_client import pg_mcp_exec
+import asyncio
+import nest_asyncio
+nest_asyncio.apply()
+
+def load_db_configs():
+    """Load database configurations from configs.yaml"""
+    configs_path = Path("configs.yaml")
+
+    if not configs_path.exists():
+        st.error("configs.yaml not found")
+        return None
+
+    with open(configs_path) as f:
+        configs = yaml.safe_load(f)
+
+    return configs.get("db_configs", {})
+
+def run_agent(message):
+    response = asyncio.run(pg_mcp_exec(message))
+    # Return in message format
+    return {"role": "assistant", "content": response}
+
+# Streamlit UI
+st.title("PostgreSQL Query Agent")
+st.write("Ask your database in natural language and get results using the smolagent executor.")
+
+user_input = st.text_input("Natural Language Request", placeholder="e.g., Show me the table of join posts and users tables.")
+
+if st.button("Run Query"):
+    if user_input.strip():
+        result = run_agent(user_input)
+        st.text_area("SQL Query / Result", value=str(result), height=300)
+    else:
+        st.warning("Please enter a natural language request.")

utils.py

@@ -8,7 +8,8 @@ def parse_mcp_output(output_dict):
     result = []
     messages = output_dict.get("messages", [])
     query_store = []
-
+    last_tool_answer = ""
+    last_answer = ""
     for msg in messages:
         role_name = msg.__class__.__name__  # Example: HumanMessage, AIMessage, ToolMessage
         content = getattr(msg, "content", "")
@@ -33,9 +34,11 @@ def parse_mcp_output(output_dict):
                 # Check for presence of "query" key
                 if "query" in arguments_dict:
                     #print("query detected!!!")
-                    print(f"=============== AI Reasoning Step ===============")
-                    print(content[0])
-                    print()
+
+                    if content:
+                        print(f"=============== AI Reasoning Step ===============")
+                        print(content[0])
+                        print()
                     print("=============== AI used the following tools ===============")
                     print(tool_name)
                     print()
@@ -73,6 +76,7 @@ def parse_mcp_output(output_dict):
                     })
 
             else:
+                final_answer = content
                 #print(f"ai final answer:{content}")
                 logger.info(f"ai final answer:{content}")
                 print("=============== AI's final answer ===============")
@@ -87,6 +91,7 @@ def parse_mcp_output(output_dict):
             tool_name = getattr(msg, "name", None)
             print("=============== The tool returned the following response ===============")
             print(content)
+            last_tool_answer = content
             logger.info(f"tool response:{content}")
             result.append({
                 "type": "tool_response",
@@ -94,7 +99,7 @@ def parse_mcp_output(output_dict):
                 "response": content
             })
 
-    return result, query_store
+    return result, final_answer, last_tool_answer, query_store