Added the map embeddings provider

Author: Mikep86

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

@@ -93,6 +93,8 @@
 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.EmbeddingsProvider;
+import org.elasticsearch.xpack.inference.queries.MapEmbeddingsProvider;
 import org.elasticsearch.xpack.inference.queries.SemanticKnnVectorQueryRewriteInterceptor;
 import org.elasticsearch.xpack.inference.queries.SemanticMatchQueryRewriteInterceptor;
 import org.elasticsearch.xpack.inference.queries.SemanticQueryBuilder;
@@ -426,6 +428,7 @@ 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(EmbeddingsProvider.class, MapEmbeddingsProvider.NAME, MapEmbeddingsProvider::new));
         return entries;
     }
 

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

@@ -0,0 +1,15 @@
+/*
+ * 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;
+
+public interface EmbeddingsProvider extends NamedWriteable {
+    InferenceResults getEmbeddings(InferenceEndpointKey key);
+}

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

@@ -0,0 +1,50 @@
+/*
+ * 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.common.io.stream.Writeable;
+import org.elasticsearch.inference.MinimalServiceSettings;
+
+import java.io.IOException;
+import java.util.Objects;
+
+public class InferenceEndpointKey implements Writeable {
+    private final String inferenceId;
+    private final MinimalServiceSettings serviceSettings;
+
+    public InferenceEndpointKey(String inferenceId, MinimalServiceSettings serviceSettings) {
+        this.inferenceId = inferenceId;
+        this.serviceSettings = serviceSettings;
+    }
+
+    public InferenceEndpointKey(StreamInput in) throws IOException {
+        this.inferenceId = in.readString();
+        this.serviceSettings = in.readNamedWriteable(MinimalServiceSettings.class);
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeString(inferenceId);
+        out.writeNamedWriteable(serviceSettings);
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (this == o) return true;
+        if (o == null || getClass() != o.getClass()) return false;
+        InferenceEndpointKey that = (InferenceEndpointKey) o;
+        return Objects.equals(inferenceId, that.inferenceId) && Objects.equals(serviceSettings, that.serviceSettings);
+    }
+
+    @Override
+    public int hashCode() {
+        return Objects.hash(inferenceId, serviceSettings);
+    }
+}

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

@@ -0,0 +1,49 @@
+/*
+ * 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.HashMap;
+import java.util.Map;
+
+public class MapEmbeddingsProvider implements EmbeddingsProvider {
+    public static final String NAME = "map_embeddings_provider";
+
+    private final Map<InferenceEndpointKey, InferenceResults> embeddings;
+
+    public MapEmbeddingsProvider() {
+        this.embeddings = new HashMap<>();
+    }
+
+    public MapEmbeddingsProvider(StreamInput in) throws IOException {
+        embeddings = in.readMap(InferenceEndpointKey::new, i -> i.readNamedWriteable(InferenceResults.class));
+    }
+
+    @Override
+    public String getWriteableName() {
+        return NAME;
+    }
+
+    @Override
+    public void writeTo(StreamOutput out) throws IOException {
+        out.writeMap(embeddings);
+    }
+
+    @Override
+    public InferenceResults getEmbeddings(InferenceEndpointKey key) {
+        return embeddings.get(key);
+    }
+
+    public void addEmbeddings(InferenceEndpointKey key, InferenceResults embeddings) {
+        this.embeddings.put(key, embeddings);
+    }
+}