Support more DBs. Simplify types.

CHANGELOG.md

@@ -1,3 +1,4 @@
+- [2024-10-10] [Support more DBs. Simplify types.](https://github.com/RubricLab/memory/commit/c9500b038646adac4a4250b9956914bafbb53bea)
 - [2024-10-10] [Add main to pkg](https://github.com/RubricLab/memory/commit/ef8082790fbd93c6ae2b4b47112bc1986dd8f0ab)
 - [2024-10-10] [Move DB to top-level](https://github.com/RubricLab/memory/commit/67913408246413448008f902963eb5a288b963d8)
 - [2024-10-10] [Decouple extraction and save](https://github.com/RubricLab/memory/commit/d5fd7c08eff9458b9775ef439b480cd38922759d)

package.json

@@ -2,7 +2,7 @@
   "name": "@rubriclab/memory",
   "module": "src/index.ts",
   "main": "src/index.ts",
-  "version": "0.0.14",
+  "version": "0.0.15",
   "private": false,
   "type": "module",
   "devDependencies": {

src/evals/index.ts

@@ -6,7 +6,7 @@ import { runOneShotExamples } from '@/evals/one-shot'
 const args = parseArgs({
 	args: Bun.argv,
 	options: {
-		fast: {
+		sota: {
 			type: 'boolean',
 			default: false
 		},
@@ -24,27 +24,29 @@ const args = parseArgs({
 })
 
 if (import.meta.path === Bun.main) {
-	const { help, fast, dataset } = args.values
+	const { help, sota, dataset } = args.values
 	if (help) {
 		console.log(`
   Usage: bun evals/index.ts [options]
 
   Options:
-    --fast    Use gpt-4o-mini instead of gpt-4o-2024-08-06
+    --sota    Use gpt-4o-2024-08-06 instead of gpt-4o-mini
     --help    Show this help message
 `)
 		process.exit(0)
 	}
 
-	const model = fast ? 'gpt-4o-mini' : 'gpt-4o-2024-08-06'
+	const model = sota ? 'gpt-4o-2024-08-06' : 'gpt-4o-mini'
 
-	const db = new Database(':memory:', { create: true, strict: true })
+	const bunDB = new Database(':memory:', { create: true, strict: true })
 
-	await db
-		.prepare(
-			'create table if not exists facts (subject text, relation text, object text, primary key (subject, object))'
-		)
-		.get()
+	const db = {
+		execute: async (cmd: string) => await bunDB.prepare(cmd).get()
+	}
+
+	await db.execute(
+		'create table if not exists facts (subject text, relation text, object text, primary key (subject, object))'
+	)
 
 	if (dataset === '1') {
 		await runOneShotExamples({ db, model })

src/evals/multi-turn/index.ts

@@ -1,14 +1,10 @@
 import { Memory } from '@/index'
-import type { Database, Fact } from '@/types'
+import type { Database, Fact, LLM } from '@/types'
 import { format } from '@/utils/string'
-import type { openai } from '@ai-sdk/openai'
 import chalk from 'chalk'
 import { EXAMPLES } from './examples'
 
-export const runMultiTurnExamples = async ({
-	db,
-	model
-}: { model: Parameters<typeof openai>[0]; db: Database }) => {
+export const runMultiTurnExamples = async ({ db, model }: { model: LLM; db: Database }) => {
 	const memory = new Memory({ model, db })
 
 	let totalFacts = 0
@@ -34,7 +30,7 @@ export const runMultiTurnExamples = async ({
 					`\n🎯 ${i + 1} of ${message.facts.length}: ${chalk.magenta(fact.subject)} ${chalk.yellow(fact.relation)} ${chalk.blue(fact.object)}`
 				)
 
-				const newFacts = db.query('select * from facts').all()
+				const newFacts = (await db.execute('select * from facts')) as Fact[]
 
 				for (const [k, newFact] of newFacts.entries()) {
 					const { subject, relation, object } = newFact as Fact

src/evals/one-shot/index.ts

@@ -1,14 +1,10 @@
 import { Memory } from '@/index'
-import type { Database } from '@/types'
+import type { Database, LLM } from '@/types'
 import { format } from '@/utils/string'
-import type { openai } from '@ai-sdk/openai'
 import chalk from 'chalk'
 import { EXAMPLES } from './examples'
 
-export const runOneShotExamples = async ({
-	db,
-	model
-}: { model: Parameters<typeof openai>[0]; db: Database }) => {
+export const runOneShotExamples = async ({ db, model }: { model: LLM; db: Database }) => {
 	const memory = new Memory({ model, db })
 
 	let totalFacts = 0

src/index.ts

@@ -1,27 +1,24 @@
-import type { Database } from 'bun:sqlite'
 import { clean } from '@/utils/string'
 import { openai } from '@ai-sdk/openai'
 import { generateObject } from 'ai'
 import { z } from 'zod'
-import type { Fact } from './types'
+import type { Database, Fact, LLM } from './types'
 
 export class Memory {
-	model: Parameters<typeof openai>[0]
+	model: LLM
 	db: Database
 
 	async createTable() {
-		await this.db
-			.prepare(
-				'create table if not exists facts (subject text, relation text, object text, primary key (subject, object))'
-			)
-			.get()
+		await this.db.execute(
+			'create table if not exists facts (subject text, relation text, object text, primary key (subject, object))'
+		)
 	}
 
 	constructor({
-		model,
+		model = 'gpt-4o-mini',
 		db
 	}: {
-		model: Parameters<typeof openai>[0]
+		model?: LLM
 		db: Database
 	}) {
 		this.model = model
@@ -38,7 +35,7 @@ export class Memory {
 			schema: z.object({
 				entities: z.array(
 					z.object({
-						subject: z.string()
+						name: z.string()
 					})
 				)
 			}),
@@ -48,13 +45,11 @@ export class Memory {
             "${content}"`
 		})
 
-		const tags = entities?.map(({ subject }) => `"${subject}"`).join(', ') || ''
+		const tags = entities?.map(({ name }) => `"${name}"`).join(', ') || ''
 
-		const query = this.db.query(
+		const relevantFacts = (await this.db.execute(
 			`select * from facts where subject in (${tags}) or object in (${tags})`
-		)
-
-		const relevantFacts = query.all() as Fact[]
+		)) as Fact[]
 
 		const {
 			object: { facts }
@@ -91,16 +86,14 @@ export class Memory {
 		for await (const fact of facts) {
 			const { subject, relation, object } = fact
 
-			this.db
-				.prepare(`
+			await this.db.execute(`
 					insert into facts (subject, relation, object)
-					values ($1, $2, $3)
-					on conflict (subject, object) do update set relation = $2
+					values (${subject}, ${relation}, ${object})
+					on conflict (subject, object) do update set relation = ${relation}
 				`)
-				.run(subject, relation, object)
 		}
 
-		const priorFacts = this.db.query('select * from facts').all()
+		const priorFacts = await this.db.execute('select * from facts')
 		console.log({ priorFacts })
 	}
 }

src/types/index.ts

@@ -1,4 +1,4 @@
-import type { Database as BunDatabase } from 'bun:sqlite'
+import type { openai } from '@ai-sdk/openai'
 
 export type Fact = {
 	subject: string
@@ -7,4 +7,8 @@ export type Fact = {
 	data?: Record<string, string>
 }
 
-export type Database = BunDatabase
+export type Database = {
+	execute: (cmd: string) => Promise<unknown>
+}
+
+export type LLM = Parameters<typeof openai>[0]