chore: add AzureOpenAI model in the model list

Author: himanshusinghs

package-lock.json

@@ -34,6 +34,7 @@
   },
   "license": "Apache-2.0",
   "devDependencies": {
+    "@ai-sdk/azure": "^1.3.23",
     "@eslint/js": "^9.24.0",
     "@himanshusinghs/google": "^1.2.11",
     "@jest/globals": "^30.0.0",

package.json

@@ -4,7 +4,7 @@ import { AccuracyTestConfig } from "./sdk/describe-accuracy-tests.js";
 
 function describeListCollectionsAccuracyTests(prompt: string): AccuracyTestConfig {
     return {
-        systemPrompt: "Assume that you're already connected.",
+        injectConnectedAssumption: true,
         prompt: prompt,
         mockedTools: {
             "list-collections": function listCollections() {

tests/accuracy/list-collections.test.ts

@@ -4,7 +4,7 @@ import { AccuracyTestConfig } from "./sdk/describe-accuracy-tests.js";
 
 function describeListDatabasesAccuracyTests(prompt: string): AccuracyTestConfig {
     return {
-        systemPrompt: "Assume that you're already connected.",
+        injectConnectedAssumption: true,
         prompt: prompt,
         mockedTools: {
             "list-databases": function listDatabases() {

tests/accuracy/list-databases.test.ts

@@ -7,6 +7,7 @@ import { appendAccuracySnapshot } from "./accuracy-snapshot.js";
 
 export interface AccuracyTestConfig {
     systemPrompt?: string;
+    injectConnectedAssumption?: boolean;
     prompt: string;
     expectedToolCalls: ExpectedToolCall[];
     mockedTools: MockedTools;
@@ -44,7 +45,10 @@ export function describeAccuracyTests(
 
         eachTest("$prompt", async function (testConfig) {
             testTools.mockTools(testConfig.mockedTools);
-            const conversation = await agent.prompt(testConfig.prompt, model, testTools.vercelAiTools());
+            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, testTools.vercelAiTools());
             const toolCalls = testTools.getToolCalls();
             const toolCallingAccuracy = toolCallingAccuracyScorer(testConfig.expectedToolCalls, toolCalls);
             const parameterMatchingAccuracy = parameterMatchingAccuracyScorer(testConfig.expectedToolCalls, toolCalls);

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

@@ -1,5 +1,6 @@
 import { LanguageModelV1 } from "ai";
 import { createGoogleGenerativeAI } from "@himanshusinghs/google";
+import { createAzure } from "@ai-sdk/azure";
 import { ollama } from "ollama-ai-provider";
 
 export interface Model<P extends LanguageModelV1 = LanguageModelV1> {
@@ -8,6 +9,22 @@ export interface Model<P extends LanguageModelV1 = LanguageModelV1> {
     getModel(): P;
 }
 
+export class OpenAIModel implements Model {
+    constructor(readonly modelName: string) {}
+
+    isAvailable(): boolean {
+        return !!process.env.MDB_AZURE_OPEN_AI_API_KEY && !!process.env.MDB_AZURE_OPEN_AI_API_URL;
+    }
+
+    getModel() {
+        return createAzure({
+            baseURL: process.env.MDB_AZURE_OPEN_AI_API_URL,
+            apiKey: process.env.MDB_AZURE_OPEN_AI_API_KEY,
+            apiVersion: "2024-12-01-preview",
+        })(this.modelName);
+    }
+}
+
 export class GeminiModel implements Model {
     constructor(readonly modelName: string) {}
 
@@ -35,9 +52,9 @@ export class OllamaModel implements Model {
 }
 
 const ALL_TESTABLE_MODELS = [
-    new GeminiModel("gemini-1.5-flash"),
     new GeminiModel("gemini-2.0-flash"),
-    new OllamaModel("qwen3:1.7b"),
+    new OpenAIModel("gpt-4o"),
+    // new OllamaModel("qwen3:1.7b"),
 ];
 
 export type TestableModels = ReturnType<typeof getAvailableModels>;