Merge pull request #2135 from odilitime/pr1531

refactor: Optimize memory fetching by moving sorting and slicing to DB (PR #1531 remake)

Author: odilitime

packages/adapter-postgres/src/index.ts

@@ -329,6 +329,7 @@ export class PostgresDatabaseAdapter
         roomIds: UUID[];
         agentId?: UUID;
         tableName: string;
+        limit?: number;
     }): Promise<Memory[]> {
         return this.withDatabase(async () => {
             if (params.roomIds.length === 0) return [];
@@ -344,6 +345,13 @@ export class PostgresDatabaseAdapter
                 queryParams = [...queryParams, params.agentId];
             }
 
+            // Add sorting, and conditionally add LIMIT if provided
+            query += ` ORDER BY "createdAt" DESC`;
+            if (params.limit) {
+                query += ` LIMIT $${queryParams.length + 1}`;
+                queryParams.push(params.limit.toString());
+            }
+
             const { rows } = await this.pool.query(query, queryParams);
             return rows.map((row) => ({
                 ...row,

packages/client-direct/src/index.ts

@@ -571,7 +571,7 @@ export class DirectClient {
                                     memory,
                                     [responseMessage],
                                     state,
-                                    async (newMessages) => {
+                                    async (_newMessages) => {
                                         // FIXME: this is supposed override what the LLM said/decided
                                         // but the promise doesn't make this possible
                                         //message = newMessages;

packages/core/src/database.ts

@@ -95,6 +95,7 @@ export abstract class DatabaseAdapter<DB = any> implements IDatabaseAdapter {
         agentId: UUID;
         roomIds: UUID[];
         tableName: string;
+        limit?: number;
     }): Promise<Memory[]>;
 
     abstract getMemoryById(id: UUID): Promise<Memory | null>;

packages/core/src/memory.ts

@@ -189,11 +189,12 @@ export class MemoryManager implements IMemoryManager {
         );
     }
 
-    async getMemoriesByRoomIds(params: { roomIds: UUID[] }): Promise<Memory[]> {
+    async getMemoriesByRoomIds(params: { roomIds: UUID[], limit?: number; }): Promise<Memory[]> {
         return await this.runtime.databaseAdapter.getMemoriesByRoomIds({
             tableName: this.tableName,
             agentId: this.runtime.agentId,
             roomIds: params.roomIds,
+            limit: params.limit
         });
     }
 

packages/core/src/runtime.ts

@@ -1099,21 +1099,11 @@ Text: ${attachment.text}
             ]);
 
             // Check the existing memories in the database
-            const existingMemories =
-                await this.messageManager.getMemoriesByRoomIds({
+            return this.messageManager.getMemoriesByRoomIds({
                     // filter out the current room id from rooms
                     roomIds: rooms.filter((room) => room !== roomId),
+                    limit: 20
                 });
-
-            // Sort messages by timestamp in descending order
-            existingMemories.sort(
-                (a, b) =>
-                    (b?.createdAt ?? Date.now()) - (a?.createdAt ?? Date.now())
-            );
-
-            // Take the most recent messages
-            const recentInteractionsData = existingMemories.slice(0, 20);
-            return recentInteractionsData;
         };
 
         const recentInteractions =

packages/core/src/types.ts

@@ -906,6 +906,7 @@ export interface IDatabaseAdapter {
         tableName: string;
         agentId: UUID;
         roomIds: UUID[];
+        limit?: number;
     }): Promise<Memory[]>;
 
     getCachedEmbeddings(params: {
@@ -1079,7 +1080,7 @@ export interface IMemoryManager {
     ): Promise<{ embedding: number[]; levenshtein_score: number }[]>;
 
     getMemoryById(id: UUID): Promise<Memory | null>;
-    getMemoriesByRoomIds(params: { roomIds: UUID[] }): Promise<Memory[]>;
+    getMemoriesByRoomIds(params: { roomIds: UUID[], limit?: number }): Promise<Memory[]>;
     searchMemoriesByEmbedding(
         embedding: number[],
         opts: {