chore: adds accuracyRunStatus to snapshot entries

The new field `accuracyRunStatus` is supposed to help guard against
cases where jest might fail in between, maybe due to LLM rate limit
errors or something else, and we then have a partially saved state of an
accuracy run. With the new field `accuracyRunStatus` we should be able
to safely look for last runs where `accuracyRunStatus` is done and have
complete state of accuracy snapshot.

Author: himanshusinghs

scripts/mark-accuracy-run-finished.ts

@@ -0,0 +1,7 @@
+import { getAccuracySnapshotStorage } from "../tests/accuracy/sdk/accuracy-snapshot-storage/get-snapshot-storage.js";
+
+console.time(`Marked accuracy run id - ${process.env.MDB_ACCURACY_RUN_ID} as finished in`);
+const storage = await getAccuracySnapshotStorage();
+await storage.accuracyRunFinished();
+await storage.close();
+console.timeEnd(`Marked accuracy run id - ${process.env.MDB_ACCURACY_RUN_ID} as finished in`);

scripts/run-accuracy-tests.sh

@@ -2,6 +2,27 @@
 # Variables necessary for the accuracy test runs
 export MDB_ACCURACY_RUN_ID=$(npx uuid v4)
 
+# For providing access tokens for different LLM providers
+# export MDB_OPEN_AI_API_KEY=""
+# export MDB_GEMINI_API_KEY=""
+# export MDB_AZURE_OPEN_AI_API_KEY=""
+# export MDB_AZURE_OPEN_AI_API_URL=""
+
+# For providing a mongodb based storage to store accuracy snapshots
+# export MDB_ACCURACY_MDB_URL=""
+# export MDB_ACCURACY_MDB_DB=""
+# export MDB_ACCURACY_MDB_COLLECTION=""
+
+# By default we run all the tests under tests/accuracy folder unless a path is
+# specified in the command line. Such as:
+# npm run test:accuracy -- tests/accuracy/some-test.test.ts
 TEST_PATH_PATTERN="${1:-tests/accuracy}"
 shift || true
-node --experimental-vm-modules node_modules/jest/bin/jest.js --testPathPattern "$TEST_PATH_PATTERN" "$@"
+node --experimental-vm-modules node_modules/jest/bin/jest.js --testPathPattern "$TEST_PATH_PATTERN" "$@"
+
+# Each test run submits an accuracy snapshot entry for each prompt with the
+# accuracyRunStatus: "in-progress". When all the tests are done and jest exits
+# with an exit code of 0, we can safely mark accuracy run as finished.
+if [ $? -eq 0 ]; then
+  npx tsx scripts/mark-accuracy-run-finished.ts
+fi

tests/accuracy/sdk/accuracy-snapshot-storage/mdb-snapshot-storage.ts

@@ -1,5 +1,10 @@
 import { Collection, MongoClient } from "mongodb";
-import { AccuracySnapshotEntry, AccuracySnapshotEntrySchema, AccuracySnapshotStorage } from "./snapshot-storage.js";
+import {
+    AccuracyRunStatus,
+    AccuracySnapshotEntry,
+    AccuracySnapshotEntrySchema,
+    AccuracySnapshotStorage,
+} from "./snapshot-storage.js";
 
 export class MongoDBSnapshotStorage implements AccuracySnapshotStorage {
     private readonly client: MongoClient;
@@ -46,6 +51,7 @@ export class MongoDBSnapshotStorage implements AccuracySnapshotStorage {
             ...snapshotEntry,
             commitSHA: this.commitSHA,
             accuracyRunId: this.accuracyRunId,
+            accuracyRunStatus: AccuracyRunStatus.InProgress,
             createdOn: Date.now(),
         };
         await this.snapshotCollection.insertOne(snapshotWithMeta);
@@ -70,6 +76,13 @@ export class MongoDBSnapshotStorage implements AccuracySnapshotStorage {
         return AccuracySnapshotEntrySchema.array().parse(snapshotEntries);
     }
 
+    async accuracyRunFinished(): Promise<void> {
+        await this.snapshotCollection.updateMany(
+            { accuracyRunId: this.accuracyRunId },
+            { $set: { accuracyRunStatus: AccuracyRunStatus.Done } }
+        );
+    }
+
     static getStorage(commitSHA: string, accuracyRunId: string): MongoDBSnapshotStorage {
         const mongodbUrl = process.env.MDB_ACCURACY_MDB_URL;
         const database = process.env.MDB_ACCURACY_MDB_DB;

tests/accuracy/sdk/accuracy-snapshot-storage/snapshot-storage.ts

@@ -4,15 +4,22 @@ const ExpectedToolCallSchema = z.object({
     toolName: z.string(),
     parameters: z.record(z.string(), z.unknown()),
 });
+export type ExpectedToolCall = z.infer<typeof ExpectedToolCallSchema>;
 
 const ActualToolCallSchema = ExpectedToolCallSchema.extend({ toolCallId: z.string() });
-
-export type ExpectedToolCall = z.infer<typeof ExpectedToolCallSchema>;
 export type ActualToolCall = z.infer<typeof ActualToolCallSchema>;
 
+export const AccuracyRunStatus = {
+    Done: "done",
+    InProgress: "in-progress",
+} as const;
+
 export const AccuracySnapshotEntrySchema = z.object({
     // Git and meta information for snapshot entries
     accuracyRunId: z.string(),
+    accuracyRunStatus: z
+        .enum([AccuracyRunStatus.Done, AccuracyRunStatus.InProgress])
+        .default(AccuracyRunStatus.InProgress),
     createdOn: z.number(),
     commitSHA: z.string(),
     // Accuracy info
@@ -60,5 +67,7 @@ export interface AccuracySnapshotStorage {
 
     getLatestSnapshotsForCommit(commit: string): Promise<AccuracySnapshotEntry[]>;
 
+    accuracyRunFinished(): Promise<void>;
+
     close(): Promise<void>;
 }