ai-utils/src/memory.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 class MemoryCache {
	private tree: KDTree<MemoryRef>;
	public memories: Memory[];

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

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

	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[] {
		const results = this.tree.knn(query, limit);
		return results.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.tree = this.buildTree();
	}
}

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 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 extractLinks(content: string): string[] {
	if (!content) return [];
	const matches = content.matchAll(/\[\[([^\]]+)\]\]/g);
	return [...new Set([...matches].map(m => m[1].trim()))];
}

export function rebuildGraph(memories: Memory[]): void {
	for (const m of memories) m.links = extractLinks(m.content).filter(l => l !== m.name);
	for (const m of memories) m.backlinks = [];
	for (const m of memories) {
		for (const link of m.links) {
			const target = memories.find(t => t.name === link);
			if (target) target.backlinks.push(m.name);
		}
	}
}

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 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);
}

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

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

	tools = {
		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 mems = this.unwrap(memories);
				const mem = mems.find(m => m.name === args.name);
				if (!mem) return 'Document not found';
				this.touch(mem.name);
				return mem.content;
			},
		}),

		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}`;
			},
		}),
	};

	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 unwrap(memories: Memory[] | MemoryCache): Memory[] {
		return memories instanceof MemoryCache ? memories.memories : memories;
	}

	private sync(memories: Memory[] | MemoryCache): void {
		if (memories instanceof MemoryCache) memories.rebuild();
	}

	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 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()}`;
	}

	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 ensureDoc(node: Memory): void {
		if (node.content) return;
		const title = node.name.split('/').pop() ?? node.name;
		node.content = this.touchHeader(node, `# ${title}\n`);
	}

	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);
	}

	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);
	}

	getTouched(): string[] {
		return [...this.recentlyTouched.keys()];
	}

	forget(name: string, memories: Memory[] | MemoryCache): boolean {
		const mem = this.unwrap(memories);
		const idx = mem.findIndex(m => m.name === name);
		if (idx === -1) return false;

		mem.splice(idx, 1);
		rebuildGraph(mem);
		this.sync(memories);
		return true;
	}

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

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

		let vectorResults: MemoryRef[];
		if (memories instanceof MemoryCache) vectorResults = memories.search(e.embedding, limit);
		else vectorResults = this.cosineSearch(e.embedding, mem, limit);
		const found = new Set<string>(vectorResults.map(r => r.name));

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

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

	private cosineSearch(query: number[], memories: Memory[], limit: number): MemoryRef[] {
		const scored = 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);
		return scored.map(s => s.ref);
	}

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

	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)}`;
		// NOTE: adjust field names below (id/tool_call_id/name) to match your LLMMessage/tool-call schema.
		const pending = {role: 'tool', name: 'memory_process', id: uid, content: 'Processing…'} as unknown as LLMMessage;
		history.push(pending);

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

		for (const {subject, facts} of buckets) {
			let node = mem.find(m => m.name === subject);
			if (!node) {
				node = {name: subject, description: '', content: '', embedding: [], links: [], backlinks: []};
				mem.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) {
			rebuildGraph(mem);
			this.sync(memories);
			(pending as any).content = `Saved to ${touched.map(n => `[[${n.name}]]`).join(', ')}`;
			for (const node of touched) this.reconcile(node, memories, options).catch(() => {});
		} else {
			(pending as any).content = 'Nothing worth remembering.';
		}

		return touched;
	}

	/** Manual/cron entry point. scope 'touched' only reconciles docs with a pending Facts inbox. */
	async reconcileVault(memories: Memory[] | MemoryCache, options: LLMRequest, scope: 'touched' | 'all' = 'touched'): Promise<void> {
		const mem = this.unwrap(memories);
		const targets = scope === 'all' ? mem : mem.filter(m => m.content.includes(FACTS_HEADING));
		await Promise.all(targets.map(node => this.reconcile(node, memories, options)));
		this.sync(memories);
	}

	/**
	 * Coalescing queue: if a doc is already reconciling, mark it dirty and abort the in-flight
	 * request. The loop below always re-reads node.content fresh, so nothing is ever dropped.
	 */
	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 mem = this.unwrap(memories);
		entry.task = (async () => {
			do {
				entry.dirty = false;
				await this.reconcileDoc(node, mem, options, entry);
			} while (entry.dirty);
		})().finally(() => {
			this.queues.delete(key);
			rebuildGraph(mem);
			this.sync(memories);
		});
		return entry.task;
	}

	private async reconcileDoc(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 async factAgent(conversation: string, memories: Memory[], options: LLMRequest, weekKey: string): Promise<FactBucket[]> {
		const buckets = new Map<string, string[]>();
		await this.llm.ask(conversation, {
			model: options.model,
			temperature: 0.2,
			system: `You are a fact extractor. Analyze this conversation and extract facts worth remembering long-term.

Rules:
- ONLY extract current facts the USER explicitly stated about themselves, their work, or their projects
- ONLY extract decisions that were MADE during this conversation
- DO NOT extract anything the AI said, its capabilities, or meta-conversation about the AI
- DO NOT extract greetings, pleasantries, or generic exchanges
- DO NOT extract deltas or changes in facts; ONLY the end fact
- If nothing worth remembering was said, do not call any tools

When extracting facts, you MUST also decide the exact destination path:
- Use an existing node name if the facts clearly belong 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(memories).filter(n => !n.name.includes('Journal')).map(n => `- ${n.name}: ${n.description}`).join('\n') || 'None yet.'}`,
			tools: [{
				name: 'facts_extract',
				description: 'Submit facts with their destination',
				args: {
					destination: {type: 'string', description: 'Exact existing node name OR new path (e.g. "People/Sarah", "Projects/Oxide")', required: true},
					facts: {type: 'string', description: 'Comma-separated facts', required: true},
				},
				fn: (args: any) => {
					const subject = args.destination.trim().toLowerCase() === 'journal'
						? `Journal/${weekKey}` : args.destination.trim();
					const facts = buckets.get(subject) ?? [];
					facts.push(...dedupeFacts(String(args.facts).split(',')));
					buckets.set(subject, facts);
					return 'Recorded';
				},
			}],
		});
		return buckets.entries().toArray().map(([subject, facts]) => ({subject, facts}));
	}
}