Skip to content

Commit 35c5546

Browse files
martin-gaievskimartin-gaievski
committed
Changing model for Judgment entity
Signed-off-by: Martin Gaievski <[email protected]>
1 parent 5492261 commit 35c5546

File tree

18 files changed

+27786
-7119
lines changed

18 files changed

+27786
-7119
lines changed

‎src/main/java/org/opensearch/searchrelevance/dao/JudgmentDao.java

Lines changed: 2 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -11,8 +11,6 @@
1111

1212
import java.io.IOException;
1313

14-
import org.apache.logging.log4j.LogManager;
15-
import org.apache.logging.log4j.Logger;
1614
import org.opensearch.action.StepListener;
1715
import org.opensearch.action.delete.DeleteResponse;
1816
import org.opensearch.action.search.SearchResponse;
@@ -28,7 +26,6 @@
2826
import org.opensearch.searchrelevance.model.Judgment;
2927

3028
public class JudgmentDao {
31-
private static final Logger LOGGER = LogManager.getLogger(JudgmentDao.class);
3229
private final SearchRelevanceIndicesManager searchRelevanceIndicesManager;
3330

3431
@Inject
@@ -56,7 +53,7 @@ public void putJudgement(final Judgment judgment, final ActionListener listener)
5653
}
5754
try {
5855
searchRelevanceIndicesManager.putDoc(
59-
judgment.id(),
56+
judgment.getId(),
6057
judgment.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS),
6158
JUDGMENT,
6259
listener
@@ -73,7 +70,7 @@ public void updateJudgment(final Judgment judgment, final ActionListener listene
7370
}
7471
try {
7572
searchRelevanceIndicesManager.updateDoc(
76-
judgment.id(),
73+
judgment.getId(),
7774
judgment.toXContent(XContentFactory.jsonBuilder(), ToXContent.EMPTY_PARAMS),
7875
JUDGMENT,
7976
listener

‎src/main/java/org/opensearch/searchrelevance/judgments/BaseJudgmentsProcessor.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
*/
88
package org.opensearch.searchrelevance.judgments;
99

10+
import java.util.List;
1011
import java.util.Map;
1112

1213
import org.opensearch.core.action.ActionListener;
@@ -27,5 +28,5 @@ public interface BaseJudgmentsProcessor {
2728
* @param metadata used to generate judgment ratings for various judgment type
2829
* @param listener async action
2930
*/
30-
public void generateJudgmentRating(Map<String, Object> metadata, ActionListener<Map<String, Map<String, String>>> listener);
31+
public void generateJudgmentRating(Map<String, Object> metadata, ActionListener<List<Map<String, Object>>> listener);
3132
}

‎src/main/java/org/opensearch/searchrelevance/judgments/ImportJudgmentsProcessor.java

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
*/
88
package org.opensearch.searchrelevance.judgments;
99

10+
import java.util.ArrayList;
1011
import java.util.HashMap;
1112
import java.util.List;
1213
import java.util.Map;
@@ -35,18 +36,19 @@ public JudgmentType getJudgmentType() {
3536
}
3637

3738
@Override
38-
public void generateJudgmentRating(Map<String, Object> metadata, ActionListener<Map<String, Map<String, String>>> listener) {
39+
public void generateJudgmentRating(Map<String, Object> metadata, ActionListener<List<Map<String, Object>>> listener) {
3940

40-
Map<String, Object> sourceJudgementRatings = (Map<String, Object>) metadata.get("judgmentRatings");
41+
List<Map<String, Object>> sourceJudgementRatings = (List<Map<String, Object>>) metadata.get("judgmentRatings");
4142
metadata.remove("judgmentRatings");
4243

4344
// Create the result map in the expected format
44-
Map<String, Map<String, String>> formattedRatings = new HashMap<>();
45+
List<Map<String, Object>> formattedRatings = new ArrayList<>();
4546

4647
// Process each query
47-
for (Map.Entry<String, Object> queryEntry : sourceJudgementRatings.entrySet()) {
48-
String queryText = queryEntry.getKey();
49-
Object ratingData = queryEntry.getValue();
48+
// for (Map.Entry<String, Object> queryEntry : sourceJudgementRatings.entrySet()) {
49+
for (Map<String, Object> queryJudgment : sourceJudgementRatings) {
50+
String queryText = queryJudgment.get("query").toString();
51+
Object ratingData = queryJudgment.get("ratings");
5052

5153
if (!(ratingData instanceof List)) {
5254
listener.onFailure(
@@ -96,7 +98,14 @@ public void generateJudgmentRating(Map<String, Object> metadata, ActionListener<
9698
}
9799

98100
// Add the formatted ratings for this query
99-
formattedRatings.put(queryText, docRatings);
101+
Map<String, Object> queryRatings = new HashMap<>();
102+
queryRatings.put("query", queryText);
103+
List<Map<String, String>> docIdRatings = docRatings.entrySet()
104+
.stream()
105+
.map(entry -> Map.of("docId", entry.getKey(), "rating", entry.getValue()))
106+
.toList();
107+
queryRatings.put("ratings", docIdRatings);
108+
formattedRatings.add(queryRatings);
100109
}
101110

102111
listener.onResponse(formattedRatings);

‎src/main/java/org/opensearch/searchrelevance/judgments/LlmJudgmentsProcessor.java

Lines changed: 14 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,7 @@ public JudgmentType getJudgmentType() {
8686
}
8787

8888
@Override
89-
public void generateJudgmentRating(Map<String, Object> metadata, ActionListener<Map<String, Map<String, String>>> listener) {
89+
public void generateJudgmentRating(Map<String, Object> metadata, ActionListener<List<Map<String, Object>>> listener) {
9090
String querySetId = (String) metadata.get("querySetId");
9191
List<String> searchConfigurationList = (List<String>) metadata.get("searchConfigurationList");
9292
int size = (int) metadata.get("size");
@@ -129,12 +129,13 @@ private void generateLLMJudgments(
129129
List<String> contextFields,
130130
Map<String, Object> results,
131131
boolean ignoreFailure,
132-
ActionListener<Map<String, Map<String, String>>> listener
132+
ActionListener<List<Map<String, Object>>> listener
133133
) {
134134
Map<String, List<String>> indexAndQueries = (Map<String, List<String>>) results.get(METRICS_INDEX_AND_QUERIES_FIELD_NAME);
135135
List<String> queryTextWithReferences = (List<String>) results.get(METRICS_QUERY_TEXT_FIELD_NAME);
136136

137-
Map<String, Map<String, String>> allJudgments = new HashMap<>();
137+
// Map<String, Map<String, String>> allJudgments = new HashMap<>();
138+
List<Map<String, Object>> allJudgments = new ArrayList<>();
138139
AtomicInteger remainingQueries = new AtomicInteger(queryTextWithReferences.size());
139140

140141
for (String queryTextWithReference : queryTextWithReferences) {
@@ -146,11 +147,19 @@ private void generateLLMJudgments(
146147
indexAndQueries,
147148
queryTextWithReference,
148149
ignoreFailure,
149-
new ActionListener<Map<String, String>>() {
150+
new ActionListener<>() {
150151
@Override
151152
public void onResponse(Map<String, String> docIdToScore) {
152153
synchronized (allJudgments) {
153-
allJudgments.put(queryTextWithReference, docIdToScore);
154+
// allJudgments.put(queryTextWithReference, docIdToScore);
155+
Map<String, Object> judgmentForQuery = new HashMap<>();
156+
judgmentForQuery.put("query", queryTextWithReference);
157+
List<Map<String, String>> docIdRatings = docIdToScore.entrySet()
158+
.stream()
159+
.map(entry -> Map.of("docId", entry.getKey(), "rating", entry.getValue()))
160+
.toList();
161+
judgmentForQuery.put("ratings", docIdRatings);
162+
allJudgments.add(judgmentForQuery);
154163
}
155164
if (remainingQueries.decrementAndGet() == 0) {
156165
listener.onResponse(allJudgments);

‎src/main/java/org/opensearch/searchrelevance/judgments/UbiJudgmentsProcessor.java

Lines changed: 4 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
*/
88
package org.opensearch.searchrelevance.judgments;
99

10+
import java.util.List;
1011
import java.util.Map;
1112

1213
import org.apache.logging.log4j.LogManager;
@@ -36,7 +37,7 @@ public JudgmentType getJudgmentType() {
3637
}
3738

3839
@Override
39-
public void generateJudgmentRating(Map<String, Object> metadata, ActionListener<Map<String, Map<String, String>>> listener) {
40+
public void generateJudgmentRating(Map<String, Object> metadata, ActionListener<List<Map<String, Object>>> listener) {
4041
String clickModel = (String) metadata.get("clickModel");
4142
int maxRank = (int) metadata.get("maxRank");
4243

@@ -48,9 +49,9 @@ public void generateJudgmentRating(Map<String, Object> metadata, ActionListener<
4849
StepListener<Map<String, Map<String, String>>> clickModelStep = new StepListener<>();
4950

5051
try {
51-
coecClickModel.calculateJudgments(new ActionListener<Map<String, Map<String, String>>>() {
52+
coecClickModel.calculateJudgments(new ActionListener<>() {
5253
@Override
53-
public void onResponse(Map<String, Map<String, String>> judgments) {
54+
public void onResponse(List<Map<String, Object>> judgments) {
5455
listener.onResponse(judgments);
5556
}
5657

‎src/main/java/org/opensearch/searchrelevance/judgments/clickmodel/ClickModel.java

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
*/
88
package org.opensearch.searchrelevance.judgments.clickmodel;
99

10+
import java.util.List;
1011
import java.util.Map;
1112

1213
import org.opensearch.core.action.ActionListener;
@@ -19,5 +20,5 @@ public abstract class ClickModel {
1920
/**
2021
* Calculate implicit judgments.
2122
*/
22-
public abstract void calculateJudgments(ActionListener<Map<String, Map<String, String>>> listener);
23+
public abstract void calculateJudgments(ActionListener<List<Map<String, Object>>> listener);
2324
}

‎src/main/java/org/opensearch/searchrelevance/judgments/clickmodel/coec/CoecClickModel.java

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,10 @@
99

1010
import static org.opensearch.searchrelevance.common.PluginConstants.UBI_EVENTS_INDEX;
1111

12+
import java.util.ArrayList;
1213
import java.util.HashMap;
1314
import java.util.HashSet;
15+
import java.util.List;
1416
import java.util.Locale;
1517
import java.util.Map;
1618
import java.util.Set;
@@ -57,7 +59,7 @@ public CoecClickModel(final Client client, final CoecClickModelParameters parame
5759
}
5860

5961
@Override
60-
public void calculateJudgments(ActionListener<Map<String, Map<String, String>>> listener) {
62+
public void calculateJudgments(ActionListener<List<Map<String, Object>>> listener) {
6163
// Step 1: Calculate rank-aggregated click-through
6264
getRankAggregatedClickThrough(ActionListener.wrap(rankAggregatedClickThrough -> {
6365
// Step 2: Get clickthrough rates
@@ -662,10 +664,11 @@ private void getQueryCount(String userQuery, String objectId, int rank, ActionLi
662664
private void calculateCoecJudgments(
663665
Map<Integer, Double> rankAggregatedClickThrough,
664666
Map<String, Set<ClickthroughRate>> clickthroughRates,
665-
ActionListener<Map<String, Map<String, String>>> listener
667+
ActionListener<List<Map<String, Object>>> listener
666668
) {
667669
LOGGER.debug("Starting COEC calculation with rank CTR: {}", rankAggregatedClickThrough);
668-
Map<String, Map<String, String>> judgmentRatings = new HashMap<>();
670+
// Map<String, Map<String, String>> judgmentRatings = new HashMap<>();
671+
List<Map<String, Object>> judgmentRatings = new ArrayList<>();
669672

670673
for (Map.Entry<String, Set<ClickthroughRate>> entry : clickthroughRates.entrySet()) {
671674
String userQuery = entry.getKey();
@@ -691,14 +694,17 @@ private void calculateCoecJudgments(
691694
}
692695

693696
if (!docScores.isEmpty()) {
694-
judgmentRatings.put(userQuery, docScores);
697+
Map<String, Object> queryRating = new HashMap<>();
698+
queryRating.put("query", userQuery);
699+
queryRating.put("ratings", docScores);
700+
judgmentRatings.add(queryRating);
695701
}
696702
LOGGER.debug(
697703
"Final judgment ratings size - Queries: {}, Total Documents: {}",
698704
judgmentRatings.size(),
699-
judgmentRatings.values().stream().mapToInt(Map::size).sum()
705+
// judgmentRatings.values().stream().mapToInt(Map::size).sum()
706+
judgmentRatings.stream().mapToInt(item -> ((Map<String, Object>) item.get("ratings")).size()).sum()
700707
);
701-
702708
listener.onResponse(judgmentRatings);
703709
}
704710
}

‎src/main/java/org/opensearch/searchrelevance/model/Judgment.java

Lines changed: 25 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -8,11 +8,17 @@
88
package org.opensearch.searchrelevance.model;
99

1010
import java.io.IOException;
11+
import java.util.List;
1112
import java.util.Map;
1213

1314
import org.opensearch.core.xcontent.ToXContentObject;
1415
import org.opensearch.core.xcontent.XContentBuilder;
1516

17+
import lombok.AllArgsConstructor;
18+
import lombok.Getter;
19+
20+
@Getter
21+
@AllArgsConstructor
1622
public class Judgment implements ToXContentObject {
1723
public static final String ID = "id";
1824
public static final String TIME_STAMP = "timestamp";
@@ -25,31 +31,13 @@ public class Judgment implements ToXContentObject {
2531
/**
2632
* Identifier of the system index
2733
*/
28-
private String id;
29-
private String timestamp;
30-
private String name;
34+
private final String id;
35+
private final String timestamp;
36+
private final String name;
3137
private final AsyncStatus status;
32-
private JudgmentType type;
33-
private Map<String, Object> metadata;
34-
private Map<String, Map<String, String>> judgmentRatings;
35-
36-
public Judgment(
37-
String id,
38-
String timestamp,
39-
String name,
40-
AsyncStatus status,
41-
JudgmentType type,
42-
Map<String, Object> metadata,
43-
Map<String, Map<String, String>> judgmentRatings
44-
) {
45-
this.id = id;
46-
this.timestamp = timestamp;
47-
this.name = name;
48-
this.status = status;
49-
this.type = type;
50-
this.metadata = metadata;
51-
this.judgmentRatings = judgmentRatings;
52-
}
38+
private final JudgmentType type;
39+
private final Map<String, Object> metadata;
40+
private final List<Map<String, Object>> judgmentRatings;
5341

5442
@Override
5543
public XContentBuilder toXContent(XContentBuilder builder, Params params) throws IOException {
@@ -61,45 +49,23 @@ public XContentBuilder toXContent(XContentBuilder builder, Params params) throws
6149
xContentBuilder.field(TYPE, this.type.name().trim());
6250
xContentBuilder.field(METADATA, this.metadata);
6351
// Start judgmentRatings object
64-
xContentBuilder.startObject(JUDGMENT_RATINGS);
65-
for (Map.Entry<String, Map<String, String>> queryEntry : this.judgmentRatings.entrySet()) {
66-
xContentBuilder.startArray(queryEntry.getKey());
67-
for (Map.Entry<String, String> docEntry : queryEntry.getValue().entrySet()) {
68-
xContentBuilder.startObject().field("docId", docEntry.getKey()).field("rating", docEntry.getValue()).endObject();
52+
xContentBuilder.startArray(JUDGMENT_RATINGS);
53+
for (Map<String, Object> judgment : this.judgmentRatings) {
54+
xContentBuilder.startObject();
55+
xContentBuilder.field("query", judgment.get("query"));
56+
xContentBuilder.startArray("ratings");
57+
for (Map<String, Object> rating : (List<Map<String, Object>>) judgment.get("ratings")) {
58+
xContentBuilder.startObject();
59+
for (Map.Entry<String, Object> entry : rating.entrySet()) {
60+
xContentBuilder.field(entry.getKey(), entry.getValue());
61+
}
62+
xContentBuilder.endObject();
6963
}
7064
xContentBuilder.endArray();
65+
xContentBuilder.endObject();
7166
}
72-
xContentBuilder.endObject();
67+
xContentBuilder.endArray();
7368
// End judgmentRatings object
7469
return xContentBuilder.endObject();
7570
}
76-
77-
public String id() {
78-
return id;
79-
}
80-
81-
public String timestamp() {
82-
return timestamp;
83-
}
84-
85-
public String name() {
86-
return name;
87-
}
88-
89-
public AsyncStatus status() {
90-
return status;
91-
}
92-
93-
public JudgmentType type() {
94-
return type;
95-
}
96-
97-
public Map<String, Object> metadata() {
98-
return metadata;
99-
}
100-
101-
public Map<String, Map<String, String>> judgmentScores() {
102-
return judgmentRatings;
103-
}
104-
10571
}

‎src/main/java/org/opensearch/searchrelevance/rest/RestPutJudgmentAction.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -114,7 +114,7 @@ protected RestChannelConsumer prepareRequest(RestRequest request, NodeClient cli
114114
createRequest = new PutUbiJudgmentRequest(type, name, description, clickModel, maxRank);
115115
}
116116
case IMPORT_JUDGMENT -> {
117-
Map<String, Object> judgmentRatings = (Map<String, Object>) source.get("judgmentRatings");
117+
List<Map<String, Object>> judgmentRatings = (List<Map<String, Object>>) source.get("judgmentRatings");
118118
createRequest = new PutImportJudgmentRequest(type, name, description, judgmentRatings);
119119
}
120120
default -> {

0 commit comments

Comments
 (0)