fix(memory): tokenize query and events alike in InMemoryMemoryService

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<String, Map<String, List<Event>>> sessionEvents;
@@ -91,8 +90,7 @@ public Single<SearchMemoryResponse> searchMemory(String appName, String userId,
 
           Map<String, List<Event>> userSessions = sessionEvents.get(key);
 
-          ImmutableSet<String> wordsInQuery =
-              ImmutableSet.copyOf(query.toLowerCase(Locale.ROOT).split("\\s+"));
+          ImmutableSet<String> wordsInQuery = extractWordsLower(query);
 
           List<MemoryEntry> matchingMemories = new ArrayList<>();
 
@@ -102,22 +100,29 @@ public Single<SearchMemoryResponse> searchMemory(String appName, String userId,
                 continue;
               }
 
-              Set<String> wordsInEvent = new HashSet<>();
+              List<String> 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<String> 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<SearchMemoryResponse> searchMemory(String appName, String userId,
         });
   }
 
+  /** Extracts words from a string and converts them to lowercase. */
+  private static ImmutableSet<String> extractWordsLower(String text) {
+    ImmutableSet.Builder<String> 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<String> extractSearchableWords(String text) {
+    ImmutableSet<String> words = extractWordsLower(text);
+    ImmutableSet.Builder<String> searchable = ImmutableSet.<String>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();
   }

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);
+  }
 }