Derived value of crossEncodingEnabled based on enableEmbeddingModelsV…

…iaSagemaker config.
azaylamba · Feb 9, 2024 · 13ce71e · 13ce71e
1 parent cf0dfc1
commit 13ce71e
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 10 deletions.
diff --git a/bin/config.ts b/bin/config.ts
@@ -25,6 +25,7 @@ export function getConfig(): SystemConfig {
     },
     rag: {
       enabled: false,
+      enableEmbeddingModelsViaSagemaker: false,
       engines: {
         aurora: {
           enabled: false,
@@ -61,14 +62,14 @@ export function getConfig(): SystemConfig {
           dimensions: 1536,
         },
       ],
+      crossEncodingEnabled: false,
       crossEncoderModels: [
         {
           provider: "sagemaker",
           name: "cross-encoder/ms-marco-MiniLM-L-12-v2",
           default: true,
         },
       ],
-      crossEncodingEnabled: false,
     },
   };
 }

diff --git a/cli/magic-config.ts b/cli/magic-config.ts
@@ -203,6 +203,15 @@ async function processCreateOptions(options: any): Promise<void> {
       message: "Do you want to enable RAG",
       initial: options.enableRag || false,
     },
+    {
+      type: "confirm",
+      name: "enableEmbeddingModelsViaSagemaker",
+      message: "Do you want to enable embedding models via SageMaker?",
+      initial: options.enableEmbeddingModelsViaSagemaker || false,
+      skip(): boolean {
+        return !(this as any).state.answers.enableRag;
+      },
+    },
     {
       type: "multiselect",
       name: "ragsToEnable",
@@ -349,10 +358,13 @@ async function processCreateOptions(options: any): Promise<void> {
         }
       : undefined,
     llms: {
+      enableSagemakerModels: answers.enableSagemakerModels,
       sagemaker: answers.sagemakerModels,
     },
     rag: {
       enabled: answers.enableRag,
+      enableEmbeddingModelsViaSagemaker:
+        answers.enableEmbeddingModelsViaSagemaker,
       engines: {
         aurora: {
           enabled: answers.ragsToEnable.includes("aurora"),
@@ -367,6 +379,7 @@ async function processCreateOptions(options: any): Promise<void> {
           enterprise: false,
         },
       },
+      crossEncodingEnabled: answers.enableEmbeddingModelsViaSagemaker,
       embeddingsModels: [{}],
       crossEncoderModels: [{}],
     },
@@ -377,12 +390,24 @@ async function processCreateOptions(options: any): Promise<void> {
     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;
+  if (answers.enableEmbeddingModelsViaSagemaker && answers.enableSagemakerModels) {
+    config.rag.crossEncoderModels[0] = {
+      provider: "sagemaker",
+      name: "cross-encoder/ms-marco-MiniLM-L-12-v2",
+      default: true,
+    };
+  } else {
+    config.rag.crossEncoderModels[0] = {
+      provider: "None",
+      name: "None",
+      default: true,
+    };
+  }
+  if (!config.rag.enableEmbeddingModelsViaSagemaker) {
+    config.rag.embeddingsModels = embeddingModels.filter(model => model.provider !== "sagemaker");
+  } else {
+    config.rag.embeddingsModels = embeddingModels;
+  }
   config.rag.embeddingsModels.forEach((m: any) => {
     if (m.name === models.defaultEmbedding) {
       m.default = true;

diff --git a/lib/aws-genai-llm-chatbot-stack.ts b/lib/aws-genai-llm-chatbot-stack.ts
@@ -152,8 +152,7 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
       api: chatBotApi,
       chatbotFilesBucket: chatBotApi.filesBucket,
       crossEncodersEnabled: props.config.rag.crossEncodingEnabled,
-      sagemakerEmbeddingsEnabled:
-        typeof ragEngines?.sageMakerRagModels?.model !== "undefined",
+      sagemakerEmbeddingsEnabled: props.config.rag.enableEmbeddingModelsViaSagemaker,
     });
 
     /**

diff --git a/lib/shared/types.ts b/lib/shared/types.ts
@@ -88,6 +88,7 @@ export interface SystemConfig {
   };
   rag: {
     enabled: boolean;
+    enableEmbeddingModelsViaSagemaker: boolean;
     engines: {
       aurora: {
         enabled: boolean;
@@ -113,12 +114,12 @@ export interface SystemConfig {
       dimensions: number;
       default?: boolean;
     }[];
+    crossEncodingEnabled: boolean;
     crossEncoderModels: {
       provider: ModelProvider;
       name: string;
       default?: boolean;
     }[];
-    crossEncodingEnabled: boolean;
   };
 }