Remove ElasticInferenceServiceRerankRequestManager

Author: timgrein

x-pack/plugin/core/src/test/java/org/elasticsearch/xpack/core/inference/results/RankedDocsResultsTests.java

@@ -66,27 +66,6 @@ protected RankedDocsResults mutateInstanceForVersion(RankedDocsResults instance,
         }
     }
 
-    public record RerankExpectation(int index, float relevanceScore) {}
-
-    public static Map<String, Object> buildExpectationRankedDocResults(List<RerankExpectation> rerankExpectations) {
-        return Map.of(
-            RankedDocsResults.RERANK,
-            rerankExpectations.stream()
-                .map(
-                    rerankExpectation -> Map.of(
-                        RankedDocsResults.RankedDoc.NAME,
-                        Map.of(
-                            RankedDocsResults.RankedDoc.INDEX,
-                            rerankExpectation.index,
-                            RankedDocsResults.RankedDoc.RELEVANCE_SCORE,
-                            rerankExpectation.relevanceScore
-                        )
-                    )
-                )
-                .toList()
-        );
-    }
-
     private List<RankedDocsResults.RankedDoc> rankedDocsNullStringToEmpty(List<RankedDocsResults.RankedDoc> rankedDocs) {
         var result = new ArrayList<RankedDocsResults.RankedDoc>(rankedDocs.size());
         for (var doc : rankedDocs) {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/elastic/rerank/ElasticInferenceServiceRerankRequest.java

@@ -15,8 +15,8 @@
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.inference.external.request.Request;
-import org.elasticsearch.xpack.inference.external.request.elastic.ElasticInferenceServiceRequest;
-import org.elasticsearch.xpack.inference.external.request.elastic.ElasticInferenceServiceRequestMetadata;
+import org.elasticsearch.xpack.inference.services.elastic.request.ElasticInferenceServiceRequest;
+import org.elasticsearch.xpack.inference.services.elastic.request.ElasticInferenceServiceRequestMetadata;
 import org.elasticsearch.xpack.inference.services.elastic.rerank.ElasticInferenceServiceRerankModel;
 import org.elasticsearch.xpack.inference.telemetry.TraceContext;
 import org.elasticsearch.xpack.inference.telemetry.TraceContextHandler;

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceModel.java

@@ -7,14 +7,15 @@
 
 package org.elasticsearch.xpack.inference.services.elastic;
 
-import org.elasticsearch.inference.Model;
 import org.elasticsearch.inference.ModelConfigurations;
 import org.elasticsearch.inference.ModelSecrets;
 import org.elasticsearch.inference.ServiceSettings;
+import org.elasticsearch.xpack.inference.services.RateLimitGroupingModel;
+import org.elasticsearch.xpack.inference.services.settings.RateLimitSettings;
 
 import java.util.Objects;
 
-public abstract class ElasticInferenceServiceModel extends Model {
+public abstract class ElasticInferenceServiceModel extends RateLimitGroupingModel {
 
     private final ElasticInferenceServiceRateLimitServiceSettings rateLimitServiceSettings;
 
@@ -35,12 +36,18 @@ public ElasticInferenceServiceModel(
     public ElasticInferenceServiceModel(ElasticInferenceServiceModel model, ServiceSettings serviceSettings) {
         super(model, serviceSettings);
 
-        this.rateLimitServiceSettings = model.rateLimitServiceSettings();
+        this.rateLimitServiceSettings = model.rateLimitServiceSettings;
         this.elasticInferenceServiceComponents = model.elasticInferenceServiceComponents();
     }
 
-    public ElasticInferenceServiceRateLimitServiceSettings rateLimitServiceSettings() {
-        return rateLimitServiceSettings;
+    @Override
+    public int rateLimitGroupingHash() {
+        // We only have one model for rerank
+        return Objects.hash(this.getServiceSettings().modelId());
+    }
+
+    public RateLimitSettings rateLimitSettings() {
+        return rateLimitServiceSettings.rateLimitSettings();
     }
 
     public ElasticInferenceServiceComponents elasticInferenceServiceComponents() {

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceRequestManager.java

@@ -20,7 +20,7 @@ public abstract class ElasticInferenceServiceRequestManager extends BaseRequestM
     private final ElasticInferenceServiceRequestMetadata requestMetadata;
 
     protected ElasticInferenceServiceRequestManager(ThreadPool threadPool, ElasticInferenceServiceModel model) {
-        super(threadPool, model.getInferenceEntityId(), RateLimitGrouping.of(model), model.rateLimitServiceSettings().rateLimitSettings());
+        super(threadPool, model.getInferenceEntityId(), RateLimitGrouping.of(model), model.rateLimitSettings());
         this.requestMetadata = extractRequestMetadataFromThreadContext(threadPool.getThreadContext());
     }
 
@@ -32,7 +32,7 @@ record RateLimitGrouping(int modelIdHash) {
         public static RateLimitGrouping of(ElasticInferenceServiceModel model) {
             Objects.requireNonNull(model);
 
-            return new RateLimitGrouping(model.rateLimitServiceSettings().modelId().hashCode());
+            return new RateLimitGrouping(model.rateLimitGroupingHash());
         }
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceRerankRequestManager.java

@@ -31,6 +31,7 @@
 import static org.elasticsearch.xpack.inference.common.Truncator.truncate;
 import static org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceService.ELASTIC_INFERENCE_SERVICE_IDENTIFIER;
 
+// TODO: remove and use GenericRequestManager in ElasticInferenceServiceActionCreator
 public class ElasticInferenceServiceSparseEmbeddingsRequestManager extends ElasticInferenceServiceRequestManager {
 
     private static final Logger logger = LogManager.getLogger(ElasticInferenceServiceSparseEmbeddingsRequestManager.class);

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/ElasticInferenceServiceSparseEmbeddingsRequestManager.java

@@ -9,9 +9,14 @@
 
 import org.elasticsearch.xpack.inference.external.action.ExecutableAction;
 import org.elasticsearch.xpack.inference.external.action.SenderExecutableAction;
+import org.elasticsearch.xpack.inference.external.http.retry.ResponseHandler;
+import org.elasticsearch.xpack.inference.external.http.sender.GenericRequestManager;
+import org.elasticsearch.xpack.inference.external.http.sender.QueryAndDocsInputs;
 import org.elasticsearch.xpack.inference.external.http.sender.Sender;
+import org.elasticsearch.xpack.inference.external.request.elastic.rerank.ElasticInferenceServiceRerankRequest;
+import org.elasticsearch.xpack.inference.external.response.elastic.ElasticInferenceServiceRerankResponseEntity;
 import org.elasticsearch.xpack.inference.services.ServiceComponents;
-import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceRerankRequestManager;
+import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceResponseHandler;
 import org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceServiceSparseEmbeddingsRequestManager;
 import org.elasticsearch.xpack.inference.services.elastic.rerank.ElasticInferenceServiceRerankModel;
 import org.elasticsearch.xpack.inference.services.elastic.sparseembeddings.ElasticInferenceServiceSparseEmbeddingsModel;
@@ -22,6 +27,7 @@
 
 import static org.elasticsearch.xpack.inference.external.action.ActionUtils.constructFailedToSendRequestMessage;
 import static org.elasticsearch.xpack.inference.services.elastic.ElasticInferenceService.ELASTIC_INFERENCE_SERVICE_IDENTIFIER;
+import static org.elasticsearch.xpack.inference.services.elastic.request.ElasticInferenceServiceRequest.extractRequestMetadataFromThreadContext;
 
 public class ElasticInferenceServiceActionCreator implements ElasticInferenceServiceActionVisitor {
 
@@ -31,6 +37,11 @@ public class ElasticInferenceServiceActionCreator implements ElasticInferenceSer
 
     private final TraceContext traceContext;
 
+    static final ResponseHandler RERANK_HANDLER = new ElasticInferenceServiceResponseHandler(
+        "elastic rerank",
+        (request, response) -> ElasticInferenceServiceRerankResponseEntity.fromResponse(response)
+    );
+
     public ElasticInferenceServiceActionCreator(Sender sender, ServiceComponents serviceComponents, TraceContext traceContext) {
         this.sender = Objects.requireNonNull(sender);
         this.serviceComponents = Objects.requireNonNull(serviceComponents);
@@ -48,7 +59,15 @@ public ExecutableAction create(ElasticInferenceServiceSparseEmbeddingsModel mode
 
     @Override
     public ExecutableAction create(ElasticInferenceServiceRerankModel model) {
-        var requestManager = new ElasticInferenceServiceRerankRequestManager(model, serviceComponents, traceContext);
+        var threadPool = serviceComponents.threadPool();
+        var requestManager = new GenericRequestManager<>(
+            threadPool,
+            model,
+            RERANK_HANDLER,
+            (rerankInput) -> new ElasticInferenceServiceRerankRequest(rerankInput.getQuery(), rerankInput.getChunks(), model,
+                traceContext, extractRequestMetadataFromThreadContext(threadPool.getThreadContext())),
+            QueryAndDocsInputs.class
+        );
         var errorMessage = constructFailedToSendRequestMessage(
             String.format(Locale.ROOT, "%s rerank", ELASTIC_INFERENCE_SERVICE_IDENTIFIER)
         );

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/elastic/action/ElasticInferenceServiceActionCreator.java

@@ -18,7 +18,7 @@
 import java.util.List;
 
 import static org.elasticsearch.xpack.inference.external.http.Utils.entityAsMap;
-import static org.elasticsearch.xpack.inference.external.request.elastic.ElasticInferenceServiceRequestTests.randomElasticInferenceServiceRequestMetadata;
+import static org.elasticsearch.xpack.inference.services.elastic.request.ElasticInferenceServiceRequestTests.randomElasticInferenceServiceRequestMetadata;
 import static org.hamcrest.Matchers.aMapWithSize;
 import static org.hamcrest.Matchers.instanceOf;
 import static org.hamcrest.Matchers.is;

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/elastic/ElasticInferenceServiceRerankRequestTests.java

@@ -225,10 +225,10 @@ public void testExecute_ReturnsSuccessfulResponse_ForRerankAction() throws IOExc
             assertThat(
                 result.asMap(),
                 equalTo(
-                    RankedDocsResultsTests.buildExpectationRankedDocResults(
+                    RankedDocsResultsTests.buildExpectationRerank(
                         List.of(
-                            new RankedDocsResultsTests.RerankExpectation(0, 0.94f),
-                            new RankedDocsResultsTests.RerankExpectation(1, 0.21f)
+                            new RankedDocsResultsTests.RerankExpectation(Map.of("index", 0, "relevance_score", 0.94f)),
+                            new RankedDocsResultsTests.RerankExpectation(Map.of("index", 1, "relevance_score", 0.21f))
                         )
                     )
                 )