Skip to content

Commit 3adeaf6

Browse files
committed
Add cohere v2 default post process function
Signed-off-by: zane-neo <zaniu@amazon.com>
1 parent 58dd965 commit 3adeaf6

File tree

3 files changed

+5
-2
lines changed

3 files changed

+5
-2
lines changed

common/src/main/java/org/opensearch/ml/common/connector/MLPostProcessFunction.java

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
public class MLPostProcessFunction {
1818

1919
public static final String COHERE_EMBEDDING = "connector.post_process.cohere.embedding";
20+
public static final String COHERE_V2_EMBEDDING = "connector.post_process.cohere_v2.embedding";
2021
public static final String OPENAI_EMBEDDING = "connector.post_process.openai.embedding";
2122
public static final String BEDROCK_EMBEDDING = "connector.post_process.bedrock.embedding";
2223
public static final String BEDROCK_BATCH_JOB_ARN = "connector.post_process.bedrock.batch_job_arn";
@@ -35,13 +36,15 @@ public class MLPostProcessFunction {
3536
CohereRerankPostProcessFunction cohereRerankPostProcessFunction = new CohereRerankPostProcessFunction();
3637
JSON_PATH_EXPRESSION.put(OPENAI_EMBEDDING, "$.data[*].embedding");
3738
JSON_PATH_EXPRESSION.put(COHERE_EMBEDDING, "$.embeddings");
39+
JSON_PATH_EXPRESSION.put(COHERE_V2_EMBEDDING, "$.embeddings.float");
3840
JSON_PATH_EXPRESSION.put(DEFAULT_EMBEDDING, "$[*]");
3941
JSON_PATH_EXPRESSION.put(BEDROCK_EMBEDDING, "$.embedding");
4042
JSON_PATH_EXPRESSION.put(BEDROCK_BATCH_JOB_ARN, "$");
4143
JSON_PATH_EXPRESSION.put(COHERE_RERANK, "$.results");
4244
JSON_PATH_EXPRESSION.put(DEFAULT_RERANK, "$[*]");
4345
POST_PROCESS_FUNCTIONS.put(OPENAI_EMBEDDING, embeddingPostProcessFunction);
4446
POST_PROCESS_FUNCTIONS.put(COHERE_EMBEDDING, embeddingPostProcessFunction);
47+
POST_PROCESS_FUNCTIONS.put(COHERE_V2_EMBEDDING, embeddingPostProcessFunction);
4548
POST_PROCESS_FUNCTIONS.put(DEFAULT_EMBEDDING, embeddingPostProcessFunction);
4649
POST_PROCESS_FUNCTIONS.put(BEDROCK_EMBEDDING, bedrockEmbeddingPostProcessFunction);
4750
POST_PROCESS_FUNCTIONS.put(BEDROCK_BATCH_JOB_ARN, batchJobArnPostProcessFunction);

plugin/src/test/java/org/opensearch/ml/rest/RestBedRockInferenceIT.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -321,7 +321,7 @@ public void test_bedrockEmbeddingTypeSupportedModel_withDifferentResponseFilters
321321
if (testCaseName.equals("response_filter_to_embedding_concrete_type")) {
322322
assertEquals(errorMsg, 1024, ((List) ((Map<?, ?>) outputList.get(0)).get("data")).size());
323323
} else {
324-
assertEquals(errorMsg, 1536, ((List) ((Map<?, ?>) outputList.get(0)).get("data")).size());
324+
assertEquals(errorMsg, 1024, ((List) ((Map<?, ?>) outputList.get(0)).get("data")).size());
325325
}
326326
}
327327
}

plugin/src/test/resources/org/opensearch/ml/rest/templates/BedRockEmbeddingTypeSupportedConnectorBodies.json

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -100,7 +100,7 @@
100100
"parameters": {
101101
"region": "%s",
102102
"service_name": "bedrock",
103-
"model_name": "amazon.titan-embed-text-v1"
103+
"model_name": "amazon.titan-embed-text-v2:0"
104104
},
105105
"credential": {
106106
"access_key": "%s",

0 commit comments

Comments
 (0)