Skip to content

Commit f6254eb

Browse files
committed
add fill for selected tools
Signed-off-by: xinyual <xinyual@amazon.com>
1 parent 5cba422 commit f6254eb

File tree

2 files changed

+13
-15
lines changed

2 files changed

+13
-15
lines changed

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/AgentUtils.java

+11
Original file line numberDiff line numberDiff line change
@@ -383,6 +383,17 @@ public static String getToolName(MLToolSpec toolSpec) {
383383
return toolSpec.getName() != null ? toolSpec.getName() : toolSpec.getType();
384384
}
385385

386+
public static void fillSelectTool(Map<String, String> params)
387+
{
388+
if (params.containsKey(SELECTED_TOOLS))
389+
{
390+
String selectedTools = params.get(SELECTED_TOOLS);
391+
StringSubstitutor substitutor = new StringSubstitutor(params, "${parameters.", "}");
392+
selectedTools = substitutor.replace(selectedTools);
393+
params.put(SELECTED_TOOLS, selectedTools);
394+
}
395+
}
396+
386397
public static List<MLToolSpec> getMlToolSpecs(MLAgent mlAgent, Map<String, String> params) {
387398
String selectedToolsStr = params.get(SELECTED_TOOLS);
388399
List<MLToolSpec> toolSpecs = mlAgent.getTools();

ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java

+2-15
Original file line numberDiff line numberDiff line change
@@ -8,21 +8,7 @@
88
import static org.opensearch.ml.common.conversation.ActionConstants.ADDITIONAL_INFO_FIELD;
99
import static org.opensearch.ml.common.conversation.ActionConstants.AI_RESPONSE_FIELD;
1010
import static org.opensearch.ml.common.utils.StringUtils.gson;
11-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.DISABLE_TRACE;
12-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.PROMPT_CHAT_HISTORY_PREFIX;
13-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.PROMPT_PREFIX;
14-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.PROMPT_SUFFIX;
15-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.RESPONSE_FORMAT_INSTRUCTION;
16-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.TOOL_RESPONSE;
17-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.VERBOSE;
18-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.constructToolParams;
19-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.createTools;
20-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.getMessageHistoryLimit;
21-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.getMlToolSpecs;
22-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.getToolName;
23-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.getToolNames;
24-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.outputToOutputString;
25-
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.parseLLMOutput;
11+
import static org.opensearch.ml.engine.algorithms.agent.AgentUtils.*;
2612
import static org.opensearch.ml.engine.algorithms.agent.PromptTemplate.CHAT_HISTORY_PREFIX;
2713

2814
import java.security.PrivilegedActionException;
@@ -170,6 +156,7 @@ public void run(MLAgent mlAgent, Map<String, String> params, ActionListener<Obje
170156
}
171157

172158
private void runAgent(MLAgent mlAgent, Map<String, String> params, ActionListener<Object> listener, Memory memory, String sessionId) {
159+
fillSelectTool(params);
173160
List<MLToolSpec> toolSpecs = getMlToolSpecs(mlAgent, params);
174161
Map<String, Tool> tools = new HashMap<>();
175162
Map<String, MLToolSpec> toolSpecMap = new HashMap<>();

0 commit comments

Comments
 (0)