chore: add disk based accuracy storage for local runs

Author: himanshusinghs

.gitignore

@@ -11,3 +11,4 @@ state.json
 
 tests/tmp
 coverage
+.accuracy-snapshots

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

@@ -0,0 +1,122 @@
+import path from "path";
+import fs from "fs/promises";
+import { fileURLToPath } from "url";
+import {
+    AccuracyRunStatus,
+    AccuracySnapshotEntry,
+    AccuracySnapshotEntrySchema,
+    AccuracySnapshotStorage,
+} from "./snapshot-storage.js";
+const __dirname = fileURLToPath(import.meta.url);
+const rootDir = path.resolve(__dirname, "..", "..", "..", "..", "..");
+const snapshotsDir = path.resolve(rootDir, ".accuracy-snapshots");
+export const snapshotFilePath = path.resolve(snapshotsDir, "snapshots.json");
+
+export class DiskSnapshotStorage implements AccuracySnapshotStorage {
+    private constructor(
+        private readonly accuracyRunId: string,
+        private readonly commitSHA: string
+    ) {}
+
+    async createSnapshotEntry(
+        snapshotEntry: Pick<
+            AccuracySnapshotEntry,
+            | "provider"
+            | "requestedModel"
+            | "test"
+            | "prompt"
+            | "toolCallingAccuracy"
+            | "expectedToolCalls"
+            | "actualToolCalls"
+            | "llmResponseTime"
+            | "tokensUsage"
+            | "respondingModel"
+            | "text"
+            | "messages"
+        >
+    ): Promise<void> {
+        const snapshotWithMeta: AccuracySnapshotEntry = {
+            ...snapshotEntry,
+            commitSHA: this.commitSHA,
+            accuracyRunId: this.accuracyRunId,
+            accuracyRunStatus: AccuracyRunStatus.InProgress,
+            createdOn: Date.now(),
+        };
+
+        await this.appendAccuracySnapshot(snapshotWithMeta);
+    }
+
+    async getLatestSnapshotsForCommit(commit: string): Promise<AccuracySnapshotEntry[]> {
+        const snapshot = await this.readSnapshot();
+        const entries = snapshot
+            .filter((entry) => {
+                return entry.commitSHA === commit && entry.accuracyRunStatus === AccuracyRunStatus.Done;
+            })
+            .sort((a, b) => b.createdOn - a.createdOn);
+        const latestRunId = entries[0]?.accuracyRunId;
+        return latestRunId ? snapshot.filter((entry) => entry.accuracyRunId === latestRunId) : [];
+    }
+
+    async accuracyRunFinished(): Promise<void> {
+        const snapshot = await this.readSnapshot();
+        const updatedSnapshot = snapshot.map((entry) => {
+            if (entry.accuracyRunId === this.accuracyRunId) {
+                return {
+                    ...entry,
+                    accuracyRunStatus: AccuracyRunStatus.Done,
+                };
+            }
+
+            return entry;
+        });
+        await this.writeSnapshot(updatedSnapshot);
+    }
+
+    close(): Promise<void> {
+        return Promise.resolve();
+    }
+
+    private async appendAccuracySnapshot(entry: AccuracySnapshotEntry): Promise<void> {
+        for (let attempt = 0; attempt < 5; attempt++) {
+            try {
+                const snapshot = await this.readSnapshot();
+                snapshot.unshift(entry);
+                await this.writeSnapshot(snapshot);
+                return;
+            } catch (e) {
+                if (attempt < 4) {
+                    await this.waitFor(100 + Math.random() * 200);
+                } else {
+                    throw e;
+                }
+            }
+        }
+    }
+
+    private async writeSnapshot(snapshot: AccuracySnapshotEntry[]): Promise<void> {
+        const tmp = `${snapshotFilePath}~${Date.now()}`;
+        await fs.writeFile(tmp, JSON.stringify(snapshot, null, 2));
+        await fs.rename(tmp, snapshotFilePath);
+    }
+
+    private async readSnapshot(): Promise<AccuracySnapshotEntry[]> {
+        try {
+            const raw = await fs.readFile(snapshotFilePath, "utf8");
+            return AccuracySnapshotEntrySchema.array().parse(JSON.parse(raw));
+        } catch (e: unknown) {
+            if ((e as { code: string }).code === "ENOENT") {
+                return [];
+            }
+            throw e;
+        }
+    }
+
+    private waitFor(ms: number) {
+        return new Promise((resolve) => setTimeout(resolve, ms));
+    }
+
+    static async getStorage(commitSHA: string, accuracyRunId: string) {
+        await fs.mkdir(snapshotsDir, { recursive: true });
+        return new DiskSnapshotStorage(commitSHA, accuracyRunId);
+    }
+}

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

@@ -1,4 +1,5 @@
 import { getCommitSHA } from "../git-info.js";
+import { DiskSnapshotStorage } from "./disk-snapshot-storage.js";
 import { MongoDBSnapshotStorage } from "./mdb-snapshot-storage.js";
 import { AccuracySnapshotStorage } from "./snapshot-storage.js";
 
@@ -15,5 +16,8 @@ export async function getAccuracySnapshotStorage(): Promise<AccuracySnapshotStor
         throw new Error("Cannot create AccuracySnapshotStorage without a commitSHA.");
     }
 
-    return MongoDBSnapshotStorage.getStorage(commitSHA, accuracyRunId);
+    return (
+        MongoDBSnapshotStorage.getStorage(commitSHA, accuracyRunId) ??
+        (await DiskSnapshotStorage.getStorage(commitSHA, accuracyRunId))
+    );
 }

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

@@ -58,13 +58,13 @@ export class MongoDBSnapshotStorage implements AccuracySnapshotStorage {
     }
 
     async getLatestSnapshotsForCommit(commit: string): Promise<AccuracySnapshotEntry[]> {
-        const latestRunId = await this.getLastRunIdForCommit(commit);
+        const latestRunId = await this.getLatestAccuracyRunForCommit(commit);
         return latestRunId ? this.getSnapshotEntriesForRunId(latestRunId) : [];
     }
 
-    private async getLastRunIdForCommit(commit: string): Promise<string | undefined> {
+    private async getLatestAccuracyRunForCommit(commit: string): Promise<string | undefined> {
         const document = await this.snapshotCollection.findOne(
-            { commit: commit },
+            { commit: commit, accuracyRunStatus: AccuracyRunStatus.Done },
             { sort: { createdOn: -1 }, projection: { accuracyRunId: 1 } }
         );
 
@@ -83,12 +83,16 @@ export class MongoDBSnapshotStorage implements AccuracySnapshotStorage {
         );
     }
 
-    static getStorage(commitSHA: string, accuracyRunId: string): MongoDBSnapshotStorage {
+    async close(): Promise<void> {
+        await this.client.close();
+    }
+
+    static getStorage(commitSHA: string, accuracyRunId: string): MongoDBSnapshotStorage | null {
         const mongodbUrl = process.env.MDB_ACCURACY_MDB_URL;
         const database = process.env.MDB_ACCURACY_MDB_DB;
         const collection = process.env.MDB_ACCURACY_MDB_COLLECTION;
         if (!mongodbUrl || !database || !collection) {
-            throw new Error("Cannot create MongoDBAccuracySnapshot storage without relevant configuration provided");
+            return null;
         }
 
         return new MongoDBSnapshotStorage({
@@ -99,8 +103,4 @@ export class MongoDBSnapshotStorage implements AccuracySnapshotStorage {
             accuracyRunId,
         });
     }
-
-    async close(): Promise<void> {
-        await this.client.close();
-    }
 }