Feature request: Option to disable cross encoder models

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 };
@@ -177,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
       );
@@ -575,7 +576,8 @@ async function processCreateOptions(options: any): Promise<void> {
     {
       type: "confirm",
       name: "enableEmbeddingModelsViaSagemaker",
-      message: "Do you want to enable embedding models via SageMaker?",
+      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;
@@ -1091,10 +1093,7 @@ async function processCreateOptions(options: any): Promise<void> {
     },
   };
 
-  if (
-    answers.enableEmbeddingModelsViaSagemaker &&
-    answers.enableSagemakerModels
-  ) {
+  if (config.rag.crossEncodingEnabled) {
     config.rag.crossEncoderModels[0] = {
       provider: "sagemaker",
       name: "cross-encoder/ms-marco-MiniLM-L-12-v2",
@@ -1108,7 +1107,9 @@ async function processCreateOptions(options: any): Promise<void> {
     };
   }
   if (!config.rag.enableEmbeddingModelsViaSagemaker) {
-    config.rag.embeddingsModels = embeddingModels.filter(model => model.provider !== "sagemaker");
+    config.rag.embeddingsModels = embeddingModels.filter(
+      (model) => model.provider !== "sagemaker"
+    );
   } else {
     config.rag.embeddingsModels = embeddingModels;
   }

lib/rag-engines/index.ts

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

lib/user-interface/index.ts

@@ -15,8 +15,8 @@ import { Shared } from "../shared";
 import { SystemConfig } from "../shared/types";
 import { Utils } from "../shared/utils";
 import { ChatBotApi } from "../chatbot-api";
-import { PrivateWebsite } from "./private-website"
-import { PublicWebsite } from "./public-website"
+import { PrivateWebsite } from "./private-website";
+import { PublicWebsite } from "./public-website";
 import { NagSuppressions } from "cdk-nag";
 
 export interface UserInterfaceProps {
@@ -52,13 +52,13 @@ export class UserInterface extends Construct {
       removalPolicy: cdk.RemovalPolicy.DESTROY,
       blockPublicAccess: s3.BlockPublicAccess.BLOCK_ALL,
       autoDeleteObjects: true,
-      bucketName: props.config.privateWebsite ? props.config.domain : undefined, 
+      bucketName: props.config.privateWebsite ? props.config.domain : undefined,
       websiteIndexDocument: "index.html",
       websiteErrorDocument: "index.html",
       enforceSSL: true,
       serverAccessLogsBucket: uploadLogsBucket,
     });
-    
+
     // Deploy either Private (only accessible within VPC) or Public facing website
     let apiEndpoint: string;
     let websocketEndpoint: string;
@@ -74,8 +74,6 @@ export class UserInterface extends Construct {
       this.publishedDomain = distribution.distributionDomainName;
     }
 
-
-
     const exportsAsset = s3deploy.Source.jsonData("aws-exports.json", {
       aws_project_region: cdk.Aws.REGION,
       aws_cognito_region: cdk.Aws.REGION,
@@ -118,6 +116,8 @@ export class UserInterface extends Construct {
         rag_enabled: props.config.rag.enabled,
         cross_encoders_enabled: props.crossEncodersEnabled,
         sagemaker_embeddings_enabled: props.sagemakerEmbeddingsEnabled,
+        enable_embedding_models_via_sagemaker:
+          props.config.rag.enableEmbeddingModelsViaSagemaker,
         default_embeddings_model: Utils.getDefaultEmbeddingsModel(props.config),
         default_cross_encoder_model: Utils.getDefaultCrossEncoderModel(
           props.config
@@ -221,21 +221,17 @@ export class UserInterface extends Construct {
       prune: false,
       sources: [asset, exportsAsset],
       destinationBucket: websiteBucket,
-      distribution: props.config.privateWebsite ? undefined : distribution
+      distribution: props.config.privateWebsite ? undefined : distribution,
     });
 
-
     /**
      * CDK NAG suppression
      */
-    NagSuppressions.addResourceSuppressions(
-      uploadLogsBucket, 
-      [
-        {
-          id: "AwsSolutions-S1",
-          reason: "Bucket is the server access logs bucket for websiteBucket.",
-        },
-      ]
-    );
+    NagSuppressions.addResourceSuppressions(uploadLogsBucket, [
+      {
+        id: "AwsSolutions-S1",
+        reason: "Bucket is the server access logs bucket for websiteBucket.",
+      },
+    ]);
   }
 }

lib/user-interface/react-app/src/common/helpers/embeddings-model-helper.ts

@@ -1,5 +1,6 @@
 import { SelectProps } from "@cloudscape-design/components";
 import { EmbeddingModel } from "../../API";
+import { AppConfig } from "../types";
 
 export abstract class EmbeddingsModelHelper {
   static getSelectOption(model?: string): SelectProps.Option | null {
@@ -32,9 +33,18 @@ export abstract class EmbeddingsModelHelper {
     };
   }
 
-  static getSelectOptions(embeddingsModels: EmbeddingModel[]) {
+  static getSelectOptions(
+    appContext: AppConfig | null,
+    embeddingsModels: EmbeddingModel[]
+  ) {
     const modelsMap = new Map<string, EmbeddingModel[]>();
     embeddingsModels.forEach((model) => {
+      if (
+        model.provider === "sagemaker" &&
+        !appContext?.config.enable_embedding_models_via_sagemaker
+      ) {
+        return;
+      }
       let items = modelsMap.get(model.provider);
       if (!items) {
         items = [];

lib/user-interface/react-app/src/common/types.ts

@@ -25,6 +25,7 @@ export interface AppConfig {
     rag_enabled: boolean;
     cross_encoders_enabled: boolean;
     sagemaker_embeddings_enabled: boolean;
+    enable_embedding_models_via_sagemaker: boolean;
     api_endpoint: string;
     websocket_endpoint: string;
     default_embeddings_model: string;

lib/user-interface/react-app/src/pages/rag/create-workspace/aurora-form.tsx

@@ -112,22 +112,20 @@ function AuroraFooter(props: {
             Create an index
           </Toggle>
         </FormField>
-        {props.data.crossEncodingEnabled && (
-          <>
-            <HybridSearchField
-              submitting={props.submitting}
-              errors={props.errors}
-              checked={props.data.hybridSearch}
-              onChange={props.onChange}
-            />
-            <CrossEncoderSelectorField
-              errors={props.errors}
-              submitting={props.submitting}
-              selectedModel={props.data.crossEncoderModel}
-              onChange={props.onChange}
-            />
-          </>
-        )}
+        <HybridSearchField
+          submitting={props.submitting}
+          disabled={!props.data.crossEncodingEnabled}
+          errors={props.errors}
+          checked={props.data.hybridSearch}
+          onChange={props.onChange}
+        />
+        <CrossEncoderSelectorField
+          errors={props.errors}
+          submitting={props.submitting}
+          disabled={!props.data.crossEncodingEnabled}
+          selectedModel={props.data.crossEncoderModel}
+          onChange={props.onChange}
+        />
         <ChunkSelectorField
           submitting={props.submitting}
           onChange={props.onChange}

lib/user-interface/react-app/src/pages/rag/create-workspace/cross-encoder-selector-field.tsx

@@ -9,6 +9,7 @@ import { Utils } from "../../../common/utils";
 
 interface CrossEncoderSelectorProps {
   submitting: boolean;
+  disabled: boolean;
   onChange: (data: Partial<{ crossEncoderModel: SelectProps.Option }>) => void;
   selectedModel: SelectProps.Option | null;
   errors: Record<string, string | string[]>;
@@ -45,7 +46,7 @@ export function CrossEncoderSelectorField(props: CrossEncoderSelectorProps) {
 
   return (
     <Select
-      disabled={props.submitting}
+      disabled={props.submitting || props.disabled}
       selectedAriaLabel="Selected"
       placeholder="Choose a cross-encoder model"
       statusType={crossEncoderModelsStatus}

lib/user-interface/react-app/src/pages/rag/create-workspace/embeddings-selector-field.tsx

@@ -40,7 +40,7 @@ export default function EmbeddingSelector(props: EmbeddingsSelectionProps) {
   }, [appContext]);
 
   const embeddingsModelOptions =
-    EmbeddingsModelHelper.getSelectOptions(embeddingsModels);
+    EmbeddingsModelHelper.getSelectOptions(appContext, embeddingsModels);
 
   return (
     <FormField

lib/user-interface/react-app/src/pages/rag/create-workspace/hybrid-search-field.tsx

@@ -2,6 +2,7 @@ import { FormField, Toggle } from "@cloudscape-design/components";
 
 interface HybridSearchProps {
   submitting: boolean;
+  disabled: boolean;
   onChange: (data: Partial<{ hybridSearch: boolean }>) => void;
   checked: boolean;
   errors: Record<string, string | string[]>;
@@ -15,7 +16,7 @@ export function HybridSearchField(props: HybridSearchProps) {
       errorText={props.errors.hybridSearch}
     >
       <Toggle
-        disabled={props.submitting}
+        disabled={props.submitting || props.disabled}
         checked={props.checked}
         onChange={({ detail: { checked } }) =>
           props.onChange({ hybridSearch: checked })

lib/user-interface/react-app/src/pages/rag/create-workspace/opensearch-form.tsx

@@ -70,22 +70,20 @@ function OpenSearchFooter(props: {
   return (
     <ExpandableSection headerText="Additional settings" variant="footer">
       <SpaceBetween size="l">
-        {props.data.crossEncodingEnabled && (
-          <>
-            <HybridSearchField
-              submitting={props.submitting}
-              errors={props.errors}
-              checked={props.data.hybridSearch}
-              onChange={props.onChange}
-            />
-            <CrossEncoderSelectorField
-              errors={props.errors}
-              submitting={props.submitting}
-              selectedModel={props.data.crossEncoderModel}
-              onChange={props.onChange}
-            />
-          </>
-        )}
+        <HybridSearchField
+          submitting={props.submitting}
+          disabled={!props.data.crossEncodingEnabled}
+          errors={props.errors}
+          checked={props.data.hybridSearch}
+          onChange={props.onChange}
+        />
+        <CrossEncoderSelectorField
+          errors={props.errors}
+          submitting={props.submitting}
+          disabled={!props.data.crossEncodingEnabled}
+          selectedModel={props.data.crossEncoderModel}
+          onChange={props.onChange}
+        />
         <ChunkSelectorField
           submitting={props.submitting}
           onChange={props.onChange}

lib/user-interface/react-app/src/pages/rag/embeddings/embeddings.tsx

@@ -186,6 +186,7 @@ export default function Embeddings() {
   };
 
   const embeddingsModelOptions = EmbeddingsModelHelper.getSelectOptions(
+    appContext,
     embeddingsModelsResults
   );