diff --git a/agent/src/index.ts b/agent/src/index.ts index 70a62da6e85..c2912b08cc6 100644 --- a/agent/src/index.ts +++ b/agent/src/index.ts @@ -303,6 +303,7 @@ function intializeDbCache(character: Character, db: IDatabaseCacheAdapter) { } async function startAgent(character: Character, directClient) { + let db: IDatabaseAdapter & IDatabaseCacheAdapter; try { character.id ??= stringToUuid(character.name); character.username ??= character.name; @@ -314,7 +315,7 @@ async function startAgent(character: Character, directClient) { fs.mkdirSync(dataDir, { recursive: true }); } - const db = initializeDatabase(dataDir); + db = initializeDatabase(dataDir); await db.init(); @@ -334,6 +335,9 @@ async function startAgent(character: Character, directClient) { error ); console.error(error); + if (db) { + await db.close(); + } throw error; } } diff --git a/packages/adapter-postgres/src/index.ts b/packages/adapter-postgres/src/index.ts index 3b84b2caf3a..ddd89458f9a 100644 --- a/packages/adapter-postgres/src/index.ts +++ b/packages/adapter-postgres/src/index.ts @@ -113,6 +113,10 @@ export class PostgresDatabaseAdapter await this.query(schema); } + async close() { + await this.pool.end(); + } + async testConnection(): Promise { let client; try { diff --git a/packages/adapter-sqlite/src/index.ts b/packages/adapter-sqlite/src/index.ts index 33b30385a82..3518eee2e03 100644 --- a/packages/adapter-sqlite/src/index.ts +++ b/packages/adapter-sqlite/src/index.ts @@ -80,6 +80,10 @@ export class SqliteDatabaseAdapter this.db.exec(sqliteTables); } + async close() { + this.db.close(); + } + async getAccountById(userId: UUID): Promise { const sql = "SELECT * FROM accounts WHERE id = ?"; const account = this.db.prepare(sql).get(userId) as Account; diff --git a/packages/adapter-sqljs/src/index.ts b/packages/adapter-sqljs/src/index.ts index a7c6588c54c..56db333ebb7 100644 --- a/packages/adapter-sqljs/src/index.ts +++ b/packages/adapter-sqljs/src/index.ts @@ -18,7 +18,8 @@ import { Database } from "./types.ts"; export class SqlJsDatabaseAdapter extends DatabaseAdapter - implements IDatabaseCacheAdapter { + implements IDatabaseCacheAdapter +{ constructor(db: Database) { super(); this.db = db; @@ -28,6 +29,10 @@ export class SqlJsDatabaseAdapter this.db.exec(sqliteTables); } + async close() { + this.db.close(); + } + async getRoom(roomId: UUID): Promise { const sql = "SELECT id FROM rooms WHERE id = ?"; const stmt = this.db.prepare(sql); @@ -77,10 +82,14 @@ export class SqlJsDatabaseAdapter const placeholders = params.roomIds.map(() => "?").join(", "); const sql = `SELECT * FROM memories WHERE 'type' = ? AND agentId = ? AND roomId IN (${placeholders})`; const stmt = this.db.prepare(sql); - const queryParams = [params.tableName, params.agentId, ...params.roomIds]; - console.log({ queryParams }) + const queryParams = [ + params.tableName, + params.agentId, + ...params.roomIds, + ]; + console.log({ queryParams }); stmt.bind(queryParams); - console.log({ queryParams }) + console.log({ queryParams }); const memories: Memory[] = []; while (stmt.step()) { diff --git a/packages/adapter-supabase/src/index.ts b/packages/adapter-supabase/src/index.ts index 1bef8a5069e..321d9e27c46 100644 --- a/packages/adapter-supabase/src/index.ts +++ b/packages/adapter-supabase/src/index.ts @@ -100,6 +100,14 @@ export class SupabaseDatabaseAdapter extends DatabaseAdapter { this.supabase = createClient(supabaseUrl, supabaseKey); } + async init() { + // noop + } + + async close() { + // noop + } + async getMemoriesByRoomIds(params: { roomIds: UUID[]; agentId?: UUID; diff --git a/packages/core/src/database.ts b/packages/core/src/database.ts index fb8f9555b31..01c80b91dc8 100644 --- a/packages/core/src/database.ts +++ b/packages/core/src/database.ts @@ -19,6 +19,19 @@ export abstract class DatabaseAdapter implements IDatabaseAdapter { * The database instance. */ db: DB; + + /** + * Optional initialization method for the database adapter. + * @returns A Promise that resolves when initialization is complete. + */ + abstract init(): Promise; + + /** + * Optional close method for the database adapter. + * @returns A Promise that resolves when closing is complete. + */ + abstract close(): Promise; + /** * Retrieves an account by its ID. * @param userId The UUID of the user account to retrieve. diff --git a/packages/core/src/types.ts b/packages/core/src/types.ts index fd5f65b50cb..6c08d9b0464 100644 --- a/packages/core/src/types.ts +++ b/packages/core/src/types.ts @@ -297,9 +297,9 @@ export interface State { formattedConversation?: string; /** Optional formatted knowledge */ - knowledge?: string, + knowledge?: string; /** Optional knowledge data */ - knowledgeData?: KnowledgeItem[], + knowledgeData?: KnowledgeItem[]; /** Additional dynamic properties */ [key: string]: unknown; @@ -714,7 +714,10 @@ export interface IDatabaseAdapter { db: any; /** Optional initialization */ - init?(): Promise; + init(): Promise; + + /** Close database connection */ + close(): Promise; /** Get account by ID */ getAccountById(userId: UUID): Promise;