484 lines
17 KiB
TypeScript
484 lines
17 KiB
TypeScript
import {LLMRequest, LLMMessage} from './llm.ts';
|
||
import {MemoryCache} from './memory-cache.ts';
|
||
import {AiTool} from './tools.ts';
|
||
|
||
export type Memory = {
|
||
name: string;
|
||
description: string;
|
||
content: string;
|
||
embedding: number[];
|
||
}
|
||
|
||
export type MemoryRef = {
|
||
name: string;
|
||
description: string;
|
||
}
|
||
|
||
export type FactBucket = {
|
||
subject: string;
|
||
facts: string[];
|
||
}
|
||
|
||
export type MemoryNode = {
|
||
name: string;
|
||
missing: boolean;
|
||
links: string[];
|
||
backlinks: string[];
|
||
}
|
||
|
||
function extractLinks(content: string): string[] {
|
||
if(!content) return [];
|
||
const matches = content.matchAll(/\[\[([^\]]+)\]\]/g);
|
||
return [...new Set([...matches].map(m => m[1].trim()))];
|
||
}
|
||
|
||
export function extractMetadata(content: string): {links: string[], backlinks: string[]} {
|
||
const match = content.match(/^---\n([\s\S]*?)\n---/);
|
||
if (!match) return {links: [], backlinks: []};
|
||
|
||
const fm = match[1];
|
||
const getList = (key: string): string[] => {
|
||
const m = fm.match(new RegExp(`^${key}:\\s*\\[(.*)\\]$`, 'm'));
|
||
if (!m || !m[1].trim()) return [];
|
||
return m[1].split(',').map(s => s.trim().replace(/^"|"$/g, '')).filter(Boolean);
|
||
};
|
||
|
||
return {
|
||
links: getList('links'),
|
||
backlinks: getList('backlinks'),
|
||
};
|
||
}
|
||
|
||
function dedupeFacts(facts: string[]): string[] {
|
||
const seen = new Map<string, string>();
|
||
for (const f of facts) {
|
||
const clean = f.trim();
|
||
if (clean) seen.set(clean.toLowerCase(), clean);
|
||
}
|
||
return [...seen.values()];
|
||
}
|
||
|
||
function cosineDistance(a: number[], b: number[]): number {
|
||
let dot = 0, normA = 0, normB = 0;
|
||
for (let i = 0; i < a.length; i++) {
|
||
dot += a[i] * b[i];
|
||
normA += a[i] * a[i];
|
||
normB += b[i] * b[i];
|
||
}
|
||
const denom = Math.sqrt(normA) * Math.sqrt(normB);
|
||
return denom === 0 ? 1 : 1 - dot / denom;
|
||
}
|
||
|
||
function getWeekMonday(date: Date = new Date()): string {
|
||
const d = new Date(Date.UTC(date.getFullYear(), date.getMonth(), date.getDate()));
|
||
const day = d.getUTCDay();
|
||
const diff = day === 0 ? -6 : 1 - day;
|
||
d.setUTCDate(d.getUTCDate() + diff);
|
||
return d.toISOString().slice(0, 10);
|
||
}
|
||
|
||
function getWeekSunday(monday: string): string {
|
||
const d = new Date(`${monday}T00:00:00Z`);
|
||
d.setUTCDate(d.getUTCDate() + 6);
|
||
return d.toISOString().slice(0, 10);
|
||
}
|
||
|
||
|
||
|
||
export class MemoryManager {
|
||
private pendingMemorizations = new Map<string, {
|
||
memories: Memory[] | MemoryCache,
|
||
tempMemoryName: string,
|
||
timestamp: number,
|
||
}>();
|
||
|
||
private queues = new Map<string, {
|
||
pending: string[],
|
||
request: {abort?: () => void} | null,
|
||
task: Promise<void>,
|
||
}>();
|
||
|
||
tools = {
|
||
read: (memories: Memory[] | MemoryCache): AiTool => ({
|
||
name: 'memory_recall',
|
||
description: 'Read the full content of a memory document',
|
||
args: {
|
||
name: {type: 'string', description: 'Exact memory name', required: true},
|
||
},
|
||
fn: (args: any) => {
|
||
const mems = memories instanceof MemoryCache ? memories.memories : memories;
|
||
const mem = mems.find(m => m.name === args.name);
|
||
if (!mem) return 'Document not found';
|
||
return mem.content;
|
||
},
|
||
}),
|
||
|
||
forget: (memories: Memory[] | MemoryCache): AiTool => ({
|
||
name: 'memory_forget',
|
||
description: 'Permanently delete a memory document and clean up all references to it',
|
||
args: {
|
||
name: {type: 'string', description: 'Exact memory name to forget', required: true}
|
||
},
|
||
fn: (args: any) => {
|
||
const result = this.forget(args.name, memories);
|
||
return result ? `Forgotten: ${args.name}` : `Not found: ${args.name}`;
|
||
},
|
||
}),
|
||
};
|
||
|
||
constructor(private llm: any) {}
|
||
|
||
private async createTempMemory(conversation: string): Promise<Memory> {
|
||
const timestamp = Date.now();
|
||
const content = `---
|
||
name: _temp_${timestamp}
|
||
description: Temporary memory - processing in background
|
||
tags: [_temporary]
|
||
links: []
|
||
backlinks: []
|
||
modified: ${new Date().toISOString()}
|
||
---
|
||
|
||
# Recent Conversation (Processing)
|
||
|
||
${conversation}`;
|
||
const [e] = await this.llm.embedding(content);
|
||
return {
|
||
name: `_temp_${timestamp}`,
|
||
description: 'Temporary memory - processing in background',
|
||
content,
|
||
embedding: e?.embedding || [],
|
||
};
|
||
}
|
||
|
||
forget(name: string, memories: Memory[] | MemoryCache): boolean {
|
||
const mem = memories instanceof MemoryCache ? memories.memories : memories;
|
||
const idx = mem.findIndex(m => m.name === name);
|
||
if (idx === -1) return false;
|
||
|
||
for (const node of mem) {
|
||
const {links, backlinks} = extractMetadata(node.content);
|
||
const newBacklinks = backlinks.filter(b => b !== name);
|
||
const newLinks = links.filter(l => l !== name);
|
||
|
||
if (newBacklinks.length !== backlinks.length || newLinks.length !== links.length) {
|
||
node.content = this.updateFrontmatter(node.content, {
|
||
links: newLinks,
|
||
backlinks: newBacklinks,
|
||
});
|
||
}
|
||
}
|
||
|
||
mem.splice(idx, 1);
|
||
|
||
if (memories instanceof MemoryCache) memories.rebuild();
|
||
return true;
|
||
}
|
||
|
||
private cosineSearch(query: number[], memories: Memory[], limit: number): MemoryRef[] {
|
||
const scored = memories
|
||
.filter(m => m.embedding?.length)
|
||
.map(m => ({
|
||
ref: {name: m.name, description: m.description},
|
||
distance: cosineDistance(query, m.embedding),
|
||
}))
|
||
.sort((a, b) => a.distance - b.distance)
|
||
.slice(0, limit);
|
||
return scored.map(s => s.ref);
|
||
}
|
||
|
||
private listNodes(memories: Memory[]): MemoryRef[] {
|
||
return memories.map(m => ({name: m.name, description: m.description}));
|
||
}
|
||
|
||
async recollect(query: string, memories: Memory[] | MemoryCache, limit = 5, graphDepth = 1): Promise<Memory[]> {
|
||
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<string>(vectorResults.map(r => r.name));
|
||
|
||
if (graphDepth > 0) {
|
||
const frontier = [...found];
|
||
for (let depth = 0; depth < graphDepth; depth++) {
|
||
const next: string[] = [];
|
||
for (const name of frontier) {
|
||
const node = mem.find(m => m.name === name);
|
||
if (!node) continue;
|
||
const {links} = extractMetadata(node.content);
|
||
for (const link of 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<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 trackingId = `${Date.now()}_${Math.random()}`;
|
||
const tempMemory = await this.createTempMemory(conversation);
|
||
const mem = memories instanceof MemoryCache ? memories.memories : memories;
|
||
mem.push(tempMemory);
|
||
if (memories instanceof MemoryCache) memories.rebuild();
|
||
this.pendingMemorizations.set(trackingId, {
|
||
memories,
|
||
tempMemoryName: tempMemory.name,
|
||
timestamp: Date.now(),
|
||
});
|
||
|
||
try {
|
||
await this._memorizeBackground(conversation, memories, options, tempMemory.name);
|
||
const finalMem = memories instanceof MemoryCache ? memories.memories : memories;
|
||
return finalMem.filter(m => !m.name.startsWith('_temp_'));
|
||
} finally {
|
||
const pending = this.pendingMemorizations.get(trackingId);
|
||
if (pending) {
|
||
const cleanMem = pending.memories instanceof MemoryCache
|
||
? pending.memories.memories
|
||
: pending.memories;
|
||
const idx = cleanMem.findIndex(m => m.name === pending.tempMemoryName);
|
||
if (idx !== -1) cleanMem.splice(idx, 1);
|
||
if (pending.memories instanceof MemoryCache) pending.memories.rebuild();
|
||
}
|
||
this.pendingMemorizations.delete(trackingId);
|
||
}
|
||
}
|
||
|
||
private async _memorizeBackground(conversation: string, memories: Memory[] | MemoryCache, options: LLMRequest, tempName: string): Promise<void> {
|
||
const mem = memories instanceof MemoryCache ? memories.memories : memories;
|
||
const monday = getWeekMonday();
|
||
const sunday = getWeekSunday(monday);
|
||
const buckets = await this.factAgent(conversation, mem, options, monday);
|
||
if(!buckets.length) return;
|
||
const jobs = [...buckets].map(({subject, facts}) => {
|
||
let node = mem.find(m => m.name === subject);
|
||
if(!node) {
|
||
node = {name: subject, description: '', content: '', embedding: [],};
|
||
mem.push(node);
|
||
}
|
||
const week = subject.startsWith('Journal/') ? {monday, sunday} : undefined;
|
||
return this.enqueue(node, facts, mem, options, tempName, week);
|
||
});
|
||
await Promise.all(jobs);
|
||
}
|
||
|
||
/**
|
||
* Coalescing queue: if a doc is already compiling, abort the in-flight run, merge its
|
||
* facts with the new ones and restart. Never blocks a pending update, never drops facts.
|
||
*/
|
||
private enqueue(node: Memory, facts: string[], memories: Memory[] | MemoryCache, options: LLMRequest, tempName: string, week?: {monday: string, sunday: string}): Promise<void> {
|
||
const key = node.name;
|
||
const existing = this.queues.get(key);
|
||
if (existing) {
|
||
existing.pending.push(...facts);
|
||
existing.request?.abort?.();
|
||
return existing.task;
|
||
}
|
||
|
||
const entry: {pending: string[], request: {abort?: () => void} | null, task: Promise<void>} = {pending: [...facts], request: null, task: Promise.resolve()};
|
||
this.queues.set(key, entry);
|
||
const m = memories instanceof MemoryCache ? memories.memories : memories;
|
||
entry.task = (async () => {
|
||
while (entry.pending.length) {
|
||
const batch = dedupeFacts(entry.pending.splice(0, entry.pending.length));
|
||
const written = await this.docAgent(node, batch, m, options, tempName, week, entry);
|
||
if (!written) entry.pending.unshift(...batch);
|
||
}
|
||
})().finally(() => {
|
||
this.queues.delete(key);
|
||
if(!this.queues.size && memories instanceof MemoryCache) memories.rebuild();
|
||
});
|
||
return entry.task;
|
||
}
|
||
|
||
private buildHeader(node: Memory, week?: {monday: string, sunday: string}, links: string[] = [], backlinks: string[] = []): string {
|
||
const tags = node.name.split('/')[0]?.toLowerCase();
|
||
const lines = [
|
||
'---',
|
||
`name: ${node.name}`,
|
||
`description: ${node.description || ''}`,
|
||
tags ? `tags: [${tags}]` : '',
|
||
links.length ? `links: [${links.map(l => `"${l}"`).join(', ')}]` : 'links: []',
|
||
backlinks.length ? `backlinks: [${backlinks.map(l => `"${l}"`).join(', ')}]` : 'backlinks: []',
|
||
week ? `week: ${week.monday} – ${week.sunday}` : '',
|
||
`modified: ${new Date().toISOString()}`,
|
||
'---',
|
||
].filter(Boolean);
|
||
return lines.join('\n');
|
||
}
|
||
|
||
private applyHeader(content: string, header: string): string {
|
||
return `${header}\n\n${this.stripHeader(content)}`;
|
||
}
|
||
|
||
private updateFrontmatter(content: string, updates: {links?: string[], backlinks?: string[]}): string {
|
||
const match = content.match(/^---\n([\s\S]*?)\n---\n\n?([\s\S]*)$/);
|
||
if (!match) return content;
|
||
|
||
const [, fm, body] = match;
|
||
let newFm = fm;
|
||
|
||
if (updates.links !== undefined) {
|
||
const linksList = updates.links.length ? `[${updates.links.map(l => `"${l}"`).join(', ')}]` : '[]';
|
||
newFm = newFm.replace(/^links:.*$/m, `links: ${linksList}`);
|
||
}
|
||
|
||
if (updates.backlinks !== undefined) {
|
||
const backlinksList = updates.backlinks.length ? `[${updates.backlinks.map(l => `"${l}"`).join(', ')}]` : '[]';
|
||
newFm = newFm.replace(/^backlinks:.*$/m, `backlinks: ${backlinksList}`);
|
||
}
|
||
|
||
newFm = newFm.replace(/^modified:.*$/m, `modified: ${new Date().toISOString()}`);
|
||
|
||
return `---\n${newFm}\n---\n\n${body}`;
|
||
}
|
||
|
||
private stripHeader(content: string): string {
|
||
return content.replace(/^---[\s\S]*?\n---\n?/, '').trimStart();
|
||
}
|
||
|
||
private async docAgent(node: Memory, facts: string[], memories: Memory[], options: LLMRequest, tempName: string, week: {monday: string, sunday: string} | undefined, entry: {request: {abort?: () => void} | null}): Promise<boolean> {
|
||
const {links: oldLinks} = extractMetadata(node.content);
|
||
const currentBody = this.stripHeader(node.content);
|
||
let update;
|
||
try {
|
||
for(let i = 0; i < 3 && !update?.content; i++) {
|
||
const request = this.llm.ask(`New Facts:\n${facts.map(f => `- ${f}`).join('\n')}`, {
|
||
model: options.model,
|
||
temperature: 0.3,
|
||
schema: {
|
||
description: {type: 'string', description: 'One-line description of what this document covers, no formatting or emojis', required: true},
|
||
content: {type: 'string', description: 'Rewritten document in markdown, without the frontmatter block', required: true},
|
||
},
|
||
system: `You are a knowledge base editor. Rewrite the current document below so it incorporates the new facts.
|
||
|
||
Formatting rules:
|
||
- Use Obsidian-style markdown: # headings, **bold** to add emphasis, __italics__ for titles, terms, etc, bullet & numbered lists for grouped 1D data and tables for 2D data
|
||
- Link related concepts with [[WikiLink]] notation using full paths like [[People/Sarah]] or [[Projects/Website]]
|
||
- Create links for specific entities (person, place, project, program) and abstract concepts (quantum mechanics, entropy) but skip generics (car, red, dog)
|
||
- Keep the document concise, factual, and human-readable
|
||
- Resolve contradictions: the new facts always win — delete the outdated statement entirely, never keep both
|
||
- Later facts in the list override earlier ones
|
||
- Do not add frontmatter blocks, filler, preamble, or AI commentary
|
||
${week ? '- This is a weekly journal entry.\n' : ''}
|
||
All nodes:
|
||
${this.listNodes(memories).map(n => n.name).join(', ') || 'none'}
|
||
|
||
Current document:
|
||
\`\`\`markdown
|
||
${currentBody}
|
||
\`\`\``}
|
||
);
|
||
entry.request = request;
|
||
update = await request;
|
||
}
|
||
} catch (err: any) {
|
||
if (err?.name === 'AbortError') return false;
|
||
throw err;
|
||
} finally {
|
||
entry.request = null;
|
||
}
|
||
|
||
if(!update?.content) return false;
|
||
const newLinks = extractLinks(update.content).filter(l => l !== node.name && l !== tempName);
|
||
const newLinkSet = new Set(newLinks);
|
||
const oldLinkSet = new Set(oldLinks);
|
||
|
||
for (const added of newLinkSet) {
|
||
if (!oldLinkSet.has(added)) {
|
||
const target = memories.find(m => m.name === added);
|
||
if (target) {
|
||
const {backlinks} = extractMetadata(target.content);
|
||
if (!backlinks.includes(node.name)) {
|
||
target.content = this.updateFrontmatter(target.content, {
|
||
backlinks: [...backlinks, node.name],
|
||
});
|
||
}
|
||
}
|
||
}
|
||
}
|
||
for (const removed of oldLinkSet) {
|
||
if (!newLinkSet.has(removed)) {
|
||
const target = memories.find(m => m.name === removed);
|
||
if (target) {
|
||
const {backlinks} = extractMetadata(target.content);
|
||
target.content = this.updateFrontmatter(target.content, {
|
||
backlinks: backlinks.filter(b => b !== node.name),
|
||
});
|
||
}
|
||
}
|
||
}
|
||
|
||
const {backlinks} = extractMetadata(node.content);
|
||
node.description = node.name !== 'Person/User' ? update.description : 'All information about the current user';
|
||
node.content = this.applyHeader(update.content, this.buildHeader(node, week, newLinks, backlinks));
|
||
const [e] = await this.llm.embedding(node.content);
|
||
if(e) node.embedding = e.embedding;
|
||
return true;
|
||
}
|
||
|
||
private async factAgent(conversation: string, memories: Memory[], options: LLMRequest, weekKey: string): Promise<FactBucket[]> {
|
||
const buckets = new Map<string, string[]>();
|
||
await this.llm.ask(conversation, {
|
||
model: options.model,
|
||
temperature: 0.2,
|
||
system: `You are a fact extractor. Analyze this conversation and extract facts worth remembering long-term.
|
||
|
||
Rules:
|
||
- ONLY extract current facts the USER explicitly stated about themselves, their work, or their projects
|
||
- ONLY extract decisions that were MADE during this conversation
|
||
- DO NOT extract anything the AI said, its capabilities, or meta-conversation about the AI
|
||
- DO NOT extract greetings, pleasantries, or generic exchanges
|
||
- DO NOT extract deltas or changes in facts; ONLY the end fact
|
||
- If nothing worth remembering was said, do not call any tools
|
||
|
||
When extracting facts, you MUST also decide the exact destination path:
|
||
- Use an existing node name if the facts clearly belong there
|
||
- All information primary about the user should go under "People/User"
|
||
- When required, create a new path following collection/subject format (e.g., People/Sarah, Projects/Oxide)
|
||
- For journal entries, use "Journal"
|
||
|
||
Available nodes:
|
||
- Journal
|
||
${this.listNodes(memories).filter(n => !n.name.includes('_temp_') && !n.name.includes('Journal')).map(n => `- ${n.name}: ${n.description}`).join('\n') || 'None yet.'}`,
|
||
tools: [{
|
||
name: 'facts_extract',
|
||
description: 'Submit facts with their destination',
|
||
args: {
|
||
destination: {type: 'string', description: 'Exact existing node name OR new path (e.g. "People/Sarah", "Projects/Oxide")', required: true},
|
||
facts: {type: 'string', description: 'Comma-separated facts', required: true},
|
||
},
|
||
fn: (args: any) => {
|
||
const subject = args.destination.trim().toLowerCase() === 'journal'
|
||
? `Journal/${weekKey}` : args.destination.trim();
|
||
const facts = buckets.get(subject) ?? [];
|
||
facts.push(...dedupeFacts(String(args.facts).split(',')));
|
||
buckets.set(subject, facts);
|
||
return 'Recorded';
|
||
},
|
||
}],
|
||
});
|
||
return buckets.entries().toArray().map(([subject, facts]) => ({subject, facts}));
|
||
}
|
||
}
|