Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add IT flow agent with search index tool #2448

Merged
merged 1 commit into from
May 17, 2024
Merged

add IT flow agent with search index tool #2448

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
76 changes: 71 additions & 5 deletions plugin/src/test/java/org/opensearch/ml/rest/RestMLFlowAgentIT.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,26 @@
import java.util.List;
import java.util.Map;

import org.apache.hc.core5.http.ParseException;
import org.junit.After;
import org.junit.Before;
import org.opensearch.client.Response;
import org.opensearch.ml.utils.TestHelper;

public class RestMLFlowAgentIT extends MLCommonsRestTestCase {

private String irisIndex = "iris_data";

@Before
public void setup() throws IOException, ParseException {
ingestIrisData(irisIndex);
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Is it possible to merge this IT with RestMLTrainAndPredictIT, to save time in the ingesting irisIndex before and delete it after?

Copy link
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I prefer to have it in agent IT, as it is for agent framework testing. And it can be reused for other agent test cases.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM for this test. But I think we can delete RestMLTrainAndPredictIT, and move the tests with new names inside this agentIT to save IT time.

Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In my opinion we still need to separate the ITs by its category unless there is obvious concern due to maintainability. Consuming same index doesn't necessarily mean they need to be merged together.

}

@After
public void deleteIndices() throws IOException {
deleteIndexWithAdminClient(irisIndex);
}

public void testAgentCatIndexTool() throws IOException {
// Register agent with CatIndexTool.
Response response = registerAgentWithCatIndexTool();
Expand All @@ -35,6 +50,27 @@ public void testAgentCatIndexTool() throws IOException {
assertTrue(result.contains(".plugins-ml-agent"));
}

public void testAgentSearchIndexTool() throws IOException {
// Register agent with SearchIndexTool.
Response response = registerAgentWithSearchIndexTool();
Map responseMap = parseResponseToMap(response);
String agentId = (String) responseMap.get("agent_id");
assertNotNull(agentId);
assertEquals(20, agentId.length());

// Execute agent.
response = executeAgentSearchIndexTool(agentId);
responseMap = parseResponseToMap(response);
List responseList = (List) responseMap.get("inference_results");
responseMap = (Map) responseList.get(0);
responseList = (List) responseMap.get("output");
responseMap = (Map) responseList.get(0);
assertEquals("response", responseMap.get("name"));
String result = (String) responseMap.get("result");
assertNotNull(result);
assertTrue(result.contains("\"_source\":{\"petal_length_in_cm\""));
}

public static Response registerAgentWithCatIndexTool() throws IOException {
String registerAgentEntity = "{\n"
+ " \"name\": \"Test_Agent_For_CatIndex_tool\",\n"
Expand All @@ -54,20 +90,50 @@ public static Response registerAgentWithCatIndexTool() throws IOException {
.makeRequest(client(), "POST", "/_plugins/_ml/agents/_register", null, TestHelper.toHttpEntity(registerAgentEntity), null);
}

public static Response registerAgentWithSearchIndexTool() throws IOException {
String registerAgentEntity = "{\n"
+ " \"name\": \"Test_Agent_For_SearchIndex_tool\",\n"
+ " \"type\": \"flow\",\n"
+ " \"description\": \"this is a test agent for the SearchIndexTool\",\n"
+ " \"tools\": [\n"
+ " {\n"
+ " \"type\": \"SearchIndexTool\""
+ " }\n"
+ " ]\n"
+ "}";
return TestHelper
.makeRequest(client(), "POST", "/_plugins/_ml/agents/_register", null, TestHelper.toHttpEntity(registerAgentEntity), null);
}

public static Response executeAgentCatIndexTool(String agentId) throws IOException {
String question = "How many indices do I have?";
return executeAgent(agentId, question);
String question = "\"How many indices do I have?\"";
return executeAgent(agentId, Map.of("question", question));
}

public static Response executeAgent(String agentId, String question) throws IOException {
String executeAgentEntity = "{\n" + " \"parameters\": {\n" + " \"question\": \"" + question + " \"\n" + " }\n" + "}";
public static Response executeAgentSearchIndexTool(String agentId) throws IOException {
String input = "{\"index\": \"iris_data\", \"query\": {\"size\": 2, \"_source\": \"petal_length_in_cm\"}}";
return executeAgent(agentId, Map.of("input", input));
}

public static Response executeAgent(String agentId, Map<String, String> args) throws IOException {
if (args == null || args.isEmpty()) {
return null;
}

// Construct parameters.
StringBuilder entityBuilder = new StringBuilder("{\"parameters\":{");
for (Map.Entry entry : args.entrySet()) {
entityBuilder.append('"').append(entry.getKey()).append("\":").append(entry.getValue()).append(',');
}
entityBuilder.replace(entityBuilder.length() - 1, entityBuilder.length(), "}}");

return TestHelper
.makeRequest(
client(),
"POST",
String.format("/_plugins/_ml/agents/%s/_execute", agentId),
null,
TestHelper.toHttpEntity(executeAgentEntity),
TestHelper.toHttpEntity(entityBuilder.toString()),
null
);
}
Expand Down
Loading