From b172ce0b247c958b5fc44713ab3e3dc661ac9d78 Mon Sep 17 00:00:00 2001
From: Pavan Kumar <66913595+ascender1729@users.noreply.github.com>
Date: Thu, 17 Apr 2025 18:14:53 +0000
Subject: [PATCH] feat: add input validation to MathAgent for math-specific
 queries

- Updated `mock_multi_agent.py` to ensure agent only responds to valid math inputs
- Prevented irrelevant queries from being processed by MathAgent
---
 examples/mcp_example/mock_multi_agent.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/examples/mcp_example/mock_multi_agent.py b/examples/mcp_example/mock_multi_agent.py
index 84576f03..e9a47cf6 100644
--- a/examples/mcp_example/mock_multi_agent.py
+++ b/examples/mcp_example/mock_multi_agent.py
@@ -25,6 +25,15 @@ class MathAgent:
 
     async def process(self, task: str):
         try:
+            # Check if input is math-related
+            math_keywords = ['add', 'plus', '+', 'multiply', 'times', '*', 'x', 'divide', '/', 'by']
+            if not any(keyword in task.lower() for keyword in math_keywords):
+                return {
+                    "agent": self.agent.agent_name,
+                    "task": task,
+                    "response": "Please provide a mathematical operation (add, multiply, or divide)"
+                }
+            
             response = await self.agent.arun(task)
             return {
                 "agent": self.agent.agent_name,