PR feedback

Author: ymao1

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

@@ -58,9 +58,8 @@ public HttpRequest createHttpRequest() {
                     query,
                     input,
                     returnDocuments,
-                    topN,
-                    model.getServiceSettings().modelId(),
-                    model.getTaskSettings().topN()
+                    topN != null ? topN : model.getTaskSettings().topN(),
+                    model.getServiceSettings().modelId()
                 )
             ).getBytes(StandardCharsets.UTF_8)
         );

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

@@ -20,8 +20,7 @@ public record GoogleVertexAiRerankRequestEntity(
     List<String> inputs,
     @Nullable Boolean returnDocuments,
     @Nullable Integer topN,
-    @Nullable String model,
-    @Nullable Integer taskSettingsTopN
+    @Nullable String model
 ) implements ToXContentObject {
 
     private static final String MODEL_FIELD = "model";
@@ -66,8 +65,6 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
         // prefer the root level top_n over task settings
         if (topN != null) {
             builder.field(TOP_N_FIELD, topN);
-        } else if (taskSettingsTopN != null) {
-            builder.field(TOP_N_FIELD, taskSettingsTopN);
         }
 
         if (returnDocuments != null) {

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiRerankRequestEntityTests.java

@@ -21,7 +21,7 @@
 
 public class GoogleVertexAiRerankRequestEntityTests extends ESTestCase {
     public void testXContent_SingleRequest_WritesAllFieldsIfDefined() throws IOException {
-        var entity = new GoogleVertexAiRerankRequestEntity("query", List.of("abc"), Boolean.TRUE, 10, "model", 8);
+        var entity = new GoogleVertexAiRerankRequestEntity("query", List.of("abc"), Boolean.TRUE, 10, "model");
 
         XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
         entity.toXContent(builder, null);
@@ -44,7 +44,7 @@ public void testXContent_SingleRequest_WritesAllFieldsIfDefined() throws IOExcep
     }
 
     public void testXContent_SingleRequest_WritesMinimalFields() throws IOException {
-        var entity = new GoogleVertexAiRerankRequestEntity("query", List.of("abc"), null, null, null, null);
+        var entity = new GoogleVertexAiRerankRequestEntity("query", List.of("abc"), null, null, null);
 
         XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
         entity.toXContent(builder, null);
@@ -64,7 +64,7 @@ public void testXContent_SingleRequest_WritesMinimalFields() throws IOException
     }
 
     public void testXContent_MultipleRequests_WritesAllFieldsIfDefined() throws IOException {
-        var entity = new GoogleVertexAiRerankRequestEntity("query", List.of("abc", "def"), Boolean.FALSE, 12, "model", 8);
+        var entity = new GoogleVertexAiRerankRequestEntity("query", List.of("abc", "def"), Boolean.FALSE, 12, "model");
 
         XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
         entity.toXContent(builder, null);
@@ -91,7 +91,7 @@ public void testXContent_MultipleRequests_WritesAllFieldsIfDefined() throws IOEx
     }
 
     public void testXContent_MultipleRequests_WritesMinimalFields() throws IOException {
-        var entity = new GoogleVertexAiRerankRequestEntity("query", List.of("abc", "def"), null, null, null, null);
+        var entity = new GoogleVertexAiRerankRequestEntity("query", List.of("abc", "def"), null, null, null);
 
         XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
         entity.toXContent(builder, null);
@@ -113,27 +113,4 @@ public void testXContent_MultipleRequests_WritesMinimalFields() throws IOExcepti
             }
             """));
     }
-
-    public void testXContent_SingleRequest_UsesTaskSettingsTopNIfRootIsNotDefined() throws IOException {
-        var entity = new GoogleVertexAiRerankRequestEntity("query", List.of("abc"), null, null, "model", 8);
-
-        XContentBuilder builder = XContentFactory.contentBuilder(XContentType.JSON);
-        entity.toXContent(builder, null);
-        String xContentResult = Strings.toString(builder);
-
-        assertThat(xContentResult, equalToIgnoringWhitespaceInJsonString("""
-            {
-                "model": "model",
-                "query": "query",
-                "records": [
-                    {
-                        "id": "0",
-                        "content": "abc"
-                    }
-                ],
-                "topN": 8
-            }
-            """));
-    }
-
 }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/request/googlevertexai/GoogleVertexAiRerankRequestTests.java

@@ -53,8 +53,31 @@ public void testCreateRequest_WithTopNSet() throws IOException {
         var input = "input";
         var query = "query";
         var topN = 1;
+        var taskSettingsTopN = 3;
 
-        var request = createRequest(query, input, null, topN, null, null);
+        var request = createRequest(query, input, null, topN, null, taskSettingsTopN);
+        var httpRequest = request.createHttpRequest();
+
+        assertThat(httpRequest.httpRequestBase(), instanceOf(HttpPost.class));
+        var httpPost = (HttpPost) httpRequest.httpRequestBase();
+
+        assertThat(httpPost.getLastHeader(HttpHeaders.CONTENT_TYPE).getValue(), is(XContentType.JSON.mediaType()));
+        assertThat(httpPost.getLastHeader(HttpHeaders.AUTHORIZATION).getValue(), is(AUTH_HEADER_VALUE));
+
+        var requestMap = entityAsMap(httpPost.getEntity().getContent());
+
+        assertThat(requestMap, aMapWithSize(3));
+        assertThat(requestMap.get("records"), is(List.of(Map.of("id", "0", "content", input))));
+        assertThat(requestMap.get("query"), is(query));
+        assertThat(requestMap.get("topN"), is(topN));
+    }
+
+    public void testCreateRequest_UsesTaskSettingsTopNWhenRootLevelIsNull() throws IOException {
+        var input = "input";
+        var query = "query";
+        var topN = 1;
+
+        var request = createRequest(query, input, null, null, null, topN);
         var httpRequest = request.createHttpRequest();
 
         assertThat(httpRequest.httpRequestBase(), instanceOf(HttpPost.class));