From 879b8fe59987cfda7d3b2a89ed5456f2894bdcdb Mon Sep 17 00:00:00 2001 From: nalbam Date: Tue, 13 Aug 2024 19:44:39 +0900 Subject: [PATCH 1/7] add env KB_ID --- handler.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/handler.py b/handler.py index 8047199..634c217 100644 --- a/handler.py +++ b/handler.py @@ -21,6 +21,9 @@ # Keep track of conversation history by thread and user DYNAMODB_TABLE_NAME = os.environ.get("DYNAMODB_TABLE_NAME", "gurumi-ai-bot-context") +# Amazon Bedrock Knowledge Base ID +KB_ID = os.environ.get("KB_ID", "None") + # Amazon Bedrock Model ID MODEL_ID_TEXT = os.environ.get("MODEL_ID_TEXT", "anthropic.claude-3") MODEL_ID_IMAGE = os.environ.get("MODEL_ID_IMAGE", "stability.stable-diffusion-xl") From 632e8938736bc9829e7c701348f981b6110c9a81 Mon Sep 17 00:00:00 2001 From: nalbam Date: Wed, 14 Aug 2024 11:53:16 +0900 Subject: [PATCH 2/7] add invoke_knowledge_base --- bedrock/invoke_knowledge_base.py | 99 ++++++++++++++++++++++++++++++ bedrock/invoke_stable_diffusion.py | 2 +- 2 files changed, 100 insertions(+), 1 deletion(-) create mode 100644 bedrock/invoke_knowledge_base.py diff --git a/bedrock/invoke_knowledge_base.py b/bedrock/invoke_knowledge_base.py new file mode 100644 index 0000000..bb92ab7 --- /dev/null +++ b/bedrock/invoke_knowledge_base.py @@ -0,0 +1,99 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- + +import argparse +import json +import boto3 +import pprint + +from botocore.client import Config + +pp = pprint.PrettyPrinter(indent=2) + +bedrock_config = Config( + connect_timeout=120, read_timeout=120, retries={"max_attempts": 0} +) +bedrock_client = boto3.client("bedrock-runtime", region_name="us-east-1") +bedrock_agent_client = boto3.client( + "bedrock-agent-runtime", region_name="us-east-1", config=bedrock_config +) +# boto3_session = boto3.session.Session() +# region_name = boto3_session.region_name + +model_id = "anthropic.claude-v2:1" # try with both claude instant as well as claude-v2. for claude v2 - "anthropic.claude-v2" +region_id = "us-east-1" # replace it with the region you're running sagemaker notebook + +SYSTEM_MESSAGE = "답변은 한국어 해요체로 해요." + + +# def parse_args(): +# p = argparse.ArgumentParser(description="invoke_knowledge_base") +# p.add_argument("-p", "--prompt", default="안녕", help="prompt") +# p.add_argument("-d", "--debug", default="False", help="debug") +# return p.parse_args() + + +def retrieve(query, kbId, numberOfResults=5): + return bedrock_agent_client.retrieve( + retrievalQuery={"text": query}, + knowledgeBaseId=kbId, + retrievalConfiguration={ + "vectorSearchConfiguration": { + "numberOfResults": numberOfResults, + # "overrideSearchType": "HYBRID", # optional + } + }, + ) + + +def retrieveAndGenerate( + input, + kbId, + sessionId=None, + model_id="anthropic.claude-v2:1", + region_id="us-east-1", +): + model_arn = f"arn:aws:bedrock:{region_id}::foundation-model/{model_id}" + if sessionId: + return bedrock_agent_client.retrieve_and_generate( + input={"text": input}, + retrieveAndGenerateConfiguration={ + "type": "KNOWLEDGE_BASE", + "knowledgeBaseConfiguration": { + "knowledgeBaseId": kbId, + "modelArn": model_arn, + }, + }, + sessionId=sessionId, + ) + else: + return bedrock_agent_client.retrieve_and_generate( + input={"text": input}, + retrieveAndGenerateConfiguration={ + "type": "KNOWLEDGE_BASE", + "knowledgeBaseConfiguration": { + "knowledgeBaseId": kbId, + "modelArn": model_arn, + }, + }, + ) + + +def main(): + # args = parse_args() + + kb_id = "knowledge-base-whitepaper" + + query = "Please tell me about the Kontrol." + + # response = retrieve(query, kb_id, 3) + # retrievalResults = response["retrievalResults"] + # pp.pprint(retrievalResults) + + response = retrieveAndGenerate(query, kb_id, model_id=model_id, region_id=region_id) + generated_text = response["output"]["text"] + pp.pprint(generated_text) + + +if __name__ == "__main__": + main() diff --git a/bedrock/invoke_stable_diffusion.py b/bedrock/invoke_stable_diffusion.py index 5de37c3..fc96bc6 100644 --- a/bedrock/invoke_stable_diffusion.py +++ b/bedrock/invoke_stable_diffusion.py @@ -11,7 +11,7 @@ def parse_args(): - p = argparse.ArgumentParser(description="invoke_claude_3") + p = argparse.ArgumentParser(description="invoke_stable_diffusion") p.add_argument("-p", "--prompt", default="Hello", help="prompt", required=True) p.add_argument("-d", "--debug", default="False", help="debug") return p.parse_args() From 34eb0587d03a70642e5377b81b4b3ba9eefb4c50 Mon Sep 17 00:00:00 2001 From: nalbam Date: Wed, 14 Aug 2024 12:36:07 +0900 Subject: [PATCH 3/7] feat: Add knowledge base retrieval functionality --- bedrock/invoke_knowledge_base.py | 87 +++++++++++++++++++++++++------- handler.py | 57 ++++++++++++++++++++- 2 files changed, 125 insertions(+), 19 deletions(-) diff --git a/bedrock/invoke_knowledge_base.py b/bedrock/invoke_knowledge_base.py index bb92ab7..12c3258 100644 --- a/bedrock/invoke_knowledge_base.py +++ b/bedrock/invoke_knowledge_base.py @@ -4,11 +4,9 @@ import argparse import json import boto3 -import pprint from botocore.client import Config -pp = pprint.PrettyPrinter(indent=2) bedrock_config = Config( connect_timeout=120, read_timeout=120, retries={"max_attempts": 0} @@ -17,8 +15,6 @@ bedrock_agent_client = boto3.client( "bedrock-agent-runtime", region_name="us-east-1", config=bedrock_config ) -# boto3_session = boto3.session.Session() -# region_name = boto3_session.region_name model_id = "anthropic.claude-v2:1" # try with both claude instant as well as claude-v2. for claude v2 - "anthropic.claude-v2" region_id = "us-east-1" # replace it with the region you're running sagemaker notebook @@ -26,11 +22,11 @@ SYSTEM_MESSAGE = "답변은 한국어 해요체로 해요." -# def parse_args(): -# p = argparse.ArgumentParser(description="invoke_knowledge_base") -# p.add_argument("-p", "--prompt", default="안녕", help="prompt") -# p.add_argument("-d", "--debug", default="False", help="debug") -# return p.parse_args() +def parse_args(): + p = argparse.ArgumentParser(description="invoke_claude_3") + p.add_argument("-p", "--prompt", default="안녕", help="prompt") + p.add_argument("-d", "--debug", default="False", help="debug") + return p.parse_args() def retrieve(query, kbId, numberOfResults=5): @@ -79,20 +75,75 @@ def retrieveAndGenerate( ) +# fetch context from the response +def get_contexts(retrievalResults): + contexts = [] + for retrievedResult in retrievalResults: + contexts.append(retrievedResult["content"]["text"]) + return contexts + + def main(): # args = parse_args() - kb_id = "knowledge-base-whitepaper" - - query = "Please tell me about the Kontrol." + kb_id = "DQXVNP05K5" + + query = "kontrol의 기능 알려줘." + + # response = retrieveAndGenerate(query, kb_id, model_id=model_id, region_id=region_id) + # generated_text = response["output"]["text"] + # print(generated_text) + + response = retrieve(query, kb_id, 3) + retrievalResults = response["retrievalResults"] + # print(retrievalResults) + + contexts = get_contexts(retrievalResults) + # print(contexts) + + prompt = f""" +Human: You are a financial advisor AI system, and provides answers to questions by using fact based and statistical information when possible. +Use the following pieces of information to provide a concise answer to the question enclosed in tags. +If you don't know the answer, just say that you don't know, don't try to make up an answer. + +{contexts} + + + +{query} + + +The response should be specific and use statistics or numbers when possible. + +Assistant:""" + + # payload with model paramters + messages = [ + { + "role": "user", + "content": [{"type": "text", "text": prompt}], + } + ] + sonnet_payload = json.dumps( + { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 512, + "messages": messages, + "temperature": 0.5, + "top_p": 1, + } + ) - # response = retrieve(query, kb_id, 3) - # retrievalResults = response["retrievalResults"] - # pp.pprint(retrievalResults) + modelId = "anthropic.claude-3-sonnet-20240229-v1:0" # change this to use a different version from the model provider + accept = "application/json" + contentType = "application/json" + response = bedrock_client.invoke_model( + body=sonnet_payload, modelId=modelId, accept=accept, contentType=contentType + ) + response_body = json.loads(response.get("body").read()) + response_text = response_body.get("content")[0]["text"] - response = retrieveAndGenerate(query, kb_id, model_id=model_id, region_id=region_id) - generated_text = response["output"]["text"] - pp.pprint(generated_text) + print(response_text) if __name__ == "__main__": diff --git a/handler.py b/handler.py index 634c217..b34c4df 100644 --- a/handler.py +++ b/handler.py @@ -9,11 +9,16 @@ import requests import io +from botocore.client import Config + from slack_bolt import App, Say from slack_bolt.adapter.aws_lambda import SlackRequestHandler + BOT_CURSOR = os.environ.get("BOT_CURSOR", ":robot_face:") +AWS_REGION = os.environ.get("AWS_REGION", "us-east-1") + # Set up Slack API credentials SLACK_BOT_TOKEN = os.environ["SLACK_BOT_TOKEN"] SLACK_SIGNING_SECRET = os.environ["SLACK_SIGNING_SECRET"] @@ -76,7 +81,14 @@ table = dynamodb.Table(DYNAMODB_TABLE_NAME) # Initialize the Amazon Bedrock runtime client -bedrock = boto3.client(service_name="bedrock-runtime", region_name="us-east-1") +bedrock = boto3.client(service_name="bedrock-runtime", region_name=AWS_REGION) + +bedrock_config = Config( + connect_timeout=120, read_timeout=120, retries={"max_attempts": 0} +) +bedrock_agent_client = boto3.client( + "bedrock-agent-runtime", region_name=AWS_REGION, config=bedrock_config +) # Get the context from DynamoDB @@ -164,6 +176,41 @@ def chat_update(say, channel, thread_ts, latest_ts, message="", continue_thread= return message, latest_ts +def invoke_knowledge_base(content): + """ + Invokes the Amazon Bedrock Knowledge Base to retrieve information using the input + provided in the request body. + + :param content: The content that you want to use for retrieval. + :return: The retrieved contexts from the knowledge base. + """ + + try: + response = bedrock_agent_client.retrieve( + retrievalQuery={"text": content}, + knowledgeBaseId=KB_ID, + retrievalConfiguration={ + "vectorSearchConfiguration": { + "numberOfResults": 3, + # "overrideSearchType": "HYBRID", # optional + } + }, + ) + + retrievalResults = response["retrievalResults"] + + contexts = [] + for retrievedResult in retrievalResults: + contexts.append(retrievedResult["content"]["text"]) + + return contexts + + except Exception as e: + print("invoke_knowledge_base: Error: {}".format(e)) + + raise e + + def invoke_claude_3(content): """ Invokes Anthropic Claude 3 Sonnet to run an inference using the input @@ -350,6 +397,14 @@ def conversation(say: Say, thread_ts, content, channel, user, client_msg_id): prompts.append(message) + if KB_ID != "None": + chat_update(say, channel, thread_ts, latest_ts, MSG_RESPONSE) + + # Get the knowledge base contexts + contexts = invoke_knowledge_base(prompt) + + prompts.extend(contexts) + if prompt: prompts.append(prompt) From 89e88e4267dd7b5d4d1d104e8530fba8c1507641 Mon Sep 17 00:00:00 2001 From: nalbam Date: Wed, 14 Aug 2024 12:39:15 +0900 Subject: [PATCH 4/7] chore: Update push.yml to include KB_ID in environment variables --- .github/workflows/push.yml | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/.github/workflows/push.yml b/.github/workflows/push.yml index c1ed4ae..f9df8c0 100644 --- a/.github/workflows/push.yml +++ b/.github/workflows/push.yml @@ -10,9 +10,10 @@ env: BOT_CURSOR: ${{ vars.BOT_CURSOR }} DYNAMODB_TABLE_NAME: ${{ vars.DYNAMODB_TABLE_NAME }} ENABLE_IMAGE: ${{ vars.ENABLE_IMAGE }} + KB_ID: ${{ vars.KB_ID }} MODEL_ID_IMAGE: ${{ vars.MODEL_ID_IMAGE }} - SYSTEM_MESSAGE: ${{ vars.SYSTEM_MESSAGE }} MODEL_ID_TEXT: ${{ vars.MODEL_ID_TEXT }} + SYSTEM_MESSAGE: ${{ vars.SYSTEM_MESSAGE }} AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }} AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }} @@ -52,11 +53,12 @@ jobs: echo "BOT_CURSOR=${BOT_CURSOR}" >> .env echo "DYNAMODB_TABLE_NAME=${DYNAMODB_TABLE_NAME}" >> .env echo "ENABLE_IMAGE=${ENABLE_IMAGE}" >> .env + echo "KB_ID=${KB_ID}" >> .env echo "MODEL_ID_IMAGE=${MODEL_ID_IMAGE}" >> .env + echo "MODEL_ID_TEXT=${MODEL_ID_TEXT}" >> .env echo "SLACK_BOT_TOKEN=${SLACK_BOT_TOKEN}" >> .env echo "SLACK_SIGNING_SECRET=${SLACK_SIGNING_SECRET}" >> .env echo "SYSTEM_MESSAGE=${SYSTEM_MESSAGE}" >> .env - echo "MODEL_ID_TEXT=${MODEL_ID_TEXT}" >> .env - name: Deploy to AWS Lambda 🚀 run: npx serverless deploy --region us-east-1 From 35d7594f96464872d20fdbf09d6533d8eee47850 Mon Sep 17 00:00:00 2001 From: nalbam Date: Wed, 14 Aug 2024 12:43:18 +0900 Subject: [PATCH 5/7] feat: Update serverless.yml with new Bedrock permissions --- serverless.yml | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/serverless.yml b/serverless.yml index a91047d..14af61a 100644 --- a/serverless.yml +++ b/serverless.yml @@ -13,13 +13,18 @@ provider: Action: - dynamodb:* Resource: - - "arn:aws:dynamodb:*:*:table/${self:provider.environment.DYNAMODB_TABLE_NAME}" + - "arn:aws:dynamodb:::table/${self:provider.environment.DYNAMODB_TABLE_NAME}" - Effect: Allow Action: - bedrock:InvokeModel Resource: - - "arn:aws:bedrock:*::foundation-model/anthropic.claude-*" - - "arn:aws:bedrock:*::foundation-model/stability.stable-diffusion-*" + - "arn:aws:bedrock:::foundation-model/anthropic.claude-*" + - "arn:aws:bedrock:::foundation-model/stability.stable-diffusion-*" + - Effect: Allow + Action: + - bedrock:Retrieve + Resource: + - "arn:aws:bedrock:::knowledge-base/*" functions: mention: From feb851562bacb1e509d0ec5bd2b64c675e230369 Mon Sep 17 00:00:00 2001 From: nalbam Date: Wed, 14 Aug 2024 13:34:37 +0900 Subject: [PATCH 6/7] feat: Update serverless.yml with new Bedrock permissions --- serverless.yml | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/serverless.yml b/serverless.yml index 14af61a..ca08b6d 100644 --- a/serverless.yml +++ b/serverless.yml @@ -13,18 +13,18 @@ provider: Action: - dynamodb:* Resource: - - "arn:aws:dynamodb:::table/${self:provider.environment.DYNAMODB_TABLE_NAME}" + - "arn:aws:dynamodb:us-east-1:*:table/${self:provider.environment.DYNAMODB_TABLE_NAME}" - Effect: Allow Action: - - bedrock:InvokeModel + - bedrock:Retrieve Resource: - - "arn:aws:bedrock:::foundation-model/anthropic.claude-*" - - "arn:aws:bedrock:::foundation-model/stability.stable-diffusion-*" + - "arn:aws:bedrock:us-east-1:*:knowledge-base/*" - Effect: Allow Action: - - bedrock:Retrieve + - bedrock:InvokeModel Resource: - - "arn:aws:bedrock:::knowledge-base/*" + - "arn:aws:bedrock:::foundation-model/anthropic.claude-*" + - "arn:aws:bedrock:::foundation-model/stability.stable-diffusion-*" functions: mention: From 46935fa3e0d59216f614a196e606f8fa4c8e2775 Mon Sep 17 00:00:00 2001 From: nalbam Date: Wed, 14 Aug 2024 13:39:40 +0900 Subject: [PATCH 7/7] Update Bedrock permissions in serverless.yml --- serverless.yml | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/serverless.yml b/serverless.yml index ca08b6d..0b0c1e8 100644 --- a/serverless.yml +++ b/serverless.yml @@ -13,18 +13,18 @@ provider: Action: - dynamodb:* Resource: - - "arn:aws:dynamodb:us-east-1:*:table/${self:provider.environment.DYNAMODB_TABLE_NAME}" + - "arn:aws:dynamodb:*:*:table/${self:provider.environment.DYNAMODB_TABLE_NAME}" - Effect: Allow Action: - bedrock:Retrieve Resource: - - "arn:aws:bedrock:us-east-1:*:knowledge-base/*" + - "arn:aws:bedrock:*:*:knowledge-base/*" - Effect: Allow Action: - bedrock:InvokeModel Resource: - - "arn:aws:bedrock:::foundation-model/anthropic.claude-*" - - "arn:aws:bedrock:::foundation-model/stability.stable-diffusion-*" + - "arn:aws:bedrock:*::foundation-model/anthropic.claude-*" + - "arn:aws:bedrock:*::foundation-model/stability.stable-diffusion-*" functions: mention: