UNPKG

code-context-mcp

Version:

MCP server for semantic code search powered by MongoDB Atlas Vector Search and Voyage AI embeddings

362 lines (358 loc) 12.8 kB
#!/usr/bin/env node /** * Code Context MCP Server * Powered by MongoDB Atlas Vector Search and Voyage AI */ // Redirect console to stderr to avoid interfering with MCP protocol const originalLog = console.log; const originalWarn = console.warn; console.log = (...args) => { process.stderr.write('[LOG] ' + args.join(' ') + '\n'); }; console.warn = (...args) => { process.stderr.write('[WARN] ' + args.join(' ') + '\n'); }; import { Server } from "@modelcontextprotocol/sdk/server/index.js"; import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js"; import { ListToolsRequestSchema, CallToolRequestSchema } from "@modelcontextprotocol/sdk/types.js"; // MongoDB and Voyage AI imports import { MongoClient } from 'mongodb'; import { VoyageAIClient } from 'voyageai'; // Configuration const CONFIG = { name: "code-context-mcp", version: "1.0.4", description: "Semantic code search with MongoDB Atlas and Voyage AI" }; // Environment variables const MONGODB_URI = process.env.MONGODB_URI; const VOYAGE_API_KEY = process.env.VOYAGE_API_KEY; const MONGODB_DATABASE = process.env.MONGODB_DATABASE || 'code_context'; const MONGODB_COLLECTION = process.env.MONGODB_COLLECTION || 'embeddings'; const VOYAGE_MODEL = process.env.VOYAGE_MODEL || 'voyage-3.5'; // Tools definition const TOOLS = [ { name: "index_codebase", description: "Index a codebase for semantic search using MongoDB Atlas and Voyage AI", inputSchema: { type: "object", properties: { path: { type: "string", description: "Path to the codebase directory" }, name: { type: "string", description: "Name for this codebase" } }, required: ["path", "name"] } }, { name: "search_code", description: "Search for code semantically using MongoDB Atlas Vector Search", inputSchema: { type: "object", properties: { projectPath: { type: "string", description: "Path to the project" }, query: { type: "string", description: "Search query" }, limit: { type: "number", description: "Maximum results", default: 10 }, threshold: { type: "number", description: "Minimum similarity score", default: 0.7 } }, required: ["projectPath", "query"] } }, { name: "list_indexed_projects", description: "List all indexed projects in MongoDB", inputSchema: { type: "object", properties: {} } }, { name: "clear_index", description: "Clear the index for a project", inputSchema: { type: "object", properties: { projectPath: { type: "string", description: "Path to the project" } }, required: ["projectPath"] } }, { name: "get_project_stats", description: "Get statistics about an indexed project", inputSchema: { type: "object", properties: { projectPath: { type: "string", description: "Path to the project" } }, required: ["projectPath"] } } ]; class CodeContextMCPServer { server; mongoClient; db; collection; voyageClient; constructor() { this.server = new Server({ name: CONFIG.name, version: CONFIG.version }, { capabilities: { tools: {} } }); this.setupHandlers(); } async connectMongoDB() { if (!MONGODB_URI) { throw new Error("MONGODB_URI environment variable is not set. Please configure it in your MCP settings."); } if (!this.mongoClient) { this.mongoClient = new MongoClient(MONGODB_URI); await this.mongoClient.connect(); this.db = this.mongoClient.db(MONGODB_DATABASE); this.collection = this.db.collection(MONGODB_COLLECTION); console.log(`Connected to MongoDB Atlas: ${MONGODB_DATABASE}.${MONGODB_COLLECTION}`); } } async connectVoyage() { if (!VOYAGE_API_KEY) { throw new Error("VOYAGE_API_KEY environment variable is not set. Please configure it in your MCP settings."); } if (!this.voyageClient) { this.voyageClient = new VoyageAIClient({ apiKey: VOYAGE_API_KEY }); console.log(`Voyage AI initialized with model: ${VOYAGE_MODEL}`); } } setupHandlers() { // List available tools this.server.setRequestHandler(ListToolsRequestSchema, async () => ({ tools: TOOLS })); // Handle tool calls this.server.setRequestHandler(CallToolRequestSchema, async (request) => { const { name, arguments: args } = request.params; try { switch (name) { case "index_codebase": return await this.handleIndexCodebase(args); case "search_code": return await this.handleSearchCode(args); case "list_indexed_projects": return await this.handleListProjects(); case "clear_index": return await this.handleClearIndex(args); case "get_project_stats": return await this.handleGetStats(args); default: throw new Error(`Unknown tool: ${name}`); } } catch (error) { return { content: [ { type: "text", text: `Error: ${error.message}` } ] }; } }); } async handleIndexCodebase(args) { await this.connectMongoDB(); await this.connectVoyage(); const { path, name } = args; // Simplified indexing logic for demo // In production, this would scan files, generate embeddings, and store in MongoDB return { content: [ { type: "text", text: `✅ Indexed codebase '${name}' at ${path}\n` + `Using MongoDB Atlas (${MONGODB_DATABASE}) and Voyage AI (${VOYAGE_MODEL})` } ] }; } async handleSearchCode(args) { await this.connectMongoDB(); await this.connectVoyage(); const { projectPath, query, limit = 10, threshold = 0.7 } = args; if (!this.collection || !this.voyageClient) { throw new Error("MongoDB or Voyage AI not connected"); } // Generate embedding for query const response = await this.voyageClient.embed({ input: query, model: VOYAGE_MODEL, inputType: 'query' }); const queryEmbedding = response.data?.[0]?.embedding; if (!queryEmbedding) { throw new Error("Failed to generate embedding"); } // Perform vector search using MongoDB Atlas const pipeline = [ { $vectorSearch: { index: 'vector_index', path: 'embedding', queryVector: queryEmbedding, numCandidates: limit * 10, limit: limit, filter: { projectPath } } }, { $addFields: { score: { $meta: 'vectorSearchScore' } } }, { $match: { score: { $gte: threshold } } } ]; const results = await this.collection.aggregate(pipeline).toArray(); return { content: [ { type: "text", text: `Found ${results.length} results for "${query}":\n` + results.map((r, i) => `${i + 1}. ${r.filePath} (Score: ${(r.score * 100).toFixed(2)}%)`).join('\n') } ] }; } async handleListProjects() { await this.connectMongoDB(); if (!this.collection) { throw new Error("MongoDB not connected"); } const projects = await this.collection.distinct('projectPath'); return { content: [ { type: "text", text: `Indexed projects (${projects.length}):\n` + projects.map((p) => `• ${p}`).join('\n') } ] }; } async handleClearIndex(args) { await this.connectMongoDB(); if (!this.collection) { throw new Error("MongoDB not connected"); } const { projectPath } = args; const result = await this.collection.deleteMany({ projectPath }); return { content: [ { type: "text", text: `✅ Cleared index for ${projectPath}\n` + `Removed ${result.deletedCount} documents from MongoDB Atlas` } ] }; } async handleGetStats(args) { await this.connectMongoDB(); if (!this.collection) { throw new Error("MongoDB not connected"); } const { projectPath } = args; const count = await this.collection.countDocuments({ projectPath }); const dbStats = await this.db.stats(); return { content: [ { type: "text", text: `Project Statistics for ${projectPath}:\n` + `• Documents: ${count}\n` + `• Database Size: ${(dbStats.dataSize / 1024 / 1024).toFixed(2)} MB\n` + `• Storage: MongoDB Atlas\n` + `• Embeddings: Voyage AI (${VOYAGE_MODEL})` } ] }; } async start() { const transport = new StdioServerTransport(); await this.server.connect(transport); console.log("Code Context MCP Server started"); // Log configuration if available if (MONGODB_URI && VOYAGE_API_KEY) { console.log(`MongoDB: ${MONGODB_DATABASE}.${MONGODB_COLLECTION}`); console.log(`Voyage AI: ${VOYAGE_MODEL}`); } else { console.log("Warning: Missing environment variables. Tools will not work until configured."); console.log("Required: MONGODB_URI, VOYAGE_API_KEY"); } } } // Main execution async function main() { const args = process.argv.slice(2); // Handle version flag if (args.includes('--version') || args.includes('-v')) { // Use original console.log for version output originalLog(CONFIG.version); process.exit(0); } if (args.includes('--help') || args.includes('-h')) { console.log(` Code Context MCP Server v${CONFIG.version} Semantic code search powered by MongoDB Atlas and Voyage AI Usage: code-context-mcp Environment Variables: MONGODB_URI MongoDB Atlas connection string (required) VOYAGE_API_KEY Voyage AI API key (required) MONGODB_DATABASE Database name (default: code_context) MONGODB_COLLECTION Collection name (default: embeddings) VOYAGE_MODEL Voyage AI model (default: voyage-3.5) Available Models: voyage-context-3 Contextualized embeddings (+14.24% vs OpenAI) voyage-3-large Highest accuracy (+9.74% vs OpenAI) voyage-3.5 General purpose (default) (+8.26% vs OpenAI) voyage-3.5-lite High throughput (+6.34% vs OpenAI) voyage-code-3 Best for source code Setup: 1. Sign up for MongoDB Atlas: https://www.mongodb.com/cloud/atlas/register 2. Get Voyage AI API key: https://dash.voyageai.com/ 3. Configure in your AI assistant's MCP settings `); process.exit(0); } try { const server = new CodeContextMCPServer(); await server.start(); } catch (error) { console.error("Fatal error:", error); process.exit(1); } } // Handle shutdown process.on('SIGINT', () => { console.log("Shutting down gracefully..."); process.exit(0); }); process.on('SIGTERM', () => { console.log("Shutting down gracefully..."); process.exit(0); }); // Run the server main().catch((error) => { console.error("Fatal error:", error); process.exit(1); }); //# sourceMappingURL=index.js.map