From 10e056a05edf80224fab66de408d4328f1a909d0 Mon Sep 17 00:00:00 2001
From: Pavan Kumar <66913595+ascender1729@users.noreply.github.com>
Date: Thu, 17 Apr 2025 16:23:34 +0000
Subject: [PATCH] test: enhance MCP integration test with tool discovery and
 user guidance

---
 .replit                                  | 17 ++++++
 examples/mcp_example/test_integration.py | 70 +++++++++++++-----------
 2 files changed, 54 insertions(+), 33 deletions(-)

diff --git a/.replit b/.replit
index 7a74e339..9b2aec89 100644
--- a/.replit
+++ b/.replit
@@ -41,3 +41,20 @@ mode = "sequential"
 [[workflows.workflow.tasks]]
 task = "shell.exec"
 args = "python examples/mcp_example/test_integration.py"
+
+[[workflows.workflow]]
+name = "Run MCP Test"
+author = 13983571
+mode = "sequential"
+
+[[workflows.workflow.tasks]]
+task = "shell.exec"
+args = "python examples/mcp_example/math_server.py & "
+
+[[workflows.workflow.tasks]]
+task = "shell.exec"
+args = "sleep 2"
+
+[[workflows.workflow.tasks]]
+task = "shell.exec"
+args = "python examples/mcp_example/test_integration.py"
diff --git a/examples/mcp_example/test_integration.py b/examples/mcp_example/test_integration.py
index 192542ae..5429c2ea 100644
--- a/examples/mcp_example/test_integration.py
+++ b/examples/mcp_example/test_integration.py
@@ -1,43 +1,47 @@
+
 from swarms import Agent
 from swarms.prompts.finance_agent_sys_prompt import FINANCIAL_AGENT_SYS_PROMPT
 from swarms.tools.mcp_integration import MCPServerSseParams
 import logging
 
-# Configure MCP server connection
-server = MCPServerSseParams(
-    url="http://0.0.0.0:6274",
-    headers={"Content-Type": "application/json"},
-    timeout=10.0,
-    sse_read_timeout=300.0
-)
+def main():
+    # Configure MCP server connection
+    server = MCPServerSseParams(
+        url="http://0.0.0.0:6274",
+        headers={"Content-Type": "application/json"},
+        timeout=10.0,
+        sse_read_timeout=300.0
+    )
 
-# Initialize agent with MCP capabilities
-agent = Agent(
-    agent_name="Math-Agent",
-    agent_description="Agent that performs math operations",
-    system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
-    max_loops=1,
-    mcp_servers=[server],
-    streaming_on=True
-)
+    # Initialize agent with MCP capabilities
+    agent = Agent(
+        agent_name="Math-Agent",
+        agent_description="Agent that performs math operations",
+        system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
+        max_loops=1,
+        mcp_servers=[server],
+        streaming_on=True
+    )
 
-def test_mcp_operations():
-    """Test basic MCP operations with error handling"""
     try:
-        # Test addition
-        print("\nTesting addition...")
-        add_result = agent.run("Use the add tool to add 5 and 3")
-        print("Addition result:", add_result)
-
-        # Test multiplication
-        print("\nTesting multiplication...")
-        mult_result = agent.run("Use the multiply tool to multiply 4 and 6") 
-        print("Multiplication result:", mult_result)
-
-        # Test error case
-        print("\nTesting error handling...")
-        error_result = agent.run("Use the add tool with invalid inputs")
-        print("Error handling result:", error_result)
+        # First get available tools from server
+        print("\nDiscovering available tools from MCP server...")
+        tools = agent.mcp_tool_handling()
+        print("\nAvailable tools:", tools)
+        
+        while True:
+            # Get user input
+            user_input = input("\nEnter a math operation (or 'exit' to quit): ")
+            
+            if user_input.lower() == 'exit':
+                break
+                
+            # Process user input through agent
+            try:
+                result = agent.run(user_input)
+                print("\nResult:", result)
+            except Exception as e:
+                print(f"Error processing request: {e}")
 
     except Exception as e:
         logging.error(f"Test failed: {e}")
@@ -45,4 +49,4 @@ def test_mcp_operations():
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO)
-    test_mcp_operations()
\ No newline at end of file
+    main()