Onboard ML inference ingest processor (#172)

Signed-off-by: Tyler Ohlsen <[email protected]>
(cherry picked from commit f077c82)

common/constants.ts

@@ -124,7 +124,7 @@ export const NEURAL_SPARSE_TOKENIZER_TRANSFORMER = {
  */
 
 export enum PROCESSOR_TYPE {
-  MODEL = 'model_processor',
+  ML = 'ml_processor',
 }
 
 export enum MODEL_TYPE {

common/interfaces.ts

@@ -6,7 +6,7 @@
 import { Node, Edge } from 'reactflow';
 import { FormikValues } from 'formik';
 import { ObjectSchema } from 'yup';
-import { COMPONENT_CLASS, PROCESSOR_TYPE, MODEL_TYPE } from './constants';
+import { COMPONENT_CLASS, PROCESSOR_TYPE } from './constants';
 
 export type Index = {
   name: string;
@@ -18,7 +18,7 @@ export type Index = {
 TODO: over time these can become less generic as the form inputs & UX becomes finalized
  */
 
-export type ConfigFieldType = 'string' | 'json' | 'select' | 'model';
+export type ConfigFieldType = 'string' | 'json' | 'select' | 'model' | 'map';
 export type ConfigSelectType = 'model';
 export type ConfigFieldValue = string | {};
 export interface IConfigField {
@@ -46,10 +46,6 @@ export interface IProcessorConfig extends IConfig {
   type: PROCESSOR_TYPE;
 }
 
-export interface IModelProcessorConfig extends IProcessorConfig {
-  modelType: MODEL_TYPE;
-}
-
 export type EnrichConfig = {
   processors: IProcessorConfig[];
 };
@@ -75,6 +71,13 @@ export type WorkflowConfig = {
   search: SearchConfig;
 };
 
+export type MapEntry = {
+  key: string;
+  value: string;
+};
+
+export type MapFormValue = MapEntry[];
+
 export type WorkflowFormValues = {
   ingest: FormikValues;
   search: FormikValues;
@@ -181,6 +184,14 @@ export type SearchRequestProcessor = SearchProcessor & {};
 export type SearchResponseProcessor = SearchProcessor & {};
 export type SearchPhaseResultsProcessor = SearchProcessor & {};
 
+export type MLInferenceProcessor = IngestProcessor & {
+  ml_inference: {
+    model_id: string;
+    input_map?: {};
+    output_map?: {};
+  };
+};
+
 export type TextEmbeddingProcessor = IngestProcessor & {
   text_embedding: {
     model_id: string;

public/component_types/transformer/ml_transformer.ts

@@ -7,13 +7,15 @@ import { COMPONENT_CLASS } from '../../../common';
 import { BaseTransformer } from './base_transformer';
 
 /**
- * A generic ML transformer UI component
+ * A generic ML inference transformer. Can be used across ingest, search request, and search response.
+ * Under the hood, using the implemented ML inference processors.
+ * Ref (ingest): https://opensearch.org/docs/latest/ingest-pipelines/processors/ml-inference/
  */
 export class MLTransformer extends BaseTransformer {
   constructor() {
     super();
     this.type = COMPONENT_CLASS.ML_TRANSFORMER;
-    this.label = 'ML Transformer';
-    this.description = 'A general ML transformer';
+    this.label = 'ML Processor';
+    this.description = 'A general ML processor';
   }
 }

public/configs/base_config.ts

@@ -10,11 +10,13 @@ import { IConfig, IConfigField } from '../../common';
  */
 export abstract class BaseConfig implements IConfig {
   id: string;
+  name: string;
   fields: IConfigField[];
 
   // No-op constructor. If there are general / defaults for field values, add in here.
   constructor() {
     this.id = '';
+    this.name = '';
     this.fields = [];
   }
 

public/configs/ingest_processors/index.ts

@@ -4,3 +4,4 @@
  */
 
 export * from './text_embedding_processor';
+export * from './ml_ingest_processor';

public/configs/ingest_processors/ml_ingest_processor.ts

@@ -0,0 +1,18 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { generateId } from '../../utils';
+import { MLProcessor } from '../ml_processor';
+
+/**
+ * The ML processor in the context of ingest
+ */
+export class MLIngestProcessor extends MLProcessor {
+  constructor() {
+    super();
+    this.id = generateId('ml_ingest_processor');
+    this.name = 'ML ingest processor';
+  }
+}

public/configs/ingest_processors/text_embedding_processor.ts

@@ -3,19 +3,17 @@
  * SPDX-License-Identifier: Apache-2.0
  */
 
-import { MODEL_TYPE } from '../../../common';
 import { generateId } from '../../utils';
-import { ModelProcessor } from './model_processor';
+import { BaseConfig } from '../base_config';
 
 /**
  * A specialized text embedding processor config
  */
-export class TextEmbeddingProcessor extends ModelProcessor {
+export class TextEmbeddingProcessor extends BaseConfig {
   constructor() {
     super();
     this.id = generateId('text_embedding_processor');
     this.name = 'Text embedding processor';
-    this.type = MODEL_TYPE.TEXT_EMBEDDING;
     this.fields = [
       {
         label: 'Text Embedding Model',

public/configs/ml_processor.ts

@@ -0,0 +1,42 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+import { BaseConfig } from './base_config';
+
+/**
+ * A generic ML processor config. Used in ingest and search flows.
+ * The interfaces are identical across ingest / search request / search response processors.
+ */
+export abstract class MLProcessor extends BaseConfig {
+  constructor() {
+    super();
+    this.fields = [
+      {
+        label: 'Model',
+        id: 'model',
+        type: 'model',
+        helpText: 'The model ID.',
+        helpLink:
+          'https://opensearch.org/docs/latest/ml-commons-plugin/integrating-ml-models/#choosing-a-model',
+      },
+      {
+        label: 'Input Map',
+        id: 'inputMap',
+        type: 'map',
+        helpText: `An array specifying how to map fields from the ingested document to the model’s input.`,
+        helpLink:
+          'https://opensearch.org/docs/latest/ingest-pipelines/processors/ml-inference/#configuration-parameters',
+      },
+      {
+        label: 'Output Map',
+        id: 'outputMap',
+        type: 'map',
+        helpText: `An array specifying how to map the model’s output to new fields.`,
+        helpLink:
+          'https://opensearch.org/docs/latest/ingest-pipelines/processors/ml-inference/#configuration-parameters',
+      },
+    ];
+  }
+}