@prometheus-ai/memory 0.5.0

package/src/core/migrations/e6-triplestore-split.ts

@@ -0,0 +1,211 @@
+import type { Database } from "bun:sqlite";
+import { existsSync, writeFileSync } from "node:fs";
+import { closeQuietly, type DatabasePath, openDatabase } from "../../db";
+
+export const ANNOTATION_KINDS = ["mentions", "fact", "occurred_on", "has_source"] as const;
+export type AnnotationKind = (typeof ANNOTATION_KINDS)[number];
+
+export interface MigrationOptions {
+	readonly dbPath: DatabasePath;
+	readonly dryRun?: boolean;
+	readonly backup?: boolean;
+	readonly logFn?: (line: string) => void;
+}
+
+export interface PendingConnection {
+	query<T = unknown>(sql: string): { get(...params: unknown[]): T | null };
+}
+type SerializableDatabase = Database & { serialize(): Uint8Array };
+
+interface TripleCandidateRow {
+	id: number;
+	subject: string;
+	predicate: AnnotationKind;
+	object: string;
+	source: string | null;
+	confidence: number | null;
+	created_at: string | null;
+}
+
+interface Classification {
+	rows: TripleCandidateRow[];
+	total: number;
+}
+
+function placeholders(count: number): string {
+	return Array.from({ length: count }, () => "?").join(",");
+}
+
+function hasTable(db: Database, name: string): boolean {
+	return db.query("SELECT 1 FROM sqlite_master WHERE type = 'table' AND name = ?").get(name) !== null;
+}
+
+function copyDatabase(source: DatabasePath, destination: string): void {
+	let db: Database | null = null;
+	try {
+		db = openDatabase(source, { create: false, readwrite: false, pragmas: false });
+		writeFileSync(destination, (db as SerializableDatabase).serialize());
+	} finally {
+		closeQuietly(db);
+	}
+}
+
+function initAnnotations(db: Database): void {
+	db.run(`
+		CREATE TABLE IF NOT EXISTS annotations (
+			id INTEGER PRIMARY KEY AUTOINCREMENT,
+			memory_id TEXT NOT NULL,
+			kind TEXT NOT NULL,
+			value TEXT NOT NULL,
+			source TEXT,
+			confidence REAL DEFAULT 1.0,
+			created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
+		)
+	`);
+	db.run("CREATE INDEX IF NOT EXISTS idx_annot_memory_kind ON annotations(memory_id, kind)");
+	db.run("CREATE INDEX IF NOT EXISTS idx_annot_kind_value ON annotations(kind, value)");
+	db.run("CREATE UNIQUE INDEX IF NOT EXISTS idx_annot_unique ON annotations(memory_id, kind, value)");
+}
+
+export function hasPendingMigration(db: Database): boolean {
+	if (!hasTable(db, "triples")) return false;
+	const marks = placeholders(ANNOTATION_KINDS.length);
+	if (!hasTable(db, "annotations")) {
+		return db.query(`SELECT 1 FROM triples WHERE predicate IN (${marks}) LIMIT 1`).get(...ANNOTATION_KINDS) !== null;
+	}
+	return (
+		db
+			.query(`
+				SELECT 1
+				FROM triples t
+				WHERE t.predicate IN (${marks})
+				  AND NOT EXISTS (
+					  SELECT 1 FROM annotations a
+					  WHERE a.memory_id = t.subject
+						AND a.kind = t.predicate
+						AND a.value = t.object
+				  )
+				LIMIT 1
+			`)
+			.get(...ANNOTATION_KINDS) !== null
+	);
+}
+function classifyRows(db: Database): Classification {
+	if (!hasTable(db, "triples")) return { rows: [], total: 0 };
+	const totalRow = db.query("SELECT COUNT(*) AS count FROM triples").get() as { count: number };
+	const marks = placeholders(ANNOTATION_KINDS.length);
+	const candidates = db
+		.query(`
+			SELECT id, subject, predicate, object, source, confidence, created_at
+			FROM triples
+			WHERE predicate IN (${marks})
+			ORDER BY id ASC
+		`)
+		.all(...ANNOTATION_KINDS) as TripleCandidateRow[];
+	if (!hasTable(db, "annotations")) return { rows: candidates, total: totalRow.count };
+	const rows = candidates.filter(row => {
+		return (
+			db
+				.query("SELECT 1 FROM annotations WHERE memory_id = ? AND kind = ? AND value = ? LIMIT 1")
+				.get(row.subject, row.predicate, row.object) === null
+		);
+	});
+	return { rows, total: totalRow.count };
+}
+
+function kindCounts(rows: readonly TripleCandidateRow[]): Record<string, number> {
+	const counts: Record<string, number> = {};
+	for (const row of rows) counts[row.predicate] = (counts[row.predicate] ?? 0) + 1;
+	return counts;
+}
+
+function migrateRows(db: Database, rows: readonly TripleCandidateRow[]): number {
+	if (rows.length === 0) return 0;
+	const insert = db.prepare(`
+		INSERT OR IGNORE INTO annotations (memory_id, kind, value, source, confidence, created_at)
+		VALUES (?, ?, ?, ?, ?, ?)
+	`);
+	let written = 0;
+	for (const row of rows) {
+		const result = insert.run(
+			row.subject,
+			row.predicate,
+			row.object,
+			row.source,
+			row.confidence ?? 1.0,
+			row.created_at,
+		) as {
+			readonly changes: number;
+		};
+		written += result.changes;
+	}
+	return written;
+}
+
+export function migrate(
+	dbPathOrOptions: DatabasePath | MigrationOptions,
+	dryRun = false,
+	backup = true,
+	logFn: (line: string) => void = console.log,
+): number {
+	const options =
+		typeof dbPathOrOptions === "string" ? { dbPath: dbPathOrOptions, dryRun, backup, logFn } : dbPathOrOptions;
+	const dbPath = options.dbPath;
+	const effectiveDryRun = options.dryRun ?? false;
+	const effectiveBackup = options.backup ?? true;
+	const effectiveLog = options.logFn ?? console.log;
+	if (dbPath === ":memory:" || !existsSync(dbPath)) {
+		effectiveLog(`ERROR: database not found: ${dbPath}`);
+		throw new Error(`database not found: ${dbPath}`);
+	}
+
+	let db = openDatabase(dbPath);
+	let classified: Classification;
+	try {
+		classified = classifyRows(db);
+	} finally {
+		closeQuietly(db);
+	}
+
+	effectiveLog(`Database: ${dbPath}`);
+	effectiveLog(`  triples rows (total):        ${classified.total}`);
+	effectiveLog(`  rows-to-migrate (this run):  ${classified.rows.length}`);
+	if (classified.rows.length > 0) {
+		const counts = kindCounts(classified.rows);
+		for (const kind of Object.keys(counts).sort()) effectiveLog(`    ${kind.padEnd(14, " ")} ${counts[kind]}`);
+	}
+	if (classified.rows.length === 0) {
+		effectiveLog("Nothing to migrate. Schema is already split or no annotation rows exist.");
+		return 0;
+	}
+	if (effectiveDryRun) {
+		effectiveLog("Dry run: no changes written.");
+		return classified.rows.length;
+	}
+	if (effectiveBackup) {
+		const backupPath = `${dbPath}.pre_e6_backup`;
+		if (existsSync(backupPath)) effectiveLog(`Backup already exists at ${backupPath}; leaving as-is.`);
+		else {
+			copyDatabase(dbPath, backupPath);
+			effectiveLog(`Backup written to ${backupPath}`);
+		}
+	}
+
+	db = openDatabase(dbPath);
+	try {
+		db.run("BEGIN IMMEDIATE");
+		try {
+			initAnnotations(db);
+			const lockedClassification = classifyRows(db);
+			const written = migrateRows(db, lockedClassification.rows);
+			db.run("COMMIT");
+			effectiveLog(`Migration complete: ${written} rows moved to annotations table.`);
+			return written;
+		} catch (error) {
+			db.run("ROLLBACK");
+			throw error;
+		}
+	} finally {
+		closeQuietly(db);
+	}
+}

