amazon-bedrock-agentcore-samples/02-use-cases/AWS-operations-agent/mcp-tool-lambda/lambda/mcp-tool-handler.py

"""
AWS Operations Agent Gateway Lambda Handler - Optimized Version
Handles AWS resource inspection tools via Strands Agent integration
Updated: 2025-07-15 - Fixed Strands import issue
"""
import json
import logging
import os
from datetime import datetime
from typing import Dict, Any, Optional

# Import Strands components at module level
try:
    from strands import Agent
    from strands.models import BedrockModel
    from strands_tools import use_aws, calculator, think, current_time
    STRANDS_AVAILABLE = True
    logging.info("Strands modules imported successfully")
except ImportError as e:
    STRANDS_AVAILABLE = False
    logging.error(f"Failed to import Strands modules: {e}")

# Configure logging
logger = logging.getLogger()
logger.setLevel(logging.INFO)

# AWS Service Tool Configurations
SERVICE_QUERIES = {
    # Core AWS Services
    'ec2_read_operations': "List and describe EC2 resources including instances, security groups, VPCs, subnets, and key pairs. Include instance states, types, and network configurations.",
    's3_read_operations': "List and describe S3 resources including buckets, bucket policies, lifecycle configurations, and access settings. Include bucket regions and creation dates.",
    'lambda_read_operations': "List and describe Lambda resources including functions, layers, aliases, and event source mappings. Include runtime, memory, timeout, and last modified information.",
    'cloudformation_read_operations': "List and describe CloudFormation resources including stacks, stack resources, stack events, and templates. Include stack status and creation times.",
    'iam_read_operations': "List and describe IAM resources including users, roles, policies, and groups. Include policy attachments and permissions (read-only operations only).",
    'rds_read_operations': "List and describe RDS and database resources including DB instances, clusters, snapshots, and parameter groups. Include engine types, versions, and status.",
    'cloudwatch_read_operations': "List and describe CloudWatch resources including metrics, alarms, log groups, and dashboards. Include alarm states and metric statistics.",
    'cost_explorer_read_operations': "Retrieve cost and billing information including cost breakdowns, usage reports, and budget information. Include cost trends and service-wise spending.",
    
    # Additional AWS Services
    'ecs_read_operations': "List and describe ECS resources including clusters, services, tasks, and task definitions. Include service status and task counts.",
    'eks_read_operations': "List and describe EKS resources including clusters, node groups, and add-ons. Include cluster status, versions, and configurations.",
    'sns_read_operations': "List and describe SNS resources including topics, subscriptions, and platform applications. Include topic ARNs and subscription counts.",
    'sqs_read_operations': "List and describe SQS resources including queues, queue attributes, and message statistics. Include queue types and visibility timeouts.",
    'dynamodb_read_operations': "List and describe DynamoDB resources including tables, indexes, and backups. Include table status, item counts, and throughput settings.",
    'route53_read_operations': "List and describe Route53 resources including hosted zones, record sets, and health checks. Include DNS configurations and routing policies.",
    'apigateway_read_operations': "List and describe API Gateway resources including REST APIs, resources, methods, and deployments. Include API stages and endpoint configurations.",
    'ses_read_operations': "List and describe SES resources including verified identities, configuration sets, and sending statistics. Include reputation metrics and quotas.",
    'bedrock_read_operations': "List and describe Bedrock resources including foundation models, model customization jobs, and inference profiles. Include model capabilities and availability.",
    'sagemaker_read_operations': "List and describe SageMaker resources including endpoints, models, training jobs, and notebook instances. Include status and configurations."
}

BASIC_TOOLS = ['hello_world', 'get_time']
AWS_SERVICE_TOOLS = list(SERVICE_QUERIES.keys())
ALL_TOOLS = BASIC_TOOLS + AWS_SERVICE_TOOLS


