graph-memory-mcp-server/src/index.ts

Graphiti Memory Graph MCP Server - Search and explore knowledge graphs using Graphiti

/**
 * 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: string) => {
    fs.appendFileSync(logFile, `[INFO] ${new Date().toISOString()} - ${message}\n`);
  },
  error: (message: string, error?: any) => {
    fs.appendFileSync(logFile, `[ERROR] ${new Date().toISOString()} - ${message}\n`);
    if (error) {
      fs.appendFileSync(logFile, `${error.stack || error}\n`);
    }
  }
};

type ZepConfig = {
  apiKey: string;
  baseUrl: string;
};

let zepConfig: ZepConfig;
let apiClient: any;
let zepClient: 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: MongoClient | null = null;

/**
 * Get MongoDB client instance
 */
async function getMongoClient(): Promise<MongoClient> {
  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<T>(args: any, defaultValues: T): T {
  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(): ZepConfig {
  // 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: ZepConfig) {
  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: any = {
          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: [] as number[]
      });
      
      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: any) => 
            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: string, limit: number) => {
                  const body: any = {
                      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: any = {
                  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: string) => {
              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: string, limit: number) => {
              const body: any = {
                  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));