Add 7 living memory tools: split, merge, dedupe, prune, reorganize, bulk_tag, stats
This commit is contained in:
256
src/mcp/index.ts
256
src/mcp/index.ts
@@ -1,8 +1,10 @@
|
||||
import { McpServer } from '@modelcontextprotocol/sdk/server/mcp.js';
|
||||
import { StdioServerTransport } from '@modelcontextprotocol/sdk/server/stdio.js';
|
||||
import { z } from 'zod/v3';
|
||||
import { query, listNodes, getNode, findNodeByPrefix, addNode, addEdge } from '../core/store';
|
||||
import { getConnections } from '../core/graph';
|
||||
import { query, listNodes, getNode, findNodeByPrefix, addNode, addEdge, removeNode, removeEdge, updateNode } from '../core/store';
|
||||
import { getConnections, getEdgesByNode } from '../core/graph';
|
||||
import { cosineSimilarity } from '../core/search/vector';
|
||||
import { getDb } from '../core/db';
|
||||
import { NodeKind, EdgeType } from '../types';
|
||||
|
||||
function serialize(data: any): string {
|
||||
@@ -113,6 +115,256 @@ server.tool(
|
||||
}
|
||||
);
|
||||
|
||||
// --- memory_split ---
|
||||
server.tool(
|
||||
'memory_split',
|
||||
'Break a large node into smaller children, updating parent with a summary',
|
||||
{
|
||||
id: z.string().describe('Node ID to split'),
|
||||
pieces: z.array(z.object({ title: z.string(), content: z.string() })).describe('Child pieces to create'),
|
||||
summary: z.string().optional().describe('Optional summary to replace parent content'),
|
||||
},
|
||||
async ({ id, pieces, summary }) => {
|
||||
const node = getNode(id) ?? findNodeByPrefix(id);
|
||||
if (!node) return { content: [{ type: 'text' as const, text: serialize({ error: 'Node not found' }) }], isError: true };
|
||||
|
||||
if (summary !== undefined) {
|
||||
await updateNode(node.id, { content: summary });
|
||||
}
|
||||
|
||||
const children = [];
|
||||
for (const piece of pieces) {
|
||||
const child = await addNode({ kind: node.kind, title: piece.title, content: piece.content, tags: [...node.tags] });
|
||||
addEdge(node.id, child.id, 'contains');
|
||||
children.push(child);
|
||||
}
|
||||
return { content: [{ type: 'text' as const, text: serialize({ parent: node.id, children: children.map(c => ({ id: c.id, title: c.title })) }) }] };
|
||||
}
|
||||
);
|
||||
|
||||
// --- memory_merge ---
|
||||
server.tool(
|
||||
'memory_merge',
|
||||
'Merge multiple nodes into a single new node, relinking edges and deleting originals',
|
||||
{
|
||||
nodeIds: z.array(z.string()).describe('Node IDs to merge'),
|
||||
title: z.string().describe('Title for merged node'),
|
||||
content: z.string().describe('Content for merged node'),
|
||||
kind: z.enum(['memory', 'component', 'task', 'decision']).optional().describe('Kind (defaults to first node\'s kind)'),
|
||||
},
|
||||
async ({ nodeIds, title, content, kind }) => {
|
||||
const nodes = nodeIds.map(id => getNode(id) ?? findNodeByPrefix(id));
|
||||
if (nodes.some(n => !n)) return { content: [{ type: 'text' as const, text: serialize({ error: 'One or more nodes not found' }) }], isError: true };
|
||||
|
||||
const validNodes = nodes as Exclude<typeof nodes[number], null>[];
|
||||
const allTags = [...new Set(validNodes.flatMap(n => n.tags))];
|
||||
const mergedKind = (kind as NodeKind) ?? validNodes[0].kind;
|
||||
|
||||
const merged = await addNode({ kind: mergedKind, title, content, tags: allTags });
|
||||
const oldIds = new Set(validNodes.map(n => n.id));
|
||||
|
||||
// Relink edges from old nodes to merged node
|
||||
const db = getDb();
|
||||
for (const n of validNodes) {
|
||||
const edges = getEdgesByNode(n.id);
|
||||
for (const e of edges) {
|
||||
const from = oldIds.has(e.fromId) ? merged.id : e.fromId;
|
||||
const to = oldIds.has(e.toId) ? merged.id : e.toId;
|
||||
if (from !== to) {
|
||||
try { addEdge(from, to, e.type, e.metadata); } catch {}
|
||||
}
|
||||
}
|
||||
removeNode(n.id, true);
|
||||
}
|
||||
|
||||
return { content: [{ type: 'text' as const, text: serialize({ merged: { id: merged.id, title: merged.title, tags: merged.tags } }) }] };
|
||||
}
|
||||
);
|
||||
|
||||
// --- memory_dedupe ---
|
||||
server.tool(
|
||||
'memory_dedupe',
|
||||
'Find similar/duplicate nodes using embedding cosine similarity',
|
||||
{
|
||||
threshold: z.number().optional().describe('Similarity threshold (default 0.85)'),
|
||||
kind: z.enum(['memory', 'component', 'task', 'decision']).optional().describe('Filter by kind'),
|
||||
limit: z.number().optional().describe('Max groups to return (default 10)'),
|
||||
},
|
||||
async ({ threshold, kind, limit }) => {
|
||||
const thresh = threshold ?? 0.85;
|
||||
const maxGroups = limit ?? 10;
|
||||
const nodes = listNodes({ kind: kind as NodeKind, includeStale: false });
|
||||
const withEmb = nodes.filter(n => n.embedding);
|
||||
|
||||
const groups: { nodes: { id: string; title: string }[]; similarity: number }[] = [];
|
||||
const seen = new Set<string>();
|
||||
|
||||
for (let i = 0; i < withEmb.length && groups.length < maxGroups; i++) {
|
||||
if (seen.has(withEmb[i].id)) continue;
|
||||
const group = [{ id: withEmb[i].id, title: withEmb[i].title }];
|
||||
let maxSim = 0;
|
||||
for (let j = i + 1; j < withEmb.length; j++) {
|
||||
if (seen.has(withEmb[j].id)) continue;
|
||||
const sim = cosineSimilarity(withEmb[i].embedding!, withEmb[j].embedding!);
|
||||
if (sim >= thresh) {
|
||||
group.push({ id: withEmb[j].id, title: withEmb[j].title });
|
||||
seen.add(withEmb[j].id);
|
||||
maxSim = Math.max(maxSim, sim);
|
||||
}
|
||||
}
|
||||
if (group.length > 1) {
|
||||
seen.add(withEmb[i].id);
|
||||
groups.push({ nodes: group, similarity: maxSim });
|
||||
}
|
||||
}
|
||||
|
||||
return { content: [{ type: 'text' as const, text: serialize({ groups, count: groups.length }) }] };
|
||||
}
|
||||
);
|
||||
|
||||
// --- memory_prune ---
|
||||
server.tool(
|
||||
'memory_prune',
|
||||
'Clean up the graph: delete stale nodes, decay old nodes, or remove orphans',
|
||||
{
|
||||
mode: z.enum(['hard_delete_stale', 'decay_then_delete', 'delete_orphans']).describe('Prune mode'),
|
||||
maxAgeDays: z.number().optional().describe('Max age in days for decay/delete modes'),
|
||||
},
|
||||
async ({ mode, maxAgeDays }) => {
|
||||
const db = getDb();
|
||||
let affected = 0;
|
||||
|
||||
if (mode === 'hard_delete_stale') {
|
||||
const result = db.prepare('DELETE FROM nodes WHERE is_stale = 1').run();
|
||||
affected = result.changes;
|
||||
} else if (mode === 'decay_then_delete') {
|
||||
const cutoff = Date.now() - (maxAgeDays ?? 90) * 86400000;
|
||||
// First mark old non-stale as stale
|
||||
const decayed = db.prepare('UPDATE nodes SET is_stale = 1, updated_at = ? WHERE is_stale = 0 AND last_accessed_at < ?').run(Date.now(), cutoff);
|
||||
// Then hard-delete already-stale that are also old
|
||||
const deleted = db.prepare('DELETE FROM nodes WHERE is_stale = 1 AND updated_at < ?').run(cutoff);
|
||||
affected = decayed.changes + deleted.changes;
|
||||
} else if (mode === 'delete_orphans') {
|
||||
const orphans = db.prepare(`
|
||||
SELECT n.id FROM nodes n
|
||||
WHERE n.is_stale = 0
|
||||
AND NOT EXISTS (SELECT 1 FROM edges e WHERE e.from_id = n.id OR e.to_id = n.id)
|
||||
`).all() as any[];
|
||||
for (const o of orphans) {
|
||||
removeNode(o.id, true);
|
||||
}
|
||||
affected = orphans.length;
|
||||
}
|
||||
|
||||
return { content: [{ type: 'text' as const, text: serialize({ mode, affected }) }] };
|
||||
}
|
||||
);
|
||||
|
||||
// --- memory_reorganize ---
|
||||
server.tool(
|
||||
'memory_reorganize',
|
||||
'Move a node under a new parent (updates contains edges)',
|
||||
{
|
||||
nodeId: z.string().describe('Node ID to move'),
|
||||
newParentId: z.string().describe('New parent node ID'),
|
||||
},
|
||||
async ({ nodeId, newParentId }) => {
|
||||
const node = getNode(nodeId) ?? findNodeByPrefix(nodeId);
|
||||
const parent = getNode(newParentId) ?? findNodeByPrefix(newParentId);
|
||||
if (!node) return { content: [{ type: 'text' as const, text: serialize({ error: 'Node not found' }) }], isError: true };
|
||||
if (!parent) return { content: [{ type: 'text' as const, text: serialize({ error: 'Parent not found' }) }], isError: true };
|
||||
|
||||
// Remove existing incoming contains edges
|
||||
const db = getDb();
|
||||
const incomingContains = db.prepare('SELECT id FROM edges WHERE to_id = ? AND type = ?').all(node.id, 'contains') as any[];
|
||||
for (const e of incomingContains) {
|
||||
removeEdge(e.id);
|
||||
}
|
||||
|
||||
const edge = addEdge(parent.id, node.id, 'contains');
|
||||
return { content: [{ type: 'text' as const, text: serialize({ moved: node.id, newParent: parent.id, edge: edge.id }) }] };
|
||||
}
|
||||
);
|
||||
|
||||
// --- memory_bulk_tag ---
|
||||
server.tool(
|
||||
'memory_bulk_tag',
|
||||
'Add or remove tags on multiple nodes at once',
|
||||
{
|
||||
action: z.enum(['add', 'remove']).describe('Whether to add or remove tags'),
|
||||
tags: z.array(z.string()).describe('Tags to add/remove'),
|
||||
nodeIds: z.array(z.string()).optional().describe('Specific node IDs'),
|
||||
filter: z.object({
|
||||
kind: z.enum(['memory', 'component', 'task', 'decision']).optional(),
|
||||
status: z.string().optional(),
|
||||
tags: z.array(z.string()).optional(),
|
||||
}).optional().describe('Filter to select nodes'),
|
||||
},
|
||||
async ({ action, tags, nodeIds, filter }) => {
|
||||
let targets: { id: string; tags: string[] }[];
|
||||
|
||||
if (nodeIds?.length) {
|
||||
targets = nodeIds.map(id => {
|
||||
const n = getNode(id) ?? findNodeByPrefix(id);
|
||||
return n ? { id: n.id, tags: n.tags } : null;
|
||||
}).filter(Boolean) as { id: string; tags: string[] }[];
|
||||
} else if (filter) {
|
||||
targets = listNodes({ kind: filter.kind as NodeKind, status: filter.status, tags: filter.tags }).map(n => ({ id: n.id, tags: n.tags }));
|
||||
} else {
|
||||
return { content: [{ type: 'text' as const, text: serialize({ error: 'Provide nodeIds or filter' }) }], isError: true };
|
||||
}
|
||||
|
||||
let modified = 0;
|
||||
for (const t of targets) {
|
||||
let newTags: string[];
|
||||
if (action === 'add') {
|
||||
newTags = [...new Set([...t.tags, ...tags])];
|
||||
} else {
|
||||
newTags = t.tags.filter(tag => !tags.includes(tag));
|
||||
}
|
||||
if (JSON.stringify(newTags) !== JSON.stringify(t.tags)) {
|
||||
await updateNode(t.id, { tags: newTags });
|
||||
modified++;
|
||||
}
|
||||
}
|
||||
|
||||
return { content: [{ type: 'text' as const, text: serialize({ action, tags, modified, total: targets.length }) }] };
|
||||
}
|
||||
);
|
||||
|
||||
// --- memory_stats ---
|
||||
server.tool(
|
||||
'memory_stats',
|
||||
'Get graph statistics: counts by kind, stale/orphan counts, tag distribution',
|
||||
{},
|
||||
async () => {
|
||||
const db = getDb();
|
||||
|
||||
const byKind = db.prepare('SELECT kind, COUNT(*) as count FROM nodes WHERE is_stale = 0 GROUP BY kind').all();
|
||||
const totalNodes = db.prepare('SELECT COUNT(*) as count FROM nodes WHERE is_stale = 0').get() as any;
|
||||
const staleCount = db.prepare('SELECT COUNT(*) as count FROM nodes WHERE is_stale = 1').get() as any;
|
||||
const edgeCount = db.prepare('SELECT COUNT(*) as count FROM edges').get() as any;
|
||||
const edgesByType = db.prepare('SELECT type, COUNT(*) as count FROM edges GROUP BY type').all();
|
||||
const orphanCount = db.prepare(`
|
||||
SELECT COUNT(*) as count FROM nodes n
|
||||
WHERE n.is_stale = 0
|
||||
AND NOT EXISTS (SELECT 1 FROM edges e WHERE e.from_id = n.id OR e.to_id = n.id)
|
||||
`).get() as any;
|
||||
const tagDist = db.prepare('SELECT tag, COUNT(*) as count FROM node_tags GROUP BY tag ORDER BY count DESC LIMIT 20').all();
|
||||
|
||||
return {
|
||||
content: [{
|
||||
type: 'text' as const,
|
||||
text: serialize({
|
||||
nodes: { total: totalNodes.count, stale: staleCount.count, orphans: orphanCount.count, byKind },
|
||||
edges: { total: edgeCount.count, byType: edgesByType },
|
||||
topTags: tagDist,
|
||||
}),
|
||||
}],
|
||||
};
|
||||
}
|
||||
);
|
||||
|
||||
async function main() {
|
||||
const transport = new StdioServerTransport();
|
||||
await server.connect(transport);
|
||||
|
||||
Reference in New Issue
Block a user