Add run ID to response

kderusso · Oct 30, 2024 · 7278d2a · 7278d2a
1 parent 58cbb6b
commit 7278d2a
Show file tree

Hide file tree

Showing 2 changed files with 23 additions and 3 deletions.
diff --git a/modules/rank-eval/src/main/java/org/elasticsearch/index/rankeval/RankEvalResponse.java b/modules/rank-eval/src/main/java/org/elasticsearch/index/rankeval/RankEvalResponse.java
@@ -10,6 +10,7 @@
 package org.elasticsearch.index.rankeval;
 
 import org.elasticsearch.ElasticsearchException;
+import org.elasticsearch.TransportVersions;
 import org.elasticsearch.action.ActionResponse;
 import org.elasticsearch.common.Strings;
 import org.elasticsearch.common.io.stream.StreamInput;
@@ -30,14 +31,22 @@
  **/
 public class RankEvalResponse extends ActionResponse implements ToXContentObject {
 
+    /** A GUID run ID used to pull history. */
+    private final String runId;
     /** The overall evaluation result. */
     private final double metricScore;
     /** details about individual ranking evaluation queries, keyed by their id */
     private final Map<String, EvalQueryQuality> details;
     /** exceptions for specific ranking evaluation queries, keyed by their id */
     private final Map<String, Exception> failures;
 
-    public RankEvalResponse(double metricScore, Map<String, EvalQueryQuality> partialResults, Map<String, Exception> failures) {
+    public RankEvalResponse(
+        String runId,
+        double metricScore,
+        Map<String, EvalQueryQuality> partialResults,
+        Map<String, Exception> failures
+    ) {
+        this.runId = runId;
         this.metricScore = metricScore;
         this.details = new HashMap<>(partialResults);
         this.failures = new HashMap<>(failures);
@@ -59,6 +68,11 @@ public RankEvalResponse(double metricScore, Map<String, EvalQueryQuality> partia
             String queryId = in.readString();
             this.failures.put(queryId, in.readException());
         }
+        if (in.getTransportVersion().onOrAfter(TransportVersions.RANK_EVAL_STORED_CORPORA)) {
+            this.runId = in.readOptionalString();
+        } else {
+            this.runId = null;
+        }
     }
 
     public double getMetricScore() {
@@ -83,11 +97,17 @@ public void writeTo(StreamOutput out) throws IOException {
         out.writeDouble(metricScore);
         out.writeMap(details, StreamOutput::writeWriteable);
         out.writeMap(failures, StreamOutput::writeException);
+        if (out.getTransportVersion().onOrAfter(TransportVersions.RANK_EVAL_STORED_CORPORA)) {
+            out.writeOptionalString(runId);
+        }
     }
 
     @Override
     public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
         builder.startObject();
+        if (runId != null) {
+            builder.field("run_id", runId);
+        }
         builder.field("metric_score", metricScore);
         builder.startObject("details");
         for (String key : details.keySet()) {

diff --git a/...les/rank-eval/src/main/java/org/elasticsearch/index/rankeval/TransportRankEvalAction.java b/...les/rank-eval/src/main/java/org/elasticsearch/index/rankeval/TransportRankEvalAction.java
@@ -303,7 +303,7 @@ public void onResponse(MultiSearchResponse multiSearchResponse) {
             }
 
             // Record history here
-            String runId = UUID.randomUUID().toString(); // TODO return this in the response too
+            String runId = UUID.randomUUID().toString();
             double metricScore = this.rankEvalSpec.getMetric().combine(responseDetails.values());
 
             HistoricalRankEvalRun.Metric historicalMetric = new HistoricalRankEvalRun.Metric(
@@ -334,7 +334,7 @@ public void onResponse(MultiSearchResponse multiSearchResponse) {
                 LogManager.getLogger(TransportRankEvalAction.class).error("Failed to store historical run", e);
             }
 
-            delegate.onResponse(new RankEvalResponse(metricScore, responseDetails, this.errors));
+            delegate.onResponse(new RankEvalResponse(runId, metricScore, responseDetails, this.errors));
         }
     }
 }