Vertex in braintrust (#769)

* vertex models thru braintrust

* clean up unused import

---------

Co-authored-by: Ben Perlmutter <mongodben@mongodb.com>

Author: mongodben

packages/benchmarks/.env.example

@@ -11,6 +11,4 @@ BRAINTRUST_TEXT_TO_DRIVER_PROJECT_NAME="<some project>"
 OPENAI_API_VERSION="2024-06-01"
 OPENAI_ENDPOINT="https://<resource_name>.openai.azure.com/"
 OPENAI_API_KEY="<api_key>"
-GCP_OPENAI_ENDPOINT="https://{LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/openapi"
-GCP_API_KEY="SOME_API_KEY"
 MONGODB_TEXT_TO_DRIVER_CONNECTION_URI="<CONNECTION URI>"

packages/benchmarks/src/benchmarkModels.ts

@@ -19,12 +19,9 @@ export const MODELS = (
     "llama-3.1-70b",
     "llama-3.2-90b",
     "llama-3.3-70b",
-    "gemini-1.5-flash-002",
     "gemini-2-flash",
-    "gemini-2.0-flash-lite",
-    "gemini-1.0-pro-002",
-    "gemini-1.5-pro-002",
-    "gemini-2.5-flash",
+    "gemini-2.0-flash-lite-001",
+    "gemini-2.5-flash-preview-04-17",
     "nova-micro-v1:0",
     "nova-lite-v1:0",
     "nova-pro-v1:0",

packages/benchmarks/src/envVars.ts

@@ -1,12 +1,9 @@
-import { GCP_VERTEX_AI_ENV_VARS } from "mongodb-rag-core";
-
 export const envVars = {
   MONGODB_DATABASE_NAME: "",
   MONGODB_CONNECTION_URI: "",
   MONGODB_QUIZ_QUESTIONS_DATABASE_NAME: "",
   MONGODB_QUIZ_QUESTIONS_CONNECTION_URI: "",
   MONGODB_QUIZ_QUESTIONS_COLLECTION_NAME: "",
-  ...GCP_VERTEX_AI_ENV_VARS,
 };
 
 export const TEXT_TO_DRIVER_ENV_VARS = {

packages/benchmarks/src/nlPromptResponse/bin/globalConfig.ts

@@ -31,7 +31,7 @@ export const models = getModelsFromLabels([
   "o3",
   "o4-mini",
   "gemini-2-flash",
-  "gemini-2.0-flash-lite",
-  "gemini-2.5-flash",
+  "gemini-2.0-flash-lite-001",
+  "gemini-2.5-flash-preview-04-17",
   "gemini-2.5-pro-preview-03-25",
 ]);

packages/benchmarks/src/openAiClients.ts

@@ -1,8 +1,4 @@
-import {
-  assertEnvVars,
-  BRAINTRUST_ENV_VARS,
-  GCP_VERTEX_AI_ENV_VARS,
-} from "mongodb-rag-core";
+import { assertEnvVars, BRAINTRUST_ENV_VARS } from "mongodb-rag-core";
 import { makeOpenAiClientFactory } from "mongodb-rag-core/models";
 import { ModelConfig, models } from "mongodb-rag-core/models";
 import { strict as assert } from "assert";
@@ -31,23 +27,13 @@ export function makeLlmOptions(
   };
 }
 
-const {
-  BRAINTRUST_API_KEY,
-  BRAINTRUST_ENDPOINT,
-  GCP_API_KEY,
-  GCP_OPENAI_ENDPOINT,
-} = assertEnvVars({
+const { BRAINTRUST_API_KEY, BRAINTRUST_ENDPOINT } = assertEnvVars({
   ...BRAINTRUST_ENV_VARS,
-  ...GCP_VERTEX_AI_ENV_VARS,
 });
 
 export const openAiClientFactory = makeOpenAiClientFactory({
   braintrust: {
     apiKey: BRAINTRUST_API_KEY,
     endpoint: BRAINTRUST_ENDPOINT,
   },
-  vertexAi: {
-    apiKey: GCP_API_KEY,
-    endpoint: GCP_OPENAI_ENDPOINT,
-  },
 });

packages/mongodb-rag-core/.env.example

@@ -13,5 +13,3 @@ VECTOR_SEARCH_INDEX_NAME=<index_name>
 OPENAI_API_VERSION="2024-06-01"
 BRAINTRUST_ENDPOINT="https://api.braintrust.dev/v1/proxy"
 BRAINTRUST_API_KEY="<some api key>"
-GCP_OPENAI_ENDPOINT="https://{LOCATION}-aiplatform.googleapis.com/v1beta1/projects/{PROJECT_ID}/locations/{LOCATION}/endpoints/openapi"
-GCP_API_KEY="SOME_API_KEY"

packages/mongodb-rag-core/src/CoreEnvVars.ts

@@ -41,8 +41,3 @@ export const BRAINTRUST_ENV_VARS = {
   BRAINTRUST_API_KEY: "",
   BRAINTRUST_ENDPOINT: "",
 };
-
-export const GCP_VERTEX_AI_ENV_VARS = {
-  GCP_OPENAI_ENDPOINT: "",
-  GCP_API_KEY: "",
-};

packages/mongodb-rag-core/src/models/getOpenAiEndpointAndApiKey.test.ts

@@ -1,8 +1,6 @@
 import { ModelConfig } from "./models";
-import { BRAINTRUST_ENV_VARS, GCP_VERTEX_AI_ENV_VARS } from "../CoreEnvVars";
+import { BRAINTRUST_ENV_VARS } from "../CoreEnvVars";
 import { assertEnvVars } from "../assertEnvVars";
-import { GoogleAuth } from "google-auth-library";
-import { strict as assert } from "assert";
 
 export async function getOpenAiEndpointAndApiKey(model: ModelConfig) {
   if (model.provider === "braintrust") {
@@ -13,28 +11,5 @@ export async function getOpenAiEndpointAndApiKey(model: ModelConfig) {
       apiKey: BRAINTRUST_API_KEY,
     };
   }
-  if (model.provider === "gcp_vertex_ai") {
-    const { GCP_OPENAI_ENDPOINT } = assertEnvVars(GCP_VERTEX_AI_ENV_VARS);
-    const accessToken = await getGcpAccessToken();
-    assert(typeof accessToken === "string", "GCP access token not found");
-    return {
-      baseURL: GCP_OPENAI_ENDPOINT,
-      apiKey: accessToken,
-    };
-  }
   throw new Error(`Unknown provider: ${model.provider}`);
 }
-
-/**
-  Get GCP access token. To set up locally, run `gcloud auth application-default login`.
- */
-export async function getGcpAccessToken() {
-  const auth = new GoogleAuth({
-    scopes: ["https://www.googleapis.com/auth/cloud-platform"],
-  });
-
-  const client = await auth.getClient();
-  const accessTokenResponse = await client.getAccessToken();
-  const accessToken = accessTokenResponse.token;
-  return accessToken;
-}

packages/mongodb-rag-core/src/models/getOpenAiEndpointAndApiKey.ts

@@ -14,7 +14,6 @@ interface MakeOpenAiClientFactoryParams {
 
 export function makeOpenAiClientFactory({
   braintrust,
-  vertexAi,
 }: MakeOpenAiClientFactoryParams) {
   return {
     makeOpenAiClient(modelConfig: ModelConfig) {
@@ -27,14 +26,6 @@ export function makeOpenAiClientFactory({
             baseURL: braintrust.endpoint,
           })
         );
-      } else if (modelConfig.provider === "gcp_vertex_ai") {
-        assert(vertexAi, "GCP Vertex AI config must be provided");
-        openAiClient = wrapOpenAI(
-          new OpenAI({
-            apiKey: vertexAi.apiKey,
-            baseURL: vertexAi.endpoint,
-          })
-        );
       } else {
         throw new Error(`Unsupported provider: ${modelConfig.provider}`);
       }

packages/mongodb-rag-core/src/models/makeOpenAiClientFactory.ts

@@ -8,7 +8,6 @@ import { models } from "./models";
 import { makeOpenAiClientFactory } from "./makeOpenAiClientFactory";
 import { OpenAI } from "openai";
 import "dotenv/config";
-import { getOpenAiEndpointAndApiKey } from "./getOpenAiEndpointAndApiKey";
 
 jest.setTimeout(60000);
 
@@ -33,23 +32,6 @@ describe.skip("Braintrust models", () => {
   );
 });
 
-describe.skip("GCP Vertex AI models", () => {
-  test.each(models.filter((m) => m.provider === "gcp_vertex_ai"))(
-    "'$label' model should generate data",
-    async (model) => {
-      const { apiKey, baseURL } = await getOpenAiEndpointAndApiKey(model);
-      const openAiClientFactory = makeOpenAiClientFactory({
-        vertexAi: {
-          apiKey,
-          endpoint: baseURL,
-        },
-      });
-      const openAiClient = openAiClientFactory.makeOpenAiClient(model);
-      await expectModelResponse(openAiClient, model.deployment);
-    }
-  );
-});
-
 async function expectModelResponse(client: OpenAI, model: string) {
   const res = await client.chat.completions.create({
     model,

packages/mongodb-rag-core/src/models/models.test.ts

@@ -8,7 +8,7 @@ export type ModelDeveloper =
   | "DeepSeek"
   | "Alibaba Cloud";
 
-export type ModelProvider = "braintrust" | "gcp_vertex_ai";
+export type ModelProvider = "braintrust";
 
 export interface ModelConfig {
   /**
@@ -321,71 +321,63 @@ const allModels = [
     generation: "nova-1",
   },
   {
-    label: "gemini-1.5-flash-002",
-    deployment: "google/gemini-1.5-flash-002",
+    label: "gemini-2-flash",
+    deployment: "publishers/google/models/gemini-2.0-flash-001",
     developer: "Google",
-    maxConcurrency: 3,
-    provider: "gcp_vertex_ai",
+    maxConcurrency: 10,
+    provider: "braintrust",
     authorized: true,
-    generation: "gemini-1",
+    parent: "gemini-1.5-flash-002",
+    generation: "gemini-2",
   },
   {
-    label: "gemini-2-flash",
-    deployment: "models/gemini-2.0-flash-001",
+    label: "gemini-2.5-flash-preview-04-17",
+    deployment: "publishers/google/models/gemini-2.5-flash-preview-04-17",
     developer: "Google",
     maxConcurrency: 10,
-    provider: "gcp_vertex_ai",
+    provider: "braintrust",
     authorized: true,
-    parent: "gemini-1.5-flash-002",
+    parent: "gemini-2-flash",
     generation: "gemini-2",
   },
   {
-    label: "gemini-2.5-flash",
-    deployment: "models/gemini-2.5-flash-preview-04-17",
+    label: "gemini-2.5-flash-preview-05-20",
+    deployment: "publishers/google/models/gemini-2.5-flash-preview-05-20",
     developer: "Google",
     maxConcurrency: 10,
-    provider: "gcp_vertex_ai",
+    provider: "braintrust",
     authorized: true,
     parent: "gemini-2-flash",
     generation: "gemini-2",
   },
   {
-    label: "gemini-2.0-flash-lite",
-    deployment: "models/gemini-2.0-flash-lite-001",
+    label: "gemini-2.0-flash-lite-001",
+    deployment: "publishers/google/models/gemini-2.0-flash-lite-001",
     developer: "Google",
     maxConcurrency: 10,
-    provider: "gcp_vertex_ai",
+    provider: "braintrust",
     authorized: true,
     generation: "gemini-2",
   },
   {
     label: "gemini-2.5-pro-preview-03-25",
-    deployment: "models/gemini-2.5-pro-preview-03-25",
+    deployment: "publishers/google/models/gemini-2.5-pro-preview-03-25",
     developer: "Google",
     maxConcurrency: 5,
-    provider: "gcp_vertex_ai",
+    provider: "braintrust",
     authorized: true,
     parent: "gemini-1.5-pro-002",
     generation: "gemini-2",
   },
   {
-    label: "gemini-1.0-pro-002",
-    deployment: "google/gemini-1.0-pro-002",
-    developer: "Google",
-    provider: "gcp_vertex_ai",
-    maxConcurrency: 1,
-    authorized: true,
-    generation: "gemini-1",
-  },
-  {
-    label: "gemini-1.5-pro-002",
-    deployment: "google/gemini-1.5-pro-002",
+    label: "gemini-2.5-pro-preview-05-06",
+    deployment: "publishers/google/models/gemini-2.5-pro-preview-05-06",
     developer: "Google",
-    provider: "gcp_vertex_ai",
-    maxConcurrency: 1,
+    maxConcurrency: 5,
+    provider: "braintrust",
     authorized: true,
-    parent: "gemini-1.0-pro-002",
-    generation: "gemini-1",
+    parent: "gemini-1.5-pro-002",
+    generation: "gemini-2",
   },
   {
     label: "deepseek-r1",