ai-utils/src/memory.ts

import {MemoryNode, rebuildGraph} from './helpers.ts';
import {LLMRequest, LLMMessage} from './llm.ts';
import {AiTool} from './tools.ts';
import {KDPoint, KDTree} from './kd-tree.ts';

const FACTS_HEADING = '## Facts';

const GENERIC_TEMPLATE = `# {{Title}}

## Summary

## Details

## Related`;

export type Memory = {
	name: string;
	description: string;
	content: string;
	embedding: number[];
	links: string[];
	backlinks: string[];
}

type MemoryRef = {
	name: string;
	description: string;
}

type FactBucket = {
	subject: string;
	facts: string[];
}

function dedupeFacts(facts: string[]): string[] {
	const seen = new Map<string, string>();
	for (const f of facts) {
		const clean = f.trim();
		if (clean) seen.set(clean.toLowerCase(), clean);
	}
	return [...seen.values()];
}

function cosineDistance(a: number[], b: number[]): number {
	let dot = 0, normA = 0, normB = 0;
	for (let i = 0; i < a.length; i++) {
		dot   += a[i] * b[i];
		normA += a[i] * a[i];
		normB += b[i] * b[i];
	}
	const denom = Math.sqrt(normA) * Math.sqrt(normB);
	return denom === 0 ? 1 : 1 - dot / denom;
}

function cosineSearch(query: number[], memories: Memory[], limit: number): MemoryRef[] {
	return memories
		.filter(m => m.embedding?.length)
		.map(m => ({ref: {name: m.name, description: m.description}, distance: cosineDistance(query, m.embedding)}))
		.sort((a, b) => a.distance - b.distance)
		.slice(0, limit)
		.map(s => s.ref);
}

export class MemoryCache {
	private tree!: KDTree<MemoryRef>;
	public memories: Memory[];
	public nodes: MemoryNode[] = [];

	get length() { return this.memories.length; }

	constructor(memories: Memory[]) {
		this.memories = memories;
		this.rebuild();
	}

	private buildTree(): KDTree<MemoryRef> {
		const embedded = this.memories.filter(m => m.embedding?.length);
		if (!embedded.length) return new KDTree<MemoryRef>(0);

		const dims = embedded[0].embedding.length;
		const points: KDPoint<MemoryRef>[] = embedded.map(m => ({
			vector: m.embedding,
			payload: {name: m.name, description: m.description},
		}));

		return new KDTree<MemoryRef>(dims, 'cosine', points);
	}

	search(query: number[], limit: number): MemoryRef[] {
		if (!this.tree || this.tree.dims === 0) return [];
		return this.tree.knn(query, limit).map(r => r.point.payload);
	}

	add(memory: Memory): void {
		this.memories.push(memory);
		this.rebuild();
	}

	update(memory: Memory): void {
		const idx = this.memories.findIndex(m => m.name === memory.name);
		if (idx !== -1) this.memories[idx] = memory;
		else this.memories.push(memory);
		this.rebuild();
	}

	remove(name: string): void {
		const idx = this.memories.findIndex(m => m.name === name);
		if (idx !== -1) {
			this.memories.splice(idx, 1);
			this.rebuild();
		}
	}

	rebuild(): void {
		this.nodes = rebuildGraph(this.memories);
		this.tree = this.buildTree();
	}
}

class MemoryAccessor {
	readonly list: Memory[];
	private readonly cache: MemoryCache | null;

	constructor(memories: Memory[] | MemoryCache) {
		this.cache = memories instanceof MemoryCache ? memories : null;
		this.list = this.cache ? this.cache.memories : <Memory[]>memories;
	}

	find(name: string): Memory | undefined {
		return this.list.find(m => m.name === name);
	}

	commit(): MemoryNode[] {
		if (this.cache) {
			this.cache.rebuild();
			return this.cache.nodes;
		}
		return rebuildGraph(this.list);
	}

	ghosts(): string[] {
		const nodes = this.cache ? this.cache.nodes : rebuildGraph(this.list);
		return nodes.filter(n => n.missing).map(n => n.name);
	}

	search(vector: number[], limit: number): MemoryRef[] {
		return this.cache ? this.cache.search(vector, limit) : cosineSearch(vector, this.list, limit);
	}

	forget(name: string): boolean {
		const idx = this.list.findIndex(m => m.name === name);
		if (idx === -1) return false;
		this.list.splice(idx, 1);
		this.commit();
		return true;
	}

	async backfillEmbeddings(llm: any): Promise<number> {
		const missing = this.list.filter(m => !m.embedding?.length);
		if (!missing.length) return 0;
		await Promise.all(missing.map(async node => {
			const [e] = await llm.embedding(node.content);
			if (e) node.embedding = e.embedding;
		}));
		this.commit();
		return missing.length;
	}
}

