[Backport 2.x] Fixing register agent tests (#596)

Fixing register agent tests (#595) * Fixing register agent tests * fixing register agent IT * Fixing agent type --------- (cherry picked from commit ebcc8f7) Signed-off-by: Joshua Palis <[email protected]> Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com> Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
opensearch-project · Mar 20, 2024 · 31edb1b · 31edb1b
1 parent d03ca83
commit 31edb1b
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 3 deletions.
diff --git a/src/test/java/org/opensearch/flowframework/workflow/RegisterAgentTests.java b/src/test/java/org/opensearch/flowframework/workflow/RegisterAgentTests.java
@@ -16,6 +16,7 @@
 import org.opensearch.flowframework.exception.FlowFrameworkException;
 import org.opensearch.flowframework.indices.FlowFrameworkIndicesHandler;
 import org.opensearch.ml.client.MachineLearningNodeClient;
+import org.opensearch.ml.common.MLAgentType;
 import org.opensearch.ml.common.agent.LLMSpec;
 import org.opensearch.ml.common.agent.MLAgent;
 import org.opensearch.ml.common.agent.MLMemorySpec;
@@ -69,7 +70,7 @@ public void setUp() throws Exception {
             Map.ofEntries(
                 Map.entry("name", "test"),
                 Map.entry("description", "description"),
-                Map.entry("type", "type"),
+                Map.entry("type", MLAgentType.FLOW.name()),
                 Map.entry("llm.model_id", "xyz"),
                 Map.entry("llm.parameters", Collections.emptyMap()),
                 Map.entry("tools", tools),

diff --git a/src/test/resources/template/agent-framework.json b/src/test/resources/template/agent-framework.json
@@ -96,7 +96,7 @@
           },
           "user_inputs": {
             "name": "Sub Agent",
-            "type": "cot",
+            "type": "CONVERSATIONAL",
             "description": "this is a test agent",
             "parameters": {
               "hello": "world"
@@ -139,7 +139,7 @@
           },
           "user_inputs": {
             "name": "DEMO-Test_Agent_For_CoT",
-            "type": "cot",
+            "type": "CONVERSATIONAL",
             "description": "this is a test agent",
             "parameters": {
               "hello": "world"