Add autocomplete hints for ESQL function parameters

Author: afoucret

docs/reference/query-languages/esql/_snippets/functions/examples/text_embedding.md

@@ -610,24 +610,52 @@ public static String normalizeName(String name) {
     }
 
     public static class ArgSignature {
+
+        public record AutocompleteHint(String entityType, Map<String, String> constraints) {}
+
         protected final String name;
         protected final String[] type;
         protected final String description;
         protected final boolean optional;
         protected final boolean variadic;
         protected final DataType targetDataType;
-
-        public ArgSignature(String name, String[] type, String description, boolean optional, boolean variadic, DataType targetDataType) {
+        protected final AutocompleteHint autocompleteHint;
+
+        public ArgSignature(
+            String name,
+            String[] type,
+            String description,
+            boolean optional,
+            boolean variadic,
+            AutocompleteHint autocompleteHint,
+            DataType targetDataType
+        ) {
             this.name = name;
             this.type = type;
             this.description = description;
             this.optional = optional;
             this.variadic = variadic;
             this.targetDataType = targetDataType;
+            this.autocompleteHint = autocompleteHint;
+        }
+
+        public ArgSignature(
+            String name,
+            String[] type,
+            String description,
+            boolean optional,
+            AutocompleteHint autocompleteHint,
+            boolean variadic
+        ) {
+            this(name, type, description, optional, variadic, autocompleteHint, UNSUPPORTED);
+        }
+
+        public ArgSignature(String name, String[] type, String description, boolean optional, boolean variadic, DataType targetDataType) {
+            this(name, type, description, optional, variadic, null, targetDataType);
         }
 
         public ArgSignature(String name, String[] type, String description, boolean optional, boolean variadic) {
-            this(name, type, description, optional, variadic, UNSUPPORTED);
+            this(name, type, description, optional, variadic, null, UNSUPPORTED);
         }
 
         public String name() {
@@ -810,7 +838,25 @@ public static ArgSignature param(Param param, boolean variadic) {
         String[] type = removeUnderConstruction(param.type());
         String desc = param.description().replace('\n', ' ');
         DataType targetDataType = getTargetType(type);
-        return new EsqlFunctionRegistry.ArgSignature(param.name(), type, desc, param.optional(), variadic, targetDataType);
+        ArgSignature.AutocompleteHint autocompleteHint = null;
+        if (param.autocompleteHint() != null && param.autocompleteHint().entityType() != Param.AutocompleteHint.ENTITY_TYPE.NONE) {
+            Map<String, String> constraints = Arrays.stream(param.autocompleteHint().constraints())
+                .collect(Collectors.toMap(Param.AutocompleteHint.Constraint::name, Param.AutocompleteHint.Constraint::value));
+            autocompleteHint = new ArgSignature.AutocompleteHint(
+                param.autocompleteHint().entityType().name().toLowerCase(Locale.ROOT),
+                constraints
+            );
+        }
+
+        return new EsqlFunctionRegistry.ArgSignature(
+            param.name(),
+            type,
+            desc,
+            param.optional(),
+            variadic,
+            autocompleteHint,
+            targetDataType
+        );
     }
 
     public static ArgSignature mapParam(MapParam mapParam) {

docs/reference/query-languages/esql/kibana/definition/functions/text_embedding.json

@@ -7,6 +7,8 @@
 
 package org.elasticsearch.xpack.esql.expression.function;
 
+import org.elasticsearch.core.Nullable;
+
 import java.lang.annotation.ElementType;
 import java.lang.annotation.Retention;
 import java.lang.annotation.RetentionPolicy;
@@ -25,4 +27,28 @@
     String description() default "";
 
     boolean optional() default false;
+
+    @Nullable
+    AutocompleteHint autocompleteHint() default @AutocompleteHint(entityType = AutocompleteHint.ENTITY_TYPE.NONE);
+
+    @Retention(RetentionPolicy.RUNTIME)
+    @Target(ElementType.PARAMETER)
+    @interface AutocompleteHint {
+        enum ENTITY_TYPE {
+            NONE,
+            INFERENCE_ENDPOINT,
+        }
+
+        ENTITY_TYPE entityType();
+
+        Constraint[] constraints() default {};
+
+        @Retention(RetentionPolicy.RUNTIME)
+        @Target(ElementType.PARAMETER)
+        @interface Constraint {
+            String name();
+
+            String value();
+        }
+    }
 }

docs/reference/query-languages/esql/kibana/docs/functions/text_embedding.md

@@ -48,17 +48,7 @@ public class TextEmbedding extends InferenceFunction<TextEmbedding> {
         preview = true,
         examples = {
             @Example(
-                description = "Basic text embedding generation from a text string using an inference endpoint.",
-                file = "text-embedding",
-                tag = "text-embedding-eval"
-            ),
-            @Example(
-                description = "Generate text embeddings and store them in a variable for reuse in KNN vector search queries.",
-                file = "text-embedding",
-                tag = "text-embedding-knn"
-            ),
-            @Example(
-                description = "Directly embed text within a KNN query for streamlined vector search without intermediate variables.",
+                description = "Generate text embeddings using the 'test_dense_inference' inference endpoint.",
                 file = "text-embedding",
                 tag = "text-embedding-knn-inline"
             ) }
@@ -75,7 +65,11 @@ public TextEmbedding(
             type = { "keyword" },
             description = "Identifier of an existing inference endpoint the that will generate the embeddings. "
                 + "The inference endpoint must have the `text_embedding` task type and should use the same model "
-                + "that was used to embed your indexed data."
+                + "that was used to embed your indexed data.",
+            autocompleteHint = @Param.AutocompleteHint(
+                entityType = Param.AutocompleteHint.ENTITY_TYPE.INFERENCE_ENDPOINT,
+                constraints = { @Param.AutocompleteHint.Constraint(name = "task_type", value = "text_embedding") }
+            )
         ) Expression inferenceId
     ) {
         super(source, List.of(inputText, inferenceId));
@@ -109,7 +103,7 @@ public boolean foldable() {
 
     @Override
     public DataType dataType() {
-        return inputText.dataType() == DataType.NULL ? DataType.NULL : DataType.DENSE_VECTOR;
+        return DataType.DENSE_VECTOR;
     }
 
     @Override

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java

@@ -622,6 +622,7 @@ private void renderDocs(FunctionDefinition definition) throws Exception {
                     trueValue.type(),
                     "The value that’s returned when no condition evaluates to `true`.",
                     true,
+                    null,
                     true
                 );
                 description = new EsqlFunctionRegistry.FunctionDescription(
@@ -1382,6 +1383,18 @@ void renderKibanaFunctionDefinition(
                         builder.field("optional", arg.optional());
                         String cleanedParamDesc = removeAppliesToBlocks(arg.description());
                         builder.field("description", cleanedParamDesc);
+                        if (arg.autocompleteHint != null) {
+                            builder.startObject("autocompleteHint");
+                            builder.field("entityType", arg.autocompleteHint.entityType());
+                            if (arg.autocompleteHint.constraints() != null && arg.autocompleteHint.constraints().size() > 0) {
+                                builder.startObject("constraints");
+                                for (Map.Entry<String, String> constraint : arg.autocompleteHint.constraints().entrySet()) {
+                                    builder.field(constraint.getKey(), constraint.getValue());
+                                }
+                                builder.endObject();
+                            }
+                            builder.endObject();
+                        }
                         builder.endObject();
                     }
                     builder.endArray();