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';
import {escapeRegex} from '@ztimson/utils';

const MERGE_THRESHOLD = 0.88;
const PENDING_HEADING = '## Pending';
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[];
}

type FactAgentResult = {
	buckets: FactBucket[];
	journal: 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 function stripHeader(content: string): string {
	return content.replace(/^---[\s\S]*?\n---\n?/, '').trimStart();
}

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 existing = this.memories.find(m => m.name === memory.name);
		if (existing) Object.assign(existing, 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.description}\n\n${stripHeader(node.content)}`.trim());
			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 stage(node: Memory, block: string): void {
		this.ensureDoc(node);
		const body = stripHeader(node.content);
		const idx = body.indexOf(PENDING_HEADING);
		const newBody = idx === -1
			? `${body.trimEnd()}\n\n${PENDING_HEADING}\n${block}\n`
			: `${body.slice(0, idx + PENDING_HEADING.length)}\n${block}${body.slice(idx + PENDING_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 sanitizeDescription(text: string): string {
		return (text ?? '').replace(/\s+/g, ' ').trim().slice(0, 240);
	}

	private relink(memories: Memory[], from: string, to: string): void {
		const pattern = new RegExp(`\\[\\[${escapeRegex(from)}\\]\\]`, 'g');
		for (const m of memories) if (pattern.test(m.content)) m.content = m.content.replace(pattern, `[[${to}]]`);
	}

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

		const response = await this.llm.ask(conversation, {
			model: options.model,
			temperature: 0.2,
			system: `You are a fact extractor for Obsidian-style knowledge vaults. Analyze the conversation and produce:

1. Journal recap (single paragraph)
   - "Captains Log" style record keeping
   - What was discussed/worked on, decisions, user's events/state/mood, general context
   - Leave empty only for trivial/empty exchanges/small talk

2. Fact buckets
   - ONLY facts the USER explicitly stated about themselves, their work, projects, or decisions made during this conversation
   - NEVER extract greetings, pleasantries, or anything the assistant itself said
   - Extract the final/end state, not deltas

Path assignment rules:
   - Reuse existing node names whenever possible
   - Documents should be grouped and named by the root subject
     - Person → People/Name
     - Project → Projects/Name
     - Concept → Concepts/Name
   - A bug report, its investigation, should be nested and attached to the same root subject node
   - Tickets/one-off tasks → file under the project/name/component they belong to
   - Only create a new top-level node when the fact belongs to a genuinely new subject (person/project/concept)\`

Available nodes:
${this.listNodes(store.list).map(n => `- ${n.name}: ${n.description}`).join('\n') || 'None yet.'}
${ghosts.length ? `${ghosts.map(g => `- ${g}: (Ghost)`).join('\n')}` : ''}`,
			schema: {
				journal: {type: 'string', description: 'Short day-to-day recap; empty if nothing happened.', required: false},
				buckets: {type: 'array', description: 'Groups of facts to remember; empty array if nothing worth storing.', items: {
						type: 'object', items: {
							subject: {type: 'string', description: 'Exact node name or new path (e.g. "People/Sarah", "Projects/Oxide")', required: true},
							facts: {type: 'array', description: 'Facts to store here', items: {type: 'string'}},
						},
					},
				},
			},
		});

		const buckets = new Map<string, string[]>();
		for (const bucket of response.buckets ?? []) {
			const subject = bucket.subject.trim();
			const facts = buckets.get(subject) ?? [];
			facts.push(...dedupeFacts(bucket.facts));
			buckets.set(subject, facts);
		}

		return {
			buckets: buckets.entries().toArray().map(([subject, facts]) => ({subject, facts})),
			journal: (response.journal ?? '').trim(),
		};
	}

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

	/** Find the nearest node above the similarity threshold and fold the smaller/less-connected one into
	 * the other. Journals are exempt — they're partitioned by date, not topic, and merging across weeks
	 * would wreck the timeline. Returns 'merged' if `node` absorbed another (caller should re-run the doc
	 * agent), 'absorbed' if `node` itself got folded away (caller should stop touching it), or null. */
	private async checkMerge(node: Memory, memories: Memory[] | MemoryCache, options: LLMRequest, threshold = MERGE_THRESHOLD): Promise<Memory | null> {
		if (!node.embedding?.length || node.name.startsWith('Journal/')) return null;
		const store = this.access(memories);

		let closest: Memory | null = null, closestDist = Infinity;
		for (const other of store.list) {
			if (other.name === node.name || other.name.startsWith('Journal/') || !other.embedding?.length) continue;
			const d = cosineDistance(node.embedding, other.embedding);
			if (d < closestDist) { closestDist = d; closest = other; }
		}
		if (!closest || closestDist > threshold) return null;

		const result = await this.mergeAgent(node, closest, options);
		const merged: Memory = {name: result.name, description: this.sanitizeDescription(result.description), content: '', embedding: [], links: [], backlinks: []};
		merged.content = this.touchHeader(merged, result.content);
		const [e] = await this.llm.embedding(`${merged.description}\n\n${result.content}`.trim());
		if (e) merged.embedding = e.embedding;

		this.relink(store.list, node.name, merged.name);
		this.relink(store.list, closest.name, merged.name);

		this.queues.get(closest.name)?.request?.abort?.();
		this.queues.delete(closest.name);

		store.forget(node.name);
		store.forget(closest.name);
		store.list.push(merged);
		store.commit();

		return merged;
	}

	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 () => {
			let current = node;
			do {
				entry.dirty = false;
				await this.docAgent(current, store.list, options, entry);
				const merged = await this.checkMerge(current, memories, options);
				if (merged) { current = merged; entry.dirty = true; }
			} 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> {
		if(!memories.includes(node)) return;
		const currentBody = 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 factual sentence describing the document\'s ENTIRE SUBJECT MATTER — for use as a search/merge fingerprint', 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 Obsidian-style document.

If it has a "## Pending" section, fold all new material into the appropriate part, resolve overlap, then remove the section entirely. If no section, just tidy per the rules below.

Use this loose structure, adapting headings to what the content needs:

# Title
## Summary
## Details
## Related

Rules:
- Contradictions: newer facts always win — delete outdated statements entirely
- Journals (Journal/...): keep entries as a chronological timeline; clean up grammar within entries but never delete history
- Use Obsidian markdown: # headings, **bold**, bullet/numbered lists, tables for 2D data
- Link specific entities and concepts with [[WikiLink]] (e.g., [[Projects/KiwixServer]]); skip generics
- Keep concise, factual, human-readable
- NO frontmatter, filler, preamble, or AI commentary

Available nodes to link to (don't duplicate 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' ? this.sanitizeDescription(update.description) : 'All information about the current user';
		node.content = this.touchHeader(node, update.content);
		const [e] = await this.llm.embedding(`${node.description}\n\n${update.content}`.trim());
		if (e) node.embedding = e.embedding;
	}

	private async mergeAgent(a: Memory, b: Memory, options: LLMRequest): Promise<{name: string, description: string, content: string}> {
		return this.llm.ask('', {
			model: options.model,
			temperature: 0.3,
			schema: {
				name: {type: 'string', description: 'New path for the merged doc, collection/subject format (e.g. Projects/Oxide) — only reuse an old title if it\'s genuinely the best fit', required: true},
				description: {type: 'string', description: 'One factual sentence describing the merged document\'s subject matter', required: true},
				content: {type: 'string', description: 'Fully reconciled body in markdown, without frontmatter', required: true},
			},
			system: `You are a knowledge base editor merging two overlapping Obsidian documents into one. Newer facts win on contradiction.

Structure loosely:

# Title
## Summary
## Details
## Related

Combine both documents, resolve duplication and contradictions.

Document A ("${a.name}"):
\`\`\`markdown
${stripHeader(a.content)}
\`\`\`

Document B ("${b.name}"):
\`\`\`markdown
${stripHeader(b.content)}
\`\`\``,
		});
	}

	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;
			const key = line.slice(0, i).trim();
			const raw = line.slice(i + 1).trim();
			let value = raw;
			try { value = JSON.parse(raw); } catch { /* legacy unquoted value, keep raw */ }
			fm.set(key, value);
		}
		return {fm, body: match[2]};
	}

	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}: ${JSON.stringify(String(v).replace(/\s+/g, ' ').trim())}`);
		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, journal} = await this.factAgent(conversation, store, options);
		const touched: Memory[] = [];

		if (journal) {
			const journalName = `Journal/${this.getWeekMonday()}`;
			let jnode = store.find(journalName);
			if (!jnode) {
				jnode = {name: journalName, description: '', content: '', embedding: [], links: [], backlinks: []};
				store.list.push(jnode);
			}
			this.stage(jnode, `### ${new Date().toISOString().slice(0, 10)}\n${journal}`);
			touched.push(jnode);
		}

		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.stage(node, facts.map(f => `- ${f}`).join('\n'));
			touched.push(node);
		}

		for (const node of touched) {
			const [e] = await this.llm.embedding(`${node.description}\n\n${stripHeader(node.content)}`.trim());
			if (e) node.embedding = e.embedding;
			this.touch(node.name);
		}

		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 reconcileAll(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(PENDING_HEADING));
		await Promise.all(targets.map(node => this.reconcile(node, memories, options)));
		store.commit();
	}
}