Support using the semantic query across multiple inference IDs

docs/changelog/133675.yaml

@@ -0,0 +1,5 @@
+pr: 133675
+summary: Support using the semantic query across multiple inference IDs
+area: Vector Search
+type: enhancement
+issues: []

server/src/main/java/org/elasticsearch/TransportVersions.java

@@ -369,6 +369,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion SCRIPT_RESCORER = def(9_143_0_00);
     public static final TransportVersion ESQL_LOOKUP_OPERATOR_EMITTED_ROWS = def(9_144_0_00);
     public static final TransportVersion ALLOCATION_DECISION_NOT_PREFERRED = def(9_145_0_00);
+    public static final TransportVersion SEMANTIC_QUERY_MULTIPLE_INFERENCE_IDS = def(9_146_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -11,6 +11,7 @@
 import org.elasticsearch.features.NodeFeature;
 import org.elasticsearch.xpack.inference.mapper.SemanticInferenceMetadataFieldsMapper;
 import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper;
+import org.elasticsearch.xpack.inference.queries.SemanticQueryBuilder;
 import org.elasticsearch.xpack.inference.rank.textsimilarity.TextSimilarityRankRetrieverBuilder;
 
 import java.util.HashSet;
@@ -82,7 +83,8 @@ public Set<NodeFeature> getTestFeatures() {
                 SEMANTIC_QUERY_REWRITE_INTERCEPTORS_PROPAGATE_BOOST_AND_QUERY_NAME_FIX,
                 SEMANTIC_TEXT_HIGHLIGHTING_FLAT,
                 SEMANTIC_TEXT_SPARSE_VECTOR_INDEX_OPTIONS,
-                SEMANTIC_TEXT_FIELDS_CHUNKS_FORMAT
+                SEMANTIC_TEXT_FIELDS_CHUNKS_FORMAT,
+                SemanticQueryBuilder.SEMANTIC_QUERY_MULTIPLE_INFERENCE_IDS
             )
         );
         if (RERANK_SNIPPETS.isEnabled()) {

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

@@ -95,10 +95,13 @@
 import org.elasticsearch.xpack.inference.mapper.OffsetSourceFieldMapper;
 import org.elasticsearch.xpack.inference.mapper.SemanticInferenceMetadataFieldsMapper;
 import org.elasticsearch.xpack.inference.mapper.SemanticTextFieldMapper;
+import org.elasticsearch.xpack.inference.queries.InferenceResultsProvider;
+import org.elasticsearch.xpack.inference.queries.MapInferenceResultsProvider;
 import org.elasticsearch.xpack.inference.queries.SemanticKnnVectorQueryRewriteInterceptor;
 import org.elasticsearch.xpack.inference.queries.SemanticMatchQueryRewriteInterceptor;
 import org.elasticsearch.xpack.inference.queries.SemanticQueryBuilder;
 import org.elasticsearch.xpack.inference.queries.SemanticSparseVectorQueryRewriteInterceptor;
+import org.elasticsearch.xpack.inference.queries.SingleInferenceResultsProvider;
 import org.elasticsearch.xpack.inference.rank.random.RandomRankBuilder;
 import org.elasticsearch.xpack.inference.rank.random.RandomRankRetrieverBuilder;
 import org.elasticsearch.xpack.inference.rank.textsimilarity.TextSimilarityRankBuilder;
@@ -430,6 +433,20 @@ public List<NamedWriteableRegistry.Entry> getNamedWriteables() {
         entries.add(new NamedWriteableRegistry.Entry(RankDoc.class, TextSimilarityRankDoc.NAME, TextSimilarityRankDoc::new));
         entries.add(new NamedWriteableRegistry.Entry(Metadata.ProjectCustom.class, ModelRegistryMetadata.TYPE, ModelRegistryMetadata::new));
         entries.add(new NamedWriteableRegistry.Entry(NamedDiff.class, ModelRegistryMetadata.TYPE, ModelRegistryMetadata::readDiffFrom));
+        entries.add(
+            new NamedWriteableRegistry.Entry(
+                InferenceResultsProvider.class,
+                MapInferenceResultsProvider.NAME,
+                MapInferenceResultsProvider::new
+            )
+        );
+        entries.add(
+            new NamedWriteableRegistry.Entry(
+                InferenceResultsProvider.class,
+                SingleInferenceResultsProvider.NAME,
+                SingleInferenceResultsProvider::new
+            )
+        );
         return entries;
     }
 

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

@@ -0,0 +1,19 @@
+/*
+ * 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.inference.queries;
+
+import org.elasticsearch.common.io.stream.NamedWriteable;
+import org.elasticsearch.inference.InferenceResults;
+
+import java.util.Collection;
+
+public interface InferenceResultsProvider extends NamedWriteable {
+    InferenceResults getInferenceResults(String inferenceId);
+
+    Collection<InferenceResults> getAllInferenceResults();
+}

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

@@ -0,0 +1,70 @@
+/*
+ * 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.inference.queries;
+
+import org.elasticsearch.common.io.stream.StreamInput;
+import org.elasticsearch.common.io.stream.StreamOutput;
+import org.elasticsearch.inference.InferenceResults;
+
+import java.io.IOException;
+import java.util.Collection;
+import java.util.Collections;
+import java.util.Map;
+import java.util.Objects;
+import java.util.concurrent.ConcurrentHashMap;
+
+public class MapInferenceResultsProvider implements InferenceResultsProvider {
+    public static final String NAME = "map_inference_results_provider";
+
+    private final Map<String, InferenceResults> inferenceResultsMap;
+
+    public MapInferenceResultsProvider() {
+        this.inferenceResultsMap = new ConcurrentHashMap<>();
+    }
+
+    public MapInferenceResultsProvider(StreamInput in) throws IOException {
+        this.inferenceResultsMap = in.readImmutableMap(i -> i.readNamedWriteable(InferenceResults.class));
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeMap(inferenceResultsMap, StreamOutput::writeNamedWriteable);
+    }
+
+    @Override
+    public InferenceResults getInferenceResults(String inferenceId) {
+        return inferenceResultsMap.get(inferenceId);
+    }
+
+    @Override
+    public Collection<InferenceResults> getAllInferenceResults() {
+        return Collections.unmodifiableCollection(inferenceResultsMap.values());
+    }
+
+    public void addInferenceResults(String inferenceId, InferenceResults inferenceResults) {
+        this.inferenceResultsMap.put(inferenceId, inferenceResults);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        MapInferenceResultsProvider that = (MapInferenceResultsProvider) o;
+        return Objects.equals(inferenceResultsMap, that.inferenceResultsMap);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hashCode(inferenceResultsMap);
+    }
+}