fix(memory): tokenize on punctuation and all whitespace in extractWords

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 != '_'
+}

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{