Skip to content

Commit 161068b

Browse files
Hailong-amHailong-am
committed
painlessScript tool poc
Signed-off-by: Hailong Cui <[email protected]>
1 parent 3f04902 commit 161068b

File tree

5 files changed

+264
-1
lines changed

5 files changed

+264
-1
lines changed

‎src/main/java/org/opensearch/agent/ToolPlugin.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import org.opensearch.agent.tools.CreateAnomalyDetectorTool;
1515
import org.opensearch.agent.tools.NeuralSparseSearchTool;
1616
import org.opensearch.agent.tools.PPLTool;
17+
import org.opensearch.agent.tools.PainlessTool;
1718
import org.opensearch.agent.tools.RAGTool;
1819
import org.opensearch.agent.tools.SearchAlertsTool;
1920
import org.opensearch.agent.tools.SearchAnomalyDetectorsTool;
@@ -71,6 +72,7 @@ public Collection<Object> createComponents(
7172
SearchMonitorsTool.Factory.getInstance().init(client);
7273
CreateAlertTool.Factory.getInstance().init(client);
7374
CreateAnomalyDetectorTool.Factory.getInstance().init(client);
75+
PainlessTool.Factory.getInstance().init(scriptService);
7476
return Collections.emptyList();
7577
}
7678

@@ -87,7 +89,8 @@ public List<Tool.Factory<? extends Tool>> getToolFactories() {
8789
SearchAnomalyResultsTool.Factory.getInstance(),
8890
SearchMonitorsTool.Factory.getInstance(),
8991
CreateAlertTool.Factory.getInstance(),
90-
CreateAnomalyDetectorTool.Factory.getInstance()
92+
CreateAnomalyDetectorTool.Factory.getInstance(),
93+
PainlessTool.Factory.getInstance()
9194
);
9295
}
9396

‎src/main/java/org/opensearch/agent/tools/PainlessTool.java‎

Lines changed: 150 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,150 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.agent.tools;
7+
8+
import java.util.Collections;
9+
import java.util.HashMap;
10+
import java.util.Map;
11+
12+
import org.opensearch.core.action.ActionListener;
13+
import org.opensearch.ml.common.spi.tools.Tool;
14+
import org.opensearch.ml.common.spi.tools.ToolAnnotation;
15+
import org.opensearch.script.Script;
16+
import org.opensearch.script.ScriptService;
17+
import org.opensearch.script.ScriptType;
18+
import org.opensearch.script.TemplateScript;
19+
20+
import com.google.gson.Gson;
21+
22+
import lombok.Getter;
23+
import lombok.Setter;
24+
import lombok.extern.log4j.Log4j2;
25+
26+
/**
27+
* use case for this tool will only focus on flow agent
28+
*/
29+
@Log4j2
30+
@Setter
31+
@Getter
32+
@ToolAnnotation(PainlessTool.TYPE)
33+
public class PainlessTool implements Tool {
34+
public static final String TYPE = "PainlessTool";
35+
private static final String DEFAULT_DESCRIPTION = "Use this tool to execute painless script";
36+
37+
@Setter
38+
@Getter
39+
private String name = TYPE;
40+
41+
@Getter
42+
private String type = TYPE;
43+
44+
@Getter
45+
@Setter
46+
private String description = DEFAULT_DESCRIPTION;
47+
48+
@Getter
49+
private String version;
50+
51+
private ScriptService scriptService;
52+
@Setter
53+
private String scriptCode;
54+
55+
public PainlessTool(ScriptService scriptEngine, String script) {
56+
this.scriptService = scriptEngine;
57+
this.scriptCode = script;
58+
}
59+
60+
private Gson gson = new Gson();
61+
62+
@Override
63+
public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
64+
Script script = new Script(ScriptType.INLINE, "painless", scriptCode, Collections.emptyMap());
65+
Map<String, Object> flattenedParameters = new HashMap<>();
66+
for (Map.Entry<String, String> entry : parameters.entrySet()) {
67+
// keep original values and flatten
68+
flattenedParameters.put(entry.getKey(), entry.getValue());
69+
// TODO default is json parser. we may support format
70+
try {
71+
String value = org.apache.commons.text.StringEscapeUtils.unescapeJson(entry.getValue());
72+
Map<String, ?> map = gson.fromJson(value, Map.class);
73+
flattenMap(map, flattenedParameters, entry.getKey());
74+
} catch (Throwable ignored) {}
75+
}
76+
TemplateScript templateScript = scriptService.compile(script, TemplateScript.CONTEXT).newInstance(flattenedParameters);
77+
try {
78+
String result = templateScript.execute();
79+
listener.onResponse(result == null ? (T) "" : (T) result);
80+
} catch (Exception e) {
81+
listener.onFailure(e);
82+
}
83+
}
84+
85+
private void flattenMap(Map<String, ?> map, Map<String, Object> flatMap, String prefix) {
86+
for (Map.Entry<String, ?> entry : map.entrySet()) {
87+
String key = entry.getKey();
88+
if (prefix != null && !prefix.isEmpty()) {
89+
key = prefix + "." + entry.getKey();
90+
}
91+
Object value = entry.getValue();
92+
if (value instanceof Map) {
93+
flattenMap((Map<String, ?>) value, flatMap, key);
94+
} else {
95+
flatMap.put(key, value);
96+
}
97+
}
98+
}
99+
100+
@Override
101+
public boolean validate(Map<String, String> map) {
102+
return true;
103+
}
104+
105+
public static class Factory implements Tool.Factory<PainlessTool> {
106+
private ScriptService scriptService;
107+
108+
private static PainlessTool.Factory INSTANCE;
109+
110+
public static PainlessTool.Factory getInstance() {
111+
if (INSTANCE != null) {
112+
return INSTANCE;
113+
}
114+
synchronized (PainlessTool.class) {
115+
if (INSTANCE != null) {
116+
return INSTANCE;
117+
}
118+
INSTANCE = new PainlessTool.Factory();
119+
return INSTANCE;
120+
}
121+
}
122+
123+
public void init(ScriptService scriptService) {
124+
this.scriptService = scriptService;
125+
}
126+
127+
@Override
128+
public PainlessTool create(Map<String, Object> map) {
129+
String script = (String) map.get("script");
130+
// TODO add script non null/empty check
131+
return new PainlessTool(scriptService, script);
132+
}
133+
134+
@Override
135+
public String getDefaultDescription() {
136+
return DEFAULT_DESCRIPTION;
137+
}
138+
139+
@Override
140+
public String getDefaultType() {
141+
return TYPE;
142+
}
143+
144+
@Override
145+
public String getDefaultVersion() {
146+
return null;
147+
}
148+
149+
}
150+
}

