Moved to core

Author: mridula-s109

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/search/vectors/SparseVectorQueryBuilder.java

@@ -25,16 +25,16 @@
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.inference.InferenceResults;
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.inference.TaskType;
 import org.elasticsearch.inference.WeightedToken;
 import org.elasticsearch.xcontent.ConstructingObjectParser;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentParser;
-import org.elasticsearch.xpack.core.ml.action.CoordinatedInferenceAction;
-import org.elasticsearch.xpack.core.ml.inference.TrainedModelPrefixStrings;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
 import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults;
 import org.elasticsearch.xpack.core.ml.inference.results.WarningInferenceResults;
-import org.elasticsearch.xpack.core.ml.inference.trainedmodel.TextExpansionConfigUpdate;
 
 import java.io.IOException;
 import java.util.ArrayList;
@@ -44,7 +44,7 @@
 
 import static org.elasticsearch.xcontent.ConstructingObjectParser.constructorArg;
 import static org.elasticsearch.xcontent.ConstructingObjectParser.optionalConstructorArg;
-import static org.elasticsearch.xpack.core.ClientHelper.ML_ORIGIN;
+import static org.elasticsearch.xpack.core.ClientHelper.INFERENCE_ORIGIN;
 import static org.elasticsearch.xpack.core.ClientHelper.executeAsyncWithOrigin;
 
 public class SparseVectorQueryBuilder extends AbstractQueryBuilder<SparseVectorQueryBuilder> {
@@ -272,27 +272,29 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) {
             throw new IllegalArgumentException("inference_id required to perform vector search on query string");
         }
 
-        // TODO move this to xpack core and use inference APIs
-        CoordinatedInferenceAction.Request inferRequest = CoordinatedInferenceAction.Request.forTextInput(
+        InferenceAction.Request inferRequest = new InferenceAction.Request(
+            TaskType.SPARSE_EMBEDDING,
             inferenceId,
-            List.of(query),
-            TextExpansionConfigUpdate.EMPTY_UPDATE,
-            false,
-            null
+            null, // query field (not needed for sparse embedding)
+            null, // returnDocuments (not needed)
+            null, // topN (not needed)
+            List.of(query), // input text
+            Map.of(), // taskSettings (empty for now)
+            null, // input type not allowed for sparse_embedding task type
+            null, // timeout (use default)
+            false // not streaming
         );
-        inferRequest.setHighPriority(true);
-        inferRequest.setPrefixType(TrainedModelPrefixStrings.PrefixType.SEARCH);
 
         SetOnce<TextExpansionResults> textExpansionResultsSupplier = new SetOnce<>();
         queryRewriteContext.registerAsyncAction(
             (client, listener) -> executeAsyncWithOrigin(
                 client,
-                ML_ORIGIN,
-                CoordinatedInferenceAction.INSTANCE,
+                INFERENCE_ORIGIN,
+                InferenceAction.INSTANCE,
                 inferRequest,
                 ActionListener.wrap(inferenceResponse -> {
 
-                    List<InferenceResults> inferenceResults = inferenceResponse.getInferenceResults();
+                    List<? extends InferenceResults> inferenceResults = inferenceResponse.getResults().transformToCoordinationFormat();
                     if (inferenceResults.isEmpty()) {
                         listener.onFailure(new IllegalStateException("inference response contain no results"));
                         return;

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/search/vectors/SparseVectorQueryBuilderTests.java

@@ -39,10 +39,9 @@
 import org.elasticsearch.test.AbstractQueryTestCase;
 import org.elasticsearch.test.index.IndexVersionUtils;
 import org.elasticsearch.xpack.core.XPackClientPlugin;
-import org.elasticsearch.xpack.core.ml.action.CoordinatedInferenceAction;
-import org.elasticsearch.xpack.core.ml.action.InferModelAction;
-import org.elasticsearch.xpack.core.ml.inference.TrainedModelPrefixStrings;
-import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults;
+import org.elasticsearch.inference.TaskType;
+import org.elasticsearch.xpack.core.inference.action.InferenceAction;
+import org.elasticsearch.xpack.core.inference.results.SparseEmbeddingResults;
 
 import java.io.IOException;
 import java.lang.reflect.Method;
@@ -118,15 +117,14 @@ protected Settings createTestIndexSettings() {
     @Override
     protected boolean canSimulateMethod(Method method, Object[] args) throws NoSuchMethodException {
         return method.equals(Client.class.getMethod("execute", ActionType.class, ActionRequest.class, ActionListener.class))
-            && (args[0] instanceof CoordinatedInferenceAction);
+            && (args[0] instanceof InferenceAction);
     }
 
     @Override
     protected Object simulateMethod(Method method, Object[] args) {
-        CoordinatedInferenceAction.Request request = (CoordinatedInferenceAction.Request) args[1];
-        assertNull(request.getInferenceTimeout());
-        assertEquals(TrainedModelPrefixStrings.PrefixType.SEARCH, request.getPrefixType());
-        assertEquals(CoordinatedInferenceAction.Request.RequestModelType.NLP_MODEL, request.getRequestModelType());
+        InferenceAction.Request request = (InferenceAction.Request) args[1];
+        assertEquals(TaskType.SPARSE_EMBEDDING, request.getTaskType());
+        assertNull(request.getInputType()); // Should be null for sparse_embedding
 
         // Randomisation cannot be used here as {@code #doAssertLuceneQuery}
         // asserts that 2 rewritten queries are the same
@@ -135,12 +133,13 @@ protected Object simulateMethod(Method method, Object[] args) {
             tokens.add(new WeightedToken(Integer.toString(i), (i + 1) * 1.0f));
         }
 
-        var response = InferModelAction.Response.builder()
-            .setId(request.getModelId())
-            .addInferenceResults(List.of(new TextExpansionResults("foo", tokens, randomBoolean())))
-            .build();
+        var embeddings = List.of(
+            new SparseEmbeddingResults.Embedding(tokens, randomBoolean())
+        );
+        var results = new SparseEmbeddingResults(embeddings);
+        var response = new InferenceAction.Response(results);
         @SuppressWarnings("unchecked")  // We matched the method above.
-        ActionListener<InferModelAction.Response> listener = (ActionListener<InferModelAction.Response>) args[2];
+        ActionListener<InferenceAction.Response> listener = (ActionListener<InferenceAction.Response>) args[2];
         listener.onResponse(response);
         return null;
     }