export type MemoryOptions = {
	/** Memory object */
	memory: Memory[] | MemoryCache;
	/** Inject N memories into the system prompt */
	inject?: boolean;
	/** expose recall tool to LLM */
	tool?: boolean;
	/** Update memory on compression */
	update?: boolean;
	/** Max context size of memories to inject to each call (removed immediately after use) */
	maxTokens?: number;
}

export class MemoryManager {
	private recentlyTouched = new Map<string, number>();

	private queues = new Map<string, {
		dirty: boolean,
		request: {abort?: () => void} | null,
		task: Promise<void>,
	}>();

	tools = {
		forget: (memories: Memory[] | MemoryCache): AiTool => ({
			name: 'memory_forget',
			description: 'Permanently delete a memory document and clean up all references to it',
			args: {
				name: {type: 'string', description: 'Exact memory name to forget', required: true}
			},
			fn: (args: any) => {
				const result = this.forget(args.name, memories);
				return result ? `Forgotten: ${args.name}` : `Not found: ${args.name}`;
			},
		}),

		read: (memories: Memory[] | MemoryCache): AiTool => ({
			name: 'memory_recall',
			description: 'Read the full content of a memory document',
			args: {
				name: {type: 'string', description: 'Exact memory name', required: true},
			},
			fn: (args: any) => {
				const mem = this.access(memories).find(args.name);
				if (!mem) return 'Document not found';
				this.touch(mem.name);
				return mem.content;
			},
		}),

		search: (memories: Memory[] | MemoryCache): AiTool => ({
			name: 'memory_search',
			description: 'Use embeddings to find the MOST relevant memories, even if NOT relevant',
			args: {
				query: {type: 'string', description: 'What to look for in the memories', required: true},
				limit: {type: 'number', description: 'Number of memories to return', default: 1},
			},
			fn: async ({query, limit}) => {
				const mem = await this.recollect(query, memories, limit)
				return mem.map(m => `Memory: ${m.name}
Description: ${m.description}
Links: ${[...m.links, ...m.backlinks].join(', ')}
\`\`\`
${m.content}
\`\`\``).join('\n\n');
			},
		}),
	};

	constructor(private llm: any) {}

	static normalize(m?: Memory[] | MemoryCache | MemoryOptions) {
		if (!m) return null;
		const raw = m instanceof MemoryCache || Array.isArray(m);
		return raw ? {memory: <Memory[] | MemoryCache>m, inject: true, tool: true, update: true} : {inject: true, tool: true, update: true, ...m};
	}

	private access(memories: Memory[] | MemoryCache): MemoryAccessor {
		return new MemoryAccessor(memories);
	}

	private appendFacts(node: Memory, facts: string[]): void {
		this.ensureDoc(node);
		const body = this.stripHeader(node.content);
		const bullets = facts.map(f => `- ${f}`).join('\n');
		const idx = body.indexOf(FACTS_HEADING);
		const newBody = idx === -1
			? `${body.trimEnd()}\n\n${FACTS_HEADING}\n${bullets}\n`
			: `${body.slice(0, idx + FACTS_HEADING.length)}\n${bullets}${body.slice(idx + FACTS_HEADING.length)}`;
		node.content = this.touchHeader(node, newBody);
	}

	private ensureDoc(node: Memory): void {
		if (node.content) return;
		const title = node.name.split('/').pop() ?? node.name;
		node.content = this.touchHeader(node, `# ${title}\n`);
	}

	private async factAgent(conversation: string, store: MemoryAccessor, options: LLMRequest, weekKey: string): Promise<FactBucket[]> {
		const ghosts = store.ghosts();

		const response = await this.llm.ask(conversation, {
			model: options.model,
			temperature: 0.2,
			system: `You are a fact extractor to build obsidian knowledge vaults.
Analyze this conversation and extract facts worth remembering long-term.

Rules:
- Always extract facts that the user explicitly told you to remember
- ONLY extract current facts the USER explicitly stated about themselves, their work, projects or decisions that were MADE during this conversation
- DO NOT extract greetings, pleasantries, or generic exchanges
- DO NOT extract deltas or changes in facts; ONLY the end fact
- DO NOT extract anything the AI/assistant itself said
- If nothing worth remembering was said, return an empty buckets array

When extracting facts, you MUST also decide the exact destination path:
- Reuse node names (including ghost) as much as possible IF the facts belongs there
- All information primarily about the user should go under "People/User"
- When required, create a new path following collection/subject format (e.g., People/Sarah, Projects/Oxide) — you are not limited to any fixed list of collections, use whatever fits
- For journal entries, use "Journal"

Available nodes:
- Journal
${this.listNodes(store.list).filter(n => !n.name.includes('Journal')).map(n => `- ${n.name}: ${n.description}`).join('\n') || 'None yet.'}
${ghosts.length ? `${ghosts.map(g => `- ${g}: (Ghost)`).join('\n')}` : ''}`,
			schema: {
				buckets: {type: 'array', description: 'Groups of facts to remember, each assigned to a different node. Return an empty array if there is nothing worth storing in an obsidian vault', items: {
					type: 'object', items: {
							subject: {type: 'string', description: 'Exact existing node name OR new path (e.g. "People/Sarah", "Projects/Oxide"), or "Journal"', required: true},
							facts: {
								type: 'array',
								description: 'Facts to store at this destination',
								items: {type: 'string', description: 'A single fact'},
							},
						},
					},
				},
			},
		});

		const buckets = new Map<string, string[]>();
		for(const bucket of response.buckets ?? []) {
			const subject = bucket.subject.trim().toLowerCase() === 'journal'
				? `Journal/${weekKey}` : bucket.subject.trim();
			const facts = buckets.get(subject) ?? [];
			facts.push(...dedupeFacts(bucket.facts));
			buckets.set(subject, facts);
		}

		return buckets.entries().toArray().map(([subject, facts]) => ({subject, facts}));
	}

