From a6fb8ae828f0f7a21bd4270ee1677c303e70143b Mon Sep 17 00:00:00 2001 From: ztimson Date: Mon, 27 Jul 2026 03:59:39 -0400 Subject: [PATCH] New memory system --- package.json | 2 +- src/kd-tree.ts | 334 ++++++++++++++++++++++++++++++++ src/llm.ts | 31 +-- src/memory.ts | 515 ++++++++++++++++++++++++++++++++++++------------- 4 files changed, 730 insertions(+), 152 deletions(-) create mode 100644 src/kd-tree.ts diff --git a/package.json b/package.json index d3f49f2..598afe3 100644 --- a/package.json +++ b/package.json @@ -1,6 +1,6 @@ { "name": "@ztimson/ai-utils", - "version": "1.2.1", + "version": "1.2.2", "description": "AI Utility library", "author": "Zak Timson", "license": "MIT", diff --git a/src/kd-tree.ts b/src/kd-tree.ts new file mode 100644 index 0000000..db32ec2 --- /dev/null +++ b/src/kd-tree.ts @@ -0,0 +1,334 @@ +export type DistanceMetric = "euclidean" | "cosine"; + +export interface KDPoint { + vector: number[]; + payload: T; +} + +export interface KNNResult { + point: KDPoint; + distance: number; +} + +interface KDNode { + point: KDPoint; + axis: number; + left: KDNode | null; + right: KDNode | null; +} + +// ─── Distance helpers ───────────────────────────────────────────────────────── + +function euclidean(a: number[], b: number[]): number { + let sum = 0; + for (let i = 0; i < a.length; i++) { + const d = a[i] - b[i]; + sum += d * d; + } + return Math.sqrt(sum); +} + +function cosine(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; // distance = 1 - similarity +} + +/** + * Keeps the k closest candidates in memory, evicts the furthest when full + */ +class BoundedMaxHeap { + private heap: KNNResult[] = []; + + constructor(private readonly k: number) {} + + get size(): number { return this.heap.length; } + + get worstDistance(): number { + return this.heap.length < this.k ? Infinity : this.heap[0].distance; + } + + push(item: KNNResult): void { + if (this.heap.length < this.k) { + this.heap.push(item); + this.bubbleUp(this.heap.length - 1); + } else if (item.distance < this.heap[0].distance) { + this.heap[0] = item; + this.sinkDown(0); + } + } + + toSortedArray(): KNNResult[] { + return [...this.heap].sort((a, b) => a.distance - b.distance); + } + + private bubbleUp(i: number): void { + while (i > 0) { + const parent = (i - 1) >> 1; + if (this.heap[parent].distance >= this.heap[i].distance) break; + [this.heap[parent], this.heap[i]] = [this.heap[i], this.heap[parent]]; + i = parent; + } + } + + private sinkDown(i: number): void { + const n = this.heap.length; + while (true) { + let largest = i; + const l = 2 * i + 1, r = 2 * i + 2; + if (l < n && this.heap[l].distance > this.heap[largest].distance) largest = l; + if (r < n && this.heap[r].distance > this.heap[largest].distance) largest = r; + if (largest === i) break; + [this.heap[largest], this.heap[i]] = [this.heap[i], this.heap[largest]]; + i = largest; + } + } +} + +/** + * K-D Tree for efficient nearest-neighbor search over high-dimensional vectors / embeddings. + * + * Supports: + * - Insertion of labeled points + * - k-nearest-neighbor (KNN) search + * - Radius search (all points within a given distance) + * - Euclidean and cosine distance metrics + * - Bulk construction (balanced tree) for best query performance + */ +export class KDTree { + private root: KDNode | null = null; + private _size = 0; + private readonly dims: number; + private readonly distanceFn: (a: number[], b: number[]) => number; + + /** + * @param dims Dimensionality of all vectors (must be consistent). + * @param metric Distance metric to use. Default: "euclidean". + * @param points Optional initial set of points. Builds a balanced tree + * in O(n log² n) — prefer this over inserting one-by-one + * when you have a large corpus. + */ + constructor( + dims: number, + metric: DistanceMetric = "euclidean", + points?: KDPoint[] + ) { + this.dims = dims; + this.distanceFn = metric === "cosine" ? cosine : euclidean; + + if (points && points.length > 0) { + this.validateAll(points); + this.root = this.buildBalanced([...points], 0); + this._size = points.length; + } + } + + /** Total number of points stored in the tree. */ + get size(): number { return this._size; } + + // ── Insertion ────────────────────────────────────────────────────────────── + + /** + * Insert a single point. O(log n) average, O(n) worst case on skewed data. + * For bulk loading prefer passing points to the constructor. + */ + insert(point: KDPoint): void { + this.validate(point); + this.root = this.insertNode(this.root, point, 0); + this._size++; + } + + // ── KNN search ───────────────────────────────────────────────────────────── + + /** + * Find the k nearest neighbors to `query`. + * Returns results sorted by distance ascending. + */ + knn(query: number[], k: number): KNNResult[] { + if (k <= 0) throw new RangeError("k must be a positive integer"); + this.validateVector(query); + + const heap = new BoundedMaxHeap(k); + this.searchKNN(this.root, query, k, heap, 0); + return heap.toSortedArray(); + } + + /** + * Nearest single neighbor. Convenience wrapper around knn(query, 1). + * Returns null if the tree is empty. + */ + nearest(query: number[]): KNNResult | null { + const results = this.knn(query, 1); + return results[0] ?? null; + } + + // ── Radius search ────────────────────────────────────────────────────────── + + /** + * Return all points whose distance to `query` is ≤ `radius`, + * sorted by distance ascending. + */ + radiusSearch(query: number[], radius: number): KNNResult[] { + if (radius < 0) throw new RangeError("radius must be non-negative"); + this.validateVector(query); + + const results: KNNResult[] = []; + this.searchRadius(this.root, query, radius, results, 0); + results.sort((a, b) => a.distance - b.distance); + return results; + } + + // ── Conversion ───────────────────────────────────────────────────────────── + + /** Collect all points in the tree (order not guaranteed). */ + toArray(): KDPoint[] { + const out: KDPoint[] = []; + this.collect(this.root, out); + return out; + } + + /** + * Rebuild the tree from its current points as a balanced tree. + * Useful after many individual insertions to restore O(log n) query time. + */ + rebalance(): void { + const points = this.toArray(); + this.root = points.length ? this.buildBalanced(points, 0) : null; + } + + // ── Private: build ───────────────────────────────────────────────────────── + + private buildBalanced(points: KDPoint[], depth: number): KDNode { + const axis = depth % this.dims; + points.sort((a, b) => a.vector[axis] - b.vector[axis]); + + const mid = Math.floor(points.length / 2); + return { + point: points[mid], + axis, + left: points.slice(0, mid).length + ? this.buildBalanced(points.slice(0, mid), depth + 1) + : null, + right: points.slice(mid + 1).length + ? this.buildBalanced(points.slice(mid + 1), depth + 1) + : null, + }; + } + + // ── Private: insert ──────────────────────────────────────────────────────── + + private insertNode( + node: KDNode | null, + point: KDPoint, + depth: number + ): KDNode { + if (node === null) { + return { point, axis: depth % this.dims, left: null, right: null }; + } + const axis = depth % this.dims; + if (point.vector[axis] < node.point.vector[axis]) { + node.left = this.insertNode(node.left, point, depth + 1); + } else { + node.right = this.insertNode(node.right, point, depth + 1); + } + return node; + } + + // ── Private: KNN traversal ───────────────────────────────────────────────── + + private searchKNN( + node: KDNode | null, + query: number[], + k: number, + heap: BoundedMaxHeap, + depth: number + ): void { + if (node === null) return; + + const dist = this.distanceFn(query, node.point.vector); + heap.push({ point: node.point, distance: dist }); + + const axis = node.axis; + const diff = query[axis] - node.point.vector[axis]; + const [near, far] = diff <= 0 + ? [node.left, node.right] + : [node.right, node.left]; + + this.searchKNN(near, query, k, heap, depth + 1); + + // Only explore the far side if it could contain a closer point. + // For cosine distance we can't prune by axis gap alone, so always explore. + const shouldExplore = + this.distanceFn === cosine + ? true + : Math.abs(diff) < heap.worstDistance; + + if (shouldExplore) { + this.searchKNN(far, query, k, heap, depth + 1); + } + } + + // ── Private: radius traversal ────────────────────────────────────────────── + + private searchRadius( + node: KDNode | null, + query: number[], + radius: number, + results: KNNResult[], + depth: number + ): void { + if (node === null) return; + + const dist = this.distanceFn(query, node.point.vector); + if (dist <= radius) { + results.push({ point: node.point, distance: dist }); + } + + const axis = node.axis; + const diff = query[axis] - node.point.vector[axis]; + const [near, far] = diff <= 0 + ? [node.left, node.right] + : [node.right, node.left]; + + this.searchRadius(near, query, radius, results, depth + 1); + + const shouldExplore = + this.distanceFn === cosine ? true : Math.abs(diff) <= radius; + + if (shouldExplore) { + this.searchRadius(far, query, radius, results, depth + 1); + } + } + + // ── Private: collect ─────────────────────────────────────────────────────── + + private collect(node: KDNode | null, out: KDPoint[]): void { + if (node === null) return; + out.push(node.point); + this.collect(node.left, out); + this.collect(node.right, out); + } + + // ── Private: validation ──────────────────────────────────────────────────── + + private validateVector(v: number[]): void { + if (v.length !== this.dims) { + throw new TypeError( + `Vector length ${v.length} does not match tree dimensionality ${this.dims}` + ); + } + } + + private validate(point: KDPoint): void { + this.validateVector(point.vector); + } + + private validateAll(points: KDPoint[]): void { + for (const p of points) this.validate(p); + } +} diff --git a/src/llm.ts b/src/llm.ts index e3f1b66..2d10168 100644 --- a/src/llm.ts +++ b/src/llm.ts @@ -6,7 +6,7 @@ import {AiTool, AiToolArg} from './tools.ts'; import {fileURLToPath} from 'url'; import {dirname, join} from 'path'; import {spawn} from 'node:child_process'; -import {Memory, MemoryManager} from './memory.ts'; +import {Memory, MemoryCache, MemoryManager} from './memory.ts'; export type AnthropicConfig = {proto: 'anthropic', token: string}; export type OpenAiConfig = {proto: 'openai', host?: string, token: string}; @@ -55,7 +55,7 @@ export type LLMRequest = { /** Compress old messages in the chat to free up context */ compress?: {max: number; min: number}; /** User's memory documents - RAG injected automatically each turn */ - memory?: Memory[]; + memory?: Memory[] | MemoryCache; /** Model to use for memory operations */ memoryModel?: string; /** Skill documents the AI can browse and read on demand */ @@ -190,19 +190,20 @@ class LLM { } // Memory - if(options.memory) { - const relevant = await this.memoryManager.recollect(message, options.memory, 1); + if (options.memory) { + const mems = options.memory instanceof MemoryCache ? options.memory.memories : options.memory; + const relevant = await this.memoryManager.recollect(message, options.memory, 5); prompts.unshift(`You have access to the following memory files: -${options.memory.map(m => `- ${m.name}: ${m.description}`).join('\n')} +${mems.map(m => `- ${m.name}: ${m.description}`).join('\n')} ${relevant.length ? ` -The closest memory has been added primitively: -\`\`\` -Name: ${relevant[0].name} -Description: ${relevant[0].description} -${relevant[0].content} -\`\`\` -`: ''}`.trim()); - tools.push(this.memoryManager.tools.read(options.memory)); +Relevant memories have been preloaded: +${relevant.map(r => ` +**${r.name}** +${r.description} +${r.content} +`).join('\n---\n')} +` : ''}`.trim()); + tools.push(this.memoryManager.tools.read(options.memory)); } prompts.unshift(options.system || this.ai.options.llm?.system || ''); @@ -215,7 +216,7 @@ ${relevant[0].content} // Auto-memorize before compressing if(options.compress && this.estimateTokens(history) >= options.compress.max) { - if(options.memory) await this.memoryManager.memorize(history, options.memory, options); + if(options.memory) await this.memoryManager.memorize(history, options.memory, {model: options.memoryModel || this.defaultModel, ...options}); const compressed = await this.compressHistory(history, options.compress.max, options.compress.min, options); if(options.history) options.history.splice(0, options.history.length, ...compressed); } @@ -228,7 +229,7 @@ ${relevant[0].content} * Digest full conversation history into memory documents. * Call on session end to persist the conversation. */ - async updateMemory(history: LLMMessage[], memories: Memory[], options: LLMRequest = {}): Promise { + async updateMemory(history: LLMMessage[], memories: Memory[] | MemoryCache, options: LLMRequest = {}): Promise { await this.memoryManager.memorize(history, memories, {model: this.defaultModel, ...options}); } diff --git a/src/memory.ts b/src/memory.ts index 542fe00..a4f0950 100644 --- a/src/memory.ts +++ b/src/memory.ts @@ -1,177 +1,420 @@ -// memory.ts import {LLMRequest, LLMMessage} from './llm.ts'; import {AiTool} from './tools.ts'; +import {KDTree, KDPoint} from './kd-tree.ts'; -/** Background information the AI will be fed as a knowledge document */ export type Memory = { - /** Memory subject */ name: string; - /** Short description of what this document contains - used for RAG retrieval */ description: string; - /** Full markdown content of the document */ content: string; - /** Embedding vector of the description - used for similarity search */ embedding: number[]; + links: string[]; + backlinks: string[]; } -export type MemoryCollection = { - /** Memory subject */ +type MemoryRef = { name: string; - /** Short description - required if isNew */ - description?: string; - /** Extracted facts to merge */ + description: string; +} + +type FactBucket = { + subject: string; facts: string[]; } +// In memory.ts - replace findGhostNodes with this: + +export type MemoryNode = { + name: string; + missing: boolean; + links: string[]; + backlinks: string[]; +} + +export function buildMemoryGraph(memories: Memory[] | MemoryCache): MemoryNode[] { + const mems = memories instanceof MemoryCache ? memories.memories : memories; + const nameSet = new Set(mems.map(m => m.name)); + const ghosts = new Set(); + + // Collect all ghost references + for (const m of mems) { + for (const link of m.links) { + if (!nameSet.has(link)) ghosts.add(link); + } + } + + // Build node list: real nodes + ghost nodes + return [ + ...mems.map(m => ({ + name: m.name, + missing: false, + links: m.links, + backlinks: m.backlinks, + })), + ...[...ghosts].map(name => ({ + name, + missing: true, + links: [], + backlinks: mems + .filter(m => m.links.includes(name)) + .map(m => m.name), + })) + ]; +} + +function extractLinks(content: string): string[] { + const matches = content.matchAll(/\[\[([^\]]+)\]\]/g); + return [...new Set([...matches].map(m => m[1].trim()))]; +} + +function rebuildBacklinks(memories: Memory[]): void { + 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 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; +} + +export class MemoryCache { + private tree: KDTree; + public memories: Memory[]; + + constructor(memories: Memory[]) { + this.memories = memories; + this.tree = this.buildTree(); + } + + private buildTree(): KDTree { + const embedded = this.memories.filter(m => m.embedding?.length); + if(!embedded.length) return new KDTree(0); + + const dims = embedded[0].embedding.length; + const points: KDPoint[] = embedded.map(m => ({ + vector: m.embedding, + payload: {name: m.name, description: m.description}, + })); + + return new KDTree(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(); + } + + rebuild(): void { + this.tree = this.buildTree(); + } + + rebuildLinks(): void { + rebuildBacklinks(this.memories); + } +} export class MemoryManager { tools = { - edit: (memory: Memory): AiTool => ({ - name: 'edit_memory', - description: 'Edit a memory. Omit start/end to append. Pass start only to replace from that line on (Note line 0 = first line of content / line AFTER description). Pass start+end to replace a specific range. start=0 replaces the whole document. Returns updated document', - args: { - content: {type: 'string', description: 'New content', required: true}, - start: {type: 'number', description: 'First line to replace (0-indexed, inclusive). Omit to append.'}, - end: {type: 'number', description: 'Last line to replace (0-indexed, inclusive). Omit to replace from start to end of doc.'}, - }, - fn: (args: any) => { - const lines = memory.content ? memory.content.split('\n') : []; - const newLines = args.content.split('\n'); - if(args.start === undefined) lines.push(...newLines); - else if(args.end === undefined) lines.splice(args.start, lines.length - args.start, ...newLines); - else lines.splice(args.start, args.end - args.start + 1, ...newLines); - memory.content = lines.join('\n'); - return memory.content; - } - }), - extract: (pools: MemoryCollection[]): AiTool => ({ - name: 'extract_facts', - description: 'Extract a list of facts to group into a single memory', - args: { - name: {type: 'string', description: 'Exact name of an existing memory, or a new name if none fits ([pro]nouns only)', required: true}, - description: {type: 'string', description: 'One sentence description of the memory subject', required: true}, - facts: {type: 'string', description: 'Comma separated list of extracted facts', required: true}, - }, - fn: (args: any) => { - pools.push({ - name: args.name, - description: args.description, - facts: args.facts.split(',').map((f: string) => f.trim()).filter(Boolean), - }); - return 'Success'; - }}), - read: (memories: Memory[]): AiTool => ({ + read: (memories: Memory[] | MemoryCache): AiTool => ({ name: 'read_memory', - description: 'Read entire memory', + description: 'Read the full content of a memory document', args: { name: {type: 'string', description: 'Exact memory name', required: true}, }, - fn: (args: any) => { - const mem = memories.find(m => m.name === args.name); + fn:(args: any) => { + const mems = memories instanceof MemoryCache ? memories.memories : memories; + const mem = mems.find(m => m.name === args.name); if(!mem) return 'Document not found'; - return `Name: ${mem.name}\nDescription: ${mem.description}\n\n${mem.content}`; + return this.formatMemory(mem); } }), + }; + + constructor(private llm: any) {} + + 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); } - constructor(private llm: any, private model?: string) {} - - /** - * Extracts facts from conversation and groups them into individual memories - * @param {string} conversation Full conversation formatted as [role]: content - * @param {Memory[]} memories The user's memory documents - * @param {LLMRequest} options LLM options - * @returns {Promise} Fact pools grouped by target document - */ - private async extract(conversation: string, memories: Memory[], options: LLMRequest): Promise { - const existingDocs = memories.map(m => `Name: ${m.name}\nDescription: ${m.description}`).join('\n\n'); - const pools: MemoryCollection[] = []; - await this.llm.ask(conversation, { - model: this.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 facts the USER explicitly stated about themselves or their business -- ONLY extract decisions that were MADE during this conversation -- DO NOT extract anything the AI said, its name, capabilities, or how it introduced itself -- DO NOT extract greetings, pleasantries or generic exchanges -- If nothing worth remembering was said, dont do anything, skip calling tools - -For each fact decide whether it belongs in an existing document or needs a new one, then call the \`extract_facts\` tool. - -Existing documents:\n${existingDocs || 'None yet.'}`, - tools: [this.tools.extract(pools)] - }); - return pools; + private createNode(name: string, memories: Memory[]): Memory { + const existing = memories.find(m => m.name === name); + if(existing) return existing; + return { + name, + description: '', + content: '', + embedding: [], + links: [], + backlinks: [], + }; } - /** - * Bot 2 - Editor: merges a pool of facts into a specific document using surgical line-based edits. - * Receives full document content and uses read + amend tools to make precise edits. - * @param {MemoryCollection} newMem The fact pool to merge - * @param {Memory[]} memories The user's memory documents - * @param {LLMRequest} options LLM options - */ - private async edit(newMem: MemoryCollection, memories: Memory[], options: LLMRequest): Promise { - const existing = memories.find(m => m.name === newMem.name); - const mem: Memory = existing || {name: newMem.name, description: newMem.description || '', content: '', embedding: []}; - const isNew = !existing; + private formatMemory(mem: Memory): string { + return [ + `# ${mem.name}`, + mem.description ? `> ${mem.description}` : '', + mem.links.length ? `**Links:** ${mem.links.map(l => `[[${l}]]`).join(', ')}` : '', + mem.backlinks.length ? `**Referenced by:** ${mem.backlinks.map(l => `[[${l}]]`).join(', ')}` : '', + '', + mem.content, + ].filter(l => l !== undefined).join('\n'); + } - await this.llm.ask(newMem.facts.map(f => `- ${f}`).join('\n'), + private listNodes(memories: Memory[]): MemoryRef[] { + return memories.map(m => ({name: m.name, description: m.description})); + } + + async recollect(query: string, memories: Memory[] | MemoryCache, limit = 5, graphDepth = 1): Promise { + const mem: Memory[] = memories instanceof MemoryCache ? memories.memories : 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(vectorResults.map(r => r.name)); + + // Graph expansion + 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); + } + + async memorize(history: LLMMessage[], memories: Memory[] | MemoryCache, options: LLMRequest): Promise { + const mem = memories instanceof MemoryCache ? memories.memories : memories; + 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 buckets = await this.factAgent(conversation, mem, options); + if(!buckets.length) return; + await Promise.all(buckets.map(async bucket => { + const node = await this.organizingAgent(bucket, mem, options); + if(!mem.find(m => m.name === node.name)) mem.push(node); + await this.docAgent(node, bucket, mem, options); + })); + + // Rebuild indexes + if (memories instanceof MemoryCache) { + memories.rebuildLinks(); + memories.rebuild(); + } else { + rebuildBacklinks(mem); + } + } + + private async docAgent(node: Memory, bucket: FactBucket, memories: Memory[], options: LLMRequest): Promise { + let finalContent = node.content; + + await this.llm.ask( + `New facts to integrate:\n${bucket.facts.map(f => `- ${f}`).join('\n')}`, { - model: this.model || options.model, - temperature: 0.2, - system: `You are a document editor. Merge the users list of facts into the following document using the \`edit_memory\` tool; call it as many times as necessary: -\`\`\` -${mem.content} + model: options.model, + temperature: 0.3, + system: `You are a knowledge base editor. Integrate the provided facts into the document below. + +Formatting rules: +- Use Obsidian-style markdown: # headings, **bold** for key terms, bullet lists for facts +- Link related concepts with [[WikiLink]] notation — only link things that are genuinely related +- You may create links to nodes that don't exist yet if the concept is important +- Keep the document concise, factual, and human-readable +- Resolve any contradictions between old content and new facts (new facts win) +- Do not add filler, preamble, or AI commentary — just clean knowledge documents + +All nodes: +${this.listNodes(memories).map(n => n.name).join(', ') || 'none'} + +Current document: +\`\`\`markdown +${node.content || '(empty — this is a new document)'} \`\`\``, - tools: [this.tools.edit(mem)] + tools: [{ + name: 'update_document', + description: 'Write the complete updated document content', + args: { + description: {type: 'string', description: 'One-line description of what this document covers, no formatting or emojis', required: true}, + content: {type: 'string', description: 'Fully updated document in markdown', required: true}, + }, + fn:(args: any) => { + node.description = args.description; + finalContent = args.content; + return 'Saved'; + } + }] } ); - if(isNew || mem.description !== existing?.description) { - const e = await this.llm.embedding(mem.description); - mem.embedding = e?.[0]?.embedding; - } - - if(isNew) memories.push(mem); - else { - const idx = memories.findIndex(m => m.name === newMem.name); - if(idx >= 0) memories[idx] = mem; + node.content = finalContent; + node.links = extractLinks(finalContent); + const needsEmbed = !node.embedding?.length || node.description !== memories.find(m => m.name === node.name)?.description; + if (needsEmbed) { + const [e] = await this.llm.embedding(node.description); + if (e) node.embedding = e.embedding; } } - /** - * Find relevant memory documents for a query using description embeddings - * @param {string} query The query to search against - * @param {Memory[]} memories The user's memory documents - * @param {number} limit Max number of results to return - * @returns {Promise} The most relevant memory documents - */ - async recollect(query: string, memories: Memory[], limit = 5): Promise { - const [e] = await this.llm.embedding(query); - return memories - .filter(m => m.embedding?.length) - .map(m => ({...m, score: this.llm.cosineSimilarity(m.embedding, e.embedding)})) - .toSorted((a: any, b: any) => b.score - a.score) - .slice(0, limit); + private async factAgent(conversation: string, memories: Memory[], options: LLMRequest): Promise { + const buckets: FactBucket[] = []; + + 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 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 +- If nothing worth remembering was said, do not call any tools + +Group facts by subject. For each group call \`extract_facts\` once. + +Known nodes (name: description): +${this.listNodes(memories).map(n => `- ${n.name}: ${n.description}`).join('\n') || 'None yet.'}`, + tools: [{ + name: 'extract_facts', + description: 'Submit a group of related facts for a specific subject', + args: { + subject: {type: 'string', description: 'Subject matter facts regard', required: true}, + facts: {type: 'string', description: 'Comma-separated list of extracted facts', required: true}, + }, + fn: (args: any) => { + buckets.push({ + subject: args.subject, + facts: args.facts.split(',').map((f: string) => f.trim()).filter(Boolean), + }); + return 'Recorded'; + } + }] + }); + + return buckets; } - /** - * Two-stage memory pipeline: classify facts from conversation history then surgically merge them into documents. - * Bot 1 (classify) extracts and groups facts cheaply. Bot 2 (edit) runs per-document in parallel with full content access. - * @param {LLMMessage[]} history Full conversation history to digest - * @param {Memory[]} memories The user's memory documents — mutated in place - * @param {LLMRequest} options LLM options - */ - async memorize(history: LLMMessage[], memories: Memory[], options: LLMRequest): Promise { - const conversation = history - .filter(h => h.role === 'user' || h.role === 'assistant') - .map(h => `[${h.role}]: ${h.content}`) - .join('\n\n'); - if(!conversation.trim()) return; - const pools = await this.extract(conversation, memories, options); - if(!pools.length) return; - await Promise.all(pools.map(pool => this.edit(pool, memories, options))); + private async organizingAgent(bucket: FactBucket, memories: Memory[], options: LLMRequest): Promise { + let candidates = this.listNodes(memories); + let attempts = 0; + const maxAttempts = 3; + + while (attempts++ < maxAttempts) { + let home = '', mode: string | null = null; + + const resp = await this.llm.ask(`Subject: ${bucket.subject}\n\nFacts:\n${bucket.facts.map(f => `- ${f}`).join('\n')}`, { + model: options.model, + temperature: 0.1, + system: `You are a knowledge organizer. Your job is to find the correct home for the supplied facts. + +1. Review the facts and the node list below. Pick the most likely match or decide if a new node is needed. +2. If you picked an existing node, use \`read\` to verify it's the right place. +- After reading, call either \`confirm\` (correct node) or \`mismatched\` (wrong node). +3. If none of the nodes match, call \`create\` to make a new node. + +Available nodes: +${candidates.map(n => `- ${n.name}: ${n.description}`).join('\n') || 'None — create a new node.'}`, + tools: [{ + name: 'read', + description: 'Read a node file to verify it is the right home for these facts', + args: {name: {type: 'string', description: 'Exact node name', required: true}}, + fn: ({name}) => { + const mem = memories.find(m => m.name === name); + if (!mem) return 'Node not found'; + home = name; + return this.formatMemory(mem); + } + }, { + name: 'confirm', + description: 'Confirm this is the correct node for the facts', + args: {}, + fn: () => { + mode = 'success'; + resp.abort(); + } + }, { + name: 'mismatched', + description: 'This is not the node you are looking for', + args: {}, + fn: () => { + mode = 'failed'; + resp.abort(); + } + }, { + name: 'create', + description: 'No existing node fits — create a new one', + args: {name: {type: 'string', description: 'Canonical name for the new node', required: true}}, + fn: ({name}) => { + home = name; + mode = 'create'; + resp.abort(); + } + }] + }); + + if(mode === 'create') { + return this.createNode(home, memories); + } else if (mode === 'failed') { + candidates = candidates.filter(c => c.name !== home); + if(!candidates.length) return this.createNode(bucket.subject, memories); + } else if (mode === 'success') { + const existing = memories.find(m => m.name === home); + return existing || this.createNode(home, memories); + } + } + return this.createNode(bucket.subject, memories); } }