Compare commits

..

2 Commits
1.6.5 ... 1.6.7

Author SHA1 Message Date
d42c240362 Memorization optimziations
All checks were successful
Publish Library / Build NPM Project (push) Successful in 1m2s
Publish Library / Tag Version (push) Successful in 10s
2026-08-31 12:38:40 -04:00
c1a16096ae Keep message progress on abort
All checks were successful
Publish Library / Build NPM Project (push) Successful in 43s
Publish Library / Tag Version (push) Successful in 14s
2026-08-29 21:18:28 -04:00
7 changed files with 300 additions and 98 deletions

12
package-lock.json generated
View File

@@ -1,19 +1,19 @@
{ {
"name": "@ztimson/ai-utils", "name": "@ztimson/ai-utils",
"version": "1.5.0", "version": "1.6.6",
"lockfileVersion": 3, "lockfileVersion": 3,
"requires": true, "requires": true,
"packages": { "packages": {
"": { "": {
"name": "@ztimson/ai-utils", "name": "@ztimson/ai-utils",
"version": "1.5.0", "version": "1.6.6",
"license": "MIT", "license": "MIT",
"dependencies": { "dependencies": {
"@anthropic-ai/sdk": "^0.102.0", "@anthropic-ai/sdk": "^0.102.0",
"@huggingface/transformers": "^4.2.0", "@huggingface/transformers": "^4.2.0",
"@tensorflow/tfjs": "^4.22.0", "@tensorflow/tfjs": "^4.22.0",
"@ztimson/node-utils": "^1.0.7", "@ztimson/node-utils": "^1.0.7",
"@ztimson/utils": "^0.29.4", "@ztimson/utils": "^0.30.8",
"cheerio": "^1.2.0", "cheerio": "^1.2.0",
"openai": "^6.42.0", "openai": "^6.42.0",
"pdf-parse": "^2.4.5", "pdf-parse": "^2.4.5",
@@ -1525,9 +1525,9 @@
"license": "MIT" "license": "MIT"
}, },
"node_modules/@ztimson/utils": { "node_modules/@ztimson/utils": {
"version": "0.29.7", "version": "0.30.8",
"resolved": "https://registry.npmjs.org/@ztimson/utils/-/utils-0.29.7.tgz", "resolved": "https://registry.npmjs.org/@ztimson/utils/-/utils-0.30.8.tgz",
"integrity": "sha512-cjQ9+RjC5X7gKNA/hJHDf7OtyYCa+5E0PDc76lIaATwNAxXCSx2IO9r2wHiHtZGV5bldjnFmw7aOV8Jmq7SgKQ==", "integrity": "sha512-+vBjcinqckqMHkP95xWiQeQz2E7Q1oS0b+Odjp+F9rvQ4z0US4JdodjqhEaqh+RAO/yP77x4Eu0a04yBB4HNdw==",
"license": "MIT", "license": "MIT",
"dependencies": { "dependencies": {
"var-persist": "^1.0.1" "var-persist": "^1.0.1"

View File

@@ -1,6 +1,6 @@
{ {
"name": "@ztimson/ai-utils", "name": "@ztimson/ai-utils",
"version": "1.6.5", "version": "1.6.7",
"description": "AI Utility library", "description": "AI Utility library",
"author": "Zak Timson", "author": "Zak Timson",
"license": "MIT", "license": "MIT",
@@ -29,7 +29,7 @@
"@huggingface/transformers": "^4.2.0", "@huggingface/transformers": "^4.2.0",
"@tensorflow/tfjs": "^4.22.0", "@tensorflow/tfjs": "^4.22.0",
"@ztimson/node-utils": "^1.0.7", "@ztimson/node-utils": "^1.0.7",
"@ztimson/utils": "^0.29.4", "@ztimson/utils": "^0.30.8",
"cheerio": "^1.2.0", "cheerio": "^1.2.0",
"openai": "^6.42.0", "openai": "^6.42.0",
"pdf-parse": "^2.4.5", "pdf-parse": "^2.4.5",

View File

@@ -4,7 +4,7 @@ import { Audio } from './audio.ts';
import {Vision} from './vision.ts'; import {Vision} from './vision.ts';
export type AbortablePromise<T> = Promise<T> & { export type AbortablePromise<T> = Promise<T> & {
abort: () => any abort: (keep?: boolean) => any
}; };
export type AiOptions = { export type AiOptions = {

View File

@@ -13,6 +13,56 @@ export function extractLinks(content: string): string[] {
return [...new Set([...matches].map(m => m[1].trim()))]; return [...new Set([...matches].map(m => m[1].trim()))];
} }
/**
* Incrementally patch the graph for a set of changed memories, instead of
* re-scanning every document. Only the changed memories' own content is
* re-parsed for links; affected targets have their backlinks patched.
* Does NOT handle node deletion — full rebuildGraph() is still required
* when a memory is removed, since that needs a backlink sweep across
* everyone who might reference it.
*/
export function patchGraph(mems: Memory[], nodes: MemoryNode[], changed: Memory[]): MemoryNode[] {
const nameSet = new Set(mems.map(m => m.name));
const byName = new Map(nodes.map(n => [n.name, n]));
const ensureNode = (name: string): MemoryNode => {
let n = byName.get(name);
if (!n) {
n = {name, missing: !nameSet.has(name), links: [], backlinks: []};
byName.set(name, n);
}
return n;
};
for (const m of changed) {
const node = ensureNode(m.name);
node.missing = false; // real memory, promotes any pre-existing ghost entry
const oldLinks = m.links ?? [];
const newLinks = extractLinks(m.content).filter(l => l !== m.name);
for (const target of oldLinks.filter(l => !newLinks.includes(l))) {
const t = byName.get(target);
if (!t) continue;
t.backlinks = t.backlinks.filter(n => n !== m.name);
if (t.missing && !t.backlinks.length) byName.delete(target); // fully dereferenced ghost
}
for (const target of newLinks.filter(l => !oldLinks.includes(l))) {
const t = ensureNode(target);
if (!t.backlinks.includes(m.name)) t.backlinks.push(m.name);
}
m.links = newLinks;
node.links = newLinks;
}
for (const m of mems) {
const n = byName.get(m.name);
if (n) m.backlinks = n.backlinks;
}
return [...byName.values()];
}
export function rebuildGraph(memories: Memory[] | MemoryCache): MemoryNode[] { export function rebuildGraph(memories: Memory[] | MemoryCache): MemoryNode[] {
const mems = memories instanceof MemoryCache ? memories.memories : memories; const mems = memories instanceof MemoryCache ? memories.memories : memories;
const nameSet = new Set(mems.map(m => m.name)); const nameSet = new Set(mems.map(m => m.name));

View File

@@ -15,6 +15,7 @@ interface KDNode<T> {
axis: number; axis: number;
left: KDNode<T> | null; left: KDNode<T> | null;
right: KDNode<T> | null; right: KDNode<T> | null;
deleted?: boolean;
} }
// ─── Distance helpers ───────────────────────────────────────────────────────── // ─── Distance helpers ─────────────────────────────────────────────────────────
@@ -95,6 +96,7 @@ class BoundedMaxHeap<T> {
* *
* Supports: * Supports:
* - Insertion of labeled points * - Insertion of labeled points
* - Lazy (tombstone) removal, physically purged on rebalance()
* - k-nearest-neighbor (KNN) search * - k-nearest-neighbor (KNN) search
* - Radius search (all points within a given distance) * - Radius search (all points within a given distance)
* - Euclidean and cosine distance metrics * - Euclidean and cosine distance metrics
@@ -103,6 +105,7 @@ class BoundedMaxHeap<T> {
export class KDTree<T = unknown> { export class KDTree<T = unknown> {
private root: KDNode<T> | null = null; private root: KDNode<T> | null = null;
private _size = 0; private _size = 0;
private _tombstones = 0;
private readonly distanceFn: (a: number[], b: number[]) => number; private readonly distanceFn: (a: number[], b: number[]) => number;
readonly dims: number; readonly dims: number;
@@ -129,9 +132,15 @@ export class KDTree<T = unknown> {
} }
} }
/** Total number of points stored in the tree. */ /** Total number of live points stored in the tree (excludes tombstoned). */
get size(): number { return this._size; } get size(): number { return this._size; }
/** Fraction of physical nodes that are tombstoned (pending removal on next rebalance). */
get tombstoneRatio(): number {
const total = this._size + this._tombstones;
return total ? this._tombstones / total : 0;
}
// ── Insertion ────────────────────────────────────────────────────────────── // ── Insertion ──────────────────────────────────────────────────────────────
/** /**
@@ -144,10 +153,36 @@ export class KDTree<T = unknown> {
this._size++; this._size++;
} }
// ── Removal ────────────────────────────────────────────────────────────────
/**
* Lazily remove all live points whose payload matches `predicate`.
* O(n) traversal, but avoids a full tree rebuild. Call `rebalance()`
* periodically (e.g. once tombstoneRatio crosses ~0.25) to reclaim space
* and restore optimal query depth.
* @returns number of points removed
*/
remove(predicate: (payload: T) => boolean): number {
let removed = 0;
const visit = (node: KDNode<T> | null): void => {
if (!node) return;
if (!node.deleted && predicate(node.point.payload)) {
node.deleted = true;
removed++;
}
visit(node.left);
visit(node.right);
};
visit(this.root);
this._size -= removed;
this._tombstones += removed;
return removed;
}
// ── KNN search ───────────────────────────────────────────────────────────── // ── KNN search ─────────────────────────────────────────────────────────────
/** /**
* Find the k nearest neighbors to `query`. * Find the k nearest live neighbors to `query`.
* Returns results sorted by distance ascending. * Returns results sorted by distance ascending.
*/ */
knn(query: number[], k: number): KNNResult<T>[] { knn(query: number[], k: number): KNNResult<T>[] {
@@ -171,7 +206,7 @@ export class KDTree<T = unknown> {
// ── Radius search ────────────────────────────────────────────────────────── // ── Radius search ──────────────────────────────────────────────────────────
/** /**
* Return all points whose distance to `query` is ≤ `radius`, * Return all live points whose distance to `query` is ≤ `radius`,
* sorted by distance ascending. * sorted by distance ascending.
*/ */
radiusSearch(query: number[], radius: number): KNNResult<T>[] { radiusSearch(query: number[], radius: number): KNNResult<T>[] {
@@ -186,7 +221,7 @@ export class KDTree<T = unknown> {
// ── Conversion ───────────────────────────────────────────────────────────── // ── Conversion ─────────────────────────────────────────────────────────────
/** Collect all points in the tree (order not guaranteed). */ /** Collect all live points in the tree (order not guaranteed). */
toArray(): KDPoint<T>[] { toArray(): KDPoint<T>[] {
const out: KDPoint<T>[] = []; const out: KDPoint<T>[] = [];
this.collect(this.root, out); this.collect(this.root, out);
@@ -194,12 +229,14 @@ export class KDTree<T = unknown> {
} }
/** /**
* Rebuild the tree from its current points as a balanced tree. * Rebuild the tree from its current live points as a balanced tree.
* Useful after many individual insertions to restore O(log n) query time. * Physically purges tombstones and restores O(log n) query time.
*/ */
rebalance(): void { rebalance(): void {
const points = this.toArray(); const points = this.toArray();
this.root = points.length ? this.buildBalanced(points, 0) : null; this.root = points.length ? this.buildBalanced(points, 0) : null;
this._size = points.length;
this._tombstones = 0;
} }
// ── Private: build ───────────────────────────────────────────────────────── // ── Private: build ─────────────────────────────────────────────────────────
@@ -251,8 +288,10 @@ export class KDTree<T = unknown> {
): void { ): void {
if (node === null) return; if (node === null) return;
const dist = this.distanceFn(query, node.point.vector); if (!node.deleted) {
heap.push({ point: node.point, distance: dist }); const dist = this.distanceFn(query, node.point.vector);
heap.push({ point: node.point, distance: dist });
}
const axis = node.axis; const axis = node.axis;
const diff = query[axis] - node.point.vector[axis]; const diff = query[axis] - node.point.vector[axis];
@@ -285,9 +324,11 @@ export class KDTree<T = unknown> {
): void { ): void {
if (node === null) return; if (node === null) return;
const dist = this.distanceFn(query, node.point.vector); if (!node.deleted) {
if (dist <= radius) { const dist = this.distanceFn(query, node.point.vector);
results.push({ point: node.point, distance: dist }); if (dist <= radius) {
results.push({ point: node.point, distance: dist });
}
} }
const axis = node.axis; const axis = node.axis;
@@ -310,7 +351,7 @@ export class KDTree<T = unknown> {
private collect(node: KDNode<T> | null, out: KDPoint<T>[]): void { private collect(node: KDNode<T> | null, out: KDPoint<T>[]): void {
if (node === null) return; if (node === null) return;
out.push(node.point); if (!node.deleted) out.push(node.point);
this.collect(node.left, out); this.collect(node.left, out);
this.collect(node.right, out); this.collect(node.right, out);
} }

View File

@@ -265,7 +265,7 @@ class LLM {
}; };
} }
private setupAgent(agents: Agent[] = [], allAgents: Agent[], history: LLMMessage[], aborts: (() => void)[], depth = 0, delegateState: {resp: string | null}): AiTool[] { private setupAgent(agents: Agent[] = [], allAgents: Agent[], history: LLMMessage[], aborts: ((keep?: boolean) => void)[], depth = 0, delegateState: {resp: string | null}): AiTool[] {
return agents.map(a => { return agents.map(a => {
const toolName = `${a.delegate ? '' : 'sub'}agent_${snakeCase(a.name)}`; const toolName = `${a.delegate ? '' : 'sub'}agent_${snakeCase(a.name)}`;
return { return {
@@ -397,11 +397,13 @@ ${a.system}`,
if(!this.models[m]) throw new Error(`Model does not exist: ${m}`); if(!this.models[m]) throw new Error(`Model does not exist: ${m}`);
let request: AbortablePromise<string> | null = null; let request: AbortablePromise<string> | null = null;
let aborted = false; let aborted = false;
const nestedAborts: (() => void)[] = []; let keepOnAbort = true;
const abort = () => { const nestedAborts: ((keep?: boolean) => void)[] = [];
const abort = (keep = true) => {
aborted = true; aborted = true;
request?.abort?.(); keepOnAbort = keep;
nestedAborts.forEach(a => a()); request?.abort?.(keep);
nestedAborts.forEach(a => a(keep));
}; };
let promise: any; let promise: any;
@@ -411,9 +413,25 @@ ${a.system}`,
let tools: AiTool[] = options.tools || this.ai.options.llm?.tools || []; let tools: AiTool[] = options.tools || this.ai.options.llm?.tools || [];
const prompts: string[] = []; const prompts: string[] = [];
let history = options.history || []; let history = options.history || [];
const historyStart = history.length;
const files = options.files || []; const files = options.files || [];
if(message || files.length) history.push({role: 'user', content: message || '', timestamp: Date.now()}); if(message || files.length) history.push({role: 'user', content: message || '', timestamp: Date.now()});
// Accumulate streamed text so it can be committed to history if aborted mid-generation
let partialText = '';
const onStream = options.stream;
const stream = (chunk: {text?: string, tool?: string, done?: true}) => {
if(chunk.text) partialText += chunk.text;
return onStream?.(chunk);
};
/** Commit (keep) or discard this turn's progress on abort, then throw */
const abortNow = (): never => {
if(keepOnAbort) { if(partialText) history.push({role: 'assistant', content: partialText, timestamp: Date.now()}); }
else history.splice(historyStart, history.length - historyStart);
throw Object.assign(new Error('Aborted'), {name: 'AbortError'});
};
// MCP // MCP
const mcp = options.mcp || this.ai.options?.llm?.mcp; const mcp = options.mcp || this.ai.options?.llm?.mcp;
if(mcp?.length) { if(mcp?.length) {
@@ -441,8 +459,8 @@ ${a.system}`,
const mems = mem.memory instanceof MemoryCache ? mem.memory.memories : mem.memory; const mems = mem.memory instanceof MemoryCache ? mem.memory.memories : mem.memory;
if(mems.length) { if(mems.length) {
if(mem.inject) { if(mem.inject) {
const pool = 15; // candidates considered, cheap since only refs are listed const pool = 15;
const budget = mem.maxTokens ?? 2000; // actual content injected const budget = mem.maxTokens ?? 2000;
const relevant = await this.memoryManager.recollect(message, mem.memory, pool); const relevant = await this.memoryManager.recollect(message, mem.memory, pool);
let used = 0; let used = 0;
@@ -481,7 +499,7 @@ Linked: ${makeUnique([...r.links, ...r.backlinks]).join(', ')}
} }
} }
if(aborted) throw Object.assign(new Error('Aborted'), {name: 'AbortError'}); if(aborted) abortNow();
const lastMsg = history[history.length - 1]; const lastMsg = history[history.length - 1];
if(files.length && lastMsg?.role === 'user') lastMsg.files = files; if(files.length && lastMsg?.role === 'user') lastMsg.files = files;
@@ -500,11 +518,17 @@ Linked: ${makeUnique([...r.links, ...r.backlinks]).join(', ')}
const toolTimings = new Map<string, {duration: number, tps: number}>(); const toolTimings = new Map<string, {duration: number, tps: number}>();
tools = this.wrapToolTiming(tools, toolTimings); tools = this.wrapToolTiming(tools, toolTimings);
if(aborted) throw Object.assign(new Error('Aborted'), {name: 'AbortError'}); if(aborted) abortNow();
prompts.unshift(options.system || this.ai.options.llm?.system || ''); prompts.unshift(options.system || this.ai.options.llm?.system || '');
request = this.models[m].ask('', {...options, tools, system: prompts.filter(Boolean).join('\n\n')}); request = this.models[m].ask('', {...options, tools, stream, system: prompts.filter(Boolean).join('\n\n')});
let resp = await request; let resp: string;
try {
resp = await request;
} catch(err: any) {
if(aborted) return abortNow();
throw err;
}
// Strip the file injection shim // Strip the file injection shim
restores.forEach(({msg, content}) => msg.content = content); restores.forEach(({msg, content}) => msg.content = content);

View File

@@ -1,11 +1,13 @@
import {MemoryNode, rebuildGraph} from './helpers.ts'; import {MemoryNode, patchGraph, rebuildGraph} from './helpers.ts';
import {LLMRequest, LLMMessage} from './llm.ts'; import {LLMRequest, LLMMessage} from './llm.ts';
import {AiTool} from './tools.ts'; import {AiTool} from './tools.ts';
import {KDPoint, KDTree} from './kd-tree.ts'; import {KDTree} from './kd-tree.ts';
import {escapeRegex} from '@ztimson/utils'; import {escapeRegex} from '@ztimson/utils';
const MERGE_THRESHOLD = 0.12; const MERGE_THRESHOLD = 0.12;
const PENDING_HEADING = '## Pending'; const PENDING_HEADING = '## Pending';
const TREE_TOMBSTONE_LIMIT = 0.25;
const ALIAS_MATCH_THRESHOLD = 0.55;
const GENERIC_TEMPLATE = `# {{Title}} const GENERIC_TEMPLATE = `# {{Title}}
## Summary ## Summary
@@ -18,7 +20,12 @@ export type Memory = {
name: string; name: string;
description: string; description: string;
content: string; content: string;
/** Description embedding — indexed in the KD tree, used for merge/ANN candidate lookup */
embedding: number[]; embedding: number[];
/** Title-only embedding, weighted heaviest during recall ranking */
titleEmbedding?: number[];
/** Chunked body embeddings, best-chunk match used during recall ranking */
bodyEmbeddings?: number[][];
links: string[]; links: string[];
backlinks: string[]; backlinks: string[];
} }
@@ -26,6 +33,8 @@ export type Memory = {
type MemoryRef = { type MemoryRef = {
name: string; name: string;
description: string; description: string;
/** Cosine distance from the query, present when returned from a search */
distance?: number;
} }
type FactBucket = { type FactBucket = {
@@ -61,10 +70,20 @@ function cosineDistance(a: number[], b: number[]): number {
function cosineSearch(query: number[], memories: Memory[], limit: number): MemoryRef[] { function cosineSearch(query: number[], memories: Memory[], limit: number): MemoryRef[] {
return memories return memories
.filter(m => m.embedding?.length) .filter(m => m.embedding?.length)
.map(m => ({ref: {name: m.name, description: m.description}, distance: cosineDistance(query, m.embedding)})) .map(m => ({name: m.name, description: m.description, distance: cosineDistance(query, m.embedding)}))
.sort((a, b) => a.distance - b.distance) .sort((a, b) => a.distance - b.distance)
.slice(0, limit) .slice(0, limit);
.map(s => s.ref); }
/** Re-embed a node's title / description / body fields. Description embedding stays the KD-tree index key. */
async function embedMemoryFields(node: Memory, llm: any): Promise<void> {
const body = stripHeader(node.content);
const [titleE] = await llm.embedding(node.name.split('/').pop() || node.name);
const [descE] = await llm.embedding(node.description || '');
const bodyChunks = body ? await llm.embedding(body) : [];
if (titleE) node.titleEmbedding = titleE.embedding;
if (descE) node.embedding = descE.embedding;
node.bodyEmbeddings = bodyChunks.map((c: any) => c.embedding).filter(Boolean);
} }
export function stripHeader(content: string): string { export function stripHeader(content: string): string {
@@ -73,6 +92,8 @@ export function stripHeader(content: string): string {
export class MemoryCache { export class MemoryCache {
private tree!: KDTree<MemoryRef>; private tree!: KDTree<MemoryRef>;
/** Tracks which memories are currently indexed in the tree, keyed by name -> embedding reference */
private indexed = new Map<string, number[]>();
public memories: Memory[]; public memories: Memory[];
public nodes: MemoryNode[] = []; public nodes: MemoryNode[] = [];
@@ -80,37 +101,48 @@ export class MemoryCache {
constructor(memories: Memory[]) { constructor(memories: Memory[]) {
this.memories = memories; this.memories = memories;
this.tree = new KDTree<MemoryRef>(0);
this.rebuild(); this.rebuild();
} }
private buildTree(): KDTree<MemoryRef> { /** Incrementally sync the KD tree against `this.memories` instead of rebuilding from scratch */
const embedded = this.memories.filter(m => m.embedding?.length); private syncTree(): void {
if (!embedded.length) return new KDTree<MemoryRef>(0); const current = new Set(this.memories.map(m => m.name));
const dims = embedded[0].embedding.length; for (const [name, emb] of [...this.indexed]) {
const points: KDPoint<MemoryRef>[] = embedded.map(m => ({ const mem = this.memories.find(m => m.name === name);
vector: m.embedding, if (!mem || !current.has(name) || mem.embedding !== emb) {
payload: {name: m.name, description: m.description}, this.tree.remove(p => p.name === name);
})); this.indexed.delete(name);
}
}
return new KDTree<MemoryRef>(dims, 'cosine', points); for (const mem of this.memories) {
if (!mem.embedding?.length || this.indexed.has(mem.name)) continue;
if (this.tree.dims === 0) this.tree = new KDTree<MemoryRef>(mem.embedding.length, 'cosine');
if (mem.embedding.length !== this.tree.dims) continue; // guard against embedding model/dim drift
this.tree.insert({vector: mem.embedding, payload: {name: mem.name, description: mem.description}});
this.indexed.set(mem.name, mem.embedding);
}
if (this.tree.tombstoneRatio > TREE_TOMBSTONE_LIMIT) this.tree.rebalance();
} }
search(query: number[], limit: number): MemoryRef[] { search(query: number[], limit: number): MemoryRef[] {
if (!this.tree || this.tree.dims === 0) return []; if (!this.tree || this.tree.dims === 0) return [];
return this.tree.knn(query, limit).map(r => r.point.payload); return this.tree.knn(query, limit).map(r => ({...r.point.payload, distance: r.distance}));
} }
add(memory: Memory): void { add(memory: Memory): void {
this.memories.push(memory); this.memories.push(memory);
this.rebuild(); this.rebuild([memory]);
} }
update(memory: Memory): void { update(memory: Memory): void {
const existing = this.memories.find(m => m.name === memory.name); const existing = this.memories.find(m => m.name === memory.name);
if (existing) Object.assign(existing, memory); if (existing) Object.assign(existing, memory);
else this.memories.push(memory); else this.memories.push(memory);
this.rebuild(); this.rebuild([existing ?? memory]);
} }
remove(name: string): void { remove(name: string): void {
@@ -121,9 +153,11 @@ export class MemoryCache {
} }
} }
rebuild(): void { rebuild(changed?: Memory[]): void {
this.nodes = rebuildGraph(this.memories); this.nodes = (changed?.length && this.nodes.length)
this.tree = this.buildTree(); ? patchGraph(this.memories, this.nodes, changed)
: rebuildGraph(this.memories);
this.syncTree();
} }
} }
@@ -140,9 +174,9 @@ class MemoryAccessor {
return this.list.find(m => m.name === name); return this.list.find(m => m.name === name);
} }
commit(): MemoryNode[] { commit(changed?: Memory[]): MemoryNode[] {
if (this.cache) { if (this.cache) {
this.cache.rebuild(); this.cache.rebuild(changed);
return this.cache.nodes; return this.cache.nodes;
} }
return rebuildGraph(this.list); return rebuildGraph(this.list);
@@ -153,6 +187,7 @@ class MemoryAccessor {
return nodes.filter(n => n.missing).map(n => n.name); return nodes.filter(n => n.missing).map(n => n.name);
} }
/** Cache path uses the KD tree's knn(); raw-array path (no cache available) falls back to a linear cosine scan */
search(vector: number[], limit: number): MemoryRef[] { search(vector: number[], limit: number): MemoryRef[] {
return this.cache ? this.cache.search(vector, limit) : cosineSearch(vector, this.list, limit); return this.cache ? this.cache.search(vector, limit) : cosineSearch(vector, this.list, limit);
} }
@@ -168,10 +203,7 @@ class MemoryAccessor {
async backfillEmbeddings(llm: any): Promise<number> { async backfillEmbeddings(llm: any): Promise<number> {
const missing = this.list.filter(m => !m.embedding?.length); const missing = this.list.filter(m => !m.embedding?.length);
if (!missing.length) return 0; if (!missing.length) return 0;
await Promise.all(missing.map(async node => { await Promise.all(missing.map(node => embedMemoryFields(node, llm)));
const [e] = await llm.embedding(`${node.description}\n\n${stripHeader(node.content)}`.trim());
if (e) node.embedding = e.embedding;
}));
this.commit(); this.commit();
return missing.length; return missing.length;
} }
@@ -282,6 +314,39 @@ ${m.content}
for (const m of memories) if (pattern.test(m.content)) m.content = m.content.replace(pattern, `[[${to}]]`); for (const m of memories) if (pattern.test(m.content)) m.content = m.content.replace(pattern, `[[${to}]]`);
} }
private normalizeLeaf(name: string): string {
return name.trim().toLowerCase().replace(/\s+/g, ' ');
}
/**
* Resolve a fact-agent proposed subject to an existing node when it's an alias/rename of one.
* Exact match is checked first (cheap, and covers the common case since node names are
* already normalized at creation time). Only falls through to fuzzy alias matching against
* same-root candidates when there's no existing hit — i.e. only on likely-new-doc creation.
*/
private resolveSubject(subject: string, store: MemoryAccessor): string {
const trimmed = subject.trim();
const exact = store.find(trimmed);
if (exact) return exact.name;
const normalized = this.normalizeLeaf(trimmed);
const caseInsensitive = store.list.find(m => this.normalizeLeaf(m.name) === normalized);
if (caseInsensitive) return caseInsensitive.name;
const root = trimmed.split('/')[0];
const leaf = trimmed.split('/').slice(1).join('/') || trimmed;
const candidates = store.list.filter(m => m.name.split('/')[0] === root && m.name !== trimmed);
if (!candidates.length) return trimmed;
// fuzzyMatch requires >=2 terms; pad with an empty string when there's only one candidate
const leaves = candidates.map(m => m.name.split('/').slice(1).join('/') || m.name);
const probe = leaves.length > 1 ? leaves : [...leaves, ''];
const {max, similarities} = this.llm.fuzzyMatch(leaf, ...probe);
if (max >= ALIAS_MATCH_THRESHOLD) return candidates[similarities.indexOf(max)].name;
return trimmed;
}
private async factAgent(conversation: string, store: MemoryAccessor, options: LLMRequest): Promise<FactAgentResult> { private async factAgent(conversation: string, store: MemoryAccessor, options: LLMRequest): Promise<FactAgentResult> {
const ghosts = store.ghosts(); const ghosts = store.ghosts();
@@ -301,7 +366,7 @@ ${m.content}
- Extract the final/end state, not deltas - Extract the final/end state, not deltas
Path assignment rules: Path assignment rules:
- Reuse existing node names whenever possible - Reuse existing node names whenever possible, including when the subject is an alias/nickname of an existing node (e.g. "Rob" referring to an existing "People/Robert")
- Documents should be grouped and named by the root subject - Documents should be grouped and named by the root subject
- Person → People/Name - Person → People/Name
- Project → Projects/Name - Project → Projects/Name
@@ -351,23 +416,21 @@ ${ghosts.length ? `${ghosts.map(g => `- ${g}: (Ghost)`).join('\n')}` : ''}`,
return memories.map(m => ({name: m.name, description: m.description})); return memories.map(m => ({name: m.name, description: m.description}));
} }
/** Finds the closest merge candidate via the KD tree's knn() instead of a manual O(n) cosine scan */
private async checkMerge(node: Memory, memories: Memory[] | MemoryCache, options: LLMRequest, threshold = MERGE_THRESHOLD): Promise<Memory | 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; if (!node.embedding?.length || node.name.startsWith('Journal/')) return null;
const store = this.access(memories); const store = this.access(memories);
let closest: Memory | null = null, closestDist = Infinity; const candidate = store.search(node.embedding, 5)
for (const other of store.list) { .find(r => r.name !== node.name && !r.name.startsWith('Journal/') && r.distance !== undefined && r.distance <= threshold);
if (other.name === node.name || other.name.startsWith('Journal/') || !other.embedding?.length) continue; if (!candidate) return null;
const d = cosineDistance(node.embedding, other.embedding); const closest = store.find(candidate.name);
if (d < closestDist) { closestDist = d; closest = other; } if (!closest) return null;
}
if (!closest || closestDist > threshold) return null;
const result = await this.mergeAgent(node, closest, options); const result = await this.mergeAgent(node, closest, options);
const merged: Memory = {name: result.name, description: this.sanitizeDescription(result.description), content: '', embedding: [], links: [], backlinks: []}; const merged: Memory = {name: result.name, description: this.sanitizeDescription(result.description), content: '', embedding: [], links: [], backlinks: []};
merged.content = this.touchHeader(merged, result.content); merged.content = this.touchHeader(merged, result.content);
const [e] = await this.llm.embedding(`${merged.description}\n\n${result.content}`.trim()); await embedMemoryFields(merged, this.llm);
if (e) merged.embedding = e.embedding;
this.relink(store.list, node.name, merged.name); this.relink(store.list, node.name, merged.name);
this.relink(store.list, closest.name, merged.name); this.relink(store.list, closest.name, merged.name);
@@ -396,18 +459,20 @@ ${ghosts.length ? `${ghosts.map(g => `- ${g}: (Ghost)`).join('\n')}` : ''}`,
this.queues.set(key, entry); this.queues.set(key, entry);
const store = this.access(memories); const store = this.access(memories);
entry.task = (async () => { entry.task = (async () => {
let current = node; let current = node, merged = false;
do { try {
entry.dirty = false; do {
await this.docAgent(current, store.list, options, entry); entry.dirty = false;
this.mergeLock = this.mergeLock.then(() => this.checkMerge(current, memories, options)); await this.docAgent(current, store.list, options, entry);
const merged = await this.mergeLock; this.mergeLock = this.mergeLock.then(() => this.checkMerge(current, memories, options));
if(merged) current = merged; const result = await this.mergeLock;
} while (entry.dirty); if (result) { current = result; merged = true; }
})().finally(() => { } while (entry.dirty);
this.queues.delete(key); } finally {
store.commit(); store.commit(merged ? undefined : [node]);
}); this.queues.delete(key);
}
})();
return entry.task; return entry.task;
} }
@@ -434,7 +499,7 @@ ${GENERIC_TEMPLATE}
\`\`\` \`\`\`
Rules: Rules:
- Contradictions: newer facts always win — delete outdated statements entirely - Contradictions: "## Pending" holds the newest information — bias toward it. Fold it in as the standing fact and drop the outdated statement, unless the old context adds meaningful nuance (e.g. "previously X, now Y"). This document should read as a source of truth, not an audit log
- Journals (Journal/...): keep entries as a chronological timeline; clean up grammar within entries but never delete history - 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 - 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 - Link specific entities and concepts with [[WikiLink]] (e.g., [[Projects/KiwixServer]]); skip generics
@@ -462,11 +527,12 @@ ${currentBody}
if (!update?.content) return; if (!update?.content) return;
node.description = node.name !== 'People/User' ? this.sanitizeDescription(update.description) : 'All information about the current user'; node.description = node.name !== 'People/User' ? this.sanitizeDescription(update.description) : 'All information about the current user';
node.content = this.touchHeader(node, update.content); node.content = this.touchHeader(node, update.content);
const [e] = await this.llm.embedding(`${node.description}\n\n${update.content}`.trim()); await embedMemoryFields(node, this.llm);
if (e) node.embedding = e.embedding;
} }
private async mergeAgent(a: Memory, b: Memory, options: LLMRequest): Promise<{name: string, description: string, content: string}> { private async mergeAgent(a: Memory, b: Memory, options: LLMRequest): Promise<{name: string, description: string, content: string}> {
const modifiedOf = (m: Memory) => this.parseFrontmatter(m.content).fm.get('modified') || 'unknown';
return this.llm.ask('', { return this.llm.ask('', {
model: options.model, model: options.model,
temperature: 0.3, temperature: 0.3,
@@ -475,21 +541,21 @@ ${currentBody}
description: {type: 'string', description: 'One factual sentence describing the merged document\'s subject matter', 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}, 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. system: `You are a knowledge base editor merging two overlapping Obsidian documents into one.
Structure loosely: Structure loosely:
\`\`\`markdown \`\`\`markdown
${GENERIC_TEMPLATE} ${GENERIC_TEMPLATE}
\`\`\` \`\`\`
Combine both documents, resolve duplication and contradictions. Combine both documents, resolve duplication. On contradictions, bias toward whichever document was modified more recently; drop the outdated statement unless the old context adds meaningful nuance.
Document A ("${a.name}"): Document A ("${a.name}", last modified ${modifiedOf(a)}):
\`\`\`markdown \`\`\`markdown
${stripHeader(a.content)} ${stripHeader(a.content)}
\`\`\` \`\`\`
Document B ("${b.name}"): Document B ("${b.name}", last modified ${modifiedOf(b)}):
\`\`\`markdown \`\`\`markdown
${stripHeader(b.content)} ${stripHeader(b.content)}
\`\`\``, \`\`\``,
@@ -512,12 +578,17 @@ ${stripHeader(b.content)}
return {fm, body: match[2]}; return {fm, body: match[2]};
} }
/**
* Writes the code-owned frontmatter block. `body` is passed through stripHeader() first so a
* model that ignores instructions and hallucinates its own `---` block can never corrupt or
* duplicate the real frontmatter — the LLM only ever gets to influence the body.
*/
private touchHeader(node: Memory, body: string): string { private touchHeader(node: Memory, body: string): string {
const {fm} = this.parseFrontmatter(node.content); const {fm} = this.parseFrontmatter(node.content);
fm.set('name', node.name); fm.set('name', node.name);
fm.set('description', node.description || ''); fm.set('description', node.description || '');
fm.set('modified', new Date().toISOString()); fm.set('modified', new Date().toISOString());
return this.writeFrontmatter(fm, body); return this.writeFrontmatter(fm, stripHeader(body));
} }
private writeFrontmatter(fm: Map<string, string>, body: string): string { private writeFrontmatter(fm: Map<string, string>, body: string): string {
@@ -540,6 +611,19 @@ ${stripHeader(b.content)}
return this.access(memories).forget(name); return this.access(memories).forget(name);
} }
/** Ranks a candidate pool by weighted title/description/body similarity against the query embedding */
private rankByFields(query: number[], candidates: Memory[], limit: number): Memory[] {
const scored = candidates.map(m => {
const titleSim = m.titleEmbedding?.length ? 1 - cosineDistance(query, m.titleEmbedding) : 0;
const descSim = m.embedding?.length ? 1 - cosineDistance(query, m.embedding) : 0;
const bodySim = m.bodyEmbeddings?.length
? Math.max(...m.bodyEmbeddings.map(b => 1 - cosineDistance(query, b)))
: 0;
return {memory: m, score: titleSim * 0.5 + descSim * 0.35 + bodySim * 0.15};
});
return scored.sort((a, b) => b.score - a.score).slice(0, limit).map(s => s.memory);
}
async recollect(query: string, memories: Memory[] | MemoryCache, limit = 5, graphDepth = 1): Promise<Memory[]> { async recollect(query: string, memories: Memory[] | MemoryCache, limit = 5, graphDepth = 1): Promise<Memory[]> {
const store = this.access(memories); const store = this.access(memories);
if (!store.list.length) return []; if (!store.list.length) return [];
@@ -549,8 +633,11 @@ ${stripHeader(b.content)}
const [e] = await this.llm.embedding(query); const [e] = await this.llm.embedding(query);
if (!e) return []; if (!e) return [];
const vectorResults = store.search(e.embedding, limit); // Description embedding is the cheap ANN index key; pull a wider pool then re-rank by field weight
const found = new Set<string>(vectorResults.map(r => r.name)); const pool = store.search(e.embedding, Math.max(limit * 3, limit));
const poolMemories = pool.map(r => store.find(r.name)).filter((m): m is Memory => !!m);
const ranked = this.rankByFields(e.embedding, poolMemories, limit);
const found = new Set<string>(ranked.map(m => m.name));
if (graphDepth > 0) { if (graphDepth > 0) {
let frontier = [...found]; let frontier = [...found];
@@ -570,9 +657,9 @@ ${stripHeader(b.content)}
} }
} }
const vectorOrder = vectorResults.map(r => r.name); const rankedOrder = ranked.map(m => m.name);
const graphExpansions = [...found].filter(n => !vectorOrder.includes(n)); const graphExpansions = [...found].filter(n => !rankedOrder.includes(n));
return [...vectorOrder, ...graphExpansions].map(n => store.find(n)!).filter(Boolean); return [...rankedOrder, ...graphExpansions].map(n => store.find(n)!).filter(Boolean);
} }
async memorize(history: LLMMessage[], memories: Memory[] | MemoryCache, options: LLMRequest): Promise<Memory[]> { async memorize(history: LLMMessage[], memories: Memory[] | MemoryCache, options: LLMRequest): Promise<Memory[]> {
@@ -601,9 +688,10 @@ ${stripHeader(b.content)}
} }
for (const {subject, facts} of buckets) { for (const {subject, facts} of buckets) {
let node = store.find(subject); const resolved = this.resolveSubject(subject, store);
let node = store.find(resolved);
if (!node) { if (!node) {
node = {name: subject, description: '', content: '', embedding: [], links: [], backlinks: []}; node = {name: resolved, description: '', content: '', embedding: [], links: [], backlinks: []};
store.list.push(node); store.list.push(node);
} }
this.stage(node, facts.map(f => `- ${f}`).join('\n')); this.stage(node, facts.map(f => `- ${f}`).join('\n'));
@@ -611,13 +699,12 @@ ${stripHeader(b.content)}
} }
await Promise.all(touched.map(async node => { await Promise.all(touched.map(async node => {
const [e] = await this.llm.embedding(`${node.description}\n\n${stripHeader(node.content)}`.trim()); await embedMemoryFields(node, this.llm);
if (e) node.embedding = e.embedding;
this.touch(node.name); this.touch(node.name);
})); }));
if (touched.length) { if (touched.length) {
store.commit(); store.commit(touched);
(pending as any).content = `Saved to ${touched.map(n => `[[${n.name}]]`).join(', ')}`; (pending as any).content = `Saved to ${touched.map(n => `[[${n.name}]]`).join(', ')}`;
Promise.all(touched.map(node => this.reconcile(node, memories, options).catch(() => {}))); Promise.all(touched.map(node => this.reconcile(node, memories, options).catch(() => {})));
} else { } else {