Skip to content

Commit ce8f38f

Browse files
committed
fix(memory): tokenize query and events alike in InMemoryMemoryService
Matches adk-python: both sides are normalized to NFC and split with a Unicode-aware \w+, so a query word followed by punctuation, a word with accented letters and a number now match. For event text, a word that mixes Latin and non-Latin characters also yields each single-script run, so a Latin word inside Japanese or Chinese text still matches. Before, the query was split on whitespace only and events kept only ASCII letters.
1 parent 4092a1f commit ce8f38f

2 files changed

Lines changed: 200 additions & 16 deletions

File tree

‎core/src/main/java/com/google/adk/memory/InMemoryMemoryService.java‎

Lines changed: 73 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -20,19 +20,18 @@
2020

2121
import com.google.adk.events.Event;
2222
import com.google.adk.sessions.Session;
23+
import com.google.common.base.CharMatcher;
2324
import com.google.common.collect.ImmutableList;
2425
import com.google.common.collect.ImmutableSet;
2526
import com.google.genai.types.Part;
2627
import io.reactivex.rxjava3.core.Completable;
2728
import io.reactivex.rxjava3.core.Single;
29+
import java.text.Normalizer;
2830
import java.time.Instant;
2931
import java.util.ArrayList;
30-
import java.util.Collections;
31-
import java.util.HashSet;
3232
import java.util.List;
3333
import java.util.Locale;
3434
import java.util.Map;
35-
import java.util.Set;
3635
import java.util.concurrent.ConcurrentHashMap;
3736
import java.util.regex.Matcher;
3837
import java.util.regex.Pattern;
@@ -44,8 +43,8 @@
4443
*/
4544
public final class InMemoryMemoryService implements BaseMemoryService {
4645

47-
// Pattern to extract words, matching the Python version.
48-
private static final Pattern WORD_PATTERN = Pattern.compile("[A-Za-z]+");
46+
// Letters, numbers and _, as Python's \w+ matches them.
47+
private static final Pattern WORD_PATTERN = Pattern.compile("[\\p{L}\\p{N}_]+");
4948

5049
/** Keys are "app_name/user_id", values are maps of "session_id" to a list of events. */
5150
private final Map<String, Map<String, List<Event>>> sessionEvents;
@@ -91,8 +90,7 @@ public Single<SearchMemoryResponse> searchMemory(String appName, String userId,
9190

9291
Map<String, List<Event>> userSessions = sessionEvents.get(key);
9392

94-
ImmutableSet<String> wordsInQuery =
95-
ImmutableSet.copyOf(query.toLowerCase(Locale.ROOT).split("\\s+"));
93+
ImmutableSet<String> wordsInQuery = extractWordsLower(query);
9694

9795
List<MemoryEntry> matchingMemories = new ArrayList<>();
9896

@@ -102,22 +100,29 @@ public Single<SearchMemoryResponse> searchMemory(String appName, String userId,
102100
continue;
103101
}
104102

105-
Set<String> wordsInEvent = new HashSet<>();
103+
List<String> eventTexts = new ArrayList<>();
106104
for (Part part : event.content().get().parts().get()) {
107-
String text = part.text().orElse("");
108-
if (!text.isEmpty()) {
109-
Matcher matcher = WORD_PATTERN.matcher(text);
110-
while (matcher.find()) {
111-
wordsInEvent.add(matcher.group().toLowerCase(Locale.ROOT));
112-
}
113-
}
105+
part.text().filter(text -> !text.isEmpty()).ifPresent(eventTexts::add);
114106
}
107+
String eventText = String.join(" ", eventTexts);
108+
ImmutableSet<String> wordsInEvent = extractSearchableWords(eventText);
115109

116110
if (wordsInEvent.isEmpty()) {
117111
continue;
118112
}
119113

120-
if (!Collections.disjoint(wordsInQuery, wordsInEvent)) {
114+
// A non-ASCII query word also matches inside the text, as Japanese and Chinese put no
115+
// spaces between words.
116+
String eventTextLower =
117+
Normalizer.normalize(eventText, Normalizer.Form.NFC).toLowerCase(Locale.ROOT);
118+
boolean matches =
119+
wordsInQuery.stream()
120+
.anyMatch(
121+
word ->
122+
wordsInEvent.contains(word)
123+
|| (!CharMatcher.ascii().matchesAllOf(word)
124+
&& eventTextLower.contains(word)));
125+
if (matches) {
121126
MemoryEntry memory =
122127
MemoryEntry.builder()
123128
.content(event.content().get())
@@ -135,6 +140,58 @@ public Single<SearchMemoryResponse> searchMemory(String appName, String userId,
135140
});
136141
}
137142

143+
/** Extracts words from a string and converts them to lowercase. */
144+
private static ImmutableSet<String> extractWordsLower(String text) {
145+
ImmutableSet.Builder<String> words = ImmutableSet.builder();
146+
Matcher matcher = WORD_PATTERN.matcher(Normalizer.normalize(text, Normalizer.Form.NFC));
147+
while (matcher.find()) {
148+
words.add(matcher.group().toLowerCase(Locale.ROOT));
149+
}
150+
return words.build();
151+
}
152+
153+
/**
154+
* Extracts the words an event can be matched on, in lowercase: the words of {@link
155+
* #extractWordsLower} plus, for a word that mixes Latin and non-Latin characters, each of its
156+
* single-script runs. Japanese and Chinese are written without spaces, so {@code 私はPythonを使う} is
157+
* a single word, and splitting it where the script changes lets a query for {@code python} match
158+
* while a partial word such as {@code thon} still does not.
159+
*/
160+
private static ImmutableSet<String> extractSearchableWords(String text) {
161+
ImmutableSet<String> words = extractWordsLower(text);
162+
ImmutableSet.Builder<String> searchable = ImmutableSet.<String>builder().addAll(words);
163+
for (String word : words) {
164+
if (CharMatcher.ascii().matchesAllOf(word)) {
165+
continue;
166+
}
167+
int runStart = 0;
168+
boolean runIsLatin = isLatin(word.codePointAt(0));
169+
for (int i = 0; i < word.length(); i += Character.charCount(word.codePointAt(i))) {
170+
boolean latin = isLatin(word.codePointAt(i));
171+
if (latin != runIsLatin) {
172+
searchable.add(word.substring(runStart, i));
173+
runStart = i;
174+
runIsLatin = latin;
175+
}
176+
}
177+
searchable.add(word.substring(runStart));
178+
}
179+
return searchable.build();
180+
}
181+
182+
/**
183+
* ASCII letters, digits and {@code _}, or a character above ASCII whose Unicode name starts with
184+
* {@code LATIN}, as Python checks it. Unlike the Latin script, this leaves out characters such as
185+
* {@code º} and {@code ª}.
186+
*/
187+
private static boolean isLatin(int codePoint) {
188+
if (codePoint < 0x80) {
189+
return Character.isLetterOrDigit(codePoint) || codePoint == '_';
190+
}
191+
String name = Character.getName(codePoint);
192+
return name != null && name.startsWith("LATIN");
193+
}
194+
138195
private String formatTimestamp(long timestamp) {
139196
return Instant.ofEpochMilli(timestamp).toString();
140197
}

‎core/src/test/java/com/google/adk/memory/InMemoryMemoryServiceTest.java‎

Lines changed: 127 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -137,4 +137,131 @@ public void searchMemory_textAndFunctionCallInOneEvent_stillMatchesOnTheText() {
137137
assertThat(response.memories().get(0).content())
138138
.isEqualTo(Content.fromParts(Part.fromText("weather in Seoul"), functionCall));
139139
}
140+
141+
@Test
142+
public void searchMemory_queryWithPunctuation_matchesWord() {
143+
InMemoryMemoryService service = new InMemoryMemoryService();
144+
service
145+
.addSessionToMemory(sessionWith(event(Part.fromText("weather in Seoul"))))
146+
.blockingAwait();
147+
148+
SearchMemoryResponse response =
149+
service.searchMemory(APP_NAME, USER_ID, "weather?").blockingGet();
150+
151+
assertThat(response.memories()).hasSize(1);
152+
}
153+
154+
@Test
155+
public void searchMemory_nonAsciiWord_matchesWholeWord() {
156+
InMemoryMemoryService service = new InMemoryMemoryService();
157+
service
158+
.addSessionToMemory(sessionWith(event(Part.fromText("Meet me at the café"))))
159+
.blockingAwait();
160+
161+
SearchMemoryResponse response = service.searchMemory(APP_NAME, USER_ID, "café").blockingGet();
162+
163+
assertThat(response.memories()).hasSize(1);
164+
assertThat(service.searchMemory(APP_NAME, USER_ID, "caf").blockingGet().memories()).isEmpty();
165+
}
166+
167+
@Test
168+
public void searchMemory_decomposedText_matchesPrecomposedQuery() {
169+
InMemoryMemoryService service = new InMemoryMemoryService();
170+
service
171+
.addSessionToMemory(sessionWith(event(Part.fromText("Meet me at the café"))))
172+
.blockingAwait();
173+
174+
SearchMemoryResponse response = service.searchMemory(APP_NAME, USER_ID, "café").blockingGet();
175+
176+
assertThat(response.memories()).hasSize(1);
177+
}
178+
179+
@Test
180+
public void searchMemory_decomposedQuery_matchesPrecomposedText() {
181+
InMemoryMemoryService service = new InMemoryMemoryService();
182+
service
183+
.addSessionToMemory(sessionWith(event(Part.fromText("Meet me at the café"))))
184+
.blockingAwait();
185+
186+
SearchMemoryResponse response = service.searchMemory(APP_NAME, USER_ID, "café").blockingGet();
187+
188+
assertThat(response.memories()).hasSize(1);
189+
}
190+
191+
@Test
192+
public void searchMemory_emojiVariationSelectorAndJoiner_areNotWords() {
193+
InMemoryMemoryService service = new InMemoryMemoryService();
194+
service
195+
.addSessionToMemory(
196+
sessionWith(
197+
event(Part.fromText("⚠️ disk full")), event(Part.fromText("👩‍🔬 lab notes"))))
198+
.blockingAwait();
199+
200+
assertThat(service.searchMemory(APP_NAME, USER_ID, "thanks ❤️").blockingGet().memories())
201+
.isEmpty();
202+
assertThat(service.searchMemory(APP_NAME, USER_ID, "👨‍💻 coding").blockingGet().memories())
203+
.isEmpty();
204+
}
205+
206+
@Test
207+
public void searchMemory_superscriptDigit_isPartOfWord() {
208+
InMemoryMemoryService service = new InMemoryMemoryService();
209+
service.addSessionToMemory(sessionWith(event(Part.fromText("Area is x²")))).blockingAwait();
210+
211+
assertThat(service.searchMemory(APP_NAME, USER_ID, "x²").blockingGet().memories()).hasSize(1);
212+
}
213+
214+
@Test
215+
public void searchMemory_latinWordInUnspacedScript_matchesWholeWord() {
216+
InMemoryMemoryService service = new InMemoryMemoryService();
217+
service
218+
.addSessionToMemory(sessionWith(event(Part.fromText("私はPythonでADKを使っています"))))
219+
.blockingAwait();
220+
221+
assertThat(service.searchMemory(APP_NAME, USER_ID, "python").blockingGet().memories())
222+
.hasSize(1);
223+
assertThat(service.searchMemory(APP_NAME, USER_ID, "adk").blockingGet().memories()).hasSize(1);
224+
assertThat(service.searchMemory(APP_NAME, USER_ID, "thon").blockingGet().memories()).isEmpty();
225+
assertThat(service.searchMemory(APP_NAME, USER_ID, "java").blockingGet().memories()).isEmpty();
226+
}
227+
228+
@Test
229+
public void searchMemory_digits_matchWord() {
230+
InMemoryMemoryService service = new InMemoryMemoryService();
231+
service
232+
.addSessionToMemory(sessionWith(event(Part.fromText("Order 12345 shipped"))))
233+
.blockingAwait();
234+
235+
SearchMemoryResponse response = service.searchMemory(APP_NAME, USER_ID, "12345").blockingGet();
236+
237+
assertThat(response.memories()).hasSize(1);
238+
}
239+
240+
@Test
241+
public void searchMemory_nonAsciiQueryWord_matchesInsideUnspacedText() {
242+
InMemoryMemoryService service = new InMemoryMemoryService();
243+
service
244+
.addSessionToMemory(
245+
sessionWith(
246+
event(Part.fromText("私の名前は太郎です")),
247+
event(Part.fromText("我喜欢机器学习")),
248+
event(Part.fromText("私はPythonでADKを使っています"))))
249+
.blockingAwait();
250+
251+
assertThat(service.searchMemory(APP_NAME, USER_ID, "太郎").blockingGet().memories()).hasSize(1);
252+
assertThat(service.searchMemory(APP_NAME, USER_ID, "天気").blockingGet().memories()).isEmpty();
253+
assertThat(service.searchMemory(APP_NAME, USER_ID, "机器学习").blockingGet().memories()).hasSize(1);
254+
assertThat(service.searchMemory(APP_NAME, USER_ID, "天气预报").blockingGet().memories()).isEmpty();
255+
assertThat(service.searchMemory(APP_NAME, USER_ID, "使って").blockingGet().memories()).hasSize(1);
256+
}
257+
258+
@Test
259+
public void searchMemory_ordinalIndicatorBeforeDigits_matchesDigits() {
260+
InMemoryMemoryService service = new InMemoryMemoryService();
261+
service.addSessionToMemory(sessionWith(event(Part.fromText("Pedido nº12345")))).blockingAwait();
262+
263+
SearchMemoryResponse response = service.searchMemory(APP_NAME, USER_ID, "12345").blockingGet();
264+
265+
assertThat(response.memories()).hasSize(1);
266+
}
140267
}

0 commit comments

Comments
 (0)