diff --git a/memory/inmemory.go b/memory/inmemory.go index b2ee8d598..d91dedacf 100644 --- a/memory/inmemory.go +++ b/memory/inmemory.go @@ -60,31 +60,9 @@ func (s *inMemoryService) AddSessionToMemory(ctx context.Context, curSession ses var values []value for event := range curSession.Events().All() { - if event.LLMResponse.Content == nil { - continue - } - - words := make(map[string]struct{}) - for _, part := range event.LLMResponse.Content.Parts { - if part.Text == "" { - continue - } - - maps.Copy(words, extractWords(part.Text)) + if v, ok := eventToValue(event); ok { + values = append(values, v) } - - if len(words) == 0 { - continue - } - - values = append(values, value{ - id: event.ID, - content: event.LLMResponse.Content, - author: event.Author, - timestamp: event.Timestamp, - customMetadata: event.CustomMetadata, - words: words, - }) } k := key{ @@ -106,6 +84,72 @@ func (s *inMemoryService) AddSessionToMemory(ctx context.Context, curSession ses return nil } +func eventToValue(event *session.Event) (value, bool) { + if event.LLMResponse.Content == nil { + return value{}, false + } + + words := make(map[string]struct{}) + for _, part := range event.LLMResponse.Content.Parts { + if part.Text == "" { + continue + } + + maps.Copy(words, extractWords(part.Text)) + } + + if len(words) == 0 { + return value{}, false + } + + return value{ + id: event.ID, + content: event.LLMResponse.Content, + author: event.Author, + timestamp: event.Timestamp, + customMetadata: event.CustomMetadata, + words: words, + }, true +} + +func (s *inMemoryService) AddEventsToMemory(ctx context.Context, req *AddEventsToMemoryRequest) error { + k := key{ + appName: req.AppName, + userID: req.UserID, + } + sid := sessionID(req.SessionID) + + s.mu.Lock() + defer s.mu.Unlock() + + sessions, ok := s.store[k] + if !ok { + sessions = map[sessionID][]value{} + s.store[k] = sessions + } + existing := sessions[sid] + + existingIDs := make(map[string]struct{}, len(existing)) + for _, v := range existing { + existingIDs[v.id] = struct{}{} + } + + for _, event := range req.Events { + if _, ok := existingIDs[event.ID]; ok { + continue + } + v, ok := eventToValue(event) + if !ok { + continue + } + existing = append(existing, v) + existingIDs[event.ID] = struct{}{} + } + + sessions[sid] = existing + return nil +} + func (s *inMemoryService) SearchMemory(ctx context.Context, req *SearchRequest) (*SearchResponse, error) { queryWords := extractWords(req.Query) diff --git a/memory/inmemory_test.go b/memory/inmemory_test.go index 5e74772a5..168a0db45 100644 --- a/memory/inmemory_test.go +++ b/memory/inmemory_test.go @@ -161,6 +161,105 @@ func Test_inMemoryService_SearchMemory(t *testing.T) { } } +func Test_inMemoryService_AddEventsToMemory(t *testing.T) { + newEvent := func(id, text string) *session.Event { + return &session.Event{ + ID: id, + Author: "user1", + LLMResponse: model.LLMResponse{ + Content: genai.NewContentFromText(text, genai.RoleUser), + }, + } + } + + t.Run("events become searchable", func(t *testing.T) { + s := memory.InMemoryService() + err := s.AddEventsToMemory(t.Context(), &memory.AddEventsToMemoryRequest{ + AppName: "app1", + UserID: "user1", + SessionID: "sess1", + Events: []*session.Event{newEvent("event1", "The quick brown fox")}, + }) + if err != nil { + t.Fatalf("AddEventsToMemory() error = %v", err) + } + + got, err := s.SearchMemory(t.Context(), &memory.SearchRequest{AppName: "app1", UserID: "user1", Query: "fox"}) + if err != nil { + t.Fatalf("SearchMemory() error = %v", err) + } + if len(got.Memories) != 1 || got.Memories[0].ID != "event1" { + t.Errorf("SearchMemory() = %+v, want a single match for event1", got.Memories) + } + }) + + t.Run("repeated calls with overlapping events are deduped by ID", func(t *testing.T) { + s := memory.InMemoryService() + req := &memory.AddEventsToMemoryRequest{ + AppName: "app1", + UserID: "user1", + SessionID: "sess1", + Events: []*session.Event{newEvent("event1", "The quick brown fox")}, + } + if err := s.AddEventsToMemory(t.Context(), req); err != nil { + t.Fatalf("AddEventsToMemory() [1] error = %v", err) + } + req.Events = []*session.Event{newEvent("event1", "The quick brown fox"), newEvent("event2", "jumps over the lazy dog")} + if err := s.AddEventsToMemory(t.Context(), req); err != nil { + t.Fatalf("AddEventsToMemory() [2] error = %v", err) + } + + got, err := s.SearchMemory(t.Context(), &memory.SearchRequest{AppName: "app1", UserID: "user1", Query: "fox dog"}) + if err != nil { + t.Fatalf("SearchMemory() error = %v", err) + } + if len(got.Memories) != 2 { + t.Errorf("SearchMemory() = %+v, want 2 deduped memories", got.Memories) + } + }) + + t.Run("does not affect other sessions or users", func(t *testing.T) { + s := memory.InMemoryService() + if err := s.AddEventsToMemory(t.Context(), &memory.AddEventsToMemoryRequest{ + AppName: "app1", + UserID: "user1", + SessionID: "sess1", + Events: []*session.Event{newEvent("event1", "unique-marker-word")}, + }); err != nil { + t.Fatalf("AddEventsToMemory() error = %v", err) + } + + got, err := s.SearchMemory(t.Context(), &memory.SearchRequest{AppName: "app1", UserID: "user2", Query: "unique-marker-word"}) + if err != nil { + t.Fatalf("SearchMemory() error = %v", err) + } + if len(got.Memories) != 0 { + t.Errorf("SearchMemory() leaked across users, got %+v", got.Memories) + } + }) + + t.Run("events without content are ignored", func(t *testing.T) { + s := memory.InMemoryService() + err := s.AddEventsToMemory(t.Context(), &memory.AddEventsToMemoryRequest{ + AppName: "app1", + UserID: "user1", + SessionID: "sess1", + Events: []*session.Event{{ID: "event1", Author: "user1"}}, + }) + if err != nil { + t.Fatalf("AddEventsToMemory() error = %v", err) + } + + got, err := s.SearchMemory(t.Context(), &memory.SearchRequest{AppName: "app1", UserID: "user1", Query: "anything"}) + if err != nil { + t.Fatalf("SearchMemory() error = %v", err) + } + if len(got.Memories) != 0 { + t.Errorf("SearchMemory() = %+v, want no memories for a contentless event", got.Memories) + } + }) +} + func makeSession(t *testing.T, appName, userID, sessionID string, events []*session.Event) session.Session { t.Helper() diff --git a/memory/service.go b/memory/service.go index 9d27b3092..e58b232f5 100644 --- a/memory/service.go +++ b/memory/service.go @@ -33,11 +33,35 @@ type Service interface { // // A session can be added multiple times during its lifetime. AddSessionToMemory(ctx context.Context, s session.Session) error + // AddEventsToMemory adds an explicit list of events to the memory service. + // + // This is intended for callers who want to persist only a subset of + // events (e.g. the latest turn) rather than re-ingesting the full + // session via AddSessionToMemory. Implementations should treat Events + // as an incremental update and must not assume it represents the full + // session. + AddEventsToMemory(ctx context.Context, req *AddEventsToMemoryRequest) error // SearchMemory returns memory entries relevant to the given query. // Empty slice is returned if there are no matches. SearchMemory(ctx context.Context, req *SearchRequest) (*SearchResponse, error) } +// AddEventsToMemoryRequest represents a request for [Service.AddEventsToMemory]. +type AddEventsToMemoryRequest struct { + AppName string + UserID string + Events []*session.Event + + // Below are optional fields. + + // SessionID scopes the events to a session. Implementations may ignore + // it if not applicable. + SessionID string + // CustomMetadata is optional, implementation-defined metadata for + // memory generation (e.g. TTL). + CustomMetadata map[string]any +} + // SearchRequest represents a request for memory search. type SearchRequest struct { Query string diff --git a/memory/vertexai/vertexai.go b/memory/vertexai/vertexai.go index ef79187e0..2451eac3d 100644 --- a/memory/vertexai/vertexai.go +++ b/memory/vertexai/vertexai.go @@ -89,6 +89,12 @@ func (v *vertexAIService) AddSessionToMemory(ctx context.Context, s session.Sess return err } +// AddEventsToMemory is not yet supported by the Vertex AI Memory Bank +// service; use AddSessionToMemory to ingest a full session instead. +func (v *vertexAIService) AddEventsToMemory(ctx context.Context, req *memory.AddEventsToMemoryRequest) error { + return fmt.Errorf("vertexAIService does not support AddEventsToMemory") +} + // SearchMemory implements [memory.Service]. func (v *vertexAIService) SearchMemory(ctx context.Context, req *memory.SearchRequest) (*memory.SearchResponse, error) { return v.client.searchMemory(ctx, req) diff --git a/memory/vertexai/vertexai_test.go b/memory/vertexai/vertexai_test.go new file mode 100644 index 000000000..4ecf92d9f --- /dev/null +++ b/memory/vertexai/vertexai_test.go @@ -0,0 +1,29 @@ +// Copyright 2026 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 vertexai + +import ( + "testing" + + "google.golang.org/adk/v2/memory" +) + +func TestVertexAIService_AddEventsToMemory_NotSupported(t *testing.T) { + v := &vertexAIService{} + err := v.AddEventsToMemory(t.Context(), &memory.AddEventsToMemoryRequest{AppName: "app", UserID: "user"}) + if err == nil { + t.Error("AddEventsToMemory() error = nil, want non-nil: vertexAIService does not support AddEventsToMemory") + } +}