-
Notifications
You must be signed in to change notification settings - Fork 72
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[Feature] Support for Default Model Id #337
Changes from 8 commits
dac5012
6903f94
88ba2e9
c06d07a
7c3a302
dee67ca
b13eb80
badced1
9c010e7
d0b70c7
c9ada83
4be198c
e9db552
e9efd72
ff10aba
c3d0649
ef5f939
29017ba
ce28954
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
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. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 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. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 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. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. 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); | ||
} | ||
} | ||
} |
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 Codecov / codecov/patchsrc/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java#L32-L33
|
||
} 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 Codecov / codecov/patchsrc/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java#L37
|
||
"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 Codecov / codecov/patchsrc/main/java/org/opensearch/neuralsearch/query/visitor/NeuralSearchQueryVisitor.java#L46
|
||
} | ||
} |
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; | ||
} | ||
} | ||
|
||
} |
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