@mhalder/qdrant-mcp-server 1.1.0

package/src/embeddings/openai.ts

@@ -0,0 +1,158 @@
+import OpenAI from "openai";
+import Bottleneck from "bottleneck";
+import {
+  EmbeddingProvider,
+  EmbeddingResult,
+  RateLimitConfig,
+  ProviderConfig,
+} from "./base.js";
+
+interface OpenAIError {
+  status?: number;
+  code?: string;
+  message?: string;
+  headers?: Record<string, string>;
+  response?: {
+    headers?: Record<string, string>;
+  };
+}
+
+export class OpenAIEmbeddings implements EmbeddingProvider {
+  private client: OpenAI;
+  private model: string;
+  private dimensions: number;
+  private limiter: Bottleneck;
+  private retryAttempts: number;
+  private retryDelayMs: number;
+
+  constructor(
+    apiKey: string,
+    model: string = "text-embedding-3-small",
+    dimensions?: number,
+    rateLimitConfig?: RateLimitConfig,
+  ) {
+    this.client = new OpenAI({ apiKey });
+    this.model = model;
+
+    // Default dimensions for different models
+    const defaultDimensions: Record<string, number> = {
+      "text-embedding-3-small": 1536,
+      "text-embedding-3-large": 3072,
+      "text-embedding-ada-002": 1536,
+    };
+
+    this.dimensions = dimensions || defaultDimensions[model] || 1536;
+
+    // Rate limiting configuration
+    const maxRequestsPerMinute = rateLimitConfig?.maxRequestsPerMinute || 3500;
+    this.retryAttempts = rateLimitConfig?.retryAttempts || 3;
+    this.retryDelayMs = rateLimitConfig?.retryDelayMs || 1000;
+
+    // Initialize bottleneck limiter
+    // Uses reservoir (token bucket) pattern for burst handling with per-minute refresh
+    // Note: Using both reservoir and minTime provides defense in depth but may be
+    // more conservative than necessary. Future optimization could use reservoir-only
+    // for better burst handling or minTime-only for simpler even distribution.
+    this.limiter = new Bottleneck({
+      reservoir: maxRequestsPerMinute,
+      reservoirRefreshAmount: maxRequestsPerMinute,
+      reservoirRefreshInterval: 60 * 1000, // 1 minute
+      maxConcurrent: 10,
+      minTime: Math.floor((60 * 1000) / maxRequestsPerMinute),
+    });
+  }
+
+  private async retryWithBackoff<T>(
+    fn: () => Promise<T>,
+    attempt: number = 0,
+  ): Promise<T> {
+    try {
+      return await fn();
+    } catch (error: unknown) {
+      const apiError = error as OpenAIError;
+      const isRateLimitError =
+        apiError?.status === 429 ||
+        apiError?.code === "rate_limit_exceeded" ||
+        apiError?.message?.toLowerCase().includes("rate limit");
+
+      if (isRateLimitError && attempt < this.retryAttempts) {
+        // Check for Retry-After header (different HTTP clients may nest differently)
+        const retryAfter =
+          apiError?.response?.headers?.["retry-after"] ||
+          apiError?.headers?.["retry-after"];
+        let delayMs: number;
+
+        if (retryAfter) {
+          // Use Retry-After header if available (in seconds)
+          const parsed = parseInt(retryAfter, 10);
+          delayMs =
+            !isNaN(parsed) && parsed > 0
+              ? parsed * 1000
+              : this.retryDelayMs * Math.pow(2, attempt);
+        } else {
+          // Exponential backoff: 1s, 2s, 4s, 8s...
+          delayMs = this.retryDelayMs * Math.pow(2, attempt);
+        }
+
+        const waitTimeSeconds = (delayMs / 1000).toFixed(1);
+        console.error(
+          `Rate limit reached. Retrying in ${waitTimeSeconds}s (attempt ${attempt + 1}/${this.retryAttempts})...`,
+        );
+
+        await new Promise((resolve) => setTimeout(resolve, delayMs));
+        return this.retryWithBackoff(fn, attempt + 1);
+      }
+
+      // If not a rate limit error or max retries exceeded, throw
+      if (isRateLimitError) {
+        throw new Error(
+          `OpenAI API rate limit exceeded after ${this.retryAttempts} retry attempts. Please try again later or reduce request frequency.`,
+        );
+      }
+
+      throw error;
+    }
+  }
+
+  async embed(text: string): Promise<EmbeddingResult> {
+    return this.limiter.schedule(() =>
+      this.retryWithBackoff(async () => {
+        const response = await this.client.embeddings.create({
+          model: this.model,
+          input: text,
+          dimensions: this.dimensions,
+        });
+
+        return {
+          embedding: response.data[0].embedding,
+          dimensions: this.dimensions,
+        };
+      }),
+    );
+  }
+
+  async embedBatch(texts: string[]): Promise<EmbeddingResult[]> {
+    return this.limiter.schedule(() =>
+      this.retryWithBackoff(async () => {
+        const response = await this.client.embeddings.create({
+          model: this.model,
+          input: texts,
+          dimensions: this.dimensions,
+        });
+
+        return response.data.map((item) => ({
+          embedding: item.embedding,
+          dimensions: this.dimensions,
+        }));
+      }),
+    );
+  }
+
+  getDimensions(): number {
+    return this.dimensions;
+  }
+
+  getModel(): string {
+    return this.model;
+  }
+}

