Skip to content

Commit 191da59

Browse files
djthorpedjthorpe
committed
Updated memory to make it recursive
1 parent 3d5f59f commit 191da59

4 files changed

Lines changed: 150 additions & 10 deletions

File tree

‎memory/manager/connector.go‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -128,7 +128,7 @@ func (*searchMemoryTool) Name() string {
128128
}
129129

130130
func (*searchMemoryTool) Description() string {
131-
return "Search memory entries for the current session using PostgreSQL web-style search syntax. Leave q empty or use * to list all memories for the session. Current UTC date, time, datetime, and timezone are returned as dynamic memory entries."
131+
return "Search memory entries for the current session and its parent sessions using PostgreSQL web-style search syntax. Leave q empty or use * to list all memories in that ancestry. Current UTC date, time, datetime, and timezone are returned as dynamic memory entries."
132132
}
133133

134134
func (*searchMemoryTool) InputSchema() *jsonschema.Schema {

‎memory/schema/memory.go‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,7 @@ type MemorySelector struct {
4444
// MemoryListRequest represents a request to list or search memory entries.
4545
type MemoryListRequest struct {
4646
pg.OffsetLimit
47-
Session *uuid.UUID `json:"session,omitzero" help:"Restrict results to a single session" optional:""`
47+
Session *uuid.UUID `json:"session,omitzero" help:"Restrict results to a session and recursively include its parent sessions" optional:""`
4848
Q string `json:"q,omitempty" help:"Web-style text query matched against memory keys and values using PostgreSQL websearch syntax; leave empty or use * to list all memories for the session" optional:""`
4949
Start *time.Time `json:"start,omitempty" help:"Return memories on or after this timestamp" optional:""`
5050
End *time.Time `json:"end,omitempty" help:"Return memories on or before this timestamp" optional:""`
@@ -128,11 +128,13 @@ func (sel MemorySelector) Select(bind *pg.Bind, op pg.Op) (string, error) {
128128

129129
func (req MemoryListRequest) Select(bind *pg.Bind, op pg.Op) (string, error) {
130130
bind.Del("where")
131+
queryName := "memory.list"
131132
if req.Session != nil {
132133
if *req.Session == uuid.Nil {
133134
return "", fmt.Errorf("memory session cannot be nil")
134135
}
135-
bind.Append("where", `memory."session" = `+bind.Set("session", *req.Session))
136+
bind.Set("session", *req.Session)
137+
queryName = "memory.list_recursive"
136138
}
137139
if q := strings.TrimSpace(req.Q); q != "" && q != "*" {
138140
bind.Append("where", `to_tsvector('simple', COALESCE(memory."key", '') || ' ' || COALESCE(memory."value", '')) @@ websearch_to_tsquery('simple', `+bind.Set("q", q)+`)`)
@@ -148,14 +150,18 @@ func (req MemoryListRequest) Select(bind *pg.Bind, op pg.Op) (string, error) {
148150
if where == "" {
149151
bind.Set("where", "")
150152
} else {
151-
bind.Set("where", "WHERE "+where)
153+
prefix := "WHERE "
154+
if queryName == "memory.list_recursive" {
155+
prefix = "AND "
156+
}
157+
bind.Set("where", prefix+where)
152158
}
153159
bind.Set("orderby", `ORDER BY memory.created_at DESC, memory."session" ASC, memory."key" ASC`)
154160
req.OffsetLimit.Bind(bind, MemoryListMax)
155161

156162
switch op {
157163
case pg.List:
158-
return bind.Query("memory.list"), nil
164+
return bind.Query(queryName), nil
159165
default:
160166
return "", fmt.Errorf("MemoryListRequest: unsupported operation %q", op)
161167
}

‎memory/schema/memory_test.go‎

Lines changed: 73 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@ import (
1111

1212
func TestMemoryListRequestSelectUsesWebsearchSyntax(t *testing.T) {
1313
session := uuid.New()
14-
bind := pg.NewBind("memory.list", "SELECT")
14+
bind := pg.NewBind("memory.list_recursive", "SELECT")
1515

1616
query, err := (MemoryListRequest{Session: &session, Q: `"david" OR berlin`}).Select(bind, pg.List)
1717
if err != nil {
@@ -21,6 +21,9 @@ func TestMemoryListRequestSelectUsesWebsearchSyntax(t *testing.T) {
2121
t.Fatalf("unexpected query: %q", query)
2222
}
2323
where, _ := bind.Get("where").(string)
24+
if !strings.HasPrefix(where, "AND ") {
25+
t.Fatalf("expected recursive query filters to be AND-prefixed, got %q", where)
26+
}
2427
if !strings.Contains(where, "websearch_to_tsquery('simple'") {
2528
t.Fatalf("expected websearch_to_tsquery in where clause, got %q", where)
2629
}
@@ -37,7 +40,7 @@ func TestMemoryListRequestSelectUsesWebsearchSyntax(t *testing.T) {
3740

3841
func TestMemoryListRequestSelectWildcardOmitsSearchFilter(t *testing.T) {
3942
session := uuid.New()
40-
bind := pg.NewBind("memory.list", "SELECT")
43+
bind := pg.NewBind("memory.list_recursive", "SELECT")
4144

4245
_, err := (MemoryListRequest{Session: &session, Q: "*"}).Select(bind, pg.List)
4346
if err != nil {
@@ -47,10 +50,76 @@ func TestMemoryListRequestSelectWildcardOmitsSearchFilter(t *testing.T) {
4750
if strings.Contains(where, "tsquery") {
4851
t.Fatalf("did not expect text search filter in where clause, got %q", where)
4952
}
50-
if !strings.Contains(where, `memory."session" = `) {
51-
t.Fatalf("expected session filter in where clause, got %q", where)
53+
if where != "" {
54+
t.Fatalf("expected no extra where clause for wildcard search, got %q", where)
5255
}
5356
if bind.Get("q") != nil {
5457
t.Fatalf("did not expect q binding for wildcard search, got %v", bind.Get("q"))
5558
}
5659
}
60+
61+
func TestMemoryListRequestSelectUsesRecursiveSessionQuery(t *testing.T) {
62+
session := uuid.New()
63+
bind := pg.NewBind("memory.list_recursive", "SELECT")
64+
65+
query, err := (MemoryListRequest{Session: &session}).Select(bind, pg.List)
66+
if err != nil {
67+
t.Fatal(err)
68+
}
69+
if query != "SELECT" {
70+
t.Fatalf("unexpected query: %q", query)
71+
}
72+
if got := bind.Get("session"); got != session {
73+
t.Fatalf("unexpected session binding: %v", got)
74+
}
75+
}
76+
77+
func TestMemoryListRequestSelectWithoutSessionUsesPlainListQuery(t *testing.T) {
78+
bind := pg.NewBind("memory.list", "SELECT")
79+
80+
query, err := (MemoryListRequest{}).Select(bind, pg.List)
81+
if err != nil {
82+
t.Fatal(err)
83+
}
84+
if query != "SELECT" {
85+
t.Fatalf("unexpected query: %q", query)
86+
}
87+
if got := bind.Get("session"); got != nil {
88+
t.Fatalf("did not expect session binding, got %v", got)
89+
}
90+
}
91+
92+
func TestMemoryRecursiveQueryPrefersChildSessionValues(t *testing.T) {
93+
if !strings.Contains(Queries, `SELECT DISTINCT ON (memory."key")`) {
94+
t.Fatalf("expected recursive memory query to de-duplicate by key")
95+
}
96+
if !strings.Contains(Queries, `session_tree.depth ASC`) {
97+
t.Fatalf("expected recursive memory query to prioritize nearest child session")
98+
}
99+
}
100+
101+
func TestMemorySelectorSelectUsesRecursiveSessionLookup(t *testing.T) {
102+
session := uuid.New()
103+
bind := pg.NewBind("memory.select", "SELECT")
104+
105+
query, err := (MemorySelector{Session: session, Key: "topic"}).Select(bind, pg.Get)
106+
if err != nil {
107+
t.Fatal(err)
108+
}
109+
if query != "SELECT" {
110+
t.Fatalf("unexpected query: %q", query)
111+
}
112+
if got := bind.Get("session"); got != session {
113+
t.Fatalf("unexpected session binding: %v", got)
114+
}
115+
if got := bind.Get("key"); got != "topic" {
116+
t.Fatalf("unexpected key binding: %v", got)
117+
}
118+
if !strings.Contains(Queries, `ORDER BY
119+
session_tree.depth ASC`) {
120+
t.Fatalf("expected memory.select to prefer nearest child session")
121+
}
122+
if !strings.Contains(Queries, `LIMIT 1;`) {
123+
t.Fatalf("expected memory.select to return a single effective value")
124+
}
125+
}

‎memory/schema/queries.sql‎

Lines changed: 66 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,10 @@ INSERT INTO ${"schema"}.memory (
44
) VALUES (
55
@session, @key, @value, NOW()
66
)
7+
ON CONFLICT ("session", "key") DO UPDATE
8+
SET
9+
"value" = EXCLUDED."value",
10+
modified_at = NOW()
711
RETURNING
812
"session",
913
"key",
@@ -12,14 +16,75 @@ RETURNING
1216
modified_at;
1317

1418
-- memory.select
19+
WITH RECURSIVE session_tree AS (
20+
SELECT
21+
session.id,
22+
session.parent,
23+
0 AS depth
24+
FROM ${"llm_schema"}.session AS session
25+
WHERE session.id = @session
26+
UNION ALL
27+
SELECT
28+
parent.id,
29+
parent.parent,
30+
session_tree.depth + 1
31+
FROM ${"llm_schema"}.session AS parent
32+
INNER JOIN session_tree ON session_tree.parent = parent.id
33+
)
1534
SELECT
1635
memory."session",
1736
memory."key",
1837
memory."value",
1938
memory.created_at,
2039
memory.modified_at
2140
FROM ${"schema"}.memory AS memory
22-
WHERE memory."session" = @session AND memory."key" = @key;
41+
INNER JOIN session_tree ON session_tree.id = memory."session"
42+
WHERE memory."key" = @key
43+
ORDER BY
44+
session_tree.depth ASC,
45+
COALESCE(memory.modified_at, memory.created_at) DESC,
46+
memory.created_at DESC
47+
LIMIT 1;
48+
49+
-- memory.list_recursive
50+
WITH RECURSIVE session_tree AS (
51+
SELECT
52+
session.id,
53+
session.parent,
54+
0 AS depth
55+
FROM ${"llm_schema"}.session AS session
56+
WHERE session.id = @session
57+
UNION ALL
58+
SELECT
59+
parent.id,
60+
parent.parent,
61+
session_tree.depth + 1
62+
FROM ${"llm_schema"}.session AS parent
63+
INNER JOIN session_tree ON session_tree.parent = parent.id
64+
), effective_memory AS (
65+
SELECT DISTINCT ON (memory."key")
66+
memory."session",
67+
memory."key",
68+
memory."value",
69+
memory.created_at,
70+
memory.modified_at
71+
FROM ${"schema"}.memory AS memory
72+
INNER JOIN session_tree ON session_tree.id = memory."session"
73+
ORDER BY
74+
memory."key" ASC,
75+
session_tree.depth ASC,
76+
COALESCE(memory.modified_at, memory.created_at) DESC,
77+
memory.created_at DESC
78+
)
79+
SELECT
80+
memory."session",
81+
memory."key",
82+
memory."value",
83+
memory.created_at,
84+
memory.modified_at
85+
FROM effective_memory AS memory
86+
${where}
87+
${orderby}
2388

2489
-- memory.list
2590
SELECT

0 commit comments

Comments
 (0)