openclaw/src/memory-host-sdk/host/batch-openai.test.ts

import { beforeEach, describe, expect, it, vi } from "vitest";

const mocks = vi.hoisted(() => ({
  uploadBatchJsonlFile: vi.fn(async () => "file_in"),
  postJsonWithRetry: vi.fn(async () => ({ id: "batch_1", status: "in_progress" })),
  resolveCompletedBatchResult: vi.fn(async () => ({ outputFileId: "file_out" })),
  withRemoteHttpResponse: vi.fn(
    async (params: { url: string; onResponse: (res: Response) => Promise<unknown> }) => {
      if (params.url.endsWith("/files/file_out/content")) {
        return await params.onResponse(
          new Response(
            [
              JSON.stringify({
                custom_id: "0",
                response: {
                  status_code: 200,
                  body: { data: [{ embedding: [1, 0, 0], index: 0 }] },
                },
              }),
              JSON.stringify({
                custom_id: "1",
                response: {
                  status_code: 200,
                  body: { data: [{ embedding: [2, 0, 0], index: 0 }] },
                },
              }),
            ].join("\n"),
            { status: 200, headers: { "Content-Type": "application/jsonl" } },
          ),
        );
      }
      return await params.onResponse(
        new Response(
          JSON.stringify({ id: "batch_1", status: "completed", output_file_id: "file_out" }),
          {
            status: 200,
            headers: { "Content-Type": "application/json" },
          },
        ),
      );
    },
  ),
}));

vi.mock("./batch-embedding-common.js", async () => {
  const actual = await vi.importActual<typeof import("./batch-embedding-common.js")>(
    "./batch-embedding-common.js",
  );
  return {
    ...actual,
    uploadBatchJsonlFile: mocks.uploadBatchJsonlFile,
    postJsonWithRetry: mocks.postJsonWithRetry,
    resolveCompletedBatchResult: mocks.resolveCompletedBatchResult,
    withRemoteHttpResponse: mocks.withRemoteHttpResponse,
  };
});

describe("runOpenAiEmbeddingBatches", () => {
  beforeEach(() => {
    vi.clearAllMocks();
  });

  it("maps uploaded batch output rows back to embeddings", async () => {
    const { runOpenAiEmbeddingBatches, OPENAI_BATCH_ENDPOINT } = await import("./batch-openai.js");

    const result = await runOpenAiEmbeddingBatches({
      openAi: {
        baseUrl: "https://api.openai.com/v1",
        headers: { Authorization: "Bearer test" },
        fetchImpl: fetch,
        model: "text-embedding-3-small",
      },
      agentId: "main",
      requests: [
        {
          custom_id: "0",
          method: "POST",
          url: OPENAI_BATCH_ENDPOINT,
          body: { model: "text-embedding-3-small", input: "hello" },
        },
        {
          custom_id: "1",
          method: "POST",
          url: OPENAI_BATCH_ENDPOINT,
          body: { model: "text-embedding-3-small", input: "world" },
        },
      ],
      wait: true,
      pollIntervalMs: 1,
      timeoutMs: 1000,
      concurrency: 3,
    });

    expect(mocks.uploadBatchJsonlFile).toHaveBeenCalled();
    expect(mocks.postJsonWithRetry).toHaveBeenCalledWith(
      expect.objectContaining({
        errorPrefix: "openai batch create failed",
        body: expect.objectContaining({
          endpoint: OPENAI_BATCH_ENDPOINT,
          metadata: { source: "openclaw-memory", agent: "main" },
        }),
      }),
    );
    expect(result.get("0")).toEqual([1, 0, 0]);
    expect(result.get("1")).toEqual([2, 0, 0]);
  });
});