Merge branch 'main' of github.com:aws-samples/aws-genai-llm-chatbot

cli/magic-create.ts

@@ -13,7 +13,8 @@ import {
 import { LIB_VERSION } from "./version.js";
 import * as fs from "fs";
 
-const versionRegExp = /\d+.\d+.\d+/;
+const iamRoleRegExp = RegExp(/arn:aws:iam::\d+:role\/[\w-_]+/);
+const kendraIdRegExp = RegExp(/^\w{8}-\w{4}-\w{4}-\w{4}-\w{12}$/);
 
 const embeddingModels = [
   {
@@ -146,7 +147,7 @@ async function processCreateOptions(options: any): Promise<void> {
       message:
         "Cross account role arn to invoke Bedrock - leave empty if Bedrock is in same account",
       validate: (v: string) => {
-        const valid = RegExp(/arn:aws:iam::\d+:role\/[\w-_]+/).test(v);
+        const valid = iamRoleRegExp.test(v);
         return v.length === 0 || valid;
       },
       initial: options.bedrockRoleArn || "",
@@ -174,7 +175,7 @@ async function processCreateOptions(options: any): Promise<void> {
         { message: "OpenSearch", name: "opensearch" },
         { message: "Kendra (managed)", name: "kendra" },
       ],
-      skip: function (): boolean {
+      skip(): boolean {
         // workaround for https://github.com/enquirer/enquirer/issues/298
         (this as any).state._choices = (this as any).state.choices;
         return !(this as any).state.answers.enableRag;
@@ -199,57 +200,76 @@ async function processCreateOptions(options: any): Promise<void> {
   const answers: any = await enquirer.prompt(questions);
   const kendraExternal = [];
   let newKendra = answers.enableRag && answers.kendra;
-
+  const existingKendraIndices = Array.from(options.kendraExternal);
   while (newKendra === true) {
+    let existingIndex: any = existingKendraIndices.pop();
+    console.log(
+      existingIndex?.region,
+      Object.values(SupportedRegion).indexOf(existingIndex?.region)
+    );
     const kendraQ = [
       {
         type: "input",
         name: "name",
         message: "Kendra source name",
+        validate(v: string) {
+          return RegExp(/^\w[\w-_]*\w$/).test(v);
+        },
+        initial: existingIndex?.name,
       },
       {
         type: "autocomplete",
         limit: 8,
         name: "region",
         choices: Object.values(SupportedRegion),
-        message: "Region of the Kendra index",
+        message: `Region of the Kendra index${
+          existingIndex?.region ? " (" + existingIndex?.region + ")" : ""
+        }`,
+        initial: Object.values(SupportedRegion).indexOf(existingIndex?.region),
       },
       {
         type: "input",
         name: "roleArn",
         message:
           "Cross account role Arn to assume to call Kendra, leave empty if not needed",
         validate: (v: string) => {
-          const valid = RegExp(/arn:aws:iam::\d+:role\/[\w-_]+/).test(v);
+          const valid = iamRoleRegExp.test(v);
           return v.length === 0 || valid;
         },
-        initial: "",
+        initial: existingIndex?.roleArn ?? "",
       },
       {
         type: "input",
         name: "kendraId",
         message: "Kendra ID",
         validate(v: string) {
-          return RegExp(/\w{8}-\w{4}-\w{4}-\w{4}-\w{12}/).test(v);
+          return kendraIdRegExp.test(v);
         },
+        initial: existingIndex?.kendraId,
+      },
+      {
+        type: "confirm",
+        name: "enabled",
+        message: "Enable this index",
+        initial: existingIndex?.enabled ?? true,
       },
       {
         type: "confirm",
         name: "newKendra",
         message: "Do you want to add another Kendra source",
-        default: false,
+        initial: false,
       },
     ];
     const kendraInstance: any = await enquirer.prompt(kendraQ);
-    const ext = (({ name, roleArn, kendraId, region }) => ({
+    const ext = (({ enabled, name, roleArn, kendraId, region }) => ({
+      enabled,
       name,
       roleArn,
       kendraId,
       region,
     }))(kendraInstance);
     if (ext.roleArn === "") ext.roleArn = undefined;
     kendraExternal.push({
-      enabled: true,
       ...ext,
     });
     newKendra = kendraInstance.newKendra;
@@ -261,6 +281,9 @@ async function processCreateOptions(options: any): Promise<void> {
       message: "Which is the default embedding model",
       choices: embeddingModels.map((m) => ({ name: m.name, value: m })),
       initial: options.defaultEmbedding || undefined,
+      skip(): boolean {
+        return !(this as any).state.answers.enableRag;
+      },
     },
   ];
   const models: any = await enquirer.prompt(modelsPrompts);
@@ -296,27 +319,33 @@ async function processCreateOptions(options: any): Promise<void> {
         },
       },
       embeddingsModels: [{}],
-      crossEncoderModels: [
-        {
-          provider: "sagemaker",
-          name: "cross-encoder/ms-marco-MiniLM-L-12-v2",
-          default: true,
-        },
-      ],
+      crossEncoderModels: [{}],
     },
   };
 
-  config.rag.engines.kendra.createIndex =
-    answers.ragsToEnable.includes("kendra");
-  config.rag.engines.kendra.enabled =
-    config.rag.engines.kendra.createIndex || kendraExternal.length > 0;
-  config.rag.engines.kendra.external = [...kendraExternal];
+  // If we have not enabled rag the default embedding is set to the first model
+  if (!answers.enableRag) {
+    models.defaultEmbedding = embeddingModels[0].name;
+  }
+
+  config.rag.crossEncoderModels[0] = {
+    provider: "sagemaker",
+    name: "cross-encoder/ms-marco-MiniLM-L-12-v2",
+    default: true,
+  };
   config.rag.embeddingsModels = embeddingModels;
   config.rag.embeddingsModels.forEach((m: any) => {
     if (m.name === models.defaultEmbedding) {
       m.default = true;
     }
   });
+
+  config.rag.engines.kendra.createIndex =
+    answers.ragsToEnable.includes("kendra");
+  config.rag.engines.kendra.enabled =
+    config.rag.engines.kendra.createIndex || kendraExternal.length > 0;
+  config.rag.engines.kendra.external = [...kendraExternal];
+
   console.log("\n✨ This is the chosen configuration:\n");
   console.log(JSON.stringify(config, undefined, 2));
   (

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

@@ -67,8 +67,7 @@ def file_upload(workspace_id: str):
     if extension not in allowed_extensions:
         raise genai_core.types.CommonError("Invalid file extension")
 
-    result = genai_core.upload.generate_presigned_post(
-        workspace_id, request.fileName)
+    result = genai_core.upload.generate_presigned_post(workspace_id, request.fileName)
 
     return {"ok": True, "data": result}
 
@@ -147,7 +146,7 @@ def add_document(workspace_id: str, document_type: str):
             crawler_properties={
                 "follow_links": request.followLinks,
                 "limit": request.limit,
-            }
+            },
         )
 
         return {

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

@@ -77,6 +77,14 @@ def workspace(workspace_id: str):
     return {"ok": True, "data": ret_value}
 
 
+@router.delete("/workspaces/<workspace_id>")
+@tracer.capture_method
+def workspace(workspace_id: str):
+    genai_core.workspaces.delete_workspace(workspace_id)
+
+    return {"ok": True}
+
+
 @router.put("/workspaces")
 @tracer.capture_method
 def create_workspace():

lib/chatbot-api/rest-api.ts

@@ -77,6 +77,8 @@ export class RestApi extends Construct {
           props.ragEngines?.documentsByCompountKeyIndexName ?? "",
         SAGEMAKER_RAG_MODELS_ENDPOINT:
           props.ragEngines?.sageMakerRagModelsEndpoint?.attrEndpointName ?? "",
+        DELETE_WORKSPACE_WORKFLOW_ARN:
+          props.ragEngines?.deleteWorkspaceWorkflow?.stateMachineArn ?? "",
         CREATE_AURORA_WORKSPACE_WORKFLOW_ARN:
           props.ragEngines?.auroraPgVector?.createAuroraWorkspaceWorkflow
             ?.stateMachineArn ?? "",
@@ -206,6 +208,10 @@ export class RestApi extends Construct {
       props.ragEngines.websiteCrawlingWorkflow.grantStartExecution(apiHandler);
     }
 
+    if (props.ragEngines?.deleteWorkspaceWorkflow) {
+      props.ragEngines.deleteWorkspaceWorkflow.grantStartExecution(apiHandler);
+    }
+
     if (props.ragEngines?.sageMakerRagModelsEndpoint) {
       apiHandler.addToRolePolicy(
         new iam.PolicyStatement({

lib/rag-engines/data-import/functions/website-crawling-workflow/website-parser/index.py

@@ -19,7 +19,7 @@ def lambda_handler(event, context: LambdaContext):
     workspace_id = event["workspace_id"]
     document_id = event["document_id"]
     response = s3_client.get_object(Bucket=bucket_name, Key=object_key)
-    file_content = response['Body'].read().decode('utf-8')
+    file_content = response["Body"].read().decode("utf-8")
     data = json.loads(file_content)
 
     iteration = data["iteration"]
@@ -31,8 +31,7 @@ def lambda_handler(event, context: LambdaContext):
     follow_links = data["follow_links"]
     limit = data["limit"]
 
-    logger.info(
-        f"Processing document {document_id} in workspace {workspace_id}")
+    logger.info(f"Processing document {document_id} in workspace {workspace_id}")
     logger.info(f"Workspace: {workspace}")
     logger.info(f"Document: {document}")
     logger.info(f"Limit: {limit}")
@@ -55,7 +54,9 @@ def lambda_handler(event, context: LambdaContext):
     result["iteration"] = iteration
     result["crawler_job_id"] = crawler_job_id
 
-    iteration_object_key = f"{workspace_id}/{document_id}/crawler/{crawler_job_id}/{iteration}.json"
+    iteration_object_key = (
+        f"{workspace_id}/{document_id}/crawler/{crawler_job_id}/{iteration}.json"
+    )
     s3_client.put_object(
         Bucket=PROCESSING_BUCKET_NAME,
         Key=iteration_object_key,

lib/rag-engines/index.ts

@@ -7,6 +7,7 @@ import { DataImport } from "./data-import";
 import { RagDynamoDBTables } from "./rag-dynamodb-tables";
 import { OpenSearchVector } from "./opensearch-vector";
 import { KendraRetrieval } from "./kendra-retrieval";
+import { Workspaces } from "./workspaces";
 import * as sagemaker from "aws-cdk-lib/aws-sagemaker";
 import * as s3 from "aws-cdk-lib/aws-s3";
 import * as dynamodb from "aws-cdk-lib/aws-dynamodb";
@@ -30,6 +31,7 @@ export class RagEngines extends Construct {
   public readonly sageMakerRagModelsEndpoint: sagemaker.CfnEndpoint;
   public readonly fileImportWorkflow?: sfn.StateMachine;
   public readonly websiteCrawlingWorkflow?: sfn.StateMachine;
+  public readonly deleteWorkspaceWorkflow?: sfn.StateMachine;
 
   constructor(scope: Construct, id: string, props: RagEnginesProps) {
     super(scope, id);
@@ -86,6 +88,16 @@ export class RagEngines extends Construct {
       kendraRetrieval: kendraRetrieval ?? undefined,
     });
 
+    const workspaces = new Workspaces(this, "Workspaces", {
+      shared: props.shared,
+      config: props.config,
+      dataImport,
+      ragDynamoDBTables: tables,
+      auroraPgVector: auroraPgVector ?? undefined,
+      openSearchVector: openSearchVector ?? undefined,
+      kendraRetrieval: kendraRetrieval ?? undefined,
+    });
+
     this.auroraPgVector = auroraPgVector;
     this.openSearchVector = openSearchVector;
     this.kendraRetrieval = kendraRetrieval;
@@ -100,5 +112,6 @@ export class RagEngines extends Construct {
       tables.documentsByCompountKeyIndexName;
     this.fileImportWorkflow = dataImport.fileImportWorkflow;
     this.websiteCrawlingWorkflow = dataImport.websiteCrawlingWorkflow;
+    this.deleteWorkspaceWorkflow = workspaces.deleteWorkspaceWorkflow;
   }
 }