Skip to content

Commit ca6bbe7

Browse files
[Backport 2.16] Add initial search request inference processor (#2731)
1 parent 32a83e3 commit ca6bbe7

File tree

10 files changed

+2315
-5
lines changed

10 files changed

+2315
-5
lines changed

plugin/src/main/java/org/opensearch/ml/plugin/MachineLearningPlugin.java

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -213,6 +213,7 @@
213213
import org.opensearch.ml.model.MLModelCacheHelper;
214214
import org.opensearch.ml.model.MLModelManager;
215215
import org.opensearch.ml.processor.MLInferenceIngestProcessor;
216+
import org.opensearch.ml.processor.MLInferenceSearchRequestProcessor;
216217
import org.opensearch.ml.repackage.com.google.common.collect.ImmutableList;
217218
import org.opensearch.ml.rest.RestMLCreateConnectorAction;
218219
import org.opensearch.ml.rest.RestMLCreateControllerAction;
@@ -977,7 +978,11 @@ public Map<String, Processor.Factory<SearchRequestProcessor>> getRequestProcesso
977978
GenerativeQAProcessorConstants.REQUEST_PROCESSOR_TYPE,
978979
new GenerativeQARequestProcessor.Factory(() -> this.ragSearchPipelineEnabled)
979980
);
980-
981+
requestProcessors
982+
.put(
983+
MLInferenceSearchRequestProcessor.TYPE,
984+
new MLInferenceSearchRequestProcessor.Factory(parameters.client, parameters.namedXContentRegistry)
985+
);
981986
return requestProcessors;
982987
}
983988

plugin/src/main/java/org/opensearch/ml/processor/MLInferenceSearchRequestProcessor.java

Lines changed: 578 additions & 0 deletions
Large diffs are not rendered by default.

plugin/src/main/java/org/opensearch/ml/processor/ModelExecutor.java

Lines changed: 7 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -189,7 +189,13 @@ default Object getModelOutputValue(MLOutput mlOutput, String modelOutputFieldNam
189189
return modelTensorOutputMap;
190190
} else {
191191
try {
192-
return JsonPath.parse(modelTensorOutputMap).read(modelOutputFieldName);
192+
Object modelOutputValue = JsonPath.parse(modelTensorOutputMap).read(modelOutputFieldName);
193+
if (modelOutputValue == null) {
194+
throw new IllegalArgumentException(
195+
"model inference output cannot find such json path: " + modelOutputFieldName + " in " + modelTensorOutputMap
196+
);
197+
}
198+
return modelOutputValue;
193199
} catch (Exception e) {
194200
if (ignoreMissing) {
195201
return modelTensorOutputMap;
@@ -313,7 +319,6 @@ default List<String> writeNewDotPathForNestedObject(Object json, String dotPath)
313319
* @return the converted dot path notation string
314320
*/
315321
default String convertToDotPath(String path) {
316-
317322
return path.replaceAll("\\[(\\d+)\\]", "$1\\.").replaceAll("\\['(.*?)']", "$1\\.").replaceAll("^\\$", "").replaceAll("\\.$", "");
318323
}
319324

plugin/src/test/java/org/opensearch/ml/plugin/MachineLearningPluginTests.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@
3838
import org.opensearch.ml.common.spi.MLCommonsExtension;
3939
import org.opensearch.ml.common.spi.tools.Tool;
4040
import org.opensearch.ml.engine.tools.MLModelTool;
41+
import org.opensearch.ml.processor.MLInferenceSearchRequestProcessor;
4142
import org.opensearch.plugins.ExtensiblePlugin;
4243
import org.opensearch.plugins.SearchPipelinePlugin;
4344
import org.opensearch.plugins.SearchPlugin;
@@ -73,10 +74,11 @@ public void testGetSearchExts() {
7374
public void testGetRequestProcessors() {
7475
SearchPipelinePlugin.Parameters parameters = mock(SearchPipelinePlugin.Parameters.class);
7576
Map<String, ?> requestProcessors = plugin.getRequestProcessors(parameters);
76-
assertEquals(1, requestProcessors.size());
77+
assertEquals(2, requestProcessors.size());
7778
assertTrue(
7879
requestProcessors.get(GenerativeQAProcessorConstants.REQUEST_PROCESSOR_TYPE) instanceof GenerativeQARequestProcessor.Factory
7980
);
81+
assertTrue(requestProcessors.get(MLInferenceSearchRequestProcessor.TYPE) instanceof MLInferenceSearchRequestProcessor.Factory);
8082
}
8183

8284
@Test

0 commit comments

Comments
 (0)