"""
Script to register tools with MCP Gateway
This script registers Lambda functions and REST APIs as MCP tools
"""
import os
import sys
import json
import boto3
from typing import Dict, Any
# Add parent directory to path
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from gateway.gateway import MCPGateway
def get_lambda_arn(function_name: str, region: str = "us-east-1") -> str:
"""Get Lambda function ARN by name"""
lambda_client = boto3.client('lambda', region_name=region)
try:
response = lambda_client.get_function(FunctionName=function_name)
return response['Configuration']['FunctionArn']
except Exception as e:
print(f"Error getting ARN for {function_name}: {e}")
# Return a placeholder ARN format
account_id = boto3.client('sts').get_caller_identity()['Account']
return f"arn:aws:lambda:{region}:{account_id}:function:{function_name}"
def register_lambda_tools(gateway: MCPGateway, environment: str = "dev", region: str = "us-east-1"):
"""Register Lambda functions as MCP tools"""
# Weather Tool
gateway.register_lambda_tool(
name="weather_lookup",
lambda_arn=get_lambda_arn(f"{environment}-weather-tool", region),
description="Get weather information for a location",
input_schema={
"type": "object",
"properties": {
"location": {
"type": "string",
"description": "City name or location"
},
"units": {
"type": "string",
"enum": ["celsius", "fahrenheit"],
"description": "Temperature units",
"default": "celsius"
}
},
"required": ["location"]
}
)
# Calculator Tool
gateway.register_lambda_tool(
name="calculator",
lambda_arn=get_lambda_arn(f"{environment}-calculator-tool", region),
description="Perform mathematical calculations",
input_schema={
"type": "object",
"properties": {
"operation": {
"type": "string",
"enum": ["add", "subtract", "multiply", "divide"],
"description": "Mathematical operation"
},
"a": {
"type": "number",
"description": "First number"
},
"b": {
"type": "number",
"description": "Second number"
}
},
"required": ["operation", "a", "b"]
}
)
# Data Lookup Tool
gateway.register_lambda_tool(
name="data_lookup",
lambda_arn=get_lambda_arn(f"{environment}-data-lookup-tool", region),
description="Lookup data by query type and value",
input_schema={
"type": "object",
"properties": {
"query": {
"type": "string",
"description": "Type of query (e.g., 'user_id', 'order_id')"
},
"value": {
"type": "string",
"description": "Value to lookup"
}
},
"required": ["query", "value"]
}
)
def register_rest_tools(gateway: MCPGateway, rest_api_url: str = "http://localhost:8000"):
"""Register REST API endpoints as MCP tools"""
# Stock Price Tool
gateway.register_rest_tool(
name="stock_price",
endpoint=f"{rest_api_url}/stock/price",
description="Get stock price information",
input_schema={
"type": "object",
"properties": {
"symbol": {
"type": "string",
"description": "Stock symbol (e.g., AAPL, GOOGL)"
}
},
"required": ["symbol"]
},
method="POST",
headers={"Content-Type": "application/json"}
)
# Current Time Tool
gateway.register_rest_tool(
name="current_time",
endpoint=f"{rest_api_url}/time/current",
description="Get current UTC time",
input_schema={
"type": "object",
"properties": {}
},
method="GET"
)
# Text Summarize Tool
gateway.register_rest_tool(
name="summarize_text",
endpoint=f"{rest_api_url}/text/summarize",
description="Summarize text content",
input_schema={
"type": "object",
"properties": {
"text": {
"type": "string",
"description": "Text to summarize"
}
},
"required": ["text"]
},
method="POST",
headers={"Content-Type": "application/json"}
)
def main():
"""Main registration function"""
import argparse
parser = argparse.ArgumentParser(description="Register tools with MCP Gateway")
parser.add_argument("--environment", default="dev", help="Environment name")
parser.add_argument("--region", default="us-east-1", help="AWS Region")
parser.add_argument("--rest-api-url", default=None, help="REST API URL (optional)")
args = parser.parse_args()
# Initialize gateway
gateway = MCPGateway(region=args.region)
print(f"Registering tools for environment: {args.environment}")
# Register Lambda tools
try:
register_lambda_tools(gateway, args.environment, args.region)
print("✓ Registered Lambda tools")
except Exception as e:
print(f"✗ Error registering Lambda tools: {e}")
# Register REST tools if URL provided
if args.rest_api_url:
try:
register_rest_tools(gateway, args.rest_api_url)
print(f"✓ Registered REST API tools from {args.rest_api_url}")
except Exception as e:
print(f"✗ Error registering REST tools: {e}")
# Print registered tools
print("\nRegistered tools:")
for tool in gateway.tool_registry.list_tools():
print(f" - {tool.name} ({tool.tool_type.value})")
print("\nNote: In production, tool registration should be persisted (e.g., in DynamoDB or SSM)")
if __name__ == "__main__":
main()
