@cue-dev/retrieval-core 0.1.3

package/test/embedding-provider.test.ts

@@ -0,0 +1,570 @@
+import { afterEach, describe, expect, it, vi } from "vitest";
+import {
+  LocalProviderRequestLimiter,
+  OpenAICompatibleEmbeddingProvider,
+  OpenAICompatibleRerankerProvider,
+  ProviderRateLimitExceededError,
+  RetrievalError
+} from "../src/index.js";
+
+describe("provider request limiter", () => {
+  it("waits for token refill in local mode", async () => {
+    let now = 0;
+    const limiter = new LocalProviderRequestLimiter({
+      now: () => now,
+      sleeper: async (ms) => {
+        now += ms;
+      }
+    });
+
+    await limiter.acquire({
+      scope: "provider:test|credential:a",
+      max_requests_per_minute: 1,
+      max_wait_ms: 0
+    });
+
+    const acquired = await limiter.acquire({
+      scope: "provider:test|credential:a",
+      max_requests_per_minute: 1,
+      max_wait_ms: 60_000
+    });
+    expect(acquired.wait_ms).toBe(60_000);
+  });
+
+  it("fails when wait budget is exhausted", async () => {
+    let now = 0;
+    const limiter = new LocalProviderRequestLimiter({
+      now: () => now,
+      sleeper: async (ms) => {
+        now += ms;
+      }
+    });
+
+    await limiter.acquire({
+      scope: "provider:test|credential:a",
+      max_requests_per_minute: 1,
+      max_wait_ms: 0
+    });
+
+    await expect(
+      limiter.acquire({
+        scope: "provider:test|credential:a",
+        max_requests_per_minute: 1,
+        max_wait_ms: 1_000
+      })
+    ).rejects.toBeInstanceOf(ProviderRateLimitExceededError);
+  });
+});
+
+describe("openai-compatible embedding provider", () => {
+  const originalFetch = globalThis.fetch;
+
+  afterEach(() => {
+    globalThis.fetch = originalFetch;
+    vi.restoreAllMocks();
+  });
+
+  it("retries on 5xx and succeeds", async () => {
+    let calls = 0;
+    globalThis.fetch = vi.fn(async () => {
+      calls += 1;
+      if (calls === 1) {
+        return new Response(JSON.stringify({ error: "temporary" }), { status: 503 });
+      }
+      return new Response(
+        JSON.stringify({
+          data: [{ index: 0, embedding: [0.1, 0.2, 0.3] }]
+        }),
+        { status: 200 }
+      );
+    }) as typeof fetch;
+
+    const provider = new OpenAICompatibleEmbeddingProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      model: "Qwen/Qwen3-Embedding-4B",
+      dimensions: 3,
+      timeout_ms: 200,
+      batch_size: 16,
+      max_retries: 2
+    });
+
+    const vectors = await provider.embed({
+      texts: ["hello"],
+      purpose: "query"
+    });
+
+    expect(vectors).toEqual([[0.1, 0.2, 0.3]]);
+    expect(calls).toBe(2);
+  });
+
+  it("does not retry auth/config 4xx failures", async () => {
+    let calls = 0;
+    globalThis.fetch = vi.fn(async () => {
+      calls += 1;
+      return new Response(JSON.stringify({ error: "bad key" }), { status: 401 });
+    }) as typeof fetch;
+
+    const provider = new OpenAICompatibleEmbeddingProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      dimensions: 3
+    });
+
+    await expect(
+      provider.embed({
+        texts: ["hello"],
+        purpose: "query"
+      })
+    ).rejects.toMatchObject({
+      code: "UPSTREAM_FAILURE"
+    } satisfies Partial<RetrievalError>);
+    expect(calls).toBe(1);
+  });
+
+  it("retries transient 403 responses and succeeds", async () => {
+    let calls = 0;
+    globalThis.fetch = vi.fn(async () => {
+      calls += 1;
+      if (calls === 1) {
+        return new Response(
+          JSON.stringify({
+            error: {
+              message: "openai_error",
+              type: "bad_response_status_code"
+            }
+          }),
+          { status: 403 }
+        );
+      }
+      return new Response(
+        JSON.stringify({
+          data: [{ index: 0, embedding: [0.1, 0.2, 0.3] }]
+        }),
+        { status: 200 }
+      );
+    }) as typeof fetch;
+
+    const provider = new OpenAICompatibleEmbeddingProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      dimensions: 3,
+      max_retries: 1
+    });
+
+    const vectors = await provider.embed({
+      texts: ["hello"],
+      purpose: "query"
+    });
+
+    expect(vectors).toEqual([[0.1, 0.2, 0.3]]);
+    expect(calls).toBe(2);
+  });
+
+  it("allows transient 403 retries beyond generic max_retries", async () => {
+    let calls = 0;
+    globalThis.fetch = vi.fn(async () => {
+      calls += 1;
+      if (calls <= 2) {
+        return new Response(
+          JSON.stringify({
+            error: {
+              message: "openai_error",
+              type: "bad_response_status_code"
+            }
+          }),
+          { status: 403 }
+        );
+      }
+      return new Response(
+        JSON.stringify({
+          data: [{ index: 0, embedding: [0.1, 0.2, 0.3] }]
+        }),
+        { status: 200 }
+      );
+    }) as typeof fetch;
+
+    const provider = new OpenAICompatibleEmbeddingProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      dimensions: 3,
+      max_retries: 0
+    });
+
+    const vectors = await provider.embed({
+      texts: ["hello"],
+      purpose: "query"
+    });
+
+    expect(vectors).toEqual([[0.1, 0.2, 0.3]]);
+    expect(calls).toBe(3);
+  });
+
+  it("does not retry explicit invalid api key 403 failures", async () => {
+    let calls = 0;
+    globalThis.fetch = vi.fn(async () => {
+      calls += 1;
+      return new Response(
+        JSON.stringify({
+          error: {
+            message: "Invalid API key provided"
+          }
+        }),
+        { status: 403 }
+      );
+    }) as typeof fetch;
+
+    const provider = new OpenAICompatibleEmbeddingProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      dimensions: 3,
+      max_retries: 4
+    });
+
+    await expect(
+      provider.embed({
+        texts: ["hello"],
+        purpose: "query"
+      })
+    ).rejects.toMatchObject({
+      code: "UPSTREAM_FAILURE"
+    } satisfies Partial<RetrievalError>);
+    expect(calls).toBe(1);
+  });
+
+  it("does not retry ambiguous 403 failures without transient signals", async () => {
+    let calls = 0;
+    globalThis.fetch = vi.fn(async () => {
+      calls += 1;
+      return new Response(
+        JSON.stringify({
+          error: {
+            message: "forbidden"
+          }
+        }),
+        { status: 403 }
+      );
+    }) as typeof fetch;
+
+    const provider = new OpenAICompatibleEmbeddingProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      dimensions: 3,
+      max_retries: 4
+    });
+
+    await expect(
+      provider.embed({
+        texts: ["hello"],
+        purpose: "query"
+      })
+    ).rejects.toMatchObject({
+      code: "UPSTREAM_FAILURE"
+    } satisfies Partial<RetrievalError>);
+    expect(calls).toBe(1);
+  });
+
+  it("retries 403 failures when retry-after is present", async () => {
+    let calls = 0;
+    globalThis.fetch = vi.fn(async () => {
+      calls += 1;
+      if (calls === 1) {
+        return new Response(
+          JSON.stringify({
+            error: {
+              message: "forbidden"
+            }
+          }),
+          {
+            status: 403,
+            headers: {
+              "retry-after": "0"
+            }
+          }
+        );
+      }
+      return new Response(
+        JSON.stringify({
+          data: [{ index: 0, embedding: [0.1, 0.2, 0.3] }]
+        }),
+        { status: 200 }
+      );
+    }) as typeof fetch;
+
+    const provider = new OpenAICompatibleEmbeddingProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      dimensions: 3,
+      max_retries: 1
+    });
+
+    const vectors = await provider.embed({
+      texts: ["hello"],
+      purpose: "query"
+    });
+
+    expect(vectors).toEqual([[0.1, 0.2, 0.3]]);
+    expect(calls).toBe(2);
+  });
+
+  it("rejects non-numeric vectors", async () => {
+    globalThis.fetch = vi.fn(async () => {
+      return new Response(
+        JSON.stringify({
+          data: [{ index: 0, embedding: [0.1, "bad", 0.3] }]
+        }),
+        { status: 200 }
+      );
+    }) as typeof fetch;
+
+    const provider = new OpenAICompatibleEmbeddingProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      dimensions: 3
+    });
+
+    await expect(
+      provider.embed({
+        texts: ["hello"],
+        purpose: "query"
+      })
+    ).rejects.toMatchObject({
+      code: "UPSTREAM_FAILURE"
+    } satisfies Partial<RetrievalError>);
+  });
+
+  it("rejects vectors with unexpected dimensions", async () => {
+    globalThis.fetch = vi.fn(async () => {
+      return new Response(
+        JSON.stringify({
+          data: [{ index: 0, embedding: [0.1, 0.2] }]
+        }),
+        { status: 200 }
+      );
+    }) as typeof fetch;
+
+    const provider = new OpenAICompatibleEmbeddingProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      dimensions: 3
+    });
+
+    await expect(
+      provider.embed({
+        texts: ["hello"],
+        purpose: "query"
+      })
+    ).rejects.toMatchObject({
+      code: "UPSTREAM_FAILURE"
+    } satisfies Partial<RetrievalError>);
+  });
+
+  it("maps local limiter saturation to RATE_LIMITED before calling upstream", async () => {
+    let calls = 0;
+    globalThis.fetch = vi.fn(async () => {
+      calls += 1;
+      return new Response(
+        JSON.stringify({
+          data: [{ index: 0, embedding: [0.1, 0.2, 0.3] }]
+        }),
+        { status: 200 }
+      );
+    }) as typeof fetch;
+
+    const limiter = new LocalProviderRequestLimiter();
+    const provider = new OpenAICompatibleEmbeddingProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      dimensions: 3,
+      request_limiter: limiter,
+      max_requests_per_minute: 1,
+      query_max_wait_ms: 0
+    });
+
+    await provider.embed({
+      texts: ["hello"],
+      purpose: "query"
+    });
+
+    await expect(
+      provider.embed({
+        texts: ["world"],
+        purpose: "query"
+      })
+    ).rejects.toMatchObject({
+      code: "RATE_LIMITED"
+    } satisfies Partial<RetrievalError>);
+    expect(calls).toBe(1);
+  });
+
+  it("retries index embeddings after transient local limiter saturation", async () => {
+    let calls = 0;
+    globalThis.fetch = vi.fn(async () => {
+      calls += 1;
+      return new Response(
+        JSON.stringify({
+          data: [{ index: 0, embedding: [0.1, 0.2, 0.3] }]
+        }),
+        { status: 200 }
+      );
+    }) as typeof fetch;
+
+    const requestLimiter = {
+      mode: "local" as const,
+      acquire: vi
+        .fn()
+        .mockRejectedValueOnce(new ProviderRateLimitExceededError("transient local limit", 1))
+        .mockResolvedValue({ wait_ms: 0 })
+    };
+
+    const provider = new OpenAICompatibleEmbeddingProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      dimensions: 3,
+      request_limiter: requestLimiter,
+      max_requests_per_minute: 1,
+      index_max_wait_ms: 0
+    });
+
+    const vectors = await provider.embed({
+      texts: ["hello"],
+      purpose: "index"
+    });
+
+    expect(vectors).toEqual([[0.1, 0.2, 0.3]]);
+    expect(requestLimiter.acquire).toHaveBeenCalledTimes(2);
+    expect(calls).toBe(1);
+  });
+
+  it("keeps query limiter saturation fail-fast", async () => {
+    let calls = 0;
+    globalThis.fetch = vi.fn(async () => {
+      calls += 1;
+      return new Response(
+        JSON.stringify({
+          data: [{ index: 0, embedding: [0.1, 0.2, 0.3] }]
+        }),
+        { status: 200 }
+      );
+    }) as typeof fetch;
+
+    const requestLimiter = {
+      mode: "local" as const,
+      acquire: vi.fn().mockRejectedValueOnce(new ProviderRateLimitExceededError("transient local limit", 1))
+    };
+
+    const provider = new OpenAICompatibleEmbeddingProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      dimensions: 3,
+      request_limiter: requestLimiter,
+      max_requests_per_minute: 1,
+      query_max_wait_ms: 0
+    });
+
+    await expect(
+      provider.embed({
+        texts: ["hello"],
+        purpose: "query"
+      })
+    ).rejects.toMatchObject({
+      code: "RATE_LIMITED"
+    } satisfies Partial<RetrievalError>);
+
+    expect(requestLimiter.acquire).toHaveBeenCalledTimes(1);
+    expect(calls).toBe(0);
+  });
+
+  it("maps upstream 429 exhaustion to RATE_LIMITED", async () => {
+    let calls = 0;
+    globalThis.fetch = vi.fn(async () => {
+      calls += 1;
+      return new Response(JSON.stringify({ error: "rate limit" }), { status: 429 });
+    }) as typeof fetch;
+
+    const provider = new OpenAICompatibleEmbeddingProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      dimensions: 3,
+      max_retries: 1
+    });
+
+    await expect(
+      provider.embed({
+        texts: ["hello"],
+        purpose: "query"
+      })
+    ).rejects.toMatchObject({
+      code: "RATE_LIMITED"
+    } satisfies Partial<RetrievalError>);
+    expect(calls).toBe(2);
+  });
+});
+
+describe("openai-compatible reranker provider", () => {
+  const originalFetch = globalThis.fetch;
+
+  afterEach(() => {
+    globalThis.fetch = originalFetch;
+    vi.restoreAllMocks();
+  });
+
+  it("returns ranked indexes from reranker response", async () => {
+    globalThis.fetch = vi.fn(async () => {
+      return new Response(
+        JSON.stringify({
+          results: [
+            { index: 2, relevance_score: 0.91 },
+            { index: 0, relevance_score: 0.82 },
+            { index: 1, relevance_score: 0.73 }
+          ]
+        }),
+        { status: 200 }
+      );
+    }) as typeof fetch;
+
+    const provider = new OpenAICompatibleRerankerProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      model: "Qwen/Qwen3-Reranker-4B",
+      timeout_ms: 200
+    });
+
+    const ranked = await provider.rerank({
+      query: "tenant auth",
+      documents: ["a", "b", "c"],
+      top_n: 2
+    });
+
+    expect(ranked).toEqual([
+      { index: 2, relevance_score: 0.91 },
+      { index: 0, relevance_score: 0.82 }
+    ]);
+  });
+
+  it("enforces timeout on stalled reranker requests", async () => {
+    globalThis.fetch = vi.fn((_, init) => {
+      return new Promise((_, reject) => {
+        init?.signal?.addEventListener("abort", () => {
+          reject(new DOMException("The operation was aborted.", "AbortError"));
+        });
+      });
+    }) as typeof fetch;
+
+    const provider = new OpenAICompatibleRerankerProvider({
+      base_url: "https://router.tumuer.me/v1",
+      api_key: "test-key",
+      timeout_ms: 10
+    });
+
+    await expect(
+      provider.rerank({
+        query: "tenant auth",
+        documents: ["a", "b"],
+        top_n: 2
+      })
+    ).rejects.toMatchObject({
+      message: expect.stringMatching(/timed out/i)
+    });
+  });
+});