diff --git a/examples/mcp_example/math_server.py b/examples/mcp_example/math_server.py index 2be29417..08be7a00 100644 --- a/examples/mcp_example/math_server.py +++ b/examples/mcp_example/math_server.py @@ -1,17 +1,77 @@ from fastmcp import FastMCP +from typing import Dict, Any, Optional +# Initialize MCP server mcp = FastMCP("Math-Server") +# Add tool documentation and type hints @mcp.tool() -def add(a: int, b: int) -> int: - """Add two numbers""" - return a + b +def add(a: int, b: int) -> Dict[str, Any]: + """Add two numbers together + + Args: + a (int): First number to add + b (int): Second number to add + + Returns: + Dict[str, Any]: Result dictionary containing sum and metadata + """ + try: + result = a + b + return { + "status": "success", + "result": result, + "message": f"Successfully added {a} and {b}" + } + except Exception as e: + return { + "status": "error", + "error": str(e), + "message": "Failed to perform addition" + } @mcp.tool() -def multiply(a: int, b: int) -> int: - """Multiply two numbers""" - return a * b +def multiply(a: int, b: int) -> Dict[str, Any]: + """Multiply two numbers together + + Args: + a (int): First number to multiply + b (int): Second number to multiply + + Returns: + Dict[str, Any]: Result dictionary containing product and metadata + """ + try: + result = a * b + return { + "status": "success", + "result": result, + "message": f"Successfully multiplied {a} and {b}" + } + except Exception as e: + return { + "status": "error", + "error": str(e), + "message": "Failed to perform multiplication" + } + +@mcp.tool() +def get_available_operations() -> Dict[str, Any]: + """Get list of available mathematical operations + + Returns: + Dict[str, Any]: Dictionary containing available operations and their descriptions + """ + return { + "status": "success", + "operations": { + "add": "Add two numbers together", + "multiply": "Multiply two numbers together" + } + } if __name__ == "__main__": + print("Starting Math Server...") + print("Available operations:", get_available_operations()) mcp.run(host="0.0.0.0", port=6274, transport="sse") diff --git a/examples/mcp_example/test_integration.py b/examples/mcp_example/test_integration.py index 1d583158..b9dbd74e 100644 --- a/examples/mcp_example/test_integration.py +++ b/examples/mcp_example/test_integration.py @@ -2,27 +2,70 @@ 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 +from typing import Dict, Any, Optional -# Configure MCP server connection -server = MCPServerSseParams( - url="http://0.0.0.0:6274", - headers={"Content-Type": "application/json"} -) - -# 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 -) - -# Test addition -result = agent.run("Use the add tool to add 5 and 3") -print("Addition result:", result) - -# Test multiplication -result = agent.run("Use the multiply tool to multiply 4 and 6") -print("Multiplication result:", result) +def handle_mcp_response(response: Dict[str, Any]) -> str: + """Handle MCP response and extract meaningful output""" + if response.get("status") == "error": + return f"Error: {response.get('message', 'Unknown error occurred')}" + return str(response.get("result", response)) + +def setup_mcp_agent(name: str, description: str) -> Agent: + """Setup an MCP-enabled agent with proper configuration""" + try: + server = MCPServerSseParams( + url="http://0.0.0.0:6274", + headers={"Content-Type": "application/json"}, + timeout=10.0, + sse_read_timeout=300.0 + ) + + return Agent( + agent_name=name, + agent_description=description, + system_prompt=FINANCIAL_AGENT_SYS_PROMPT, + max_loops=1, + mcp_servers=[server], + streaming_on=True + ) + except Exception as e: + logging.error(f"Failed to setup agent: {e}") + raise + +def test_mcp_operations(): + """Test basic MCP operations with error handling""" + try: + # Initialize agent + agent = setup_mcp_agent( + "Math-Agent", + "Agent that performs math operations" + ) + + # Get available operations + print("\nQuerying available operations...") + result = agent.run("What operations are available?") + print("Available operations:", result) + + # Test addition + print("\nTesting addition...") + add_result = agent.run("Use the add tool to add 5 and 3") + print("Addition result:", handle_mcp_response(add_result)) + + # Test multiplication + print("\nTesting multiplication...") + mult_result = agent.run("Use the multiply tool to multiply 4 and 6") + print("Multiplication result:", handle_mcp_response(mult_result)) + + # Test error case + print("\nTesting error handling...") + error_result = agent.run("Use the add tool with invalid inputs") + print("Error handling result:", handle_mcp_response(error_result)) + + except Exception as e: + logging.error(f"Test failed: {e}") + raise + +if __name__ == "__main__": + logging.basicConfig(level=logging.INFO) + test_mcp_operations()