graph-memory-mcp-server
Version:
Graphiti Memory Graph MCP Server - Search and explore knowledge graphs using Graphiti
710 lines (709 loc) • 30.3 kB
JavaScript
/**
* Zep Graphiti MCP Server
* A comprehensive server implementation using Model Context Protocol for Zep's Knowledge Graph
*/
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 * as fs from 'fs';
import * as path from 'path';
import * as os from 'os';
import axios from 'axios';
import { ZepClient } from '@getzep/zep-cloud';
import { MongoClient, ObjectId } from 'mongodb';
import { tools } from "./resources.js";
const logFile = path.join(os.tmpdir(), 'zep-graphiti-mcp.log');
fs.writeFileSync(logFile, `[INFO] ${new Date().toISOString()} - Starting Zep Graphiti MCP Server...\n`);
const logger = {
log: (message) => {
fs.appendFileSync(logFile, `[INFO] ${new Date().toISOString()} - ${message}\n`);
},
error: (message, error) => {
fs.appendFileSync(logFile, `[ERROR] ${new Date().toISOString()} - ${message}\n`);
if (error) {
fs.appendFileSync(logFile, `${error.stack || error}\n`);
}
}
};
let zepConfig;
let apiClient;
let zepClient;
// MongoDB Configuration
const mongoConfig = {
mongoUri: process.env.MONGO_URI || process.env.MONGODB_URI || '',
dbName: process.env.MONGO_DB_NAME || process.env.DB_NAME || 'graphiti-memory'
};
let mongoClient = null;
/**
* Get MongoDB client instance
*/
async function getMongoClient() {
if (!mongoConfig.mongoUri) {
throw new Error('MongoDB URI is required. Set MONGO_URI or MONGODB_URI environment variable.');
}
if (!mongoClient) {
mongoClient = new MongoClient(mongoConfig.mongoUri);
await mongoClient.connect();
logger.log('MongoDB client connected successfully');
}
return mongoClient;
}
/**
* Helper function to safely extract parameters from request arguments
*/
function safeGetArgs(args, defaultValues) {
if (!args || typeof args !== 'object') {
return defaultValues;
}
const result = { ...defaultValues };
for (const key in defaultValues) {
if (args[key] !== undefined) {
result[key] = args[key];
}
}
return result;
}
function parseArgs() {
// Configuration is now handled via environment variables set by CLI
const apiKey = process.env.ZEP_API_KEY;
const baseUrl = process.env.ZEP_API_BASE_URL || 'https://api.getzep.com';
if (!apiKey) {
throw new Error('Zep API key is required. Set ZEP_API_KEY environment variable.');
}
return {
apiKey,
baseUrl
};
}
function initApiClient(config) {
logger.log(`Initializing API client with base URL: ${config.baseUrl}`);
logger.log(`API Key length: ${config.apiKey ? config.apiKey.length : 0}`);
logger.log(`API Key starts with: ${config.apiKey ? config.apiKey.substring(0, 10) + '...' : 'undefined'}`);
// Initialize axios client for direct API calls
const client = axios.create({
baseURL: config.baseUrl,
headers: {
'Authorization': `Api-Key ${config.apiKey}`,
'Content-Type': 'application/json'
},
timeout: 30000
});
// Initialize official Zep Cloud SDK client
zepClient = new ZepClient({
apiKey: config.apiKey
});
// Add request interceptor for logging
client.interceptors.request.use(request => {
logger.log(`Making request to: ${request.method?.toUpperCase()} ${request.url}`);
logger.log(`Request headers: ${JSON.stringify(request.headers)}`);
logger.log(`Request body: ${JSON.stringify(request.data)}`);
return request;
}, error => {
logger.error('Request error:', error);
return Promise.reject(error);
});
// Add response interceptor for logging
client.interceptors.response.use(response => {
logger.log(`Response status: ${response.status}`);
return response;
}, error => {
logger.error(`Response error: ${error.response?.status} ${error.response?.statusText}`);
logger.error(`Response data: ${JSON.stringify(error.response?.data)}`);
return Promise.reject(error);
});
return client;
}
const server = new Server({
name: "memory_search",
version: "1.0.0"
}, {
capabilities: {
tools: {
list: true,
call: true
}
}
});
/**
* List available tools for interacting with Zep Graphiti.
*/
server.setRequestHandler(ListToolsRequestSchema, async () => {
return {
tools: tools
};
});
/**
* Handle tool calls to the Zep API.
*/
server.setRequestHandler(CallToolRequestSchema, async (request) => {
logger.log('Received call tool request: ' + JSON.stringify(request));
// Ensure API client is initialized
if (!apiClient) {
if (!zepConfig) {
throw new Error("Zep API client is not initialized. Please configure it before querying.");
}
apiClient = initApiClient(zepConfig);
}
switch (request.params.name) {
case "get_entity_details": {
const args = safeGetArgs(request.params.arguments, {
node_uuid: '',
include_relationships: true,
relationship_depth: 1
});
if (!args.node_uuid) {
throw new Error("Node UUID is required");
}
try {
// Get the node details
const nodeResponse = await apiClient.get(`/api/v2/graph/node/${args.node_uuid}`);
const result = {
node: nodeResponse.data
};
if (args.include_relationships) {
// Search for edges connected to this node
const edgeSearchBody = {
query: result.node.name || args.node_uuid,
center_node_uuid: args.node_uuid,
reranker: "node_distance",
scope: "edges",
limit: 50
};
try {
const edgeResponse = await apiClient.post('/api/v2/graph/search', edgeSearchBody);
result.connected_edges = edgeResponse.data;
}
catch (edgeError) {
logger.error('Error fetching connected edges:', edgeError);
result.connected_edges = [];
}
}
return {
content: [{
type: "text",
text: JSON.stringify(result, null, 2)
}]
};
}
catch (error) {
throw new Error(`Failed to get entity details: ${error instanceof Error ? error.message : 'Unknown error'}`);
}
}
case "get_users": {
const args = safeGetArgs(request.params.arguments, {
pageNumber: 1,
pageSize: 20
});
try {
const queryParams = new URLSearchParams();
queryParams.append('pageNumber', args.pageNumber.toString());
queryParams.append('pageSize', args.pageSize.toString());
const response = await apiClient.get(`/api/v2/users-ordered?${queryParams.toString()}`);
return {
content: [{
type: "text",
text: JSON.stringify(response.data, null, 2)
}]
};
}
catch (error) {
throw new Error(`Failed to get users: ${error instanceof Error ? error.message : 'Unknown error'}`);
}
}
case "get_casefile_page_by_uuid": {
const args = safeGetArgs(request.params.arguments, {
uuid: ''
});
if (!args.uuid) {
throw new Error("Page UUID is required");
}
try {
const response = await apiClient.get(`/api/v2/graph/episodes/${args.uuid}`);
return {
content: [{
type: "text",
text: JSON.stringify(response.data, null, 2)
}]
};
}
catch (error) {
throw new Error(`Failed to get casefile page details: ${error instanceof Error ? error.message : 'Unknown error'}`);
}
}
case "get_casefile_index": {
const args = safeGetArgs(request.params.arguments, {
casefile_id: '',
limit: 10,
next_iter: 0
});
if (!args.casefile_id) {
throw new Error("casefile_id is required");
}
try {
const mongoClientInstance = await getMongoClient();
const db = mongoClientInstance.db(mongoConfig.dbName);
const collection = db.collection("casefiles");
// Get total entries count
const pipelineLength = [
{ $match: { _id: new ObjectId(args.casefile_id) } },
{ $project: { index: { $size: "$index" } } }
];
const result = await collection.aggregate(pipelineLength).toArray();
if (!result || result.length === 0) {
return {
content: [{
type: "text",
text: JSON.stringify({ error: `No casefile found with ID: ${args.casefile_id}` }, null, 2)
}]
};
}
const total_entries = result[0].index;
let skip;
if (args.next_iter === 0) {
skip = Math.max(0, total_entries - args.limit);
}
else {
skip = Math.max(0, total_entries - (args.limit * (args.next_iter + 1)));
if (skip < 0)
skip = 0;
}
const entries_to_take = Math.min(args.limit, total_entries - skip);
if (entries_to_take <= 0) {
return {
content: [{
type: "text",
text: JSON.stringify({ message: `No more index entries to retrieve for casefile ID: ${args.casefile_id}` }, null, 2)
}]
};
}
const pipelineFetch = [
{ $match: { _id: new ObjectId(args.casefile_id) } },
{
$project: {
index: { $slice: ["$index", skip, entries_to_take] },
category: 1,
imo: 1,
link: 1,
casefile: 1,
_id: 1
}
}
];
const fetchResult = await collection.aggregate(pipelineFetch).toArray();
if (!fetchResult || fetchResult.length === 0) {
return {
content: [{
type: "text",
text: JSON.stringify({ error: `Failed to retrieve index entries for casefile ID: ${args.casefile_id}` }, null, 2)
}]
};
}
const document = fetchResult[0];
const response = {
casefile_id: document._id.toString(),
index: document.index || [],
category: document.category,
imo: document.imo,
link: document.link,
casefile: document.casefile,
total_entries: total_entries,
current_batch: args.next_iter,
total_batches: Math.max(1, Math.ceil(total_entries / args.limit)),
has_more: skip > 0
};
return {
content: [{
type: "text",
text: JSON.stringify(response, null, 2)
}]
};
}
catch (error) {
throw new Error(`Error getting casefile index: ${error instanceof Error ? error.message : 'Unknown error'}`);
}
}
case "get_casefile_pages_by_number": {
const args = safeGetArgs(request.params.arguments, {
casefile_id: '',
pages: []
});
if (!args.casefile_id) {
throw new Error("casefile_id is required");
}
let page_list = args.pages || [];
if (!Array.isArray(page_list) || page_list.length === 0) {
throw new Error("pages must be a non-empty list");
}
try {
page_list = [...new Set(page_list.map(p => parseInt(p.toString())))].sort((a, b) => a - b);
}
catch (error) {
throw new Error("All entries in 'pages' must be integers or castable to integers");
}
try {
const mongoClientInstance = await getMongoClient();
const db = mongoClientInstance.db(mongoConfig.dbName);
const collection = db.collection("casefiles");
const doc = await collection.findOne({ _id: new ObjectId(args.casefile_id) }, { projection: { link: 1, casefile: 1 } });
if (!doc) {
return {
content: [{
type: "text",
text: JSON.stringify({ error: `No casefile found with ID: ${args.casefile_id}` }, null, 2)
}]
};
}
const page_indices = page_list.map(p => p - 1); // convert to 0-based
const min_index = Math.min(...page_indices);
const max_index = Math.max(...page_indices);
const slice_start = Math.max(0, min_index);
const slice_count = max_index - min_index + 1;
const pipeline = [
{ $match: { _id: new ObjectId(args.casefile_id) } },
{
$project: {
pages: { $slice: ["$pages", slice_start, slice_count] },
_id: 0
}
}
];
const result = await collection.aggregate(pipeline).toArray();
if (!result || result.length === 0) {
return {
content: [{
type: "text",
text: JSON.stringify({ error: "No pages found for the given casefile ID." }, null, 2)
}]
};
}
const sliced_pages = result[0].pages || [];
const requested_pages = [];
for (const page_num of page_list) {
const actual_index = page_num - 1; // 0-based
const i = actual_index - slice_start;
if (i >= 0 && i < sliced_pages.length) {
requested_pages.push(sliced_pages[i]);
}
}
if (requested_pages.length === 0) {
return {
content: [{
type: "text",
text: JSON.stringify({ message: "No pages matched the requested indices" }, null, 2)
}]
};
}
const response = {
casefile_id: args.casefile_id,
casefile: doc.casefile,
pages_requested: page_list,
pages_returned: requested_pages,
link: doc.link
};
return {
content: [{
type: "text",
text: JSON.stringify(response, null, 2)
}]
};
}
catch (error) {
throw new Error(`Error getting casefile pages: ${error instanceof Error ? error.message : 'Unknown error'}`);
}
}
case "get_recent_user_emails": {
const args = safeGetArgs(request.params.arguments, {
user_id: '',
lastn: 20,
role_filter: 'all'
});
if (!args.user_id) {
throw new Error("user_id is required");
}
try {
let episodesResponse;
// Use official Zep Cloud SDK
if (args.user_id) {
episodesResponse = await zepClient.graph.episode.getByUserId(args.user_id);
}
else {
// For group episodes, we'll need to handle this differently
// For now, throw an error as the SDK method for groups might be different
throw new Error("Group episodes retrieval not yet implemented with SDK. Please use user_id.");
}
let episodes = episodesResponse.episodes || [];
// Limit to requested number
if (args.lastn && args.lastn < episodes.length) {
episodes = episodes.slice(0, args.lastn);
}
// Apply role filter if specified
if (args.role_filter !== 'all') {
episodes = episodes.filter((ep) => ep.roleType === args.role_filter || ep.role === args.role_filter);
}
return {
content: [{
type: "text",
text: JSON.stringify(episodes, null, 2)
}]
};
}
catch (error) {
throw new Error(`Failed to get recent user emails: ${error instanceof Error ? error.message : 'Unknown error'}`);
}
}
case "search_casefile_content": {
const args = safeGetArgs(request.params.arguments, {
query: '',
user_id: '',
search_focus: 'comprehensive',
max_results: 15,
center_around: ''
});
if (!args.query) {
throw new Error("Query is required");
}
try {
if (args.search_focus === 'comprehensive') {
const limitPerScope = Math.floor(args.max_results / 3);
const remainder = args.max_results % 3;
const searchPromises = [];
const createSearchBody = (scope, limit) => {
const body = {
query: args.query,
limit: limit,
scope: scope
};
if (args.user_id)
body.user_id = args.user_id;
if (args.center_around)
body.center_node_uuid = args.center_around;
return body;
};
if (limitPerScope > 0) {
searchPromises.push(apiClient.post('/api/v2/graph/search', createSearchBody('edges', limitPerScope + remainder)));
searchPromises.push(apiClient.post('/api/v2/graph/search', createSearchBody('nodes', limitPerScope)));
searchPromises.push(apiClient.post('/api/v2/graph/search', createSearchBody('episodes', limitPerScope)));
}
else if (args.max_results > 0) {
searchPromises.push(apiClient.post('/api/v2/graph/search', createSearchBody('edges', args.max_results)));
}
else {
return {
content: [{
type: "text",
text: JSON.stringify({ edges: [], nodes: [], episodes: [], message: "max_results is 0" }, null, 2)
}]
};
}
const responses = await Promise.all(searchPromises);
const combinedResult = responses.reduce((acc, response) => {
const data = response.data;
if (data.edges)
acc.edges.push(...data.edges);
if (data.nodes)
acc.nodes.push(...data.nodes);
if (data.episodes)
acc.episodes.push(...data.episodes);
return acc;
}, { edges: [], nodes: [], episodes: [] });
return {
content: [{
type: "text",
text: JSON.stringify(combinedResult, null, 2)
}]
};
}
else {
let scope = 'edges';
if (args.search_focus === 'facts_only')
scope = 'edges';
else if (args.search_focus === 'entities_only')
scope = 'nodes';
else if (args.search_focus === 'emails_only')
scope = 'episodes';
const searchBody = {
query: args.query,
limit: args.max_results,
scope: scope,
};
if (args.user_id)
searchBody.user_id = args.user_id;
if (args.center_around)
searchBody.center_node_uuid = args.center_around;
const response = await apiClient.post('/api/v2/graph/search', searchBody);
return {
content: [{
type: "text",
text: JSON.stringify(response.data, null, 2)
}]
};
}
}
catch (error) {
throw new Error(`Failed to perform casefile content search: ${error instanceof Error ? error.message : 'Unknown error'}`);
}
}
case "advanced_casefile_search": {
const args = safeGetArgs(request.params.arguments, {
query: '',
user_id: '',
ranking_strategy: 'balanced_ranking',
search_focus: 'comprehensive',
reference_point: '',
max_results: 15
});
if (!args.query) {
throw new Error("Query is required");
}
try {
const getSearchConfig = (strategy) => {
switch (strategy) {
case "most_accurate":
return { reranker: "cross_encoder" };
case "proximity_focused":
return { reranker: "node_distance" };
case "recent_emails":
return { reranker: "episode_mentions" };
case "diverse_perspectives":
return { reranker: "mmr" };
case "balanced_ranking":
default:
return { reranker: "rrf" };
}
};
const config = getSearchConfig(args.ranking_strategy);
const createSearchBody = (scope, limit) => {
const body = {
query: args.query,
limit: limit,
scope: scope,
reranker: config.reranker,
};
if (args.user_id)
body.user_id = args.user_id;
if (args.reference_point)
body.center_node_uuid = args.reference_point;
return body;
};
if (args.search_focus === 'comprehensive') {
const limitPerScope = Math.floor(args.max_results / 3);
const remainder = args.max_results % 3;
const searchPromises = [];
if (limitPerScope > 0) {
searchPromises.push(apiClient.post('/api/v2/graph/search', createSearchBody('edges', limitPerScope + remainder)));
searchPromises.push(apiClient.post('/api/v2/graph/search', createSearchBody('nodes', limitPerScope)));
searchPromises.push(apiClient.post('/api/v2/graph/search', createSearchBody('episodes', limitPerScope)));
}
else if (args.max_results > 0) {
searchPromises.push(apiClient.post('/api/v2/graph/search', createSearchBody('edges', args.max_results)));
}
else {
return {
content: [{
type: "text",
text: JSON.stringify({ edges: [], nodes: [], episodes: [], message: "max_results is 0" }, null, 2)
}]
};
}
const responses = await Promise.all(searchPromises);
const combinedResult = responses.reduce((acc, response) => {
const data = response.data;
if (data.edges)
acc.edges.push(...data.edges);
if (data.nodes)
acc.nodes.push(...data.nodes);
if (data.episodes)
acc.episodes.push(...data.episodes);
return acc;
}, { edges: [], nodes: [], episodes: [] });
return {
content: [{
type: "text",
text: JSON.stringify(combinedResult, null, 2)
}]
};
}
else {
let scope = 'edges';
if (args.search_focus === 'facts_only')
scope = 'edges';
else if (args.search_focus === 'entities_only')
scope = 'nodes';
else if (args.search_focus === 'emails_only')
scope = 'episodes';
const searchBody = createSearchBody(scope, args.max_results);
const response = await apiClient.post('/api/v2/graph/search', searchBody);
return {
content: [{
type: "text",
text: JSON.stringify(response.data, null, 2)
}]
};
}
}
catch (error) {
throw new Error(`Failed to perform advanced casefile search: ${error instanceof Error ? error.message : 'Unknown error'}`);
}
}
default:
throw new Error(`Unknown tool: ${request.params.name}`);
}
});
/**
* Main function to initialize and run the MCP server
*/
async function main() {
try {
zepConfig = parseArgs();
logger.log('Zep API configuration loaded');
apiClient = initApiClient(zepConfig);
logger.log('Zep API client initialized');
try {
// Test connection to API
const testResponse = await apiClient.get('/api/v2/users-ordered?pageNumber=1&pageSize=1');
logger.log('Zep API connection test successful');
}
catch (error) {
logger.error('Zep API connection test failed:', error);
logger.log('Continuing anyway - connection will be tested on first use');
}
// Test MongoDB connection
try {
await getMongoClient();
logger.log('MongoDB connection test successful');
}
catch (error) {
logger.error('MongoDB connection test failed:', error);
logger.log('Continuing anyway - connection will be tested on first use');
}
logger.log('Connecting to stdio transport...');
const transport = new StdioServerTransport();
await server.connect(transport);
logger.log('Zep Graphiti MCP server connected and ready');
// Handle process termination
process.on('SIGINT', async () => {
logger.log('Received SIGINT, shutting down gracefully...');
if (mongoClient) {
await mongoClient.close();
logger.log('MongoDB connection closed');
}
process.exit(0);
});
process.on('SIGTERM', async () => {
logger.log('Received SIGTERM, shutting down gracefully...');
if (mongoClient) {
await mongoClient.close();
logger.log('MongoDB connection closed');
}
process.exit(0);
});
}
catch (error) {
logger.error('Error running MCP server:', error);
if (mongoClient) {
await mongoClient.close();
}
process.exit(1);
}
}
main().catch(err => logger.error('Unhandled error:', err));