def extract_tool_name(context, event: Dict[str, Any]) -> Optional[str]:
    """Extract tool name from Gateway context or event."""
    
    # Try Gateway context first
    if hasattr(context, 'client_context') and context.client_context:
        if hasattr(context.client_context, 'custom') and context.client_context.custom:
            tool_name = context.client_context.custom.get('bedrockAgentCoreToolName')
            if tool_name and '___' in tool_name:
                # Remove namespace prefix (e.g., "aws-tools___hello_world" -> "hello_world")
                return tool_name.split('___', 1)[1]
            elif tool_name:
                return tool_name
    
    # Fallback to event-based extraction
    for field in ['tool_name', 'toolName', 'name', 'method', 'action', 'function']:
        if field in event:
            return event[field]
    
    # Infer from event structure
    if isinstance(event, dict):
        if 'name' in event and len(event) == 1:
            return 'hello_world'  # Typical hello_world structure
        elif len(event) == 0:
            return 'get_time'  # Empty args typically means get_time
    
    return None


def handle_hello_world(event: Dict[str, Any]) -> Dict[str, Any]:
    """Handle hello_world tool."""
    name = event.get('name', 'World')
    message = f"Hello, {name}! This message is from a Lambda function via AWS Operations Agent Gateway."
    
    return {
        'success': True,
        'result': message,
        'tool': 'hello_world',
        'timestamp': datetime.utcnow().isoformat() + 'Z'
    }


def handle_get_time(event: Dict[str, Any]) -> Dict[str, Any]:
    """Handle get_time tool."""
    current_time = datetime.utcnow().isoformat() + 'Z'
    
    return {
        'success': True,
        'result': f"Current UTC time: {current_time}",
        'tool': 'get_time',
        'timestamp': current_time
    }


def handle_aws_service_tool(tool_name: str, event: Dict[str, Any]) -> Dict[str, Any]:
    """Handle AWS service tools using Strands Agent."""
    
    # Check if Strands is available
    if not STRANDS_AVAILABLE:
        return {
            'success': False,
            'error': f"Strands modules not available for {tool_name}. Please check Lambda dependencies.",
            'tool': tool_name,
            'timestamp': datetime.utcnow().isoformat() + 'Z'
        }
    
    try:
        # Get the natural language query from the simplified schema
        user_query = event.get('query', '')
        if not user_query:
            return {
                'success': False,
                'error': f"Missing required 'query' parameter for {tool_name}",
                'tool': tool_name,
                'timestamp': datetime.utcnow().isoformat() + 'Z'
            }
        
        logger.info(f"Initializing Strands Agent for: {tool_name}")
        
        # Initialize Bedrock model
        bedrock_model = BedrockModel(
            region_name='us-east-1',
            model_id='us.anthropic.claude-3-7-sonnet-20250219-v1:0',
            temperature=0.1,
            system_prompt="""You are an AWS Operational Support Agent executing read-only AWS operations. You receive natural language queries and must perform the requested AWS operations efficiently.

CRITICAL: For any date-related queries, ALWAYS use the current_time tool first to get the current date before calculating date ranges.

For each query:
1. If the query involves dates or time periods (like "last month", "last 14 days", "current month"):
   - FIRST call current_time to get the current date
   - Calculate the correct date ranges based on the current date
   - Use YYYY-MM-DD format for all AWS API calls
2. Identify the specific AWS service and operation needed
3. Execute only the minimum required AWS calls to answer the query
4. Return structured, actionable results in a clear format
5. Focus on the specific information requested, avoid unnecessary details

Date Calculation Examples:
- Query: 'last month expenses' → Call current_time, then calculate previous month from current date
- Query: 'last 14 days costs' → Call current_time, then subtract 14 days from current date
- Query: 'current month breakdown' → Call current_time, then use current month start to current date

Service Operation Examples:
- Query: 'list running instances' → Use EC2 describe-instances with running state filter
- Query: 'count S3 buckets' → Use S3 list-buckets and return count
- Query: 'show failed stacks' → Use CloudFormation list-stacks with failed status filter

Always optimize for speed and relevance. Return concise, well-structured responses."""
        )
        
        # Create Strands Agent
        agent = Agent(model=bedrock_model, tools=[use_aws, calculator, think, current_time])
        
        # Build query
        # Get the natural language query from the simplified schema
        user_query = event.get('query', '')
        if not user_query:
            return {
                'success': False,
                'error': f"Missing required 'query' parameter for {tool_name}",
                'tool': tool_name,
                'timestamp': datetime.utcnow().isoformat() + 'Z'
            }
        
        # Build the final query combining service context with user query
        service_context = SERVICE_QUERIES.get(tool_name, f"AWS {tool_name.replace('_read_operations', '').upper()} service operations")
        final_query = f"AWS Service: {tool_name.replace('_read_operations', '').upper()}\nUser Request: {user_query}\nContext: {service_context}\n\nExecute this AWS operation and return structured results."
        
        logger.info(f"Executing simplified query for {tool_name}: {user_query}")
        
        # Execute query
        response = agent(final_query)
        
        # Extract response text
        response_text = ""
        if hasattr(response, 'message') and 'content' in response.message:
            for content_block in response.message['content']:
                if content_block.get('type') == 'text' or 'text' in content_block:
                    response_text += content_block.get('text', '')
        else:
            response_text = str(response)
        
        logger.info(f"Response length: {len(response_text)} characters")
        
        return {
            'success': True,
            'result': response_text,
            'tool': tool_name,
            'service': tool_name.replace('_read_operations', '').replace('_', '-'),
            'user_query': user_query,
            'timestamp': datetime.utcnow().isoformat() + 'Z'
        }
        
    except Exception as e:
        logger.error(f"AWS service tool error: {str(e)}")
        return {
            'success': False,
            'error': f"AWS Service Tool Error: {str(e)}",
            'tool': tool_name,
            'service': tool_name.replace('_read_operations', '').replace('_', '-'),
            'timestamp': datetime.utcnow().isoformat() + 'Z'
        }


