[ML] Switch TextExpansionQueryBuilder and TextEmbeddingQueryVectorBuilder to return 400 instead of 500 errors

docs/changelog/135800.yaml

@@ -0,0 +1,6 @@
+pr: 135800
+summary: Switch `TextExpansionQueryBuilder` and `TextEmbeddingQueryVectorBuilder`
+  to return 400 instead of 500 errors
+area: Machine Learning
+type: bug
+issues: []

test/framework/src/main/java/org/elasticsearch/test/AbstractQueryVectorBuilderTestCase.java

@@ -172,7 +172,7 @@ public final void testVectorFetch() throws Exception {
      */
     protected abstract ActionResponse createResponse(float[] array, T builder);
 
-    protected static float[] randomVector(int dim) {
+    public static float[] randomVector(int dim) {
         float[] vector = new float[dim];
         for (int i = 0; i < vector.length; i++) {
             vector[i] = randomFloat();

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/search/TextExpansionQueryBuilder.java

@@ -9,6 +9,7 @@
 
 import org.apache.lucene.search.Query;
 import org.apache.lucene.util.SetOnce;
+import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.action.ActionListener;
@@ -24,6 +25,7 @@
 import org.elasticsearch.index.query.QueryBuilders;
 import org.elasticsearch.index.query.QueryRewriteContext;
 import org.elasticsearch.index.query.SearchExecutionContext;
+import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.xcontent.ParseField;
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentParser;
@@ -188,14 +190,15 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) {
                         listener.onFailure(new IllegalStateException(warning.getWarning()));
                     } else {
                         listener.onFailure(
-                            new IllegalStateException(
+                            new ElasticsearchStatusException(
                                 "expected a result of type ["
                                     + TextExpansionResults.NAME
                                     + "] received ["
                                     + inferenceResponse.getInferenceResults().get(0).getWriteableName()
                                     + "]. Is ["
                                     + modelId
-                                    + "] a compatible model?"
+                                    + "] a compatible model?",
+                                RestStatus.BAD_REQUEST
                             )
                         );
                     }

x-pack/plugin/core/src/main/java/org/elasticsearch/xpack/core/ml/vectors/TextEmbeddingQueryVectorBuilder.java

@@ -7,12 +7,14 @@
 
 package org.elasticsearch.xpack.core.ml.vectors;
 
+import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.TransportVersion;
 import org.elasticsearch.TransportVersions;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.search.vectors.QueryVectorBuilder;
 import org.elasticsearch.xcontent.ConstructingObjectParser;
 import org.elasticsearch.xcontent.ParseField;
@@ -132,14 +134,15 @@ public void buildVector(Client client, ActionListener<float[]> listener) {
             } else if (response.getInferenceResults().get(0) instanceof WarningInferenceResults warning) {
                 listener.onFailure(new IllegalStateException(warning.getWarning()));
             } else {
-                throw new IllegalStateException(
+                throw new ElasticsearchStatusException(
                     "expected a result of type ["
                         + MlTextEmbeddingResults.NAME
                         + "] received ["
                         + response.getInferenceResults().get(0).getWriteableName()
                         + "]. Is ["
                         + modelId
-                        + "] a text embedding model?"
+                        + "] a text embedding model?",
+                    RestStatus.BAD_REQUEST
                 );
             }
         }, listener::onFailure));

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/ml/search/TextExpansionQueryBuilderTests.java

@@ -16,44 +16,64 @@
 import org.apache.lucene.search.Query;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.tests.index.RandomIndexWriter;
+import org.elasticsearch.ElasticsearchStatusException;
 import org.elasticsearch.action.ActionListener;
 import org.elasticsearch.action.ActionRequest;
 import org.elasticsearch.action.ActionType;
 import org.elasticsearch.action.admin.indices.mapping.put.PutMappingRequest;
+import org.elasticsearch.action.support.PlainActionFuture;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.compress.CompressedXContent;
+import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.mapper.MapperService;
 import org.elasticsearch.index.mapper.extras.MapperExtrasPlugin;
 import org.elasticsearch.index.mapper.vectors.TokenPruningConfig;
 import org.elasticsearch.index.query.BoolQueryBuilder;
 import org.elasticsearch.index.query.QueryBuilder;
+import org.elasticsearch.index.query.Rewriteable;
 import org.elasticsearch.index.query.SearchExecutionContext;
 import org.elasticsearch.inference.WeightedToken;
 import org.elasticsearch.plugins.Plugin;
+import org.elasticsearch.rest.RestStatus;
 import org.elasticsearch.search.vectors.SparseVectorQueryWrapper;
 import org.elasticsearch.test.AbstractQueryTestCase;
 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.MlTextEmbeddingResults;
 import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults;
+import org.junit.Before;
 
 import java.io.IOException;
 import java.lang.reflect.Method;
 import java.util.ArrayList;
 import java.util.Collection;
 import java.util.List;
+import java.util.concurrent.atomic.AtomicBoolean;
 
+import static org.elasticsearch.test.AbstractQueryVectorBuilderTestCase.randomVector;
 import static org.hamcrest.CoreMatchers.instanceOf;
+import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.either;
 import static org.hamcrest.Matchers.hasSize;
+import static org.hamcrest.Matchers.is;
 
 public class TextExpansionQueryBuilderTests extends AbstractQueryTestCase<TextExpansionQueryBuilder> {
 
     private static final String RANK_FEATURES_FIELD = "rank";
     private static final int NUM_TOKENS = 10;
 
+    // This is a hack so that we can control when the client returns valid and invalid results
+    // to test both the success and failure paths
+    private final AtomicBoolean shouldProduceDenseVectorResults = new AtomicBoolean();
+
+    @Before
+    public void setUpFailure() {
+        shouldProduceDenseVectorResults.set(false);
+    }
+
     @Override
     protected TextExpansionQueryBuilder doCreateTestQueryBuilder() {
         TokenPruningConfig tokenPruningConfig = randomBoolean()
@@ -100,21 +120,40 @@ protected Object simulateMethod(Method method, Object[] args) {
         assertEquals(TrainedModelPrefixStrings.PrefixType.SEARCH, request.getPrefixType());
         assertEquals(CoordinatedInferenceAction.Request.RequestModelType.NLP_MODEL, request.getRequestModelType());
 
+        InferModelAction.Response response;
+        if (shouldProduceDenseVectorResults.get()) {
+            response = createDenseVectorResults(request.getModelId());
+        } else {
+            response = createTextExpansionResults(request.getModelId());
+        }
+
+        @SuppressWarnings("unchecked")  // We matched the method above.
+        ActionListener<InferModelAction.Response> listener = (ActionListener<InferModelAction.Response>) args[2];
+        listener.onResponse(response);
+        return null;
+    }
+
+    private static InferModelAction.Response createDenseVectorResults(String modelId) {
+        float[] resultFloats = randomVector(randomIntBetween(10, 1024));
+        double[] resultDoubles = new double[resultFloats.length];
+        for (int i = 0; i < resultDoubles.length; i++) {
+            resultDoubles[i] = resultFloats[i];
+        }
+        return new InferModelAction.Response(List.of(new MlTextEmbeddingResults("foo", resultDoubles, randomBoolean())), modelId, true);
+    }
+
+    private static InferModelAction.Response createTextExpansionResults(String modelId) {
         // Randomisation cannot be used here as {@code #doAssertLuceneQuery}
         // asserts that 2 rewritten queries are the same
         var tokens = new ArrayList<WeightedToken>();
         for (int i = 0; i < NUM_TOKENS; i++) {
             tokens.add(new WeightedToken(Integer.toString(i), (i + 1) * 1.0f));
         }
 
-        var response = InferModelAction.Response.builder()
-            .setId(request.getModelId())
+        return InferModelAction.Response.builder()
+            .setId(modelId)
             .addInferenceResults(List.of(new TextExpansionResults("foo", tokens, randomBoolean())))
             .build();
-        @SuppressWarnings("unchecked")  // We matched the method above.
-        ActionListener<InferModelAction.Response> listener = (ActionListener<InferModelAction.Response>) args[2];
-        listener.onResponse(response);
-        return null;
     }
 
     @Override
@@ -294,4 +333,32 @@ public void testThatTokensAreCorrectlyPruned() {
             assertTrue(rewrittenQueryBuilder instanceof WeightedTokensQueryBuilder);
         }
     }
+
+    public void testFailure() {
+        // Tells the client to return invalid results to trigger a failure
+        shouldProduceDenseVectorResults.set(true);
+
+        SearchExecutionContext searchExecutionContext = createSearchExecutionContext();
+        TextExpansionQueryBuilder queryBuilder = createTestQueryBuilderNoPruning();
+
+        PlainActionFuture<QueryBuilder> future = new PlainActionFuture<>();
+        Rewriteable.rewriteAndFetch(queryBuilder, searchExecutionContext, future);
+
+        var exception = expectThrows(ElasticsearchStatusException.class, () -> future.actionGet(TimeValue.timeValueSeconds(30)));
+        assertThat(
+            exception.getMessage(),
+            containsString("expected a result of type [text_expansion_result] received [text_embedding_result]")
+        );
+        assertThat(exception.status(), is(RestStatus.BAD_REQUEST));
+    }
+
+    private static TextExpansionQueryBuilder createTestQueryBuilderNoPruning() {
+        return new TextExpansionQueryBuilder(
+            RANK_FEATURES_FIELD,
+            randomAlphaOfLength(4),
+            randomAlphaOfLength(4),
+            new TokenPruningConfig(randomIntBetween(1, 100), randomFloat(), randomBoolean())
+        );
+    }
+
 }

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/vectors/TextEmbeddingQueryVectorBuilderFailureTests.java

@@ -0,0 +1,80 @@
+/*
+ * Copyright Elasticsearch B.V. and/or licensed to Elasticsearch B.V. under one
+ * or more contributor license agreements. Licensed under the Elastic License
+ * 2.0; you may not use this file except in compliance with the Elastic License
+ * 2.0.
+ */
+
+package org.elasticsearch.xpack.ml.vectors;
+
+import org.elasticsearch.ElasticsearchStatusException;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.action.ActionRequest;
+import org.elasticsearch.action.ActionResponse;
+import org.elasticsearch.action.ActionType;
+import org.elasticsearch.action.support.PlainActionFuture;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.inference.WeightedToken;
+import org.elasticsearch.rest.RestStatus;
+import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.test.client.NoOpClient;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.core.ml.action.InferModelAction;
+import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults;
+import org.elasticsearch.xpack.core.ml.vectors.TextEmbeddingQueryVectorBuilder;
+
+import java.util.List;
+
+import static org.hamcrest.Matchers.containsString;
+import static org.hamcrest.Matchers.is;
+
+public class TextEmbeddingQueryVectorBuilderFailureTests extends ESTestCase {
+
+    private static class AssertingClient extends NoOpClient {
+        private final TextEmbeddingQueryVectorBuilder queryVectorBuilder;
+
+        AssertingClient(ThreadPool threadPool, TextEmbeddingQueryVectorBuilder queryVectorBuilder) {
+            super(threadPool);
+            this.queryVectorBuilder = queryVectorBuilder;
+        }
+
+        @Override
+        @SuppressWarnings("unchecked")
+        protected <Request extends ActionRequest, Response extends ActionResponse> void doExecute(
+            ActionType<Response> action,
+            Request request,
+            ActionListener<Response> listener
+        ) {
+            listener.onResponse((Response) createResponse(queryVectorBuilder.getModelId()));
+        }
+    }
+
+    public void testReceiving_TextExpansionResults_ThrowsBadRequestException() {
+        var queryVectorBuilder = createTestInstance();
+
+        try (var threadPool = createThreadPool()) {
+            final var client = new AssertingClient(threadPool, queryVectorBuilder);
+            PlainActionFuture<float[]> future = new PlainActionFuture<>();
+            queryVectorBuilder.buildVector(client, future);
+
+            var exception = expectThrows(ElasticsearchStatusException.class, () -> future.actionGet(TimeValue.timeValueSeconds(30)));
+            assertThat(
+                exception.getMessage(),
+                containsString("expected a result of type [text_embedding_result] received [text_expansion_result]")
+            );
+            assertThat(exception.status(), is(RestStatus.BAD_REQUEST));
+        }
+    }
+
+    private static ActionResponse createResponse(String modelId) {
+        return new InferModelAction.Response(
+            List.of(new TextExpansionResults("foo", List.of(new WeightedToken("toke", 0.1f)), randomBoolean())),
+            modelId,
+            true
+        );
+    }
+
+    private static TextEmbeddingQueryVectorBuilder createTestInstance() {
+        return new TextEmbeddingQueryVectorBuilder(randomAlphaOfLength(4), randomAlphaOfLength(4));
+    }
+}

x-pack/plugin/ml/src/test/java/org/elasticsearch/xpack/ml/vectors/TextEmbeddingQueryVectorBuilderTests.java

@@ -46,6 +46,7 @@ protected void doAssertClientRequest(ActionRequest request, TextEmbeddingQueryVe
         assertEquals(CoordinatedInferenceAction.Request.RequestModelType.NLP_MODEL, inferRequest.getRequestModelType());
     }
 
+    @Override
     public ActionResponse createResponse(float[] array, TextEmbeddingQueryVectorBuilder builder) {
         double[] embedding = new double[array.length];
         for (int i = 0; i < embedding.length; i++) {