package/src/core/migrations/index.ts

@@ -0,0 +1 @@
+export * from "./e6-triplestore-split";

package/src/core/mmr.ts

@@ -0,0 +1,71 @@
+export interface MmrResult {
+	readonly content?: string;
+	readonly score?: number;
+	readonly [key: string]: unknown;
+}
+
+export type SimilarityFn = (textA: string, textB: string) => number;
+
+export function jaccardSimilarity(textA: string, textB: string): number {
+	const wordsA = new Set(textA.toLowerCase().split(/\s+/).filter(Boolean));
+	const wordsB = new Set(textB.toLowerCase().split(/\s+/).filter(Boolean));
+
+	if (wordsA.size === 0 || wordsB.size === 0) return 0.0;
+
+	let intersection = 0;
+	for (const word of wordsA) {
+		if (wordsB.has(word)) intersection += 1;
+	}
+
+	return intersection / (wordsA.size + wordsB.size - intersection);
+}
+export function mmrRerank<T extends MmrResult>(
+	results: readonly T[],
+	lambdaParam = 0.7,
+	topK = 10,
+	similarityFn: SimilarityFn = jaccardSimilarity,
+): T[] {
+	const limit = Math.max(0, Math.trunc(topK));
+	if (limit <= 0) return [];
+	if (results.length <= 1) return results.slice(0, limit);
+
+	const sortedResults = results.slice().sort((left, right) => (right.score ?? 0) - (left.score ?? 0));
+	const first = sortedResults[0];
+	if (first === undefined) return [];
+
+	const selected: T[] = [first];
+	const remaining = sortedResults.slice(1);
+
+	while (remaining.length > 0 && selected.length < limit) {
+		let bestIdx = 0;
+		let bestScore = Number.NEGATIVE_INFINITY;
+
+		for (let idx = 0; idx < remaining.length; idx += 1) {
+			const candidate = remaining[idx];
+			if (candidate === undefined) continue;
+
+			let maxSimilarity = 0.0;
+			const candidateContent = candidate.content ?? "";
+			for (const selectedResult of selected) {
+				const similarity = similarityFn(candidateContent, selectedResult.content ?? "");
+				if (similarity > maxSimilarity) maxSimilarity = similarity;
+			}
+
+			const relevance = candidate.score ?? 0;
+			const mmrScore = lambdaParam * relevance - (1.0 - lambdaParam) * maxSimilarity;
+			if (mmrScore > bestScore) {
+				bestScore = mmrScore;
+				bestIdx = idx;
+			}
+		}
+
+		const chosen = remaining.splice(bestIdx, 1)[0];
+		if (chosen !== undefined) selected.push(chosen);
+	}
+
+	if (selected.length < limit) {
+		selected.push(...remaining.slice(0, limit - selected.length));
+	}
+
+	return selected;
+}

