chore: more consolidated prompt tests

Author: himanshusinghs

tests/accuracy/list-databases.test.ts

@@ -1,26 +1,38 @@
 import { describeAccuracyTests } from "./sdk/describe-accuracy-tests.js";
 import { getAvailableModels } from "./sdk/models.js";
+import { AccuracyTestConfig } from "./sdk/describe-accuracy-tests.js";
 
-describeAccuracyTests("list-databases", getAvailableModels(), [
-    {
-        prompt: "Assume that you're already connected. How many collections are there in sample_mflix database",
+function describeListDatabasesAccuracyTests(prompt: string): AccuracyTestConfig {
+    return {
+        systemPrompt: "Assume that you're already connected.",
+        prompt: prompt,
         mockedTools: {
-            "list-collections": function listCollections() {
+            "list-databases": function listDatabases() {
                 return {
                     content: [
                         {
                             type: "text",
-                            text: "Name: coll1",
+                            text: "Name: db1",
+                        },
+                        {
+                            type: "text",
+                            text: "Name: db2",
                         },
                     ],
                 };
             },
         },
         expectedToolCalls: [
             {
-                toolName: "list-collections",
-                parameters: { database: "sample_mflix" },
+                toolName: "list-databases",
+                parameters: {},
             },
         ],
-    },
+    };
+}
+
+describeAccuracyTests("list-databases", getAvailableModels(), [
+    describeListDatabasesAccuracyTests("How many databases do I have?"),
+    describeListDatabasesAccuracyTests("List all the databases in my cluster."),
+    describeListDatabasesAccuracyTests("Is there a sample_mflix database in my cluster?"),
 ]);

tests/accuracy/sdk/agent.ts

@@ -8,23 +8,20 @@ const systemPrompt = [
     "When calling a tool, you MUST strictly follow its input schema and MUST provide all required arguments",
     "If a task requires multiple steps, you MUST call the necessary tools in sequence",
     'If you do not know the answer or the request cannot be fulfilled, you MUST reply with "I don\'t know"',
-    "You SHOULD assume that you are already connected to a MongoDB connection",
-].join("\n");
+];
 
 export interface Agent<M = unknown, T = unknown, R = unknown> {
     prompt(prompt: string, model: M, tools: T): Promise<R>;
 }
 
-export function getVercelToolCallingAgent(): Agent<
-    Model<LanguageModelV1>,
-    Record<string, Tool<Schema<unknown>>>,
-    { text: string; messages: unknown[] }
-> {
+export function getVercelToolCallingAgent(
+    requestedSystemPrompt?: string
+): Agent<Model<LanguageModelV1>, Record<string, Tool<Schema<unknown>>>, { text: string; messages: unknown[] }> {
     return {
         async prompt(prompt: string, model: Model<LanguageModelV1>, tools: Record<string, Tool<Schema<unknown>>>) {
             const result = await generateText({
                 model: model.getModel(),
-                system: systemPrompt,
+                system: [...systemPrompt, requestedSystemPrompt].join("\n"),
                 prompt,
                 tools,
                 maxSteps: 100,

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

@@ -5,7 +5,8 @@ import { ExpectedToolCall, parameterMatchingAccuracyScorer, toolCallingAccuracyS
 import { Agent, getVercelToolCallingAgent } from "./agent.js";
 import { appendAccuracySnapshot } from "./accuracy-snapshot.js";
 
-interface AccuracyTestConfig {
+export interface AccuracyTestConfig {
+    systemPrompt?: string;
     prompt: string;
     expectedToolCalls: ExpectedToolCall[];
     mockedTools: MockedTools;
@@ -17,13 +18,7 @@ export function describeAccuracyTests(
     accuracyTestConfigs: AccuracyTestConfig[]
 ) {
     const accuracyDatetime = process.env.MDB_ACCURACY_DATETIME;
-    if (!accuracyDatetime) {
-        throw new Error("MDB_ACCURACY_DATETIME environment variable is not set");
-    }
     const accuracyCommit = process.env.MDB_ACCURACY_COMMIT;
-    if (!accuracyCommit) {
-        throw new Error("MDB_ACCURACY_COMMIT environment variable is not set");
-    }
 
     if (!models.length) {
         console.warn(`No models available to test ${suiteName}`);
@@ -53,25 +48,31 @@ export function describeAccuracyTests(
             const toolCalls = testTools.getToolCalls();
             const toolCallingAccuracy = toolCallingAccuracyScorer(testConfig.expectedToolCalls, toolCalls);
             const parameterMatchingAccuracy = parameterMatchingAccuracyScorer(testConfig.expectedToolCalls, toolCalls);
-            await appendAccuracySnapshot({
-                datetime: accuracyDatetime,
-                commit: accuracyCommit,
-                model: model.modelName,
-                suite: suiteName,
-                test: testConfig.prompt,
-                toolCallingAccuracy,
-                parameterAccuracy: parameterMatchingAccuracy,
-            });
+            if (accuracyDatetime && accuracyCommit) {
+                await appendAccuracySnapshot({
+                    datetime: accuracyDatetime,
+                    commit: accuracyCommit,
+                    model: model.modelName,
+                    suite: suiteName,
+                    test: testConfig.prompt,
+                    toolCallingAccuracy,
+                    parameterAccuracy: parameterMatchingAccuracy,
+                });
+            } else {
+                console.info(
+                    `Skipping accuracy snapshot update for ${model.modelName} - ${suiteName} - ${testConfig.prompt}`
+                );
+            }
 
             try {
                 expect(toolCallingAccuracy).not.toEqual(0);
                 expect(parameterMatchingAccuracy).toBeGreaterThanOrEqual(0.5);
             } catch (error) {
                 console.warn(`Accuracy test failed for ${model.modelName} - ${suiteName} - ${testConfig.prompt}`);
-                console.warn(`Conversation`, JSON.stringify(conversation, null, 2));
-                console.warn(`Tool calls`, JSON.stringify(toolCalls, null, 2));
-                console.warn(`Tool calling accuracy`, toolCallingAccuracy);
-                console.warn(`Parameter matching accuracy`, parameterMatchingAccuracy);
+                console.debug(`Conversation`, JSON.stringify(conversation, null, 2));
+                console.debug(`Tool calls`, JSON.stringify(toolCalls, null, 2));
+                console.debug(`Tool calling accuracy`, toolCallingAccuracy);
+                console.debug(`Parameter matching accuracy`, parameterMatchingAccuracy);
                 throw error;
             }
         });

tests/accuracy/sdk/models.ts

@@ -36,8 +36,8 @@ export class OllamaModel implements Model {
 
 const ALL_TESTABLE_MODELS = [
     new GeminiModel("gemini-1.5-flash"),
-    new GeminiModel("gemini-2.0-flash"),
-    new OllamaModel("qwen3:latest"),
+    // new GeminiModel("gemini-2.0-flash"),
+    // new OllamaModel("qwen3:latest"),
 ];
 
 export type TestableModels = ReturnType<typeof getAvailableModels>;