# Sample code, software libraries, command line tools, proofs of concept, templates, or other related technology are provided as AWS Content or Third-Party Content under the AWS Customer Agreement, or the relevant written agreement between you and AWS (whichever applies). You should not use this AWS Content or Third-Party Content in your production accounts, or on production or other critical data. You are responsible for testing, securing, and optimizing the AWS Content or Third-Party Content, such as sample code, as appropriate for production grade use based on your specific quality control practices and standards. Deploying AWS Content or Third-Party Content may incur AWS charges for creating or using AWS chargeable resources, such as running Amazon EC2 instances or using Amazon S3 storage.

import json, boto3, os

# Read the KB ID and DataSource ID from Configuration
kb = os.environ["KB_ID"]
dsid = os.environ["KB_DATA_SOURCE_ID"]

# global variables - avoid creating a new client for every request
bedrock_agent = None

# This function is triggered by an Object modification event triggered in a 
# specific bucket in S3. Once triggered this function invokes the Start 
# Ingestion Job for the Knowledge base configured in the environment variables.
def handler(event, context):
    global bedrock_agent

    try:
        if bedrock_agent is None:
            # Instantiate the runtime object.
            bedrock_agent = boto3.client(
                service_name="bedrock-agent"
            )

        response = bedrock_agent.start_ingestion_job(
            knowledgeBaseId=kb,
            dataSourceId=dsid
        )

        response["ingestionJob"]["startedAt"] = str(
            response["ingestionJob"]["startedAt"])
        response["ingestionJob"]["updatedAt"] = str(
            response["ingestionJob"]["updatedAt"])
        
        print("Successfully invoked KB injestion job:")
        print(response)
        

        return {"statusCode": 200,
                "headers": {
                    "Content-Type": "application/json"},
                "body": json.dumps(response)
                }
    except Exception as e:
        print(e)
        print('Error getting object {} from bucket {}. Make sure they exist and your bucket is in the same region as this function.'.format(key, bucket))
        raise e