From cb6eca6784efa73915a1c65b418ce60ef41e458c Mon Sep 17 00:00:00 2001
From: Pavan Kumar <66913595+ascender1729@users.noreply.github.com>
Date: Thu, 17 Apr 2025 17:39:59 +0000
Subject: [PATCH] fix: update MCP server config and test script for
 single-agent multi-server setup

- Fixed `math_server.py` to support clean server initialization
- Updated `multi_server_test.py` to validate single agent interactions with multiple MCP servers
---
 examples/mcp_example/math_server.py       |  2 +-
 examples/mcp_example/multi_server_test.py | 88 ++++-------------------
 2 files changed, 16 insertions(+), 74 deletions(-)

diff --git a/examples/mcp_example/math_server.py b/examples/mcp_example/math_server.py
index a4500cd6..66c976c2 100644
--- a/examples/mcp_example/math_server.py
+++ b/examples/mcp_example/math_server.py
@@ -29,4 +29,4 @@ def divide(a: float, b: float) -> float:
 
 if __name__ == "__main__":
     print("Starting Math Server on port 6274...")
-    mcp.run(transport="sse", port=6274)
+    mcp.run(transport="sse", host="0.0.0.0", port=6274)
diff --git a/examples/mcp_example/multi_server_test.py b/examples/mcp_example/multi_server_test.py
index b6d316de..8c08d086 100644
--- a/examples/mcp_example/multi_server_test.py
+++ b/examples/mcp_example/multi_server_test.py
@@ -1,8 +1,7 @@
 from swarms import Agent
 from swarms.tools.mcp_integration import MCPServerSseParams
-from swarms.prompts.finance_agent_sys_prompt import FINANCIAL_AGENT_SYS_PROMPT
 
-# Configure multiple MCP servers
+# Configure servers
 math_server = MCPServerSseParams(
     url="http://0.0.0.0:6274",
     headers={"Content-Type": "application/json"},
@@ -11,40 +10,17 @@ math_server = MCPServerSseParams(
 )
 
 calc_server = MCPServerSseParams(
-    url="http://0.0.0.0:6275", 
+    url="http://0.0.0.0:6275",
     headers={"Content-Type": "application/json"},
     timeout=10.0,
     sse_read_timeout=300.0
 )
 
-# Create specialized agents with different server access
-math_agent = Agent(
-    agent_name="Math-Specialist",
-    agent_description="Advanced mathematics expert",
-    system_prompt="You are a mathematics expert. Use available math operations.",
-    max_loops=1,
-    mcp_servers=[math_server],
-    interactive=True,
-    streaming_on=True,
-    model_name="gpt-4o-mini"
-)
-
-finance_agent = Agent(
-    agent_name="Finance-Specialist",
-    agent_description="Financial calculations expert", 
-    system_prompt=FINANCIAL_AGENT_SYS_PROMPT,
-    max_loops=1,
-    mcp_servers=[calc_server],
-    interactive=True,
-    streaming_on=True,
-    model_name="gpt-4o-mini"
-)
-
-# Multi-server agent with access to all operations
-super_agent = Agent(
-    agent_name="Super-Calculator",
-    agent_description="Multi-capable calculation expert",
-    system_prompt="You have access to multiple calculation servers. Use them appropriately.",
+# Create agent with access to both servers
+calculator = Agent(
+    agent_name="Calculator",
+    agent_description="Math and calculation expert",
+    system_prompt="You are a math expert. Use the available calculation tools to solve problems.",
     max_loops=1,
     mcp_servers=[math_server, calc_server],
     interactive=True,
@@ -52,62 +28,28 @@ super_agent = Agent(
     model_name="gpt-4o-mini"
 )
 
-def format_agent_output(agent_name: str, output: str) -> str:
-    return f"""
-╭─── {agent_name} Response ───╮
-{output}
-╰───────────────────────────╯
-"""
-
 def main():
-    print("\nMulti-Agent MCP Test Environment")
+    print("\nMulti-Server Calculator Test")
     print("Type 'exit' to quit\n")
-    print("Available commands:")
-    print("- math: <calculation> (e.g., 'math: add 5 and 3')")
-    print("- finance: <calculation> (e.g., 'finance: calculate compound interest on 1000 at 5% for 3 years')")
-    print("- Any other calculation will use the super agent\n")
+    print("Example commands:")
+    print("- add 5 and 3")
+    print("- multiply 4 by 7\n")
 
     while True:
         try:
-            user_input = input("\nEnter your calculation request: ")
+            user_input = input("\nEnter calculation: ")
 
             if user_input.lower() == 'exit':
                 break
 
-            # Route request to appropriate agent based on keywords
-            if 'finance' in user_input.lower():
-                response = finance_agent.run(user_input)
-                print("\nFinance Agent Response:")
-                print("-" * 50)
-                print(f"Response: {response}")
-                print("-" * 50)
-            elif 'math' in user_input.lower():
-                response = math_agent.run(user_input)
-                print("\nMath Agent Response:")
-                print("-" * 50) 
-                print(f"Response: {response}")
-                print("-" * 50)
-            else:
-                try:
-                    response = super_agent.run(user_input)
-                    print("\nSuper Agent Response:")
-                    print("-" * 50)
-                    if isinstance(response, dict):
-                        result = response.get('result', response)
-                        print(f"Result: {result}")
-                    elif hasattr(response, 'content'):
-                        print(f"Result: {response.content}")
-                    else:
-                        print(f"Result: {response}")
-                    print("-" * 50)
-                except Exception as e:
-                    print(f"Error processing calculation: {str(e)}")
+            response = calculator.run(user_input)
+            print(f"\nCalculation result: {response}")
 
         except KeyboardInterrupt:
             print("\nExiting gracefully...")
             break
         except Exception as e:
-            print(f"Error processing request: {e}")
+            print(f"Error: {e}")
 
 if __name__ == "__main__":
     main()
\ No newline at end of file