Skip to content

Commit 0595813

Browse files
add IT for flow agent with CatIndexTool (opensearch-project#2425) (opensearch-project#2447)
Signed-off-by: Jing Zhang <jngz@amazon.com> (cherry picked from commit d74c623) Co-authored-by: Jing Zhang <jngz@amazon.com>
1 parent 5d083c8 commit 0595813

File tree

1 file changed

+74
-0
lines changed

1 file changed

+74
-0
lines changed
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.ml.rest;
7+
8+
import java.io.IOException;
9+
import java.util.List;
10+
import java.util.Map;
11+
12+
import org.opensearch.client.Response;
13+
import org.opensearch.ml.utils.TestHelper;
14+
15+
public class RestMLFlowAgentIT extends MLCommonsRestTestCase {
16+
17+
public void testAgentCatIndexTool() throws IOException {
18+
// Register agent with CatIndexTool.
19+
Response response = registerAgentWithCatIndexTool();
20+
Map responseMap = parseResponseToMap(response);
21+
String agentId = (String) responseMap.get("agent_id");
22+
assertNotNull(agentId);
23+
assertEquals(20, agentId.length());
24+
25+
// Execute agent.
26+
response = executeAgentCatIndexTool(agentId);
27+
responseMap = parseResponseToMap(response);
28+
List responseList = (List) responseMap.get("inference_results");
29+
responseMap = (Map) responseList.get(0);
30+
responseList = (List) responseMap.get("output");
31+
responseMap = (Map) responseList.get(0);
32+
assertEquals("response", responseMap.get("name"));
33+
String result = (String) responseMap.get("result");
34+
assertNotNull(result);
35+
assertTrue(result.contains(".plugins-ml-agent"));
36+
}
37+
38+
public static Response registerAgentWithCatIndexTool() throws IOException {
39+
String registerAgentEntity = "{\n"
40+
+ " \"name\": \"Test_Agent_For_CatIndex_tool\",\n"
41+
+ " \"type\": \"flow\",\n"
42+
+ " \"description\": \"this is a test agent for the CatIndexTool\",\n"
43+
+ " \"tools\": [\n"
44+
+ " {\n"
45+
+ " \"type\": \"CatIndexTool\",\n"
46+
+ " \"name\": \"DemoCatIndexTool\",\n"
47+
+ " \"parameters\": {\n"
48+
+ " \"input\": \"${parameters.question}\"\n"
49+
+ " }\n"
50+
+ " }\n"
51+
+ " ]\n"
52+
+ "}";
53+
return TestHelper
54+
.makeRequest(client(), "POST", "/_plugins/_ml/agents/_register", null, TestHelper.toHttpEntity(registerAgentEntity), null);
55+
}
56+
57+
public static Response executeAgentCatIndexTool(String agentId) throws IOException {
58+
String question = "How many indices do I have?";
59+
return executeAgent(agentId, question);
60+
}
61+
62+
public static Response executeAgent(String agentId, String question) throws IOException {
63+
String executeAgentEntity = "{\n" + " \"parameters\": {\n" + " \"question\": \"" + question + " \"\n" + " }\n" + "}";
64+
return TestHelper
65+
.makeRequest(
66+
client(),
67+
"POST",
68+
String.format("/_plugins/_ml/agents/%s/_execute", agentId),
69+
null,
70+
TestHelper.toHttpEntity(executeAgentEntity),
71+
null
72+
);
73+
}
74+
}

0 commit comments

Comments
 (0)