server.ts•2.86 kB
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import {
CallToolRequestSchema,
ListToolsRequestSchema,
} from "@modelcontextprotocol/sdk/types.js";
import { KnowledgeGraphManager } from './knowledge-graph.js';
import { MEMORY_TOOLS } from './tools.js';
import { Entity, Relation, ObservationInput, DeletionInput } from './types.js';
export function createMemoryServer(): { server: Server; cleanup: () => Promise<void> } {
const knowledgeGraphManager = new KnowledgeGraphManager();
const server = new Server({
name: "memory-server",
version: "0.6.3",
}, {
capabilities: {
tools: {},
},
});
server.setRequestHandler(ListToolsRequestSchema, async () => {
return {
tools: MEMORY_TOOLS,
};
});
server.setRequestHandler(CallToolRequestSchema, async (request) => {
const { name, arguments: args } = request.params;
if (name === "read_graph") {
return { content: [{ type: "text", text: JSON.stringify(await knowledgeGraphManager.readGraph(), null, 2) }] };
}
if (!args) {
throw new Error(`No arguments provided for tool: ${name}`);
}
switch (name) {
case "create_entities":
return { content: [{ type: "text", text: JSON.stringify(await knowledgeGraphManager.createEntities(args.entities as Entity[]), null, 2) }] };
case "create_relations":
return { content: [{ type: "text", text: JSON.stringify(await knowledgeGraphManager.createRelations(args.relations as Relation[]), null, 2) }] };
case "add_observations":
return { content: [{ type: "text", text: JSON.stringify(await knowledgeGraphManager.addObservations(args.observations as ObservationInput[]), null, 2) }] };
case "delete_entities":
await knowledgeGraphManager.deleteEntities(args.entityNames as string[]);
return { content: [{ type: "text", text: "Entities deleted successfully" }] };
case "delete_observations":
await knowledgeGraphManager.deleteObservations(args.deletions as DeletionInput[]);
return { content: [{ type: "text", text: "Observations deleted successfully" }] };
case "delete_relations":
await knowledgeGraphManager.deleteRelations(args.relations as Relation[]);
return { content: [{ type: "text", text: "Relations deleted successfully" }] };
case "search_nodes":
return { content: [{ type: "text", text: JSON.stringify(await knowledgeGraphManager.searchNodes(args.query as string), null, 2) }] };
case "open_nodes":
return { content: [{ type: "text", text: JSON.stringify(await knowledgeGraphManager.openNodes(args.names as string[]), null, 2) }] };
default:
throw new Error(`Unknown tool: ${name}`);
}
});
const cleanup = async () => {
// Any cleanup logic if needed
};
return { server, cleanup };
}