feat: improve performance by moving limit and sorting to db query

iotexproject · Dec 27, 2024 · 4c31dc3 · 4c31dc3
1 parent 3f0ad47
commit 4c31dc3
Show file tree

Hide file tree

Showing 5 changed files with 18 additions and 10 deletions.
diff --git a/packages/adapter-postgres/src/index.ts b/packages/adapter-postgres/src/index.ts
@@ -286,6 +286,7 @@ export class PostgresDatabaseAdapter
         roomIds: UUID[];
         agentId?: UUID;
         tableName: string;
+        limit?: number;
     }): Promise<Memory[]> {
         return this.withDatabase(async () => {
             if (params.roomIds.length === 0) return [];
@@ -301,6 +302,15 @@ 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());
+            }
+
+            console.log(query, queryParams);
+
             const { rows } = await this.pool.query(query, queryParams);
             return rows.map((row) => ({
                 ...row,

diff --git a/packages/core/src/database.ts b/packages/core/src/database.ts
@@ -94,6 +94,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>;

diff --git a/packages/core/src/memory.ts b/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
         });
     }
 

diff --git a/packages/core/src/runtime.ts b/packages/core/src/runtime.ts
@@ -923,19 +923,14 @@ Text: ${attachment.text}
             // Check the existing memories in the database
             console.time("existing-memories");
             console.timeLog("existing-memories", rooms.filter((room) => room !== roomId));
-            const existingMemories =
+            const recentInteractionsData =
             await this.messageManager.getMemoriesByRoomIds({
                 // filter out the current room id from rooms
                 roomIds: rooms.filter((room) => room !== roomId),
+                limit: 20
             });
-            console.timeLog("existing-memories", existingMemories.length);
+            console.timeLog("existing-memories", recentInteractionsData.length);
             console.timeEnd("existing-memories");
-
-            // Sort messages by timestamp in descending order
-            existingMemories.sort((a, b) => b.createdAt - a.createdAt);
-
-            // Take the most recent messages
-            const recentInteractionsData = existingMemories.slice(0, 20);
             return recentInteractionsData;
         };
 

diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts
@@ -816,6 +816,7 @@ export interface IDatabaseAdapter {
         tableName: string;
         agentId: UUID;
         roomIds: UUID[];
+        limit?: number;
     }): Promise<Memory[]>;
 
     getCachedEmbeddings(params: {
@@ -969,7 +970,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: {