move ai accuracy tests to generative-ai package, add p-queue and make…

… them run in parallel

packages/compass-generative-ai/package.json

@@ -48,6 +48,7 @@
     "depcheck": "compass-scripts check-peer-deps && depcheck",
     "check": "npm run typecheck && npm run lint && npm run depcheck",
     "check-ci": "npm run check",
+    "ai-accuracy-tests": "node ./scripts/ai-accuracy-tests.js",
     "test": "mocha",
     "test-electron": "xvfb-maybe electron-mocha --no-sandbox",
     "test-cov": "nyc --compact=false --produce-source-map=false -x \"**/*.spec.*\" --reporter=lcov --reporter=text --reporter=html npm run test",
@@ -73,15 +74,25 @@
     "@testing-library/user-event": "^13.5.0",
     "@types/chai": "^4.2.21",
     "@types/chai-dom": "^0.0.10",
+    "@types/decomment": "^0.9.4",
     "@types/mocha": "^9.0.0",
     "@types/react": "^17.0.5",
     "@types/react-dom": "^17.0.10",
     "@types/sinon-chai": "^3.2.5",
+    "bson": "^6.0.0",
     "chai": "^4.3.6",
+    "decomment": "^0.9.5",
     "depcheck": "^1.4.1",
+    "digest-fetch": "^2.0.3",
+    "ejson-shell-parser": "^1.2.4",
     "eslint": "^7.25.0",
     "mocha": "^10.2.0",
+    "mongodb": "^6.1.0",
+    "mongodb-runner": "^5.4.4",
+    "mongodb-schema": "^11.2.2",
+    "node-fetch": "^2.7.0",
     "nyc": "^15.1.0",
+    "p-queue": "^7.4.1",
     "prettier": "^2.7.1",
     "react": "^17.0.2",
     "react-dom": "^17.0.2",

scripts/ai-accuracy-tests/index.js → packages/compass-generative-ai/scripts/ai-accuracy-tests.js

@@ -1,4 +1,7 @@
 'use strict';
+
+/* eslint-disable no-console */
+
 // To run these tests against cloud-dev:
 // > ATLAS_PUBLIC_KEY="..." \
 //   ATLAS_PRIVATE_KEY="..." \
@@ -20,6 +23,11 @@ const os = require('os');
 const assert = require('assert');
 const ejsonShellParser = require('ejson-shell-parser');
 const { MongoClient } = require('mongodb');
+let PQueue;
+(async () => {
+  // ESM package only. This could introduce race conditions.
+  PQueue = (await import('p-queue')).default;
+})();
 const { EJSON } = require('bson');
 const { getSimplifiedSchema } = require('mongodb-schema');
 const path = require('path');
@@ -33,6 +41,11 @@ const DEFAULT_MIN_ACCURACY = 0.8;
 
 const MAX_TIMEOUTS_PER_TEST = 10;
 
+// There are a limited amount of resources available both on the Atlas
+// and on the ai service side of things, so we want to limit how many
+// requests can be happening at a time.
+const TESTS_TO_RUN_CONCURRENTLY = 4;
+
 const ATTEMPTS_PER_TEST = process.env.AI_TESTS_ATTEMPTS_PER_TEST
   ? +process.env.AI_TESTS_ATTEMPTS_PER_TEST
   : DEFAULT_ATTEMPTS_PER_TEST;
@@ -487,34 +500,41 @@ const tests = [
     ]),
   },
 ];
-
 async function main() {
   try {
     await setup();
     const table = [];
 
     let anyFailed = false;
 
-    for (const test of tests) {
-      const {
-        accuracy,
-        // usageStats
-      } = await runTest(test);
-      const minAccuracy = test.minAccuracy ?? DEFAULT_MIN_ACCURACY;
-      const failed = accuracy < minAccuracy;
-
-      table.push({
-        Type: test.type.slice(0, 1).toUpperCase(),
-        'User Input': test.userInput.slice(0, 50),
-        Namespace: `${test.databaseName}.${test.collectionName}`,
-        Accuracy: accuracy,
-        // 'Prompt Tokens': usageStats[0]?.promptTokens,
-        // 'Completion Tokens': usageStats[0]?.completionTokens,
-        Pass: failed ? '✗' : '✓',
-      });
-
-      anyFailed = anyFailed || failed;
-    }
+    const testPromiseQueue = new PQueue({
+      concurrency: TESTS_TO_RUN_CONCURRENTLY,
+    });
+
+    tests.map((test) =>
+      testPromiseQueue.add(async () => {
+        const {
+          accuracy,
+          // usageStats
+        } = await runTest(test);
+        const minAccuracy = test.minAccuracy ?? DEFAULT_MIN_ACCURACY;
+        const failed = accuracy < minAccuracy;
+
+        table.push({
+          Type: test.type.slice(0, 1).toUpperCase(),
+          'User Input': test.userInput.slice(0, 50),
+          Namespace: `${test.databaseName}.${test.collectionName}`,
+          Accuracy: accuracy,
+          // 'Prompt Tokens': usageStats[0]?.promptTokens,
+          // 'Completion Tokens': usageStats[0]?.completionTokens,
+          Pass: failed ? '✗' : '✓',
+        });
+
+        anyFailed = anyFailed || failed;
+      })
+    );
+
+    await testPromiseQueue.onIdle();
 
     console.table(table, [
       'Type',
@@ -536,4 +556,5 @@ async function main() {
   }
 }
 
+// eslint-disable-next-line @typescript-eslint/no-floating-promises
 main();

scripts/package.json

@@ -42,19 +42,12 @@
     "@mongodb-js/monorepo-tools": "^1.1.1",
     "@mongodb-js/webpack-config-compass": "^1.2.3",
     "commander": "^11.0.0",
-    "decomment": "^0.9.5",
-    "digest-fetch": "^2.0.3",
-    "ejson-shell-parser": "^1.2.4",
     "electron": "^25.8.4",
     "glob": "^10.2.5",
     "jsdom": "^21.1.0",
     "keytar": "^7.9.0",
     "make-fetch-happen": "^8.0.14",
-    "mongodb": "^6.1.0",
     "mongodb-connection-string-url": "^2.6.0",
-    "mongodb-runner": "^5.4.4",
-    "mongodb-schema": "^11.2.2",
-    "node-fetch": "^2.7.0",
     "pacote": "^11.3.5",
     "pkg-up": "^3.1.0",
     "prompts": "^2.4.1",