feat: Add InMemoryMemoryService to ADK.

cornellgit · copybara-github · commit c8fed2db9b14 · 2025-06-28T11:02:13.000-07:00
PiperOrigin-RevId: 776988199
diff --git a/core/src/main/java/com/google/adk/memory/BaseMemoryService.java b/core/src/main/java/com/google/adk/memory/BaseMemoryService.java
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.adk.memory;
+
+import com.google.adk.sessions.Session;
+import io.reactivex.rxjava3.core.Completable;
+import io.reactivex.rxjava3.core.Single;
+
+/**
+ * Base contract for memory services.
+ *
+ * <p>The service provides functionalities to ingest sessions into memory so that the memory can be
+ * used for user queries.
+ */
+public interface BaseMemoryService {
+
+  /**
+   * Adds a session to the memory service.
+   *
+   * <p>A session may be added multiple times during its lifetime.
+   *
+   * @param session The session to add.
+   */
+  Completable addSessionToMemory(Session session);
+
+  /**
+   * Searches for sessions that match the query asynchronously.
+   *
+   * @param appName The name of the application.
+   * @param userId The id of the user.
+   * @param query The query to search for.
+   * @return A {@link SearchMemoryResponse} containing the matching memories.
+   */
+  Single<SearchMemoryResponse> searchMemory(String appName, String userId, String query);
+}
diff --git a/core/src/main/java/com/google/adk/memory/InMemoryMemoryService.java b/core/src/main/java/com/google/adk/memory/InMemoryMemoryService.java
@@ -0,0 +1,139 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.adk.memory;
+
+import static com.google.common.collect.ImmutableList.toImmutableList;
+
+import com.google.adk.events.Event;
+import com.google.adk.sessions.Session;
+import com.google.common.base.Strings;
+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.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;
+
+/**
+ * An in-memory memory service for prototyping purposes only.
+ *
+ * <p>Uses keyword matching instead of semantic search.
+ */
+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]+");
+
+  /** 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;
+
+  public InMemoryMemoryService() {
+    this.sessionEvents = new ConcurrentHashMap<>();
+  }
+
+  private static String userKey(String appName, String userId) {
+    return appName + "/" + userId;
+  }
+
+  @Override
+  public Completable addSessionToMemory(Session session) {
+    return Completable.fromAction(
+        () -> {
+          String key = userKey(session.appName(), session.userId());
+          Map<String, List<Event>> userSessions =
+              sessionEvents.computeIfAbsent(key, k -> new ConcurrentHashMap<>());
+          ImmutableList<Event> nonEmptyEvents =
+              session.events().stream()
+                  .filter(
+                      event ->
+                          event.content().isPresent()
+                              && event.content().get().parts().isPresent()
+                              && !event.content().get().parts().get().isEmpty())
+                  .collect(toImmutableList());
+          userSessions.put(session.id(), nonEmptyEvents);
+        });
+  }
+
+  @Override
+  public Single<SearchMemoryResponse> searchMemory(String appName, String userId, String query) {
+    return Single.fromCallable(
+        () -> {
+          String key = userKey(appName, userId);
+
+          if (!sessionEvents.containsKey(key)) {
+            return SearchMemoryResponse.builder().build();
+          }
+
+          Map<String, List<Event>> userSessions = sessionEvents.get(key);
+
+          ImmutableSet<String> wordsInQuery =
+              ImmutableSet.copyOf(query.toLowerCase(Locale.ROOT).split("\\s+"));
+
+          List<MemoryEntry> matchingMemories = new ArrayList<>();
+
+          for (List<Event> eventsInSession : userSessions.values()) {
+            for (Event event : eventsInSession) {
+              if (event.content().isEmpty() || event.content().get().parts().isEmpty()) {
+                continue;
+              }
+
+              Set<String> wordsInEvent = new HashSet<>();
+              for (Part part : event.content().get().parts().get()) {
+                if (!Strings.isNullOrEmpty(part.text().get())) {
+                  Matcher matcher = WORD_PATTERN.matcher(part.text().get());
+                  while (matcher.find()) {
+                    wordsInEvent.add(matcher.group().toLowerCase(Locale.ROOT));
+                  }
+                }
+              }
+
+              if (wordsInEvent.isEmpty()) {
+                continue;
+              }
+
+              if (!Collections.disjoint(wordsInQuery, wordsInEvent)) {
+                MemoryEntry memory =
+                    MemoryEntry.builder()
+                        .setContent(event.content().get())
+                        .setAuthor(event.author())
+                        .setTimestamp(formatTimestamp(event.timestamp()))
+                        .build();
+                matchingMemories.add(memory);
+              }
+            }
+          }
+
+          return SearchMemoryResponse.builder()
+              .setMemories(ImmutableList.copyOf(matchingMemories))
+              .build();
+        });
+  }
+
+  private String formatTimestamp(long timestamp) {
+    return Instant.ofEpochSecond(timestamp).toString();
+  }
+}
diff --git a/core/src/main/java/com/google/adk/memory/MemoryEntry.java b/core/src/main/java/com/google/adk/memory/MemoryEntry.java
@@ -0,0 +1,85 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.adk.memory;
+
+import com.google.auto.value.AutoValue;
+import com.google.genai.types.Content;
+import java.time.Instant;
+import javax.annotation.Nullable;
+
+/** Represents one memory entry. */
+@AutoValue
+public abstract class MemoryEntry {
+
+  /** Returns the main content of the memory. */
+  public abstract Content content();
+
+  /** Returns the author of the memory, or null if not set. */
+  @Nullable
+  public abstract String author();
+
+  /**
+   * Returns the timestamp when the original content of this memory happened, or null if not set.
+   *
+   * <p>This string will be forwarded to LLM. Preferred format is ISO 8601 format
+   */
+  @Nullable
+  public abstract String timestamp();
+
+  /** Returns a new builder for creating a {@link MemoryEntry}. */
+  public static Builder builder() {
+    return new AutoValue_MemoryEntry.Builder();
+  }
+
+  /**
+   * Creates a new builder with a copy of this entry's values.
+   *
+   * @return a new {@link Builder} instance.
+   */
+  public abstract Builder toBuilder();
+
+  /** Builder for {@link MemoryEntry}. */
+  @AutoValue.Builder
+  public abstract static class Builder {
+
+    /**
+     * Sets the main content of the memory.
+     *
+     * <p>This is a required field.
+     */
+    public abstract Builder setContent(Content content);
+
+    /** Sets the author of the memory. */
+    public abstract Builder setAuthor(@Nullable String author);
+
+    /** Sets the timestamp when the original content of this memory happened. */
+    public abstract Builder setTimestamp(@Nullable String timestamp);
+
+    /**
+     * A convenience method to set the timestamp from an {@link Instant} object, formatted as an ISO
+     * 8601 string.
+     *
+     * @param instant The timestamp as an Instant object.
+     */
+    public Builder setTimestamp(Instant instant) {
+      return setTimestamp(instant.toString());
+    }
+
+    /** Builds the immutable {@link MemoryEntry} object. */
+    public abstract MemoryEntry build();
+  }
+}
diff --git a/core/src/main/java/com/google/adk/memory/SearchMemoryResponse.java b/core/src/main/java/com/google/adk/memory/SearchMemoryResponse.java
@@ -0,0 +1,49 @@
+/*
+ * Copyright 2025 Google LLC
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.google.adk.memory;
+
+import com.google.auto.value.AutoValue;
+import com.google.common.collect.ImmutableList;
+import java.util.List;
+
+/** Represents the response from a memory search. */
+@AutoValue
+public abstract class SearchMemoryResponse {
+
+  /** Returns a list of memory entries that relate to the search query. */
+  public abstract ImmutableList<MemoryEntry> memories();
+
+  /** Creates a new builder for {@link SearchMemoryResponse}. */
+  public static Builder builder() {
+    return new AutoValue_SearchMemoryResponse.Builder().setMemories(ImmutableList.of());
+  }
+
+  /** Builder for {@link SearchMemoryResponse}. */
+  @AutoValue.Builder
+  public abstract static class Builder {
+
+    abstract Builder setMemories(ImmutableList<MemoryEntry> memories);
+
+    /** Sets the list of memory entries using a list. */
+    public Builder setMemories(List<MemoryEntry> memories) {
+      return setMemories(ImmutableList.copyOf(memories));
+    }
+
+    /** Builds the immutable {@link SearchMemoryResponse} object. */
+    public abstract SearchMemoryResponse build();
+  }
+}
