add postprocess function for batch inferernce jobArn

Signed-off-by: Yaliang Wu <[email protected]>
ylwu-amzn · Sep 13, 2024 · 31f7634 · 31f7634
1 parent c1d56ea
commit 31f7634
Show file tree

Hide file tree

Showing 2 changed files with 52 additions and 0 deletions.
diff --git a/common/src/main/java/org/opensearch/ml/common/connector/MLPostProcessFunction.java b/common/src/main/java/org/opensearch/ml/common/connector/MLPostProcessFunction.java
@@ -10,6 +10,7 @@
 import java.util.Map;
 import java.util.function.Function;
 
+import org.opensearch.ml.common.connector.functions.postprocess.BedrockBatchJobArnPostProcessFunction;
 import org.opensearch.ml.common.connector.functions.postprocess.BedrockEmbeddingPostProcessFunction;
 import org.opensearch.ml.common.connector.functions.postprocess.CohereRerankPostProcessFunction;
 import org.opensearch.ml.common.connector.functions.postprocess.EmbeddingPostProcessFunction;
@@ -20,6 +21,7 @@ public class MLPostProcessFunction {
     public static final String COHERE_EMBEDDING = "connector.post_process.cohere.embedding";
     public static final String OPENAI_EMBEDDING = "connector.post_process.openai.embedding";
     public static final String BEDROCK_EMBEDDING = "connector.post_process.bedrock.embedding";
+    public static final String BEDROCK_BATCH_JOB_ARN = "connector.post_process.bedrock.batch_job_arn";
     public static final String COHERE_RERANK = "connector.post_process.cohere.rerank";
     public static final String DEFAULT_EMBEDDING = "connector.post_process.default.embedding";
     public static final String DEFAULT_RERANK = "connector.post_process.default.rerank";
@@ -31,17 +33,20 @@ public class MLPostProcessFunction {
     static {
         EmbeddingPostProcessFunction embeddingPostProcessFunction = new EmbeddingPostProcessFunction();
         BedrockEmbeddingPostProcessFunction bedrockEmbeddingPostProcessFunction = new BedrockEmbeddingPostProcessFunction();
+        BedrockBatchJobArnPostProcessFunction batchJobArnPostProcessFunction = new BedrockBatchJobArnPostProcessFunction();
         CohereRerankPostProcessFunction cohereRerankPostProcessFunction = new CohereRerankPostProcessFunction();
         JSON_PATH_EXPRESSION.put(OPENAI_EMBEDDING, "$.data[*].embedding");
         JSON_PATH_EXPRESSION.put(COHERE_EMBEDDING, "$.embeddings");
         JSON_PATH_EXPRESSION.put(DEFAULT_EMBEDDING, "$[*]");
         JSON_PATH_EXPRESSION.put(BEDROCK_EMBEDDING, "$.embedding");
+        JSON_PATH_EXPRESSION.put(BEDROCK_BATCH_JOB_ARN, "$");
         JSON_PATH_EXPRESSION.put(COHERE_RERANK, "$.results");
         JSON_PATH_EXPRESSION.put(DEFAULT_RERANK, "$[*]");
         POST_PROCESS_FUNCTIONS.put(OPENAI_EMBEDDING, embeddingPostProcessFunction);
         POST_PROCESS_FUNCTIONS.put(COHERE_EMBEDDING, embeddingPostProcessFunction);
         POST_PROCESS_FUNCTIONS.put(DEFAULT_EMBEDDING, embeddingPostProcessFunction);
         POST_PROCESS_FUNCTIONS.put(BEDROCK_EMBEDDING, bedrockEmbeddingPostProcessFunction);
+        POST_PROCESS_FUNCTIONS.put(BEDROCK_BATCH_JOB_ARN, batchJobArnPostProcessFunction);
         POST_PROCESS_FUNCTIONS.put(COHERE_RERANK, cohereRerankPostProcessFunction);
         POST_PROCESS_FUNCTIONS.put(DEFAULT_RERANK, cohereRerankPostProcessFunction);
     }

diff --git a/...arch/ml/common/connector/functions/postprocess/BedrockBatchJobArnPostProcessFunction.java b/...arch/ml/common/connector/functions/postprocess/BedrockBatchJobArnPostProcessFunction.java
@@ -0,0 +1,47 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.ml.common.connector.functions.postprocess;
+
+import org.opensearch.ml.common.output.model.ModelTensor;
+
+import java.util.ArrayList;
+import java.util.HashMap;
+import java.util.List;
+import java.util.Map;
+
+public class BedrockBatchJobArnPostProcessFunction extends ConnectorPostProcessFunction<Map<String, String>> {
+    public static final String JOB_ARN = "jobArn";
+    public static final String PROCESSED_JOB_ARN = "processedJobArn";
+
+    @Override
+    public void validate(Object input) {
+        if (!(input instanceof Map)) {
+            throw new IllegalArgumentException("Post process function input is not a Map.");
+        }
+        Map<String, String> jobInfo = (Map<String, String>) input;
+        if (!(jobInfo.containsKey(JOB_ARN))) {
+            throw new IllegalArgumentException("Bedrock batch job arn missing.");
+        }
+    }
+
+    @Override
+    public List<ModelTensor> process(Map<String, String> jobInfo) {
+        List<ModelTensor> modelTensors = new ArrayList<>();
+        Map<String, String> processedResult = new HashMap<>();
+        processedResult.putAll(jobInfo);
+        String jobArn = jobInfo.get(JOB_ARN);
+        processedResult.put(PROCESSED_JOB_ARN, jobArn.replace("/", "%2F"));
+        modelTensors
+            .add(
+                ModelTensor
+                    .builder()
+                    .name("response")
+                    .dataAsMap(processedResult)
+                    .build()
+            );
+        return modelTensors;
+    }
+}