Fix broken tests.

Signed-off-by: Austin Lee <[email protected]>
austintlee · Dec 20, 2023 · 27d1cab · 27d1cab
1 parent b73d381
commit 27d1cab
Showing 1 changed file with 12 additions and 0 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
@@ -171,6 +171,8 @@ public class RestMLRAGSearchProcessorIT extends RestMLRemoteInferenceIT {
         + "      \"generative_qa_parameters\": {\n"
         + "        \"llm_model\": \"%s\",\n"
         + "        \"llm_question\": \"%s\",\n"
+        + "        \"system_prompt\": \"%s\",\n"
+        + "        \"user_instructions\": \"%s\",\n"
         + "        \"context_size\": %d,\n"
         + "        \"interaction_size\": %d,\n"
         + "        \"timeout\": %d\n"
@@ -188,6 +190,8 @@ public class RestMLRAGSearchProcessorIT extends RestMLRemoteInferenceIT {
         + "        \"llm_model\": \"%s\",\n"
         + "        \"llm_question\": \"%s\",\n"
         + "        \"conversation_id\": \"%s\",\n"
+        + "        \"system_prompt\": \"%s\",\n"
+        + "        \"user_instructions\": \"%s\",\n"
         + "        \"context_size\": %d,\n"
         + "        \"interaction_size\": %d,\n"
         + "        \"timeout\": %d\n"
@@ -306,6 +310,8 @@ public void testBM25WithOpenAI() throws Exception {
         requestParameters.match = "president";
         requestParameters.llmModel = OPENAI_MODEL;
         requestParameters.llmQuestion = "who is lincoln";
+        requestParameters.systemPrompt = "You are great at answering questions";
+        requestParameters.userInstructions = "Follow my instructions as best you can";
         requestParameters.contextSize = 5;
         requestParameters.interactionSize = 5;
         requestParameters.timeout = 60;
@@ -525,6 +531,8 @@ 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
@@ -539,6 +547,8 @@ private Response performSearch(String indexName, String pipeline, int size, Sear
                     requestParameters.llmModel,
                     requestParameters.llmQuestion,
                     requestParameters.conversationId,
+                    requestParameters.systemPrompt,
+                    requestParameters.userInstructions,
                     requestParameters.contextSize,
                     requestParameters.interactionSize,
                     requestParameters.timeout
@@ -579,6 +589,8 @@ static class SearchRequestParameters {
         String match;
         String llmModel;
         String llmQuestion;
+        String systemPrompt;
+        String userInstructions;
         int contextSize;
         int interactionSize;
         int timeout;