def lambda_handler(event, context):
    """
    AWS Operations Agent Gateway Lambda Handler - Optimized Version
    
    Handles basic tools (hello_world, get_time) and AWS service tools
    via Strands Agent integration with comprehensive error handling.
    """
    logger.info("AWS Operations Agent Gateway Lambda Handler - START")
    logger.info(f"Event: {json.dumps(event, default=str)}")
    
    try:
        # Extract tool name
        tool_name = extract_tool_name(context, event)
        logger.info(f"Tool: {tool_name}")
        
        if not tool_name:
            return {
                'success': False,
                'error': 'Unable to determine tool name from context or event',
                'available_tools': ALL_TOOLS,
                'timestamp': datetime.utcnow().isoformat() + 'Z'
            }
        
        # Route to appropriate handler
        if tool_name == 'hello_world':
            return handle_hello_world(event)
        
        elif tool_name == 'get_time':
            return handle_get_time(event)
        
        elif tool_name in AWS_SERVICE_TOOLS:
            return handle_aws_service_tool(tool_name, event)
        
        else:
            # Unknown tool
            return {
                'success': False,
                'error': f"Unknown tool: {tool_name}",
                'available_tools': ALL_TOOLS,
                'total_tools': len(ALL_TOOLS),
                'categories': {
                    'basic': BASIC_TOOLS,
                    'aws_services': AWS_SERVICE_TOOLS
                },
                'tool': tool_name,
                'timestamp': datetime.utcnow().isoformat() + 'Z'
            }
    
    except Exception as e:
        logger.error(f"Handler error: {str(e)}")
        return {
            'success': False,
            'error': f"Internal error: {str(e)}",
            'timestamp': datetime.utcnow().isoformat() + 'Z'
        }
    
    finally:
        logger.info("AWS Operations Agent Gateway Lambda Handler - END")