evOo-local/packages/adapters/src/openai-compatible-embeddings.test.ts

import { afterEach, describe, expect, it, vi } from "vitest";
import { OpenAiCompatibleEmbeddingServerEmulator } from "./embedding-emulator.js";
import {
  OpenAiCompatibleEmbeddingProvider,
  parseEmbeddingResponse,
} from "./openai-compatible-embeddings.js";

afterEach(() => vi.unstubAllEnvs());

function provider(
  server = new OpenAiCompatibleEmbeddingServerEmulator(),
  config: { baseUrl?: string; model?: string; apiKey?: string } = {},
) {
  return new OpenAiCompatibleEmbeddingProvider(
    {
      baseUrl: config.baseUrl ?? "http://127.0.0.1:8080",
      model: config.model ?? "nomic-embed-text",
      ...(config.apiKey ? { apiKey: config.apiKey } : {}),
    },
    server.fetch,
  );
}

describe("OpenAiCompatibleEmbeddingProvider", () => {
  it("posts the OpenAI embeddings request to the versioned endpoint in batches", async () => {
    const server = new OpenAiCompatibleEmbeddingServerEmulator();
    const texts = Array.from({ length: 70 }, (_, index) => `text ${index}`);

    await provider(server).embed(texts);

    expect(server.requests.map((request) => (request.input as string[]).length)).toEqual([
      32, 32, 6,
    ]);
    expect(server.requests[0]).toMatchObject({
      url: "http://127.0.0.1:8080/v1/embeddings",
      authorization: null,
      model: "nomic-embed-text",
    });
    expect(server.requests.flatMap((request) => request.input as string[])).toEqual(texts);
  });

  it("names the configured model as its vector space", () => {
    expect(provider().describe()).toMatchObject({
      id: "openai-compatible",
      capabilities: { model: "nomic-embed-text" },
    });
    expect(() => provider(undefined, { model: " " })).toThrow(/model is required/);
  });

  it("sends the API key as a bearer token", async () => {
    const server = new OpenAiCompatibleEmbeddingServerEmulator({ apiKey: "sk-test-key" });
    await provider(server, { apiKey: "sk-test-key" }).embed(["hello"]);
    expect(server.requests[0]?.authorization).toBe("Bearer sk-test-key");
  });

  it("applies the OpenAI-compatible endpoint policy", () => {
    expect(() => provider(undefined, { baseUrl: "http://169.254.169.254/v1" })).toThrow(/metadata/);
    expect(() => provider(undefined, { baseUrl: "https://embeddings.example.com/v1" })).toThrow(
      /Public model endpoints are blocked/,
    );
    vi.stubEnv("RAKAZO_OPENAI_COMPAT_ALLOW_PUBLIC", "1");
    expect(() =>
      provider(undefined, { baseUrl: "http://embeddings.example.com/v1", apiKey: "sk-test-key" }),
    ).toThrow(/must use HTTPS/);
  });

  it("reports HTTP failures without echoing the response or the key", async () => {
    const server = new OpenAiCompatibleEmbeddingServerEmulator({ apiKey: "sk-test-key" });
    server.failNextStatus = 503;
    const failure = provider(server, { apiKey: "sk-test-key" }).embed(["hello"]);
    await expect(failure).rejects.toThrow("Embeddings server returned 503");
    await expect(failure).rejects.not.toThrow(/sk-test-key|Upstream/);
    await expect(provider(server).embed(["hello"])).rejects.toThrow("returned 401");
  });

  it("reports a refused connection as an unreachable server", async () => {
    const embeddings = new OpenAiCompatibleEmbeddingProvider(
      { baseUrl: "http://127.0.0.1:8080", model: "m" },
      async () => {
        throw new TypeError("fetch failed");
      },
    );
    await expect(embeddings.embed(["hello"])).rejects.toThrow("Embeddings server is unreachable");
  });

  it("rejects vectors of different lengths across batches", async () => {
    let call = 0;
    const fetchImpl: typeof fetch = async (_input, init) => {
      const texts = JSON.parse(String(init?.body)).input as string[];
      call += 1;
      return Response.json({
        data: texts.map((_, index) => ({ index, embedding: call === 1 ? [1, 0] : [1, 0, 0] })),
      });
    };
    const embeddings = new OpenAiCompatibleEmbeddingProvider(
      { baseUrl: "http://127.0.0.1:8080", model: "m" },
      fetchImpl,
    );
    await expect(embeddings.embed(Array.from({ length: 40 }, () => "x"))).rejects.toThrow(
      /different lengths/,
    );
  });
});

describe("parseEmbeddingResponse", () => {
  it("orders vectors by index", () => {
    expect(
      parseEmbeddingResponse(
        {
          data: [
            { index: 1, embedding: [0, 1] },
            { index: 0, embedding: [1, 0] },
          ],
        },
        2,
      ),
    ).toEqual([
      [1, 0],
      [0, 1],
    ]);
  });

  it("falls back to position when a server omits the index", () => {
    expect(parseEmbeddingResponse({ data: [{ embedding: [0.5] }] }, 1)).toEqual([[0.5]]);
  });

  it.each([
    ["a missing data list", {}],
    ["the wrong count", { data: [{ index: 0, embedding: [1] }] }],
    [
      "a duplicate index",
      {
        data: [
          { index: 0, embedding: [1] },
          { index: 0, embedding: [1] },
        ],
      },
    ],
    [
      "a non-finite value",
      {
        data: [
          { index: 0, embedding: [1] },
          { index: 1, embedding: [Number.NaN] },
        ],
      },
    ],
    [
      "an empty vector",
      {
        data: [
          { index: 0, embedding: [] },
          { index: 1, embedding: [1] },
        ],
      },
    ],
  ])("rejects %s", (_name, body) => {
    expect(() => parseEmbeddingResponse(body, 2)).toThrow(/Embeddings server returned/);
  });
});