[Inference API] Add Google Vertex AI Rerank support (elastic#110273)

Author: timgrein

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

@@ -206,6 +206,7 @@ static TransportVersion def(int id) {
  public static final TransportVersion SECURITY_MIGRATIONS_MIGRATION_NEEDED_ADDED = def(8_697_00_0);
  public static final TransportVersion K_FOR_KNN_QUERY_ADDED = def(8_698_00_0);
  public static final TransportVersion TEXT_SIMILARITY_RERANKER_RETRIEVER = def(8_699_00_0);
+ public static final TransportVersion ML_INFERENCE_GOOGLE_VERTEX_AI_RERANKING_ADDED = def(8_700_00_0);
 
  /*
  * STOP! READ THIS FIRST! No, really,

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

@@ -53,6 +53,8 @@
 import org.elasticsearch.xpack.inference.services.googlevertexai.GoogleVertexAiSecretSettings;
 import org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsServiceSettings;
 import org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsTaskSettings;
+import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankServiceSettings;
+import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankTaskSettings;
 import org.elasticsearch.xpack.inference.services.huggingface.HuggingFaceServiceSettings;
 import org.elasticsearch.xpack.inference.services.huggingface.elser.HuggingFaceElserServiceSettings;
 import org.elasticsearch.xpack.inference.services.mistral.embeddings.MistralEmbeddingsServiceSettings;
@@ -314,6 +316,22 @@ private static void addGoogleVertexAiNamedWriteables(List<NamedWriteableRegistry
  GoogleVertexAiEmbeddingsTaskSettings::new
  )
  );
+
+ namedWriteables.add(
+ new NamedWriteableRegistry.Entry(
+ ServiceSettings.class,
+ GoogleVertexAiRerankServiceSettings.NAME,
+ GoogleVertexAiRerankServiceSettings::new
+ )
+ );
+
+ namedWriteables.add(
+ new NamedWriteableRegistry.Entry(
+ TaskSettings.class,
+ GoogleVertexAiRerankTaskSettings.NAME,
+ GoogleVertexAiRerankTaskSettings::new
+ )
+ );
  }
 
  private static void addInternalElserNamedWriteables(List<NamedWriteableRegistry.Entry> namedWriteables) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/googlevertexai/GoogleVertexAiActionCreator.java

@@ -11,6 +11,7 @@
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
 import org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsModel;
+import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankModel;
 
 import java.util.Map;
 import java.util.Objects;
@@ -30,4 +31,9 @@ public GoogleVertexAiActionCreator(Sender sender, ServiceComponents serviceCompo
  public ExecutableAction create(GoogleVertexAiEmbeddingsModel model, Map<String, Object> taskSettings) {
  return new GoogleVertexAiEmbeddingsAction(sender, model, serviceComponents);
  }
+
+ @Override
+ public ExecutableAction create(GoogleVertexAiRerankModel model, Map<String, Object> taskSettings) {
+ return new GoogleVertexAiRerankAction(sender, model, serviceComponents.threadPool());
+ }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/googlevertexai/GoogleVertexAiActionVisitor.java

@@ -9,11 +9,13 @@
 
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.services.googlevertexai.embeddings.GoogleVertexAiEmbeddingsModel;
+import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankModel;
 
 import java.util.Map;
 
 public interface GoogleVertexAiActionVisitor {
 
  ExecutableAction create(GoogleVertexAiEmbeddingsModel model, Map<String, Object> taskSettings);
 
+ ExecutableAction create(GoogleVertexAiRerankModel model, Map<String, Object> taskSettings);
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/action/googlevertexai/GoogleVertexAiRerankAction.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.inference.external.action.googlevertexai;
+
+import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.core.TimeValue;
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
+import org.elasticsearch.xpack.inference.external.http.sender.GoogleVertexAiRerankRequestManager;
+import org.elasticsearch.xpack.inference.external.http.sender.InferenceInputs;
+import org.elasticsearch.xpack.inference.external.http.sender.Sender;
+import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankModel;
+
+import java.util.Objects;
+
+import static org.elasticsearch.xpack.inference.external.action.ActionUtils.constructFailedToSendRequestMessage;
+import static org.elasticsearch.xpack.inference.external.action.ActionUtils.createInternalServerError;
+import static org.elasticsearch.xpack.inference.external.action.ActionUtils.wrapFailuresInElasticsearchException;
+
+public class GoogleVertexAiRerankAction implements ExecutableAction {
+
+ private final String failedToSendRequestErrorMessage;
+
+ private final Sender sender;
+
+ private final GoogleVertexAiRerankRequestManager requestManager;
+
+ public GoogleVertexAiRerankAction(Sender sender, GoogleVertexAiRerankModel model, ThreadPool threadPool) {
+ Objects.requireNonNull(model);
+ this.sender = Objects.requireNonNull(sender);
+ this.failedToSendRequestErrorMessage = constructFailedToSendRequestMessage(model.uri(), "Google Vertex AI rerank");
+ this.requestManager = GoogleVertexAiRerankRequestManager.of(model, threadPool);
+ }
+
+ @Override
+ public void execute(InferenceInputs inferenceInputs, TimeValue timeout, ActionListener<InferenceServiceResults> listener) {
+ try {
+ ActionListener<InferenceServiceResults> wrappedListener = wrapFailuresInElasticsearchException(
+ failedToSendRequestErrorMessage,
+ listener
+ );
+ sender.send(requestManager, inferenceInputs, timeout, wrappedListener);
+ } catch (ElasticsearchException e) {
+ listener.onFailure(e);
+ } catch (Exception e) {
+ listener.onFailure(createInternalServerError(e, failedToSendRequestErrorMessage));
+ }
+ }
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/GoogleVertexAiEmbeddingsRequestManager.java

@@ -41,11 +41,19 @@ private static ResponseHandler createEmbeddingsHandler() {
  private final Truncator truncator;
 
  public GoogleVertexAiEmbeddingsRequestManager(GoogleVertexAiEmbeddingsModel model, Truncator truncator, ThreadPool threadPool) {
- super(threadPool, model);
+ super(threadPool, model, RateLimitGrouping.of(model));
  this.model = Objects.requireNonNull(model);
  this.truncator = Objects.requireNonNull(truncator);
  }
 
+ record RateLimitGrouping(int modelIdHash) {
+ public static RateLimitGrouping of(GoogleVertexAiEmbeddingsModel model) {
+ Objects.requireNonNull(model);
+
+ return new RateLimitGrouping(model.rateLimitServiceSettings().modelId().hashCode());
+ }
+ }
+
  @Override
  public void execute(
  String query,

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/GoogleVertexAiRequestManager.java

@@ -10,19 +10,9 @@
 import org.elasticsearch.threadpool.ThreadPool;
 import org.elasticsearch.xpack.inference.services.googlevertexai.GoogleVertexAiModel;
 
-import java.util.Objects;
-
 public abstract class GoogleVertexAiRequestManager extends BaseRequestManager {
 
- GoogleVertexAiRequestManager(ThreadPool threadPool, GoogleVertexAiModel model) {
- super(threadPool, model.getInferenceEntityId(), RateLimitGrouping.of(model), model.rateLimitServiceSettings().rateLimitSettings());
- }
-
- record RateLimitGrouping(int modelIdHash) {
- public static RateLimitGrouping of(GoogleVertexAiModel model) {
- Objects.requireNonNull(model);
-
- return new RateLimitGrouping(model.rateLimitServiceSettings().modelId().hashCode());
- }
+ GoogleVertexAiRequestManager(ThreadPool threadPool, GoogleVertexAiModel model, Object rateLimitGroup) {
+ super(threadPool, model.getInferenceEntityId(), rateLimitGroup, model.rateLimitServiceSettings().rateLimitSettings());
  }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/http/sender/GoogleVertexAiRerankRequestManager.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.external.http.sender;
+
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
+import org.elasticsearch.action.ActionListener;
+import org.elasticsearch.inference.InferenceServiceResults;
+import org.elasticsearch.threadpool.ThreadPool;
+import org.elasticsearch.xpack.inference.external.googlevertexai.GoogleVertexAiResponseHandler;
+import org.elasticsearch.xpack.inference.external.http.retry.RequestSender;
+import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler;
+import org.elasticsearch.xpack.inference.external.request.googlevertexai.GoogleVertexAiRerankRequest;
+import org.elasticsearch.xpack.inference.external.response.googlevertexai.GoogleVertexAiRerankResponseEntity;
+import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankModel;
+
+import java.util.List;
+import java.util.Objects;
+import java.util.function.Supplier;
+
+public class GoogleVertexAiRerankRequestManager extends GoogleVertexAiRequestManager {
+
+ private static final Logger logger = LogManager.getLogger(GoogleVertexAiRerankRequestManager.class);
+
+ private static final ResponseHandler HANDLER = createGoogleVertexAiResponseHandler();
+
+ private static ResponseHandler createGoogleVertexAiResponseHandler() {
+ return new GoogleVertexAiResponseHandler(
+ "Google Vertex AI rerank",
+ (request, response) -> GoogleVertexAiRerankResponseEntity.fromResponse(response)
+ );
+ }
+
+ public static GoogleVertexAiRerankRequestManager of(GoogleVertexAiRerankModel model, ThreadPool threadPool) {
+ return new GoogleVertexAiRerankRequestManager(Objects.requireNonNull(model), Objects.requireNonNull(threadPool));
+ }
+
+ private final GoogleVertexAiRerankModel model;
+
+ private GoogleVertexAiRerankRequestManager(GoogleVertexAiRerankModel model, ThreadPool threadPool) {
+ super(threadPool, model, RateLimitGrouping.of(model));
+ this.model = model;
+ }
+
+ record RateLimitGrouping(int projectIdHash) {
+ public static RateLimitGrouping of(GoogleVertexAiRerankModel model) {
+ Objects.requireNonNull(model);
+
+ return new RateLimitGrouping(model.rateLimitServiceSettings().projectId().hashCode());
+ }
+ }
+
+ @Override
+ public void execute(
+ String query,
+ List<String> input,
+ RequestSender requestSender,
+ Supplier<Boolean> hasRequestCompletedFunction,
+ ActionListener<InferenceServiceResults> listener
+ ) {
+ GoogleVertexAiRerankRequest request = new GoogleVertexAiRerankRequest(query, input, model);
+
+ execute(new ExecutableInferenceRequest(requestSender, logger, request, HANDLER, hasRequestCompletedFunction, listener));
+ }
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiRerankRequest.java

@@ -0,0 +1,83 @@
+/*
+ * 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.external.request.googlevertexai;
+
+import org.apache.http.HttpHeaders;
+import org.apache.http.client.methods.HttpPost;
+import org.apache.http.entity.ByteArrayEntity;
+import org.elasticsearch.common.Strings;
+import org.elasticsearch.xcontent.XContentType;
+import org.elasticsearch.xpack.inference.external.request.HttpRequest;
+import org.elasticsearch.xpack.inference.external.request.Request;
+import org.elasticsearch.xpack.inference.services.googlevertexai.rerank.GoogleVertexAiRerankModel;
+
+import java.net.URI;
+import java.nio.charset.StandardCharsets;
+import java.util.List;
+import java.util.Objects;
+
+public class GoogleVertexAiRerankRequest implements GoogleVertexAiRequest {
+
+ private final GoogleVertexAiRerankModel model;
+
+ private final String query;
+
+ private final List<String> input;
+
+ public GoogleVertexAiRerankRequest(String query, List<String> input, GoogleVertexAiRerankModel model) {
+ this.model = Objects.requireNonNull(model);
+ this.query = Objects.requireNonNull(query);
+ this.input = Objects.requireNonNull(input);
+ }
+
+ @Override
+ public HttpRequest createHttpRequest() {
+ HttpPost httpPost = new HttpPost(model.uri());
+
+ ByteArrayEntity byteEntity = new ByteArrayEntity(
+ Strings.toString(
+ new GoogleVertexAiRerankRequestEntity(query, input, model.getServiceSettings().modelId(), model.getTaskSettings().topN())
+ ).getBytes(StandardCharsets.UTF_8)
+ );
+
+ httpPost.setEntity(byteEntity);
+ httpPost.setHeader(HttpHeaders.CONTENT_TYPE, XContentType.JSON.mediaType());
+
+ decorateWithAuth(httpPost);
+
+ return new HttpRequest(httpPost, getInferenceEntityId());
+ }
+
+ public void decorateWithAuth(HttpPost httpPost) {
+ GoogleVertexAiRequest.decorateWithBearerToken(httpPost, model.getSecretSettings());
+ }
+
+ public GoogleVertexAiRerankModel model() {
+ return model;
+ }
+
+ @Override
+ public String getInferenceEntityId() {
+ return model.getInferenceEntityId();
+ }
+
+ @Override
+ public URI getURI() {
+ return model.uri();
+ }
+
+ @Override
+ public Request truncate() {
+ return this;
+ }
+
+ @Override
+ public boolean[] getTruncationInfo() {
+ return null;
+ }
+}

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiRerankRequestEntity.java

@@ -0,0 +1,68 @@
+/*
+ * 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.external.request.googlevertexai;
+
+import org.elasticsearch.core.Nullable;
+import org.elasticsearch.xcontent.ToXContentObject;
+import org.elasticsearch.xcontent.XContentBuilder;
+
+import java.io.IOException;
+import java.util.List;
+import java.util.Objects;
+
+public record GoogleVertexAiRerankRequestEntity(String query, List<String> inputs, @Nullable String model, @Nullable Integer topN)
+ implements
+ ToXContentObject {
+
+ private static final String MODEL_FIELD = "model";
+ private static final String QUERY_FIELD = "query";
+ private static final String RECORDS_FIELD = "records";
+ private static final String ID_FIELD = "id";
+
+ private static final String CONTENT_FIELD = "content";
+ private static final String TOP_N_FIELD = "topN";
+
+ public GoogleVertexAiRerankRequestEntity {
+ Objects.requireNonNull(query);
+ Objects.requireNonNull(inputs);
+ }
+
+ @Override
+ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
+ builder.startObject();
+
+ if (model != null) {
+ builder.field(MODEL_FIELD, model);
+ }
+
+ builder.field(QUERY_FIELD, query);
+
+ builder.startArray(RECORDS_FIELD);
+
+ for (int recordId = 0; recordId < inputs.size(); recordId++) {
+ builder.startObject();
+
+ {
+ builder.field(ID_FIELD, String.valueOf(recordId));
+ builder.field(CONTENT_FIELD, inputs.get(recordId));
+ }
+
+ builder.endObject();
+ }
+
+ builder.endArray();
+
+ if (topN != null) {
+ builder.field(TOP_N_FIELD, topN);
+ }
+
+ builder.endObject();
+
+ return builder;
+ }
+}