index.ts•25.7 kB
#!/usr/bin/env node
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
import {
CallToolRequestSchema,
ListToolsRequestSchema,
} from "@modelcontextprotocol/sdk/types.js";
import { promises as fs } from 'fs';
import path from 'path';
import { fileURLToPath } from 'url';
import express from 'express';
import { randomUUID } from 'node:crypto';
// 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[];
}
interface LocationMatch {
text: string;
start: number;
end: number;
type: 'city' | 'state' | 'country' | 'address' | 'landmark';
}
// Location extraction utilities
class LocationExtractor {
private static readonly LOCATION_PATTERNS = [
// Cities with state/country: "New York, NY", "Paris, France"
/\b([A-Z][a-z]+(?:\s+[A-Z][a-z]+)*),\s*([A-Z]{2}|[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*)\b/g,
// Addresses: "123 Main St", "456 Oak Avenue"
/\b\d+\s+[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\s+(?:St|Street|Ave|Avenue|Rd|Road|Blvd|Boulevard|Dr|Drive|Way|Lane|Ln)\b/gi,
// Landmarks/Places: "Central Park", "Golden Gate Bridge"
/\b(?:Mount|Mt\.?|Lake|River|Park|Bridge|University|Hospital|Airport|Station)\s+[A-Z][a-z]+(?:\s+[A-Z][a-z]+)*\b/g,
// States/Provinces: "California", "Ontario"
/\b(?:Alabama|Alaska|Arizona|Arkansas|California|Colorado|Connecticut|Delaware|Florida|Georgia|Hawaii|Idaho|Illinois|Indiana|Iowa|Kansas|Kentucky|Louisiana|Maine|Maryland|Massachusetts|Michigan|Minnesota|Mississippi|Missouri|Montana|Nebraska|Nevada|New\s+Hampshire|New\s+Jersey|New\s+Mexico|New\s+York|North\s+Carolina|North\s+Dakota|Ohio|Oklahoma|Oregon|Pennsylvania|Rhode\s+Island|South\s+Carolina|South\s+Dakota|Tennessee|Texas|Utah|Vermont|Virginia|Washington|West\s+Virginia|Wisconsin|Wyoming)\b/g,
// Countries: "United States", "United Kingdom", etc.
/\b(?:United\s+States|United\s+Kingdom|Canada|Mexico|France|Germany|Italy|Spain|Japan|China|India|Australia|Brazil|Argentina)\b/g
];
static extractLocations(text: string): LocationMatch[] {
const matches: LocationMatch[] = [];
this.LOCATION_PATTERNS.forEach((pattern, index) => {
let match;
const regex = new RegExp(pattern.source, pattern.flags);
while ((match = regex.exec(text)) !== null) {
const locationText = match[0].trim();
let type: LocationMatch['type'] = 'landmark';
// Determine location type based on pattern
if (index === 0) type = 'city'; // City, State pattern
else if (index === 1) type = 'address'; // Street address
else if (index === 2) type = 'landmark'; // Named places
else if (index === 3) type = 'state'; // US States
else if (index === 4) type = 'country'; // Countries
matches.push({
text: locationText,
start: match.index,
end: match.index + locationText.length,
type
});
}
});
// Remove duplicates and overlapping matches
return this.deduplicateMatches(matches);
}
private static deduplicateMatches(matches: LocationMatch[]): LocationMatch[] {
// Sort by start position
matches.sort((a, b) => a.start - b.start);
const filtered: LocationMatch[] = [];
for (const match of matches) {
// Check if this match overlaps with any previous match
const overlaps = filtered.some(existing =>
(match.start >= existing.start && match.start < existing.end) ||
(match.end > existing.start && match.end <= existing.end)
);
if (!overlaps) {
filtered.push(match);
}
}
return filtered;
}
}
// The KnowledgeGraphManager class contains all operations to interact with the knowledge graph
class KnowledgeGraphManager {
private async loadGraph(): Promise<KnowledgeGraph> {
try {
const data = await fs.readFile(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",
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(MEMORY_FILE_PATH, 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;
}
async extractAndAddLocations(text: string, sourceEntity?: string): Promise<{ entities: Entity[], relations: Relation[] }> {
const locations = LocationExtractor.extractLocations(text);
const newEntities: Entity[] = [];
const newRelations: Relation[] = [];
for (const location of locations) {
// Create location entity
const locationEntity: Entity = {
name: location.text,
entityType: 'location',
observations: [
`Location type: ${location.type}`,
`Extracted from text: "${text.substring(Math.max(0, location.start - 20), Math.min(text.length, location.end + 20))}"`,
`Original context: positions ${location.start}-${location.end}`
]
};
newEntities.push(locationEntity);
// If source entity provided, create relation
if (sourceEntity) {
const relation: Relation = {
from: sourceEntity,
to: location.text,
relationType: 'mentions_location'
};
newRelations.push(relation);
}
// Create hierarchical relations for compound locations (e.g., "New York, NY")
if (location.type === 'city' && location.text.includes(',')) {
const parts = location.text.split(',').map(p => p.trim());
if (parts.length === 2) {
const [city, stateOrCountry] = parts;
// Create state/country entity if it doesn't exist
const parentEntity: Entity = {
name: stateOrCountry,
entityType: 'location',
observations: [
`Location type: ${stateOrCountry.length <= 3 ? 'state' : 'country'}`,
`Parent location of: ${city}`
]
};
newEntities.push(parentEntity);
// Create hierarchical relation
const hierarchyRelation: Relation = {
from: city,
to: stateOrCountry,
relationType: 'located_in'
};
newRelations.push(hierarchyRelation);
}
}
}
// Add to graph
const createdEntities = await this.createEntities(newEntities);
const createdRelations = await this.createRelations(newRelations);
return {
entities: createdEntities,
relations: createdRelations
};
}
}
const knowledgeGraphManager = new KnowledgeGraphManager();
// The server instance shared across all sessions
const server = new Server({
name: "memory-server",
version: "0.6.3",
}, {
capabilities: {
tools: {},
},
},);
server.setRequestHandler(ListToolsRequestSchema, async () => {
return {
tools: [
{
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"],
additionalProperties: false,
},
},
},
required: ["entities"],
additionalProperties: false,
},
},
{
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"],
additionalProperties: false,
},
},
},
required: ["relations"],
additionalProperties: false,
},
},
{
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"],
additionalProperties: false,
},
},
},
required: ["observations"],
additionalProperties: false,
},
},
{
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"
},
},
required: ["entityNames"],
additionalProperties: false,
},
},
{
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"],
additionalProperties: false,
},
},
},
required: ["deletions"],
additionalProperties: false,
},
},
{
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"],
additionalProperties: false,
},
description: "An array of relations to delete"
},
},
required: ["relations"],
additionalProperties: false,
},
},
{
name: "read_graph",
description: "Read the entire knowledge graph",
inputSchema: {
type: "object",
properties: {},
additionalProperties: false,
},
},
{
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" },
},
required: ["query"],
additionalProperties: false,
},
},
{
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",
},
},
required: ["names"],
additionalProperties: false,
},
},
{
name: "extract_locations",
description: "Extract locations from text and add them to the knowledge graph as entities with geographic relationships",
inputSchema: {
type: "object",
properties: {
text: {
type: "string",
description: "The text to extract locations from"
},
sourceEntity: {
type: "string",
description: "Optional: name of source entity that mentions these locations (creates 'mentions_location' relations)"
},
},
required: ["text"],
additionalProperties: false,
},
},
],
};
});
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 { entityName: string; contents: string[] }[]), 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 { entityName: string; observations: string[] }[]);
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) }] };
case "extract_locations":
const result = await knowledgeGraphManager.extractAndAddLocations(
args.text as string,
args.sourceEntity as string | undefined
);
return { content: [{ type: "text", text: JSON.stringify(result, null, 2) }] };
default:
throw new Error(`Unknown tool: ${name}`);
}
});
async function main() {
// Check if running in stdio mode (for local testing)
const isStdio = process.argv.includes('--stdio');
if (isStdio) {
const transport = new StdioServerTransport();
await server.connect(transport);
console.error("Knowledge Graph MCP Server running on stdio");
return;
}
// HTTP mode for Smithery deployment
const app = express();
const port = parseInt(process.env.PORT || '8081');
// CORS middleware
app.use((req, res, next) => {
res.header('Access-Control-Allow-Origin', '*');
res.header('Access-Control-Allow-Credentials', 'true');
res.header('Access-Control-Allow-Methods', 'GET, POST, OPTIONS');
res.header('Access-Control-Allow-Headers', '*');
res.header('Access-Control-Expose-Headers', 'mcp-session-id, mcp-protocol-version');
if (req.method === 'OPTIONS') {
return res.sendStatus(200);
}
next();
});
// Store transports by session ID
const transports: { [sessionId: string]: StreamableHTTPServerTransport } = {};
// NOTE: Avoid body-parsing middleware before the MCP route; the transport
// needs access to the raw request stream to process multipart messages.
// MCP Streamable HTTP endpoint
app.post('/mcp', async (req, res) => {
const sessionId = req.headers['mcp-session-id'] as string | undefined;
try {
let transport: StreamableHTTPServerTransport;
if (sessionId && transports[sessionId]) {
// Existing session - reuse transport
transport = transports[sessionId];
} else {
// New session - create new transport
transport = new StreamableHTTPServerTransport({
sessionIdGenerator: () => randomUUID(),
enableDnsRebindingProtection: false,
onsessioninitialized: (newSessionId: string) => {
transports[newSessionId] = transport;
console.error(`Session initialized: ${newSessionId}`);
},
onsessionclosed: (closedSessionId: string) => {
delete transports[closedSessionId];
console.error(`Session closed: ${closedSessionId}`);
},
});
// Clean up on transport close
transport.onclose = () => {
if (transport.sessionId && transports[transport.sessionId]) {
delete transports[transport.sessionId];
console.error(`Transport closed for session: ${transport.sessionId}`);
}
};
// Connect the shared server to this transport
await server.connect(transport);
}
// Handle the request
await transport.handleRequest(req, res);
} catch (error) {
console.error('Error handling MCP request:', error);
if (!res.headersSent) {
res.status(500).json({
jsonrpc: '2.0',
error: {
code: -32603,
message: 'Internal server error',
data: String(error)
},
id: null
});
}
}
});
// Health check endpoint
app.get('/health', (_req, res) => {
res.json({ status: 'ok', sessions: Object.keys(transports).length });
});
app.listen(port, '0.0.0.0', () => {
console.error(`MCP Server listening on port ${port}`);
console.error(`MCP endpoint: http://0.0.0.0:${port}/mcp`);
console.error(`Health endpoint: http://0.0.0.0:${port}/health`);
});
}
main().catch((error) => {
console.error("Fatal error in main():", error);
process.exit(1);
});