Prefill dimension if known remote model found; onboard search connect…

…ors API (#330) * onboard search connectors api; add auto-fetching of dimensions Signed-off-by: Tyler Ohlsen <[email protected]> * Reset to undefined if no values found Signed-off-by: Tyler Ohlsen <[email protected]> * add comment Signed-off-by: Tyler Ohlsen <[email protected]> * minor tuning of select autofill in maps Signed-off-by: Tyler Ohlsen <[email protected]> --------- Signed-off-by: Tyler Ohlsen <[email protected]>
opensearch-project · Aug 30, 2024 · 8461369 · 8461369
1 parent 29eb275
commit 8461369
Show file tree

Hide file tree

Showing 16 changed files with 272 additions and 34 deletions.
diff --git a/common/constants.ts b/common/constants.ts
@@ -22,7 +22,9 @@ export const FLOW_FRAMEWORK_SEARCH_WORKFLOW_STATE_ROUTE = `${FLOW_FRAMEWORK_WORK
  */
 export const ML_API_ROUTE_PREFIX = '/_plugins/_ml';
 export const ML_MODEL_ROUTE_PREFIX = `${ML_API_ROUTE_PREFIX}/models`;
+export const ML_CONNECTOR_ROUTE_PREFIX = `${ML_API_ROUTE_PREFIX}/connectors`;
 export const ML_SEARCH_MODELS_ROUTE = `${ML_MODEL_ROUTE_PREFIX}/_search`;
+export const ML_SEARCH_CONNECTORS_ROUTE = `${ML_CONNECTOR_ROUTE_PREFIX}/_search`;
 
 /**
  * NODE APIs
@@ -51,7 +53,32 @@ export const GET_PRESET_WORKFLOWS_NODE_API_PATH = `${BASE_WORKFLOW_NODE_API_PATH
 
 // ML Plugin node APIs
 export const BASE_MODEL_NODE_API_PATH = `${BASE_NODE_API_PATH}/model`;
+export const BASE_CONNECTOR_NODE_API_PATH = `${BASE_NODE_API_PATH}/connector`;
 export const SEARCH_MODELS_NODE_API_PATH = `${BASE_MODEL_NODE_API_PATH}/search`;
+export const SEARCH_CONNECTORS_NODE_API_PATH = `${BASE_CONNECTOR_NODE_API_PATH}/search`;
+
+/**
+ * Remote model dimensions. Used for attempting to pre-fill dimension size
+ * based on the specified remote model from a remote service, if found
+ */
+
+// Cohere
+export const COHERE_DIMENSIONS = {
+  [`embed-english-v3.0`]: 1024,
+  [`embed-english-light-v3.0`]: 384,
+  [`embed-multilingual-v3.0`]: 1024,
+  [`embed-multilingual-light-v3.0`]: 384,
+  [`embed-english-v2.0`]: 4096,
+  [`embed-english-light-v2.0`]: 1024,
+  [`embed-multilingual-v2.0`]: 768,
+};
+
+// OpenAI
+export const OPENAI_DIMENSIONS = {
+  [`text-embedding-3-small`]: 1536,
+  [`text-embedding-3-large`]: 3072,
+  [`text-embedding-ada-002`]: 1536,
+};
 
 /**
  * Various constants pertaining to Workflow configs

diff --git a/common/interfaces.ts b/common/interfaces.ts
@@ -393,19 +393,35 @@ export type ModelInterface = {
   output: { [key: string]: ModelOutput };
 };
 
+export type ConnectorParameters = {
+  model?: string;
+  dimensions?: number;
+};
+
 export type Model = {
   id: string;
   name: string;
   algorithm: MODEL_ALGORITHM;
   state: MODEL_STATE;
   modelConfig?: ModelConfig;
   interface?: ModelInterface;
+  connectorId?: string;
+};
+
+export type Connector = {
+  id: string;
+  name: string;
+  parameters?: ConnectorParameters;
 };
 
 export type ModelDict = {
   [modelId: string]: Model;
 };
 
+export type ConnectorDict = {
+  [connectorId: string]: Connector;
+};
+
 export type ModelFormValue = {
   id: string;
   algorithm?: MODEL_ALGORITHM;

diff --git a/public/pages/workflow_detail/workflow_inputs/input_fields/map_field.tsx b/public/pages/workflow_detail/workflow_inputs/input_fields/map_field.tsx
@@ -108,6 +108,9 @@ export function MapField(props: MapFieldProps) {
                                   fieldPath={`${props.fieldPath}.${idx}.key`}
                                   options={props.keyOptions as any[]}
                                   placeholder={props.keyPlaceholder || 'Input'}
+                                  autofill={
+                                    props.keyOptions?.length === 1 && idx === 0
+                                  }
                                 />
                               ) : (
                                 <TextField
@@ -134,6 +137,10 @@ export function MapField(props: MapFieldProps) {
                                   placeholder={
                                     props.valuePlaceholder || 'Output'
                                   }
+                                  autofill={
+                                    props.valueOptions?.length === 1 &&
+                                    idx === 0
+                                  }
                                 />
                               ) : (
                                 <TextField

diff --git a/public/pages/workflow_detail/workflow_inputs/input_fields/model_field.tsx b/public/pages/workflow_detail/workflow_inputs/input_fields/model_field.tsx
@@ -43,7 +43,7 @@ export function ModelField(props: ModelFieldProps) {
   // Initial store is fetched when loading base <DetectorDetail /> page. We don't
   // re-fetch here as it could overload client-side if user clicks back and forth /
   // keeps re-rendering this component (and subsequently re-fetching data) as they're building flows
-  const models = useSelector((state: AppState) => state.models.models);
+  const models = useSelector((state: AppState) => state.ml.models);
 
   const { errors, touched } = useFormikContext<WorkspaceFormValues>();
 

diff --git a/public/pages/workflow_detail/workflow_inputs/input_fields/select_with_custom_options.tsx b/public/pages/workflow_detail/workflow_inputs/input_fields/select_with_custom_options.tsx
@@ -13,6 +13,7 @@ interface SelectWithCustomOptionsProps {
   fieldPath: string;
   placeholder: string;
   options: any[];
+  autofill: boolean;
 }
 
 /**
@@ -30,13 +31,14 @@ export function SelectWithCustomOptions(props: SelectWithCustomOptionsProps) {
   // default to the top option. by default, this will re-trigger this hook with a populated
   // value, to then finally update the displayed option.
   useEffect(() => {
-    const formValue = getIn(values, props.fieldPath);
-    if (!isEmpty(formValue)) {
-      setSelectedOption([{ label: getIn(values, props.fieldPath) }]);
-    } else {
-      if (props.options.length > 0) {
-        setFieldTouched(props.fieldPath, true);
-        setFieldValue(props.fieldPath, props.options[0].label);
+    if (props.autofill) {
+      const formValue = getIn(values, props.fieldPath);
+      if (!isEmpty(formValue)) {
+        setSelectedOption([{ label: getIn(values, props.fieldPath) }]);
+      } else {
+        if (props.options.length > 0) {
+          setFieldValue(props.fieldPath, props.options[0].label);
+        }
       }
     }
   }, [getIn(values, props.fieldPath)]);

diff --git a/public/pages/workflow_detail/workflow_inputs/processor_inputs/ml_processor_inputs.tsx b/public/pages/workflow_detail/workflow_inputs/processor_inputs/ml_processor_inputs.tsx
@@ -53,7 +53,7 @@ interface MLProcessorInputsProps {
  * output map configuration forms, respectively.
  */
 export function MLProcessorInputs(props: MLProcessorInputsProps) {
-  const models = useSelector((state: AppState) => state.models.models);
+  const models = useSelector((state: AppState) => state.ml.models);
   const { values, setFieldValue, setFieldTouched } = useFormikContext<
     WorkflowFormValues
   >();

diff --git a/public/pages/workflow_detail/workflow_inputs/workflow_inputs.tsx b/public/pages/workflow_detail/workflow_inputs/workflow_inputs.tsx
@@ -9,7 +9,6 @@ import { debounce, isEmpty, isEqual } from 'lodash';
 import {
   EuiButton,
   EuiButtonEmpty,
-  EuiCallOut,
   EuiFlexGroup,
   EuiFlexItem,
   EuiHorizontalRule,
@@ -100,7 +99,6 @@ export function WorkflowInputs(props: WorkflowInputsProps) {
     setFieldValue,
     values,
     touched,
-    dirty,
   } = useFormikContext<WorkflowFormValues>();
   const dispatch = useAppDispatch();
   const dataSourceId = getDataSourceId();
@@ -650,15 +648,6 @@ export function WorkflowInputs(props: WorkflowInputsProps) {
                 </EuiModalFooter>
               </EuiModal>
             )}
-            {onIngest &&
-              dirty &&
-              hasProvisionedSearchResources(props.workflow) && (
-                <EuiCallOut
-                  title="Making changes to ingest may affect your configured search flow"
-                  iconType={'alert'}
-                  color="warning"
-                />
-              )}
             {onIngestAndUnprovisioned && (
               <>
                 <EuiSpacer size="m" />

diff --git a/public/pages/workflows/new_workflow/new_workflow.tsx b/public/pages/workflows/new_workflow/new_workflow.tsx
@@ -24,6 +24,7 @@ import {
   useAppDispatch,
   getWorkflowPresets,
   searchModels,
+  searchConnectors,
 } from '../../../store';
 import { enrichPresetWorkflowWithUiMetadata } from './utils';
 import { getDataSourceId } from '../../../utils';
@@ -56,11 +57,12 @@ export function NewWorkflow(props: NewWorkflowProps) {
 
   // on initial load:
   // 1. fetch the workflow presets persisted on server-side
-  // 2. fetch the ML models. these may be used in quick-create views when selecting a preset,
+  // 2. fetch the ML models and connectors. these may be used in quick-create views when selecting a preset,
   //    so we optimize by fetching once at the top-level here.
   useEffect(() => {
     dispatch(getWorkflowPresets());
     dispatch(searchModels({ apiBody: FETCH_ALL_QUERY, dataSourceId }));
+    dispatch(searchConnectors({ apiBody: FETCH_ALL_QUERY, dataSourceId }));
   }, []);
 
   // initial hook to populate all workflows

diff --git a/public/pages/workflows/new_workflow/quick_configure_inputs.tsx b/public/pages/workflows/new_workflow/quick_configure_inputs.tsx
@@ -16,8 +16,10 @@ import {
   EuiCompressedFieldNumber,
 } from '@elastic/eui';
 import {
+  COHERE_DIMENSIONS,
   MODEL_STATE,
   Model,
+  OPENAI_DIMENSIONS,
   QuickConfigureFields,
   WORKFLOW_TYPE,
 } from '../../../../common';
@@ -35,7 +37,7 @@ const DEFAULT_IMAGE_FIELD = 'my_image';
 // Dynamic component to allow optional input configuration fields for different use cases.
 // Hooks back to the parent component with such field values
 export function QuickConfigureInputs(props: QuickConfigureInputsProps) {
-  const models = useSelector((state: AppState) => state.models.models);
+  const { models, connectors } = useSelector((state: AppState) => state.ml);
 
   // Deployed models state
   const [deployedModels, setDeployedModels] = useState<Model[]>([]);
@@ -88,6 +90,45 @@ export function QuickConfigureInputs(props: QuickConfigureInputsProps) {
     props.setFields(fieldValues);
   }, [fieldValues]);
 
+  // Try to pre-fill the dimensions based on the chosen model
+  useEffect(() => {
+    const selectedModel = deployedModels.find(
+      (model) => model.id === fieldValues.embeddingModelId
+    );
+    if (selectedModel?.connectorId !== undefined) {
+      const connector = connectors[selectedModel.connectorId];
+      if (connector !== undefined) {
+        // some APIs allow specifically setting the dimensions at runtime,
+        // so we check for that first.
+        if (connector.parameters?.dimensions !== undefined) {
+          setFieldValues({
+            ...fieldValues,
+            embeddingLength: connector.parameters?.dimensions,
+          });
+        } else if (connector.parameters?.model !== undefined) {
+          const dimensions =
+            // @ts-ignore
+            COHERE_DIMENSIONS[connector.parameters?.model] ||
+            // @ts-ignore
+            (OPENAI_DIMENSIONS[connector.parameters?.model] as
+              | number
+              | undefined);
+          if (dimensions !== undefined) {
+            setFieldValues({
+              ...fieldValues,
+              embeddingLength: dimensions,
+            });
+          }
+        } else {
+          setFieldValues({
+            ...fieldValues,
+            embeddingLength: undefined,
+          });
+        }
+      }
+    }
+  }, [fieldValues.embeddingModelId, deployedModels, connectors]);
+
   return (
     <>
       {(props.workflowType === WORKFLOW_TYPE.SEMANTIC_SEARCH ||
@@ -196,7 +237,7 @@ export function QuickConfigureInputs(props: QuickConfigureInputsProps) {
             <EuiCompressedFormRow
               label={'Embedding length'}
               isInvalid={false}
-              helpText="The length / dimension of the generated vector embeddings"
+              helpText="The length / dimension of the generated vector embeddings. Autofilled values may be inaccurate."
             >
               <EuiCompressedFieldNumber
                 value={fieldValues?.embeddingLength || ''}

diff --git a/public/route_service.ts b/public/route_service.ts
@@ -24,6 +24,7 @@ import {
   SimulateIngestPipelineDoc,
   BULK_NODE_API_PATH,
   BASE_NODE_API_PATH,
+  SEARCH_CONNECTORS_NODE_API_PATH,
 } from '../common';
 
 /**
@@ -108,6 +109,11 @@ export interface RouteService {
     body: {},
     dataSourceId?: string
   ) => Promise<any | HttpFetchError>;
+  searchConnectors: (
+    body: {},
+    dataSourceId?: string
+  ) => Promise<any | HttpFetchError>;
+
   simulatePipeline: (
     body: {
       pipeline: IngestPipelineConfig;
@@ -342,6 +348,19 @@ export function configureRoutes(core: CoreStart): RouteService {
         return e as HttpFetchError;
       }
     },
+    searchConnectors: async (body: {}, dataSourceId?: string) => {
+      try {
+        const url = dataSourceId
+          ? `${BASE_NODE_API_PATH}/${dataSourceId}/connector/search`
+          : SEARCH_CONNECTORS_NODE_API_PATH;
+        const response = await core.http.post<{ respString: string }>(url, {
+          body: JSON.stringify(body),
+        });
+        return response;
+      } catch (e: any) {
+        return e as HttpFetchError;
+      }
+    },
     simulatePipeline: async (
       body: {
         pipeline: IngestPipelineConfig;

diff --git a/public/store/reducers/index.ts b/public/store/reducers/index.ts
@@ -6,4 +6,4 @@
 export * from './opensearch_reducer';
 export * from './workflows_reducer';
 export * from './presets_reducer';
-export * from './models_reducer';
+export * from './ml_reducer';