mcp_server_local_rag_chromadb_ollama_document_search.ts

1import { Server } from "@modelcontextprotocol/sdk/server/index.js";
2import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
3import {
4  CallToolRequestSchema,
5  ListToolsRequestSchema,
6} from "@modelcontextprotocol/sdk/types.js";
7import { ChromaClient, OpenAIEmbeddingFunction } from "chromadb";
8import { Ollama } from "ollama";
9
10const ollama = new Ollama();
11const client = new ChromaClient();
12const embedder = new OpenAIEmbeddingFunction({
13  openai_api_key: "ollama", // Used for local embedding compatibility
14});
15
16const server = new Server(
17  {
18    name: "mcp-local-rag",
19    version: "0.1.0",
20  },
21  {
22    capabilities: {
23      tools: {},
24    },
25  }
26);
27
28server.setRequestHandler(ListToolsRequestSchema, async () => {
29  return {
30    tools: [
31      {
32        name: "add_document",
33        description: "Add a document to the local RAG database",
34        inputSchema: {
35          type: "object",
36          properties: {
37            content: { type: "string" },
38            metadata: { type: "object" },
39          },
40          required: ["content"],
41        },
42      },
43      {
44        name: "search_documents",
45        description: "Search for documents in the local RAG database",
46        inputSchema: {
47          type: "object",
48          properties: {
49            query: { type: "string" },
50            n_results: { type: "number", default: 5 },
51          },
52          required: ["query"],
53        },
54      },
55    ],
56  };
57});
58
59server.setRequestHandler(CallToolRequestSchema, async (request) => {
60  const collection = await client.getOrCreateCollection({
61    name: "mcp-documents",
62    embeddingFunction: embedder,
63  });
64
65  switch (request.params.name) {
66    case "add_document": {
67      const { content, metadata } = request.params.arguments as any;
68      await collection.add({
69        ids: [Date.now().toString()],
70        metadatas: [metadata || {}],
71        documents: [content],
72      });
73      return { content: [{ type: "text", text: "Document added successfully" }] };
74    }
75    case "search_documents": {
76      const { query, n_results } = request.params.arguments as any;
77      const results = await collection.query({
78        queryTexts: [query],
79        nResults: n_results || 5,
80      });
81      return {
82        content: [{ type: "text", text: JSON.stringify(results.documents) }],
83      };
84    }
85    default:
86      throw new Error("Tool not found");
87  }
88});
89
90async function main() {
91  const transport = new StdioServerTransport();
92  await server.connect(transport);
93}
94
95main().catch(console.error);