chore: MongoDB based snapshot storage for accuracy runs

introduces the following necessary env variables:
- MDB_ACCURACY_RUN_ID: The accuracy run id
- MDB_ACCURACY_MDB_URL: The connection string to mongodb instance where the snapshots will be stored
- MDB_ACCURACY_MDB_DB: The database for snapshots
- MDB_ACCURACY_MDB_COLLECTION: The collection for snapshots

Author: himanshusinghs

package-lock.json

@@ -31,7 +31,7 @@
     "generate": "./scripts/generate.sh",
     "test": "node --experimental-vm-modules node_modules/jest/bin/jest.js --coverage --testPathIgnorePatterns=/tests/accuracy/",
     "test:accuracy": "node --experimental-vm-modules node_modules/jest/bin/jest.js --testPathPattern tests/accuracy",
-    "test:accuracy-file": "node --experimental-vm-modules node_modules/jest/bin/jest.js"
+    "test:accuracy-file": "MDB_ACCURACY_RUN_ID=$(npx uuid v4) node --experimental-vm-modules node_modules/jest/bin/jest.js"
   },
   "license": "Apache-2.0",
   "devDependencies": {
@@ -62,6 +62,7 @@
     "openapi-types": "^12.1.3",
     "openapi-typescript": "^7.6.1",
     "prettier": "^3.5.3",
+    "simple-git": "^3.28.0",
     "ts-jest": "^29.3.1",
     "tsx": "^4.19.3",
     "typescript": "^5.8.2",

package.json

@@ -0,0 +1,19 @@
+import { getCommitSHA } from "../git-info.js";
+import { MongoDBSnapshotStorage } from "./mdb-snapshot-storage.js";
+import { AccuracySnapshotStorage } from "./snapshot-storage.js";
+
+export async function getAccuracySnapshotStorage(): Promise<AccuracySnapshotStorage> {
+    const accuracyRunId = process.env.MDB_ACCURACY_RUN_ID;
+    if (!accuracyRunId) {
+        throw new Error(
+            "Cannot create AccuracySnapshotStorage without an accuracyRunId - ensure that the relevant env variable is present."
+        );
+    }
+
+    const commitSHA = await getCommitSHA();
+    if (!commitSHA) {
+        throw new Error("Cannot create AccuracySnapshotStorage without a commitSHA.");
+    }
+
+    return MongoDBSnapshotStorage.getStorage(commitSHA, accuracyRunId);
+}

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

@@ -0,0 +1,86 @@
+import { Collection, MongoClient } from "mongodb";
+import { AccuracySnapshotEntry, AccuracySnapshotEntrySchema, AccuracySnapshotStorage } from "./snapshot-storage.js";
+
+export class MongoDBSnapshotStorage implements AccuracySnapshotStorage {
+    private readonly client: MongoClient;
+    private readonly snapshotCollection: Collection;
+    private readonly accuracyRunId: string;
+    private readonly commitSHA: string;
+    private constructor({
+        mongodbUrl,
+        database,
+        collection,
+        accuracyRunId,
+        commitSHA,
+    }: {
+        mongodbUrl: string;
+        database: string;
+        collection: string;
+        accuracyRunId: string;
+        commitSHA: string;
+    }) {
+        this.client = new MongoClient(mongodbUrl);
+        this.snapshotCollection = this.client.db(database).collection(collection);
+        this.accuracyRunId = accuracyRunId;
+        this.commitSHA = commitSHA;
+    }
+
+    async createSnapshotEntry(
+        snapshotEntry: Pick<
+            AccuracySnapshotEntry,
+            | "requestedModel"
+            | "test"
+            | "prompt"
+            | "toolCallingAccuracy"
+            | "parameterAccuracy"
+            | "llmResponseTime"
+            | "tokensUsage"
+            | "respondingModel"
+            | "text"
+            | "messages"
+        >
+    ): Promise<void> {
+        const snapshotWithMeta: AccuracySnapshotEntry = {
+            ...snapshotEntry,
+            commitSHA: this.commitSHA,
+            accuracyRunId: this.accuracyRunId,
+            createdOn: Date.now(),
+        };
+        await this.snapshotCollection.insertOne(snapshotWithMeta);
+    }
+
+    async getLastRunIdForCommit(commit: string): Promise<string | undefined> {
+        const document = await this.snapshotCollection.findOne(
+            { commit: commit },
+            { sort: { createdOn: -1 }, projection: { accuracyRunId: 1 } }
+        );
+
+        return document?.accuracyRunId ? `${document?.accuracyRunId}` : undefined;
+    }
+
+    async getSnapshotEntriesForRunId(accuracyRunId: string): Promise<AccuracySnapshotEntry[]> {
+        const snapshotEntries = await this.snapshotCollection.find({ accuracyRunId }).toArray();
+        return AccuracySnapshotEntrySchema.array().parse(snapshotEntries);
+    }
+
+    static getStorage(commitSHA: string, accuracyRunId: string): MongoDBSnapshotStorage {
+        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 new MongoDBSnapshotStorage({
+            mongodbUrl,
+            database,
+            collection,
+            commitSHA,
+            accuracyRunId,
+        });
+    }
+
+    async close(): Promise<void> {
+        await this.client.close();
+    }
+}

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

@@ -0,0 +1,51 @@
+import z from "zod";
+
+export const AccuracySnapshotEntrySchema = z.object({
+    // Git and meta information for snapshot entries
+    accuracyRunId: z.string(),
+    createdOn: z.number(),
+    commitSHA: z.string(),
+    // Accuracy info
+    requestedModel: z.string(),
+    test: z.string(),
+    prompt: z.string(),
+    toolCallingAccuracy: z.number(),
+    parameterAccuracy: z.number(),
+    llmResponseTime: z.number(),
+    tokensUsage: z
+        .object({
+            promptTokens: z.number().optional(),
+            completionTokens: z.number().optional(),
+            totalTokens: z.number().optional(),
+        })
+        .optional(),
+    respondingModel: z.string(),
+    text: z.string(),
+    messages: z.array(z.record(z.string(), z.unknown())),
+});
+
+export type AccuracySnapshotEntry = z.infer<typeof AccuracySnapshotEntrySchema>;
+
+export interface AccuracySnapshotStorage {
+    createSnapshotEntry(
+        snapshotEntry: Pick<
+            AccuracySnapshotEntry,
+            | "requestedModel"
+            | "test"
+            | "prompt"
+            | "toolCallingAccuracy"
+            | "parameterAccuracy"
+            | "llmResponseTime"
+            | "tokensUsage"
+            | "respondingModel"
+            | "text"
+            | "messages"
+        >
+    ): Promise<void>;
+
+    getLastRunIdForCommit(commit: string): Promise<string | undefined>;
+
+    getSnapshotEntriesForRunId(accuracyRunId: string): Promise<AccuracySnapshotEntry[]>;
+
+    close(): Promise<void>;
+}

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

@@ -1,4 +1,4 @@
-import { generateText, Tool, Schema, LanguageModelV1 } from "ai";
+import { generateText, LanguageModelV1, experimental_createMCPClient } from "ai";
 import { Model } from "./models.js";
 
 const systemPrompt = [
@@ -10,15 +10,32 @@ const systemPrompt = [
     'If you do not know the answer or the request cannot be fulfilled, you MUST reply with "I don\'t know"',
 ];
 
-export interface Agent<M = unknown, T = unknown, R = unknown> {
-    prompt(prompt: string, model: M, tools: T): Promise<R>;
+// Some necessary types from Vercel SDK
+export type VercelMCPClient = Awaited<ReturnType<typeof experimental_createMCPClient>>;
+export type VercelMCPClientTools = Awaited<ReturnType<VercelMCPClient["tools"]>>;
+export type VercelAgent = ReturnType<typeof getVercelToolCallingAgent>;
+
+// Generic interface for Agent, in case we need to switch to some other agent
+// development SDK
+export interface AgentPromptResult {
+    respondingModel: string;
+    tokensUsage?: {
+        promptTokens?: number;
+        completionTokens?: number;
+        totalTokens?: number;
+    };
+    text: string;
+    messages: Record<string, unknown>[];
+}
+export interface Agent<Model = unknown, Tools = unknown, Result = unknown> {
+    prompt(prompt: string, model: Model, tools: Tools): Promise<Result>;
 }
 
 export function getVercelToolCallingAgent(
     requestedSystemPrompt?: string
-): Agent<Model<LanguageModelV1>, Record<string, Tool<Schema<unknown>>>, { text: string; messages: unknown[] }> {
+): Agent<Model<LanguageModelV1>, VercelMCPClientTools, AgentPromptResult> {
     return {
-        async prompt(prompt: string, model: Model<LanguageModelV1>, tools: Record<string, Tool<Schema<unknown>>>) {
+        async prompt(prompt: string, model: Model<LanguageModelV1>, tools: VercelMCPClientTools) {
             const result = await generateText({
                 model: model.getModel(),
                 system: [...systemPrompt, requestedSystemPrompt].join("\n"),
@@ -29,6 +46,8 @@ export function getVercelToolCallingAgent(
             return {
                 text: result.text,
                 messages: result.response.messages,
+                respondingModel: result.response.modelId,
+                tokensUsage: result.usage,
             };
         },
     };

tests/accuracy/sdk/agent.ts

@@ -1,8 +1,10 @@
 import { TestableModels } from "./models.js";
 import { ExpectedToolCall, parameterMatchingAccuracyScorer, toolCallingAccuracyScorer } from "./accuracy-scorers.js";
-import { Agent, getVercelToolCallingAgent } from "./agent.js";
+import { getVercelToolCallingAgent, VercelAgent } from "./agent.js";
 import { prepareTestData, setupMongoDBIntegrationTest } from "../../integration/tools/mongodb/mongodbHelpers.js";
 import { AccuracyTestingClient, MockedTools } from "./accuracy-testing-client.js";
+import { getAccuracySnapshotStorage } from "./accuracy-snapshot-storage/get-snapshot-storage.js";
+import { AccuracySnapshotStorage } from "./accuracy-snapshot-storage/snapshot-storage.js";
 
 export interface AccuracyTestConfig {
     systemPrompt?: string;
@@ -35,10 +37,12 @@ export function describeAccuracyTests(
         const mdbIntegration = setupMongoDBIntegrationTest();
         const { populateTestData, cleanupTestDatabases } = prepareTestData(mdbIntegration);
 
+        let accuracySnapshotStorage: AccuracySnapshotStorage;
         let testMCPClient: AccuracyTestingClient;
-        let agent: Agent;
+        let agent: VercelAgent;
 
         beforeAll(async () => {
+            accuracySnapshotStorage = await getAccuracySnapshotStorage();
             testMCPClient = await AccuracyTestingClient.initializeClient(mdbIntegration.connectionString());
             agent = getVercelToolCallingAgent();
         });
@@ -50,6 +54,7 @@ export function describeAccuracyTests(
         });
 
         afterAll(async () => {
+            await accuracySnapshotStorage.close();
             await testMCPClient.close();
         });
 
@@ -62,21 +67,27 @@ export function describeAccuracyTests(
                 const promptForModel = testConfig.injectConnectedAssumption
                     ? [testConfig.prompt, "(Assume that you are already connected to a MongoDB cluster!)"].join(" ")
                     : testConfig.prompt;
-                const conversation = await agent.prompt(promptForModel, model, toolsForModel);
+
+                const timeBeforePrompt = Date.now();
+                const result = await agent.prompt(promptForModel, model, toolsForModel);
+                const timeAfterPrompt = Date.now();
                 const toolCalls = testMCPClient.getToolCalls();
                 const toolCallingAccuracy = toolCallingAccuracyScorer(testConfig.expectedToolCalls, toolCalls);
                 const parameterMatchingAccuracy = parameterMatchingAccuracyScorer(
                     testConfig.expectedToolCalls,
                     toolCalls
                 );
-                console.debug(testConfig.prompt);
-                // console.debug(`Conversation`, JSON.stringify(conversation, null, 2));
-                // console.debug(`Tool calls`, JSON.stringify(toolCalls, null, 2));
-                console.debug(
-                    "Tool calling accuracy: %s, Parameter Accuracy: %s",
+
+                const responseTime = timeAfterPrompt - timeBeforePrompt;
+                await accuracySnapshotStorage.createSnapshotEntry({
+                    requestedModel: model.modelName,
+                    test: suiteName,
+                    prompt: testConfig.prompt,
+                    llmResponseTime: responseTime,
                     toolCallingAccuracy,
-                    parameterMatchingAccuracy
-                );
+                    parameterAccuracy: parameterMatchingAccuracy,
+                    ...result,
+                });
             });
         });
     });

tests/accuracy/sdk/describe-accuracy-tests.ts

@@ -0,0 +1,7 @@
+import { simpleGit } from "simple-git";
+
+export async function getCommitSHA(): Promise<string | undefined> {
+    const commitLogs = await simpleGit().log();
+    const lastCommit = commitLogs.latest;
+    return lastCommit?.hash;
+}