Memory Custom

#!/usr/bin/env node import { Server } from "@modelcontextprotocol/sdk/server/index.js"; import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; import { CallToolRequestSchema, ListToolsRequestSchema, } from "@modelcontextprotocol/sdk/types.js"; import { promises as fs } from "fs"; import path from "path"; import { fileURLToPath } from "url"; // Define memory file path using environment variable with fallback const defaultMemoryPath = path.join( path.dirname(fileURLToPath(import.meta.url)), "memory.json" ); // If MEMORY_FILE_PATH is just a filename, put it in the same directory as the script const MEMORY_FILE_PATH = process.env.MEMORY_FILE_PATH ? path.isAbsolute(process.env.MEMORY_FILE_PATH) ? process.env.MEMORY_FILE_PATH : path.join( path.dirname(fileURLToPath(import.meta.url)), process.env.MEMORY_FILE_PATH ) : defaultMemoryPath; // We are storing our memory using entities, relations, and observations in a graph structure interface Entity { name: string; entityType: string; observations: string[]; } interface Relation { from: string; to: string; relationType: string; } interface KnowledgeGraph { entities: Entity[]; relations: Relation[]; } // The KnowledgeGraphManager class contains all operations to interact with the knowledge graph class KnowledgeGraphManager { private memoryFilePath: string; constructor(memoryFilePath: string) { this.memoryFilePath = memoryFilePath; } async getCurrentTime() { return new Date().toISOString(); } async setMemoryFilePath(memoryFilePath: string) { // check if path is valid if (!path.isAbsolute(memoryFilePath)) { throw new Error("Memory file path must be an absolute path"); } memoryFilePath = path.normalize(memoryFilePath); try { await fs.stat(memoryFilePath); } catch (error) { // If the error is because the file does not exist, create it if ((error as any).code === "ENOENT") { await fs.writeFile(memoryFilePath, ""); } else { // Handle other potential errors throw error; } } this.memoryFilePath = memoryFilePath; } private async loadGraph(): Promise<KnowledgeGraph> { try { const data = await fs.readFile( this.memoryFilePath ?? MEMORY_FILE_PATH, "utf-8" ); const lines = data.split("\n").filter((line) => line.trim() !== ""); return lines.reduce( (graph: KnowledgeGraph, line) => { const item = JSON.parse(line); if (item.type === "entity") graph.entities.push(item as Entity); if (item.type === "relation") graph.relations.push(item as Relation); return graph; }, { entities: [], relations: [] } ); } catch (error) { if ( error instanceof Error && "code" in error && (error as any).code === "ENOENT" ) { return { entities: [], relations: [] }; } throw error; } } private async saveGraph(graph: KnowledgeGraph): Promise<void> { const lines = [ ...graph.entities.map((e) => JSON.stringify({ type: "entity", ...e })), ...graph.relations.map((r) => JSON.stringify({ type: "relation", ...r })), ]; await fs.writeFile( this.memoryFilePath ?? MEMORY_FILE_PATH, lines.join("\n") ); } async createEntities( entities: Entity[], filepath: string ): Promise<Entity[]> { await this.setMemoryFilePath(filepath); const graph = await this.loadGraph(); const newEntities = entities.filter( (e) => !graph.entities.some((existingEntity) => existingEntity.name === e.name) ); graph.entities.push(...newEntities); await this.saveGraph(graph); return newEntities; } async createRelations( relations: Relation[], filepath: string ): Promise<Relation[]> { await this.setMemoryFilePath(filepath); const graph = await this.loadGraph(); const newRelations = relations.filter( (r) => !graph.relations.some( (existingRelation) => existingRelation.from === r.from && existingRelation.to === r.to && existingRelation.relationType === r.relationType ) ); graph.relations.push(...newRelations); await this.saveGraph(graph); return newRelations; } async addObservations( observations: { entityName: string; contents: string[] }[], filepath: string ): Promise<{ entityName: string; addedObservations: string[] }[]> { await this.setMemoryFilePath(filepath); const graph = await this.loadGraph(); const results = observations.map((o) => { const entity = graph.entities.find((e) => e.name === o.entityName); if (!entity) { throw new Error(`Entity with name ${o.entityName} not found`); } const newObservations = o.contents.filter( (content) => !entity.observations.includes(content) ); entity.observations.push(...newObservations); return { entityName: o.entityName, addedObservations: newObservations }; }); await this.saveGraph(graph); return results; } async deleteEntities(entityNames: string[], filepath: string): Promise<void> { await this.setMemoryFilePath(filepath); const graph = await this.loadGraph(); graph.entities = graph.entities.filter( (e) => !entityNames.includes(e.name) ); graph.relations = graph.relations.filter( (r) => !entityNames.includes(r.from) && !entityNames.includes(r.to) ); await this.saveGraph(graph); } async deleteObservations( deletions: { entityName: string; observations: string[] }[], filepath: string ): Promise<void> { await this.setMemoryFilePath(filepath); const graph = await this.loadGraph(); deletions.forEach((d) => { const entity = graph.entities.find((e) => e.name === d.entityName); if (entity) { entity.observations = entity.observations.filter( (o) => !d.observations.includes(o) ); } }); await this.saveGraph(graph); } async deleteRelations( relations: Relation[], filepath: string ): Promise<void> { await this.setMemoryFilePath(filepath); const graph = await this.loadGraph(); graph.relations = graph.relations.filter( (r) => !relations.some( (delRelation) => r.from === delRelation.from && r.to === delRelation.to && r.relationType === delRelation.relationType ) ); await this.saveGraph(graph); } async readGraph(filepath: string): Promise<KnowledgeGraph> { await this.setMemoryFilePath(filepath); return this.loadGraph(); } // Very basic search function async searchNodes(query: string, filepath: string): Promise<KnowledgeGraph> { await this.setMemoryFilePath(filepath); const graph = await this.loadGraph(); // Filter entities const filteredEntities = graph.entities.filter( (e) => e.name.toLowerCase().includes(query.toLowerCase()) || e.entityType.toLowerCase().includes(query.toLowerCase()) || e.observations.some((o) => o.toLowerCase().includes(query.toLowerCase()) ) ); // Create a Set of filtered entity names for quick lookup const filteredEntityNames = new Set(filteredEntities.map((e) => e.name)); // Filter relations to only include those between filtered entities const filteredRelations = graph.relations.filter( (r) => filteredEntityNames.has(r.from) && filteredEntityNames.has(r.to) ); const filteredGraph: KnowledgeGraph = { entities: filteredEntities, relations: filteredRelations, }; return filteredGraph; } async openNodes(names: string[], filepath: string): Promise<KnowledgeGraph> { await this.setMemoryFilePath(filepath); const graph = await this.loadGraph(); // Filter entities const filteredEntities = graph.entities.filter((e) => names.includes(e.name) ); // Create a Set of filtered entity names for quick lookup const filteredEntityNames = new Set(filteredEntities.map((e) => e.name)); // Filter relations to only include those between filtered entities const filteredRelations = graph.relations.filter( (r) => filteredEntityNames.has(r.from) && filteredEntityNames.has(r.to) ); const filteredGraph: KnowledgeGraph = { entities: filteredEntities, relations: filteredRelations, }; return filteredGraph; } } const knowledgeGraphManager = new KnowledgeGraphManager(MEMORY_FILE_PATH); // The server instance and tools exposed to Claude const server = new Server( { name: "memory-server", version: "1.0.0", }, { capabilities: { tools: {}, }, } ); server.setRequestHandler(ListToolsRequestSchema, async () => { return { tools: [ { name: "get_current_time", description: "Get the current time", inputSchema: { type: "object", properties: {}, }, }, // { // name: "set_memory_file_path", // description: "Set the memory file path", // inputSchema: { // type: "object", // properties: { // memoryFilePath: { // type: "string", // description: "Absolute path to the memory file", // }, // }, // required: ["memoryFilePath"], // }, // }, { name: "create_entities", description: "Create multiple new entities in the knowledge graph", inputSchema: { type: "object", properties: { entities: { type: "array", items: { type: "object", properties: { name: { type: "string", description: "The name of the entity", }, entityType: { type: "string", description: "The type of the entity", }, observations: { type: "array", items: { type: "string" }, description: "An array of observation contents associated with the entity", }, }, required: ["name", "entityType", "observations"], }, }, memoryFilePath: { type: "string", description: "The path to the memory file", }, }, required: ["entities", "memoryFilePath"], }, }, { name: "create_relations", description: "Create multiple new relations between entities in the knowledge graph. Relations should be in active voice", inputSchema: { type: "object", properties: { relations: { type: "array", items: { type: "object", properties: { from: { type: "string", description: "The name of the entity where the relation starts", }, to: { type: "string", description: "The name of the entity where the relation ends", }, relationType: { type: "string", description: "The type of the relation", }, }, required: ["from", "to", "relationType"], }, }, memoryFilePath: { type: "string", description: "The path to the memory file", }, }, required: ["relations", "memoryFilePath"], }, }, { name: "add_observations", description: "Add new observations to existing entities in the knowledge graph", inputSchema: { type: "object", properties: { observations: { type: "array", items: { type: "object", properties: { entityName: { type: "string", description: "The name of the entity to add the observations to", }, contents: { type: "array", items: { type: "string" }, description: "An array of observation contents to add", }, }, required: ["entityName", "contents"], }, }, memoryFilePath: { type: "string", description: "The path to the memory file", }, }, required: ["observations", "memoryFilePath"], }, }, { name: "delete_entities", description: "Delete multiple entities and their associated relations from the knowledge graph", inputSchema: { type: "object", properties: { entityNames: { type: "array", items: { type: "string" }, description: "An array of entity names to delete", }, memoryFilePath: { type: "string", description: "The path to the memory file", }, }, required: ["entityNames", "memoryFilePath"], }, }, { name: "delete_observations", description: "Delete specific observations from entities in the knowledge graph", inputSchema: { type: "object", properties: { deletions: { type: "array", items: { type: "object", properties: { entityName: { type: "string", description: "The name of the entity containing the observations", }, observations: { type: "array", items: { type: "string" }, description: "An array of observations to delete", }, }, required: ["entityName", "observations"], }, }, memoryFilePath: { type: "string", description: "The path to the memory file", }, }, required: ["deletions", "memoryFilePath"], }, }, { name: "delete_relations", description: "Delete multiple relations from the knowledge graph", inputSchema: { type: "object", properties: { relations: { type: "array", items: { type: "object", properties: { from: { type: "string", description: "The name of the entity where the relation starts", }, to: { type: "string", description: "The name of the entity where the relation ends", }, relationType: { type: "string", description: "The type of the relation", }, }, required: ["from", "to", "relationType"], }, description: "An array of relations to delete", }, memoryFilePath: { type: "string", description: "The path to the memory file", }, }, required: ["relations", "memoryFilePath"], }, }, { name: "read_graph", description: "Read the entire knowledge graph", inputSchema: { type: "object", properties: { memoryFilePath: { type: "string", description: "The path to the memory file", }, }, required: ["memoryFilePath"], }, }, { name: "search_nodes", description: "Search for nodes in the knowledge graph based on a query", inputSchema: { type: "object", properties: { query: { type: "string", description: "The search query to match against entity names, types, and observation content", }, memoryFilePath: { type: "string", description: "The path to the memory file", }, }, required: ["query", "memoryFilePath"], }, }, { name: "open_nodes", description: "Open specific nodes in the knowledge graph by their names", inputSchema: { type: "object", properties: { names: { type: "array", items: { type: "string" }, description: "An array of entity names to retrieve", }, memoryFilePath: { type: "string", description: "The path to the memory file", }, }, required: ["names", "memoryFilePath"], }, }, ], }; }); server.setRequestHandler(CallToolRequestSchema, async (request) => { const { name, arguments: args } = request.params; if (!args) { throw new Error(`No arguments provided for tool: ${name}`); } switch (name) { case "get_current_time": return { content: [ { type: "text", text: await knowledgeGraphManager.getCurrentTime() }, ], }; // case "set_memory_file_path": // knowledgeGraphManager.setMemoryFilePath(args.memoryFilePath as string); // return { // content: [{ type: "text", text: "Memory file path set successfully" }], // }; case "create_entities": return { content: [ { type: "text", text: JSON.stringify( await knowledgeGraphManager.createEntities( args.entities as Entity[], args.memoryFilePath as string ), null, 2 ), }, ], }; case "create_relations": return { content: [ { type: "text", text: JSON.stringify( await knowledgeGraphManager.createRelations( args.relations as Relation[], args.memoryFilePath as string ), null, 2 ), }, ], }; case "add_observations": return { content: [ { type: "text", text: JSON.stringify( await knowledgeGraphManager.addObservations( args.observations as { entityName: string; contents: string[]; }[], args.memoryFilePath as string ), null, 2 ), }, ], }; case "delete_entities": await knowledgeGraphManager.deleteEntities( args.entityNames as string[], args.memoryFilePath as string ); return { content: [{ type: "text", text: "Entities deleted successfully" }], }; case "delete_observations": await knowledgeGraphManager.deleteObservations( args.deletions as { entityName: string; observations: string[] }[], args.memoryFilePath as string ); return { content: [{ type: "text", text: "Observations deleted successfully" }], }; case "delete_relations": await knowledgeGraphManager.deleteRelations( args.relations as Relation[], args.memoryFilePath as string ); return { content: [{ type: "text", text: "Relations deleted successfully" }], }; case "read_graph": return { content: [ { type: "text", text: JSON.stringify( await knowledgeGraphManager.readGraph( args.memoryFilePath as string ), null, 2 ), }, ], }; case "search_nodes": return { content: [ { type: "text", text: JSON.stringify( await knowledgeGraphManager.searchNodes( args.query as string, args.memoryFilePath as string ), null, 2 ), }, ], }; case "open_nodes": return { content: [ { type: "text", text: JSON.stringify( await knowledgeGraphManager.openNodes( args.names as string[], args.memoryFilePath as string ), null, 2 ), }, ], }; default: throw new Error(`Unknown tool: ${name}`); } }); async function main() { const transport = new StdioServerTransport(); await server.connect(transport); console.error("Knowledge Graph MCP Server running on stdio"); } main().catch((error) => { console.error("Fatal error in main():", error); process.exit(1); });
ID: w6hi2myrxq