‎src/test/java/org/opensearch/integTest/PainlessToolIT.java‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,79 @@
1+
/*
2+
* Copyright OpenSearch Contributors
3+
* SPDX-License-Identifier: Apache-2.0
4+
*/
5+
6+
package org.opensearch.integTest;
7+
8+
import java.io.IOException;
9+
import java.net.URISyntaxException;
10+
import java.nio.file.Files;
11+
import java.nio.file.Path;
12+
13+
import org.junit.Assert;
14+
import org.junit.Before;
15+
16+
import lombok.SneakyThrows;
17+
import lombok.extern.log4j.Log4j2;
18+
19+
@Log4j2
20+
public class PainlessToolIT extends BaseAgentToolsIT {
21+
22+
private String registerAgentRequestBody;
23+
24+
@Before
25+
@SneakyThrows
26+
public void setUp() {
27+
super.setUp();
28+
registerAgentRequestBody = Files
29+
.readString(
30+
Path.of(this.getClass().getClassLoader().getResource("org/opensearch/agent/tools/register_painless_agent.json").toURI())
31+
);
32+
}
33+
34+
public void test_execute() {
35+
String script = "def x = new HashMap(); x.abc = '5'; return x.abc;";
36+
String agentRequestBody = registerAgentRequestBody.replaceAll("<SCRIPT>", script);
37+
String agentId = createAgent(agentRequestBody);
38+
String agentInput = "{\"parameters\":{}}";
39+
String result = executeAgent(agentId, agentInput);
40+
Assert.assertEquals("5", result);
41+
}
42+
43+
public void test_execute_with_parameter() {
44+
String script = "params.x + params.y";
45+
String agentRequestBody = registerAgentRequestBody.replaceAll("<SCRIPT>", script);
46+
String agentId = createAgent(agentRequestBody);
47+
String agentInput = "{\"parameters\":{\"x\":1,\"y\":2}}";
48+
String result = executeAgent(agentId, agentInput);
49+
Assert.assertEquals("12", result);
50+
}
51+
52+
public void test_execute_with_parameter2() throws URISyntaxException, IOException {
53+
String script =
54+
"return 'An example output: with ppl:<ppl>' + params.get('PPL.output.ppl') + '</ppl>, and this is ppl result: <ppl_result>' + params.get('PPL.output.executionResult') + '</ppl_result>'";
55+
String mockPPLOutput = "return '{\\\\\"executionResult\\\\\":\\\\\"result\\\\\",\\\\\"ppl\\\\\":\\\\\"source=demo| head 1\\\\\"}'";
56+
String registerAgentRequestBody2 = Files
57+
.readString(
58+
Path
59+
.of(
60+
this
61+
.getClass()
62+
.getClassLoader()
63+
.getResource("org/opensearch/agent/tools/register_painless_agent_with_multiple_tools.json")
64+
.toURI()
65+
)
66+
);
67+
String agentRequestBody = registerAgentRequestBody2.replaceAll("<SCRIPT1>", mockPPLOutput).replaceAll("<SCRIPT2>", script);
68+
69+
log.info("agentRequestBody = {}", agentRequestBody);
70+
String agentId = createAgent(agentRequestBody);
71+
String agentInput = "{\"parameters\":{}}";
72+
String result = executeAgent(agentId, agentInput);
73+
Assert
74+
.assertEquals(
75+
"An example output: with ppl:<ppl>source=demo| head 1</ppl>, and this is ppl result: <ppl_result>result</ppl_result>",
76+
result
77+
);
78+
}
79+
}

‎src/test/resources/org/opensearch/agent/tools/register_painless_agent.json‎

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
{
2+
"name": "Test_PainlessTool",
3+
"type": "flow",
4+
"tools": [
5+
{
6+
"type": "PainlessTool",
7+
"parameters": {
8+
"script": "<SCRIPT>"
9+
}
10+
}
11+
]
12+
}

‎src/test/resources/org/opensearch/agent/tools/register_painless_agent_with_multiple_tools.json‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,19 @@
1+
{
2+
"name": "Test_PainlessTool",
3+
"type": "flow",
4+
"tools": [
5+
{
6+
"type": "PainlessTool",
7+
"name": "PPL",
8+
"parameters": {
9+
"script": "<SCRIPT1>"
10+
}
11+
},
12+
{
13+
"type": "PainlessTool",
14+
"parameters": {
15+
"script": "<SCRIPT2>"
16+
}
17+
}
18+
]
19+
}

0 commit comments

Comments
 (0)