Feature request: Option to disable cross encoder models

bin/config.ts

@@ -23,10 +23,12 @@ export function getConfig(): SystemConfig {
     },
     llms: {
       // sagemaker: [SupportedSageMakerModels.FalconLite]
+      enableSagemakerModels: false,
       sagemaker: [],
     },
     rag: {
       enabled: false,
+      enableEmbeddingModelsViaSagemaker: false,
       engines: {
         aurora: {
           enabled: false,
@@ -45,11 +47,13 @@ export function getConfig(): SystemConfig {
           provider: "sagemaker",
           name: "intfloat/multilingual-e5-large",
           dimensions: 1024,
+          default: false,
         },
         {
           provider: "sagemaker",
           name: "sentence-transformers/all-MiniLM-L6-v2",
           dimensions: 384,
+          default: false,
         },
         {
           provider: "bedrock",
@@ -77,8 +81,10 @@ export function getConfig(): SystemConfig {
           provider: "openai",
           name: "text-embedding-ada-002",
           dimensions: 1536,
+          default: false,
         },
       ],
+      crossEncodingEnabled: false,
       crossEncoderModels: [
         {
           provider: "sagemaker",

cli/magic-config.ts

@@ -32,7 +32,6 @@ function getTimeZonesWithCurrentTime(): { message: string; name: string }[] {
 function getCountryCodesAndNames(): { message: string; name: string }[] {
   // Use country-list to get an array of countries with their codes and names
   const countries = getData();
-
   // Map the country data to match the desired output structure
   const countryInfo = countries.map(({ code, name }) => {
     return { message: `${name} (${code})`, name: code };
@@ -91,16 +90,19 @@ const embeddingModels = [
     provider: "sagemaker",
     name: "intfloat/multilingual-e5-large",
     dimensions: 1024,
+    default: false,
   },
   {
     provider: "sagemaker",
     name: "sentence-transformers/all-MiniLM-L6-v2",
     dimensions: 384,
+    default: false,
   },
   {
     provider: "bedrock",
     name: "amazon.titan-embed-text-v1",
     dimensions: 1536,
+    default: false,
   },
   //Support for inputImage is not yet implemented for amazon.titan-embed-image-v1
   {
@@ -122,6 +124,7 @@ const embeddingModels = [
     provider: "openai",
     name: "text-embedding-ada-002",
     dimensions: 1536,
+    default: false,
   },
 ];
 
@@ -173,6 +176,8 @@ const embeddingModels = [
       options.startScheduleEndDate =
         config.llms?.sagemakerSchedule?.startScheduleEndDate;
       options.enableRag = config.rag.enabled;
+      options.enableEmbeddingModelsViaSagemaker =
+        config.rag.enableEmbeddingModelsViaSagemaker;
       options.ragsToEnable = Object.keys(config.rag.engines ?? {}).filter(
         (v: string) => (config.rag.engines as any)[v].enabled
       );
@@ -568,6 +573,16 @@ 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 and cross-encoder models via SageMaker?",
+      initial: options.enableEmbeddingModelsViaSagemaker || false,
+      skip(): boolean {
+        return !(this as any).state.answers.enableRag;
+      },
+    },
     {
       type: "multiselect",
       name: "ragsToEnable",
@@ -696,7 +711,6 @@ async function processCreateOptions(options: any): Promise<void> {
         if ((this as any).state.answers.enableRag) {
           return value ? true : "Select a default embedding model";
         }
-
         return true;
       },
       skip() {
@@ -1037,6 +1051,7 @@ async function processCreateOptions(options: any): Promise<void> {
         }
       : undefined,
     llms: {
+      enableSagemakerModels: answers.enableSagemakerModels,
       sagemaker: answers.sagemakerModels,
       huggingfaceApiSecretArn: answers.huggingfaceApiSecretArn,
       sagemakerSchedule: answers.enableSagemakerModelsSchedule
@@ -1056,6 +1071,8 @@ async function processCreateOptions(options: any): Promise<void> {
     },
     rag: {
       enabled: answers.enableRag,
+      enableEmbeddingModelsViaSagemaker:
+        answers.enableEmbeddingModelsViaSagemaker,
       engines: {
         aurora: {
           enabled: answers.ragsToEnable.includes("aurora"),
@@ -1070,28 +1087,43 @@ async function processCreateOptions(options: any): Promise<void> {
           enterprise: false,
         },
       },
+      crossEncodingEnabled: answers.enableEmbeddingModelsViaSagemaker,
       embeddingsModels: [{}],
       crossEncoderModels: [{}],
     },
   };
 
+  if (config.rag.crossEncodingEnabled) {
+    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;
+  }
   // 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.embeddingsModels[0] as any).default = true;
+  } else {
+    config.rag.embeddingsModels.forEach((m: any) => {
+      if (m.name === models.defaultEmbedding) {
+        m.default = true;
+      }
+    });
   }
 
-  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 =

lib/aws-genai-llm-chatbot-stack.ts

@@ -154,10 +154,9 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
       identityPool: authentication.identityPool,
       api: chatBotApi,
       chatbotFilesBucket: chatBotApi.filesBucket,
-      crossEncodersEnabled:
-        typeof ragEngines?.sageMakerRagModels?.model !== "undefined",
+      crossEncodersEnabled: props.config.rag.crossEncodingEnabled,
       sagemakerEmbeddingsEnabled:
-        typeof ragEngines?.sageMakerRagModels?.model !== "undefined",
+        props.config.rag.enableEmbeddingModelsViaSagemaker,
     });
 
     if (props.config.cognitoFederation?.enabled) {
@@ -317,10 +316,7 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
         ]
       );
 
-      if (
-        props.config.rag.engines.aurora.enabled ||
-        props.config.rag.engines.opensearch.enabled
-      ) {
+      if (props.config.llms.enableSagemakerModels) {
         NagSuppressions.addResourceSuppressionsByPath(
           this,
           [
@@ -356,59 +352,59 @@ export class AwsGenAILLMChatbotStack extends cdk.Stack {
             },
           ]
         );
-        if (props.config.rag.engines.aurora.enabled) {
-          NagSuppressions.addResourceSuppressionsByPath(
-            this,
-            `/${this.stackName}/RagEngines/AuroraPgVector/AuroraDatabase/Secret/Resource`,
-            [
-              {
-                id: "AwsSolutions-SMG4",
-                reason: "Secret created implicitly by CDK.",
-              },
-            ]
-          );
-          NagSuppressions.addResourceSuppressionsByPath(
-            this,
-            [
-              `/${this.stackName}/RagEngines/AuroraPgVector/DatabaseSetupFunction/ServiceRole/Resource`,
-              `/${this.stackName}/RagEngines/AuroraPgVector/DatabaseSetupProvider/framework-onEvent/ServiceRole/Resource`,
-              `/${this.stackName}/RagEngines/AuroraPgVector/DatabaseSetupProvider/framework-onEvent/ServiceRole/DefaultPolicy/Resource`,
-              `/${this.stackName}/RagEngines/AuroraPgVector/CreateAuroraWorkspace/CreateAuroraWorkspaceFunction/ServiceRole/Resource`,
-              `/${this.stackName}/RagEngines/AuroraPgVector/CreateAuroraWorkspace/CreateAuroraWorkspaceFunction/ServiceRole/DefaultPolicy/Resource`,
-              `/${this.stackName}/RagEngines/AuroraPgVector/CreateAuroraWorkspace/CreateAuroraWorkspace/Role/DefaultPolicy/Resource`,
-            ],
-            [
-              {
-                id: "AwsSolutions-IAM4",
-                reason: "IAM role implicitly created by CDK.",
-              },
-              {
-                id: "AwsSolutions-IAM5",
-                reason: "IAM role implicitly created by CDK.",
-              },
-            ]
-          );
-        }
-        if (props.config.rag.engines.opensearch.enabled) {
-          NagSuppressions.addResourceSuppressionsByPath(
-            this,
-            [
-              `/${this.stackName}/RagEngines/OpenSearchVector/CreateOpenSearchWorkspace/CreateOpenSearchWorkspaceFunction/ServiceRole/Resource`,
-              `/${this.stackName}/RagEngines/OpenSearchVector/CreateOpenSearchWorkspace/CreateOpenSearchWorkspaceFunction/ServiceRole/DefaultPolicy/Resource`,
-              `/${this.stackName}/RagEngines/OpenSearchVector/CreateOpenSearchWorkspace/CreateOpenSearchWorkspace/Role/DefaultPolicy/Resource`,
-            ],
-            [
-              {
-                id: "AwsSolutions-IAM4",
-                reason: "IAM role implicitly created by CDK.",
-              },
-              {
-                id: "AwsSolutions-IAM5",
-                reason: "IAM role implicitly created by CDK.",
-              },
-            ]
-          );
-        }
+      }
+      if (props.config.rag.engines.aurora.enabled) {
+        NagSuppressions.addResourceSuppressionsByPath(
+          this,
+          `/${this.stackName}/RagEngines/AuroraPgVector/AuroraDatabase/Secret/Resource`,
+          [
+            {
+              id: "AwsSolutions-SMG4",
+              reason: "Secret created implicitly by CDK.",
+            },
+          ]
+        );
+        NagSuppressions.addResourceSuppressionsByPath(
+          this,
+          [
+            `/${this.stackName}/RagEngines/AuroraPgVector/DatabaseSetupFunction/ServiceRole/Resource`,
+            `/${this.stackName}/RagEngines/AuroraPgVector/DatabaseSetupProvider/framework-onEvent/ServiceRole/Resource`,
+            `/${this.stackName}/RagEngines/AuroraPgVector/DatabaseSetupProvider/framework-onEvent/ServiceRole/DefaultPolicy/Resource`,
+            `/${this.stackName}/RagEngines/AuroraPgVector/CreateAuroraWorkspace/CreateAuroraWorkspaceFunction/ServiceRole/Resource`,
+            `/${this.stackName}/RagEngines/AuroraPgVector/CreateAuroraWorkspace/CreateAuroraWorkspaceFunction/ServiceRole/DefaultPolicy/Resource`,
+            `/${this.stackName}/RagEngines/AuroraPgVector/CreateAuroraWorkspace/CreateAuroraWorkspace/Role/DefaultPolicy/Resource`,
+          ],
+          [
+            {
+              id: "AwsSolutions-IAM4",
+              reason: "IAM role implicitly created by CDK.",
+            },
+            {
+              id: "AwsSolutions-IAM5",
+              reason: "IAM role implicitly created by CDK.",
+            },
+          ]
+        );
+      }
+      if (props.config.rag.engines.opensearch.enabled) {
+        NagSuppressions.addResourceSuppressionsByPath(
+          this,
+          [
+            `/${this.stackName}/RagEngines/OpenSearchVector/CreateOpenSearchWorkspace/CreateOpenSearchWorkspaceFunction/ServiceRole/Resource`,
+            `/${this.stackName}/RagEngines/OpenSearchVector/CreateOpenSearchWorkspace/CreateOpenSearchWorkspaceFunction/ServiceRole/DefaultPolicy/Resource`,
+            `/${this.stackName}/RagEngines/OpenSearchVector/CreateOpenSearchWorkspace/CreateOpenSearchWorkspace/Role/DefaultPolicy/Resource`,
+          ],
+          [
+            {
+              id: "AwsSolutions-IAM4",
+              reason: "IAM role implicitly created by CDK.",
+            },
+            {
+              id: "AwsSolutions-IAM5",
+              reason: "IAM role implicitly created by CDK.",
+            },
+          ]
+        );
       }
       if (props.config.rag.engines.kendra.enabled) {
         NagSuppressions.addResourceSuppressionsByPath(

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

@@ -30,14 +30,19 @@ def models():
 @tracer.capture_method
 def cross_encoders(input: dict):
     request = CrossEncodersRequest(**input)
-    selected_model = genai_core.cross_encoder.get_cross_encoder_model(
-        request.provider, request.model
-    )
-
-    if selected_model is None:
-        raise genai_core.types.CommonError("Model not found")
-
-    ret_value = genai_core.cross_encoder.rank_passages(
-        selected_model, request.reference, request.passages
-    )
-    return [{"score": v, "passage": p} for v, p in zip(ret_value, request.passages)]
+    config = genai_core.parameters.get_config()
+    crossEncodingEnabled = config["rag"]["crossEncodingEnabled"]
+    if (crossEncodingEnabled):
+        selected_model = genai_core.cross_encoder.get_cross_encoder_model(
+            request.provider, request.model
+        )
+
+        if selected_model is None:
+            raise genai_core.types.CommonError("Model not found")
+
+        ret_value = genai_core.cross_encoder.rank_passages(
+            selected_model, request.reference, request.passages
+        )
+        return [{"score": v, "passage": p} for v, p in zip(ret_value, request.passages)]
+
+    return [{"score": 0, "passage": p} for p in request.passages]

lib/rag-engines/data-import/index.ts

@@ -132,7 +132,7 @@ export class DataImport extends Construct {
         processingBucket,
         auroraDatabase: props.auroraDatabase,
         ragDynamoDBTables: props.ragDynamoDBTables,
-        sageMakerRagModelsEndpoint: props.sageMakerRagModels?.model.endpoint,
+        sageMakerRagModelsEndpoint: props.sageMakerRagModels?.model?.endpoint,
         openSearchVector: props.openSearchVector,
       }
     );

lib/rag-engines/index.ts

@@ -42,10 +42,7 @@ export class RagEngines extends Construct {
     const tables = new RagDynamoDBTables(this, "RagDynamoDBTables");
 
     let sageMakerRagModels: SageMakerRagModels | null = null;
-    if (
-      props.config.rag.engines.aurora.enabled ||
-      props.config.rag.engines.opensearch.enabled
-    ) {
+    if (props.config.rag.crossEncodingEnabled) {
       sageMakerRagModels = new SageMakerRagModels(this, "SageMaker", {
         shared: props.shared,
         config: props.config,