Refactor code

Author: saikatsarkar056

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

@@ -170,6 +170,7 @@ static TransportVersion def(int id) {
     public static final TransportVersion RESOURCE_DEPRECATION_CHECKS = def(8_836_00_0);
     public static final TransportVersion LINEAR_RETRIEVER_SUPPORT = def(8_837_00_0);
     public static final TransportVersion TIMEOUT_GET_PARAM_FOR_RESOLVE_CLUSTER = def(8_838_00_0);
+    public static final TransportVersion ML_INFERENCE_IBM_WATSONX_RERANK_ADDED = def(8_839_00_0);
 
     /*
      * STOP! READ THIS FIRST! No, really,

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

@@ -19,6 +19,7 @@
 import org.elasticsearch.xpack.inference.external.response.ibmwatsonx.IbmWatsonxRankedResponseEntity;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.rerank.IbmWatsonxRerankModel;
 
+import java.util.List;
 import java.util.Objects;
 import java.util.function.Supplier;
 
@@ -40,7 +41,7 @@ public static IbmWatsonxRerankRequestManager of(IbmWatsonxRerankModel model, Thr
 
     private final IbmWatsonxRerankModel model;
 
-    private IbmWatsonxRerankRequestManager(IbmWatsonxRerankModel model, ThreadPool threadPool) {
+    public IbmWatsonxRerankRequestManager(IbmWatsonxRerankModel model, ThreadPool threadPool) {
         super(threadPool, model);
         this.model = model;
     }
@@ -54,8 +55,19 @@ public void execute(
     ) {
         var rerankInput = QueryAndDocsInputs.of(inferenceInputs);
 
-        IbmWatsonxRerankRequest request = new IbmWatsonxRerankRequest(rerankInput.getQuery(), rerankInput.getChunks(), model);
+        execute(new ExecutableInferenceRequest(requestSender,
+            logger,
+            getRerankRequest(rerankInput.getQuery(), rerankInput.getChunks(), model),
+            HANDLER,
+            hasRequestCompletedFunction,
+            listener));
+    }
 
-        execute(new ExecutableInferenceRequest(requestSender, logger, request, HANDLER, hasRequestCompletedFunction, listener));
+    protected IbmWatsonxRerankRequest getRerankRequest(
+       String query,
+       List<String> chunks,
+       IbmWatsonxRerankModel model
+    ) {
+        return new IbmWatsonxRerankRequest(query, chunks, model);
     }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/request/ibmwatsonx/IbmWatsonxRerankRequest.java

@@ -9,13 +9,11 @@
 
 import org.apache.http.HttpHeaders;
 import org.apache.http.client.methods.HttpPost;
-import org.apache.http.client.utils.URIBuilder;
 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.external.request.cohere.CohereUtils;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.rerank.IbmWatsonxRerankModel;
 import org.elasticsearch.xpack.inference.services.ibmwatsonx.rerank.IbmWatsonxRerankTaskSettings;
 
@@ -31,7 +29,6 @@ public class IbmWatsonxRerankRequest implements IbmWatsonxRequest {
     private final List<String> input;
     private final IbmWatsonxRerankTaskSettings taskSettings;
     private final IbmWatsonxRerankModel model;
-    private final String inferenceEntityId;
 
     public IbmWatsonxRerankRequest(String query, List<String> input, IbmWatsonxRerankModel model) {
         Objects.requireNonNull(model);
@@ -40,7 +37,6 @@ public IbmWatsonxRerankRequest(String query, List<String> input, IbmWatsonxReran
         this.query = Objects.requireNonNull(query);
         taskSettings = model.getTaskSettings();
         this.model = model;
-        inferenceEntityId = model.getInferenceEntityId();
     }
 
     @Override
@@ -81,7 +77,7 @@ public void decorateWithAuth(HttpPost httpPost) {
 
     @Override
     public String getInferenceEntityId() {
-        return inferenceEntityId;
+        return model.getInferenceEntityId();
     }
 
     @Override
@@ -111,10 +107,4 @@ public boolean[] getTruncationInfo() {
         return null;
     }
 
-    public static URI buildDefaultUri() throws URISyntaxException {
-        return new URIBuilder().setScheme("https")
-            .setHost(CohereUtils.HOST)
-            .setPathSegments(CohereUtils.VERSION_1, CohereUtils.RERANK_PATH)
-            .build();
-    }
 }

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/ibmwatsonx/IbmWatsonxServiceSettings.java

@@ -40,13 +40,17 @@
 import static org.elasticsearch.xpack.inference.services.ServiceUtils.removeAsType;
 import static org.elasticsearch.xpack.inference.services.ibmwatsonx.IbmWatsonxServiceFields.API_VERSION;
 import static org.elasticsearch.xpack.inference.services.ibmwatsonx.IbmWatsonxServiceFields.PROJECT_ID;
-import static org.elasticsearch.xpack.inference.services.ibmwatsonx.IbmWatsonxServiceSettings.DEFAULT_RATE_LIMIT_SETTINGS;
-import static org.elasticsearch.xpack.inference.services.ibmwatsonx.IbmWatsonxServiceSettings.MODEL_ID;
+import static org.elasticsearch.xpack.inference.services.ServiceFields.MODEL_ID;
 
 public class IbmWatsonxRerankServiceSettings extends FilteredXContentObject implements ServiceSettings, IbmWatsonxRateLimitServiceSettings {
     public static final String NAME = "ibm_watsonx_rerank_service_settings";
 
-    private static final Logger logger = LogManager.getLogger(IbmWatsonxRerankServiceSettings.class);
+    /**
+     * Rate limits are defined at
+     * <a href="https://www.ibm.com/docs/en/watsonx/saas?topic=learning-watson-machine-plans">Watson Machine Learning plans</a>.
+     * For Lite plan, you've 120 requests per minute.
+     */
+    private static final RateLimitSettings DEFAULT_RATE_LIMIT_SETTINGS = new RateLimitSettings(120);
 
     public static IbmWatsonxRerankServiceSettings fromMap(Map<String, Object> map, ConfigurationParseContext context) {
         ValidationException validationException = new ValidationException();
@@ -177,7 +181,7 @@ protected XContentBuilder toXContentFragmentOfExposedFields(XContentBuilder buil
 
     @Override
     public TransportVersion getMinimalSupportedVersion() {
-        return TransportVersions.V_8_14_0;
+        return TransportVersions.ML_INFERENCE_IBM_WATSONX_RERANK_ADDED;
     }
 
     @Override
@@ -189,9 +193,7 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeOptionalString(modelId);
         out.writeOptionalString(projectId);
 
-        if (out.getTransportVersion().onOrAfter(TransportVersions.V_8_15_0)) {
-            rateLimitSettings.writeTo(out);
-        }
+        rateLimitSettings.writeTo(out);
     }
 
     @Override