From 84d1478b668b42e180ce9075e8997361f09a3439 Mon Sep 17 00:00:00 2001
From: zachary62 <zhuang333@wisc.edu>
Date: Mon, 7 Apr 2025 16:50:07 -0400
Subject: [PATCH] add more ops

---
 cookbook/pocketflow-mcp/simple_server.py | 21 +++++++++-
 cookbook/pocketflow-mcp/utils.py         | 50 ++++++++++++++++++++++++
 2 files changed, 69 insertions(+), 2 deletions(-)

diff --git a/cookbook/pocketflow-mcp/simple_server.py b/cookbook/pocketflow-mcp/simple_server.py
index 166e34d..329e915 100644
--- a/cookbook/pocketflow-mcp/simple_server.py
+++ b/cookbook/pocketflow-mcp/simple_server.py
@@ -1,14 +1,31 @@
 from fastmcp import FastMCP
 
 # Create a named server
-mcp = FastMCP("Addition Server")
+mcp = FastMCP("Math Operations Server")
 
-# Define an addition tool
+# Define mathematical operation tools
 @mcp.tool()
 def add(a: int, b: int) -> int:
     """Add two numbers together"""
     return a + b
 
+@mcp.tool()
+def subtract(a: int, b: int) -> int:
+    """Subtract b from a"""
+    return a - b
+
+@mcp.tool()
+def multiply(a: int, b: int) -> int:
+    """Multiply two numbers together"""
+    return a * b
+
+@mcp.tool()
+def divide(a: int, b: int) -> float:
+    """Divide a by b"""
+    if b == 0:
+        raise ValueError("Division by zero is not allowed")
+    return a / b
+
 # Start the server
 if __name__ == "__main__":
     mcp.run()
\ No newline at end of file
diff --git a/cookbook/pocketflow-mcp/utils.py b/cookbook/pocketflow-mcp/utils.py
index 9d54f71..2bd4349 100644
--- a/cookbook/pocketflow-mcp/utils.py
+++ b/cookbook/pocketflow-mcp/utils.py
@@ -42,6 +42,39 @@ def local_get_tools(server_script_path=None):
                 },
                 "required": ["a", "b"]
             }
+        },
+        {
+            "name": "subtract",
+            "description": "Subtract b from a",
+            "inputSchema": {
+                "properties": {
+                    "a": {"type": "integer"},
+                    "b": {"type": "integer"}
+                },
+                "required": ["a", "b"]
+            }
+        },
+        {
+            "name": "multiply",
+            "description": "Multiply two numbers together",
+            "inputSchema": {
+                "properties": {
+                    "a": {"type": "integer"},
+                    "b": {"type": "integer"}
+                },
+                "required": ["a", "b"]
+            }
+        },
+        {
+            "name": "divide",
+            "description": "Divide a by b",
+            "inputSchema": {
+                "properties": {
+                    "a": {"type": "integer"},
+                    "b": {"type": "integer"}
+                },
+                "required": ["a", "b"]
+            }
         }
     ]
 
@@ -88,6 +121,23 @@ def local_call_tool(server_script_path=None, tool_name=None, arguments=None):
             return arguments["a"] + arguments["b"]
         else:
             return "Error: Missing required arguments 'a' or 'b'"
+    elif tool_name == "subtract":
+        if "a" in arguments and "b" in arguments:
+            return arguments["a"] - arguments["b"]
+        else:
+            return "Error: Missing required arguments 'a' or 'b'"
+    elif tool_name == "multiply":
+        if "a" in arguments and "b" in arguments:
+            return arguments["a"] * arguments["b"]
+        else:
+            return "Error: Missing required arguments 'a' or 'b'"
+    elif tool_name == "divide":
+        if "a" in arguments and "b" in arguments:
+            if arguments["b"] == 0:
+                return "Error: Division by zero is not allowed"
+            return arguments["a"] / arguments["b"]
+        else:
+            return "Error: Missing required arguments 'a' or 'b'"
     else:
         return f"Error: Unknown tool '{tool_name}'"