Enhance documentation for OpenShift AI models and add task settings handling in rerank model

Author: Jan-Kazlouski-elastic

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openshiftai/completion/OpenShiftAiChatCompletionModel.java

@@ -23,7 +23,7 @@
 
 /**
  * Represents an OpenShift AI chat completion model.
- * This class extends the OpenShiftAiModel and provides specific configurations for chat completion tasks.
+ * This class extends the {@link OpenShiftAiModel} and provides specific configurations for chat completion tasks.
  */
 public class OpenShiftAiChatCompletionModel extends OpenShiftAiModel {
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openshiftai/embeddings/OpenShiftAiEmbeddingsModel.java

@@ -23,7 +23,7 @@
 
 /**
  * Represents an OpenShift AI embeddings model for inference.
- * This class extends the OpenShiftAiModel and provides specific configurations and settings for embeddings tasks.
+ * This class extends the {@link OpenShiftAiModel} and provides specific configurations and settings for embeddings tasks.
  */
 public class OpenShiftAiEmbeddingsModel extends OpenShiftAiModel {
 

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/services/openshiftai/rerank/OpenShiftAiRerankModel.java

@@ -19,9 +19,23 @@
 
 import java.util.Map;
 
+/**
+ * Represents an OpenShift AI rerank model.
+ * This class extends the {@link OpenShiftAiModel} and provides specific configurations for rerank tasks.
+ */
 public class OpenShiftAiRerankModel extends OpenShiftAiModel {
+
+    /**
+     * Creates a new {@link OpenShiftAiRerankModel} with updated task settings if they differ from the existing ones.
+     * @param model the existing OpenShift AI rerank model
+     * @param taskSettings the new task settings to apply
+     * @return a new {@link OpenShiftAiRerankModel} with updated task settings, or the original model if settings are unchanged
+     */
     public static OpenShiftAiRerankModel of(OpenShiftAiRerankModel model, Map<String, Object> taskSettings) {
         var requestTaskSettings = OpenShiftAiRerankTaskSettings.fromMap(taskSettings);
+        if (requestTaskSettings.isEmpty() || requestTaskSettings.equals(model.getTaskSettings())) {
+            return model;
+        }
         return new OpenShiftAiRerankModel(model, OpenShiftAiRerankTaskSettings.of(model.getTaskSettings(), requestTaskSettings));
     }
 

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openshiftai/rerank/OpenShiftAiRerankModelTests.java

@@ -13,6 +13,13 @@
 import org.elasticsearch.test.ESTestCase;
 import org.elasticsearch.xpack.inference.services.settings.DefaultSecretSettings;
 
+import java.util.HashMap;
+import java.util.Map;
+
+import static org.elasticsearch.xpack.inference.services.openshiftai.rerank.OpenShiftAiRerankTaskSettings.RETURN_DOCUMENTS;
+import static org.elasticsearch.xpack.inference.services.openshiftai.rerank.OpenShiftAiRerankTaskSettings.TOP_N;
+import static org.hamcrest.Matchers.is;
+
 public class OpenShiftAiRerankModelTests extends ESTestCase {
 
     public static OpenShiftAiRerankModel createModel(String url, String apiKey, @Nullable String modelId) {
@@ -35,4 +42,66 @@ public static OpenShiftAiRerankModel createModel(
             new DefaultSecretSettings(new SecureString(apiKey.toCharArray()))
         );
     }
+
+    public void testOverrideWith_SameParams_KeepsSameModel() {
+        testOverrideWith_KeepsSameModel(buildTaskSettingsMap(2, true));
+    }
+
+    public void testOverrideWith_EmptyParams_KeepsSameModel() {
+        testOverrideWith_KeepsSameModel(buildTaskSettingsMap(null, null));
+    }
+
+    private static void testOverrideWith_KeepsSameModel(Map<String, Object> taskSettings) {
+        var model = createModel("url", "api_key", "model_name", 2, true);
+        var overriddenModel = OpenShiftAiRerankModel.of(model, taskSettings);
+
+        assertThat(overriddenModel.getTaskSettings().getTopN(), is(2));
+        assertThat(overriddenModel.getTaskSettings().getReturnDocuments(), is(true));
+    }
+
+    public void testOverrideWith_DifferentParams_OverridesAllTaskSettings() {
+        testOverrideWith_DifferentParams(buildTaskSettingsMap(4, false), 4, false);
+    }
+
+    public void testOverrideWith_DifferentParams_OverridesOnlyReturnDocuments() {
+        testOverrideWith_DifferentParams(buildTaskSettingsMap(null, false), 2, false);
+    }
+
+    public void testOverrideWith_DifferentParams_OverridesOnlyTopN() {
+        testOverrideWith_DifferentParams(buildTaskSettingsMap(4, null), 4, true);
+    }
+
+    public void testOverrideWith_DifferentParams_OverridesNullValues() {
+        var model = createModel("url", "api_key", "model_name", null, null);
+        var overriddenModel = OpenShiftAiRerankModel.of(model, buildTaskSettingsMap(4, false));
+
+        assertThat(overriddenModel.getTaskSettings().getTopN(), is(4));
+        assertThat(overriddenModel.getTaskSettings().getReturnDocuments(), is(false));
+    }
+
+    private static void testOverrideWith_DifferentParams(
+        Map<String, Object> taskSettings,
+        int expectedTopN,
+        boolean expectedReturnDocuments
+    ) {
+        var model = createModel("url", "api_key", "model_name", 2, true);
+        var overriddenModel = OpenShiftAiRerankModel.of(model, taskSettings);
+
+        assertThat(overriddenModel.getTaskSettings().getTopN(), is(expectedTopN));
+        assertThat(overriddenModel.getTaskSettings().getReturnDocuments(), is(expectedReturnDocuments));
+    }
+
+    private static Map<String, Object> buildTaskSettingsMap(@Nullable Integer topN, @Nullable Boolean returnDocuments) {
+        final var map = new HashMap<String, Object>();
+
+        if (returnDocuments != null) {
+            map.put(RETURN_DOCUMENTS, returnDocuments);
+        }
+
+        if (topN != null) {
+            map.put(TOP_N, topN);
+        }
+
+        return map;
+    }
 }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/services/openshiftai/rerank/OpenShiftAiRerankTaskSettingsTests.java

@@ -8,15 +8,20 @@
 package org.elasticsearch.xpack.inference.services.openshiftai.rerank;
 
 import org.elasticsearch.TransportVersion;
+import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.ValidationException;
 import org.elasticsearch.common.io.stream.Writeable;
 import org.elasticsearch.test.ESTestCase;
+import org.elasticsearch.xcontent.XContentBuilder;
+import org.elasticsearch.xcontent.XContentFactory;
+import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.core.ml.AbstractBWCWireSerializationTestCase;
 
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.Map;
 
+import static org.elasticsearch.xpack.inference.MatchersUtils.equalToIgnoringWhitespaceInJsonString;
 import static org.hamcrest.Matchers.containsString;
 
 public class OpenShiftAiRerankTaskSettingsTests extends AbstractBWCWireSerializationTestCase<OpenShiftAiRerankTaskSettings> {
@@ -97,6 +102,59 @@ public void testUpdatedTaskSettings_WithMultipleNewValues_ReturnsUpdatedSettings
         assertEquals(7, updatedSettings.getTopN().intValue());
     }
 
+    public void testToXContent_WritesAllValues() throws IOException {
+        Integer topN = 2;
+        Boolean doReturnDocuments = true;
+
+        testToXContent(topN, doReturnDocuments, """
+            {
+                "top_n":2,
+                "return_documents":true
+            }
+            """);
+    }
+
+    public void testToXContent_EmptyValues() throws IOException {
+        Integer topN = null;
+        Boolean doReturnDocuments = null;
+
+        testToXContent(topN, doReturnDocuments, """
+            {}
+            """);
+    }
+
+    public void testToXContent_OnlyTopN() throws IOException {
+        Integer topN = 2;
+        Boolean doReturnDocuments = null;
+
+        testToXContent(topN, doReturnDocuments, """
+            {
+                "top_n":2
+            }
+            """);
+    }
+
+    public void testToXContent_OnlyReturnDocuments() throws IOException {
+        Integer topN = null;
+        Boolean doReturnDocuments = true;
+
+        testToXContent(topN, doReturnDocuments, """
+            {
+                "return_documents":true
+            }
+            """);
+    }
+
+    private static void testToXContent(Integer topN, Boolean doReturnDocuments, String expectedString) throws IOException {
+        var taskSettings = new OpenShiftAiRerankTaskSettings(topN, doReturnDocuments);
+
+        XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
+        taskSettings.toXContent(builder, null);
+        String xContentResult = Strings.toString(builder);
+
+        assertThat(xContentResult, equalToIgnoringWhitespaceInJsonString(expectedString));
+    }
+
     @Override
     protected Writeable.Reader<OpenShiftAiRerankTaskSettings> instanceReader() {
         return OpenShiftAiRerankTaskSettings::new;