@mhalder/qdrant-mcp-server 1.1.1 → 1.3.0

package/src/index.ts

@@ -1,31 +1,43 @@
 #!/usr/bin/env node
 
+import { readFileSync } from "node:fs";
+import { dirname, join } from "node:path";
+import { fileURLToPath } from "node:url";
 import { Server } from "@modelcontextprotocol/sdk/server/index.js";
 import { StdioServerTransport } from "@modelcontextprotocol/sdk/server/stdio.js";
+import { StreamableHTTPServerTransport } from "@modelcontextprotocol/sdk/server/streamableHttp.js";
 import {
   CallToolRequestSchema,
   ListResourcesRequestSchema,
   ListToolsRequestSchema,
   ReadResourceRequestSchema,
 } from "@modelcontextprotocol/sdk/types.js";
-import { QdrantManager } from "./qdrant/client.js";
-import { EmbeddingProviderFactory } from "./embeddings/factory.js";
+import Bottleneck from "bottleneck";
+import express from "express";
 import { z } from "zod";
-import { readFileSync } from "fs";
-import { fileURLToPath } from "url";
-import { dirname, join } from "path";
+import { EmbeddingProviderFactory } from "./embeddings/factory.js";
+import { BM25SparseVectorGenerator } from "./embeddings/sparse.js";
+import { QdrantManager } from "./qdrant/client.js";
 
 // Read package.json for version
 const __dirname = dirname(fileURLToPath(import.meta.url));
-const pkg = JSON.parse(
-  readFileSync(join(__dirname, "../package.json"), "utf-8"),
-);
+const pkg = JSON.parse(readFileSync(join(__dirname, "../package.json"), "utf-8"));
 
 // Validate environment variables
 const QDRANT_URL = process.env.QDRANT_URL || "http://localhost:6333";
-const EMBEDDING_PROVIDER = (
-  process.env.EMBEDDING_PROVIDER || "ollama"
-).toLowerCase();
+const EMBEDDING_PROVIDER = (process.env.EMBEDDING_PROVIDER || "ollama").toLowerCase();
+const TRANSPORT_MODE = (process.env.TRANSPORT_MODE || "stdio").toLowerCase();
+const HTTP_PORT = parseInt(process.env.HTTP_PORT || "3000", 10);
+
+// Validate HTTP_PORT when HTTP mode is selected
+if (TRANSPORT_MODE === "http") {
+  if (Number.isNaN(HTTP_PORT) || HTTP_PORT < 1 || HTTP_PORT > 65535) {
+    console.error(
+      `Error: Invalid HTTP_PORT "${process.env.HTTP_PORT}". Must be a number between 1 and 65535.`
+    );
+    process.exit(1);
+  }
+}
 
 // Check for required API keys based on provider
 if (EMBEDDING_PROVIDER !== "ollama") {
@@ -47,15 +59,13 @@ if (EMBEDDING_PROVIDER !== "ollama") {
       break;
     default:
       console.error(
-        `Error: Unknown embedding provider "${EMBEDDING_PROVIDER}". Supported providers: openai, cohere, voyage, ollama.`,
+        `Error: Unknown embedding provider "${EMBEDDING_PROVIDER}". Supported providers: openai, cohere, voyage, ollama.`
       );
       process.exit(1);
   }
 
   if (!apiKey) {
-    console.error(
-      `Error: ${requiredKeyName} is required for ${EMBEDDING_PROVIDER} provider.`,
-    );
+    console.error(`Error: ${requiredKeyName} is required for ${EMBEDDING_PROVIDER} provider.`);
     process.exit(1);
   }
 }
