[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
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,104 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.neuralsearch.processor; | ||
|
|
||
| import static org.opensearch.ingest.ConfigurationUtils.*; | ||
| import static org.opensearch.neuralsearch.processor.TextEmbeddingProcessor.TYPE; | ||
|
|
||
| import java.util.Map; | ||
|
|
||
| import lombok.Getter; | ||
| import lombok.Setter; | ||
|
|
||
| import org.opensearch.action.search.SearchRequest; | ||
| import org.opensearch.common.Nullable; | ||
| 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; | ||
|
|
||
| /** | ||
| * Neural Search Query Request Processor, It modifies the search request with neural query clause | ||
| * and adds model Id if not present in the search query. | ||
| */ | ||
| @Setter | ||
| @Getter | ||
| public class NeuralQueryEnricherProcessor extends AbstractProcessor implements SearchRequestProcessor { | ||
|
|
||
| /** | ||
| * Key to reference this processor type from a search pipeline. | ||
| */ | ||
| public static final String TYPE = "neural_query_enricher"; | ||
|
|
||
| private final String modelId; | ||
|
|
||
| private final Map<String, Object> neuralFieldDefaultIdMap; | ||
|
|
||
| /** | ||
| * Returns the type of the processor. | ||
| * | ||
| * @return The processor type. | ||
| */ | ||
| @Override | ||
| public String getType() { | ||
| return TYPE; | ||
| } | ||
|
|
||
| private NeuralQueryEnricherProcessor( | ||
| String tag, | ||
| String description, | ||
| boolean ignoreFailure, | ||
| @Nullable String modelId, | ||
| @Nullable Map<String, Object> neuralFieldDefaultIdMap | ||
| ) { | ||
| super(tag, description, ignoreFailure); | ||
| this.modelId = modelId; | ||
| this.neuralFieldDefaultIdMap = neuralFieldDefaultIdMap; | ||
| } | ||
|
|
||
| /** | ||
| * Processes the Search Request. | ||
| * | ||
| * @return The Search Request. | ||
| */ | ||
| @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> { | ||
| private static final String DEFAULT_MODEL_ID = "default_model_id"; | ||
| private static final String NEURAL_FIELD_DEFAULT_ID = "neural_field_default_id"; | ||
|
|
||
| /** | ||
| * Create the processor object. | ||
| * | ||
| * @return NeuralQueryEnricherProcessor | ||
| */ | ||
| @Override | ||
| public NeuralQueryEnricherProcessor create( | ||
| Map<String, Processor.Factory<SearchRequestProcessor>> processorFactories, | ||
| String tag, | ||
| String description, | ||
| boolean ignoreFailure, | ||
| Map<String, Object> config, | ||
| PipelineContext pipelineContext | ||
| ) throws IllegalArgumentException { | ||
| String modelId = readOptionalStringProperty(TYPE, tag, config, DEFAULT_MODEL_ID); | ||
| 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 NeuralQueryEnricherProcessor(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,60 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.neuralsearch.query.visitor; | ||
|
|
||
| import java.util.Map; | ||
|
|
||
| import lombok.AllArgsConstructor; | ||
|
|
||
| import org.apache.lucene.search.BooleanClause; | ||
| import org.opensearch.index.query.QueryBuilder; | ||
| import org.opensearch.index.query.QueryBuilderVisitor; | ||
| import org.opensearch.neuralsearch.query.NeuralQueryBuilder; | ||
|
|
||
| /** | ||
| * Neural Search Query Visitor. It visits each and every component of query buikder tree. | ||
| */ | ||
| @AllArgsConstructor | ||
| public class NeuralSearchQueryVisitor implements QueryBuilderVisitor { | ||
|
|
||
| private final String modelId; | ||
| private final Map<String, Object> neuralFieldMap; | ||
|
|
||
| /** | ||
| * Accept method accepts every query builder from the search request, | ||
| * and processes it if the required conditions in accept method are satisfied. | ||
| */ | ||
| @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 (neuralQueryBuilder.modelId() == null) { | ||
| if (neuralFieldMap != null | ||
| && neuralQueryBuilder.fieldName() != null | ||
| && neuralFieldMap.get(neuralQueryBuilder.fieldName()) != null) { | ||
| String fieldDefaultModelId = (String) neuralFieldMap.get(neuralQueryBuilder.fieldName()); | ||
| neuralQueryBuilder.modelId(fieldDefaultModelId); | ||
| } else if (modelId != null) { | ||
| neuralQueryBuilder.modelId(modelId); | ||
| } else { | ||
| throw new IllegalArgumentException( | ||
| "model id must be provided in neural query or a default model id must be set in search request processor" | ||
| ); | ||
| } | ||
| } | ||
| } | ||
| } | ||
|
|
||
| /** | ||
| * Retrieves the child visitor from the Visitor object. | ||
| * | ||
| * @return The sub Query Visitor | ||
| */ | ||
| @Override | ||
| public QueryBuilderVisitor getChildVisitor(BooleanClause.Occur occur) { | ||
| return this; | ||
| } | ||
| } | ||
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,52 @@ | ||
| /* | ||
| * Copyright OpenSearch Contributors | ||
| * SPDX-License-Identifier: Apache-2.0 | ||
| */ | ||
|
|
||
| package org.opensearch.neuralsearch.util; | ||
|
|
||
| import lombok.AccessLevel; | ||
| import lombok.NoArgsConstructor; | ||
| import lombok.extern.log4j.Log4j2; | ||
|
|
||
| import org.opensearch.Version; | ||
| import org.opensearch.cluster.service.ClusterService; | ||
|
|
||
| /** | ||
| * Class abstracts information related to underlying OpenSearch cluster | ||
| */ | ||
| @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() { | ||
| return this.clusterService.state().getNodes().getMinNodeVersion(); | ||
| } | ||
|
|
||
| } | ||
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