536 lines
17 KiB
JavaScript
536 lines
17 KiB
JavaScript
#!/usr/bin/env node
|
||
|
||
import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js";
|
||
import { SSEServerTransport } from "@modelcontextprotocol/sdk/server/sse.js";
|
||
import { z } from "zod";
|
||
import { promises as fs } from 'fs';
|
||
import path from 'path';
|
||
import { fileURLToPath } from 'url';
|
||
import * as http from "http";
|
||
|
||
// Define memory file path using environment variable with fallback
|
||
export const defaultMemoryPath = path.join(path.dirname(fileURLToPath(import.meta.url)), 'memory.jsonl');
|
||
|
||
// Handle backward compatibility: migrate memory.json to memory.jsonl if needed
|
||
export async function ensureMemoryFilePath(): Promise<string> {
|
||
if (process.env.MEMORY_FILE_PATH) {
|
||
// Custom path provided, use it as-is (with absolute path resolution)
|
||
return 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);
|
||
}
|
||
|
||
// No custom path set, check for backward compatibility migration
|
||
const oldMemoryPath = path.join(path.dirname(fileURLToPath(import.meta.url)), 'memory.json');
|
||
const newMemoryPath = defaultMemoryPath;
|
||
|
||
try {
|
||
// Check if old file exists and new file doesn't
|
||
await fs.access(oldMemoryPath);
|
||
try {
|
||
await fs.access(newMemoryPath);
|
||
// Both files exist, use new one (no migration needed)
|
||
return newMemoryPath;
|
||
} catch {
|
||
// Old file exists, new file doesn't - migrate
|
||
console.error('DETECTED: Found legacy memory.json file, migrating to memory.jsonl for JSONL format compatibility');
|
||
await fs.rename(oldMemoryPath, newMemoryPath);
|
||
console.error('COMPLETED: Successfully migrated memory.json to memory.jsonl');
|
||
return newMemoryPath;
|
||
}
|
||
} catch {
|
||
// Old file doesn't exist, use new path
|
||
return newMemoryPath;
|
||
}
|
||
}
|
||
|
||
// Initialize memory file path (will be set during startup)
|
||
let MEMORY_FILE_PATH: string;
|
||
|
||
// We are storing our memory using entities, relations, and observations in a graph structure
|
||
export interface Entity {
|
||
name: string;
|
||
entityType: string;
|
||
observations: string[];
|
||
}
|
||
|
||
export interface Relation {
|
||
from: string;
|
||
to: string;
|
||
relationType: string;
|
||
}
|
||
|
||
export interface KnowledgeGraph {
|
||
entities: Entity[];
|
||
relations: Relation[];
|
||
}
|
||
|
||
// The KnowledgeGraphManager class contains all operations to interact with the knowledge graph
|
||
export class KnowledgeGraphManager {
|
||
constructor(private memoryFilePath: string) {}
|
||
|
||
private async loadGraph(): Promise<KnowledgeGraph> {
|
||
try {
|
||
const data = await fs.readFile(this.memoryFilePath, "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({
|
||
name: item.name,
|
||
entityType: item.entityType,
|
||
observations: item.observations
|
||
});
|
||
}
|
||
if (item.type === "relation") {
|
||
graph.relations.push({
|
||
from: item.from,
|
||
to: item.to,
|
||
relationType: item.relationType
|
||
});
|
||
}
|
||
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",
|
||
name: e.name,
|
||
entityType: e.entityType,
|
||
observations: e.observations
|
||
})),
|
||
...graph.relations.map(r => JSON.stringify({
|
||
type: "relation",
|
||
from: r.from,
|
||
to: r.to,
|
||
relationType: r.relationType
|
||
})),
|
||
];
|
||
await fs.writeFile(this.memoryFilePath, lines.join("\n"));
|
||
}
|
||
|
||
async createEntities(entities: Entity[]): Promise<Entity[]> {
|
||
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[]): Promise<Relation[]> {
|
||
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[] }[]): Promise<{ entityName: string; addedObservations: string[] }[]> {
|
||
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[]): Promise<void> {
|
||
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[] }[]): Promise<void> {
|
||
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[]): Promise<void> {
|
||
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(): Promise<KnowledgeGraph> {
|
||
return this.loadGraph();
|
||
}
|
||
|
||
// Very basic search function
|
||
async searchNodes(query: string): Promise<KnowledgeGraph> {
|
||
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[]): Promise<KnowledgeGraph> {
|
||
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;
|
||
}
|
||
}
|
||
|
||
let knowledgeGraphManager: KnowledgeGraphManager;
|
||
|
||
// Zod schemas for entities and relations
|
||
const EntitySchema = z.object({
|
||
name: z.string().describe("The name of the entity"),
|
||
entityType: z.string().describe("The type of the entity"),
|
||
observations: z.array(z.string()).describe("An array of observation contents associated with the entity")
|
||
});
|
||
|
||
const RelationSchema = z.object({
|
||
from: z.string().describe("The name of the entity where the relation starts"),
|
||
to: z.string().describe("The name of the entity where the relation ends"),
|
||
relationType: z.string().describe("The type of the relation")
|
||
});
|
||
|
||
// The server instance and tools exposed to Claude
|
||
const server = new McpServer({
|
||
name: "memory-server",
|
||
version: "0.6.3",
|
||
});
|
||
|
||
// Register create_entities tool
|
||
server.registerTool(
|
||
"create_entities",
|
||
{
|
||
title: "Create Entities",
|
||
description: "Create multiple new entities in the knowledge graph",
|
||
inputSchema: {
|
||
entities: z.array(EntitySchema)
|
||
},
|
||
outputSchema: {
|
||
entities: z.array(EntitySchema)
|
||
}
|
||
},
|
||
async ({ entities }) => {
|
||
const result = await knowledgeGraphManager.createEntities(entities);
|
||
return {
|
||
content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }],
|
||
structuredContent: { entities: result }
|
||
};
|
||
}
|
||
);
|
||
|
||
// Register create_relations tool
|
||
server.registerTool(
|
||
"create_relations",
|
||
{
|
||
title: "Create Relations",
|
||
description: "Create multiple new relations between entities in the knowledge graph. Relations should be in active voice",
|
||
inputSchema: {
|
||
relations: z.array(RelationSchema)
|
||
},
|
||
outputSchema: {
|
||
relations: z.array(RelationSchema)
|
||
}
|
||
},
|
||
async ({ relations }) => {
|
||
const result = await knowledgeGraphManager.createRelations(relations);
|
||
return {
|
||
content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }],
|
||
structuredContent: { relations: result }
|
||
};
|
||
}
|
||
);
|
||
|
||
// Register add_observations tool
|
||
server.registerTool(
|
||
"add_observations",
|
||
{
|
||
title: "Add Observations",
|
||
description: "Add new observations to existing entities in the knowledge graph",
|
||
inputSchema: {
|
||
observations: z.array(z.object({
|
||
entityName: z.string().describe("The name of the entity to add the observations to"),
|
||
contents: z.array(z.string()).describe("An array of observation contents to add")
|
||
}))
|
||
},
|
||
outputSchema: {
|
||
results: z.array(z.object({
|
||
entityName: z.string(),
|
||
addedObservations: z.array(z.string())
|
||
}))
|
||
}
|
||
},
|
||
async ({ observations }) => {
|
||
const result = await knowledgeGraphManager.addObservations(observations);
|
||
return {
|
||
content: [{ type: "text" as const, text: JSON.stringify(result, null, 2) }],
|
||
structuredContent: { results: result }
|
||
};
|
||
}
|
||
);
|
||
|
||
// Register delete_entities tool
|
||
server.registerTool(
|
||
"delete_entities",
|
||
{
|
||
title: "Delete Entities",
|
||
description: "Delete multiple entities and their associated relations from the knowledge graph",
|
||
inputSchema: {
|
||
entityNames: z.array(z.string()).describe("An array of entity names to delete")
|
||
},
|
||
outputSchema: {
|
||
success: z.boolean(),
|
||
message: z.string()
|
||
}
|
||
},
|
||
async ({ entityNames }) => {
|
||
await knowledgeGraphManager.deleteEntities(entityNames);
|
||
return {
|
||
content: [{ type: "text" as const, text: "Entities deleted successfully" }],
|
||
structuredContent: { success: true, message: "Entities deleted successfully" }
|
||
};
|
||
}
|
||
);
|
||
|
||
// Register delete_observations tool
|
||
server.registerTool(
|
||
"delete_observations",
|
||
{
|
||
title: "Delete Observations",
|
||
description: "Delete specific observations from entities in the knowledge graph",
|
||
inputSchema: {
|
||
deletions: z.array(z.object({
|
||
entityName: z.string().describe("The name of the entity containing the observations"),
|
||
observations: z.array(z.string()).describe("An array of observations to delete")
|
||
}))
|
||
},
|
||
outputSchema: {
|
||
success: z.boolean(),
|
||
message: z.string()
|
||
}
|
||
},
|
||
async ({ deletions }) => {
|
||
await knowledgeGraphManager.deleteObservations(deletions);
|
||
return {
|
||
content: [{ type: "text" as const, text: "Observations deleted successfully" }],
|
||
structuredContent: { success: true, message: "Observations deleted successfully" }
|
||
};
|
||
}
|
||
);
|
||
|
||
// Register delete_relations tool
|
||
server.registerTool(
|
||
"delete_relations",
|
||
{
|
||
title: "Delete Relations",
|
||
description: "Delete multiple relations from the knowledge graph",
|
||
inputSchema: {
|
||
relations: z.array(RelationSchema).describe("An array of relations to delete")
|
||
},
|
||
outputSchema: {
|
||
success: z.boolean(),
|
||
message: z.string()
|
||
}
|
||
},
|
||
async ({ relations }) => {
|
||
await knowledgeGraphManager.deleteRelations(relations);
|
||
return {
|
||
content: [{ type: "text" as const, text: "Relations deleted successfully" }],
|
||
structuredContent: { success: true, message: "Relations deleted successfully" }
|
||
};
|
||
}
|
||
);
|
||
|
||
// Register read_graph tool
|
||
server.registerTool(
|
||
"read_graph",
|
||
{
|
||
title: "Read Graph",
|
||
description: "Read the entire knowledge graph",
|
||
inputSchema: {},
|
||
outputSchema: {
|
||
entities: z.array(EntitySchema),
|
||
relations: z.array(RelationSchema)
|
||
}
|
||
},
|
||
async () => {
|
||
const graph = await knowledgeGraphManager.readGraph();
|
||
return {
|
||
content: [{ type: "text" as const, text: JSON.stringify(graph, null, 2) }],
|
||
structuredContent: { ...graph }
|
||
};
|
||
}
|
||
);
|
||
|
||
// Register search_nodes tool
|
||
server.registerTool(
|
||
"search_nodes",
|
||
{
|
||
title: "Search Nodes",
|
||
description: "Search for nodes in the knowledge graph based on a query",
|
||
inputSchema: {
|
||
query: z.string().describe("The search query to match against entity names, types, and observation content")
|
||
},
|
||
outputSchema: {
|
||
entities: z.array(EntitySchema),
|
||
relations: z.array(RelationSchema)
|
||
}
|
||
},
|
||
async ({ query }) => {
|
||
const graph = await knowledgeGraphManager.searchNodes(query);
|
||
return {
|
||
content: [{ type: "text" as const, text: JSON.stringify(graph, null, 2) }],
|
||
structuredContent: { ...graph }
|
||
};
|
||
}
|
||
);
|
||
|
||
// Register open_nodes tool
|
||
server.registerTool(
|
||
"open_nodes",
|
||
{
|
||
title: "Open Nodes",
|
||
description: "Open specific nodes in the knowledge graph by their names",
|
||
inputSchema: {
|
||
names: z.array(z.string()).describe("An array of entity names to retrieve")
|
||
},
|
||
outputSchema: {
|
||
entities: z.array(EntitySchema),
|
||
relations: z.array(RelationSchema)
|
||
}
|
||
},
|
||
async ({ names }) => {
|
||
const graph = await knowledgeGraphManager.openNodes(names);
|
||
return {
|
||
content: [{ type: "text" as const, text: JSON.stringify(graph, null, 2) }],
|
||
structuredContent: { ...graph }
|
||
};
|
||
}
|
||
);
|
||
|
||
// SSE transport session routing (sessionId -> transport)
|
||
const sseTransportsBySessionId = new Map<string, SSEServerTransport>();
|
||
|
||
function runServer() {
|
||
const port = Number(process.env.MCP_PORT ?? process.env.SSE_PORT ?? 3000);
|
||
|
||
const httpServer = http.createServer(async (req, res) => {
|
||
const url = new URL(req.url ?? "/", `http://${req.headers.host ?? "localhost"}`);
|
||
const pathname = url.pathname;
|
||
|
||
if (req.method === "GET" && (pathname === "/sse" || pathname === "/")) {
|
||
try {
|
||
const transport = new SSEServerTransport("/messages", res);
|
||
sseTransportsBySessionId.set(transport.sessionId, transport);
|
||
transport.onclose = () => {
|
||
sseTransportsBySessionId.delete(transport.sessionId);
|
||
};
|
||
await server.connect(transport);
|
||
console.error("Knowledge Graph MCP Server: new SSE client connected");
|
||
} catch (error) {
|
||
console.error("SSE connection error:", error);
|
||
if (!res.headersSent) {
|
||
res.writeHead(500).end("Internal server error");
|
||
}
|
||
}
|
||
return;
|
||
}
|
||
|
||
if (req.method === "POST" && pathname === "/messages") {
|
||
const sessionId = url.searchParams.get("sessionId");
|
||
if (!sessionId) {
|
||
res.writeHead(400).end("Missing sessionId query parameter");
|
||
return;
|
||
}
|
||
const transport = sseTransportsBySessionId.get(sessionId);
|
||
if (!transport) {
|
||
res.writeHead(404).end("Unknown session");
|
||
return;
|
||
}
|
||
await transport.handlePostMessage(req, res);
|
||
return;
|
||
}
|
||
|
||
res.writeHead(404).end("Not found");
|
||
});
|
||
|
||
httpServer.listen(port, () => {
|
||
console.error(`Knowledge Graph MCP Server running on SSE at http://localhost:${port}`);
|
||
console.error(" GET /sse – open SSE stream (then POST to /messages?sessionId=...)");
|
||
console.error(" POST /messages?sessionId=<id> – send MCP messages");
|
||
});
|
||
}
|
||
|
||
async function main() {
|
||
// Initialize memory file path with backward compatibility
|
||
MEMORY_FILE_PATH = await ensureMemoryFilePath();
|
||
|
||
// Initialize knowledge graph manager with the memory file path
|
||
knowledgeGraphManager = new KnowledgeGraphManager(MEMORY_FILE_PATH);
|
||
|
||
runServer();
|
||
}
|
||
|
||
main().catch((error) => {
|
||
console.error("Fatal error in main():", error);
|
||
process.exit(1);
|
||
});
|