always limit fetchMessagesAfterSummaryPoint to memoryWindow (#238)

getzep · Oct 19, 2023 · 04d4698 · 04d4698
1 parent 6b6627d
commit 04d4698
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/pkg/store/postgres/memorystore_test.go b/pkg/store/postgres/memorystore_test.go
@@ -242,10 +242,10 @@ func TestGetMessages(t *testing.T) {
 		withSummary    bool
 	}{
 		{
-			name:           "Get all messages",
+			name:           "Get all messages within messageWindow",
 			sessionID:      sessionID,
 			lastNMessages:  0,
-			expectedLength: len(messages),
+			expectedLength: messageWindow,
 			withSummary:    false,
 		},
 		{

diff --git a/pkg/store/postgres/messages.go b/pkg/store/postgres/messages.go
@@ -176,7 +176,7 @@ func getMessages(
 	if lastNMessages > 0 {
 		messages, err = fetchLastNMessages(ctx, db, sessionID, lastNMessages)
 	} else {
-		messages, err = fetchMessagesAfterSummaryPoint(ctx, db, sessionID, summary)
+		messages, err = fetchMessagesAfterSummaryPoint(ctx, db, sessionID, summary, memoryWindow)
 	}
 	if err != nil {
 		return nil, store.NewStorageError("failed to get messages", err)
@@ -201,6 +201,7 @@ func fetchMessagesAfterSummaryPoint(
 	db *bun.DB,
 	sessionID string,
 	summary *models.Summary,
+	memoryWindow int,
 ) ([]MessageStoreSchema, error) {
 	var summaryPointIndex int64
 	var err error
@@ -221,6 +222,9 @@ func fetchMessagesAfterSummaryPoint(
 		query.Where("id > ?", summaryPointIndex)
 	}
 
+	// Always limit to the memory window
+	query.Limit(memoryWindow)
+
 	return messages, query.Scan(ctx)
 }