diff --git a/memory/inmemory.go b/memory/inmemory.go index 2e260e48c..f4e018032 100644 --- a/memory/inmemory.go +++ b/memory/inmemory.go @@ -20,6 +20,7 @@ import ( "strings" "sync" "time" + "unicode" "google.golang.org/genai" @@ -167,12 +168,19 @@ func checkMapsIntersect(m1, m2 map[string]struct{}) bool { func extractWords(text string) map[string]struct{} { res := make(map[string]struct{}) - for s := range strings.SplitSeq(text, " ") { - if s == "" { - continue - } + // Split on runs of non-word characters so that punctuation and non-space + // whitespace are not folded into a token: without this, "great!" is stored + // as-is and a search for "great" misses it. This mirrors the `\w+` + // tokenization used by adk-python's in-memory memory service. + for _, s := range strings.FieldsFunc(text, isNotWord) { res[strings.ToLower(s)] = struct{}{} } return res } + +// isNotWord reports whether r separates two words, i.e. whether it is neither a +// letter, a digit nor an underscore. +func isNotWord(r rune) bool { + return !unicode.IsLetter(r) && !unicode.IsNumber(r) && r != '_' +} diff --git a/memory/inmemory_test.go b/memory/inmemory_test.go index a31370c09..51121b360 100644 --- a/memory/inmemory_test.go +++ b/memory/inmemory_test.go @@ -90,6 +90,102 @@ func Test_inMemoryService_SearchMemory(t *testing.T) { }, }, }, + { + name: "find events next to punctuation", + initSessions: []session.Session{ + makeSession(t, "app1", "user1", "sess1", []*session.Event{ + { + ID: "event1", + Author: "test-bot", + LLMResponse: model.LLMResponse{Content: genai.NewContentFromText("The agent works great!", genai.RoleModel)}, + Timestamp: must(time.Parse(time.RFC3339, "2023-10-01T10:00:00Z")), + }, + { + ID: "event2", + Author: "test-bot", + LLMResponse: model.LLMResponse{Content: genai.NewContentFromText("regions: us-east1,us-west1", genai.RoleModel)}, + Timestamp: must(time.Parse(time.RFC3339, "2023-10-02T10:00:00Z")), + }, + }), + }, + req: &memory.SearchRequest{ + AppName: "app1", + UserID: "user1", + Query: "great us-west1", + }, + wantResp: &memory.SearchResponse{ + Memories: []memory.Entry{ + { + ID: "event1", + Content: genai.NewContentFromText("The agent works great!", genai.RoleModel), + Author: "test-bot", + Timestamp: must(time.Parse(time.RFC3339, "2023-10-01T10:00:00Z")), + }, + { + ID: "event2", + Content: genai.NewContentFromText("regions: us-east1,us-west1", genai.RoleModel), + Author: "test-bot", + Timestamp: must(time.Parse(time.RFC3339, "2023-10-02T10:00:00Z")), + }, + }, + }, + }, + { + name: "find events separated by non-space whitespace", + initSessions: []session.Session{ + makeSession(t, "app1", "user1", "sess1", []*session.Event{ + { + ID: "event1", + Author: "test-bot", + LLMResponse: model.LLMResponse{Content: genai.NewContentFromText("first line\nsecond\tline", genai.RoleModel)}, + Timestamp: must(time.Parse(time.RFC3339, "2023-10-01T10:00:00Z")), + }, + }), + }, + req: &memory.SearchRequest{ + AppName: "app1", + UserID: "user1", + Query: "second", + }, + wantResp: &memory.SearchResponse{ + Memories: []memory.Entry{ + { + ID: "event1", + Content: genai.NewContentFromText("first line\nsecond\tline", genai.RoleModel), + Author: "test-bot", + Timestamp: must(time.Parse(time.RFC3339, "2023-10-01T10:00:00Z")), + }, + }, + }, + }, + { + name: "find events for a query containing punctuation", + initSessions: []session.Session{ + makeSession(t, "app1", "user1", "sess1", []*session.Event{ + { + ID: "event1", + Author: "test-bot", + LLMResponse: model.LLMResponse{Content: genai.NewContentFromText("the deploy step timed out", genai.RoleModel)}, + Timestamp: must(time.Parse(time.RFC3339, "2023-10-01T10:00:00Z")), + }, + }), + }, + req: &memory.SearchRequest{ + AppName: "app1", + UserID: "user1", + Query: "Why timed-out?", + }, + wantResp: &memory.SearchResponse{ + Memories: []memory.Entry{ + { + ID: "event1", + Content: genai.NewContentFromText("the deploy step timed out", genai.RoleModel), + Author: "test-bot", + Timestamp: must(time.Parse(time.RFC3339, "2023-10-01T10:00:00Z")), + }, + }, + }, + }, { name: "no leakage for different appName", initSessions: []session.Session{