From 40cd5248d8a194420702b2b85701d8cde81626d6 Mon Sep 17 00:00:00 2001
From: Pavan Kumar <66913595+ascender1729@users.noreply.github.com>
Date: Thu, 17 Apr 2025 18:04:13 +0000
Subject: [PATCH] fix: align math server tools and agent prompt with available
 capabilities

- Updated system prompt in `mock_multi_agent.py` to reflect only supported tools
- Added error handling in `mock_math_server.py` to reject unsupported operations
- Ensured agent behavior stays within defined context and MCP protocol limits
---
 examples/mcp_example/mock_math_server.py | 25 ++++++++++++++++--------
 examples/mcp_example/mock_multi_agent.py |  2 +-
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/examples/mcp_example/mock_math_server.py b/examples/mcp_example/mock_math_server.py
index 1f8df921..d2b40b71 100644
--- a/examples/mcp_example/mock_math_server.py
+++ b/examples/mcp_example/mock_math_server.py
@@ -9,22 +9,31 @@ mcp = FastMCP("Math-Mock-Server")
 @mcp.tool()
 def add(a: int, b: int) -> int:
     """Add two numbers together"""
-    time.sleep(0.1)  # Simulate processing time
-    return a + b
+    try:
+        time.sleep(0.1)  # Simulate processing time
+        return a + b
+    except Exception as e:
+        return {"error": f"Error adding numbers: {str(e)}"}
 
 @mcp.tool()
 def multiply(a: int, b: int) -> int:
     """Multiply two numbers together"""
-    time.sleep(0.1)  # Simulate processing time
-    return a * b
+    try:
+        time.sleep(0.1)  # Simulate processing time
+        return a * b
+    except Exception as e:
+        return {"error": f"Error multiplying numbers: {str(e)}"}
 
 @mcp.tool()
 def divide(a: int, b: int) -> float:
     """Divide two numbers"""
-    if b == 0:
-        raise ValueError("Cannot divide by zero")
-    time.sleep(0.1)  # Simulate processing time
-    return a / b
+    try:
+        if b == 0:
+            return {"error": "Cannot divide by zero"}
+        time.sleep(0.1)  # Simulate processing time
+        return a / b
+    except Exception as e:
+        return {"error": f"Error dividing numbers: {str(e)}"}
 
 if __name__ == "__main__":
     print("Starting Mock Math Server on port 8000...")
diff --git a/examples/mcp_example/mock_multi_agent.py b/examples/mcp_example/mock_multi_agent.py
index 80c0381c..effb7760 100644
--- a/examples/mcp_example/mock_multi_agent.py
+++ b/examples/mcp_example/mock_multi_agent.py
@@ -14,7 +14,7 @@ class MathAgent:
         self.agent = Agent(
             agent_name=name,
             agent_description="Math processing agent",
-            system_prompt=f"You are {name}, a math processing agent. Use the provided tools to solve math problems.",
+            system_prompt=f"You are {name}, a math processing agent. You have access to these mathematical operations ONLY: addition, multiplication, and division. Only suggest calculations using these available tools.",
             max_loops=1,
             mcp_servers=[self.server],
             streaming_on=False,