Skip to content

Fix scratch pad when return_history is true #4310

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

Merged
merged 1 commit into from
Oct 14, 2025
+146 −85
Merged

Fix scratch pad when return_history is true #4310

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
18 changes: 6 additions & 12 deletions ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/ReadFromScratchPadTool.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -100,26 +100,20 @@ public boolean validate(Map<String, String> parameters) {
*/
@Override
public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
// Handle both List<String> and String (JSON) formats for existing notes

List<String> notes;
Map rawParameters = parameters;
Object existingNotes = rawParameters.get(SCRATCHPAD_NOTES_KEY);
if (existingNotes instanceof List) {
notes = new ArrayList<>((List<String>) existingNotes);
} else if (existingNotes instanceof String) {
List<String> parsedNotes = StringUtils.parseStringArrayToList((String) existingNotes);
notes = parsedNotes != null ? new ArrayList<>(parsedNotes) : new ArrayList<>();
} else {
notes = new ArrayList<>();
}
String existingNotes = parameters.getOrDefault(SCRATCHPAD_NOTES_KEY, "[]");

List<String> parsedNotes = StringUtils.parseStringArrayToList(existingNotes);
notes = parsedNotes != null ? new ArrayList<>(parsedNotes) : new ArrayList<>();

String persistentNotes = parameters.getOrDefault(PERSISTENT_NOTES_KEY, "");

if (persistentNotes != null && !persistentNotes.isEmpty() && !notes.contains(persistentNotes)) {
notes.add(persistentNotes);
}

rawParameters.put(SCRATCHPAD_NOTES_KEY, notes);
parameters.put(SCRATCHPAD_NOTES_KEY, StringUtils.toJson(notes));

if (notes.isEmpty()) {
listener.onResponse((T) "Scratchpad is empty.");
Expand Down
18 changes: 6 additions & 12 deletions ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/WriteToScratchPadTool.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -111,21 +111,15 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
return;
}

// Handle both List<String> and String (JSON) formats for existing notes
List<String> notes;
Map rawParameters = parameters;
Object existingNotes = rawParameters.get(SCRATCHPAD_NOTES_KEY);
if (existingNotes instanceof List) {
notes = new ArrayList<>((List<String>) existingNotes);
} else if (existingNotes instanceof String) {
List<String> parsedNotes = StringUtils.parseStringArrayToList((String) existingNotes);
notes = parsedNotes != null ? new ArrayList<>(parsedNotes) : new ArrayList<>();
} else {
notes = new ArrayList<>();
}
;
String existingNotes = parameters.getOrDefault(SCRATCHPAD_NOTES_KEY, "[]");

List<String> parsedNotes = StringUtils.parseStringArrayToList(existingNotes);
notes = parsedNotes != null ? new ArrayList<>(parsedNotes) : new ArrayList<>();

notes.add(currentNote);
rawParameters.put(SCRATCHPAD_NOTES_KEY, notes);
parameters.put(SCRATCHPAD_NOTES_KEY, StringUtils.toJson(notes));

if (returnHistory) {
String fullNotesFormatted = "- " + String.join("\n- ", notes);
Expand Down
101 changes: 68 additions & 33 deletions ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/ReadFromScratchPadToolTests.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,8 +10,6 @@
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.verify;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.Map;

Expand Down Expand Up @@ -83,105 +81,142 @@ public void testValidate() {

@Test
public void testRun_NoNotes() {
Map<String, Object> parameters = new HashMap<>();
tool.run((Map) parameters, listener);
Map<String, String> parameters = new HashMap<>();
tool.run(parameters, listener);
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
verify(listener).onResponse(captor.capture());
assertEquals("Scratchpad is empty.", captor.getValue());
assertEquals(new ArrayList<>(), parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
assertEquals("[]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
}

@Test
public void testRun_WithScratchpadNotes() {
Map<String, Object> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, new ArrayList<>(Arrays.asList("existing note")));
tool.run((Map) parameters, listener);
Map<String, String> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"existing note\"]");
tool.run(parameters, listener);
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
verify(listener).onResponse(captor.capture());
assertEquals("Notes from scratchpad:\n- existing note", captor.getValue());
assertEquals(Arrays.asList("existing note"), parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
assertEquals("[\"existing note\"]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
}

@Test
public void testRun_WithPersistentNotes() {
Map<String, Object> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, new ArrayList<>()); // Initialize with empty list
Map<String, String> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[]");
parameters.put(ReadFromScratchPadTool.PERSISTENT_NOTES_KEY, "persistent note");
tool.run((Map) parameters, listener);
tool.run(parameters, listener);
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
verify(listener).onResponse(captor.capture());
assertEquals("Notes from scratchpad:\n- persistent note", captor.getValue());
assertEquals(Arrays.asList("persistent note"), parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
assertEquals("[\"persistent note\"]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
}

@Test
public void testRun_WithBothNotes() {
Map<String, Object> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, new ArrayList<>(Arrays.asList("existing note")));
Map<String, String> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"existing note\"]");
parameters.put(ReadFromScratchPadTool.PERSISTENT_NOTES_KEY, "persistent note");
tool.run((Map) parameters, listener);
tool.run(parameters, listener);
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
verify(listener).onResponse(captor.capture());
assertEquals("Notes from scratchpad:\n- existing note\n- persistent note", captor.getValue());
assertEquals(Arrays.asList("existing note", "persistent note"), parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
assertEquals("[\"existing note\",\"persistent note\"]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
}

@Test
public void testRun_WithDuplicatePersistentNotes() {
Map<String, Object> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, new ArrayList<>(Arrays.asList("existing note", "persistent note")));
Map<String, String> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"existing note\",\"persistent note\"]");
parameters.put(ReadFromScratchPadTool.PERSISTENT_NOTES_KEY, "persistent note");
tool.run((Map) parameters, listener);
tool.run(parameters, listener);
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
verify(listener).onResponse(captor.capture());
assertEquals("Notes from scratchpad:\n- existing note\n- persistent note", captor.getValue());
assertEquals(Arrays.asList("existing note", "persistent note"), parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
assertEquals("[\"existing note\",\"persistent note\"]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
}

@Test
public void testRun_WithNonListScratchpadNotes() {
Map<String, Object> parameters = new HashMap<>();
Map<String, String> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "not a list");
tool.run((Map) parameters, listener);
tool.run(parameters, listener);
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
verify(listener).onResponse(captor.capture());
assertEquals("Scratchpad is empty.", captor.getValue());
assertEquals(new ArrayList<>(), parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
assertEquals("[]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
}

@Test
public void testRun_WithJsonStringNotes() {
Map<String, Object> parameters = new HashMap<>();
Map<String, String> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"json note\"]");
tool.run((Map) parameters, listener);
tool.run(parameters, listener);
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
verify(listener).onResponse(captor.capture());
assertEquals("Notes from scratchpad:\n- json note", captor.getValue());
assertEquals(Arrays.asList("json note"), parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
assertEquals("[\"json note\"]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
}

@Test
public void testRun_WithEmptyPersistentNotes() {
Map<String, Object> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, new ArrayList<>(Arrays.asList("existing note")));
Map<String, String> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"existing note\"]");
parameters.put(ReadFromScratchPadTool.PERSISTENT_NOTES_KEY, "");
tool.run((Map) parameters, listener);
tool.run(parameters, listener);
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
verify(listener).onResponse(captor.capture());
assertEquals("Notes from scratchpad:\n- existing note", captor.getValue());
}

@Test
public void testRun_WithNullPersistentNotes() {
Map<String, Object> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, new ArrayList<>(Arrays.asList("existing note")));
Map<String, String> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"existing note\"]");
parameters.put(ReadFromScratchPadTool.PERSISTENT_NOTES_KEY, null);
tool.run((Map) parameters, listener);
tool.run(parameters, listener);
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
verify(listener).onResponse(captor.capture());
assertEquals("Notes from scratchpad:\n- existing note", captor.getValue());
}

@Test
public void testRun_StringConversion_EmptyArray() {
Map<String, String> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[]");
tool.run(parameters, listener);

ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
verify(listener).onResponse(captor.capture());
assertEquals("Scratchpad is empty.", captor.getValue());
assertEquals("[]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
}

@Test
public void testRun_StringConversion_WithJsonArray() {
Map<String, String> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"note1\",\"note2\"]");
tool.run(parameters, listener);

ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
verify(listener).onResponse(captor.capture());
assertEquals("Notes from scratchpad:\n- note1\n- note2", captor.getValue());
assertEquals("[\"note1\",\"note2\"]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
}

@Test
public void testRun_StringConversion_AddPersistentNote() {
Map<String, String> parameters = new HashMap<>();
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"existing\"]");
parameters.put(ReadFromScratchPadTool.PERSISTENT_NOTES_KEY, "new note");
tool.run(parameters, listener);

ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
verify(listener).onResponse(captor.capture());
assertEquals("Notes from scratchpad:\n- existing\n- new note", captor.getValue());
assertEquals("[\"existing\",\"new note\"]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
}

@Test
public void testFactory() {
ReadFromScratchPadTool.Factory factory = ReadFromScratchPadTool.Factory.getInstance();
Expand Down
Loading
Loading