lambda_function.py

# lambda_handler.py
import json
import boto3
import os
import uuid
import base64
import yaml
from botocore.exceptions import ClientError, NoCredentialsError
from util import process_response, invoke_agent_helper, add_file_to_session_state, json_serializable

def load_configurations():
    """Load file type mappings and S3 configuration from external files"""
    try:
        # Load file type mappings
        with open('file_type_mappings.json', 'r') as f:
            file_type_mappings = json.load(f)
        
        # Load S3 configuration
        with open('s3_config.yaml', 'r') as f:
            s3_config = yaml.safe_load(f)
        
        return file_type_mappings, s3_config
    except Exception as e:
        print(f"Error loading configurations: {str(e)}")
        raise

def lambda_handler(event, context):
    try:
        print("=== Starting Lambda execution ===")
        print(f"Event type: {type(event)}")
        print(f"Raw event: {json.dumps(event, default=json_serializable)}")
        
        # Handle different input formats
        try:
            if isinstance(event, str):
                print("Event is string, parsing as JSON")
                body = json.loads(event)
            elif isinstance(event, dict):
                if 'body' in event and isinstance(event['body'], str):
                    print("Event has string body, parsing as JSON")
                    body = json.loads(event['body'])
                else:
                    print("Using event as body directly")
                    body = event
            else:
                print(f"Unexpected event type: {type(event)}")
                return {
                    'statusCode': 400,
                    'body': json.dumps({
                        'error': 'Invalid request format',
                        'details': 'Request must be valid JSON or contain a JSON body field.'
                    })
                }
        except json.JSONDecodeError as e:
            print(f"JSON decode error: {str(e)}")
            return {
                'statusCode': 400,
                'body': json.dumps({
                    'error': 'Invalid JSON format',
                    'details': f'Request body contains invalid JSON: {str(e)}'
                })
            }
            
        print(f"Parsed body: {json.dumps(body, default=json_serializable)}")
        
        # Get query from the parsed body
        query = body.get('query', '').strip()
        print(f"Extracted query: {query}")
        
        # Get business_id from the parsed body - handle both string and list formats
        business_id_input = body.get('business_id', '')
        if isinstance(business_id_input, list):
            business_ids = [str(bid).strip() for bid in business_id_input if str(bid).strip()]
        elif isinstance(business_id_input, str):
            business_ids = [business_id_input.strip()] if business_id_input.strip() else []
        else:
            business_ids = []
            
        print(f"Extracted business_ids: {business_ids}")
        
        # Get information_type parameter - default to both Orders and Menus if not specified
        information_type_input = body.get('information_type', ['Payment', 'Menus'])
        
        # Handle different input formats for information_type
        if isinstance(information_type_input, str):
            # Single string input - convert to list
            information_types = [information_type_input.strip().title()]
        elif isinstance(information_type_input, list):
            # List input - normalize case
            information_types = [item.strip().title() for item in information_type_input if item.strip()]
        else:
            # Invalid format - use default
            information_types = ['Payment', 'Menus']
        
        # Validate information_types
        valid_types = ['Payment', 'Menus']
        information_types = [t for t in information_types if t in valid_types]
        
        # If no valid types after filtering, use default
        if not information_types:
            information_types = ['Payment', 'Menus']
            
        print(f"Information types to load: {information_types}")
        
        # Validate business_ids
        if not business_ids:
            print("Error: No business IDs provided")
            return {
                'statusCode': 400,
                'body': json.dumps({
                    'error': 'Missing business ID',
                    'details': 'No business IDs provided. Please contact support to set up your account.'
                })
            }
        
        # Validate query
        if not query:
            print("Error: Empty query")
            return {
                'statusCode': 400,
                'body': json.dumps({
                    'error': 'Empty query',
                    'details': 'Query cannot be empty. Please provide a valid question or request.'
                })
            }
        
        # Check query length (reasonable limit)
        if len(query) > 10000:
            print("Error: Query too long")
            return {
                'statusCode': 400,
                'body': json.dumps({
                    'error': 'Query too long',
                    'details': 'Query must be less than 10,000 characters. Please shorten your request.'
                })
            }
        
        # Initialize Bedrock Agent Runtime client
        try:
            print("Initializing Bedrock Agent Runtime client")
            bedrock_agent_runtime = boto3.client('bedrock-agent-runtime')
        except NoCredentialsError:
            print("Error: AWS credentials not found")
            return {
                'statusCode': 500,
                'body': json.dumps({
                    'error': 'Service configuration error',
                    'details': 'Unable to authenticate with AWS services. Please contact support.'
                })
            }
        except Exception as e:
            print(f"Error initializing Bedrock client: {str(e)}")
            return {
                'statusCode': 500,
                'body': json.dumps({
                    'error': 'Service initialization error',
                    'details': 'Unable to initialize AI service. Please try again or contact support.'
                })
            }
        
        # Get parameters from request body first, then fall back to environment variables
        agent_id = body.get('agent_id', '').strip() or os.environ.get('AGENT_ID', 'B6GA2N90WU')
        agent_alias_id = body.get('agent_alias_id', '').strip() or os.environ.get('AGENT_ALIAS_ID', 'VTRJVBOA8A')
        
        # Validate agent parameters
        if not agent_id or not agent_alias_id:
            print("Error: Missing required agent parameters")
            return {
                'statusCode': 500,
                'body': json.dumps({
                    'error': 'Service configuration error',
                    'details': 'AI agent configuration is missing. Please contact support.'
                })
            }
        
        print(f"Using agent_id: {agent_id}, agent_alias_id: {agent_alias_id}")
        
        # Generate a session ID if not provided or invalid
        session_id = body.get('session_id', '').strip()
        if not session_id or len(session_id) < 2:
            session_id = str(uuid.uuid4())
        print(f"Using session_id: {session_id}")
            
        enable_trace = body.get('enable_trace', False)
        end_session = body.get('end_session', False)
        print(f"enable_trace: {enable_trace}, end_session: {end_session}")
        
        # Only include memoryId if it's explicitly provided
        memory_id = body.get('memory_id', '').strip()
        if memory_id:
            print(f"Using memory_id: {memory_id}")
        
        # Load configurations
        try:
            file_type_mappings, s3_config = load_configurations()
            print(f"Loaded file type mappings: {file_type_mappings}")
            print(f"Loaded S3 config: {s3_config}")
        except Exception as e:
            print(f"Error loading configurations: {str(e)}")
            return {
                'statusCode': 500,
                'body': json.dumps({
                    'error': 'Configuration error',
                    'details': 'Unable to load required configuration files. Please contact support.'
                })
            }
        
        # Initialize session state
        session_state = {'files': []}
        
        # Add business-specific files to session state for all business IDs
        print(f"Adding business-specific files for {len(business_ids)} business(es)")
        total_files_added = 0
        failed_businesses = []
        
        # Define the constant date for file naming
        file_date = "20250528"
        
        for business_id in business_ids:
            print(f"Processing files for business ID: {business_id}")
            bucket_name = s3_config['s3']['bucket_name']
            base_path = s3_config['s3']['base_path']
            base_s3_path = f"s3://{bucket_name}/{base_path}/{business_id}"
            print(f"Using S3 base path: {base_s3_path}")
            
            # Define business-specific files based on information_type
            business_files = {}
            
            # Add files for each requested information type
            for info_type in information_types:
                if info_type in file_type_mappings:
                    file_prefix = file_type_mappings[info_type]
                    file_name = f"{file_prefix}_{file_date}.csv"
                    custom_name = f"{business_id}_{file_name}"
                    file_path = f"{base_s3_path}/{file_name}"
                    
                    business_files[custom_name] = file_path
                    print(f"Mapped {info_type} -> {file_name}")
                else:
                    print(f"Warning: No file mapping found for information type: {info_type}")
            
            print(f"Files to load for business {business_id}: {list(business_files.keys())}")
            
            business_files_added = 0
            for custom_file_name, file_path in business_files.items():
                try:
                    print(f"Adding file to session state: {custom_file_name} - {file_path}")
                    session_state = add_file_to_session_state(file_path, session_state=session_state, custom_name=custom_file_name)
                    business_files_added += 1
                    total_files_added += 1
                except Exception as e:
                    print(f"Warning: Could not add file {custom_file_name} for business {business_id}: {str(e)}")
                    # Continue with other files rather than failing completely
            
            if business_files_added == 0:
                print(f"Warning: No data files could be loaded for business ID: {business_id}")
                failed_businesses.append(business_id)
            else:
                print(f"Successfully added {business_files_added} out of {len(business_files)} files for business {business_id}")
        
        # Check if we have any files at all
        if total_files_added == 0:
            print("Error: No business data files could be loaded for any business")
            return {
                'statusCode': 404,
                'body': json.dumps({
                    'error': 'Business data not found',
                    'details': f'No data files found for business IDs {business_ids} with information types {information_types}. Please contact support to set up your data.'
                })
            }
        
        # Log summary
        successful_businesses = [bid for bid in business_ids if bid not in failed_businesses]
        print(f"Summary: Successfully loaded files for {len(successful_businesses)} business(es): {successful_businesses}")
        print(f"Information types loaded: {information_types}")
        if failed_businesses:
            print(f"Warning: Failed to load files for {len(failed_businesses)} business(es): {failed_businesses}")
        
        # For the primary business_id (used for file storage paths), use the first successful one
        primary_business_id = successful_businesses[0] if successful_businesses else business_ids[0]
        print(f"Using primary business_id for response processing: {primary_business_id}")
        
        # Prepare invoke_agent parameters
        invoke_params = {
            'inputText': query,
            'agentId': agent_id,
            'agentAliasId': agent_alias_id,
            'sessionId': session_id,
            'enableTrace': enable_trace,
            'endSession': end_session,
            'sessionState': session_state
        }
        
        # Only add memoryId if it's provided and not empty
        if memory_id:
            invoke_params['memoryId'] = memory_id
        
        print(f"Invoke parameters: {json.dumps(invoke_params, default=json_serializable)}")
        
        # Invoke Bedrock agent with timeout handling
        print("Invoking Bedrock agent")
        try:
            response = bedrock_agent_runtime.invoke_agent(**invoke_params)
        except ClientError as e:
            error_code = e.response['Error']['Code']
            print(f"Bedrock client error: {error_code} - {e.response['Error']['Message']}")
            
            if error_code == 'AccessDeniedException':
                return {
                    'statusCode': 403,
                    'body': json.dumps({
                        'error': 'Access denied',
                        'details': 'Your account does not have permission to access the AI service. Please contact support.'
                    })
                }
            elif error_code == 'ResourceNotFoundException':
                return {
                    'statusCode': 404,
                    'body': json.dumps({
                        'error': 'AI service not available',
                        'details': 'The AI agent is currently not available. Please contact support.'
                    })
                }
            elif error_code == 'ThrottlingException':
                return {
                    'statusCode': 429,
                    'body': json.dumps({
                        'error': 'Service busy',
                        'details': 'The AI service is currently busy. Please wait a moment and try again.'
                    })
                }
            elif error_code == 'ValidationException':
                return {
                    'statusCode': 400,
                    'body': json.dumps({
                        'error': 'Invalid request',
                        'details': 'Your request contains invalid parameters. Please check your input and try again.'
                    })
                }
            elif error_code == 'ServiceQuotaExceededException':
                return {
                    'statusCode': 429,
                    'body': json.dumps({
                        'error': 'Service quota exceeded',
                        'details': 'You have exceeded your usage quota. Please contact support or try again later.'
                    })
                }
            else:
                return {
                    'statusCode': 502,
                    'body': json.dumps({
                        'error': 'AI service error',
                        'details': f'The AI service encountered an error. Please try again or contact support if the issue persists.'
                    })
                }
        except Exception as e:
            print(f"Unexpected error invoking Bedrock agent: {str(e)}")
            return {
                'statusCode': 500,
                'body': json.dumps({
                    'error': 'Service unavailable',
                    'details': 'The AI service is currently unavailable. Please try again in a few moments.'
                })
            }
        
        print("Processing Bedrock agent response")
        # Pass primary_business_id to process_response
        try:
            return process_response(response, business_id=primary_business_id, enable_trace=enable_trace)
        except Exception as e:
            print(f"Error processing response: {str(e)}")
            return {
                'statusCode': 500,
                'body': json.dumps({
                    'error': 'Response processing error',
                    'details': 'An error occurred while processing the AI response. Please try again.'
                })
            }
        
    except Exception as e:
        print(f"Unexpected error in lambda_handler: {str(e)}")
        import traceback
        print(f"Traceback: {traceback.format_exc()}")
        return {
            'statusCode': 500,
            'body': json.dumps({
                'error': 'Internal server error',
                'details': 'An unexpected error occurred. Please try again or contact support if the issue persists.'
            })
        }