Add autocomplete hints for ESQL function parameters

afoucret · afoucret · commit bc2319d1aede · 2025-10-07T11:52:12.000+02:00
diff --git a/docs/reference/query-languages/esql/_snippets/functions/parameters/text_embedding.md b/docs/reference/query-languages/esql/_snippets/functions/parameters/text_embedding.md
diff --git a/docs/reference/query-languages/esql/kibana/definition/functions/text_embedding.json b/docs/reference/query-languages/esql/kibana/definition/functions/text_embedding.json
diff --git a/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/LookupJoinTypesIT.java b/x-pack/plugin/esql/src/internalClusterTest/java/org/elasticsearch/xpack/esql/action/LookupJoinTypesIT.java
@@ -1032,8 +1032,8 @@ private static void saveJoinTypes(Supplier<Set<DocsV3Support.TypeSignature>> sig
             return;
         }
         ArrayList<EsqlFunctionRegistry.ArgSignature> args = new ArrayList<>();
-        args.add(new EsqlFunctionRegistry.ArgSignature("field from the left index", null, null, false, false));
-        args.add(new EsqlFunctionRegistry.ArgSignature("field from the lookup index", null, null, false, false));
+        args.add(new EsqlFunctionRegistry.ArgSignature("field from the left index", null, null, false, null, false));
+        args.add(new EsqlFunctionRegistry.ArgSignature("field from the lookup index", null, null, false, null, false));
         DocsV3Support.CommandsDocsSupport docs = new DocsV3Support.CommandsDocsSupport(
             "lookup-join",
             LookupJoinTypesIT.class,
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/EsqlFunctionRegistry.java
@@ -577,24 +577,44 @@ 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, boolean variadic) {
-            this(name, type, description, optional, variadic, UNSUPPORTED);
+        public ArgSignature(
+            String name,
+            String[] type,
+            String description,
+            boolean optional,
+            AutocompleteHint autocompleteHint,
+            boolean variadic
+        ) {
+            this(name, type, description, optional, variadic, autocompleteHint, UNSUPPORTED);
         }
 
         public String name() {
@@ -650,7 +670,7 @@ public static class MapArgSignature extends ArgSignature {
         private final Map<String, MapEntryArgSignature> mapParams;
 
         public MapArgSignature(String name, String description, boolean optional, Map<String, MapEntryArgSignature> mapParams) {
-            super(name, new String[] { "map" }, description, optional, false);
+            super(name, new String[] { "map" }, description, optional, null, false);
             this.mapParams = mapParams;
         }
 
@@ -769,7 +789,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) {
@@ -787,7 +825,7 @@ public static ArgSignature mapParam(MapParam mapParam) {
     }
 
     public static ArgSignature paramWithoutAnnotation(String name) {
-        return new EsqlFunctionRegistry.ArgSignature(name, new String[] { "?" }, "", false, false, UNSUPPORTED);
+        return new EsqlFunctionRegistry.ArgSignature(name, new String[] { "?" }, "", false, false, null, UNSUPPORTED);
     }
 
     /**
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/Param.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/Param.java
@@ -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();
+        }
+    }
 }
diff --git a/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/inference/TextEmbedding.java b/x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/expression/function/inference/TextEmbedding.java
@@ -56,7 +56,11 @@ public TextEmbedding(
         @Param(
             name = InferenceFunction.INFERENCE_ID_PARAMETER_NAME,
             type = { "keyword" },
-            description = "Identifier of the inference endpoint"
+            description = "Identifier of the inference endpoint (must be of type text_embedding).",
+            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));
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/DocsV3Support.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/DocsV3Support.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();
diff --git a/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/grouping/BucketTests.java b/x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/expression/function/grouping/BucketTests.java
@@ -351,6 +351,7 @@ public static EsqlFunctionRegistry.ArgSignature patchKibanaSignature(EsqlFunctio
             arg.description(),
             false,
             arg.variadic(),
+            null,
             arg.targetDataType()
         );
     }

Original file line number	Diff line number	Diff line change
`@@ -351,6 +351,7 @@ public static EsqlFunctionRegistry.ArgSignature patchKibanaSignature(EsqlFunctio`
`351`	`351`	`arg.description(),`
`352`	`352`	`false,`
`353`	`353`	`arg.variadic(),`
	`354`	`+ null,`
`354`	`355`	`arg.targetDataType()`
`355`	`356`	`);`
`356`	`357`	`}`