Implementing RAG App Using Knowledge Base from Amazon Bedrock and Streamlit
In this article I will provide a step-by-step guidance on how to build a full-stack document Q&A application by employing RAG(Retrieval Augmented Generation) technique using Knowledge Base(KB) from Amazon Bedrock & Streamlit. The key-points that I will discussing here are as follows:
- Building the back-end of the app,
- Building a front-end to interact and
- Implementing autoSync functionality for KB.
The example screenshots I’ve showed here are all done in us-west-2 region. A demo of the application can be found towards the end of the article under Application Demo section and the link for the GitHub repo can be found from Some Useful Links section. So without wasting more time let’s deep dive into the details:
Building the back-end of the application on AWS
In order to build the back-end we will need to interact with 2 services which are S3 and Amazon Bedrock. In S3 bucket we will upload the documents which will act as knowledge source for the KB and then we will create the knowledge base in Bedrock by connecting it the S3 bucket.
Creating the Knowledge Source
In this step we will create a S3 bucket and upload document(s) to it. For this we will go to the service called S3 and will choose the region as us-west-2 and keep everything as it is. After we upload some documents to the bucket it will look like this:
Creating the Knowledge Base in Bedrock
In this step we will create the knowledge base which will be the bread and butter of this application. To do this we will go the service called Amazon Bedrock, there after clicking on the sandwich icon from top-left corner we will find the Knowledge base option under Orchestration.
Once we click on the Create Knowledge base, we’ll land on this page:
Here, we will have to provide the name of the KB and an optional description. After that, unless we have a specific requirements for IAM permissions and TAGs we can keep everything else as it is and go to the next step:
In this step we will connect the S3 bucket that we created earlier. As we are creating the KB in the us-west-2 region only the buckets from that region will be visible here. Here we have the flexibility of choosing the entire bucket or a specific document from it. Here under the Advanced settings
options we can customize few things such as Chunking Strategy, for which we get these below options:
so unless you have a specific requirement for chunking or you haven’t pre-processes your documents already we can go ahead with the Default chunking option and go ahead to the next step:
In this step we get to choose the embedding model for the documents and a vector data base to store, manage and update those embeddings. The vector database will be created automatically by AWS.
For the vector store, we can let Bedrock create a new vector store in Amazon OpenSearch or we can connect a previously created vector store in OpenSearch, Aurora, Pinecone or Redis. Again, unless you have a specific requirement you can go ahead with creating new vector store, which I did for this demo.
The next step is Review and create. So after reviewing all the info we can go ahead and click in Create Knowledge base which will first take some time to prepare the vector database in Amazon OpenSearch Serverless:
The Vector store can be found under Collections from Amazon OpenSearch Service:
After that we will have to sync the data source of the knowledge base in the vector store:
Sync is the most important part of creating the KB. In this process, it lookup the documents in the S3 bucket to pre-process and extracts the text. After that it creates smaller chunks of data using the pre-defined chunking strategy, pass it to the pre-defined embedding model and then store it in the vector store so that we are able to query.
Each time any data in the S3 bucket is modified, added o removed the the data source has to be re-indexed thus we need to sync the data either manually from console or automate the process some how(we discussed this in the last step).
Once the sync is done we can test it from the console itself or we can query the documents using bedrock clients from an application.
Building a front-end to interact with the app
Now in order to take full-advantage of what we built so far we need an UI where user can go and interact with the data, basically be able to retrieve the info they want from the document. That’s why we are going to build a multi-page app using Streamlit. A brief explanation of the code is given below as well.
Q&A Page:
On this page user will be able to ask questions from the documents. For this we are going to use retrieve_and_generate API from bedrock-agent-runtime client.
When a user enters a question from the UI, this API helps to retrieve the most relevant information from the available documents by querying the knowledge base based on the user query and generates the response.
This is what the AWS documentation says:
Queries a knowledge base and generates responses based on the retrieved results. The response cites up to five sources but only selects the ones that are relevant to the query.
What it does is, it retrieves the info as a context and pass it to the LLM along with the user query to generate the required response.
bedrockClient = boto3.client('bedrock-agent-runtime', 'us-west-2')
def getAnswers(questions):
knowledgeBaseResponse = bedrockClient.retrieve_and_generate(
input={'text': questions},
retrieveAndGenerateConfiguration={
'knowledgeBaseConfiguration': {
'knowledgeBaseId': 'YOUR_KNOWLEDGE_BASE_ID',
'modelArn': 'arn:aws:bedrock:us-west-2::foundation-model/anthropic.claude-instant-v1'
},
'type': 'KNOWLEDGE_BASE'
})
return knowledgeBaseResponseHere the response contains a lot of info but the ones we are going to use are the actual answer, the context it used to generate that answer and the URL of the document from where it retrieved that context.
For the latest question and it’s answer only we will show this info i.e. context and the document URL:
if len(response['citations'][0]['retrievedReferences']) != 0:
context = response['citations'][0]['retrievedReferences'][0]['content']['text']
doc_url = response['citations'][0]['retrievedReferences'][0]['location']['s3Location']['uri']
#Below lines are used to show the context and the document source for the latest Question Answer
st.markdown(f"<span style='color:#FFDA33'>Context used: </span>{context}", unsafe_allow_html=True)
st.markdown(f"<span style='color:#FFDA33'>Source Document: </span>{doc_url}", unsafe_allow_html=True)
else:
st.markdown(f"<span style='color:red'>No Context</span>", unsafe_allow_html=True)and we will maintain the chat history of the user q&a using session_state from Streamlit:
if 'chat_history' not in st.session_state:
st.session_state.chat_history = []
for message in st.session_state.chat_history:
with st.chat_message(message['role']):
st.markdown(message['text'])Upload Document Page:
Other querying the existing document in the knowledge base we can also upload document from this application and ask questions from there.
As you might have guessed this page is to add documents to the Knowledge base. From this page we can upload documents to the S3 bucket taht is connected to the KB. The code for that is as follows:
def upload_file(file_name, renamed_file_name):
bucket_name = 'YOUR_S3_BUCKET_NAME'
try:
# Upload file to an S3 object
s3_client.upload_file(file_name, bucket_name, renamed_file_name)
st.markdown(f"Successfully uploaded the file!!! 😃", unsafe_allow_html=True)
except Exception as e:
st.markdown(f"Error: {str(e)}")Here, before uploading the document to the S3 bucket, little pre-processing is done to modify the file name when it will be uploaded in the bucket. Full code can be found from the above mentioned GitHub repo.
Now other than while creating the KB for the 1st time whenever we modify, add or remove documents from the S3 bucket we again need to sync the data course for the knowledge.
This can be done in two ways. One, manually doing it by clicking on the Sync button from the console(as shown in Screenshot 8) or two, by triggering a data ingestion job whenever there is a new data in the S3 bucket. Next, we will find out how to implement the second option.
Implementing AutoSync functionality for KB
If we think about from the user experience perspective, it is not expected that after the user uploads the document from the application the person will go to AWS console to start the sync process. Once the user uploads the document the rest of the process should be taken care of automatically. This the the motivation behind doing this is.
To do that we will write a lambda function that will start the data ingestion job and when we upload the document in the S3 bucket it will trigger that lambda.
Implementing the lambda
The lambda function will facilitate the data sync process using start_ingestion_job API from bedrock-agent client. The code is given below:
import os
import json
import boto3
bedrockClient = boto3.client('bedrock-agent')
def lambda_handler(event, context):
# TODO implement
print('Inside Lambda Handler')
print('event: ', event)
dataSourceId = os.environ['DATASOURCEID']
knowledgeBaseId = os.environ['KNOWLEDGEBASEID']
print('knowledgeBaseId: ', knowledgeBaseId)
print('dataSourceId: ', dataSourceId)
response = bedrockClient.start_ingestion_job(
knowledgeBaseId=knowledgeBaseId,
dataSourceId=dataSourceId
)
print('Ingestion Job Response: ', response)
return {
'statusCode': 200,
'body': json.dumps('response')
}Two required parameters for the API are knowledgeBaseId and dataSourceId, which can be found from the KB console:
We can use the string mentioned below Knowledge base ID for knowledgeBaseId parameter and the and once we click on the available Data source name below we will the Data source ID which we can use for the dataSourceId parameter.
Adding Trigger For The lambda
In order to automatically trigger this lambda lambda function whenever there’s new data in the bucket we need to add a trigger for this lambda.
The trigger can be added from here:
Once we click on the Add trigger button it will take us to page where we can select the service which in our case is the S3, from there we can select the bucket, for which we want to trigger the alarm. We can specify for which events we want to trigger this lambda such as create event, delete event etc. Here we will only want to trigger the lambda when there is any new document in the bucket so we will choose All object create events and as of now we limiting the trigger for the new files with extension .pdf after that we can go ahead click on Add to add the trigger.
Once we add, it will look like this:
By doing this after creation of the knowledge base, whenever user uploads a new document in that bucket, it will trigger the lambda to start the data ingestion jobs to re-index the data source and once the job is completed user can query the knowledge base for the new documents as well.
One thing to keep in mind, the sync process can take take from few seconds to several minutes depending on the size of file, so user might have to wait for some time before start asking question from the newly uploaded documents.
