Add Amazon Bedrock Knowledge base as a RAG Engine (retriever)

bin/config.ts

@@ -39,6 +39,9 @@ export function getConfig(): SystemConfig {
           createIndex: false,
           enterprise: false,
         },
+        knowledgeBase: {
+          enabled: false,
+        },
       },
       embeddingsModels: [
         {

cli/magic-config.ts

@@ -156,6 +156,8 @@ const embeddingModels = [
         ? config.llms?.sagemaker.length > 0
         : false;
       options.huggingfaceApiSecretArn = config.llms?.huggingfaceApiSecretArn;
+      options.enableSagemakerModelsSchedule =
+        config.llms?.sagemakerSchedule?.enabled;
       options.enableSagemakerModelsSchedule =
         config.llms?.sagemakerSchedule?.enabled;
       options.timezonePicker = config.llms?.sagemakerSchedule?.timezonePicker;
@@ -194,6 +196,7 @@ const embeddingModels = [
         (m) => m.default
       )[0].name;
       options.kendraExternal = config.rag.engines.kendra.external;
+      options.kbExternal = config.rag.engines.knowledgeBase?.external ?? [];
       options.kendraEnterprise = config.rag.engines.kendra.enterprise;
 
       // Advanced settings
@@ -586,6 +589,7 @@ async function processCreateOptions(options: any): Promise<void> {
         { message: "Aurora", name: "aurora" },
         { message: "OpenSearch", name: "opensearch" },
         { message: "Kendra (managed)", name: "kendra" },
+        { message: "Bedrock KnowldgeBase", name: "knowledgeBase" },
       ],
       validate(choices: any) {
         return (this as any).skipped || choices.length > 0
@@ -694,6 +698,82 @@ async function processCreateOptions(options: any): Promise<void> {
     });
     newKendra = kendraInstance.newKendra;
   }
+
+  // Knowledge Bases
+  let newKB =
+    answers.enableRag && answers.ragsToEnable.includes("knowledgeBase");
+  const kbExternal: any[] = [];
+  const existingKBIndices = Array.from(options.kbExternal || []);
+  while (newKB === true) {
+    const existingIndex: any = existingKBIndices.pop();
+    const kbQ = [
+      {
+        type: "input",
+        name: "name",
+        message: "KnowledgeBase source name",
+        validate(v: string) {
+          return RegExp(/^\w[\w-_]*\w$/).test(v);
+        },
+        initial: existingIndex?.name,
+      },
+      {
+        type: "autocomplete",
+        limit: 8,
+        name: "region",
+        choices: ["us-east-1", "us-west-2"],
+        message: `Region of the Bedrock Knowledge Base index${
+          existingIndex?.region ? " (" + existingIndex?.region + ")" : ""
+        }`,
+        initial: ["us-east-1", "us-west-2"].indexOf(existingIndex?.region),
+      },
+      {
+        type: "input",
+        name: "roleArn",
+        message:
+          "Cross account role Arn to assume to call the Bedrock KnowledgeBase, leave empty if not needed",
+        validate: (v: string) => {
+          const valid = iamRoleRegExp.test(v);
+          return v.length === 0 || valid;
+        },
+        initial: existingIndex?.roleArn ?? "",
+      },
+      {
+        type: "input",
+        name: "knowledgeBaseId",
+        message: "Bedrock KnowledgeBase ID",
+        validate(v: string) {
+          return /[A-Z0-9]{10}/.test(v);
+        },
+        initial: existingIndex?.knowledgeBaseId,
+      },
+      {
+        type: "confirm",
+        name: "enabled",
+        message: "Enable this knowledge base",
+        initial: existingIndex?.enabled ?? true,
+      },
+      {
+        type: "confirm",
+        name: "newKB",
+        message: "Do you want to add another Bedrock KnowledgeBase source",
+        initial: false,
+      },
+    ];
+    const kbInstance: any = await enquirer.prompt(kbQ);
+    const ext = (({ enabled, name, roleArn, knowledgeBaseId, region }) => ({
+      enabled,
+      name,
+      roleArn,
+      knowledgeBaseId,
+      region,
+    }))(kbInstance);
+    if (ext.roleArn === "") ext.roleArn = undefined;
+    kbExternal.push({
+      ...ext,
+    });
+    newKB = kbInstance.newKB;
+  }
+
   const modelsPrompts = [
     {
       type: "select",
@@ -1078,6 +1158,10 @@ async function processCreateOptions(options: any): Promise<void> {
           external: [{}],
           enterprise: false,
         },
+        knowledgeBase: {
+          enabled: false,
+          external: [{}],
+        },
       },
       embeddingsModels: [{}],
       crossEncoderModels: [{}],
@@ -1107,6 +1191,9 @@ async function processCreateOptions(options: any): Promise<void> {
     config.rag.engines.kendra.createIndex || kendraExternal.length > 0;
   config.rag.engines.kendra.external = [...kendraExternal];
   config.rag.engines.kendra.enterprise = answers.kendraEnterprise;
+  config.rag.engines.knowledgeBase.enabled =
+    config.rag.engines.knowledgeBase.external.length > 0;
+  config.rag.engines.knowledgeBase.external = [...kbExternal];
 
   console.log("\n✨ This is the chosen configuration:\n");
   console.log(JSON.stringify(config, undefined, 2));

lib/chatbot-api/functions/api-handler/index.py

@@ -15,6 +15,7 @@
 from routes.documents import router as documents_router
 from routes.kendra import router as kendra_router
 from routes.user_feedback import router as user_feedback_router
+from routes.bedrock_kb import router as bedrock_kb_router
 
 tracer = Tracer()
 logger = Logger()
@@ -32,6 +33,7 @@
 app.include_router(documents_router)
 app.include_router(kendra_router)
 app.include_router(user_feedback_router)
+app.include_router(bedrock_kb_router)
 
 
 @logger.inject_lambda_context(

lib/chatbot-api/functions/api-handler/routes/bedrock_kb.py

@@ -0,0 +1,21 @@
+import genai_core.parameters
+import genai_core.bedrock_kb
+from pydantic import BaseModel
+from aws_lambda_powertools import Logger, Tracer
+from aws_lambda_powertools.event_handler.appsync import Router
+
+tracer = Tracer()
+router = Router()
+logger = Logger()
+
+
+class KendraDataSynchRequest(BaseModel):
+    workspaceId: str
+
+
+@router.resolver(field_name="listBedrockKnowledgeBases")
+@tracer.capture_method
+def list_bedrock_kbs():
+    indexes = genai_core.bedrock_kb.list_bedrock_kbs()
+
+    return indexes

lib/chatbot-api/functions/api-handler/routes/rag.py

@@ -35,6 +35,11 @@ def engines():
             "name": "Amazon Kendra",
             "enabled": engines.get("kendra", {}).get("enabled", False) == True,
         },
+        {
+            "id": "bedrock_kb",
+            "name": "Bedrock Knowledge Bases",
+            "enabled": engines.get("knowledgeBase", {}).get("enabled", False) == True,
+        },
     ]
 
     return ret_value

lib/chatbot-api/functions/api-handler/routes/workspaces.py

@@ -1,6 +1,7 @@
 import re
 import genai_core.types
 import genai_core.kendra
+import genai_core.bedrock_kb
 import genai_core.parameters
 import genai_core.workspaces
 from pydantic import BaseModel
@@ -55,6 +56,13 @@ class CreateWorkspaceKendraRequest(BaseModel):
     useAllData: bool
 
 
+class CreateWorkspaceBedrockKBRequest(BaseModel):
+    kind: str
+    name: str
+    knowledgeBaseId: str
+    hybridSearch: bool
+
+
 @router.resolver(field_name="listWorkspaces")
 @tracer.capture_method
 def list_workspaces():
@@ -115,6 +123,16 @@ def create_kendra_workspace(input: dict):
     return ret_value
 
 
+@router.resolver(field_name="createBedrockKBWorkspace")
+@tracer.capture_method
+def create_bedrock_kb_workspace(input: dict):
+    config = genai_core.parameters.get_config()
+
+    request = CreateWorkspaceBedrockKBRequest(**input)
+    ret_value = _create_workspace_bedrock_kb(request, config)
+    return ret_value
+
+
 def _create_workspace_aurora(request: CreateWorkspaceAuroraRequest, config: dict):
     workspace_name = request.name.strip()
     embedding_models = config["rag"]["embeddingsModels"]
@@ -291,6 +309,39 @@ def _create_workspace_kendra(request: CreateWorkspaceKendraRequest, config: dict
     )
 
 
+def _create_workspace_bedrock_kb(
+    request: CreateWorkspaceBedrockKBRequest, config: dict
+):
+    workspace_name = request.name.strip()
+    kbs = genai_core.bedrock_kb.list_bedrock_kbs()
+
+    workspace_name_match = name_regex.match(workspace_name)
+    workspace_name_is_match = bool(workspace_name_match)
+    if (
+        len(workspace_name) == 0
+        or len(workspace_name) > 100
+        or not workspace_name_is_match
+    ):
+        raise genai_core.types.CommonError("Invalid workspace name")
+
+    knowledge_base = None
+    for current in kbs:
+        if current["id"] == request.knowledgeBaseId:
+            knowledge_base = current
+            break
+
+    if knowledge_base is None:
+        raise genai_core.types.CommonError("Knowledge Base id not found")
+
+    return _convert_workspace(
+        genai_core.workspaces.create_workspace_bedrock_kb(
+            workspace_name=workspace_name,
+            knowledge_base=knowledge_base,
+            hybrid_search=request.hybridSearch,
+        )
+    )
+
+
 def _convert_workspace(workspace: dict):
     kendra_index_external = workspace.get("kendra_index_external")
 
@@ -320,6 +371,8 @@ def _convert_workspace(workspace: dict):
         "kendraIndexId": workspace.get("kendra_index_id"),
         "kendraIndexExternal": kendra_index_external,
         "kendraUseAllData": workspace.get("kendra_use_all_data", kendra_index_external),
+        "knowledgeBaseId": workspace.get("knowledge_base_id"),
+        "knowledgeBaseExternal": workspace.get("knowledge_base_external"),
         "createdAt": workspace.get("created_at"),
         "updatedAt": workspace.get("updated_at"),
     }

lib/chatbot-api/rest-api.ts

@@ -188,6 +188,33 @@ export class ApiResolvers extends Construct {
           );
         }
 
+        if (props.config.rag.engines.knowledgeBase.enabled) {
+          for (const item of props.config.rag.engines.knowledgeBase.external ||
+            []) {
+            if (item.roleArn) {
+              apiHandler.addToRolePolicy(
+                new iam.PolicyStatement({
+                  actions: ["sts:AssumeRole"],
+                  resources: [item.roleArn],
+                })
+              );
+            } else {
+              apiHandler.addToRolePolicy(
+                new iam.PolicyStatement({
+                  actions: ["bedrock:Retrieve"],
+                  resources: [
+                    `arn:${cdk.Aws.PARTITION}:bedrock:${
+                      item.region ?? cdk.Aws.REGION
+                    }:${cdk.Aws.ACCOUNT_ID}:knowledge-base/${
+                      item.knowledgeBaseId
+                    }`,
+                  ],
+                })
+              );
+            }
+          }
+        }
+
         for (const item of props.config.rag.engines.kendra.external ?? []) {
           if (item.roleArn) {
             apiHandler.addToRolePolicy(

lib/chatbot-api/schema/schema.graphql

@@ -23,6 +23,13 @@ input CreateWorkspaceKendraInput {
   useAllData: Boolean!
 }
 
+input CreateWorkspaceBedrockKBInput {
+  name: String!
+  kind: String!
+  knowledgeBaseId: String!
+  hybridSearch: Boolean!
+}
+
 input CreateWorkspaceOpenSearchInput {
   name: String!
   kind: String!
@@ -150,6 +157,12 @@ type KendraIndex @aws_cognito_user_pools {
   external: Boolean!
 }
 
+type BedrockKB @aws_cognito_user_pools {
+  id: String!
+  name: String!
+  external: Boolean!
+}
+
 input ListDocumentsInput {
   workspaceId: String!
   documentType: String!
@@ -302,6 +315,8 @@ type Workspace @aws_cognito_user_pools {
   kendraIndexId: String
   kendraIndexExternal: Boolean
   kendraUseAllData: Boolean
+  knowledgeBaseId: String
+  knowledgeBaseExternal: Boolean
   createdAt: AWSDateTime!
   updatedAt: AWSDateTime!
 }
@@ -315,6 +330,8 @@ type Channel @aws_iam @aws_cognito_user_pools {
 type Mutation {
   createKendraWorkspace(input: CreateWorkspaceKendraInput!): Workspace!
     @aws_cognito_user_pools
+  createBedrockKBWorkspace(input: CreateWorkspaceBedrockKBInput!): Workspace!
+    @aws_cognito_user_pools
   createOpenSearchWorkspace(input: CreateWorkspaceOpenSearchInput!): Workspace!
     @aws_cognito_user_pools
   createAuroraWorkspace(input: CreateWorkspaceAuroraInput!): Workspace!
@@ -358,6 +375,7 @@ type Query {
     @aws_cognito_user_pools
   getSession(id: String!): Session @aws_cognito_user_pools
   listKendraIndexes: [KendraIndex!]! @aws_cognito_user_pools
+  listBedrockKnowledgeBases: [BedrockKB!]! @aws_cognito_user_pools
   isKendraDataSynching(workspaceId: String!): Boolean @aws_cognito_user_pools
   listDocuments(input: ListDocumentsInput!): DocumentsResult!
     @aws_cognito_user_pools

lib/model-interfaces/langchain/functions/request-handler/adapters/shared/meta/llama3_instruct.py

@@ -14,7 +14,7 @@
 
 You are an helpful assistant that provides concise answers to user questions with as little sentences as possible and at maximum 3 sentences. You do not repeat yourself. You avoid bulleted list or emojis.{EOD}{{chat_history}}{USER_HEADER}
 
-{{input}}{EOD}{ASSISTANT_HEADER}"""
+Context: {{input}}{EOD}{ASSISTANT_HEADER}"""
 
 Llama3QAPrompt = f"""{BEGIN_OF_TEXT}{SYSTEM_HEADER}