package/src/embeddings/voyage.test.ts

@@ -0,0 +1,520 @@
+import { describe, it, expect, vi, beforeEach } from "vitest";
+import { VoyageEmbeddings } from "./voyage.js";
+
+// Mock fetch globally
+global.fetch = vi.fn();
+
+describe("VoyageEmbeddings", () => {
+  let embeddings: VoyageEmbeddings;
+  let mockFetch: any;
+
+  beforeEach(() => {
+    mockFetch = global.fetch as any;
+    mockFetch.mockReset();
+
+    embeddings = new VoyageEmbeddings("test-api-key");
+  });
+
+  describe("constructor", () => {
+    it("should use default model and dimensions", () => {
+      expect(embeddings.getModel()).toBe("voyage-2");
+      expect(embeddings.getDimensions()).toBe(1024);
+    });
+
+    it("should use custom model", () => {
+      const customEmbeddings = new VoyageEmbeddings(
+        "test-api-key",
+        "voyage-large-2",
+      );
+      expect(customEmbeddings.getModel()).toBe("voyage-large-2");
+      expect(customEmbeddings.getDimensions()).toBe(1536);
+    });
+
+    it("should use custom dimensions", () => {
+      const customEmbeddings = new VoyageEmbeddings(
+        "test-api-key",
+        "voyage-2",
+        512,
+      );
+      expect(customEmbeddings.getDimensions()).toBe(512);
+    });
+
+    it("should use default base URL", () => {
+      const defaultEmbeddings = new VoyageEmbeddings("test-api-key");
+      expect(defaultEmbeddings).toBeDefined();
+    });
+
+    it("should use custom base URL", () => {
+      const customEmbeddings = new VoyageEmbeddings(
+        "test-api-key",
+        "voyage-2",
+        undefined,
+        undefined,
+        "https://custom.voyage.com",
+      );
+      expect(customEmbeddings).toBeDefined();
+    });
+
+    it("should default to 1024 for unknown models", () => {
+      const unknownEmbeddings = new VoyageEmbeddings(
+        "test-api-key",
+        "custom-model",
+      );
+      expect(unknownEmbeddings.getDimensions()).toBe(1024);
+    });
+
+    it("should accept custom input type", () => {
+      const queryEmbeddings = new VoyageEmbeddings(
+        "test-api-key",
+        "voyage-2",
+        undefined,
+        undefined,
+        undefined,
+        "query",
+      );
+      expect(queryEmbeddings).toBeInstanceOf(VoyageEmbeddings);
+    });
+  });
+
+  describe("embed", () => {
+    it("should generate embedding for single text", async () => {
+      const mockEmbedding = Array(1024)
+        .fill(0)
+        .map((_, i) => i * 0.001);
+      mockFetch.mockResolvedValue({
+        ok: true,
+        json: async () => ({
+          data: [{ embedding: mockEmbedding }],
+          model: "voyage-2",
+          usage: { total_tokens: 10 },
+        }),
+      });
+
+      const result = await embeddings.embed("test text");
+
+      expect(result).toEqual({
+        embedding: mockEmbedding,
+        dimensions: 1024,
+      });
+      expect(mockFetch).toHaveBeenCalledWith(
+        "https://api.voyageai.com/v1/embeddings",
+        {
+          method: "POST",
+          headers: {
+            "Content-Type": "application/json",
+            Authorization: "Bearer test-api-key",
+          },
+          body: JSON.stringify({
+            input: ["test text"],
+            model: "voyage-2",
+          }),
+        },
+      );
+    });
+
+    it("should include input_type when specified", async () => {
+      const queryEmbeddings = new VoyageEmbeddings(
+        "test-api-key",
+        "voyage-2",
+        undefined,
+        undefined,
+        undefined,
+        "query",
+      );
+
+      const mockEmbedding = Array(1024).fill(0.5);
+      mockFetch.mockResolvedValue({
+        ok: true,
+        json: async () => ({
+          data: [{ embedding: mockEmbedding }],
+          model: "voyage-2",
+          usage: { total_tokens: 10 },
+        }),
+      });
+
+      await queryEmbeddings.embed("test text");
+
+      const callBody = JSON.parse(mockFetch.mock.calls[0][1].body);
+      expect(callBody.input_type).toBe("query");
+    });
+
+    it("should handle long text", async () => {
+      const longText = "word ".repeat(1000);
+      const mockEmbedding = Array(1024).fill(0.5);
+      mockFetch.mockResolvedValue({
+        ok: true,
+        json: async () => ({
+          data: [{ embedding: mockEmbedding }],
+          model: "voyage-2",
+          usage: { total_tokens: 1000 },
+        }),
+      });
+
+      const result = await embeddings.embed(longText);
+
+      expect(result.embedding).toEqual(mockEmbedding);
+    });
+
+    it("should use custom base URL", async () => {
+      const customEmbeddings = new VoyageEmbeddings(
+        "test-api-key",
+        "voyage-2",
+        undefined,
+        undefined,
+        "https://custom.voyage.com/v1",
+      );
+
+      const mockEmbedding = Array(1024).fill(0.1);
+      mockFetch.mockResolvedValue({
+        ok: true,
+        json: async () => ({
+          data: [{ embedding: mockEmbedding }],
+          model: "voyage-2",
+          usage: { total_tokens: 5 },
+        }),
+      });
+
+      await customEmbeddings.embed("test");
+
+      expect(mockFetch).toHaveBeenCalledWith(
+        "https://custom.voyage.com/v1/embeddings",
+        expect.any(Object),
+      );
+    });
+
+    it("should throw error if no embedding returned", async () => {
+      mockFetch.mockResolvedValue({
+        ok: true,
+        json: async () => ({
+          data: [],
+          model: "voyage-2",
+          usage: { total_tokens: 0 },
+        }),
+      });
+
+      await expect(embeddings.embed("test")).rejects.toThrow(
+        "No embedding returned from Voyage AI API",
+      );
+    });
+
+    it("should handle API errors", async () => {
+      mockFetch.mockResolvedValue({
+        ok: false,
+        status: 401,
+        text: async () => "Unauthorized: Invalid API key",
+      });
+
+      await expect(embeddings.embed("test")).rejects.toThrow();
+    });
+
+    it("should propagate network errors", async () => {
+      mockFetch.mockRejectedValue(new Error("Network Error"));
+
+      await expect(embeddings.embed("test")).rejects.toThrow("Network Error");
+    });
+  });
+
+  describe("embedBatch", () => {
+    it("should generate embeddings for multiple texts", async () => {
+      const mockEmbeddings = [
+        Array(1024).fill(0.1),
+        Array(1024).fill(0.2),
+        Array(1024).fill(0.3),
+      ];
+      mockFetch.mockResolvedValue({
+        ok: true,
+        json: async () => ({
+          data: mockEmbeddings.map((embedding) => ({ embedding })),
+          model: "voyage-2",
+          usage: { total_tokens: 30 },
+        }),
+      });
+
+      const texts = ["text1", "text2", "text3"];
+      const results = await embeddings.embedBatch(texts);
+
+      expect(results).toEqual([
+        { embedding: mockEmbeddings[0], dimensions: 1024 },
+        { embedding: mockEmbeddings[1], dimensions: 1024 },
+        { embedding: mockEmbeddings[2], dimensions: 1024 },
+      ]);
+
+      const callBody = JSON.parse(mockFetch.mock.calls[0][1].body);
+      expect(callBody.input).toEqual(texts);
+    });
+
+    it("should handle empty batch", async () => {
+      mockFetch.mockResolvedValue({
+        ok: true,
+        json: async () => ({
+          data: [],
+          model: "voyage-2",
+          usage: { total_tokens: 0 },
+        }),
+      });
+
+      const results = await embeddings.embedBatch([]);
+
+      expect(results).toEqual([]);
+    });
+
+    it("should handle single item in batch", async () => {
+      const mockEmbedding = Array(1024).fill(0.5);
+      mockFetch.mockResolvedValue({
+        ok: true,
+        json: async () => ({
+          data: [{ embedding: mockEmbedding }],
+          model: "voyage-2",
+          usage: { total_tokens: 10 },
+        }),
+      });
+
+      const results = await embeddings.embedBatch(["single text"]);
+
+      expect(results).toHaveLength(1);
+      expect(results[0].embedding).toEqual(mockEmbedding);
+    });
+
+    it("should handle large batches", async () => {
+      const batchSize = 100;
+      const mockEmbeddings = Array(batchSize)
+        .fill(null)
+        .map(() => Array(1024).fill(Math.random()));
+
+      mockFetch.mockResolvedValue({
+        ok: true,
+        json: async () => ({
+          data: mockEmbeddings.map((embedding) => ({ embedding })),
+          model: "voyage-2",
+          usage: { total_tokens: batchSize * 10 },
+        }),
+      });
+
+      const texts = Array(batchSize)
+        .fill(null)
+        .map((_, i) => `text ${i}`);
+      const results = await embeddings.embedBatch(texts);
+
+      expect(results).toHaveLength(batchSize);
+    });
+
+    it("should throw error if no embeddings returned", async () => {
+      mockFetch.mockResolvedValue({
+        ok: true,
+        json: async () => ({
+          model: "voyage-2",
+          usage: { total_tokens: 0 },
+        }),
+      });
+
+      await expect(embeddings.embedBatch(["text1"])).rejects.toThrow(
+        "No embeddings returned from Voyage AI API",
+      );
+    });
+
+    it("should propagate errors in batch", async () => {
+      mockFetch.mockRejectedValue(new Error("Batch API Error"));
+
+      await expect(embeddings.embedBatch(["text1", "text2"])).rejects.toThrow(
+        "Batch API Error",
+      );
+    });
+  });
+
+  describe("getDimensions", () => {
+    it("should return configured dimensions", () => {
+      expect(embeddings.getDimensions()).toBe(1024);
+    });
+
+    it("should return custom dimensions", () => {
+      const customEmbeddings = new VoyageEmbeddings(
+        "test-api-key",
+        "voyage-2",
+        512,
+      );
+      expect(customEmbeddings.getDimensions()).toBe(512);
+    });
+  });
+
+  describe("getModel", () => {
+    it("should return configured model", () => {
+      expect(embeddings.getModel()).toBe("voyage-2");
+    });
+
+    it("should return custom model", () => {
+      const customEmbeddings = new VoyageEmbeddings(
+        "test-api-key",
+        "voyage-large-2",
+      );
+      expect(customEmbeddings.getModel()).toBe("voyage-large-2");
+    });
+  });
+
+  describe("rate limiting", () => {
+    it("should retry on rate limit error (429 status)", async () => {
+      const mockEmbedding = Array(1024).fill(0.5);
+
+      mockFetch
+        .mockResolvedValueOnce({
+          ok: false,
+          status: 429,
+          text: async () => "Rate limit exceeded",
+        })
+        .mockResolvedValueOnce({
+          ok: false,
+          status: 429,
+          text: async () => "Rate limit exceeded",
+        })
+        .mockResolvedValue({
+          ok: true,
+          json: async () => ({
+            data: [{ embedding: mockEmbedding }],
+            model: "voyage-2",
+            usage: { total_tokens: 10 },
+          }),
+        });
+
+      const result = await embeddings.embed("test text");
+
+      expect(result.embedding).toEqual(mockEmbedding);
+      expect(mockFetch).toHaveBeenCalledTimes(3);
+    });
+
+    it("should retry on rate limit message", async () => {
+      const mockEmbedding = Array(1024).fill(0.5);
+
+      mockFetch
+        .mockRejectedValueOnce({
+          message: "You have exceeded the rate limit",
+        })
+        .mockResolvedValue({
+          ok: true,
+          json: async () => ({
+            data: [{ embedding: mockEmbedding }],
+            model: "voyage-2",
+            usage: { total_tokens: 10 },
+          }),
+        });
+
+      const result = await embeddings.embed("test text");
+
+      expect(result.embedding).toEqual(mockEmbedding);
+      expect(mockFetch).toHaveBeenCalledTimes(2);
+    });
+
+    it("should use exponential backoff", async () => {
+      const rateLimitEmbeddings = new VoyageEmbeddings(
+        "test-api-key",
+        "voyage-2",
+        undefined,
+        {
+          retryAttempts: 3,
+          retryDelayMs: 100,
+        },
+      );
+
+      const mockEmbedding = Array(1024).fill(0.5);
+
+      mockFetch
+        .mockResolvedValueOnce({
+          ok: false,
+          status: 429,
+          text: async () => "Rate limit",
+        })
+        .mockResolvedValueOnce({
+          ok: false,
+          status: 429,
+          text: async () => "Rate limit",
+        })
+        .mockResolvedValue({
+          ok: true,
+          json: async () => ({
+            data: [{ embedding: mockEmbedding }],
+            model: "voyage-2",
+            usage: { total_tokens: 10 },
+          }),
+        });
+
+      const startTime = Date.now();
+      await rateLimitEmbeddings.embed("test text");
+      const duration = Date.now() - startTime;
+
+      // Should wait: 100ms (first retry) + 200ms (second retry) = 300ms
+      expect(duration).toBeGreaterThanOrEqual(250);
+    });
+
+    it("should throw error after max retries exceeded", async () => {
+      const rateLimitEmbeddings = new VoyageEmbeddings(
+        "test-api-key",
+        "voyage-2",
+        undefined,
+        {
+          retryAttempts: 2,
+          retryDelayMs: 100,
+        },
+      );
+
+      mockFetch.mockResolvedValue({
+        ok: false,
+        status: 429,
+        text: async () => "Rate limit exceeded",
+      });
+
+      await expect(rateLimitEmbeddings.embed("test text")).rejects.toThrow(
+        "Voyage AI API rate limit exceeded after 2 retry attempts",
+      );
+
+      expect(mockFetch).toHaveBeenCalledTimes(3);
+    });
+
+    it("should handle rate limit errors in batch operations", async () => {
+      const mockEmbeddings = [Array(1024).fill(0.1), Array(1024).fill(0.2)];
+
+      mockFetch
+        .mockResolvedValueOnce({
+          ok: false,
+          status: 429,
+          text: async () => "Rate limit",
+        })
+        .mockResolvedValue({
+          ok: true,
+          json: async () => ({
+            data: mockEmbeddings.map((embedding) => ({ embedding })),
+            model: "voyage-2",
+            usage: { total_tokens: 20 },
+          }),
+        });
+
+      const results = await embeddings.embedBatch(["text1", "text2"]);
+
+      expect(results).toHaveLength(2);
+      expect(mockFetch).toHaveBeenCalledTimes(2);
+    });
+
+    it("should not retry on non-rate-limit errors", async () => {
+      mockFetch.mockResolvedValue({
+        ok: false,
+        status: 401,
+        text: async () => "Unauthorized",
+      });
+
+      await expect(embeddings.embed("test text")).rejects.toThrow();
+      expect(mockFetch).toHaveBeenCalledTimes(1);
+    });
+
+    it("should accept custom rate limit configuration", () => {
+      const customEmbeddings = new VoyageEmbeddings(
+        "test-api-key",
+        "voyage-2",
+        undefined,
+        {
+          maxRequestsPerMinute: 500,
+          retryAttempts: 5,
+          retryDelayMs: 2000,
+        },
+      );
+
+      expect(customEmbeddings).toBeDefined();
+    });
+  });
+});