diff --git a/CHANGELOG.md b/CHANGELOG.md index 42e6fe8..81d13d5 100644 --- a/CHANGELOG.md +++ b/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) diff --git a/package.json b/package.json index 8ce22bb..3fd0d2f 100644 --- a/package.json +++ b/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": { diff --git a/src/evals/index.ts b/src/evals/index.ts index fd770f6..986f099 100644 --- a/src/evals/index.ts +++ b/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 }) diff --git a/src/evals/multi-turn/index.ts b/src/evals/multi-turn/index.ts index e13d561..06e8ba8 100644 --- a/src/evals/multi-turn/index.ts +++ b/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[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 diff --git a/src/evals/one-shot/index.ts b/src/evals/one-shot/index.ts index b152967..4a70b93 100644 --- a/src/evals/one-shot/index.ts +++ b/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[0]; db: Database }) => { +export const runOneShotExamples = async ({ db, model }: { model: LLM; db: Database }) => { const memory = new Memory({ model, db }) let totalFacts = 0 diff --git a/src/index.ts b/src/index.ts index 4059df6..bba48a2 100644 --- a/src/index.ts +++ b/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[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[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 }) } } diff --git a/src/types/index.ts b/src/types/index.ts index 6a01e8d..a23f3f8 100644 --- a/src/types/index.ts +++ b/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 } -export type Database = BunDatabase +export type Database = { + execute: (cmd: string) => Promise +} + +export type LLM = Parameters[0]