Using record ID as index value when parsing Google Vertex AI rerank results

Author: ymao1

x-pack/plugin/inference/src/main/java/org/elasticsearch/xpack/inference/external/response/googlevertexai/GoogleVertexAiRerankResponseEntity.java

@@ -17,6 +17,8 @@
 import org.elasticsearch.xcontent.XContentType;
 import org.elasticsearch.xpack.core.inference.results.RankedDocsResults;
 import org.elasticsearch.xpack.inference.external.http.HttpResult;
+import org.elasticsearch.xpack.inference.external.request.Request;
+import org.elasticsearch.xpack.inference.external.request.googlevertexai.GoogleVertexAiRerankRequest;
 
 import java.io.IOException;
 import java.util.List;
@@ -109,14 +111,19 @@ private static List<RankedDocsResults.RankedDoc> doParse(XContentParser parser)
  throw new IllegalStateException(format(FAILED_TO_FIND_FIELD_TEMPLATE, RankedDoc.SCORE.getPreferredName()));
  }
 
- return new RankedDocsResults.RankedDoc(index, parsedRankedDoc.score, parsedRankedDoc.content);
+ if (parsedRankedDoc.id == null) {
+ throw new IllegalStateException(format(FAILED_TO_FIND_FIELD_TEMPLATE, RankedDoc.ID.getPreferredName()));
+ }
+
+ return new RankedDocsResults.RankedDoc(Integer.parseInt(parsedRankedDoc.id), parsedRankedDoc.score, parsedRankedDoc.content);
  });
  }
 
- private record RankedDoc(@Nullable Float score, @Nullable String content) {
+ private record RankedDoc(@Nullable Float score, @Nullable String content, @Nullable String id) {
 
  private static final ParseField CONTENT = new ParseField("content");
  private static final ParseField SCORE = new ParseField("score");
+ private static final ParseField ID = new ParseField("id");
  private static final ObjectParser<Builder, Void> PARSER = new ObjectParser<>(
  "google_vertex_ai_rerank_response",
  true,
@@ -126,6 +133,7 @@ private record RankedDoc(@Nullable Float score, @Nullable String content) {
  static {
  PARSER.declareString(Builder::setContent, CONTENT);
  PARSER.declareFloat(Builder::setScore, SCORE);
+ PARSER.declareString(Builder::setId, ID);
  }
 
  public static RankedDoc parse(XContentParser parser) {
@@ -137,6 +145,7 @@ private static final class Builder {
 
  private String content;
  private Float score;
+ private String id;
 
  private Builder() {}
 
@@ -150,8 +159,13 @@ public Builder setContent(String content) {
  return this;
  }
 
+ public Builder setId(String id) {
+ this.id = id;
+ return this;
+ }
+
  public RankedDoc build() {
- return new RankedDoc(score, content);
+ return new RankedDoc(score, content, id);
  }
  }
  }

x-pack/plugin/inference/src/test/java/org/elasticsearch/xpack/inference/external/response/googlevertexai/GoogleVertexAiRerankResponseEntityTests.java

@@ -39,7 +39,7 @@ public void testFromResponse_CreatesResultsForASingleItem() throws IOException {
  new HttpResult(mock(HttpResponse.class), responseJson.getBytes(StandardCharsets.UTF_8))
  );
 
- assertThat(parsedResults.getRankedDocs(), is(List.of(new RankedDocsResults.RankedDoc(0, 0.97F, "content 2"))));
+ assertThat(parsedResults.getRankedDocs(), is(List.of(new RankedDocsResults.RankedDoc(2, 0.97F, "content 2"))));
  }
 
  public void testFromResponse_CreatesResultsForMultipleItems() throws IOException {
@@ -68,7 +68,7 @@ public void testFromResponse_CreatesResultsForMultipleItems() throws IOException
 
  assertThat(
  parsedResults.getRankedDocs(),
- is(List.of(new RankedDocsResults.RankedDoc(0, 0.97F, "content 2"), new RankedDocsResults.RankedDoc(1, 0.90F, "content 1")))
+ is(List.of(new RankedDocsResults.RankedDoc(2, 0.97F, "content 2"), new RankedDocsResults.RankedDoc(1, 0.90F, "content 1")))
  );
  }