Add multi-turn examples

CHANGELOG.md

@@ -1,3 +1,4 @@
+- [2024-10-09] [Add multi-turn examples](https://github.com/RubricLab/memory/commit/a0e32260e510a3ef1312030fcb5e2685ac2bb300)
 - [2024-10-09] [Scaffold multi-turn evals](https://github.com/RubricLab/memory/commit/ecb5531acef6a924b81684660150eeb71d93e704)
 - [2024-10-09] [Add TSConfig](https://github.com/RubricLab/memory/commit/ed521824cc492e46adff6d38a994e18cc08166b2)
 - [2024-10-09] [Extract memory to class](https://github.com/RubricLab/memory/commit/5e165608ffad822c5b77ee03f1dfc308dcb1787a)

package.json

@@ -1,7 +1,7 @@
 {
   "name": "@rubriclab/memory",
   "module": "src/index.ts",
-  "version": "0.0.10",
+  "version": "0.0.11",
   "private": false,
   "type": "module",
   "devDependencies": {

src/evals/multi-turn/examples.ts

@@ -18,12 +18,70 @@ export const EXAMPLES: Example[] = [
 				]
 			},
 			{
-				content: 'I am not vegan',
+				content: 'I now eat meat',
 				facts: [
 					{
 						subject: 'user',
 						relation: 'is not',
 						object: 'vegan'
+					},
+					{
+						subject: 'user',
+						relation: 'eats',
+						object: 'meat'
+					}
+				]
+			}
+		]
+	},
+	{
+		messages: [
+			{
+				content: 'I went to Balthazar with George on the 10th of March 2024',
+				facts: [
+					{
+						subject: 'user',
+						relation: 'went to',
+						object: 'Balthazar'
+					},
+					{
+						subject: 'user',
+						relation: 'went with',
+						object: 'George'
+					}
+				]
+			},
+			{
+				content: 'George liked the food',
+				facts: [
+					{
+						subject: 'George',
+						relation: 'liked',
+						object: 'Balthazar'
+					}
+				]
+			}
+		]
+	},
+	{
+		messages: [
+			{
+				content: 'I have a cousin named Suzy',
+				facts: [
+					{
+						subject: 'user',
+						relation: 'has a cousin named',
+						object: 'Suzy'
+					}
+				]
+			},
+			{
+				content: 'Suzy does not like cranberries',
+				facts: [
+					{
+						subject: 'Suzy',
+						relation: 'does not like',
+						object: 'cranberries'
 					}
 				]
 			}

src/evals/multi-turn/index.ts

@@ -15,7 +15,7 @@ await db
 	.get()
 
 export const runMultiTurnExamples = async ({ model }: { model: Parameters<typeof openai>[0] }) => {
-	const memory = new Memory({ model })
+	const memory = new Memory({ model, db })
 
 	let totalFacts = 0
 	let totalRecall = 0
@@ -40,20 +40,7 @@ export const runMultiTurnExamples = async ({ model }: { model: Parameters<typeof
 					`\n🎯 ${i + 1} of ${message.facts.length}: ${chalk.magenta(fact.subject)} ${chalk.yellow(fact.relation)} ${chalk.blue(fact.object)}`
 				)
 
-				for (const attempt of attempts) {
-					const { subject, relation, object } = attempt
-
-					db
-						.prepare(`
-							insert into facts (subject, relation, object)
-							values ($1, $2, $3)
-							on conflict (subject, object) do update set relation = $2
-						`)
-						.run(subject, relation, object)
-				}
-
 				const newFacts = db.query('select * from facts').all()
-				console.log({ newFacts })
 
 				for (const [k, newFact] of newFacts.entries()) {
 					const { subject, relation, object } = newFact as Fact
@@ -77,6 +64,7 @@ export const runMultiTurnExamples = async ({ model }: { model: Parameters<typeof
 						break
 					}
 				}
+
 				totalRecall += correctFacts
 			}
 

src/evals/one-shot/index.ts

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

src/index.ts

@@ -1,19 +1,61 @@
+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'
 
 export class Memory {
 	model: Parameters<typeof openai>[0]
+	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()
+	}
+
 	constructor({
-		model
+		model,
+		db
 	}: {
 		model: Parameters<typeof openai>[0]
+		db: Database
 	}) {
 		this.model = model
+		this.db = db
+
+		this.createTable()
 	}
 
 	async extract({ content }: { content: string }) {
+		const {
+			object: { entities }
+		} = await generateObject({
+			model: openai(this.model),
+			schema: z.object({
+				entities: z.array(
+					z.object({
+						subject: z.string()
+					})
+				)
+			}),
+			prompt: clean`Please extract all entities, subject or objects, from the following passage.
+            Portray the first-person as "user".
+            Passage:
+            "${content}"`
+		})
+
+		const tags = entities?.map(({ subject }) => `"${subject}"`).join(', ') || ''
+
+		const query = this.db.query(
+			`select * from facts where subject in (${tags}) or object in (${tags})`
+		)
+
+		const relevantFacts = query.all() as Fact[]
+
 		const {
 			object: { facts }
 		} = await generateObject({
@@ -32,23 +74,29 @@ export class Memory {
             Portray the first-person as "user".
             Capture new relationships.
             Try to capture the most up-to-date state of affairs in present tense.
+						Existing facts to consider and optionally update. A (subject,object) pair is unique, so use relation for negation:
+						${relevantFacts
+							?.map(({ subject, relation, object }) => `${subject} ${relation} ${object}`)
+							.join('\n ')}
             Passage:
             "${content}"`
-			// messages: [
-			// 	{
-			// 		role: 'system',
-			// 		content: clean`Please extract all probable and implicit facts from the following passage.
-			//       Portray the first-person as "user".
-			//       Capture new relationships.
-			//       Try to capture the most up-to-date state of affairs in present tense.`
-			// 	},
-			// 	{
-			// 		role: 'user',
-			// 		content: eg.content
-			// 	}
-			// ]
 		})
 
+		for await (const fact of facts) {
+			const { subject, relation, object } = fact
+
+			this.db
+				.prepare(`
+					insert into facts (subject, relation, object)
+					values ($1, $2, $3)
+					on conflict (subject, object) do update set relation = $2
+				`)
+				.run(subject, relation, object)
+		}
+
+		const priorFacts = this.db.query('select * from facts').all()
+		console.log({ priorFacts })
+
 		return { facts }
 	}
 }