response = bedrock_client.create_guardrail(
    name='BankingAssistantGuardrail',
    description='Guardrail for online banking assistant to help users with banking and account related questions',
    topicPolicyConfig={
        'topicsConfig': [
            {
                'name': 'Investment Advice',
                'definition': 'Investment advice refers to professional guidance or recommendations provided to individuals or entities regarding the management and allocation of their financial assets.',
                'examples': [
                    'Should I buy gold?',
                    'Is investing in stocks better than bonds?',
                    'When is it a good idea to invest in gold?',
                ],
                'type': 'DENY'
            },
        ]
    },
    blockedInputMessaging='Sorry, your query violates our usage policies. We do not provide investment advices. To discuss the best investment advice for your current situation, please contact us on (XXX) XXX-XXXX and we will be happy to support you.',
    blockedOutputsMessaging='Sorry, I am unable to reply. Please contact us on (XXX) XXX-XXXX and we will be happy to support you.',
)

!pip install --upgrade -q -r requirements.txt

import os
import time
import boto3
import logging
import pprint
import json

from knowledge_base import KnowledgeBasesForAmazonBedrock
from agent import AgentsForAmazonBedrock

#Clients
s3_client = boto3.client('s3')
sts_client = boto3.client('sts')
session = boto3.session.Session()
region = session.region_name
account_id = sts_client.get_caller_identity()["Account"]
bedrock_agent_client = boto3.client('bedrock-agent')
bedrock_agent_runtime_client = boto3.client('bedrock-agent-runtime')
bedrock_client = boto3.client('bedrock')
logging.basicConfig(format='[%(asctime)s] p%(process)s {%(filename)s:%(lineno)d} %(levelname)s - %(message)s', level=logging.INFO)
logger = logging.getLogger(__name__)
region, account_id

suffix = f"{region}-{account_id}"
agent_name = 'banking-assistant1q11'
knowledge_base_name = f'{agent_name}-kb'
knowledge_base_description = "Knowledge Base that provides FAQ documentation for the banking assistant agent"
agent_alias_name = "banking-agent-alias"
bucket_name = f'{agent_name}-{suffix}'
agent_bedrock_allow_policy_name = f"{agent_name}-ba"
agent_role_name = f'AmazonBedrockExecutionRoleForAgents_{agent_name}'
agent_foundation_model = "anthropic.claude-3-sonnet-20240229-v1:0"

agent_description = "Online Banking assistant agents"
agent_instruction = """
You are an agent designed to assist customers from the ANY_BANK Corporation with online banking queries. 
You ALWAYS reply politely and concise using ONLY the available information in the ba_kb knowledge base or the data retrieved via the banking-assistant action group.

You also add the name ANY_BANK Corporation to your first answer in a session. You should start with an acknowledgement of the customer's query and thanking the customer for contacting you.

Introduce yourself as the "ANY_BANK Corporation AI Assistant".

NEVER provide account balances or book appointments without first confirming the customer's user_id
"""

knowledge_base = KnowledgeBasesForAmazonBedrock()
kb_id, ds_id = knowledge_base.create_or_retrieve_knowledge_base(
    knowledge_base_name, knowledge_base_description, data_bucket_name=bucket_name
)

def upload_directory(path, bucket_name):
        for root,dirs,files in os.walk(path):
            for file in files:
                file_to_upload = os.path.join(root,file)
                print(f"uploading file {file_to_upload} to {bucket_name}")
                s3_client.upload_file(file_to_upload,bucket_name,file)

upload_directory("dataset", bucket_name)

# ensure that the kb is available
time.sleep(30)
# sync knowledge base
knowledge_base.synchronize_data(kb_id, ds_id)

time.sleep(30)

response = bedrock_agent_runtime_client.retrieve_and_generate(
    input={
        "text": "Should I invest in bitcoin"
    },
    retrieveAndGenerateConfiguration={
        "type": "KNOWLEDGE_BASE",
        "knowledgeBaseConfiguration": {
            'knowledgeBaseId': kb_id,
            "modelArn": "arn:aws:bedrock:{}::foundation-model/{}".format(region, agent_foundation_model),
            "retrievalConfiguration": {
                "vectorSearchConfiguration": {
                    "numberOfResults":1
                } 
            }
        }
    }
)

print(response['output']['text'],end='\n'*2)

response_ret = bedrock_agent_runtime_client.retrieve(
    knowledgeBaseId=kb_id, 
    nextToken='string',
    retrievalConfiguration={
        "vectorSearchConfiguration": {
            "numberOfResults":5,
        } 
    },
    retrievalQuery={
        'text': 'What is my account value?'
    }
)

def response_print(retrieve_resp):
#structure 'retrievalResults': list of contents. Each list has content, location, score, metadata
    for num,chunk in enumerate(response_ret['retrievalResults'],1):
        print(f'Chunk {num}: ',chunk['content']['text'],end='\n'*2)
        print(f'Chunk {num} Location: ',chunk['location'],end='\n'*2)
        print(f'Chunk {num} Score: ',chunk['score'],end='\n'*2)
        print(f'Chunk {num} Metadata: ',chunk['metadata'],end='\n'*2)

response_print(response_ret)

kb_arn =  f"arn:aws:bedrock:{region}:{account_id}:knowledge-base/{kb_id}"
agents = AgentsForAmazonBedrock()
agent_id = agents.create_agent(agent_name, agent_description, agent_instruction, [agent_foundation_model], kb_arns=[kb_arn])
print(agent_id)