[Feature] Support for Default Model Id #337
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,81 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.neuralsearch.processor; | ||
|
|
||
| import java.util.Map; | ||
|
|
||
| import org.opensearch.action.search.SearchRequest; | ||
| import org.opensearch.index.query.QueryBuilder; | ||
| import org.opensearch.ingest.ConfigurationUtils; | ||
| import org.opensearch.neuralsearch.query.visitor.NeuralSearchQueryVisitor; | ||
| import org.opensearch.search.pipeline.AbstractProcessor; | ||
| import org.opensearch.search.pipeline.Processor; | ||
| import org.opensearch.search.pipeline.SearchRequestProcessor; | ||
|
|
||
| public class NeuralQueryProcessor extends AbstractProcessor implements SearchRequestProcessor { | ||
|
There was a problem hiding this comment. I would rename class to something closer to a processor type name, although I left the final call to you |
||
|
|
||
| /** | ||
| * Key to reference this processor type from a search pipeline. | ||
| */ | ||
| public static final String TYPE = "neural_query"; | ||
navneet1v marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| final String modelId; | ||
vibrantvarun marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| final Map<String, Object> neuralFieldDefaultIdMap; | ||
|
There was a problem hiding this comment. private? Also, could we add a comment for this member? Also, why is this map <String, Object> and not <String, String> There was a problem hiding this comment. Because ConfigurationUtils The class from Opensearch which is used to retrieve map from processor is generic and it Map<String, Object> defined in it. ConfigurationUtils is the designed to get values from the processor in OS. So I am using the same class |
||
|
|
||
| /** | ||
| * Returns the type of the processor. | ||
| * | ||
| * @return The processor type. | ||
| */ | ||
| @Override | ||
| public String getType() { | ||
| return TYPE; | ||
| } | ||
|
|
||
| protected NeuralQueryProcessor( | ||
vibrantvarun marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| String tag, | ||
| String description, | ||
| boolean ignoreFailure, | ||
| String modelId, | ||
| Map<String, Object> neuralFieldDefaultIdMap | ||
vibrantvarun marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| ) { | ||
| super(tag, description, ignoreFailure); | ||
| this.modelId = modelId; | ||
| this.neuralFieldDefaultIdMap = neuralFieldDefaultIdMap; | ||
| } | ||
|
|
||
| @Override | ||
| public SearchRequest processRequest(SearchRequest searchRequest) { | ||
| QueryBuilder queryBuilder = searchRequest.source().query(); | ||
| queryBuilder.visit(new NeuralSearchQueryVisitor(modelId, neuralFieldDefaultIdMap)); | ||
| return searchRequest; | ||
| } | ||
|
|
||
| public static class Factory implements Processor.Factory<SearchRequestProcessor> { | ||
vibrantvarun marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| private static final String DEFAULT_MODEL_ID = "default_model_id"; | ||
| private static final String NEURAL_FIELD_DEFAULT_ID = "neural_field_default_id"; | ||
|
|
||
| @Override | ||
| public NeuralQueryProcessor create( | ||
| Map<String, Processor.Factory<SearchRequestProcessor>> processorFactories, | ||
| String tag, | ||
| String description, | ||
| boolean ignoreFailure, | ||
| Map<String, Object> config, | ||
| PipelineContext pipelineContext | ||
| ) throws Exception { | ||
vibrantvarun marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| String modelId = (String) config.remove(DEFAULT_MODEL_ID); | ||
vibrantvarun marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| Map<String, Object> neuralInfoMap = ConfigurationUtils.readOptionalMap(TYPE, tag, config, NEURAL_FIELD_DEFAULT_ID); | ||
|
|
||
| if (modelId == null && neuralInfoMap == null) { | ||
| throw new IllegalArgumentException("model Id or neural info map either of them should be provided"); | ||
| } | ||
|
|
||
| return new NeuralQueryProcessor(tag, description, ignoreFailure, modelId, neuralInfoMap); | ||
| } | ||
| } | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,48 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.neuralsearch.query.visitor; | ||
|
|
||
| import java.util.Map; | ||
|
|
||
| import org.apache.lucene.search.BooleanClause; | ||
| import org.opensearch.index.query.QueryBuilder; | ||
| import org.opensearch.index.query.QueryBuilderVisitor; | ||
| import org.opensearch.neuralsearch.query.NeuralQueryBuilder; | ||
|
|
||
| public class NeuralSearchQueryVisitor implements QueryBuilderVisitor { | ||
|
|
||
| private String modelId; | ||
| private Map<String, Object> neuralFieldMap; | ||
vibrantvarun marked this conversation as resolved.
Show resolved
Hide resolved
|
||
|
|
||
| public NeuralSearchQueryVisitor(String modelId, Map<String, Object> neuralFieldMap) { | ||
vibrantvarun marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| this.modelId = modelId; | ||
| this.neuralFieldMap = neuralFieldMap; | ||
| } | ||
|
|
||
| @Override | ||
| public void accept(QueryBuilder queryBuilder) { | ||
vibrantvarun marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| if (queryBuilder instanceof NeuralQueryBuilder) { | ||
| NeuralQueryBuilder neuralQueryBuilder = (NeuralQueryBuilder) queryBuilder; | ||
| if (neuralFieldMap != null | ||
| && neuralQueryBuilder.fieldName() != null | ||
| && neuralFieldMap.get(neuralQueryBuilder.fieldName()) != null) { | ||
| String fieldDefaultModelId = (String) neuralFieldMap.get(neuralQueryBuilder.fieldName()); | ||
| neuralQueryBuilder.modelId(fieldDefaultModelId); | ||
|
Check warning on line 33 in src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java
|
||
| } else if (modelId != null) { | ||
| neuralQueryBuilder.modelId(modelId); | ||
| } else { | ||
| throw new IllegalArgumentException( | ||
|
Check warning on line 37 in src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java
|
||
| "model id must be provided in neural query or a default model id must be set in search request processor" | ||
| ); | ||
| } | ||
| } | ||
| } | ||
|
|
||
| @Override | ||
| public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { | ||
| return this; | ||
|
Check warning on line 46 in src/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java
|
||
| } | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,63 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.neuralsearch.util; | ||
|
|
||
| import java.util.Locale; | ||
|
|
||
| import lombok.AccessLevel; | ||
| import lombok.NoArgsConstructor; | ||
| import lombok.extern.log4j.Log4j2; | ||
|
|
||
| import org.opensearch.Version; | ||
| import org.opensearch.cluster.service.ClusterService; | ||
|
|
||
| @NoArgsConstructor(access = AccessLevel.PRIVATE) | ||
| @Log4j2 | ||
| public class NeuralSearchClusterUtil { | ||
vibrantvarun marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| private ClusterService clusterService; | ||
|
|
||
| private static NeuralSearchClusterUtil instance; | ||
|
|
||
| /** | ||
| * Return instance of the cluster context, must be initialized first for proper usage | ||
| * @return instance of cluster context | ||
| */ | ||
| public static synchronized NeuralSearchClusterUtil instance() { | ||
| if (instance == null) { | ||
| instance = new NeuralSearchClusterUtil(); | ||
| } | ||
| return instance; | ||
| } | ||
|
|
||
| /** | ||
| * Initializes instance of cluster context by injecting dependencies | ||
| * @param clusterService | ||
| */ | ||
| public void initialize(final ClusterService clusterService) { | ||
| this.clusterService = clusterService; | ||
| } | ||
|
|
||
| /** | ||
| * Return minimal OpenSearch version based on all nodes currently discoverable in the cluster | ||
| * @return minimal installed OpenSearch version, default to Version.CURRENT which is typically the latest version | ||
| */ | ||
| public Version getClusterMinVersion() { | ||
| try { | ||
| return this.clusterService.state().getNodes().getMinNodeVersion(); | ||
| } catch (Exception exception) { | ||
| log.error( | ||
vibrantvarun marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| String.format( | ||
| Locale.ROOT, | ||
| "Failed to get cluster minimum node version, returning current node version %s instead.", | ||
| Version.CURRENT | ||
| ), | ||
| exception | ||
| ); | ||
| return Version.CURRENT; | ||
| } | ||
| } | ||
|
|
||
| } | ||
Oops, something went wrong.
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
We should start getting rid of these kind of inits and move towards singleton pattern without this kind of inits.
May be an AI for maintainers