From e6da44dd527458ac0952060d9a36868a398e74c8 Mon Sep 17 00:00:00 2001
From: Kathleen DeRusso <kathleen.derusso@elastic.co>
Date: Tue, 29 Oct 2024 13:57:09 -0400
Subject: [PATCH] Read queries from templates specified in request

---
 .../rankeval/TransportRankEvalAction.java     | 36 +++++++++----------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/modules/rank-eval/src/main/java/org/elasticsearch/index/rankeval/TransportRankEvalAction.java b/modules/rank-eval/src/main/java/org/elasticsearch/index/rankeval/TransportRankEvalAction.java
index a68e8c82e7d0..404b6ce6d8a6 100644
--- a/modules/rank-eval/src/main/java/org/elasticsearch/index/rankeval/TransportRankEvalAction.java
+++ b/modules/rank-eval/src/main/java/org/elasticsearch/index/rankeval/TransportRankEvalAction.java
@@ -30,7 +30,6 @@
 import org.elasticsearch.common.xcontent.LoggingDeprecationHandler;
 import org.elasticsearch.features.FeatureService;
 import org.elasticsearch.features.NodeFeature;
-import org.elasticsearch.index.query.QueryStringQueryBuilder;
 import org.elasticsearch.index.query.TermQueryBuilder;
 import org.elasticsearch.injection.guice.Inject;
 import org.elasticsearch.script.Script;
@@ -101,21 +100,25 @@ protected void doExecute(Task task, RankEvalRequest request, ActionListener<Rank
 
         if (evaluationSpecification.hasStoredCorpus()) {
 
-            loadRequestsFromStoredCorpus(evaluationSpecification.getStoredCorpus(), new ActionListener<>() {
-                @Override
-                public void onResponse(List<RatedRequest> ratedRequests) {
-                    if (ratedRequests.isEmpty()) {
-                        listener.onFailure(new IllegalArgumentException("No rated requests found in stored corpus"));
-                        return;
+            loadRequestsFromStoredCorpus(
+                evaluationSpecification.getStoredCorpus(),
+                evaluationSpecification.getTemplates().keySet().iterator().next(),
+                new ActionListener<>() {
+                    @Override
+                    public void onResponse(List<RatedRequest> ratedRequests) {
+                        if (ratedRequests.isEmpty()) {
+                            listener.onFailure(new IllegalArgumentException("No rated requests found in stored corpus"));
+                            return;
+                        }
+                        processRatedRequests(request, evaluationSpecification, List.copyOf(ratedRequests), listener);
                     }
-                    processRatedRequests(request, evaluationSpecification, List.copyOf(ratedRequests), listener);
-                }
 
-                @Override
-                public void onFailure(Exception e) {
-                    listener.onFailure(e);
+                    @Override
+                    public void onFailure(Exception e) {
+                        listener.onFailure(e);
+                    }
                 }
-            });
+            );
         } else {
             processRatedRequests(request, evaluationSpecification, evaluationSpecification.getRatedRequests(), listener);
         }
@@ -193,7 +196,7 @@ private void processRatedRequests(
         );
     }
 
-    private void loadRequestsFromStoredCorpus(String storedCorpusId, ActionListener<List<RatedRequest>> actionListener) {
+    private void loadRequestsFromStoredCorpus(String storedCorpusId, String templateId, ActionListener<List<RatedRequest>> actionListener) {
         // load the stored corpus from the index
         client.get(new GetRequest("search_relevance_dataset", storedCorpusId), new ActionListener<>() {
             @Override
@@ -234,10 +237,7 @@ public void onResponse(SearchResponse response) {
                             RatedDocument ratedDocument = new RatedDocument(index, documentId, score);
                             ratedDocuments.add(ratedDocument);
 
-                            // TODO remove placeholder here
-                            SearchSourceBuilder query = new SearchSourceBuilder().query(new QueryStringQueryBuilder(queryString));
-
-                            ratedRequests.add(new RatedRequest(queryId, ratedDocuments, query));
+                            ratedRequests.add(new RatedRequest(queryId, ratedDocuments, Map.of("query_string", queryString), templateId));
                         }
                         actionListener.onResponse(ratedRequests);
                     }