GenAI App using Proprietery Data, Spring Boot, Amazon Bedrock - Part 3

This is part of a series of articles to introduce Generative AI and step-by-step process with code sample to build GenAI Application; Part-1 (GenAI Introduction), Part-2 (Building GenAI App using Spring Boot and Amazon Bedrock) and Part-3 (Building GenAI App with Proprietary data).

In previous article, we have built GenAI Application (simple ChatBot) using Amazon Bedrock Model (Claude 3 Haiku) and Spring Boot. This App generates answers based on training dataset (public information) of the Model.
What about if we could restrict or refine this Model with proprietary data (like Company details, Sales, Customer records) as add-on Knowledge Base instead of retraining the Model. This would allow to serve personalized and contextual response, e.g. generating response with Customer's last order details. In this article we are going to enhance our GenAI application to accept proprietary data.

Prerequisites

1. Spring Boot v3, JDK 17
2. AWS SDK v2.x
3. AWS IAM account with API credentials and necessary permissions ("AmazonBedRockFullAccess")
4. Amazon Bedrock with access granted for following Base Models-
• Claude 3 Haiku by Anthropic (LLM / Response Generation)
• Embed English by Cohere (Vector Embedding of Proprietary Data)
  

Let's first understand overall concept and complete Amazon Bedrock Knowledge Base setup process before jumping into code.

Knowledge Base
GenAI includes an architecture concept called "RAG (Retrieval Augmented Generation)" which helps an already trained Model to retrieve contextual information from proprietary dataset and augment the same context into generated response.
To achieve this, we'll extend Amazon Bedrock with Knowledge Base (Proprietary data store) and ask Model to refer this Knowledge Base while generating Response. 

Knowledge Base Setup
This is 2-step workflow process, Data Ingestion and Response Generation.
Step-1: Data Ingestion workflow (Import proprietary data)

1. Setup Amazon Bedrock Knowledge Base. Refer this article (Don't create Amazon Agent, we'll access this Knowledge Base from Spring Boot App)
   
2. Prepare proprietary information (company details, sales records, customer information etc.) which are relevant for your App in supported format (Excel, PDF, Text etc.).
3. Upload content to staging area (S3 bucket which is configured with Knowledge Base)
4. Go to Knowledge Base and Sync data store. This would detect changes in S3 bucket and read changed files
5. Sync process will call Embedding Model (Embed English) to generate Vector Embeddings which are required to store proprietary data in Vector Database.
6. This Vector Database is going to be used to retrieve contextual information during response generation workflow (Step-2)
   

Once Sync (Step-1) is complete, Test Knowledge Base using the Bedrock Console before jumping into code. In the below diagram, Bedrock Console Test (Model: Claude 3 Haiku) answer the query using shared proprietary data (dummy company: Pingo).

Step-2: Response generation workflow (using Knowledge Base and Spring Boot)

Invoke Knowledge Base from Spring Boot App
Now let's look into code changes to invoke Model (along with Knowledge Base) and retrieve contextual response.

Include dependencies for Amazon Bedrock, Bedrock Runtime and Bedrock Agent Runtime in pom.xml

Create Bedrock Agent Runtime Client using your API Credentials and Region

public AwsCredentialsProvider awsCredentialsProvider() {
	AwsCredentialsProvider awsCredentialsProvider = StaticCredentialsProvider.create(
			AwsBasicCredentials.create(accessKey, secretKey));
	return awsCredentialsProvider;
}
@Bean
public BedrockAgentRuntimeClient bedrockAgentRuntime() {
	return BedrockAgentRuntimeClient.builder()
			.region(Region.of(region))
			.credentialsProvider(awsCredentialsProvider())
			.build();
}

Compose Request Input, Retrieve and Generate Response

private String invokeModelWithRAG(String query) {
	// Config
	KnowledgeBaseRetrieveAndGenerateConfiguration kbConfig = KnowledgeBaseRetrieveAndGenerateConfiguration.builder()
			.knowledgeBaseId(KNOWLEDGE_BASE_ID)
			.modelArn(MODEL_ARN)
			.build();
	RetrieveAndGenerateConfiguration config = RetrieveAndGenerateConfiguration.builder()
			.knowledgeBaseConfiguration(kbConfig)
			.type(RetrieveAndGenerateType.KNOWLEDGE_BASE) // If not defined (Error: Invalid input provided)
			.build();
	
	// Input
	RetrieveAndGenerateInput input = RetrieveAndGenerateInput.builder().text(query).build();
	
	RetrieveAndGenerateRequest request = RetrieveAndGenerateRequest.builder()
			.input(input)
			.retrieveAndGenerateConfiguration(config)
			.build();
	
	// Response from Knowledge Base
	RetrieveAndGenerateResponse response = bedrockAgentClient.retrieveAndGenerate(request);
	// Response includes output (generated text) and citations (reference to Proprietary content and S3 links)
	return response.output().text();
}

Now hit the Chat API on Browser to send query and check generated responses (it should be from your shared proprietary data). Here I'm using a dummy company (Pingo).

This Application source code is shared on GitHub