[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

@@ -188,7 +188,7 @@ protected QueryBuilder doRewrite(QueryRewriteContext queryRewriteContext) {
                         listener.onFailure(new IllegalStateException(warning.getWarning()));
                     } else {
                         listener.onFailure(
-                            new IllegalStateException(
+                            new IllegalArgumentException(
                                 "expected a result of type ["
                                     + TextExpansionResults.NAME
                                     + "] received ["

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

@@ -132,7 +132,7 @@ 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 IllegalArgumentException(
                     "expected a result of type ["
                         + MlTextEmbeddingResults.NAME
                         + "] received ["

x-pack/plugin/inference/src/yamlRestTest/resources/rest-api-spec/test/inference/47_semantic_text_knn.yml

@@ -624,3 +624,43 @@ setup:
 
   - match: { error.root_cause.0.type: "illegal_argument_exception" }
   - match: { error.root_cause.0.reason: "Field [inference_field] does not use a [text_embedding] model" }
+
+---
+"Text expansion query against semantic_text field using a dense vector model returns an failure":
+  - requires:
+      cluster_features: "search.new_semantic_query_interceptors"
+      reason: New semantic query interceptors
+      test_runner_features: [ "allowed_warnings" ]
+
+  - do:
+      indices.create:
+        index: test-semantic-text-index-using-dense-vector
+        body:
+          mappings:
+            properties:
+              inference_field:
+                type: semantic_text
+                inference_id: dense-inference-id
+
+  - do:
+      index:
+        index: test-semantic-text-index-using-dense-vector
+        id: doc_1
+        body:
+          inference_field: [ "inference test", "another inference test" ]
+        refresh: true
+
+  - do:
+      catch: bad_request
+      search:
+        index: test-semantic-text-index-using-dense-vector
+        body:
+          query:
+            text_expansion:
+              embedding:
+                model_id: dense-inference-id
+                model_text: "octopus comforter smells"
+      allowed_warnings:
+        - "text_expansion is deprecated. Use sparse_vector instead."
+  - match: { error.root_cause.0.type: "illegal_argument_exception" }
+  - match: { error.root_cause.0.reason: "expected a result of type [text_expansion_result] received [text_embedding_result]. Is [dense-inference-id] a compatible model?" }

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

@@ -0,0 +1,76 @@
+/*
+ * 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.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.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;
+
+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(IllegalArgumentException.class, () -> future.actionGet(TimeValue.timeValueSeconds(30)));
+            assertThat(
+                exception.getMessage(),
+                containsString("expected a result of type [text_embedding_result] received [text_expansion_result]")
+            );
+        }
+    }
+
+    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++) {