Add typesafe SQL setup

.neon

@@ -0,0 +1,4 @@
+{
+	"projectId": "quiet-cloud-45113218",
+	"orgId": "org-damp-cherry-40821001"
+}

CHANGELOG.md

@@ -1,3 +1,4 @@
+- [2024-10-10] [Add typesafe SQL setup](https://github.com/RubricLab/memory/commit/9acd10a3be5e0335970ca8f551a03847e04fcdda)
 - [2024-10-10] [Use relative import](https://github.com/RubricLab/memory/commit/2161e8cc919ac63da4b092b61ff811566417bc08)
 - [2024-10-10] [Fix post-insert return](https://github.com/RubricLab/memory/commit/a79a936cff2baadbd5411caec354579b47d8c846)
 - [2024-10-10] [Support more DBs. Simplify types.](https://github.com/RubricLab/memory/commit/c9500b038646adac4a4250b9956914bafbb53bea)

package.json

@@ -2,7 +2,7 @@
   "name": "@rubriclab/memory",
   "module": "src/index.ts",
   "main": "src/index.ts",
-  "version": "0.0.17",
+  "version": "0.0.18",
   "private": false,
   "type": "module",
   "devDependencies": {
@@ -13,11 +13,14 @@
   },
   "dependencies": {
     "@ai-sdk/openai": "^0.0.66",
+    "@prisma/client": "^5.20.0",
     "@rubriclab/config": "*",
     "@rubriclab/package": "*",
     "@t3-oss/env-nextjs": "^0.11.1",
     "ai": "^3.4.9",
     "chalk": "^5.3.0",
+    "nanoid": "^5.0.7",
+    "prisma": "^5.20.0",
     "zod": "^3.23.8"
   },
   "simple-git-hooks": {
@@ -26,9 +29,13 @@
   "publishConfig": {
     "access": "public"
   },
+  "prisma": {
+    "schema": "./prisma/schema"
+  },
   "scripts": {
     "prepare": "bun x simple-git-hooks",
     "bleed": "bun x npm-check-updates -u",
+    "db:push": "prisma db push && prisma generate --sql",
     "clean": "rm -rf .next && rm -rf node_modules",
     "format": "bun x biome format --write .",
     "eval": "bun src/evals/index.ts",

prisma/schema/facts.prisma

@@ -0,0 +1,36 @@
+model fact {
+  id String @id @default(nanoid())
+
+  body String
+
+  tags relationship[]
+
+  createdAt DateTime @default(now())
+  updatedAt DateTime @default(now()) @updatedAt
+}
+
+model tag {
+  id String @id @default(nanoid())
+
+  body String @unique
+
+  vector Unsupported("vector(768)")?
+
+  facts relationship[]
+
+  createdAt DateTime @default(now())
+  updatedAt DateTime @default(now()) @updatedAt
+}
+
+model relationship {
+  id String @id @default(nanoid())
+
+  factId String
+  tagId  String
+
+  fact fact @relation(fields: [factId], references: [id], onDelete: Cascade)
+  tag  tag  @relation(fields: [tagId], references: [id], onDelete: Cascade)
+
+  createdAt DateTime @default(now())
+  updatedAt DateTime @default(now()) @updatedAt
+}

prisma/schema/schema.prisma

@@ -0,0 +1,10 @@
+generator client {
+  provider        = "prisma-client-js"
+  previewFeatures = ["postgresqlExtensions", "prismaSchemaFolder", "typedSql"]
+}
+
+datasource db {
+  provider   = "postgresql"
+  url        = env("DATABASE_URL")
+  extensions = [vector]
+}

prisma/sql/createVectorExtension.sql

@@ -0,0 +1 @@
+create extension if not exists vector;

prisma/sql/createVectorIndex.sql

@@ -0,0 +1 @@
+create index on tag using hnsw (vector vector_cosine_ops);

prisma/sql/insertVector.sql

@@ -0,0 +1,7 @@
+-- @param {String} $1:id
+-- @param {String} $2:body
+-- @param $3:vector
+insert into tag (id, body, vector)
+values ($1, $2, $3)
+on conflict (body) do nothing
+returning id

prisma/sql/searchVector.sql

@@ -0,0 +1,11 @@
+-- @param $1:vector
+-- @param {Int} $2:threshold
+-- @param {Float} $3:limit
+select id, body, similarity
+from (
+  select id, body, 1 - (vector <=> $1::vector) as similarity
+  from tag
+) as subquery
+where similarity > $2
+order by similarity desc
+limit $3;

prisma/sql/setMaxWorkers.sql

@@ -0,0 +1 @@
+set max_parallel_maintenance_workers = 7;

prisma/sql/setWorkerMemory.sql

@@ -0,0 +1 @@
+set maintenance_work_mem='10 GB';

src/evals/index.ts

@@ -1,7 +1,9 @@
-import Database from 'bun:sqlite'
 import { parseArgs } from 'node:util'
 import { runMultiTurnExamples } from '@/evals/multi-turn'
-import { runOneShotExamples } from '@/evals/one-shot'
+import type { Database } from '@/types'
+import { PrismaClient } from '@prisma/client'
+
+const db: Database = new PrismaClient()
 
 const args = parseArgs({
 	args: Bun.argv,
@@ -10,11 +12,7 @@ const args = parseArgs({
 			type: 'boolean',
 			default: false
 		},
-		dataset: {
-			type: 'string',
-			default: '1',
-			choices: ['1', '2']
-		},
+
 		help: {
 			type: 'boolean',
 			default: false
@@ -24,7 +22,7 @@ const args = parseArgs({
 })
 
 if (import.meta.path === Bun.main) {
-	const { help, sota, dataset } = args.values
+	const { help, sota } = args.values
 
 	if (help) {
 		console.log(`
@@ -39,19 +37,5 @@ if (import.meta.path === Bun.main) {
 
 	const model = sota ? 'gpt-4o-2024-08-06' : 'gpt-4o-mini'
 
-	const bunDB = new Database(':memory:', { create: true, strict: true })
-
-	const db = {
-		execute: async (cmd: string) => await bunDB.prepare(cmd).all()
-	}
-
-	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 })
-	} else if (dataset === '2') {
-		await runMultiTurnExamples({ db, model })
-	}
+	await runMultiTurnExamples({ db, model })
 }

src/evals/multi-turn/index.ts

@@ -1,6 +1,6 @@
 import { Memory } from '@/index'
-import type { Database, Fact, LLM } from '@/types'
-import { format } from '@/utils/string'
+import type { Database, LLM } from '@/types'
+// import { format } from '@/utils/string'
 import chalk from 'chalk'
 import { EXAMPLES } from './examples'
 
@@ -11,61 +11,61 @@ export const runMultiTurnExamples = async ({ db, model }: { model: LLM; db: Data
 	let totalRecall = 0
 	let totalAttempts = 0
 
-	for await (const eg of EXAMPLES) {
+	for await (const eg of EXAMPLES.slice(0, 1)) {
 		for await (const message of eg.messages) {
 			totalFacts += message.facts.length
 
 			console.log(chalk.yellow(`\n\n"${message.content}"`))
 
 			// Clean up DB in between conversations
 			const omitted: number[] = []
-			await db.execute('delete from facts')
+			// await db.fact.deleteMany()
 
-			const { facts: attempts } = await memory.insert({
+			const { facts: attempts } = await memory.ingest({
 				content: message.content
 			})
 
-			for (const [i, fact] of message.facts.entries()) {
-				let correctFacts = 0
+			// for (const [i, fact] of message.facts.entries()) {
+			let correctFacts = 0
 
-				console.log(
-					`\n🎯 ${i + 1} of ${message.facts.length}: ${chalk.magenta(fact.subject)} ${chalk.yellow(fact.relation)} ${chalk.blue(fact.object)}`
-				)
+			// 	console.log(
+			// 		`\n🎯 ${i + 1} of ${message.facts.length}: ${chalk.magenta(fact.subject)} ${chalk.yellow(fact.relation)} ${chalk.blue(fact.object)}`
+			// 	)
 
-				const newFacts = (await db.execute('select * from facts')) as Fact[]
+			// 	const newFacts = await db.fact.findMany()
 
-				for (const [k, newFact] of newFacts.entries()) {
-					const { subject, relation, object } = newFact as Fact
+			// 	for (const [k, newFact] of newFacts.entries()) {
+			// 		const { subject, relation, object } = newFact
 
-					const correctSubject = fact.subject === subject
-					const correctRelation = fact.relation === relation
-					const correctObject = fact.object === object
+			// 		const correctSubject = fact.subject === subject
+			// 		const correctRelation = fact.relation === relation
+			// 		const correctObject = fact.object === object
 
-					if (omitted.includes(k)) {
-						console.log(
-							chalk.blackBright.italic(
-								`🤖 ${k + 1} of ${newFacts.length}: ${subject} ${relation} ${object}`
-							)
-						)
-						continue
-					}
+			// 		if (omitted.includes(k)) {
+			// 			console.log(
+			// 				chalk.blackBright.italic(
+			// 					`🤖 ${k + 1} of ${newFacts.length}: ${subject} ${relation} ${object}`
+			// 				)
+			// 			)
+			// 			continue
+			// 		}
 
-					console.log(
-						`🤖 ${k + 1} of ${newFacts.length}: ${chalk.magenta(format(subject, correctSubject))} ${chalk.yellow(
-							format(relation, correctRelation)
-						)} ${chalk.blue(format(object, correctObject))}`
-					)
+			// 		console.log(
+			// 			`🤖 ${k + 1} of ${newFacts.length}: ${chalk.magenta(format(subject, correctSubject))} ${chalk.yellow(
+			// 				format(relation, correctRelation)
+			// 			)} ${chalk.blue(format(object, correctObject))}`
+			// 		)
 
-					correctFacts += Number(correctSubject && correctRelation && correctObject)
+			correctFacts += 0 //Number(correctSubject && correctRelation && correctObject)
 
-					if (correctFacts) {
-						omitted.push(k)
-						break
-					}
-				}
+			// 		if (correctFacts) {
+			// 			omitted.push(k)
+			// 			break
+			// 		}
+			// 	}
 
-				totalRecall += correctFacts
-			}
+			totalRecall += correctFacts
+			// }
 
 			totalAttempts += attempts.length
 		}