Add settings for completion and rerank ES|QL commands.

Author: afoucret

x-pack/plugin/esql/qa/testFixtures/src/main/java/org/elasticsearch/xpack/esql/EsqlTestUtils.java

@@ -565,7 +565,7 @@ public static LogicalOptimizerContext unboundLogicalOptimizerContext() {
         mock(ProjectResolver.class),
         mock(IndexNameExpressionResolver.class),
         null,
-        new InferenceService(mock(Client.class)),
+        new InferenceService(mock(Client.class), Settings.EMPTY),
         new BlockFactoryProvider(PlannerUtils.NON_BREAKING_BLOCK_FACTORY),
         TEST_PLANNER_SETTINGS,
         new CrossProjectModeDecider(Settings.EMPTY)

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceService.java

@@ -8,11 +8,15 @@
 package org.elasticsearch.xpack.esql.inference;
 
 import org.elasticsearch.client.internal.Client;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
 import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRunner;
 import org.elasticsearch.xpack.esql.inference.bulk.BulkInferenceRunnerConfig;
 
 public class InferenceService {
+
+    private final InferenceSettings inferenceSettings;
+
     private final InferenceResolver.Factory inferenceResolverFactory;
 
     private final BulkInferenceRunner.Factory bulkInferenceRunnerFactory;
@@ -21,14 +25,29 @@ public class InferenceService {
      * Creates a new inference service with the given client.
      *
      * @param client the Elasticsearch client for inference operations
+     * @param settings the node settings
      */
-    public InferenceService(Client client) {
-        this(InferenceResolver.factory(client), BulkInferenceRunner.factory(client));
+    public InferenceService(Client client, Settings settings) {
+        this(InferenceResolver.factory(client), BulkInferenceRunner.factory(client), new InferenceSettings(settings));
     }
 
-    private InferenceService(InferenceResolver.Factory inferenceResolverFactory, BulkInferenceRunner.Factory bulkInferenceRunnerFactory) {
+    private InferenceService(
+        InferenceResolver.Factory inferenceResolverFactory,
+        BulkInferenceRunner.Factory bulkInferenceRunnerFactory,
+        InferenceSettings inferenceSettings
+    ) {
         this.inferenceResolverFactory = inferenceResolverFactory;
         this.bulkInferenceRunnerFactory = bulkInferenceRunnerFactory;
+        this.inferenceSettings = inferenceSettings;
+    }
+
+    /**
+     * Returns the settings for ES|QL inference features.
+     *
+     * @return the inference settings
+     */
+    public InferenceSettings inferenceSettings() {
+        return inferenceSettings;
     }
 
     /**

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/inference/InferenceSettings.java

@@ -0,0 +1,56 @@
+/*
+ * 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.esql.inference;
+
+import org.elasticsearch.common.settings.Setting;
+import org.elasticsearch.common.settings.Settings;
+
+/**
+ * Settings for inference features such as completion and rerank.
+ */
+public record InferenceSettings(boolean completionEnabled, int completionMaxSize, boolean rerankEnabled, int rerankMaxSize) {
+
+    public static final Setting<Boolean> COMPLETION_ENABLED_SETTING = Setting.boolSetting(
+        "esql.command.completion.enabled",
+        true,
+        Setting.Property.Dynamic,
+        Setting.Property.NodeScope
+    );
+
+    public static final Setting<Integer> COMPLETION_ROW_LIMIT_SETTING = Setting.intSetting(
+        "esql.command.completion.limit",
+        100,
+        1,
+        Setting.Property.Dynamic,
+        Setting.Property.NodeScope
+    );
+
+    public static final Setting<Boolean> RERANK_ENABLED_SETTING = Setting.boolSetting(
+        "esql.command.rerank.enabled",
+        true,
+        Setting.Property.Dynamic,
+        Setting.Property.NodeScope
+    );
+
+    public static final Setting<Integer> RERANK_ROW_LIMIT_SETTING = Setting.intSetting(
+        "esql.command.rerank.limit",
+        1000,
+        1,
+        Setting.Property.Dynamic,
+        Setting.Property.NodeScope
+    );
+
+    public InferenceSettings(Settings settings) {
+        this(
+            COMPLETION_ENABLED_SETTING.get(settings),
+            COMPLETION_ROW_LIMIT_SETTING.get(settings),
+            RERANK_ENABLED_SETTING.get(settings),
+            RERANK_ROW_LIMIT_SETTING.get(settings)
+        );
+    }
+}

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/EsqlParser.java

@@ -16,10 +16,12 @@
 import org.antlr.v4.runtime.TokenSource;
 import org.antlr.v4.runtime.VocabularyImpl;
 import org.antlr.v4.runtime.atn.PredictionMode;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.logging.LogManager;
 import org.elasticsearch.logging.Logger;
 import org.elasticsearch.xpack.esql.core.util.StringUtils;
 import org.elasticsearch.xpack.esql.expression.function.EsqlFunctionRegistry;
+import org.elasticsearch.xpack.esql.inference.InferenceSettings;
 import org.elasticsearch.xpack.esql.plan.EsqlStatement;
 import org.elasticsearch.xpack.esql.plan.logical.LogicalPlan;
 import org.elasticsearch.xpack.esql.telemetry.PlanTelemetry;
@@ -105,14 +107,14 @@ public LogicalPlan createStatement(String query) {
 
     // testing utility
     public LogicalPlan createStatement(String query, QueryParams params) {
-        return createStatement(query, params, new PlanTelemetry(new EsqlFunctionRegistry()));
+        return createStatement(query, params, new PlanTelemetry(new EsqlFunctionRegistry()), new InferenceSettings(Settings.EMPTY));
     }
 
-    public LogicalPlan createStatement(String query, QueryParams params, PlanTelemetry metrics) {
+    public LogicalPlan createStatement(String query, QueryParams params, PlanTelemetry metrics, InferenceSettings inferenceSettings) {
         if (log.isDebugEnabled()) {
             log.debug("Parsing as statement: {}", query);
         }
-        return invokeParser(query, params, metrics, EsqlBaseParser::singleStatement, AstBuilder::plan);
+        return invokeParser(query, params, metrics, inferenceSettings, EsqlBaseParser::singleStatement, AstBuilder::plan);
     }
 
     // testing utility
@@ -122,20 +124,21 @@ public EsqlStatement createQuery(String query) {
 
     // testing utility
     public EsqlStatement createQuery(String query, QueryParams params) {
-        return createQuery(query, params, new PlanTelemetry(new EsqlFunctionRegistry()));
+        return createQuery(query, params, new PlanTelemetry(new EsqlFunctionRegistry()), new InferenceSettings(Settings.EMPTY));
     }
 
-    public EsqlStatement createQuery(String query, QueryParams params, PlanTelemetry metrics) {
+    public EsqlStatement createQuery(String query, QueryParams params, PlanTelemetry metrics, InferenceSettings inferenceSettings) {
         if (log.isDebugEnabled()) {
             log.debug("Parsing as statement: {}", query);
         }
-        return invokeParser(query, params, metrics, EsqlBaseParser::statements, AstBuilder::statement);
+        return invokeParser(query, params, metrics, inferenceSettings, EsqlBaseParser::statements, AstBuilder::statement);
     }
 
     private <T> T invokeParser(
         String query,
         QueryParams params,
         PlanTelemetry metrics,
+        InferenceSettings inferenceSettings,
         Function<EsqlBaseParser, ParserRuleContext> parseFunction,
         BiFunction<AstBuilder, ParserRuleContext, T> result
     ) {
@@ -169,7 +172,7 @@ private <T> T invokeParser(
                 log.trace("Parse tree: {}", tree.toStringTree());
             }
 
-            return result.apply(new AstBuilder(new ExpressionBuilder.ParsingContext(params, metrics)), tree);
+            return result.apply(new AstBuilder(new ExpressionBuilder.ParsingContext(params, metrics, inferenceSettings)), tree);
         } catch (StackOverflowError e) {
             throw new ParsingException("ESQL statement is too large, causing stack overflow when generating the parsing tree: [{}]", query);
             // likely thrown by an invalid popMode (such as extra closing parenthesis)

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/parser/ExpressionBuilder.java

@@ -67,6 +67,7 @@
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.InsensitiveEquals;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThan;
 import org.elasticsearch.xpack.esql.expression.predicate.operator.comparison.LessThanOrEqual;
+import org.elasticsearch.xpack.esql.inference.InferenceSettings;
 import org.elasticsearch.xpack.esql.telemetry.PlanTelemetry;
 import org.elasticsearch.xpack.esql.type.EsqlDataTypeConverter;
 
@@ -124,7 +125,7 @@ public abstract class ExpressionBuilder extends IdentifierBuilder {
 
     protected final ParsingContext context;
 
-    public record ParsingContext(QueryParams params, PlanTelemetry telemetry) {}
+    public record ParsingContext(QueryParams params, PlanTelemetry telemetry, InferenceSettings inferenceSettings) {}
 
     ExpressionBuilder(ParsingContext context) {
         this.context = context;

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/EsqlPlugin.java

@@ -73,6 +73,7 @@
 import org.elasticsearch.xpack.esql.enrich.LookupFromIndexOperator;
 import org.elasticsearch.xpack.esql.execution.PlanExecutor;
 import org.elasticsearch.xpack.esql.expression.ExpressionWritables;
+import org.elasticsearch.xpack.esql.inference.InferenceSettings;
 import org.elasticsearch.xpack.esql.io.stream.ExpressionQueryBuilder;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamWrapperQueryBuilder;
 import org.elasticsearch.xpack.esql.plan.PlanWritables;
@@ -236,7 +237,11 @@ public List<Setting<?>> getSettings() {
             PlannerSettings.REDUCTION_LATE_MATERIALIZATION,
             STORED_FIELDS_SEQUENTIAL_PROPORTION,
             EsqlFlags.ESQL_STRING_LIKE_ON_INDEX,
-            EsqlFlags.ESQL_ROUNDTO_PUSHDOWN_THRESHOLD
+            EsqlFlags.ESQL_ROUNDTO_PUSHDOWN_THRESHOLD,
+            InferenceSettings.COMPLETION_ENABLED_SETTING,
+            InferenceSettings.COMPLETION_ROW_LIMIT_SETTING,
+            InferenceSettings.RERANK_ENABLED_SETTING,
+            InferenceSettings.RERANK_ROW_LIMIT_SETTING
         );
     }
 

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/plugin/TransportEsqlQueryAction.java

@@ -171,7 +171,7 @@ public TransportEsqlQueryAction(
             projectResolver,
             indexNameExpressionResolver,
             usageService,
-            new InferenceService(client),
+            new InferenceService(client, clusterService.getSettings()),
             blockFactoryProvider,
             new PlannerSettings(clusterService),
             new CrossProjectModeDecider(clusterService.getSettings())

x-pack/plugin/esql/src/main/java/org/elasticsearch/xpack/esql/session/EsqlSession.java

@@ -473,7 +473,7 @@ private static void releaseLocalRelationBlocks(AtomicReference<Page> localRelati
     }
 
     private EsqlStatement parse(String query, QueryParams params) {
-        var parsed = new EsqlParser().createQuery(query, params, planTelemetry);
+        var parsed = new EsqlParser().createQuery(query, params, planTelemetry, inferenceService.inferenceSettings());
         if (LOGGER.isDebugEnabled()) {
             LOGGER.debug("Parsed logical plan:\n{}", parsed.plan());
             LOGGER.debug("Parsed settings:\n[{}]", parsed.settings().stream().map(QuerySetting::toString).collect(joining("; ")));

x-pack/plugin/esql/src/test/java/org/elasticsearch/xpack/esql/inference/InferenceOperatorTestCase.java

@@ -14,6 +14,7 @@
 import org.elasticsearch.action.ActionType;
 import org.elasticsearch.client.internal.Client;
 import org.elasticsearch.common.logging.LoggerMessageFormat;
+import org.elasticsearch.common.settings.Settings;
 import org.elasticsearch.compute.data.Block;
 import org.elasticsearch.compute.data.BlockFactory;
 import org.elasticsearch.compute.data.BooleanBlock;
@@ -140,7 +141,7 @@ private void runWithRandomDelay(Runnable runnable) {
             }
         };
 
-        return new InferenceService(mockClient);
+        return new InferenceService(mockClient, Settings.EMPTY);
     }
 
     protected abstract InferenceResultsType mockInferenceResult(InferenceAction.Request request);