From 008d409599e9be48dc9afdc936e3a71561a7cdf7 Mon Sep 17 00:00:00 2001 From: Laurianti Date: Mon, 28 Sep 2026 21:40:04 +0200 Subject: [PATCH] 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. --- .../adk/memory/InMemoryMemoryService.java | 89 +++++++++--- .../adk/memory/InMemoryMemoryServiceTest.java | 127 ++++++++++++++++++ 2 files changed, 200 insertions(+), 16 deletions(-) diff --git a/core/src/main/java/com/google/adk/memory/InMemoryMemoryService.java b/core/src/main/java/com/google/adk/memory/InMemoryMemoryService.java index a10260240..45c3ecd80 100644 --- a/core/src/main/java/com/google/adk/memory/InMemoryMemoryService.java +++ b/core/src/main/java/com/google/adk/memory/InMemoryMemoryService.java @@ -20,19 +20,18 @@ import com.google.adk.events.Event; import com.google.adk.sessions.Session; +import com.google.common.base.CharMatcher; import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; import com.google.genai.types.Part; import io.reactivex.rxjava3.core.Completable; import io.reactivex.rxjava3.core.Single; +import java.text.Normalizer; import java.time.Instant; import java.util.ArrayList; -import java.util.Collections; -import java.util.HashSet; import java.util.List; import java.util.Locale; import java.util.Map; -import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.regex.Matcher; import java.util.regex.Pattern; @@ -44,8 +43,8 @@ */ public final class InMemoryMemoryService implements BaseMemoryService { - // Pattern to extract words, matching the Python version. - private static final Pattern WORD_PATTERN = Pattern.compile("[A-Za-z]+"); + // Letters, numbers and _, as Python's \w+ matches them. + private static final Pattern WORD_PATTERN = Pattern.compile("[\\p{L}\\p{N}_]+"); /** Keys are "app_name/user_id", values are maps of "session_id" to a list of events. */ private final Map>> sessionEvents; @@ -91,8 +90,7 @@ public Single searchMemory(String appName, String userId, Map> userSessions = sessionEvents.get(key); - ImmutableSet wordsInQuery = - ImmutableSet.copyOf(query.toLowerCase(Locale.ROOT).split("\\s+")); + ImmutableSet wordsInQuery = extractWordsLower(query); List matchingMemories = new ArrayList<>(); @@ -102,22 +100,29 @@ public Single searchMemory(String appName, String userId, continue; } - Set wordsInEvent = new HashSet<>(); + List eventTexts = new ArrayList<>(); for (Part part : event.content().get().parts().get()) { - String text = part.text().orElse(""); - if (!text.isEmpty()) { - Matcher matcher = WORD_PATTERN.matcher(text); - while (matcher.find()) { - wordsInEvent.add(matcher.group().toLowerCase(Locale.ROOT)); - } - } + part.text().filter(text -> !text.isEmpty()).ifPresent(eventTexts::add); } + String eventText = String.join(" ", eventTexts); + ImmutableSet wordsInEvent = extractSearchableWords(eventText); if (wordsInEvent.isEmpty()) { continue; } - if (!Collections.disjoint(wordsInQuery, wordsInEvent)) { + // A non-ASCII query word also matches inside the text, as Japanese and Chinese put no + // spaces between words. + String eventTextLower = + Normalizer.normalize(eventText, Normalizer.Form.NFC).toLowerCase(Locale.ROOT); + boolean matches = + wordsInQuery.stream() + .anyMatch( + word -> + wordsInEvent.contains(word) + || (!CharMatcher.ascii().matchesAllOf(word) + && eventTextLower.contains(word))); + if (matches) { MemoryEntry memory = MemoryEntry.builder() .content(event.content().get()) @@ -135,6 +140,58 @@ public Single searchMemory(String appName, String userId, }); } + /** Extracts words from a string and converts them to lowercase. */ + private static ImmutableSet extractWordsLower(String text) { + ImmutableSet.Builder words = ImmutableSet.builder(); + Matcher matcher = WORD_PATTERN.matcher(Normalizer.normalize(text, Normalizer.Form.NFC)); + while (matcher.find()) { + words.add(matcher.group().toLowerCase(Locale.ROOT)); + } + return words.build(); + } + + /** + * Extracts the words an event can be matched on, in lowercase: the words of {@link + * #extractWordsLower} plus, for a word that mixes Latin and non-Latin characters, each of its + * single-script runs. Japanese and Chinese are written without spaces, so {@code 私はPythonを使う} is + * a single word, and splitting it where the script changes lets a query for {@code python} match + * while a partial word such as {@code thon} still does not. + */ + private static ImmutableSet extractSearchableWords(String text) { + ImmutableSet words = extractWordsLower(text); + ImmutableSet.Builder searchable = ImmutableSet.builder().addAll(words); + for (String word : words) { + if (CharMatcher.ascii().matchesAllOf(word)) { + continue; + } + int runStart = 0; + boolean runIsLatin = isLatin(word.codePointAt(0)); + for (int i = 0; i < word.length(); i += Character.charCount(word.codePointAt(i))) { + boolean latin = isLatin(word.codePointAt(i)); + if (latin != runIsLatin) { + searchable.add(word.substring(runStart, i)); + runStart = i; + runIsLatin = latin; + } + } + searchable.add(word.substring(runStart)); + } + return searchable.build(); + } + + /** + * ASCII letters, digits and {@code _}, or a character above ASCII whose Unicode name starts with + * {@code LATIN}, as Python checks it. Unlike the Latin script, this leaves out characters such as + * {@code º} and {@code ª}. + */ + private static boolean isLatin(int codePoint) { + if (codePoint < 0x80) { + return Character.isLetterOrDigit(codePoint) || codePoint == '_'; + } + String name = Character.getName(codePoint); + return name != null && name.startsWith("LATIN"); + } + private String formatTimestamp(long timestamp) { return Instant.ofEpochMilli(timestamp).toString(); } diff --git a/core/src/test/java/com/google/adk/memory/InMemoryMemoryServiceTest.java b/core/src/test/java/com/google/adk/memory/InMemoryMemoryServiceTest.java index 075492d76..1a9fc6600 100644 --- a/core/src/test/java/com/google/adk/memory/InMemoryMemoryServiceTest.java +++ b/core/src/test/java/com/google/adk/memory/InMemoryMemoryServiceTest.java @@ -137,4 +137,131 @@ public void searchMemory_textAndFunctionCallInOneEvent_stillMatchesOnTheText() { assertThat(response.memories().get(0).content()) .isEqualTo(Content.fromParts(Part.fromText("weather in Seoul"), functionCall)); } + + @Test + public void searchMemory_queryWithPunctuation_matchesWord() { + InMemoryMemoryService service = new InMemoryMemoryService(); + service + .addSessionToMemory(sessionWith(event(Part.fromText("weather in Seoul")))) + .blockingAwait(); + + SearchMemoryResponse response = + service.searchMemory(APP_NAME, USER_ID, "weather?").blockingGet(); + + assertThat(response.memories()).hasSize(1); + } + + @Test + public void searchMemory_nonAsciiWord_matchesWholeWord() { + InMemoryMemoryService service = new InMemoryMemoryService(); + service + .addSessionToMemory(sessionWith(event(Part.fromText("Meet me at the café")))) + .blockingAwait(); + + SearchMemoryResponse response = service.searchMemory(APP_NAME, USER_ID, "café").blockingGet(); + + assertThat(response.memories()).hasSize(1); + assertThat(service.searchMemory(APP_NAME, USER_ID, "caf").blockingGet().memories()).isEmpty(); + } + + @Test + public void searchMemory_decomposedText_matchesPrecomposedQuery() { + InMemoryMemoryService service = new InMemoryMemoryService(); + service + .addSessionToMemory(sessionWith(event(Part.fromText("Meet me at the café")))) + .blockingAwait(); + + SearchMemoryResponse response = service.searchMemory(APP_NAME, USER_ID, "café").blockingGet(); + + assertThat(response.memories()).hasSize(1); + } + + @Test + public void searchMemory_decomposedQuery_matchesPrecomposedText() { + InMemoryMemoryService service = new InMemoryMemoryService(); + service + .addSessionToMemory(sessionWith(event(Part.fromText("Meet me at the café")))) + .blockingAwait(); + + SearchMemoryResponse response = service.searchMemory(APP_NAME, USER_ID, "café").blockingGet(); + + assertThat(response.memories()).hasSize(1); + } + + @Test + public void searchMemory_emojiVariationSelectorAndJoiner_areNotWords() { + InMemoryMemoryService service = new InMemoryMemoryService(); + service + .addSessionToMemory( + sessionWith( + event(Part.fromText("⚠️ disk full")), event(Part.fromText("👩‍🔬 lab notes")))) + .blockingAwait(); + + assertThat(service.searchMemory(APP_NAME, USER_ID, "thanks ❤️").blockingGet().memories()) + .isEmpty(); + assertThat(service.searchMemory(APP_NAME, USER_ID, "👨‍💻 coding").blockingGet().memories()) + .isEmpty(); + } + + @Test + public void searchMemory_superscriptDigit_isPartOfWord() { + InMemoryMemoryService service = new InMemoryMemoryService(); + service.addSessionToMemory(sessionWith(event(Part.fromText("Area is x²")))).blockingAwait(); + + assertThat(service.searchMemory(APP_NAME, USER_ID, "x²").blockingGet().memories()).hasSize(1); + } + + @Test + public void searchMemory_latinWordInUnspacedScript_matchesWholeWord() { + InMemoryMemoryService service = new InMemoryMemoryService(); + service + .addSessionToMemory(sessionWith(event(Part.fromText("私はPythonでADKを使っています")))) + .blockingAwait(); + + assertThat(service.searchMemory(APP_NAME, USER_ID, "python").blockingGet().memories()) + .hasSize(1); + assertThat(service.searchMemory(APP_NAME, USER_ID, "adk").blockingGet().memories()).hasSize(1); + assertThat(service.searchMemory(APP_NAME, USER_ID, "thon").blockingGet().memories()).isEmpty(); + assertThat(service.searchMemory(APP_NAME, USER_ID, "java").blockingGet().memories()).isEmpty(); + } + + @Test + public void searchMemory_digits_matchWord() { + InMemoryMemoryService service = new InMemoryMemoryService(); + service + .addSessionToMemory(sessionWith(event(Part.fromText("Order 12345 shipped")))) + .blockingAwait(); + + SearchMemoryResponse response = service.searchMemory(APP_NAME, USER_ID, "12345").blockingGet(); + + assertThat(response.memories()).hasSize(1); + } + + @Test + public void searchMemory_nonAsciiQueryWord_matchesInsideUnspacedText() { + InMemoryMemoryService service = new InMemoryMemoryService(); + service + .addSessionToMemory( + sessionWith( + event(Part.fromText("私の名前は太郎です")), + event(Part.fromText("我喜欢机器学习")), + event(Part.fromText("私はPythonでADKを使っています")))) + .blockingAwait(); + + assertThat(service.searchMemory(APP_NAME, USER_ID, "太郎").blockingGet().memories()).hasSize(1); + assertThat(service.searchMemory(APP_NAME, USER_ID, "天気").blockingGet().memories()).isEmpty(); + assertThat(service.searchMemory(APP_NAME, USER_ID, "机器学习").blockingGet().memories()).hasSize(1); + assertThat(service.searchMemory(APP_NAME, USER_ID, "天气预报").blockingGet().memories()).isEmpty(); + assertThat(service.searchMemory(APP_NAME, USER_ID, "使って").blockingGet().memories()).hasSize(1); + } + + @Test + public void searchMemory_ordinalIndicatorBeforeDigits_matchesDigits() { + InMemoryMemoryService service = new InMemoryMemoryService(); + service.addSessionToMemory(sessionWith(event(Part.fromText("Pedido nº12345")))).blockingAwait(); + + SearchMemoryResponse response = service.searchMemory(APP_NAME, USER_ID, "12345").blockingGet(); + + assertThat(response.memories()).hasSize(1); + } }