Add pre and post process functions for Bedrock Rerank API opensearch-…

…project#3254 Signed-off-by: tkykenmt <[email protected]>
tkykenmt · Jan 7, 2025 · 5a90037 · 5a90037
1 parent fd7ba2a
commit 5a90037
Showing 1 changed file with 9 additions and 2 deletions.
diff --git a/...pensearch/ml/common/connector/functions/postprocess/BedrockRerankPostProcessFunction.java b/...pensearch/ml/common/connector/functions/postprocess/BedrockRerankPostProcessFunction.java
@@ -28,7 +28,7 @@ public void validate(Object input) {
             Map innerMap = (Map) outerList.get(0);
 
             if (innerMap.isEmpty() || !innerMap.containsKey("index") || !innerMap.containsKey("relevanceScore")) {
-                throw new IllegalArgumentException("The rerank result should contain index and relevance_score.");
+                throw new IllegalArgumentException("The rerank result should contain index and relevanceScore.");
             }
         }
     }
@@ -41,7 +41,14 @@ public List<ModelTensor> process(List<Map<String, Object>> rerankResults) {
             Double[] scores = new Double[rerankResults.size()];
             for (int i = 0; i < rerankResults.size(); i++) {
                 Integer index = (Integer) rerankResults.get(i).get("index");
-                scores[index] = (Double) rerankResults.get(i).get("relevanceScore");
+                Object relevanceScore = rerankResults.get(i).get("relevanceScore");
+                scores[index] = switch (relevanceScore) {
+                    case BigDecimal bd -> bd.doubleValue();
+                    case Double d -> d;
+                    case null -> throw new IllegalArgumentException("relevanceScore is null");
+                    default -> throw new IllegalArgumentException("Unexpected type for relevanceScore: " +
+                            relevanceScore.getClass().getName());
+                };
             }
 
             for (int i = 0; i < scores.length; i++) {