Knowledge Base

Amazon Bedrock knowledge bases are a great way to implement Retrieval Augmented Generation, or RAG for short. RAG is a common pattern in building generative AI applications that involves storing a lot of content, like documentation, in a vector database like Postgres with pg_vector or OpenSearch.

The default setup for an Amazon Bedrock knowledge base is OpenSearch Serverless which has a default cost whether or not you use it. You can get a large AWS bill if you are not careful. If you are just testing this out make sure to turn off the OpenSearch Serverless instance when you are done.

Create a knowledge base

To integrate Bedrock knowledge base with your conversation route, first create an Amazon Bedrock knowledge base in the console, CLI, or with CDK.

Create a custom query and tool

amplify/data/resource.ts
import { type ClientSchema, a, defineData } from "@aws-amplify/backend";
const schema = a.schema({
  knowledgeBase: a
    .query()
    .arguments({ input: a.string() })
    .handler(
      a.handler.custom({
        dataSource: "KnowledgeBaseDataSource",
        entry: "./resolvers/kbResolver.js",
    .returns(a.string())
    .authorization((allow) => allow.authenticated()),
  chat: a.conversation({
    aiModel: a.ai.model("Claude 3.5 Sonnet"),
    systemPrompt: `You are a helpful assistant.`,
    tools: [
      a.ai.dataTool({
        name: 'searchDocumentation',
        description: 'Performs a similarity search over the documentation for ...',
        query: a.ref('knowledgeBase'),

Write an AWS AppSync resolver

Then you'll need to create a JavaScript AWS AppSync resolver to connect the query to the knowledge base. You'll need to know the ID of the knowledge base you want to use, which you can find in the Amazon Bedrock console or with the AWS CLI.

amplify/data/resolvers/kbResolver.js
export function request(ctx) {
  const { input } = ctx.args;
    resourcePath: "/knowledgebases/[KNOWLEDGE_BASE_ID]/retrieve",
    method: "POST",
    params: {
      headers: {
        "Content-Type": "application/json",
      body: JSON.stringify({
        retrievalQuery: {
          text: input,
export function response(ctx) {
  return JSON.stringify(ctx.result.body);

Define the data source

Then in the amplify backend file you will need to create the data source for the knowledge base query and give it permission to call the knowledge base.

amplify/backend.ts
import { defineBackend } from '@aws-amplify/backend';
import { auth } from './auth/resource';
import { data } from './data/resource';
import { PolicyStatement } from 'aws-cdk-lib/aws-iam';
const backend = defineBackend({
const KnowledgeBaseDataSource =
  backend.data.resources.graphqlApi.addHttpDataSource(
    "KnowledgeBaseDataSource",
    `https://bedrock-runtime.${cdk.Stack.of(backend.data).region}.amazonaws.com`,
      authorizationConfig: {
        signingRegion: cdk.Stack.of(backend.data).region,
        signingServiceName: "bedrock",
KnowledgeBaseDataSource.grantPrincipal.addToPrincipalPolicy(
  new PolicyStatement({
    resources: [
      `arn:aws:bedrock:${cdk.Stack.of(backend.data).region}:[account ID]:knowledge-base/[knowledge base ID]`
    actions: ["bedrock:Retrieve"],

That's it!