package/src/core/orchestrator.ts

@@ -0,0 +1,62 @@
+import type { BeamMemoryState, RecallOptions, RecallResult } from "./beam/types";
+import { embedQuery } from "./embeddings";
+import {
+	type PolyphonicMemoryResult,
+	type PolyphonicRecallOptions,
+	polyphonicRecall,
+	polyphonicRecallIsEnabled,
+} from "./polyphonic-recall";
+
+export interface OrchestratorBeam extends BeamMemoryState {
+	recall?: (query: string, topK?: number, options?: RecallOptions) => Promise<RecallResult[]>;
+	recallEnhanced?: (query: string, topK?: number, options?: RecallOptions) => Promise<RecallResult[]>;
+}
+
+export interface OrchestrateRecallOptions
+	extends Omit<RecallOptions, "queryEmbedding">,
+		Omit<PolyphonicRecallOptions, "queryEmbedding"> {
+	readonly queryEmbedding?: readonly number[] | Float32Array | null;
+	readonly enhanced?: boolean;
+	readonly forcePolyphonic?: boolean;
+	readonly forceLinear?: boolean;
+}
+
+export interface OrchestratedRecallResult extends Omit<RecallResult, "metadata" | "score" | "tier"> {
+	score?: number;
+	metadata?: RecallResult["metadata"];
+	tier?: RecallResult["tier"] | PolyphonicMemoryResult["tier"];
+	combined_score?: PolyphonicMemoryResult["combined_score"];
+	voice_scores?: PolyphonicMemoryResult["voice_scores"];
+}
+
+function toLinearRecallOptions(options: OrchestrateRecallOptions): RecallOptions {
+	if (options.queryEmbedding instanceof Float32Array) {
+		return { ...options, queryEmbedding: Array.from(options.queryEmbedding) };
+	}
+	return options as RecallOptions;
+}
+
+export async function orchestrateRecall(
+	beam: OrchestratorBeam,
+	query: string,
+	topK = 20,
+	options: OrchestrateRecallOptions = {},
+): Promise<OrchestratedRecallResult[]> {
+	const polyphonic = !options.forceLinear && (options.forcePolyphonic === true || polyphonicRecallIsEnabled());
+	let queryEmbedding: readonly number[] | Float32Array | null | undefined = options.queryEmbedding;
+	if (queryEmbedding === undefined && query.length > 0) {
+		// Auto-derive when the caller did not pass one. `embedQuery()` returns null when
+		// embeddings are disabled or no provider is configured, so this is a no-op for
+		// FTS-only deployments. `null` (explicit "no embedding") is preserved untouched.
+		queryEmbedding = await embedQuery(query);
+	}
+	if (polyphonic) {
+		return polyphonicRecall(beam, query, topK, { ...options, queryEmbedding });
+	}
+	const linearOptions = toLinearRecallOptions({ ...options, queryEmbedding });
+	if (options.enhanced === true && typeof beam.recallEnhanced === "function") {
+		return beam.recallEnhanced(query, topK, linearOptions);
+	}
+	if (typeof beam.recall === "function") return beam.recall(query, topK, linearOptions);
+	return [];
+}