Adding settings for query time inference config and applying during semantic query build

Author: Samiul-TheSoccerFan

server/src/main/java/org/elasticsearch/index/mapper/InferenceMetadataFieldsMapper.java

@@ -15,11 +15,13 @@
 import org.elasticsearch.cluster.metadata.IndexMetadata;
 import org.elasticsearch.common.settings.Setting;
 import org.elasticsearch.common.settings.Settings;
+import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.index.IndexVersion;
 import org.elasticsearch.index.IndexVersions;
 import org.elasticsearch.index.query.SearchExecutionContext;
 
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 import java.util.function.Function;
 
 /**
@@ -28,6 +30,7 @@
  * the field name for removal from _source.
  */
 public abstract class InferenceMetadataFieldsMapper extends MetadataFieldMapper {
+    public static final TimeValue DEFAULT_SEMANTIC_TEXT_INFERENCE_TIMEOUT = TimeValue.timeValueSeconds(TimeUnit.SECONDS.toSeconds(10));
     /**
      * Internal index setting to control the format used for semantic text fields.
      * Determines whether to use the legacy format (default: true).
@@ -41,6 +44,12 @@ public abstract class InferenceMetadataFieldsMapper extends MetadataFieldMapper
         Setting.Property.IndexScope,
         Setting.Property.InternalIndex
     );
+    public static final Setting<TimeValue> SEMANTIC_TEXT_INFERENCE_TIMEOUT = Setting.timeSetting(
+        "index.semantic_text.inference_timeout",
+        DEFAULT_SEMANTIC_TEXT_INFERENCE_TIMEOUT,
+        Setting.Property.IndexScope,
+        Setting.Property.Dynamic
+    );
 
     // Check index version SOURCE_MAPPER_MODE_ATTRIBUTE_NOOP because that index version was added in the same serverless promotion
     // where the new format was enabled by default

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

@@ -25,6 +25,7 @@
 import org.elasticsearch.core.IOUtils;
 import org.elasticsearch.core.TimeValue;
 import org.elasticsearch.features.NodeFeature;
+import org.elasticsearch.index.mapper.InferenceMetadataFieldsMapper;
 import org.elasticsearch.index.mapper.Mapper;
 import org.elasticsearch.index.mapper.MetadataFieldMapper;
 import org.elasticsearch.indices.SystemIndexDescriptor;
@@ -495,7 +496,7 @@ public List<Setting<?>> getSettings() {
         settings.add(SKIP_VALIDATE_AND_START);
         settings.add(INDICES_INFERENCE_BATCH_SIZE);
         settings.addAll(ElasticInferenceServiceSettings.getSettingsDefinitions());
-
+        settings.add(InferenceMetadataFieldsMapper.SEMANTIC_TEXT_INFERENCE_TIMEOUT);
         return settings;
     }
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/queries/SemanticQueryBuilder.java

@@ -16,6 +16,8 @@
 import org.elasticsearch.cluster.metadata.InferenceFieldMetadata;
 import org.elasticsearch.common.io.stream.StreamInput;
 import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.index.mapper.InferenceMetadataFieldsMapper;
 import org.elasticsearch.index.mapper.MappedFieldType;
 import org.elasticsearch.index.query.AbstractQueryBuilder;
 import org.elasticsearch.index.query.MatchNoneQueryBuilder;
@@ -31,7 +33,6 @@
 import org.elasticsearch.xcontent.XContentBuilder;
 import org.elasticsearch.xcontent.XContentParser;
 import org.elasticsearch.xpack.core.inference.action.InferenceAction;
-import org.elasticsearch.xpack.core.ml.action.InferModelAction;
 import org.elasticsearch.xpack.core.ml.inference.results.ErrorInferenceResults;
 import org.elasticsearch.xpack.core.ml.inference.results.MlTextEmbeddingResults;
 import org.elasticsearch.xpack.core.ml.inference.results.TextExpansionResults;
@@ -225,6 +226,10 @@ private SemanticQueryBuilder doRewriteGetInferenceResults(QueryRewriteContext qu
         }
 
         String inferenceId = getInferenceIdForForField(resolvedIndices.getConcreteLocalIndicesMetadata().values(), fieldName);
+        TimeValue inferenceTimeout = getInferenceTimeeoutForSemanticField(
+            resolvedIndices.getConcreteLocalIndicesMetadata().values(),
+            fieldName
+        );
         SetOnce<InferenceServiceResults> inferenceResultsSupplier = new SetOnce<>();
         boolean noInferenceResults = false;
         if (inferenceId != null) {
@@ -237,7 +242,7 @@ private SemanticQueryBuilder doRewriteGetInferenceResults(QueryRewriteContext qu
                 List.of(query),
                 Map.of(),
                 InputType.INTERNAL_SEARCH,
-                InferModelAction.Request.DEFAULT_TIMEOUT_FOR_API,
+                inferenceTimeout,
                 false
             );
 
@@ -264,6 +269,33 @@ private SemanticQueryBuilder doRewriteGetInferenceResults(QueryRewriteContext qu
         return new SemanticQueryBuilder(this, noInferenceResults ? null : inferenceResultsSupplier, null, noInferenceResults);
     }
 
+    @SuppressWarnings("unchecked")
+    private TimeValue getInferenceTimeeoutForSemanticField(Collection<IndexMetadata> indexMetadataCollection, String fieldName) {
+        TimeValue inferenceTimeout = InferenceMetadataFieldsMapper.DEFAULT_SEMANTIC_TEXT_INFERENCE_TIMEOUT;
+        for (IndexMetadata indexMetadata : indexMetadataCollection) {
+            boolean fieldExistsInIndex = indexMetadata.mapping()
+                .getSourceAsMap()
+                .values()
+                .stream()
+                .filter(v -> v instanceof Map)
+                .map(v -> (Map<String, Object>) v)
+                .anyMatch(m -> m.containsKey(fieldName));
+
+            if (fieldExistsInIndex == false) {
+                continue;
+            }
+
+            TimeValue currentInferenceTimeout = indexMetadata.getSettings()
+                .getAsTime("index.semantic_text.inference_timeout", InferenceMetadataFieldsMapper.DEFAULT_SEMANTIC_TEXT_INFERENCE_TIMEOUT);
+
+            if (currentInferenceTimeout.compareTo(inferenceTimeout) < 0) {
+                inferenceTimeout = currentInferenceTimeout;
+            }
+        }
+
+        return inferenceTimeout;
+    }
+
     private static InferenceResults validateAndConvertInferenceResults(
         SetOnce<InferenceServiceResults> inferenceResultsSupplier,
         String fieldName