	private getWeekMonday(date: Date = new Date()): string {
		const d = new Date(Date.UTC(date.getFullYear(), date.getMonth(), date.getDate()));
		const day = d.getUTCDay();
		const diff = day === 0 ? -6 : 1 - day;
		d.setUTCDate(d.getUTCDate() + diff);
		return d.toISOString().slice(0, 10);
	}

	private listNodes(memories: Memory[]): MemoryRef[] {
		return memories.map(m => ({name: m.name, description: m.description}));
	}

	private reconcile(node: Memory, memories: Memory[] | MemoryCache, options: LLMRequest): Promise<void> {
		const key = node.name;
		const existing = this.queues.get(key);
		if (existing) {
			existing.dirty = true;
			existing.request?.abort?.();
			return existing.task;
		}

		const entry = {dirty: false, request: null, task: Promise.resolve()};
		this.queues.set(key, entry);
		const store = this.access(memories);
		entry.task = (async () => {
			do {
				entry.dirty = false;
				await this.docAgent(node, store.list, options, entry);
			} while (entry.dirty);
		})().finally(() => {
			this.queues.delete(key);
			store.commit();
		});
		return entry.task;
	}

	private async docAgent(node: Memory, memories: Memory[], options: LLMRequest, entry: {request: {abort?: () => void} | null}): Promise<void> {
		const currentBody = this.stripHeader(node.content);
		let update;
		try {
			for (let i = 0; i < 2 && !update?.content; i++) {
				const request = this.llm.ask(currentBody, {
					model: options.model,
					temperature: 0.3,
					schema: {
						description: {type: 'string', description: 'One-line description of what this document covers, no formatting or emojis', required: true},
						content: {type: 'string', description: 'Rewritten document body in markdown, without the frontmatter block', required: true},
					},
					system: `You are a knowledge base editor maintaining one document in an Obsidian-style vault.

If the document has a "${FACTS_HEADING}" section, integrate every bullet under it into the appropriate part of the document, then remove the "${FACTS_HEADING}" section entirely. If there is no such section, just tidy the document per the rules below.

Structure: follow this generic shape loosely, adapting section names/order to what the content actually needs (e.g. journal-style docs may want a timeline instead of "Details"):
\`\`\`markdown
${GENERIC_TEMPLATE}
\`\`\`

Formatting rules:
- Use Obsidian-style markdown: # headings, **bold** for emphasis, bullet & numbered lists for grouped 1D data, tables for 2D data
- Link related concepts with [[WikiLink]] notation using full paths like [[People/Sarah]] or [[Projects/Website]]
- Create links for specific entities (person, place, project, program) and abstract concepts, but skip generics (car, red, dog)
- Keep the document concise, factual, and human-readable
- Resolve contradictions: newer facts always win — delete the outdated statement entirely, never keep both
- Do not add frontmatter blocks, filler, preamble, or AI commentary

Other nodes in the vault (link to these instead of duplicating their content):
${this.listNodes(memories).filter(n => n.name !== node.name).map(n => n.name).join(', ') || 'none'}

Current document:
\`\`\`markdown
${currentBody}
\`\`\``,
				});
				entry.request = request;
				update = await request;
			}
		} catch (err: any) {
			if (err?.name === 'AbortError') return;
			throw err;
		} finally {
			entry.request = null;
		}

		if (!update?.content) return;
		node.description = node.name !== 'People/User' ? update.description : 'All information about the current user';
		node.content = this.touchHeader(node, update.content);
		const [e] = await this.llm.embedding(node.content);
		if (e) node.embedding = e.embedding;
	}

