12_add_lambda_to_gateway.py

#!/usr/bin/env python3
"""
Script to add Lambda target to AgentCore Gateway.

This script registers the OrderLookupFunction as a Gateway target,
making it available as an MCP tool for agents.

Prerequisites:
- gateway_config.json (from gateway creation)
- lambda_config.json (from Lambda creation)
"""

import json
import boto3
import sys

print("=" * 80)
print("ADD LAMBDA TARGET TO GATEWAY")
print("=" * 80)
print()

try:
    # Step 1: Load configuration files
    print("Step 1: Loading configuration files...")
    print("-" * 80)
    
    try:
        with open('gateway_config.json') as f:
            gateway_config = json.load(f)
        print("✓ Loaded gateway_config.json")
    except FileNotFoundError:
        print("❌ Error: gateway_config.json not found")
        print("   Run 11_create_gateway.py first")
        sys.exit(1)
    
    try:
        with open('lambda_config.json') as f:
            lambda_config = json.load(f)
        print("✓ Loaded lambda_config.json")
    except FileNotFoundError:
        print("❌ Error: lambda_config.json not found")
        print("   Run 10_create_lambda.py first")
        sys.exit(1)
    
    print()
    
    # Step 2: Initialize AgentCore client
    print("Step 2: Initializing AgentCore client...")
    print("-" * 80)
    
    gateway_client = boto3.client("bedrock-agentcore-control", region_name='us-west-2')
    print("✓ AgentCore control plane client initialized")
    print()
    
    # Step 3: Prepare Lambda target configuration
    print("Step 3: Preparing Lambda target configuration...")
    print("-" * 80)
    
    lambda_arn = lambda_config["function_arn"]
    tool_schema = lambda_config["tool_schema"]
    
    # Build Lambda target configuration with MCP protocol
    lambda_target_config = {
        "mcp": {
            "lambda": {
                "lambdaArn": lambda_arn,
                "toolSchema": {
                    "inlinePayload": tool_schema
                }
            }
        }
    }
    
    # Use gateway's IAM role for Lambda invocation
    credential_config = [{"credentialProviderType": "GATEWAY_IAM_ROLE"}]
    
    print(f"✓ Target configuration prepared:")
    print(f"  Lambda ARN: {lambda_arn}")
    print(f"  Tool Name: {tool_schema[0]['name']}")
    print(f"  Credential: Gateway IAM Role")
    print()
    
    # Step 4: Add Lambda permission for Gateway to invoke
    print("Step 4: Adding Lambda permission for Gateway...")
    print("-" * 80)
    
    lambda_client = boto3.client('lambda', region_name='us-west-2')
    
    try:
        lambda_client.add_permission(
            FunctionName=lambda_config["function_name"],
            StatementId=f'AllowGatewayInvoke-{gateway_config["gateway_id"]}',
            Action='lambda:InvokeFunction',
            Principal='bedrock-agentcore.amazonaws.com',
            SourceArn=gateway_config["gateway_arn"]
        )
        print("✓ Lambda permission added for Gateway")
    except lambda_client.exceptions.ResourceConflictException:
        print("⚠️  Permission already exists, continuing...")
    
    print()
    
    # Step 5: Create Gateway target
    print("Step 5: Adding Lambda as Gateway target...")
    print("-" * 80)
    print(f"  Gateway ID: {gateway_config['gateway_id']}")
    print(f"  Target Name: OrderLookup")
    print("⏳ This may take a few seconds...")
    
    create_response = gateway_client.create_gateway_target(
        gatewayIdentifier=gateway_config["gateway_id"],
        name="OrderLookup",
        description="Order lookup service for returns and refunds agent",
        targetConfiguration=lambda_target_config,
        credentialProviderConfigurations=credential_config
    )
    
    target_id = create_response["targetId"]
    
    print(f"✓ Lambda target added successfully!")
    print(f"  Target ID: {target_id}")
    print()
    
    # Step 6: Update gateway config with target info
    print("Step 6: Updating gateway_config.json...")
    print("-" * 80)
    
    gateway_config["target_id"] = target_id
    gateway_config["target_name"] = "OrderLookup"
    gateway_config["lambda_arn"] = lambda_arn
    
    with open('gateway_config.json', 'w') as f:
        json.dump(gateway_config, f, indent=2)
    
    print("✓ Configuration updated")
    print()
    
    # Summary
    print("=" * 80)
    print("✅ LAMBDA TARGET ADDED TO GATEWAY!")
    print("=" * 80)
    print("\nTarget Details:")
    print(f"  Target Name: OrderLookup")
    print(f"  Target ID: {target_id}")
    print(f"  Lambda Function: {lambda_config['function_name']}")
    print(f"  Lambda ARN: {lambda_arn}")
    print()
    print("Tool Exposed:")
    print(f"  Tool Name: {tool_schema[0]['name']}")
    print(f"  Description: {tool_schema[0]['description']}")
    print(f"  Input: order_id (string)")
    print()
    print("Gateway Configuration:")
    print(f"  Gateway URL: {gateway_config['gateway_url']}")
    print(f"  Protocol: MCP")
    print(f"  Authentication: OAuth 2.0")
    print()
    print("What This Means:")
    print("  ✓ Gateway now exposes lookup_order as an MCP tool")
    print("  ✓ Agents can discover this tool via Gateway")
    print("  ✓ Tool calls are translated to Lambda invocations")
    print("  ✓ Gateway handles authentication and authorization")
    print()
    print("Next Steps:")
    print("  1. Create an agent that connects to this Gateway")
    print("  2. Agent will discover the lookup_order tool")
    print("  3. Agent can call the tool to look up orders")
    print("=" * 80)
    
except Exception as e:
    print(f"\n❌ Error adding Lambda target: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)