10_create_lambda.py

#!/usr/bin/env python3
"""
Script to create Lambda function for order lookup.

This Lambda function will be exposed as a tool through AgentCore Gateway,
allowing agents to look up order details by order ID.
"""

import boto3
import json
import time
import zipfile
import io
import sys

# Configuration
REGION = "us-west-2"
FUNCTION_NAME = "OrderLookupFunction"
ROLE_NAME = f"OrderLookupLambdaRole-{int(time.time())}"

print("=" * 80)
print("LAMBDA FUNCTION SETUP FOR ORDER LOOKUP")
print("=" * 80)
print(f"Region: {REGION}")
print(f"Function Name: {FUNCTION_NAME}")
print("=" * 80)
print()

# Initialize clients
lambda_client = boto3.client('lambda', region_name=REGION)
iam_client = boto3.client('iam', region_name=REGION)
sts_client = boto3.client('sts', region_name=REGION)

try:
    # Get account ID
    account_id = sts_client.get_caller_identity()['Account']
    print(f"AWS Account: {account_id}")
    print()
    
    # Step 1: Create Lambda execution role
    print("Step 1: Creating Lambda execution role...")
    print("-" * 80)
    
    trust_policy = {
        "Version": "2012-10-17",
        "Statement": [
            {
                "Effect": "Allow",
                "Principal": {
                    "Service": "lambda.amazonaws.com"
                },
                "Action": "sts:AssumeRole"
            }
        ]
    }
    
    try:
        role_response = iam_client.create_role(
            RoleName=ROLE_NAME,
            AssumeRolePolicyDocument=json.dumps(trust_policy),
            Description="Execution role for OrderLookupFunction Lambda"
        )
        role_arn = role_response['Role']['Arn']
        print(f"✓ Role created: {ROLE_NAME}")
        print(f"✓ Role ARN: {role_arn}")
        
        # Attach basic Lambda execution policy
        iam_client.attach_role_policy(
            RoleName=ROLE_NAME,
            PolicyArn='arn:aws:iam::aws:policy/service-role/AWSLambdaBasicExecutionRole'
        )
        print("✓ Basic execution policy attached")
        
        # Wait for role propagation
        print("⏳ Waiting 10 seconds for role propagation...")
        time.sleep(10)
        
    except iam_client.exceptions.EntityAlreadyExistsException:
        role = iam_client.get_role(RoleName=ROLE_NAME)
        role_arn = role['Role']['Arn']
        print(f"⚠️  Role already exists: {ROLE_NAME}")
        print(f"✓ Using existing role ARN: {role_arn}")
    
    print()
    
    # Step 2: Create Lambda function code
    print("Step 2: Creating Lambda function code...")
    print("-" * 80)
    
    lambda_code = '''import json
from datetime import datetime, timedelta

# Mock order database
ORDERS = {
    "ORD-001": {
        "order_id": "ORD-001",
        "product_name": "Dell XPS 15 Laptop",
        "purchase_date": (datetime.now() - timedelta(days=15)).strftime("%Y-%m-%d"),
        "amount": 1299.99,
        "status": "delivered",
        "category": "electronics"
    },
    "ORD-002": {
        "order_id": "ORD-002",
        "product_name": "iPhone 12",
        "purchase_date": (datetime.now() - timedelta(days=45)).strftime("%Y-%m-%d"),
        "amount": 799.99,
        "status": "delivered",
        "category": "electronics"
    },
    "ORD-003": {
        "order_id": "ORD-003",
        "product_name": "Samsung Galaxy Tab (Defective)",
        "purchase_date": (datetime.now() - timedelta(days=10)).strftime("%Y-%m-%d"),
        "amount": 449.99,
        "status": "delivered",
        "category": "electronics",
        "defective": True
    }
}

def calculate_return_eligibility(order):
    """Calculate if order is eligible for return"""
    purchase_date = datetime.strptime(order["purchase_date"], "%Y-%m-%d")
    days_since_purchase = (datetime.now() - purchase_date).days
    
    # 30-day return window for electronics
    return_window = 30
    days_remaining = return_window - days_since_purchase
    
    # Defective items always eligible
    if order.get("defective", False):
        return {
            "eligible": True,
            "reason": "Defective item - eligible for full refund",
            "days_remaining": max(0, days_remaining)
        }
    
    if days_remaining > 0:
        return {
            "eligible": True,
            "reason": f"Within {return_window}-day return window",
            "days_remaining": days_remaining
        }
    else:
        return {
            "eligible": False,
            "reason": f"Return window ({return_window} days) has expired",
            "days_remaining": 0
        }

def lambda_handler(event, context):
    """
    Lambda handler for order lookup.
    
    Expected event format:
    {
        "order_id": "ORD-001"
    }
    """
    try:
        # Extract order_id from event
        order_id = event.get("order_id", "").strip().upper()
        
        if not order_id:
            return {
                "statusCode": 400,
                "body": json.dumps({
                    "error": "Missing required parameter: order_id"
                })
            }
        
        # Look up order
        order = ORDERS.get(order_id)
        
        if not order:
            return {
                "statusCode": 404,
                "body": json.dumps({
                    "error": f"Order not found: {order_id}",
                    "available_orders": list(ORDERS.keys())
                })
            }
        
        # Calculate return eligibility
        eligibility = calculate_return_eligibility(order)
        
        # Build response
        response_data = {
            "order_id": order["order_id"],
            "product_name": order["product_name"],
            "purchase_date": order["purchase_date"],
            "amount": order["amount"],
            "status": order["status"],
            "category": order["category"],
            "return_eligibility": eligibility
        }
        
        if order.get("defective"):
            response_data["notes"] = "Item marked as defective"
        
        return {
            "statusCode": 200,
            "body": json.dumps(response_data)
        }
        
    except Exception as e:
        return {
            "statusCode": 500,
            "body": json.dumps({
                "error": f"Internal error: {str(e)}"
            })
        }
'''
    
    # Create deployment package
    zip_buffer = io.BytesIO()
    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zip_file:
        zip_file.writestr('lambda_function.py', lambda_code)
    zip_buffer.seek(0)
    
    print("✓ Lambda code created with mock order data:")
    print("  - ORD-001: Recent laptop (15 days old)")
    print("  - ORD-002: Old phone (45 days old)")
    print("  - ORD-003: Defective tablet (10 days old)")
    print()
    
    # Step 3: Create Lambda function
    print("Step 3: Creating Lambda function...")
    print("-" * 80)
    
    try:
        function_response = lambda_client.create_function(
            FunctionName=FUNCTION_NAME,
            Runtime='python3.12',
            Role=role_arn,
            Handler='lambda_function.lambda_handler',
            Code={'ZipFile': zip_buffer.read()},
            Description='Order lookup function for AgentCore Gateway',
            Timeout=30,
            MemorySize=128
        )
        
        function_arn = function_response['FunctionArn']
        print(f"✓ Lambda function created: {FUNCTION_NAME}")
        print(f"✓ Function ARN: {function_arn}")
        
    except lambda_client.exceptions.ResourceConflictException:
        # Function already exists, update it
        print(f"⚠️  Function already exists, updating code...")
        zip_buffer.seek(0)
        lambda_client.update_function_code(
            FunctionName=FUNCTION_NAME,
            ZipFile=zip_buffer.read()
        )
        
        function_response = lambda_client.get_function(FunctionName=FUNCTION_NAME)
        function_arn = function_response['Configuration']['FunctionArn']
        print(f"✓ Function updated: {FUNCTION_NAME}")
        print(f"✓ Function ARN: {function_arn}")
    
    print()
    
    # Step 4: Create tool schema for Gateway
    print("Step 4: Creating tool schema for Gateway...")
    print("-" * 80)
    
    tool_schema = [
        {
            "name": "lookup_order",
            "description": "Look up order details by order ID. Returns order information including product name, purchase date, amount, and return eligibility status.",
            "inputSchema": {
                "type": "object",
                "properties": {
                    "order_id": {
                        "type": "string",
                        "description": "The order ID to look up (e.g., ORD-001, ORD-002, ORD-003)"
                    }
                },
                "required": ["order_id"]
            }
        }
    ]
    
    print("✓ Tool schema created:")
    print(f"  Tool name: lookup_order")
    print(f"  Input: order_id (string)")
    print(f"  Output: order details + return eligibility")
    print()
    
    # Step 5: Save configuration
    print("Step 5: Saving configuration to lambda_config.json...")
    print("-" * 80)
    
    config = {
        "function_name": FUNCTION_NAME,
        "function_arn": function_arn,
        "role_name": ROLE_NAME,
        "role_arn": role_arn,
        "tool_schema": tool_schema,
        "region": REGION,
        "sample_orders": ["ORD-001", "ORD-002", "ORD-003"]
    }
    
    with open('lambda_config.json', 'w') as f:
        json.dump(config, f, indent=2)
    
    print("✓ Configuration saved to lambda_config.json")
    print()
    
    # Summary
    print("=" * 80)
    print("✅ LAMBDA FUNCTION SETUP COMPLETE!")
    print("=" * 80)
    print("\nConfiguration Summary:")
    print(f"  Function Name: {FUNCTION_NAME}")
    print(f"  Function ARN: {function_arn}")
    print(f"  Role ARN: {role_arn}")
    print(f"  Tool Name: lookup_order")
    print()
    print("Sample Orders:")
    print("  - ORD-001: Dell XPS 15 Laptop (recent, eligible)")
    print("  - ORD-002: iPhone 12 (old, not eligible)")
    print("  - ORD-003: Samsung Galaxy Tab (defective, always eligible)")
    print()
    print("Tool Capabilities:")
    print("  ✓ Look up order by ID")
    print("  ✓ Return product details")
    print("  ✓ Calculate return eligibility")
    print("  ✓ Check days remaining in return window")
    print("  ✓ Handle defective items")
    print()
    print("Next Steps:")
    print("  1. Add this Lambda as a target to your Gateway")
    print("  2. Gateway will expose lookup_order as a tool")
    print("  3. Agents can call this tool to look up orders")
    print("=" * 80)
    
except Exception as e:
    print(f"\n❌ Error during Lambda setup: {e}")
    import traceback
    traceback.print_exc()
    sys.exit(1)