	private parseFrontmatter(content: string): {fm: Map<string, string>, body: string} {
		const match = content.match(/^---\n([\s\S]*?)\n---\n?([\s\S]*)$/);
		if (!match) return {fm: new Map(), body: content};
		const fm = new Map<string, string>();
		for (const line of match[1].split('\n')) {
			const i = line.indexOf(':');
			if (i === -1) continue;
			fm.set(line.slice(0, i).trim(), line.slice(i + 1).trim());
		}
		return {fm, body: match[2]};
	}

	private stripHeader(content: string): string {
		return content.replace(/^---[\s\S]*?\n---\n?/, '').trimStart();
	}

	private touchHeader(node: Memory, body: string): string {
		const {fm} = this.parseFrontmatter(node.content);
		fm.set('name', node.name);
		fm.set('description', node.description || '');
		fm.set('modified', new Date().toISOString());
		return this.writeFrontmatter(fm, body);
	}

	private writeFrontmatter(fm: Map<string, string>, body: string): string {
		const lines = [...fm.entries()].map(([k, v]) => `${k}: ${v}`);
		return `---\n${lines.join('\n')}\n---\n\n${body.trimStart()}`;
	}

	decay() {
		for (const [name, ttl] of this.recentlyTouched) {
			if (ttl <= 1) this.recentlyTouched.delete(name);
			else this.recentlyTouched.set(name, ttl - 1);
		}
	}

	touch(name: string, ttl = 2) {
		this.recentlyTouched.set(name, ttl);
	}

	forget(name: string, memories: Memory[] | MemoryCache): boolean {
		return this.access(memories).forget(name);
	}

	async recollect(query: string, memories: Memory[] | MemoryCache, limit = 5, graphDepth = 1): Promise<Memory[]> {
		const store = this.access(memories);
		if (!store.list.length) return [];

		await store.backfillEmbeddings(this.llm);

		const [e] = await this.llm.embedding(query);
		if (!e) return [];

		const vectorResults = store.search(e.embedding, limit);
		const found = new Set<string>(vectorResults.map(r => r.name));

		if (graphDepth > 0) {
			let frontier = [...found];
			for (let depth = 0; depth < graphDepth && frontier.length; depth++) {
				const next: string[] = [];
				for (const name of frontier) {
					const node = store.find(name);
					if (!node) continue;
					for (const link of node.links) {
						if (!found.has(link) && store.find(link)) {
							found.add(link);
							next.push(link);
						}
					}
				}
				frontier = next;
			}
		}

		const vectorOrder = vectorResults.map(r => r.name);
		const graphExpansions = [...found].filter(n => !vectorOrder.includes(n));
		return [...vectorOrder, ...graphExpansions].map(n => store.find(n)!).filter(Boolean);
	}

	async memorize(history: LLMMessage[], memories: Memory[] | MemoryCache, options: LLMRequest): Promise<Memory[]> {
		const conversation = history
			.filter(h => h.role === 'user' || h.role === 'assistant')
			.map(h => `[${h.role}]: ${h.content}`).join('\n\n').trim();
		if (!conversation) return [];

		const uid = `${Date.now()}_${Math.random().toString(36).slice(2)}`;
		const pending = {role: 'tool', name: 'memory_process', id: uid, content: conversation} as unknown as LLMMessage;
		history.push(pending);

		const store = this.access(memories);
		const buckets = await this.factAgent(conversation, store, options, this.getWeekMonday());
		const touched: Memory[] = [];

		for (const {subject, facts} of buckets) {
			let node = store.find(subject);
			if (!node) {
				node = {name: subject, description: '', content: '', embedding: [], links: [], backlinks: []};
				store.list.push(node);
			}
			this.appendFacts(node, facts);
			const [e] = await this.llm.embedding(node.content);
			if (e) node.embedding = e.embedding;
			this.touch(node.name);
			touched.push(node);
		}

		if (touched.length) {
			store.commit();
			(pending as any).content = `Saved to ${touched.map(n => `[[${n.name}]]`).join(', ')}`;
			await Promise.all(touched.map(node => this.reconcile(node, memories, options).catch(() => {})));
		} else {
			(pending as any).content = 'Nothing worth remembering.';
		}

		(touched as any).uid = uid;
		return touched;
	}

	async reconcileVault(memories: Memory[] | MemoryCache, options: LLMRequest, scope: 'touched' | 'all' = 'touched'): Promise<void> {
		const store = this.access(memories);
		const targets = scope === 'all' ? store.list : store.list.filter(m => m.content.includes(FACTS_HEADING));
		await Promise.all(targets.map(node => this.reconcile(node, memories, options)));
		store.commit();
	}
}