Skip to content

Commit 60ce697

Browse files
terceltercel
committed
feat: Add pagination support and account status management across database adapters
1 parent 81a3528 commit 60ce697

File tree

19 files changed

+745
-30
lines changed

19 files changed

+745
-30
lines changed

‎agent/src/index.ts

+16-17
Original file line numberDiff line numberDiff line change
@@ -42,6 +42,7 @@ import {
4242
type ICacheManager,
4343
type IDatabaseAdapter,
4444
type IDatabaseCacheAdapter,
45+
type TypeDatabaseAdapter,
4546
ModelProviderName,
4647
parseBooleanFromText,
4748
settings,
@@ -1389,24 +1390,14 @@ function initializeCache(
13891390

13901391
async function startAgent(
13911392
character: Character,
1392-
directClient: DirectClient
1393+
directClient: DirectClient,
1394+
db: TypeDatabaseAdapter
13931395
): Promise<AgentRuntime> {
1394-
let db: IDatabaseAdapter & IDatabaseCacheAdapter;
13951396
try {
13961397
character.id ??= stringToUuid(character.name);
13971398
character.username ??= character.name;
13981399

13991400
const token = getTokenForProvider(character.modelProvider, character);
1400-
const dataDir = path.join(__dirname, "../data");
1401-
1402-
if (!fs.existsSync(dataDir)) {
1403-
fs.mkdirSync(dataDir, { recursive: true });
1404-
}
1405-
1406-
db = initializeDatabase(dataDir) as IDatabaseAdapter &
1407-
IDatabaseCacheAdapter;
1408-
1409-
await db.init();
14101401

14111402
const cache = initializeCache(
14121403
process.env.CACHE_STORE ?? CacheStore.DATABASE,
@@ -1440,9 +1431,6 @@ async function startAgent(
14401431
error
14411432
);
14421433
elizaLogger.error(error);
1443-
if (db) {
1444-
await db.close();
1445-
}
14461434
throw error;
14471435
}
14481436
}
@@ -1476,6 +1464,16 @@ const startAgents = async () => {
14761464
let serverPort = Number.parseInt(settings.SERVER_PORT || "3000");
14771465
const args = parseArguments();
14781466
const charactersArg = args.characters || args.character;
1467+
const dataDir = path.join(__dirname, "../data");
1468+
1469+
if (!fs.existsSync(dataDir)) {
1470+
fs.mkdirSync(dataDir, { recursive: true });
1471+
}
1472+
1473+
const db = initializeDatabase(dataDir) as TypeDatabaseAdapter;
1474+
1475+
await db.init();
1476+
14791477
let characters = [defaultCharacter];
14801478

14811479
if (process.env.IQ_WALLET_ADDRESS && process.env.IQSOlRPC) {
@@ -1493,7 +1491,7 @@ const startAgents = async () => {
14931491

14941492
try {
14951493
for (const character of characters) {
1496-
await startAgent(character, directClient);
1494+
await startAgent(character, directClient, db);
14971495
}
14981496
} catch (error) {
14991497
elizaLogger.error("Error starting agents:", error);
@@ -1513,11 +1511,12 @@ const startAgents = async () => {
15131511
character.plugins = await handlePluginImporting(character.plugins);
15141512

15151513
// wrap it so we don't have to inject directClient later
1516-
return startAgent(character, directClient);
1514+
return startAgent(character, directClient, db);
15171515
};
15181516

15191517
directClient.loadCharacterTryPath = loadCharacterTryPath;
15201518
directClient.jsonToCharacter = jsonToCharacter;
1519+
directClient.db = db;
15211520

15221521
directClient.start(serverPort);
15231522

‎packages/adapter-mongodb/src/index.ts

+62
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,8 @@ import {
1111
type Memory,
1212
type Relationship,
1313
type UUID, elizaLogger,
14+
type PaginationParams,
15+
type PaginationResult,
1416
} from "@elizaos/core";
1517
import { v4 } from "uuid";
1618

@@ -322,6 +324,11 @@ export class MongoDBDatabaseAdapter
322324
}
323325
}
324326

327+
async updateAccount(account: Account): Promise<void> {
328+
await this.ensureConnection();
329+
await this.database.collection('accounts').updateOne({ id: account.id }, { $set: account });
330+
}
331+
325332
async getActorDetails(params: { roomId: UUID }): Promise<Actor[]> {
326333
await this.ensureConnection();
327334
const actors = await this.database.collection('participants')
@@ -1440,6 +1447,61 @@ export class MongoDBDatabaseAdapter
14401447
return [];
14411448
}
14421449
}
1450+
async paginate(collectionName: string, params: PaginationParams): Promise<PaginationResult> {
1451+
await this.ensureConnection();
1452+
1453+
const {
1454+
page = 1,
1455+
pageSize = 10,
1456+
where = {},
1457+
order = { createdAt: -1 } // MongoDB uses 1 for ASC, -1 for DESC
1458+
} = params;
1459+
1460+
const skip = (page - 1) * pageSize;
1461+
1462+
// Convert where conditions for MongoDB
1463+
const whereQuery: any = {};
1464+
// biome-ignore lint/complexity/noForEach: <explanation>
1465+
Object.entries(where).forEach(([key, value]) => {
1466+
if (value === null || value === undefined) return;
1467+
1468+
if (typeof value === 'object') {
1469+
if (key === 'createdAt') {
1470+
whereQuery[key] = {};
1471+
if (value.gte) whereQuery[key].$gte = new Date(value.gte);
1472+
if (value.lte) whereQuery[key].$lte = new Date(value.lte);
1473+
} else {
1474+
whereQuery[key] = value;
1475+
}
1476+
} else {
1477+
whereQuery[key] = value;
1478+
}
1479+
});
14431480

1481+
try {
1482+
// Get total count
1483+
const total = await this.database.collection(collectionName)
1484+
.countDocuments(whereQuery);
1485+
1486+
// Get paginated data
1487+
const list = await this.database.collection(collectionName)
1488+
.find(whereQuery)
1489+
.sort(order)
1490+
.skip(skip)
1491+
.limit(pageSize)
1492+
.toArray();
1493+
1494+
return {
1495+
list,
1496+
total,
1497+
page,
1498+
pageSize,
1499+
totalPages: Math.ceil(total / pageSize),
1500+
};
1501+
} catch (error) {
1502+
elizaLogger.error(`Error in paginate for collection ${collectionName}:`, error);
1503+
throw error;
1504+
}
1505+
}
14441506
}
14451507

‎packages/adapter-pglite/src/index.ts

+121-2
Original file line numberDiff line numberDiff line change
@@ -15,6 +15,8 @@ import {
1515
DatabaseAdapter,
1616
EmbeddingProvider,
1717
type RAGKnowledgeItem,
18+
type PaginationParams,
19+
type PaginationResult,
1820
} from "@elizaos/core";
1921
import fs from "fs";
2022
import { fileURLToPath } from "url";
@@ -241,15 +243,16 @@ export class PGLiteDatabaseAdapter
241243
try {
242244
const accountId = account.id ?? v4();
243245
await this.query(
244-
`INSERT INTO accounts (id, name, username, email, "avatarUrl", details)
245-
VALUES ($1, $2, $3, $4, $5, $6)`,
246+
`INSERT INTO accounts (id, name, username, email, "avatarUrl", details, status)
247+
VALUES ($1, $2, $3, $4, $5, $6, $7)`,
246248
[
247249
accountId,
248250
account.name,
249251
account.username || "",
250252
account.email || "",
251253
account.avatarUrl || "",
252254
JSON.stringify(account.details),
255+
account.status || "paused",
253256
]
254257
);
255258
elizaLogger.debug("Account created successfully:", {
@@ -268,6 +271,33 @@ export class PGLiteDatabaseAdapter
268271
}, "createAccount");
269272
}
270273

274+
async updateAccount(account: Account): Promise<void> {
275+
return this.withDatabase(async () => {
276+
try {
277+
await this.query(
278+
`UPDATE accounts SET name = $1, username = $2, email = $3, "avatarUrl" = $4, status = $5, details = $6 WHERE id = $7`,
279+
[
280+
account.name,
281+
account.username,
282+
account.email,
283+
account.avatarUrl,
284+
account.status,
285+
JSON.stringify(account.details),
286+
account.id,
287+
]
288+
);
289+
} catch (error) {
290+
elizaLogger.error("Failed to update accounts:", {
291+
account: account.id,
292+
error:
293+
error instanceof Error ? error.message : String(error),
294+
status: account.status,
295+
});
296+
throw error;
297+
}
298+
}, "updateAccount");
299+
}
300+
271301
async getActorById(params: { roomId: UUID }): Promise<Actor[]> {
272302
return this.withDatabase(async () => {
273303
const { rows } = await this.query<Actor>(
@@ -1558,6 +1588,95 @@ export class PGLiteDatabaseAdapter
15581588
}
15591589
}, "clearKnowledge");
15601590
}
1591+
1592+
async paginate(table: string, params: PaginationParams): Promise<PaginationResult> {
1593+
return this.withDatabase(async () => {
1594+
const {
1595+
page = 1,
1596+
pageSize = 10,
1597+
where = {},
1598+
order = { createdAt: 'DESC' }
1599+
} = params;
1600+
1601+
const offset = (page - 1) * pageSize;
1602+
const whereClause: string[] = [];
1603+
const values: any[] = [];
1604+
let paramCount = 0;
1605+
1606+
// Build where clause with proper parameter indexing
1607+
// biome-ignore lint/complexity/noForEach: <explanation>
1608+
Object.entries(where).forEach(([key, value]) => {
1609+
if (value === null || value === undefined) return;
1610+
1611+
if (typeof value === 'object') {
1612+
if (key === 'createdAt') {
1613+
if (value.gte) {
1614+
paramCount++;
1615+
whereClause.push(`"${key}" >= $${paramCount}`);
1616+
values.push(value.gte);
1617+
}
1618+
if (value.lte) {
1619+
paramCount++;
1620+
whereClause.push(`"${key}" <= $${paramCount}`);
1621+
values.push(value.lte);
1622+
}
1623+
}
1624+
} else {
1625+
paramCount++;
1626+
whereClause.push(`"${key}" = $${paramCount}`);
1627+
values.push(value);
1628+
}
1629+
});
1630+
1631+
const whereStr = whereClause.length > 0
1632+
? `WHERE ${whereClause.join(' AND ')}`
1633+
: '';
1634+
1635+
const orderClause = Object.entries(order)
1636+
.map(([key, direction]) => `"${key}" ${direction}`)
1637+
.join(', ');
1638+
1639+
// Count total records
1640+
const countQuery = `
1641+
SELECT COUNT(*) as total
1642+
FROM "${table}"
1643+
${whereStr}
1644+
`;
1645+
1646+
const { rows: countRows } = await this.query<{ total: number }>(
1647+
countQuery,
1648+
values
1649+
);
1650+
const total = Number(countRows[0]?.total || 0);
1651+
1652+
// Get paginated data
1653+
paramCount++;
1654+
const limitParam = `$${paramCount}`;
1655+
paramCount++;
1656+
const offsetParam = `$${paramCount}`;
1657+
1658+
const dataQuery = `
1659+
SELECT *
1660+
FROM "${table}"
1661+
${whereStr}
1662+
ORDER BY ${orderClause}
1663+
LIMIT ${limitParam} OFFSET ${offsetParam}
1664+
`;
1665+
1666+
const { rows: list } = await this.query(
1667+
dataQuery,
1668+
[...values, pageSize, offset]
1669+
);
1670+
1671+
return {
1672+
list,
1673+
total,
1674+
page,
1675+
pageSize,
1676+
totalPages: Math.ceil(total / pageSize),
1677+
};
1678+
}, "paginate");
1679+
}
15611680
}
15621681

15631682
export default PGLiteDatabaseAdapter;

‎packages/adapter-postgres/schema.sql

+2-1
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,8 @@ CREATE TABLE IF NOT EXISTS accounts (
4343
"username" TEXT,
4444
"email" TEXT NOT NULL,
4545
"avatarUrl" TEXT,
46-
"details" JSONB DEFAULT '{}'::jsonb
46+
"details" JSONB DEFAULT '{}'::jsonb,
47+
"status" TEXT
4748
);
4849

4950
CREATE TABLE IF NOT EXISTS rooms (

0 commit comments

Comments
 (0)