Add tests for @gel/ai (#1264)

There were many yaks to shave which is why I've held of on this for so
long, but after releasing an embarrassing bug into a stable release,
it's time to actually tighten up here. 📈

Author: scotttrinh

packages/ai/jest.config.js

@@ -0,0 +1,23 @@
+export const TS_EXT_TO_TREAT_AS_ESM = [".ts", ".tsx", ".mts"];
+export const JS_EXT_TO_TREAT_AS_ESM = [".jsx"];
+export const ESM_TS_JS_TRANSFORM_PATTERN = "^.+\\.m?[tj]sx?$";
+
+export default {
+  testEnvironment: "node",
+  testPathIgnorePatterns: ["./dist"],
+  globalSetup: "./test/globalSetup.ts",
+  globalTeardown: "./test/globalTeardown.ts",
+  extensionsToTreatAsEsm: [
+    ...JS_EXT_TO_TREAT_AS_ESM,
+    ...TS_EXT_TO_TREAT_AS_ESM,
+  ],
+  transform: {
+    [ESM_TS_JS_TRANSFORM_PATTERN]: [
+      "ts-jest",
+      {
+        tsconfig: "./tsconfig.test.json",
+        useESM: true,
+      },
+    ],
+  },
+};

packages/ai/package.json

@@ -28,18 +28,19 @@
   ],
   "scripts": {
     "typecheck": "tsc --project tsconfig.json --noEmit",
-    "test": "jest --detectOpenHandles --passWithNoTests",
+    "test": "NODE_OPTIONS='--experimental-global-webcrypto --experimental-vm-modules' jest --detectOpenHandles",
     "build": "tsc --project tsconfig.json",
     "lint": "eslint --quiet",
     "lint:fix": "eslint --fix"
   },
   "devDependencies": {
+    "@repo/test-utils": "*",
     "@repo/tsconfig": "*",
     "@types/jest": "^29.5.12",
     "@types/node": "^20.12.13",
     "gel": "^2.0.0",
     "jest": "29.7.0",
-    "ts-jest": "29.1.4",
+    "ts-jest": "^29.1.4",
     "typescript": "^5.5.2"
   },
   "peerDependencies": {

packages/ai/test/core.test.ts

@@ -0,0 +1,146 @@
+import { type Client } from "gel";
+import { createRAGClient } from "../dist/index.js";
+import { getClient, waitFor, getAvailableExtensions } from "@repo/test-utils";
+import { createMockHttpServer, type MockHttpServer } from "./mockHttpServer";
+
+const availableExtensions = getAvailableExtensions();
+
+if (availableExtensions.has("ai")) {
+  let mockServer: MockHttpServer;
+
+  beforeAll(async () => {
+    // Start the mock server
+    mockServer = createMockHttpServer();
+
+    const client = getClient();
+    await client.ensureConnected();
+    try {
+      await client.execute(`
+create extension pgvector;
+create extension ai;
+
+create type TestEmbeddingModel extending ext::ai::EmbeddingModel {
+  alter annotation ext::ai::model_name := "text-embedding-test";
+  alter annotation ext::ai::model_provider := "custom::test";
+  alter annotation ext::ai::embedding_model_max_input_tokens := "8191";
+  alter annotation ext::ai::embedding_model_max_batch_tokens := "16384";
+  alter annotation ext::ai::embedding_model_max_output_dimensions := "10";
+  alter annotation ext::ai::embedding_model_supports_shortening := "true";
+};
+
+create type TestTextGenerationModel extending ext::ai::TextGenerationModel {
+  alter annotation ext::ai::model_name := "text-generation-test";
+  alter annotation ext::ai::model_provider := "custom::test";
+  alter annotation ext::ai::text_gen_model_context_window := "16385";
+};
+
+create type Astronomy {
+  create required property content: str;
+
+  create deferred index ext::ai::index(embedding_model := "text-embedding-test") on (.content);
+};
+
+configure current branch insert ext::ai::CustomProviderConfig {
+    name := "custom::test",
+    secret := "dummy-key",
+    api_url := "${mockServer.url}/v1",
+    api_style := ext::ai::ProviderAPIStyle.OpenAI,
+};
+
+configure current branch set ext::ai::Config::indexer_naptime := <duration>"100ms";
+    `);
+    } finally {
+      await client.close();
+    }
+  }, 25_000);
+
+  afterAll(async () => {
+    // Stop the mock server
+    if (mockServer) {
+      await mockServer.close();
+    }
+  });
+
+  describe("@gel/ai", () => {
+    let client: Client;
+    beforeEach(() => {
+      mockServer.resetRequests();
+    });
+
+    afterEach(async () => {
+      await client?.close();
+    });
+
+    test("RAG query", async () => {
+      client = getClient({
+        tlsSecurity: "insecure",
+      });
+      await client.execute(`
+insert Astronomy { content := 'Skies on Mars are red' };
+insert Astronomy { content := 'Skies on Earth are blue' };
+      `);
+      await waitFor(async () =>
+        expect(mockServer.getEmbeddingsRequests().length).toBe(1),
+      );
+
+      const ragClient = createRAGClient(client, {
+        model: "text-generation-test",
+      }).withContext({
+        query: "select Astronomy",
+      });
+
+      const result = await ragClient.queryRag({
+        prompt: "What color are the skies on Mars?",
+      });
+
+      expect(result).toEqual("This is a mock response.");
+
+      const streamedResult = ragClient.streamRag({
+        prompt: "What color are the skies on Mars?",
+      });
+
+      let streamedResultString = "";
+      for await (const message of streamedResult) {
+        if (
+          message.type === "content_block_start" &&
+          message.content_block.type === "text"
+        ) {
+          streamedResultString += message.content_block.text;
+        } else if (
+          message.type === "content_block_delta" &&
+          message.delta.type === "text_delta"
+        ) {
+          streamedResultString += message.delta.text;
+        }
+      }
+
+      expect(streamedResultString).toEqual("This is a mock response.");
+    }, 25_000);
+
+    test("embedding request", async () => {
+      client = getClient({
+        tlsSecurity: "insecure",
+      });
+      const ragClient = createRAGClient(client, {
+        model: "text-generation-test",
+      });
+
+      const result = await ragClient.generateEmbeddings({
+        inputs: ["beep boop!"],
+        model: "text-embedding-test",
+      });
+
+      await waitFor(
+        async () => expect(mockServer.getEmbeddingsRequests().length).toBe(1),
+        10_000,
+        500,
+      );
+
+      expect(result).toEqual(
+        // The number of occurences of the first ten letters of the alphabet.
+        // two 'b's and two 'e's.
+        [0, 2, 0, 0, 2, 0, 0, 0, 0, 0],
+      );
+    });
+  });
+}

packages/ai/test/globalSetup.ts

@@ -0,0 +1,50 @@
+import * as process from "node:process";
+import {
+  connectToServer,
+  generateStatusFileName,
+  getServerCommand,
+  getWSLPath,
+  startServer,
+  ConnectConfig,
+} from "@repo/test-utils";
+
+export default async () => {
+  // tslint:disable-next-line
+  console.log("\nStarting Gel test cluster...");
+
+  const statusFile = generateStatusFileName("node");
+  console.log("Node status file:", statusFile);
+
+  const { args, availableFeatures } = getServerCommand(getWSLPath(statusFile));
+  console.log(`Starting server...`);
+  const { proc, config } = await startServer(args, statusFile);
+
+  const { client, version } = await connectToServer(config);
+
+  const jestConfig: ConnectConfig = {
+    ...config,
+    user: version.major >= 6 ? "admin" : "edgedb",
+  };
+
+  // @ts-ignore
+  global.gelProc = proc;
+
+  process.env._JEST_GEL_CONNECT_CONFIG = JSON.stringify(jestConfig);
+  process.env._JEST_GEL_AVAILABLE_FEATURES = JSON.stringify(availableFeatures);
+
+  // @ts-ignore
+  global.gelConn = client;
+  process.env._JEST_GEL_VERSION = JSON.stringify(version);
+
+  const availableExtensions = (
+    await client.query<{
+      name: string;
+      version: { major: number; minor: number };
+    }>(`select sys::ExtensionPackage {name, version}`)
+  ).map(({ name, version }) => [name, version]);
+  process.env._JEST_GEL_AVAILABLE_EXTENSIONS =
+    JSON.stringify(availableExtensions);
+
+  // tslint:disable-next-line
+  console.log(`Gel test cluster is up [port: ${jestConfig.port}]...`);
+};

packages/ai/test/globalTeardown.ts

@@ -0,0 +1,14 @@
+import { shutdown } from "@repo/test-utils";
+
+export default async () => {
+  // tslint:disable-next-line
+  console.log("Shutting down Gel test cluster...");
+
+  try {
+    // @ts-ignore
+    await shutdown(global.gelProc, global.gelConn);
+  } finally {
+    // tslint:disable-next-line
+    console.log("Gel test cluster is down...");
+  }
+};