Add settings to control rerank and completion command.

Author: afoucret

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

@@ -0,0 +1,55 @@
+/*
+ * 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;
+
+import java.util.Map;
+
+public record InferenceCommandConfig(boolean enabled, int rowLimit) {
+
+    public static final Setting<Boolean> COMPLETION_ENABLED_SETTING = commandEnabledSetting("completion");
+    public static final Setting<Integer> COMPLETION_ROW_LIMIT_SETTING = rowLimitSetting("completion", 100);
+    public static final Setting<Boolean> RERANK_ENABLED_SETTING = commandEnabledSetting("rerank");
+    public static final Setting<Integer> RERANK_ROW_LIMIT_SETTING = rowLimitSetting("rerank", 1000);
+
+    public static InferenceCommandConfig completionCommandConfig(Settings settings) {
+        return new InferenceCommandConfig(COMPLETION_ENABLED_SETTING.get(settings), COMPLETION_ROW_LIMIT_SETTING.get(settings));
+    }
+
+    public static InferenceCommandConfig rerankCommandConfig(Settings settings) {
+        return new InferenceCommandConfig(RERANK_ENABLED_SETTING.get(settings), RERANK_ROW_LIMIT_SETTING.get(settings));
+    }
+
+    public static Map<String, InferenceCommandConfig> fromSettings(Settings settings) {
+        return Map.of("completion", completionCommandConfig(settings), "rerank", rerankCommandConfig(settings));
+    }
+
+    private static final String ENABLED_SETTING_PATTERN = "inference.command.%s.enabled";
+    private static final String ROW_LIMIT_PATTERN = "inference.command.%s.row_limit";
+
+    private static Setting<Boolean> commandEnabledSetting(String commandName) {
+        return Setting.boolSetting(
+            String.format(ENABLED_SETTING_PATTERN, commandName),
+            true,
+            Setting.Property.NodeScope,
+            Setting.Property.Dynamic
+        );
+    }
+
+    private static Setting<Integer> rowLimitSetting(String commandName, int defaultValue) {
+        return Setting.intSetting(
+            String.format(ROW_LIMIT_PATTERN, commandName),
+            defaultValue,
+            -1,
+            Setting.Property.NodeScope,
+            Setting.Property.Dynamic
+        );
+    }
+}

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

@@ -8,9 +8,16 @@
 package org.elasticsearch.xpack.esql.parser;
 
 import org.elasticsearch.Build;
+import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.xpack.esql.inference.InferenceCommandConfig;
+
+import java.util.Locale;
+import java.util.Map;
 
 class EsqlConfig {
 
+    private final Map<String, InferenceCommandConfig> inferenceCommandConfigs;
+
     // versioning information
     boolean devVersion = Build.current().isSnapshot();
 
@@ -25,4 +32,18 @@ boolean isReleaseVersion() {
     public void setDevVersion(boolean dev) {
         this.devVersion = dev;
     }
+
+    EsqlConfig(Settings settings) {
+        inferenceCommandConfigs = InferenceCommandConfig.fromSettings(settings);
+    }
+
+    public InferenceCommandConfig inferenceCommandConfig(String commandName) {
+        InferenceCommandConfig config = inferenceCommandConfigs.get(commandName.toLowerCase(Locale.ROOT));
+
+        if (config == null) {
+            throw new IllegalArgumentException("Unknown inference command: " + commandName);
+        }
+
+        return config;
+    }
 }

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

@@ -16,6 +16,7 @@
 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;
@@ -88,14 +89,19 @@ private static void replaceSymbolWithLiteral(Map<String, String> symbolReplaceme
         replaceSymbolWithLiteral(symbolReplacements, lexerVocab.getLiteralNames(), lexerVocab.getSymbolicNames());
     }
 
-    private EsqlConfig config = new EsqlConfig();
+    private final EsqlConfig config;
 
-    public EsqlConfig config() {
-        return config;
+    public EsqlParser() {
+        config = new EsqlConfig(Settings.EMPTY);
     }
 
-    public void setEsqlConfig(EsqlConfig config) {
-        this.config = config;
+    public EsqlParser(Settings settings) {
+        config = new EsqlConfig(settings);
+
+    }
+
+    public EsqlConfig config() {
+        return config;
     }
 
     // testing utility

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

@@ -74,6 +74,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.InferenceCommandConfig;
 import org.elasticsearch.xpack.esql.io.stream.ExpressionQueryBuilder;
 import org.elasticsearch.xpack.esql.io.stream.PlanStreamWrapperQueryBuilder;
 import org.elasticsearch.xpack.esql.plan.PlanWritables;
@@ -248,7 +249,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,
+            InferenceCommandConfig.COMPLETION_ENABLED_SETTING,
+            InferenceCommandConfig.COMPLETION_ROW_LIMIT_SETTING,
+            InferenceCommandConfig.RERANK_ENABLED_SETTING,
+            InferenceCommandConfig.RERANK_ROW_LIMIT_SETTING
         );
     }
 

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

@@ -97,6 +97,7 @@
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Supplier;
 
 import static java.util.stream.Collectors.joining;
 import static java.util.stream.Collectors.toSet;
@@ -142,6 +143,7 @@ public interface PlanRunner {
     private boolean explainMode;
     private String parsedPlanString;
     private String optimizedLogicalPlanString;
+    private Supplier<EsqlParser> parserSupplier;
 
     public EsqlSession(
         String sessionId,
@@ -173,6 +175,7 @@ public EsqlSession(
         this.intermediateLocalRelationMaxSize = services.plannerSettings().intermediateLocalRelationMaxSize();
         this.crossProjectModeDecider = services.crossProjectModeDecider();
         this.clusterName = services.clusterService().getClusterName().value();
+        this.parserSupplier = () -> new EsqlParser(services.clusterService().getSettings());
     }
 
     public String sessionId() {
@@ -451,7 +454,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 = parserSupplier.get().createQuery(query, params, planTelemetry);
         if (LOGGER.isDebugEnabled()) {
             LOGGER.debug("Parsed logical plan:\n{}", parsed.plan());
             LOGGER.debug("Parsed settings:\n[{}]", parsed.settings().stream().map(QuerySetting::toString).collect(joining("; ")));