@@ -64,8 +74,7 @@ if (EMBEDDING_PROVIDER !== "ollama") {
 async function checkOllamaAvailability() {
   if (EMBEDDING_PROVIDER === "ollama") {
     const baseUrl = process.env.EMBEDDING_BASE_URL || "http://localhost:11434";
-    const isLocalhost =
-      baseUrl.includes("localhost") || baseUrl.includes("127.0.0.1");
+    const isLocalhost = baseUrl.includes("localhost") || baseUrl.includes("127.0.0.1");
 
     try {
       const response = await fetch(`${baseUrl}/api/version`);
@@ -78,7 +87,7 @@ async function checkOllamaAvailability() {
       const { models } = await tagsResponse.json();
       const modelName = process.env.EMBEDDING_MODEL || "nomic-embed-text";
       const modelExists = models.some(
-        (m: any) => m.name === modelName || m.name.startsWith(`${modelName}:`),
+        (m: any) => m.name === modelName || m.name.startsWith(`${modelName}:`)
       );
 
       if (!modelExists) {
@@ -141,7 +150,7 @@ const server = new Server(
       tools: {},
       resources: {},
     },
-  },
+  }
 );
 
 // Tool schemas
@@ -151,6 +160,10 @@ const CreateCollectionSchema = z.object({
     .enum(["Cosine", "Euclid", "Dot"])
     .optional()
     .describe("Distance metric (default: Cosine)"),
+  enableHybrid: z
+    .boolean()
+    .optional()
+    .describe("Enable hybrid search with sparse vectors (default: false)"),
 });
 
 const AddDocumentsSchema = z.object({
@@ -158,15 +171,13 @@ const AddDocumentsSchema = z.object({
   documents: z
     .array(
       z.object({
-        id: z
-          .union([z.string(), z.number()])
-          .describe("Unique identifier for the document"),
+        id: z.union([z.string(), z.number()]).describe("Unique identifier for the document"),
         text: z.string().describe("Text content to embed and store"),
         metadata: z
           .record(z.any())
           .optional()
           .describe("Optional metadata to store with the document"),
-      }),
+      })
     )
     .describe("Array of documents to add"),
 });
@@ -174,10 +185,7 @@ const AddDocumentsSchema = z.object({
 const SemanticSearchSchema = z.object({
   collection: z.string().describe("Name of the collection to search"),
   query: z.string().describe("Search query text"),
-  limit: z
-    .number()
-    .optional()
-    .describe("Maximum number of results (default: 5)"),
+  limit: z.number().optional().describe("Maximum number of results (default: 5)"),
   filter: z.record(z.any()).optional().describe("Optional metadata filter"),
 });
 
@@ -191,9 +199,14 @@ const GetCollectionInfoSchema = z.object({
 
 const DeleteDocumentsSchema = z.object({
   collection: z.string().describe("Name of the collection"),
-  ids: z
-    .array(z.union([z.string(), z.number()]))
-    .describe("Array of document IDs to delete"),
+  ids: z.array(z.union([z.string(), z.number()])).describe("Array of document IDs to delete"),
+});
+
+const HybridSearchSchema = z.object({
+  collection: z.string().describe("Name of the collection to search"),
+  query: z.string().describe("Search query text"),
+  limit: z.number().optional().describe("Maximum number of results (default: 5)"),
+  filter: z.record(z.any()).optional().describe("Optional metadata filter"),
 });
 
 // List available tools
@@ -203,7 +216,7 @@ server.setRequestHandler(ListToolsRequestSchema, async () => {
       {
         name: "create_collection",
         description:
-          "Create a new vector collection in Qdrant. The collection will be configured with the embedding provider's dimensions automatically.",
+          "Create a new vector collection in Qdrant. The collection will be configured with the embedding provider's dimensions automatically. Set enableHybrid to true to enable hybrid search combining semantic and keyword search.",
         inputSchema: {
           type: "object",
           properties: {
@@ -216,6 +229,10 @@ server.setRequestHandler(ListToolsRequestSchema, async () => {
               enum: ["Cosine", "Euclid", "Dot"],
               description: "Distance metric (default: Cosine)",
             },
+            enableHybrid: {
+              type: "boolean",
+              description: "Enable hybrid search with sparse vectors (default: false)",
+            },
           },
           required: ["name"],
         },
@@ -323,8 +340,7 @@ server.setRequestHandler(ListToolsRequestSchema, async () => {
       },
       {
         name: "delete_documents",
-        description:
-          "Delete specific documents from a collection by their IDs.",
+        description: "Delete specific documents from a collection by their IDs.",
         inputSchema: {
           type: "object",
           properties: {
@@ -343,6 +359,33 @@ server.setRequestHandler(ListToolsRequestSchema, async () => {
           required: ["collection", "ids"],
         },
       },
+      {
+        name: "hybrid_search",
+        description:
+          "Perform hybrid search combining semantic vector search with keyword search using BM25. This provides better results by combining the strengths of both approaches. The collection must be created with enableHybrid set to true.",
+        inputSchema: {
+          type: "object",
+          properties: {
+            collection: {
+              type: "string",
+              description: "Name of the collection to search",
+            },
+            query: {
+              type: "string",
+              description: "Search query text",
+            },
+            limit: {
+              type: "number",
+              description: "Maximum number of results (default: 5)",
+            },
+            filter: {
+              type: "object",
+              description: "Optional metadata filter",
+            },
+          },
+          required: ["collection", "query"],
+        },
+      },
     ],
   };
 });
@@ -354,14 +397,20 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => {
   try {
     switch (name) {
       case "create_collection": {
-        const { name, distance } = CreateCollectionSchema.parse(args);
+        const { name, distance, enableHybrid } = CreateCollectionSchema.parse(args);
         const vectorSize = embeddings.getDimensions();
-        await qdrant.createCollection(name, vectorSize, distance);
+        await qdrant.createCollection(name, vectorSize, distance, enableHybrid || false);
+
+        let message = `Collection "${name}" created successfully with ${vectorSize} dimensions and ${distance || "Cosine"} distance metric.`;
+        if (enableHybrid) {
+          message += " Hybrid search is enabled for this collection.";
+        }
+
         return {
           content: [
             {
               type: "text",
-              text: `Collection "${name}" created successfully with ${vectorSize} dimensions and ${distance || "Cosine"} distance metric.`,
+              text: message,
             },
           ],
         };
@@ -370,7 +419,7 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => {
       case "add_documents": {
         const { collection, documents } = AddDocumentsSchema.parse(args);
 
-        // Check if collection exists
+        // Check if collection exists and get info
         const exists = await qdrant.collectionExists(collection);
         if (!exists) {
           return {
@@ -384,21 +433,41 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => {
           };
         }
 
+        const collectionInfo = await qdrant.getCollectionInfo(collection);
+
         // Generate embeddings for all documents
         const texts = documents.map((doc) => doc.text);
         const embeddingResults = await embeddings.embedBatch(texts);
 
-        // Prepare points for insertion
-        const points = documents.map((doc, index) => ({
-          id: doc.id,
-          vector: embeddingResults[index].embedding,
-          payload: {
-            text: doc.text,
-            ...doc.metadata,
-          },
-        }));
+        // If hybrid search is enabled, generate sparse vectors and use appropriate method
+        if (collectionInfo.hybridEnabled) {
+          const sparseGenerator = new BM25SparseVectorGenerator();
+
+          // Prepare points with both dense and sparse vectors
+          const points = documents.map((doc, index) => ({
+            id: doc.id,
+            vector: embeddingResults[index].embedding,
+            sparseVector: sparseGenerator.generate(doc.text),
+            payload: {
+              text: doc.text,
+              ...doc.metadata,
+            },
+          }));
+
+          await qdrant.addPointsWithSparse(collection, points);
+        } else {
+          // Standard dense-only vectors
+          const points = documents.map((doc, index) => ({
+            id: doc.id,
+            vector: embeddingResults[index].embedding,
+            payload: {
+              text: doc.text,
+              ...doc.metadata,
+            },
+          }));
 
-        await qdrant.addPoints(collection, points);
+          await qdrant.addPoints(collection, points);
+        }
 
         return {
           content: [
@@ -411,8 +480,7 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => {
       }
 
       case "semantic_search": {
-        const { collection, query, limit, filter } =
-          SemanticSearchSchema.parse(args);
+        const { collection, query, limit, filter } = SemanticSearchSchema.parse(args);
 
         // Check if collection exists
         const exists = await qdrant.collectionExists(collection);
@@ -432,12 +500,7 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => {
         const { embedding } = await embeddings.embed(query);
 
         // Search
-        const results = await qdrant.search(
-          collection,
-          embedding,
-          limit || 5,
-          filter,
-        );
+        const results = await qdrant.search(collection, embedding, limit || 5, filter);
 
         return {
           content: [
@@ -500,6 +563,63 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => {
         };
       }
 
+      case "hybrid_search": {
+        const { collection, query, limit, filter } = HybridSearchSchema.parse(args);
+
+        // Check if collection exists
+        const exists = await qdrant.collectionExists(collection);
+        if (!exists) {
+          return {
+            content: [
+              {
+                type: "text",
+                text: `Error: Collection "${collection}" does not exist.`,
+              },
+            ],
+            isError: true,
+          };
+        }
+
+        // Check if collection has hybrid search enabled
+        const collectionInfo = await qdrant.getCollectionInfo(collection);
+        if (!collectionInfo.hybridEnabled) {
+          return {
+            content: [
+              {
+                type: "text",
+                text: `Error: Collection "${collection}" does not have hybrid search enabled. Create a new collection with enableHybrid set to true.`,
+              },
+            ],
+            isError: true,
+          };
+        }
+
+        // Generate dense embedding for query
+        const { embedding } = await embeddings.embed(query);
+
+        // Generate sparse vector for query
+        const sparseGenerator = new BM25SparseVectorGenerator();
+        const sparseVector = sparseGenerator.generate(query);
+
+        // Perform hybrid search
+        const results = await qdrant.hybridSearch(
+          collection,
+          embedding,
+          sparseVector,
+          limit || 5,
+          filter
+        );
+
+        return {
+          content: [
+            {
+              type: "text",
+              text: JSON.stringify(results, null, 2),
+            },
+          ],
+        };
+      }
+
       default:
         return {
           content: [
@@ -513,8 +633,7 @@ server.setRequestHandler(CallToolRequestSchema, async (request) => {
     }
   } catch (error: any) {
     // Enhanced error details for debugging
-    const errorDetails =
-      error instanceof Error ? error.message : JSON.stringify(error, null, 2);
+    const errorDetails = error instanceof Error ? error.message : JSON.stringify(error, null, 2);
 
     console.error("Tool execution error:", {
       tool: name,
@@ -600,14 +719,229 @@ server.setRequestHandler(ReadResourceRequestSchema, async (request) => {
   };
 });
 
-// Start server
-async function main() {
+// Start server with stdio transport
+async function startStdioServer() {
   await checkOllamaAvailability();
   const transport = new StdioServerTransport();
   await server.connect(transport);
   console.error("Qdrant MCP server running on stdio");
 }
 
+// Start server with HTTP transport
+async function startHttpServer() {
+  await checkOllamaAvailability();
+
+  const app = express();
+  app.use(express.json({ limit: "10mb" }));
+
+  // Configure Express to trust proxy for correct IP detection
+  app.set("trust proxy", true);
+
+  // Rate limiter group: max 100 requests per 15 minutes per IP, max 10 concurrent per IP
+  const rateLimiterGroup = new Bottleneck.Group({
+    reservoir: 100, // initial capacity per IP
+    reservoirRefreshAmount: 100, // refresh back to 100
+    reservoirRefreshInterval: 15 * 60 * 1000, // every 15 minutes
+    maxConcurrent: 10, // max concurrent requests per IP
+  });
+
+  // Periodic cleanup of inactive rate limiters to prevent memory leaks
+  // Track last access time for each IP
+  const ipLastAccess = new Map<string, number>();
+  const INACTIVE_TIMEOUT = 60 * 60 * 1000; // 1 hour
+
+  const cleanupIntervalId = setInterval(() => {
+    const now = Date.now();
+    const keysToDelete: string[] = [];
+
+    ipLastAccess.forEach((lastAccess, ip) => {
+      if (now - lastAccess > INACTIVE_TIMEOUT) {
+        keysToDelete.push(ip);
+      }
+    });
+
+    keysToDelete.forEach((ip) => {
+      rateLimiterGroup.deleteKey(ip);
+      ipLastAccess.delete(ip);
+    });
+
+    if (keysToDelete.length > 0) {
+      console.error(`Cleaned up ${keysToDelete.length} inactive rate limiters`);
+    }
+  }, INACTIVE_TIMEOUT);
+
+  // Rate limiting middleware
+  const rateLimitMiddleware = async (
+    req: express.Request,
+    res: express.Response,
+    next: express.NextFunction
+  ) => {
+    const clientIp = req.ip || req.socket.remoteAddress || "unknown";
+
+    try {
+      // Update last access time for this IP
+      ipLastAccess.set(clientIp, Date.now());
+
+      // Get or create a limiter for this specific IP
+      const limiter = rateLimiterGroup.key(clientIp);
+      await limiter.schedule(() => Promise.resolve());
+      next();
+    } catch (error) {
+      // Differentiate between rate limit errors and unexpected errors
+      if (error instanceof Bottleneck.BottleneckError) {
+        // Rate limit exceeded or Bottleneck operational error
+        console.error(`Rate limit exceeded for IP ${clientIp}:`, error.message);
+      } else {
+        // Unexpected error in rate limiting logic
+        console.error("Unexpected rate limiting error:", error);
+      }
+      res.status(429).json({
+        jsonrpc: "2.0",
+        error: {
+          code: -32000,
+          message: "Too many requests",
+        },
+        id: null,
+      });
+    }
+  };
+
+  // Health check endpoint
+  app.get("/health", (_req, res) => {
+    res.json({
+      status: "ok",
+      mode: TRANSPORT_MODE,
+      version: pkg.version,
+      embedding_provider: EMBEDDING_PROVIDER,
+    });
+  });
+
+  app.post("/mcp", rateLimitMiddleware, async (req, res) => {
+    const REQUEST_TIMEOUT = 60000; // 60 seconds
+    let timeoutId: NodeJS.Timeout | undefined;
+    let isTimedOut = false;
+    let transportClosed = false;
+
+    // Create a new transport for each request in stateless mode.
+    // This prevents request ID collisions when different clients use the same JSON-RPC request IDs.
+    const transport = new StreamableHTTPServerTransport({
+      sessionIdGenerator: undefined, // Stateless mode
+      enableJsonResponse: true,
+    });
+
+    // Helper to safely close transport only once
+    const closeTransport = async () => {
+      if (!transportClosed) {
+        transportClosed = true;
+        await transport.close().catch((e) => console.error("Error closing transport:", e));
+      }
+    };
+
+    try {
+      // Set request timeout
+      timeoutId = setTimeout(async () => {
+        isTimedOut = true;
+        // Close transport on timeout to prevent resource leaks
+        await closeTransport();
+        if (!res.headersSent) {
+          res.status(408).json({
+            jsonrpc: "2.0",
+            error: {
+              code: -32000,
+              message: "Request timeout",
+            },
+            id: null,
+          });
+        }
+      }, REQUEST_TIMEOUT);
+
+      // Clean up transport when response closes
+      res.on("close", async () => {
+        await closeTransport();
+        if (timeoutId) clearTimeout(timeoutId);
+      });
+
+      // Connect the transport to the shared server instance.
+      // In stateless mode, each request gets a new transport connection.
+      await server.connect(transport);
+      await transport.handleRequest(req, res, req.body);
+
+      // Clear timeout immediately on success to prevent race condition
+      if (timeoutId) {
+        clearTimeout(timeoutId);
+        timeoutId = undefined;
+      }
+    } catch (error) {
+      console.error("Error handling MCP request:", error);
+      if (!res.headersSent && !isTimedOut) {
+        res.status(500).json({
+          jsonrpc: "2.0",
+          error: {
+            code: -32603,
+            message: "Internal server error",
+          },
+          id: null,
+        });
+      }
+    } finally {
+      if (timeoutId) clearTimeout(timeoutId);
+      // Ensure transport is closed even if an error occurs
+      await closeTransport();
+    }
+  });
+
+  const httpServer = app
+    .listen(HTTP_PORT, () => {
+      console.error(`Qdrant MCP server running on http://localhost:${HTTP_PORT}/mcp`);
+    })
+    .on("error", (error) => {
+      console.error("HTTP server error:", error);
+      process.exit(1);
+    });
+
+  // Graceful shutdown handling
+  let isShuttingDown = false;
+
+  const shutdown = () => {
+    if (isShuttingDown) return;
+    isShuttingDown = true;
+
+    console.error("Shutdown signal received, closing HTTP server gracefully...");
+
+    // Clear the cleanup interval to allow graceful shutdown
+    clearInterval(cleanupIntervalId);
+
+    // Force shutdown after 10 seconds
+    const forceTimeout = setTimeout(() => {
+      console.error("Forcing shutdown after timeout");
+      process.exit(1);
+    }, 10000);
+
+    httpServer.close(() => {
+      clearTimeout(forceTimeout);
+      console.error("HTTP server closed");
+      process.exit(0);
+    });
+  };
+
+  process.on("SIGTERM", shutdown);
+  process.on("SIGINT", shutdown);
+}
+
+// Main entry point
+async function main() {
+  if (TRANSPORT_MODE === "http") {
+    await startHttpServer();
+  } else if (TRANSPORT_MODE === "stdio") {
+    await startStdioServer();
+  } else {
+    console.error(
+      `Error: Invalid TRANSPORT_MODE "${TRANSPORT_MODE}". Supported modes: stdio, http.`
+    );
+    process.exit(1);
+  }
+}
+
 main().catch((error) => {
   console.error("Fatal error:", error);
   process.exit(1);