From 29e3a3ce00f08bc294eab245fa6440eaded2ba13 Mon Sep 17 00:00:00 2001
From: Tyler Ohlsen <ohltyler@amazon.com>
Date: Fri, 27 Dec 2024 10:44:09 -0800
Subject: [PATCH] Update index settings based on ingest processors

Signed-off-by: Tyler Ohlsen <ohltyler@amazon.com>
---
 .../ingest_inputs/advanced_settings.tsx       | 66 ++++++++++++++++---
 public/utils/utils.ts                         | 30 ++++++++-
 2 files changed, 85 insertions(+), 11 deletions(-)

diff --git a/public/pages/workflow_detail/workflow_inputs/ingest_inputs/advanced_settings.tsx b/public/pages/workflow_detail/workflow_inputs/ingest_inputs/advanced_settings.tsx
index 7402ba3e..a9bf8f42 100644
--- a/public/pages/workflow_detail/workflow_inputs/ingest_inputs/advanced_settings.tsx
+++ b/public/pages/workflow_detail/workflow_inputs/ingest_inputs/advanced_settings.tsx
@@ -3,7 +3,9 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import React from 'react';
+import React, { useEffect } from 'react';
+import { useSelector } from 'react-redux';
+import { isEmpty } from 'lodash';
 import {
   EuiAccordion,
   EuiFlexGroup,
@@ -11,6 +13,14 @@ import {
   EuiSpacer,
 } from '@elastic/eui';
 import { JsonField } from '../input_fields';
+import { getIn, useFormikContext } from 'formik';
+import { WorkflowFormValues } from '../../../../../common';
+import { AppState } from '../../../../store';
+import {
+  getEmbeddingDimensions,
+  getUpdatedIndexSettings,
+  isKnnIndex,
+} from '../../../../utils';
 
 interface AdvancedSettingsProps {}
 
@@ -18,6 +28,50 @@ interface AdvancedSettingsProps {}
  * Input component for configuring ingest-side advanced settings
  */
 export function AdvancedSettings(props: AdvancedSettingsProps) {
+  const { values, setFieldValue } = useFormikContext<WorkflowFormValues>();
+  const { models, connectors } = useSelector((state: AppState) => state.ml);
+  const ingestProcessors = Object.values(values?.ingest?.enrich) as [];
+  const ingestProcessorModelIds = ingestProcessors
+    .map((ingestProcessor) => ingestProcessor?.model?.id as string | undefined)
+    .filter((modelId) => !isEmpty(modelId));
+  const indexMappingsPath = 'ingest.index.mappings';
+  const indexSettingsPath = 'ingest.index.settings';
+  const curMappings = getIn(values, indexMappingsPath);
+  const curSettings = getIn(values, indexSettingsPath);
+
+  // listen on when processor with models are added / removed. dynamically update index
+  // mappings and settings, if applicable.
+  useEffect(() => {
+    if (ingestProcessorModelIds.length > 0) {
+      ingestProcessorModelIds.forEach((ingestProcessorModelId) => {
+        const processorModel = Object.values(models).find(
+          (model) => model.id === ingestProcessorModelId
+        );
+        if (processorModel?.connectorId !== undefined) {
+          const processorConnector = connectors[processorModel?.connectorId];
+          const dimension = getEmbeddingDimensions(processorConnector);
+          if (dimension !== undefined) {
+            // TODO: update mappings
+            if (!isKnnIndex(curSettings)) {
+              setFieldValue(
+                indexSettingsPath,
+                getUpdatedIndexSettings(curSettings, true)
+              );
+            }
+          }
+        }
+      });
+    } else {
+      // TODO: update mappings
+      if (isKnnIndex(curSettings)) {
+        setFieldValue(
+          indexSettingsPath,
+          getUpdatedIndexSettings(curSettings, false)
+        );
+      }
+    }
+  }, [ingestProcessorModelIds.length]);
+
   return (
     <EuiFlexGroup direction="column">
       <EuiFlexItem grow={false}>
@@ -25,16 +79,10 @@ export function AdvancedSettings(props: AdvancedSettingsProps) {
           <EuiSpacer size="s" />
           <EuiFlexGroup direction="column">
             <EuiFlexItem>
-              <JsonField
-                label="Index mappings"
-                fieldPath={'ingest.index.mappings'}
-              />
+              <JsonField label="Index mappings" fieldPath={indexMappingsPath} />
             </EuiFlexItem>
             <EuiFlexItem>
-              <JsonField
-                label="Index settings"
-                fieldPath={'ingest.index.settings'}
-              />
+              <JsonField label="Index settings" fieldPath={indexSettingsPath} />
             </EuiFlexItem>
           </EuiFlexGroup>
         </EuiAccordion>
diff --git a/public/utils/utils.ts b/public/utils/utils.ts
index 647910ac..acf5aaf4 100644
--- a/public/utils/utils.ts
+++ b/public/utils/utils.ts
@@ -5,7 +5,7 @@
 
 import yaml from 'js-yaml';
 import jsonpath from 'jsonpath';
-import { escape, get, isEmpty } from 'lodash';
+import { escape, get, isEmpty, set } from 'lodash';
 import semver from 'semver';
 import queryString from 'query-string';
 import { useLocation } from 'react-router-dom';
@@ -30,6 +30,7 @@ import {
   BEDROCK_DIMENSIONS,
   COHERE_DIMENSIONS,
   OPENAI_DIMENSIONS,
+  customStringify,
 } from '../../common';
 import { getCore, getDataSourceEnabled } from '../services';
 import {
@@ -608,7 +609,7 @@ export function injectParameters(
   return finalQueryString;
 }
 
-// fetch embedding dimensions, if the selected model is a known one
+// Fetch embedding dimensions, if the selected model is a known one
 export function getEmbeddingDimensions(
   connector: Connector
 ): number | undefined {
@@ -629,3 +630,28 @@ export function getEmbeddingDimensions(
     return undefined;
   }
 }
+
+// Check if an index is a knn index
+export function isKnnIndex(existingSettings: string): boolean {
+  try {
+    return get(JSON.parse(existingSettings), 'index.knn', false);
+  } catch (error) {
+    console.error('Could not parse index settings: ', error);
+    return false;
+  }
+}
+
+// Update the index settings based on parameters passed.
+// Currently just used for updating the `knn` flag.
+export function getUpdatedIndexSettings(
+  existingSettings: string,
+  knnBool: boolean
+): string {
+  try {
+    return customStringify(
+      set(JSON.parse(existingSettings), 'index.knn', knnBool)
+    );
+  } catch {
+    return existingSettings;
+  }
+}