Skip to content

Commit 78435aa

Browse files
authored
Merge branch 'main' into create-pull-request/3.4.0-SNAPSHOT
2 parents 13fdd96 + 9b11665 commit 78435aa

File tree

4 files changed

+146
-85
lines changed

4 files changed

+146
-85
lines changed

ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/ReadFromScratchPadTool.java

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -100,26 +100,20 @@ public boolean validate(Map<String, String> parameters) {
100100
*/
101101
@Override
102102
public <T> void run(Map<String, String> parameters, ActionListener<T> listener) {
103-
// Handle both List<String> and String (JSON) formats for existing notes
103+
104104
List<String> notes;
105-
Map rawParameters = parameters;
106-
Object existingNotes = rawParameters.get(SCRATCHPAD_NOTES_KEY);
107-
if (existingNotes instanceof List) {
108-
notes = new ArrayList<>((List<String>) existingNotes);
109-
} else if (existingNotes instanceof String) {
110-
List<String> parsedNotes = StringUtils.parseStringArrayToList((String) existingNotes);
111-
notes = parsedNotes != null ? new ArrayList<>(parsedNotes) : new ArrayList<>();
112-
} else {
113-
notes = new ArrayList<>();
114-
}
105+
String existingNotes = parameters.getOrDefault(SCRATCHPAD_NOTES_KEY, "[]");
106+
107+
List<String> parsedNotes = StringUtils.parseStringArrayToList(existingNotes);
108+
notes = parsedNotes != null ? new ArrayList<>(parsedNotes) : new ArrayList<>();
115109

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

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

122-
rawParameters.put(SCRATCHPAD_NOTES_KEY, notes);
116+
parameters.put(SCRATCHPAD_NOTES_KEY, StringUtils.toJson(notes));
123117

124118
if (notes.isEmpty()) {
125119
listener.onResponse((T) "Scratchpad is empty.");

ml-algorithms/src/main/java/org/opensearch/ml/engine/tools/WriteToScratchPadTool.java

Lines changed: 6 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -111,21 +111,15 @@ public <T> void run(Map<String, String> parameters, ActionListener<T> listener)
111111
return;
112112
}
113113

114-
// Handle both List<String> and String (JSON) formats for existing notes
115114
List<String> notes;
116-
Map rawParameters = parameters;
117-
Object existingNotes = rawParameters.get(SCRATCHPAD_NOTES_KEY);
118-
if (existingNotes instanceof List) {
119-
notes = new ArrayList<>((List<String>) existingNotes);
120-
} else if (existingNotes instanceof String) {
121-
List<String> parsedNotes = StringUtils.parseStringArrayToList((String) existingNotes);
122-
notes = parsedNotes != null ? new ArrayList<>(parsedNotes) : new ArrayList<>();
123-
} else {
124-
notes = new ArrayList<>();
125-
}
115+
;
116+
String existingNotes = parameters.getOrDefault(SCRATCHPAD_NOTES_KEY, "[]");
117+
118+
List<String> parsedNotes = StringUtils.parseStringArrayToList(existingNotes);
119+
notes = parsedNotes != null ? new ArrayList<>(parsedNotes) : new ArrayList<>();
126120

127121
notes.add(currentNote);
128-
rawParameters.put(SCRATCHPAD_NOTES_KEY, notes);
122+
parameters.put(SCRATCHPAD_NOTES_KEY, StringUtils.toJson(notes));
129123

130124
if (returnHistory) {
131125
String fullNotesFormatted = "- " + String.join("\n- ", notes);

ml-algorithms/src/test/java/org/opensearch/ml/engine/tools/ReadFromScratchPadToolTests.java

Lines changed: 68 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,6 @@
1010
import static org.junit.Assert.assertTrue;
1111
import static org.mockito.Mockito.verify;
1212

13-
import java.util.ArrayList;
14-
import java.util.Arrays;
1513
import java.util.HashMap;
1614
import java.util.Map;
1715

@@ -83,105 +81,142 @@ public void testValidate() {
8381

8482
@Test
8583
public void testRun_NoNotes() {
86-
Map<String, Object> parameters = new HashMap<>();
87-
tool.run((Map) parameters, listener);
84+
Map<String, String> parameters = new HashMap<>();
85+
tool.run(parameters, listener);
8886
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
8987
verify(listener).onResponse(captor.capture());
9088
assertEquals("Scratchpad is empty.", captor.getValue());
91-
assertEquals(new ArrayList<>(), parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
89+
assertEquals("[]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
9290
}
9391

9492
@Test
9593
public void testRun_WithScratchpadNotes() {
96-
Map<String, Object> parameters = new HashMap<>();
97-
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, new ArrayList<>(Arrays.asList("existing note")));
98-
tool.run((Map) parameters, listener);
94+
Map<String, String> parameters = new HashMap<>();
95+
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"existing note\"]");
96+
tool.run(parameters, listener);
9997
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
10098
verify(listener).onResponse(captor.capture());
10199
assertEquals("Notes from scratchpad:\n- existing note", captor.getValue());
102-
assertEquals(Arrays.asList("existing note"), parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
100+
assertEquals("[\"existing note\"]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
103101
}
104102

105103
@Test
106104
public void testRun_WithPersistentNotes() {
107-
Map<String, Object> parameters = new HashMap<>();
108-
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, new ArrayList<>()); // Initialize with empty list
105+
Map<String, String> parameters = new HashMap<>();
106+
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[]");
109107
parameters.put(ReadFromScratchPadTool.PERSISTENT_NOTES_KEY, "persistent note");
110-
tool.run((Map) parameters, listener);
108+
tool.run(parameters, listener);
111109
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
112110
verify(listener).onResponse(captor.capture());
113111
assertEquals("Notes from scratchpad:\n- persistent note", captor.getValue());
114-
assertEquals(Arrays.asList("persistent note"), parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
112+
assertEquals("[\"persistent note\"]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
115113
}
116114

117115
@Test
118116
public void testRun_WithBothNotes() {
119-
Map<String, Object> parameters = new HashMap<>();
120-
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, new ArrayList<>(Arrays.asList("existing note")));
117+
Map<String, String> parameters = new HashMap<>();
118+
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"existing note\"]");
121119
parameters.put(ReadFromScratchPadTool.PERSISTENT_NOTES_KEY, "persistent note");
122-
tool.run((Map) parameters, listener);
120+
tool.run(parameters, listener);
123121
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
124122
verify(listener).onResponse(captor.capture());
125123
assertEquals("Notes from scratchpad:\n- existing note\n- persistent note", captor.getValue());
126-
assertEquals(Arrays.asList("existing note", "persistent note"), parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
124+
assertEquals("[\"existing note\",\"persistent note\"]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
127125
}
128126

129127
@Test
130128
public void testRun_WithDuplicatePersistentNotes() {
131-
Map<String, Object> parameters = new HashMap<>();
132-
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, new ArrayList<>(Arrays.asList("existing note", "persistent note")));
129+
Map<String, String> parameters = new HashMap<>();
130+
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"existing note\",\"persistent note\"]");
133131
parameters.put(ReadFromScratchPadTool.PERSISTENT_NOTES_KEY, "persistent note");
134-
tool.run((Map) parameters, listener);
132+
tool.run(parameters, listener);
135133
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
136134
verify(listener).onResponse(captor.capture());
137135
assertEquals("Notes from scratchpad:\n- existing note\n- persistent note", captor.getValue());
138-
assertEquals(Arrays.asList("existing note", "persistent note"), parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
136+
assertEquals("[\"existing note\",\"persistent note\"]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
139137
}
140138

141139
@Test
142140
public void testRun_WithNonListScratchpadNotes() {
143-
Map<String, Object> parameters = new HashMap<>();
141+
Map<String, String> parameters = new HashMap<>();
144142
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "not a list");
145-
tool.run((Map) parameters, listener);
143+
tool.run(parameters, listener);
146144
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
147145
verify(listener).onResponse(captor.capture());
148146
assertEquals("Scratchpad is empty.", captor.getValue());
149-
assertEquals(new ArrayList<>(), parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
147+
assertEquals("[]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
150148
}
151149

152150
@Test
153151
public void testRun_WithJsonStringNotes() {
154-
Map<String, Object> parameters = new HashMap<>();
152+
Map<String, String> parameters = new HashMap<>();
155153
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"json note\"]");
156-
tool.run((Map) parameters, listener);
154+
tool.run(parameters, listener);
157155
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
158156
verify(listener).onResponse(captor.capture());
159157
assertEquals("Notes from scratchpad:\n- json note", captor.getValue());
160-
assertEquals(Arrays.asList("json note"), parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
158+
assertEquals("[\"json note\"]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
161159
}
162160

163161
@Test
164162
public void testRun_WithEmptyPersistentNotes() {
165-
Map<String, Object> parameters = new HashMap<>();
166-
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, new ArrayList<>(Arrays.asList("existing note")));
163+
Map<String, String> parameters = new HashMap<>();
164+
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"existing note\"]");
167165
parameters.put(ReadFromScratchPadTool.PERSISTENT_NOTES_KEY, "");
168-
tool.run((Map) parameters, listener);
166+
tool.run(parameters, listener);
169167
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
170168
verify(listener).onResponse(captor.capture());
171169
assertEquals("Notes from scratchpad:\n- existing note", captor.getValue());
172170
}
173171

174172
@Test
175173
public void testRun_WithNullPersistentNotes() {
176-
Map<String, Object> parameters = new HashMap<>();
177-
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, new ArrayList<>(Arrays.asList("existing note")));
174+
Map<String, String> parameters = new HashMap<>();
175+
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"existing note\"]");
178176
parameters.put(ReadFromScratchPadTool.PERSISTENT_NOTES_KEY, null);
179-
tool.run((Map) parameters, listener);
177+
tool.run(parameters, listener);
180178
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
181179
verify(listener).onResponse(captor.capture());
182180
assertEquals("Notes from scratchpad:\n- existing note", captor.getValue());
183181
}
184182

183+
@Test
184+
public void testRun_StringConversion_EmptyArray() {
185+
Map<String, String> parameters = new HashMap<>();
186+
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[]");
187+
tool.run(parameters, listener);
188+
189+
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
190+
verify(listener).onResponse(captor.capture());
191+
assertEquals("Scratchpad is empty.", captor.getValue());
192+
assertEquals("[]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
193+
}
194+
195+
@Test
196+
public void testRun_StringConversion_WithJsonArray() {
197+
Map<String, String> parameters = new HashMap<>();
198+
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"note1\",\"note2\"]");
199+
tool.run(parameters, listener);
200+
201+
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
202+
verify(listener).onResponse(captor.capture());
203+
assertEquals("Notes from scratchpad:\n- note1\n- note2", captor.getValue());
204+
assertEquals("[\"note1\",\"note2\"]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
205+
}
206+
207+
@Test
208+
public void testRun_StringConversion_AddPersistentNote() {
209+
Map<String, String> parameters = new HashMap<>();
210+
parameters.put(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY, "[\"existing\"]");
211+
parameters.put(ReadFromScratchPadTool.PERSISTENT_NOTES_KEY, "new note");
212+
tool.run(parameters, listener);
213+
214+
ArgumentCaptor<String> captor = ArgumentCaptor.forClass(String.class);
215+
verify(listener).onResponse(captor.capture());
216+
assertEquals("Notes from scratchpad:\n- existing\n- new note", captor.getValue());
217+
assertEquals("[\"existing\",\"new note\"]", parameters.get(ReadFromScratchPadTool.SCRATCHPAD_NOTES_KEY));
218+
}
219+
185220
@Test
186221
public void testFactory() {
187222
ReadFromScratchPadTool.Factory factory = ReadFromScratchPadTool.Factory.getInstance();

0 commit comments

Comments
 (0)