Fix Cohere test. (#2831)

Signed-off-by: Austin Lee <[email protected]>
opensearch-project · Aug 15, 2024 · 659a836 · 659a836
1 parent f64e3f3
commit 659a836
Showing 1 changed file with 4 additions and 6 deletions.
diff --git a/plugin/src/test/java/org/opensearch/ml/rest/RestMLRAGSearchProcessorIT.java b/plugin/src/test/java/org/opensearch/ml/rest/RestMLRAGSearchProcessorIT.java
@@ -149,15 +149,15 @@ public class RestMLRAGSearchProcessorIT extends RestMLRemoteInferenceIT {
         ? BEDROCK_CONNECTOR_BLUEPRINT2
         : BEDROCK_CONNECTOR_BLUEPRINT1;
 
-    private static final String COHERE_API_KEY = System.getenv("COHERE_API_KEY");
+    private static final String COHERE_KEY = System.getenv("COHERE_KEY");
     private static final String COHERE_CONNECTOR_BLUEPRINT = "{\n"
         + "    \"name\": \"Cohere Chat Model\",\n"
         + "    \"description\": \"The connector to Cohere's public chat API\",\n"
         + "    \"version\": \"1\",\n"
         + "    \"protocol\": \"http\",\n"
         + "    \"credential\": {\n"
         + "        \"cohere_key\": \""
-        + COHERE_API_KEY
+        + COHERE_KEY
         + "\"\n"
         + "    },\n"
         + "    \"parameters\": {\n"
@@ -521,7 +521,7 @@ public void testBM25WithBedrockWithConversation() throws Exception {
 
     public void testBM25WithCohere() throws Exception {
         // Skip test if key is null
-        if (COHERE_API_KEY == null) {
+        if (COHERE_KEY == null) {
             return;
         }
         Response response = createConnector(COHERE_CONNECTOR_BLUEPRINT);
@@ -573,7 +573,7 @@ public void testBM25WithCohere() throws Exception {
 
     public void testBM25WithCohereUsingLlmResponseField() throws Exception {
         // Skip test if key is null
-        if (COHERE_API_KEY == null) {
+        if (COHERE_KEY == null) {
             return;
         }
         Response response = createConnector(COHERE_CONNECTOR_BLUEPRINT);
@@ -660,8 +660,6 @@ private Response performSearch(String indexName, String pipeline, int size, Sear
                     requestParameters.match,
                     requestParameters.llmModel,
                     requestParameters.llmQuestion,
-                    requestParameters.systemPrompt,
-                    requestParameters.userInstructions,
                     requestParameters.contextSize,
                     requestParameters.interactionSize,
                     requestParameters.timeout,