Skip to content

make memory optional in conversational agent #3626

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

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

make memory optional in conversational agent #3626

Show file tree
Hide file tree
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
44 changes: 26 additions & 18 deletions ml-algorithms/src/main/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunner.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -168,8 +168,7 @@ public void run(MLAgent mlAgent, Map<String, String> inputParams, ActionListener
if (functionCalling != null) {
functionCalling.configure(params);
}

String memoryType = mlAgent.getMemory().getType();
String memoryType = mlAgent.getMemory() == null ? null : mlAgent.getMemory().getType();
Copy link
Contributor

Choose a reason for hiding this comment

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

does the current behavior throw NPE here? wondering if there is some validation in the REST/Transport layer to check for memory, didn't see any tho

Copy link
Contributor

Choose a reason for hiding this comment

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

Memory is created by default in MLAgentExecutor right? so should we change some code there as well

String memoryId = params.get(MLAgentExecutor.MEMORY_ID);
String appType = mlAgent.getAppType();
String title = params.get(MLAgentExecutor.QUESTION);
Expand All @@ -178,6 +177,11 @@ public void run(MLAgent mlAgent, Map<String, String> inputParams, ActionListener
String chatHistoryResponseTemplate = params.get(CHAT_HISTORY_RESPONSE_TEMPLATE);
int messageHistoryLimit = getMessageHistoryLimit(params);

if (memoryType == null || memoryType.isBlank()) {
runAgent(mlAgent, params, listener, null, null, functionCalling);
return;
}

ConversationIndexMemory.Factory conversationIndexMemoryFactory = (ConversationIndexMemory.Factory) memoryFactoryMap.get(memoryType);
conversationIndexMemoryFactory.create(title, memoryId, appType, ActionListener.<ConversationIndexMemory>wrap(memory -> {
// TODO: call runAgent directly if messageHistoryLimit == 0
Expand Down Expand Up @@ -317,7 +321,9 @@ private void runReAct(
AtomicReference<String> newPrompt = new AtomicReference<>(tmpSubstitutor.replace(prompt));
tmpParameters.put(PROMPT, newPrompt.get());

List<ModelTensors> traceTensors = createModelTensors(sessionId, parentInteractionId);
List<ModelTensors> traceTensors = (conversationIndexMemory == null)
? new ArrayList<>()
: createModelTensors(sessionId, parentInteractionId);
Comment on lines +324 to +326
Copy link
Contributor

@pyek-bot pyek-bot Mar 17, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

what about passing conversationIndexMemory to this method and returning empty list rather than using ternary operator here?

*ignore, seems like it will affect a lot of other methods using this

int maxIterations = Integer.parseInt(tmpParameters.getOrDefault(MAX_ITERATION, DEFAULT_MAX_ITERATIONS));
for (int i = 0; i < maxIterations; i++) {
int finalI = i;
Expand Down Expand Up @@ -513,8 +519,8 @@ private void runReAct(
client.execute(MLPredictionTaskAction.INSTANCE, request, firstListener);
}

private static List<ModelTensors> createFinalAnswerTensors(List<ModelTensors> sessionId, List<ModelTensor> lastThought) {
List<ModelTensors> finalModelTensors = sessionId;
private static List<ModelTensors> createFinalAnswerTensors(List<ModelTensors> modelTensorsList, List<ModelTensor> lastThought) {
List<ModelTensors> finalModelTensors = modelTensorsList;
finalModelTensors.add(ModelTensors.builder().mlModelTensors(lastThought).build());
return finalModelTensors;
}
Expand Down Expand Up @@ -719,19 +725,21 @@ private void sendFinalAnswer(
public static List<ModelTensors> createModelTensors(String sessionId, String parentInteractionId) {
List<ModelTensors> cotModelTensors = new ArrayList<>();

cotModelTensors
.add(
ModelTensors
.builder()
.mlModelTensors(
List
.of(
ModelTensor.builder().name(MLAgentExecutor.MEMORY_ID).result(sessionId).build(),
ModelTensor.builder().name(MLAgentExecutor.PARENT_INTERACTION_ID).result(parentInteractionId).build()
)
)
.build()
);
if (sessionId != null && !sessionId.isBlank()) {
cotModelTensors
.add(
ModelTensors
.builder()
.mlModelTensors(
List
.of(
ModelTensor.builder().name(MLAgentExecutor.MEMORY_ID).result(sessionId).build(),
ModelTensor.builder().name(MLAgentExecutor.PARENT_INTERACTION_ID).result(parentInteractionId).build()
)
)
.build()
);
}
return cotModelTensors;
}

Expand Down
53 changes: 53 additions & 0 deletions ...rithms/src/test/java/org/opensearch/ml/engine/algorithms/agent/MLChatAgentRunnerTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -880,6 +880,42 @@ public void testToolExecutionWithChatHistoryParameter() {
Assert.assertTrue(toolParamsCapture.getValue().containsKey(MLChatAgentRunner.CHAT_HISTORY));
}

@Test
public void testParsingJsonBlockFromResponseNoMemory() {
// Prepare the response with JSON block
String jsonBlock = "{\"thought\":\"parsed thought\", \"action\":\"parsed action\", "
+ "\"action_input\":\"parsed action input\", \"final_answer\":\"parsed final answer\"}";
String responseWithJsonBlock = "Some text```json" + jsonBlock + "```More text";

// Mock LLM response to not contain "thought" but contain "response" with JSON block
Map<String, String> llmResponse = new HashMap<>();
llmResponse.put("response", responseWithJsonBlock);
doAnswer(getLLMAnswer(llmResponse))
.when(client)
.execute(any(ActionType.class), any(ActionRequest.class), isA(ActionListener.class));

// Create an MLAgent and run the MLChatAgentRunner
MLAgent mlAgent = createMLAgentNoMemory();
Map<String, String> params = new HashMap<>();
params.put("verbose", "true");
mlChatAgentRunner.run(mlAgent, params, agentActionListener);

// Capture the response passed to the listener
ArgumentCaptor<Object> responseCaptor = ArgumentCaptor.forClass(Object.class);
verify(agentActionListener).onResponse(responseCaptor.capture());

// Extract the captured response
Object capturedResponse = responseCaptor.getValue();
assertTrue(capturedResponse instanceof ModelTensorOutput);
ModelTensorOutput modelTensorOutput = (ModelTensorOutput) capturedResponse;

ModelTensor modelTensor = modelTensorOutput.getMlModelOutputs().get(0).getMlModelTensors().get(0);

assertEquals(1, modelTensorOutput.getMlModelOutputs().size());
assertEquals(1, modelTensorOutput.getMlModelOutputs().get(0).getMlModelTensors().size());
assertEquals("parsed final answer", modelTensor.getResult());
}

// Helper methods to create MLAgent and parameters
private MLAgent createMLAgentWithTools() {
LLMSpec llmSpec = LLMSpec.builder().modelId("MODEL_ID").build();
Expand Down Expand Up @@ -918,6 +954,23 @@ private MLAgent createMLAgentWithToolsConfig(Map<String, String> configMap) {
.build();
}

private MLAgent createMLAgentNoMemory() {
LLMSpec llmSpec = LLMSpec.builder().modelId("MODEL_ID").build();
MLToolSpec firstToolSpec = MLToolSpec
.builder()
.name(FIRST_TOOL)
.type(FIRST_TOOL)
.parameters(ImmutableMap.of("key1", "value1", "key2", "value2"))
.build();
return MLAgent
.builder()
.name("TestAgent")
.type(MLAgentType.CONVERSATIONAL.name())
.tools(Arrays.asList(firstToolSpec))
.llm(llmSpec)
.build();
}

private Map<String, String> createAgentParamsWithAction(String action, String actionInput) {
Map<String, String> params = new HashMap<>();
params.put("action", action);
Expand Down