#!/usr/bin/env node
import { config } from "dotenv";
import { fileURLToPath } from "url";
import { dirname, join } from "path";
// Load .env from the package directory
const __filename = fileURLToPath(import.meta.url);
const __dirname = dirname(__filename);
config({ path: join(__dirname, "..", ".env") });
import { Server } from "@modelcontextprotocol/sdk/server/index.js";
import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
import {
CallToolRequestSchema,
ListToolsRequestSchema,
ListResourcesRequestSchema,
ReadResourceRequestSchema,
} from "@modelcontextprotocol/sdk/types.js";
import { initEmbeddings, generateEmbedding } from "./embeddings.js";
import {
initVectorDB,
upsertDocument,
upsertDocuments,
queryDocuments,
deleteDocument,
deleteDocuments,
getDocumentInfo,
} from "./vectordb.js";
// Initialize services from environment variables
const UPSTASH_VECTOR_REST_URL = process.env.UPSTASH_VECTOR_REST_URL;
const UPSTASH_VECTOR_REST_TOKEN = process.env.UPSTASH_VECTOR_REST_TOKEN;
const GOOGLE_AI_API_KEY = process.env.GOOGLE_AI_API_KEY;
if (!UPSTASH_VECTOR_REST_URL || !UPSTASH_VECTOR_REST_TOKEN) {
console.error("Error: UPSTASH_VECTOR_REST_URL and UPSTASH_VECTOR_REST_TOKEN must be set");
process.exit(1);
}
if (!GOOGLE_AI_API_KEY) {
console.error("Error: GOOGLE_AI_API_KEY must be set");
process.exit(1);
}
// Initialize embeddings and vector DB
initEmbeddings(GOOGLE_AI_API_KEY);
initVectorDB(UPSTASH_VECTOR_REST_URL, UPSTASH_VECTOR_REST_TOKEN);
// Create the MCP server
const server = new Server(
{
name: "context-mcp",
version: "1.0.0",
},
{
capabilities: {
tools: {},
resources: {},
},
}
);
// Define available tools
server.setRequestHandler(ListToolsRequestSchema, async () => {
return {
tools: [
{
name: "add_context",
description:
"Add a piece of context/knowledge to the vector database. Use this to store information that can be retrieved later for relevant queries.",
inputSchema: {
type: "object",
properties: {
id: {
type: "string",
description: "Unique identifier for this context entry",
},
content: {
type: "string",
description: "The text content to store and index",
},
metadata: {
type: "object",
description: "Optional metadata to associate with the context (e.g., source, category, timestamp)",
additionalProperties: true,
},
},
required: ["id", "content"],
},
},
{
name: "add_contexts_batch",
description:
"Add multiple context entries to the vector database in a single operation. More efficient for bulk indexing.",
inputSchema: {
type: "object",
properties: {
contexts: {
type: "array",
description: "Array of context entries to add",
items: {
type: "object",
properties: {
id: {
type: "string",
description: "Unique identifier for this context entry",
},
content: {
type: "string",
description: "The text content to store and index",
},
metadata: {
type: "object",
description: "Optional metadata",
additionalProperties: true,
},
},
required: ["id", "content"],
},
},
},
required: ["contexts"],
},
},
{
name: "query_context",
description:
"Search for relevant context based on a natural language query. Returns the most semantically similar stored contexts.",
inputSchema: {
type: "object",
properties: {
query: {
type: "string",
description: "The search query in natural language",
},
topK: {
type: "number",
description: "Number of results to return (default: 5, max: 20)",
default: 5,
},
filter: {
type: "string",
description: "Optional filter expression for metadata (Upstash filter syntax)",
},
},
required: ["query"],
},
},
{
name: "delete_context",
description: "Delete a specific context entry by its ID",
inputSchema: {
type: "object",
properties: {
id: {
type: "string",
description: "The ID of the context to delete",
},
},
required: ["id"],
},
},
{
name: "delete_contexts_batch",
description: "Delete multiple context entries by their IDs",
inputSchema: {
type: "object",
properties: {
ids: {
type: "array",
items: { type: "string" },
description: "Array of context IDs to delete",
},
},
required: ["ids"],
},
},
{
name: "get_stats",
description: "Get statistics about the vector database (number of stored contexts, dimensions)",
inputSchema: {
type: "object",
properties: {},
},
},
],
};
});
// Handle tool calls
server.setRequestHandler(CallToolRequestSchema, async (request) => {
const { name, arguments: args } = request.params;
try {
switch (name) {
case "add_context": {
const { id, content, metadata } = args as {
id: string;
content: string;
metadata?: Record<string, string | number | boolean>;
};
const embedding = await generateEmbedding(content);
await upsertDocument(id, embedding, content, metadata);
return {
content: [
{
type: "text",
text: `Successfully added context with ID: ${id}`,
},
],
};
}
case "add_contexts_batch": {
const { contexts } = args as {
contexts: Array<{
id: string;
content: string;
metadata?: Record<string, string | number | boolean>;
}>;
};
const documentsWithEmbeddings = await Promise.all(
contexts.map(async (ctx) => ({
id: ctx.id,
embedding: await generateEmbedding(ctx.content),
content: ctx.content,
metadata: ctx.metadata,
}))
);
await upsertDocuments(documentsWithEmbeddings);
return {
content: [
{
type: "text",
text: `Successfully added ${contexts.length} context entries`,
},
],
};
}
case "query_context": {
const { query, topK = 5, filter } = args as {
query: string;
topK?: number;
filter?: string;
};
const clampedTopK = Math.min(Math.max(topK, 1), 20);
const queryEmbedding = await generateEmbedding(query);
const results = await queryDocuments(queryEmbedding, clampedTopK, filter);
if (results.length === 0) {
return {
content: [
{
type: "text",
text: "No relevant context found for the query.",
},
],
};
}
const formattedResults = results
.map(
(r, i) =>
`[${i + 1}] (Score: ${r.score.toFixed(4)}, ID: ${r.id})\n${r.content}${
r.metadata && Object.keys(r.metadata).length > 1
? `\nMetadata: ${JSON.stringify(
Object.fromEntries(
Object.entries(r.metadata).filter(([k]) => k !== "content")
)
)}`
: ""
}`
)
.join("\n\n---\n\n");
return {
content: [
{
type: "text",
text: `Found ${results.length} relevant contexts:\n\n${formattedResults}`,
},
],
};
}
case "delete_context": {
const { id } = args as { id: string };
await deleteDocument(id);
return {
content: [
{
type: "text",
text: `Successfully deleted context with ID: ${id}`,
},
],
};
}
case "delete_contexts_batch": {
const { ids } = args as { ids: string[] };
await deleteDocuments(ids);
return {
content: [
{
type: "text",
text: `Successfully deleted ${ids.length} context entries`,
},
],
};
}
case "get_stats": {
const info = await getDocumentInfo();
return {
content: [
{
type: "text",
text: `Vector Database Stats:\n- Total contexts stored: ${info.vectorCount}\n- Embedding dimensions: ${info.dimension}`,
},
],
};
}
default:
return {
content: [
{
type: "text",
text: `Unknown tool: ${name}`,
},
],
isError: true,
};
}
} catch (error) {
const errorMessage = error instanceof Error ? error.message : String(error);
return {
content: [
{
type: "text",
text: `Error executing ${name}: ${errorMessage}`,
},
],
isError: true,
};
}
});
// List resources (optional - expose stats as a resource)
server.setRequestHandler(ListResourcesRequestSchema, async () => {
return {
resources: [
{
uri: "context://stats",
name: "Context Database Statistics",
description: "Current statistics about the context vector database",
mimeType: "application/json",
},
],
};
});
// Read resource
server.setRequestHandler(ReadResourceRequestSchema, async (request) => {
const { uri } = request.params;
if (uri === "context://stats") {
try {
const info = await getDocumentInfo();
return {
contents: [
{
uri,
mimeType: "application/json",
text: JSON.stringify(
{
vectorCount: info.vectorCount,
dimension: info.dimension,
},
null,
2
),
},
],
};
} catch (error) {
const errorMessage = error instanceof Error ? error.message : String(error);
return {
contents: [
{
uri,
mimeType: "text/plain",
text: `Error fetching stats: ${errorMessage}`,
},
],
};
}
}
throw new Error(`Unknown resource: ${uri}`);
});
// Start the server
async function main() {
const transport = new StdioServerTransport();
await server.connect(transport);
console.error("Context MCP server running on stdio");
}
main().catch((error) => {
console.error("Fatal error:", error);
process.exit(1);
});
